update tests to use bigquery_magics module

.kokoro/presubmit/snippets-3.8.cfg → .kokoro/presubmit/system-3.11.cfg

@@ -3,5 +3,5 @@
 # Only run this nox session.
 env_vars: {
     key: "NOX_SESSION"
-    value: "snippets-3.8"
-}
+    value: "system-3.11"
+}

bigquery_magics/__init__.py

@@ -12,9 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from google.cloud.bigquery.magics.magics import context
+import bigquery_magics.config
+import bigquery_magics.version
 
+context = bigquery_magics.config.context
+__version__ = bigquery_magics.version.__version__
 
-# For backwards compatibility we need to make the context available in the path
-# google.cloud.bigquery.magics.context
-__all__ = ("context",)
+
+def load_ipython_extension(ipython):
+    """Called by IPython when this module is loaded as an IPython extension."""
+    # Import here to avoid circular imports.
+    from bigquery_magics.bigquery import _cell_magic
+
+    ipython.register_magic_function(
+        _cell_magic, magic_kind="cell", magic_name="bigquery"
+    )
+
+
+__all__ = (
+    # For backwards compatibility we need to make the context available in
+    # the path google.cloud.bigquery.magics.context.
+    "context",
+    "__version__",
+    "load_ipython_extension",
+)

bigquery_magics/_versions_helpers.py

@@ -0,0 +1,91 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared helper functions for verifying versions of installed modules."""
+
+from typing import Any
+
+from google.cloud.bigquery import exceptions
+import packaging.version
+
+_MIN_BQ_STORAGE_VERSION = packaging.version.Version("2.0.0")
+
+
+class BQStorageVersions:
+    """Version comparisons for google-cloud-bigqueyr-storage package."""
+
+    def __init__(self):
+        self._installed_version = None
+
+    @property
+    def installed_version(self) -> packaging.version.Version:
+        """Return the parsed version of google-cloud-bigquery-storage."""
+        if self._installed_version is None:
+            from google.cloud import bigquery_storage
+
+            self._installed_version = packaging.version.parse(
+                # Use 0.0.0, since it is earlier than any released version.
+                # Legacy versions also have the same property, but
+                # creating a LegacyVersion has been deprecated.
+                # https://github.com/pypa/packaging/issues/321
+                getattr(bigquery_storage, "__version__", "0.0.0")
+            )
+
+        return self._installed_version  # type: ignore
+
+    def try_import(self, raise_if_error: bool = False) -> Any:
+        """Tries to import the bigquery_storage module, and returns results
+        accordingly. It also verifies the module version is recent enough.
+
+        If the import succeeds, returns the ``bigquery_storage`` module.
+
+        If the import fails,
+        returns ``None`` when ``raise_if_error == False``,
+        raises Error when ``raise_if_error == True``.
+
+        Returns:
+            The ``bigquery_storage`` module or ``None``.
+
+        Raises:
+            exceptions.BigQueryStorageNotFoundError:
+                If google-cloud-bigquery-storage is not installed
+            exceptions.LegacyBigQueryStorageError:
+                If google-cloud-bigquery-storage package is outdated
+        """
+        try:
+            from google.cloud import bigquery_storage  # type: ignore
+        except ImportError:
+            if raise_if_error:
+                msg = (
+                    "Package google-cloud-bigquery-storage not found. "
+                    "Install google-cloud-bigquery-storage version >= "
+                    f"{_MIN_BQ_STORAGE_VERSION}."
+                )
+                raise exceptions.BigQueryStorageNotFoundError(msg)
+            return None
+
+        if self.installed_version < _MIN_BQ_STORAGE_VERSION:
+            if raise_if_error:
+                msg = (
+                    "Dependency google-cloud-bigquery-storage is outdated, "
+                    f"please upgrade it to version >= {_MIN_BQ_STORAGE_VERSION} "
+                    f"(version found: {self.installed_version})."
+                )
+                raise exceptions.LegacyBigQueryStorageError(msg)
+            return None
+
+        return bigquery_storage
+
+
+BQ_STORAGE_VERSIONS = BQStorageVersions()