From ae076e52d91ea43e2d41fbf7535222e15c96f4ed Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Thu, 1 Jul 2021 14:11:19 -0500
Subject: [PATCH] refactor: use BigQuery ibis backend from ibis-bigquery
 package (#244)

---
 data_validation/clients.py                      | 7 ++++---
 setup.py                                        | 1 +
 tests/system/ibis_addon/test_operations.py      | 4 ++--
 third_party/ibis/ibis_addon/datatypes.py        | 2 +-
 third_party/ibis/ibis_addon/operations.py       | 2 +-
 third_party/ibis/ibis_cloud_spanner/compiler.py | 5 +++--
 third_party/ibis/ibis_snowflake/udf/api.py      | 4 ++--
 third_party/ibis/ibis_teradata/api.py           | 7 +------
 third_party/ibis/ibis_teradata/client.py        | 3 ++-
 third_party/ibis/ibis_teradata/compiler.py      | 2 +-
 10 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/data_validation/clients.py b/data_validation/clients.py
index 4b606453a..734488647 100644
--- a/data_validation/clients.py
+++ b/data_validation/clients.py
@@ -20,7 +20,8 @@
 import google.oauth2.service_account
 
 from google.cloud import bigquery
-from ibis.backends.bigquery.client import BigQueryClient
+import ibis
+import ibis_bigquery
 import ibis.backends.pandas
 from ibis.backends.pandas.client import PandasClient
 from ibis.backends.mysql.client import MySQLClient
@@ -88,7 +89,7 @@ def get_bigquery_client(project_id, dataset_id=None, credentials=None):
     google_client = bigquery.Client(
         project=project_id, client_info=info, credentials=credentials
     )
-    ibis_client = BigQueryClient(
+    ibis_client = ibis_bigquery.connect(
         project_id, dataset_id=dataset_id, credentials=credentials
     )
 
@@ -182,7 +183,7 @@ def get_data_client(connection_config):
     connection_config = copy.deepcopy(connection_config)
     source_type = connection_config.pop(consts.SOURCE_TYPE)
 
-    # The BigQueryClient expects a credentials object, not a string.
+    # The ibis_bigquery.connect expects a credentials object, not a string.
     if consts.GOOGLE_SERVICE_ACCOUNT_KEY_PATH in connection_config:
         key_path = connection_config.pop(consts.GOOGLE_SERVICE_ACCOUNT_KEY_PATH)
         if key_path:
diff --git a/setup.py b/setup.py
index 5a40f43f5..e384f0751 100644
--- a/setup.py
+++ b/setup.py
@@ -33,6 +33,7 @@
     # Core dependencies
     "google-api-python-client==1.12.8",
     "ibis-framework==1.4.0",
+    "ibis-bigquery==0.1.1",
     "impyla==0.16.3",
     "SQLAlchemy==1.3.22",
     "PyMySQL==1.0.2",
diff --git a/tests/system/ibis_addon/test_operations.py b/tests/system/ibis_addon/test_operations.py
index 15a49486f..3312fe72f 100644
--- a/tests/system/ibis_addon/test_operations.py
+++ b/tests/system/ibis_addon/test_operations.py
@@ -15,7 +15,7 @@
 import textwrap
 
 import pytest
-import ibis.backends.bigquery
+import ibis_bigquery
 import ibis.expr.datatypes as dt
 
 # Import required in order to register operations.
@@ -24,7 +24,7 @@
 
 @pytest.fixture
 def bigquery_client():
-    return ibis.backends.bigquery.connect()
+    return ibis_bigquery.connect()
 
 
 def test_bit_xor_bigquery(bigquery_client):
diff --git a/third_party/ibis/ibis_addon/datatypes.py b/third_party/ibis/ibis_addon/datatypes.py
index 61804f614..a2bf23b08 100644
--- a/third_party/ibis/ibis_addon/datatypes.py
+++ b/third_party/ibis/ibis_addon/datatypes.py
@@ -14,7 +14,7 @@
 
 
 from google.cloud import bigquery
-from ibis.backends.bigquery.client import _DTYPE_TO_IBIS_TYPE
+from ibis_bigquery.client import _DTYPE_TO_IBIS_TYPE
 import ibis.expr.datatypes as dt
 from ibis.backends.pandas.client import _inferable_pandas_dtypes
 import pyarrow
diff --git a/third_party/ibis/ibis_addon/operations.py b/third_party/ibis/ibis_addon/operations.py
index 66c6bff33..3c192a3a0 100644
--- a/third_party/ibis/ibis_addon/operations.py
+++ b/third_party/ibis/ibis_addon/operations.py
@@ -27,7 +27,7 @@
 import sqlalchemy
 
 import ibis.expr.api
-from ibis.backends.bigquery.compiler import (
+from ibis_bigquery.compiler import (
     reduction as bq_reduction,
     BigQueryExprTranslator
 )
diff --git a/third_party/ibis/ibis_cloud_spanner/compiler.py b/third_party/ibis/ibis_cloud_spanner/compiler.py
index 2f0044a6d..be53eb69d 100644
--- a/third_party/ibis/ibis_cloud_spanner/compiler.py
+++ b/third_party/ibis/ibis_cloud_spanner/compiler.py
@@ -13,7 +13,8 @@
 # limitations under the License.
 
 import ibis.expr.operations as ops
-from ibis.backends.bigquery import compiler as bigquery_compiler
+import ibis_bigquery
+from ibis_bigquery import compiler as bigquery_compiler
 
 
 def build_ast(expr, context):
@@ -54,4 +55,4 @@ def _regex_extract(translator, expr):
 rewrites = bigquery_compiler.BigQueryExprTranslator.rewrites
 
 
-dialect = bigquery_compiler.BigQueryDialect
+dialect = ibis_bigquery.Backend().dialect
diff --git a/third_party/ibis/ibis_snowflake/udf/api.py b/third_party/ibis/ibis_snowflake/udf/api.py
index d30d2690a..a7480daf4 100644
--- a/third_party/ibis/ibis_snowflake/udf/api.py
+++ b/third_party/ibis/ibis_snowflake/udf/api.py
@@ -22,7 +22,7 @@
 import ibis.expr.datatypes as dt
 from ibis.expr.signature import Argument as Arg
 from ibis.snowflake.datatypes import UDFContext, ibis_type_to_snowflake_type
-from ibis.backends.bigquery.udf.core import PythonToJavaScriptTranslator
+from ibis_bigquery.udf.core import PythonToJavaScriptTranslator
 
 import ibis.expr.operations as ops
 
@@ -94,7 +94,7 @@ def wrapper(f):
         )
 
         udf_node = create_udf_node(f.__name__, udf_node_fields)
-        source = PythonToJavaScriptTranslator(f).compile() # Uses function translator implemented in ibis.backends.bigquery.udf
+        source = PythonToJavaScriptTranslator(f).compile() # Uses function translator implemented in ibis_bigquery.udf
 
         javascript_template = """\
 CREATE {to_replace} FUNCTION {external_name}({signature})
diff --git a/third_party/ibis/ibis_teradata/api.py b/third_party/ibis/ibis_teradata/api.py
index 145d5f4ff..d1aef90cf 100644
--- a/third_party/ibis/ibis_teradata/api.py
+++ b/third_party/ibis/ibis_teradata/api.py
@@ -10,13 +10,8 @@
 from compiler import dialect  # TODO make non local
 from ibis.config import options  # noqa: F401
 
-try:
-    from ibis.backends.bigquery.udf import udf  # noqa: F401 # TODO is this required?
-except ImportError:
-    pass
 
-
-__all__ = ("compile", "connect", "verify", "udf")
+__all__ = ("compile", "connect", "verify")
 
 
 def compile(expr, params=None):
diff --git a/third_party/ibis/ibis_teradata/client.py b/third_party/ibis/ibis_teradata/client.py
index 16c70eb45..c62459bbb 100644
--- a/third_party/ibis/ibis_teradata/client.py
+++ b/third_party/ibis/ibis_teradata/client.py
@@ -67,7 +67,8 @@ class TeradataClient(SQLClient):
     dialect = compiler.TeradataDialect
 
     def __init__(self, host, user_name, password, port=1025, use_no_lock_tables=False):
-        """Construct a BigQueryClient.
+        """Construct a TeradataClient.
+
         Parameters
         ----------
         host : str
diff --git a/third_party/ibis/ibis_teradata/compiler.py b/third_party/ibis/ibis_teradata/compiler.py
index 3acad54ec..7b7bb8e7a 100644
--- a/third_party/ibis/ibis_teradata/compiler.py
+++ b/third_party/ibis/ibis_teradata/compiler.py
@@ -16,7 +16,7 @@
 from ibis.backends.base_sql.identifiers import base_identifiers
 
 from ibis.common.exceptions import UnsupportedOperationError
-from ibis.backends.bigquery.datatypes import ibis_type_to_bigquery_type
+from ibis_bigquery.datatypes import ibis_type_to_bigquery_type
 from ibis.backends.base_sql import fixed_arity, literal, reduction, unary
 from ibis.backends.base_sql.compiler import (
     BaseExprTranslator,