fix: compatibility with ibis 1.4.0 (and possibly 1.2, 1.3) (#31)

* fix: compatibility with ibis 1.4.0

* add module-level connect method

* workaround for circular import

* allow test_scalar_param tests to work on ibis 1.4

* mskip tests on 1.4

* fix compatibility for 1.2

Author: tswast

ibis_bigquery/__init__.py

@@ -5,14 +5,18 @@
 import google.auth.credentials
 import google.cloud.bigquery  # noqa: F401, fail early if bigquery is missing
 import pydata_google_auth
-from ibis.backends.base import BaseBackend
 from pydata_google_auth import cache
 
 from . import version as ibis_bigquery_version
 from .client import (BigQueryClient, BigQueryDatabase, BigQueryQuery,
                      BigQueryTable)
 from .compiler import BigQueryExprTranslator, BigQueryQueryBuilder
 
+try:
+    from ibis.backends.base import BaseBackend
+except ImportError:
+    from .backcompat import BaseBackend
+
 try:
     from .udf import udf  # noqa F401
 except ImportError:
@@ -42,6 +46,13 @@ class Backend(BaseBackend):
     database_class = BigQueryDatabase
     table_class = BigQueryTable
 
+    # These were moved from TestConf for use in common test suite.
+    # TODO: Indicate RoundAwayFromZero and UnorderedComparator.
+    # https://github.com/ibis-project/ibis-bigquery/issues/30
+    supports_divide_by_zero = True
+    supports_floating_modulus = False
+    returned_timestamp_unit = 'us'
+
     def connect(
         self,
         project_id: Optional[str] = None,

ibis_bigquery/backcompat.py

@@ -0,0 +1,134 @@
+"""Helpers to make this backend compatible with Ibis versions < 2.0.
+
+Keep in sync with:
+https://github.com/ibis-project/ibis/blob/master/ibis/backends/base/__init__.py
+
+TODO: Remove this after Ibis 2.0 release and support for earlier versions of
+      Ibis < 2.0 is dropped.
+"""
+
+import abc
+
+try:
+    from ibis.common.exceptions import TranslationError
+except ImportError:
+    # 1.2
+    from ibis.common import TranslationError
+
+__all__ = ('BaseBackend',)
+
+
+class BaseBackend(abc.ABC):
+    """
+    Base backend class.
+    All Ibis backends are expected to subclass this `Backend` class,
+    and implement all the required methods.
+    """
+
+    @property
+    @abc.abstractmethod
+    def name(self) -> str:
+        """
+        Name of the backend, for example 'sqlite'.
+        """
+        pass
+
+    @property
+    @abc.abstractmethod
+    def kind(self):
+        """
+        Backend kind. One of:
+        sqlalchemy
+            Backends using a SQLAlchemy dialect.
+        sql
+            SQL based backends, not based on a SQLAlchemy dialect.
+        pandas
+            Backends using pandas to store data and perform computations.
+        spark
+            Spark based backends.
+        """
+        pass
+
+    @property
+    @abc.abstractmethod
+    def builder(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def translator(self):
+        pass
+
+    @property
+    def dialect(self):
+        """
+        Dialect class of the backend.
+        We generate it dynamically to avoid repeating the code for each
+        backend.
+        """
+        # TODO importing dialects inside the function to avoid circular
+        # imports. In the future instead of this if statement we probably
+        # want to create subclasses for each of the kinds
+        # (e.g. `BaseSQLAlchemyBackend`)
+        # TODO check if the below dialects can be merged into a single one
+        if self.kind == 'sqlalchemy':
+            from ibis.backends.base_sqlalchemy.alchemy import AlchemyDialect
+
+            dialect_class = AlchemyDialect
+        elif self.kind in ('sql', 'pandas'):
+            try:
+                from ibis.backends.base_sqlalchemy.compiler import Dialect
+            except ImportError:
+                from ibis.sql.compiler import Dialect
+
+            dialect_class = Dialect
+        elif self.kind == 'spark':
+            from ibis.backends.base_sql.compiler import BaseDialect
+
+            dialect_class = BaseDialect
+        else:
+            raise ValueError(
+                f'Backend class "{self.kind}" unknown. '
+                'Expected one of "sqlalchemy", "sql", '
+                '"pandas" or "spark".'
+            )
+
+        dialect_class.translator = self.translator
+        return dialect_class
+
+    @abc.abstractmethod
+    def connect(connection_string, **options):
+        """
+        Connect to the underlying database and return a client object.
+        """
+        pass
+
+    def register_options(self):
+        """
+        If the backend has custom options, register them here.
+        They will be prefixed with the name of the backend.
+        """
+        pass
+
+    def compile(self, expr, params=None):
+        """
+        Compile the expression.
+        """
+        context = self.dialect.make_context(params=params)
+        builder = self.builder(expr, context=context)
+        query_ast = builder.get_result()
+        # TODO make all builders return a QueryAST object
+        if isinstance(query_ast, list):
+            query_ast = query_ast[0]
+        compiled = query_ast.compile()
+        return compiled
+
+    def verify(self, expr, params=None):
+        """
+        Verify `expr` is an expression that can be compiled.
+        """
+        try:
+            self.compile(expr, params=params)
+            return True
+        except TranslationError:
+            return False

ibis_bigquery/client.py

@@ -6,7 +6,12 @@
 
 import google.cloud.bigquery as bq
 import ibis
-import ibis.common.exceptions as com
+
+try:
+    import ibis.common.exceptions as com
+except ImportError:
+    import ibis.common as com
+
 import ibis.expr.datatypes as dt
 import ibis.expr.lineage as lin
 import ibis.expr.operations as ops

ibis_bigquery/compiler.py

@@ -5,19 +5,54 @@
 from functools import partial
 
 import ibis
-import ibis.backends.base_sqlalchemy.compiler as comp
-import ibis.common.exceptions as com
+
+try:
+    import ibis.backends.base_sqlalchemy.compiler as comp
+except ImportError:
+    import ibis.sql.compiler as comp
+try:
+    import ibis.common.exceptions as com
+except ImportError:
+    import ibis.common as com
+
 import ibis.expr.datatypes as dt
 import ibis.expr.lineage as lin
 import ibis.expr.operations as ops
 import ibis.expr.types as ir
 import numpy as np
 import regex as re
 import toolz
-from ibis.backends.base.sql import (fixed_arity, literal, operation_registry,
-                                    reduction, unary)
-from ibis.backends.base_sql.compiler import (BaseExprTranslator, BaseSelect,
-                                             BaseTableSetFormatter)
+
+try:
+    # 2.x
+    from ibis.backends.base.sql import (fixed_arity, literal,
+                                        operation_registry, reduction, unary)
+except ImportError:
+    try:
+        # 1.4
+        from ibis.backends.base_sql import (fixed_arity, literal,
+                                            operation_registry, reduction,
+                                            unary)
+    except ImportError:
+        # 1.2
+        from ibis.impala.compiler import _literal as literal
+        from ibis.impala.compiler import \
+            _operation_registry as operation_registry
+        from ibis.impala.compiler import _reduction as reduction
+        from ibis.impala.compiler import fixed_arity, unary
+
+try:
+    from ibis.backends.base_sql.compiler import (BaseExprTranslator,
+                                                 BaseSelect,
+                                                 BaseTableSetFormatter)
+except ImportError:
+    # 1.2
+    from ibis.impala.compiler import ImpalaExprTranslator as BaseExprTranslator
+    from ibis.impala.compiler import ImpalaSelect as BaseSelect
+    from ibis.impala.compiler import (
+        ImpalaTableSetFormatter as BaseTableSetFormatter
+    )
+
 from multipledispatch import Dispatcher
 
 from .datatypes import ibis_type_to_bigquery_type
@@ -369,18 +404,13 @@ def _formatter(translator, expr):
 }
 _operation_registry.update(
     {
-        ops.BitAnd: reduction('BIT_AND'),
-        ops.BitOr: reduction('BIT_OR'),
-        ops.BitXor: reduction('BIT_XOR'),
         ops.ExtractYear: _extract_field('year'),
-        ops.ExtractQuarter: _extract_field('quarter'),
         ops.ExtractMonth: _extract_field('month'),
         ops.ExtractDay: _extract_field('day'),
         ops.ExtractHour: _extract_field('hour'),
         ops.ExtractMinute: _extract_field('minute'),
         ops.ExtractSecond: _extract_field('second'),
         ops.ExtractMillisecond: _extract_field('millisecond'),
-        ops.ExtractEpochSeconds: _extract_field('epochseconds'),
         ops.Hash: _hash,
         ops.StringReplace: fixed_arity('REPLACE', 3),
         ops.StringSplit: fixed_arity('SPLIT', 2),
@@ -427,6 +457,25 @@ def _formatter(translator, expr):
     }
 )
 
+
+def _try_register_op(op_name: str, value):
+    """Register operation if it exists in Ibis.
+
+    This allows us to decouple slightly from ibis-framework releases.
+    """
+    if hasattr(ops, op_name):
+        _operation_registry[getattr(ops, op_name)] = value
+
+
+# 2.x
+_try_register_op('BitAnd', reduction('BIT_AND'))
+_try_register_op('BitOr', reduction('BIT_OR'))
+_try_register_op('BitXor', reduction('BIT_XOR'))
+# 1.4
+_try_register_op('ExtractQuarter', _extract_field('quarter'))
+_try_register_op('ExtractEpochSeconds', _extract_field('epochseconds'))
+
+
 _invalid_operations = {
     ops.Translate,
     ops.FindInSet,

ibis_bigquery/udf/__init__.py

@@ -6,7 +6,14 @@
 
 import ibis.expr.datatypes as dt
 import ibis.expr.rules as rlz
-import ibis.udf.validate as v
+
+try:
+    from ibis.udf.validate import validate_output_type
+except ImportError:
+    # 1.2
+    def validate_output_type(*args):
+        pass
+
 from ibis.expr.signature import Argument as Arg
 
 from ..compiler import BigQueryUDFNode, compiles
@@ -163,7 +170,7 @@ class Rectangle {
     return my_rectangle(width, height);
     """;
     '''
-    v.validate_output_type(output_type)
+    validate_output_type(output_type)
 
     if libraries is None:
         libraries = []

tests/system/conftest.py

@@ -1,12 +1,8 @@
 import os
-from pathlib import Path
 
-import ibis
-import ibis.expr.types as ir
+import ibis  # noqa: F401
 import pytest
 from google.oauth2 import service_account
-from ibis.backends.tests.base import (BackendTest, RoundAwayFromZero,
-                                      UnorderedComparator)
 
 import ibis_bigquery
 
@@ -42,38 +38,6 @@ def _credentials():
     )
 
 
-class TestConf(UnorderedComparator, BackendTest, RoundAwayFromZero):
-    supports_divide_by_zero = True
-    supports_floating_modulus = False
-    returned_timestamp_unit = 'us'
-
-    @staticmethod
-    def connect(data_directory: Path) -> ibis.client.Client:
-        project_id = os.environ.get('GOOGLE_BIGQUERY_PROJECT_ID')
-        if project_id is None:
-            pytest.skip(
-                'Environment variable GOOGLE_BIGQUERY_PROJECT_ID '
-                'not defined'
-            )
-        elif not project_id:
-            pytest.skip(
-                'Environment variable GOOGLE_BIGQUERY_PROJECT_ID is empty'
-            )
-        return bq.connect(
-            project_id=project_id,
-            dataset_id=DATASET_ID,
-            credentials=_credentials(),
-        )
-
-    @property
-    def batting(self) -> ir.TableExpr:
-        return None
-
-    @property
-    def awards_players(self) -> ir.TableExpr:
-        return None
-
-
 @pytest.fixture(scope='session')
 def project_id():
     return PROJECT_ID