Repository: mabel-dev/opteryx

Issue:
✨ add a bigquery test to the CI suite
https://googleapis.dev/python/pybigquery/0.7.0/

I've created a mabeldev/public/planets dataset (1.7Kb) for testing


Discussion:

Resolving patch:
diff --git a/opteryx/__init__.py b/opteryx/__init__.py
--- a/opteryx/__init__.py
+++ b/opteryx/__init__.py
@@ -47,9 +47,6 @@
 if os.environ.get("OPTERYX_DEBUG") is not None:
     from opteryx.debugging import OpteryxOrsoImportFinder
 
-from orso.logging import get_logger
-from orso.logging import set_log_name
-
 from opteryx import config
 from opteryx.managers.cache.cache_manager import CacheManager  # isort:skip
 
@@ -60,6 +57,9 @@
 from opteryx.connectors import register_df
 from opteryx.connectors import register_store
 
+from opteryx.__author__ import __author__
+from opteryx.__build__ import __build__
+from opteryx.__version__ import __version__
 
 __all__ = [
     "apilevel",
@@ -71,6 +71,9 @@
     "register_arrow",
     "register_df",
     "register_store",
+    "__author__",
+    "__build__",
+    "__version__",
 ]
 
 # PEP-249 specifies these attributes for a Python Database API 2.0 compliant interface
diff --git a/opteryx/components/binder/binder_visitor.py b/opteryx/components/binder/binder_visitor.py
--- a/opteryx/components/binder/binder_visitor.py
+++ b/opteryx/components/binder/binder_visitor.py
@@ -517,7 +517,7 @@ def visit_scan(self, node: Node, context: BindingContext) -> Tuple[Node, Binding
         return node, context
 
     def visit_set(self, node: Node, context: BindingContext) -> Tuple[Node, BindingContext]:
-        node.variables = context["connection"].variables
+        node.variables = context.connection.variables
         return node, context
 
     def visit_show_columns(


Test patch:
diff --git a/tests/requirements.txt b/tests/requirements.txt
--- a/tests/requirements.txt
+++ b/tests/requirements.txt
@@ -29,5 +29,7 @@ polars
 duckdb
 duckdb-engine
 sqlalchemy-cockroachdb
+sqlalchemy-bigquery
+google-cloud-bigquery-storage
 
 setuptools_rust
\ No newline at end of file
diff --git a/tests/sql_battery/test_shapes_and_errors_battery.py b/tests/sql_battery/test_shapes_and_errors_battery.py
--- a/tests/sql_battery/test_shapes_and_errors_battery.py
+++ b/tests/sql_battery/test_shapes_and_errors_battery.py
@@ -611,8 +611,7 @@
         ("SELECT name FROM $planets WHERE id IN (SELECT * FROM UNNEST((1,2,3)) as id)", 3, 1, UnsupportedSyntaxError),  # temp
         ("SELECT count(planetId) FROM (SELECT DISTINCT planetId FROM $satellites) AS SQ", 1, 1, None),
         ("SELECT COUNT(*) FROM (SELECT planetId FROM $satellites WHERE planetId < 7) AS SQ GROUP BY planetId", 4, 1, None),
-]
-A = [
+
         ("EXPLAIN SELECT * FROM $satellites", 1, 3, None),
         ("EXPLAIN SELECT * FROM $satellites WHERE id = 8", 2, 3, None),
         ("SET version = '1.0';", None, None, PermissionsError),
@@ -673,7 +672,8 @@
         ("SELECT P_1.* FROM $planets AS P_1 INNER JOIN $satellites AS P_2 ON P_2.id = P_1.id AND P_2.name = P_1.name", 0, 20, None),
         ("SELECT P_1.* FROM $planets AS P_1 INNER JOIN $satellites AS P_2 ON P_1.id = P_2.id AND P_2.name = P_1.name", 0, 20, None),
         ("SELECT P_1.* FROM $planets AS P_1 INNER JOIN $satellites AS P_2 ON P_2.id = P_1.id AND P_1.name = P_2.name", 0, 20, None),
-
+]
+A = [
         ("SELECT * FROM $planets AS P_1 INNER JOIN $planets AS P_2 ON P_1.id = P_2.id AND P_2.name = P_1.name", 9, 40, None),
         ("SELECT * FROM $planets NATURAL JOIN generate_series(1, 5) as id", 5, 20, None),
         ("SELECT * FROM $planets NATURAL JOIN $planets", 5, 20, None),
diff --git a/tests/storage/test_sql_biq_query.py b/tests/storage/test_sql_biq_query.py
new file mode 100644
--- /dev/null
+++ b/tests/storage/test_sql_biq_query.py
@@ -0,0 +1,36 @@
+"""
+Test we can read from Sqlite - this is a basic exercise of the SQL Connector
+"""
+import os
+import sys
+
+sys.path.insert(1, os.path.join(sys.path[0], "../.."))
+
+import opteryx
+from opteryx.connectors import SqlConnector
+from tests.tools import is_arm, is_mac, is_windows, skip_if
+
+BIG_QUERY_PROJECT: str = "mabeldev"
+
+
+# skip to reduce billing
+# @skip_if(is_arm() or is_windows() or is_mac())
+def test_bigquery_storage():
+    from sqlalchemy.engine import create_engine
+
+    engine = create_engine(f"bigquery://{BIG_QUERY_PROJECT}")
+
+    opteryx.register_store("bq", SqlConnector, remove_prefix=True, engine=engine)
+
+    results = opteryx.query("SELECT * FROM bq.public.planets")
+    assert results.rowcount == 9, results.rowcount
+
+    # PROCESS THE DATA IN SOME WAY
+    results = opteryx.query("SELECT COUNT(*) FROM bq.public.planets;")
+    assert results.rowcount == 1, results.rowcount
+
+
+if __name__ == "__main__":  # pragma: no cover
+    from tests.tools import run_tests
+
+    run_tests()