feat: add ibis_bigquery.connect and ibis_bigquery.compile functions (#37)

* feat: add `ibis_bigquery.connect` and `ibis_bigquery.compile` functions

This should make it easier to use the `ibis_bigquery` module directly.
Also, fix the tests to use the methods from this module.

* fix lint errors

Author: tswast

README.rst

@@ -51,7 +51,7 @@ Using this library directly:
     import ibis
     import ibis_bigquery
 
-    conn = ibis_bigquery.Backend().connect(
+    conn = ibis_bigquery.connect(
         project_id=YOUR_PROJECT_ID,
         dataset_id='bigquery-public-data.stackoverflow'
     )

docs/bigquery.rst

@@ -1,4 +1,4 @@
-.. currentmodule:: ibis.bigquery.api
+.. currentmodule:: ibis_bigquery
 
 .. _backends.bigquery:
 
@@ -40,8 +40,8 @@ project.
    will still be billed for any and all queries**.
 
 If you want to query data that lives in a different project than the billing
-project you can use the :meth:`ibis.bigquery.client.BigQueryClient.database`
-method of :class:`ibis.bigquery.client.BigQueryClient` objects:
+project you can use the :meth:`ibis_bigquery.client.BigQueryClient.database`
+method of :class:`ibis_bigquery.client.BigQueryClient` objects:
 
 .. code-block:: python
 

ibis_bigquery/__init__.py

@@ -4,7 +4,6 @@
 
 import google.auth.credentials
 import google.cloud.bigquery  # noqa: F401, fail early if bigquery is missing
-import ibis.config
 import pydata_google_auth
 from ibis.backends.base import BaseBackend
 from pydata_google_auth import cache
@@ -52,6 +51,7 @@ def connect(
         auth_local_webserver: bool = False,
         auth_external_data: bool = False,
         auth_cache: str = "default",
+        partition_column: Optional[str] = "PARTITIONTIME",
     ) -> BigQueryClient:
         """Create a BigQueryClient for use with Ibis.
 
@@ -91,6 +91,9 @@ def connect(
                 Authenticates and does **not** cache credentials.
 
             Defaults to ``'default'``.
+        partition_column : str
+            Identifier to use instead of default ``_PARTITIONTIME`` partition
+            column. Defaults to ``'PARTITIONTIME'``.
 
         Returns
         -------
@@ -136,7 +139,96 @@ def connect(
             dataset_id=dataset_id,
             credentials=credentials,
             application_name=application_name,
+            partition_column=partition_column,
         )
 
-    def register_options(self):
-        ibis.config.register_option('partition_col', 'PARTITIONTIME')
+
+def compile(expr, params=None):
+    """Compile an expression for BigQuery.
+    Returns
+    -------
+    compiled : str
+    See Also
+    --------
+    ibis.expr.types.Expr.compile
+    """
+    backend = Backend()
+    return backend.compile(expr, params=params)
+
+
+def connect(
+    project_id: Optional[str] = None,
+    dataset_id: Optional[str] = None,
+    credentials: Optional[google.auth.credentials.Credentials] = None,
+    application_name: Optional[str] = None,
+    auth_local_webserver: bool = False,
+    auth_external_data: bool = False,
+    auth_cache: str = "default",
+    partition_column: Optional[str] = "PARTITIONTIME",
+) -> BigQueryClient:
+    """Create a BigQueryClient for use with Ibis.
+
+    Parameters
+    ----------
+    project_id : str
+        A BigQuery project id.
+    dataset_id : str
+        A dataset id that lives inside of the project indicated by
+        `project_id`.
+    credentials : google.auth.credentials.Credentials
+    application_name : str
+        A string identifying your application to Google API endpoints.
+    auth_local_webserver : bool
+        Use a local webserver for the user authentication.  Binds a
+        webserver to an open port on localhost between 8080 and 8089,
+        inclusive, to receive authentication token. If not set, defaults
+        to False, which requests a token via the console.
+    auth_external_data : bool
+        Authenticate using additional scopes required to `query external
+        data sources
+        <https://cloud.google.com/bigquery/external-data-sources>`_,
+        such as Google Sheets, files in Google Cloud Storage, or files in
+        Google Drive. If not set, defaults to False, which requests the
+        default BigQuery scopes.
+    auth_cache : str
+        Selects the behavior of the credentials cache.
+
+        ``'default'``
+            Reads credentials from disk if available, otherwise
+            authenticates and caches credentials to disk.
+
+        ``'reauth'``
+            Authenticates and caches credentials to disk.
+
+        ``'none'``
+            Authenticates and does **not** cache credentials.
+
+        Defaults to ``'default'``.
+    partition_column : str
+        Identifier to use instead of default ``_PARTITIONTIME`` partition
+        column. Defaults to ``'PARTITIONTIME'``.
+
+    Returns
+    -------
+    BigQueryClient
+
+    """
+    backend = Backend()
+    return backend.connect(
+        project_id=project_id,
+        dataset_id=dataset_id,
+        credentials=credentials,
+        application_name=application_name,
+        auth_local_webserver=auth_local_webserver,
+        auth_external_data=auth_external_data,
+        auth_cache=auth_cache,
+        partition_column=partition_column,
+    )
+
+
+__all__ = [
+    "__version__",
+    "Backend",
+    "compile",
+    "connect",
+]

ibis_bigquery/client.py

@@ -290,7 +290,8 @@ class BigQueryTable(ops.DatabaseTable):
     pass
 
 
-def rename_partitioned_column(table_expr, bq_table):
+def rename_partitioned_column(table_expr, bq_table, partition_col):
+    """Rename native partition column to user-defined name."""
     partition_info = bq_table._properties.get('timePartitioning', None)
 
     # If we don't have any partiton information, the table isn't partitioned
@@ -304,14 +305,11 @@ def rename_partitioned_column(table_expr, bq_table):
     # The partition field must be in table_expr columns
     assert partition_field in table_expr.columns
 
-    # User configured partition column name default
-    col = ibis.options.bigquery.partition_col
-
     # No renaming if the config option is set to None or the partition field
     # is not _PARTITIONTIME
-    if col is None or partition_field != NATIVE_PARTITION_COL:
+    if partition_col is None or partition_field != NATIVE_PARTITION_COL:
         return table_expr
-    return table_expr.relabel({NATIVE_PARTITION_COL: col})
+    return table_expr.relabel({NATIVE_PARTITION_COL: partition_col})
 
 
 def parse_project_and_dataset(
@@ -379,6 +377,7 @@ def __init__(
         dataset_id=None,
         credentials=None,
         application_name=None,
+        partition_column=None,
     ):
         """Construct a BigQueryClient.
 
@@ -391,6 +390,9 @@ def __init__(
         credentials : google.auth.credentials.Credentials
         application_name : str
             A string identifying your application to Google API endpoints.
+        partition_column : str
+            Identifier to use instead of default ``_PARTITIONTIME`` partition
+            column.
 
         """
         self.query_class = backend.query_class
@@ -407,6 +409,7 @@ def __init__(
             credentials=credentials,
             client_info=_create_client_info(application_name),
         )
+        self.partition_column = partition_column
 
     def _parse_project_and_dataset(self, dataset):
         if not dataset and not self.dataset:
@@ -431,7 +434,7 @@ def table(self, name, database=None):
         dataset_ref = self.client.dataset(dataset, project=project)
         table_ref = dataset_ref.table(name)
         bq_table = self.client.get_table(table_ref)
-        return rename_partitioned_column(t, bq_table)
+        return rename_partitioned_column(t, bq_table, self.partition_column)
 
     def _build_ast(self, expr, context):
         result = comp.build_ast(expr, context)

ibis_bigquery/udf/__init__.py

@@ -7,7 +7,6 @@
 import ibis.expr.datatypes as dt
 import ibis.expr.rules as rlz
 import ibis.udf.validate as v
-from ibis.compat import PY38  # noqa: F401
 from ibis.expr.signature import Argument as Arg
 
 from ..compiler import BigQueryUDFNode, compiles
@@ -73,7 +72,7 @@ def udf(input_type, output_type, strict=True, libraries=None):
     --------
     >>> if PY38:
     ...     import pytest; pytest.skip("Issue #2085")
-    >>> from ibis.bigquery import udf
+    >>> from ibis_bigquery import udf
     >>> import ibis.expr.datatypes as dt
     >>> @udf(input_type=[dt.double], output_type=dt.double)
     ... def add_one(x):

tests/system/test_client.py

@@ -12,6 +12,7 @@
 import pytz
 from google.api_core import exceptions
 
+import ibis_bigquery
 from ibis_bigquery.client import bigquery_param
 
 pytestmark = pytest.mark.bigquery
@@ -70,7 +71,7 @@ def test_compile_toplevel():
 
     # it works!
     expr = t.foo.sum()
-    result = ibis.bigquery.compile(expr)
+    result = ibis_bigquery.compile(expr)
     # FIXME: remove quotes because bigquery can't use anythig that needs
     # quoting?
     expected = """\
@@ -192,16 +193,16 @@ def test_cast_string_to_date(alltypes, df, type):
 
 
 def test_has_partitions(alltypes, parted_alltypes, client):
-    col = ibis.options.bigquery.partition_col
+    col = client.partition_column
     assert col not in alltypes.columns
     assert col in parted_alltypes.columns
 
 
-def test_different_partition_col_name(client):
+def test_different_partition_col_name(monkeypatch, client):
     col = 'FOO_BAR'
-    with ibis.config.option_context('bigquery.partition_col', col):
-        alltypes = client.table('functional_alltypes')
-        parted_alltypes = client.table('functional_alltypes_parted')
+    monkeypatch.setattr(client, 'partition_column', col)
+    alltypes = client.table('functional_alltypes')
+    parted_alltypes = client.table('functional_alltypes_parted')
     assert col not in alltypes.columns
     assert col in parted_alltypes.columns
 

tests/system/test_connect.py

@@ -10,11 +10,9 @@
 
 pytestmark = pytest.mark.bigquery
 
-bq_backend = ibis_bigquery.Backend()
-
 
 def test_repeated_project_name(project_id, credentials):
-    con = bq_backend.connect(
+    con = ibis_bigquery.connect(
         project_id=project_id,
         dataset_id='{}.testing'.format(project_id),
         credentials=credentials,
@@ -29,12 +27,12 @@ def mock_credentials(*args, **kwargs):
         return creds, 'default-project-id'
 
     monkeypatch.setattr(pydata_google_auth, 'default', mock_credentials)
-    con = bq_backend.connect(project_id='explicit-project-id',)
+    con = ibis_bigquery.connect(project_id='explicit-project-id',)
     assert con.billing_project == 'explicit-project-id'
 
 
 def test_without_dataset(project_id, credentials):
-    con = bq_backend.connect(
+    con = ibis_bigquery.connect(
         project_id=project_id, dataset_id=None, credentials=credentials,
     )
     with pytest.raises(ValueError, match="Unable to determine BigQuery"):
@@ -46,7 +44,7 @@ def test_application_name_sets_user_agent(
 ):
     mock_client = mock.create_autospec(bq.Client)
     monkeypatch.setattr(bq, 'Client', mock_client)
-    bq_backend.connect(
+    ibis_bigquery.connect(
         project_id=project_id,
         dataset_id='bigquery-public-data.stackoverflow',
         application_name='my-great-app/0.7.0',
@@ -67,7 +65,7 @@ def mock_default(*args, **kwargs):
 
     monkeypatch.setattr(pydata_google_auth, "default", mock_default)
 
-    bq_backend.connect(
+    ibis_bigquery.connect(
         project_id=project_id, dataset_id='bigquery-public-data.stackoverflow',
     )
 
@@ -93,7 +91,7 @@ def mock_default(*args, **kwargs):
 
     monkeypatch.setattr(pydata_google_auth, "default", mock_default)
 
-    bq_backend.connect(
+    ibis_bigquery.connect(
         project_id=project_id,
         dataset_id='bigquery-public-data.stackoverflow',
         auth_local_webserver=True,
@@ -114,7 +112,7 @@ def mock_default(*args, **kwargs):
 
     monkeypatch.setattr(pydata_google_auth, "default", mock_default)
 
-    bq_backend.connect(
+    ibis_bigquery.connect(
         project_id=project_id,
         dataset_id='bigquery-public-data.stackoverflow',
         auth_external_data=True,
@@ -136,7 +134,7 @@ def mock_default(*args, **kwargs):
 
     monkeypatch.setattr(pydata_google_auth, "default", mock_default)
 
-    bq_backend.connect(
+    ibis_bigquery.connect(
         project_id=project_id,
         dataset_id="bigquery-public-data.stackoverflow",
         auth_cache="reauth",
@@ -159,7 +157,7 @@ def mock_default(*args, **kwargs):
 
     monkeypatch.setattr(pydata_google_auth, "default", mock_default)
 
-    bq_backend.connect(
+    ibis_bigquery.connect(
         project_id=project_id,
         dataset_id="bigquery-public-data.stackoverflow",
         auth_cache="none",
@@ -173,7 +171,7 @@ def mock_default(*args, **kwargs):
 
 def test_auth_cache_unknown(project_id):
     with pytest.raises(ValueError, match="unexpected value for auth_cache"):
-        bq_backend.connect(
+        ibis_bigquery.connect(
             project_id=project_id,
             dataset_id="bigquery-public-data.stackoverflow",
             auth_cache="not_a_real_cache",