refactor: use BigQuery ibis backend from ibis-bigquery package (#244)

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:

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",

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):

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

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
 )

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

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})

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):

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

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,