chore: enable BQ tests (#17017)

2021-10-14 08:07:47 -07:00 · 2021-10-14 08:07:47 -07:00 · 83a2f8346e
parent d944503873
commit 83a2f8346e
5 changed files with 90 additions and 58 deletions
--- a/requirements/testing.in
+++ b/requirements/testing.in
@ -34,3 +34,5 @@ pytest-cov
 statsd
 pytest-mock
 packaging==21.0
+# DB dependencies
+-e file:.[bigquery]
--- a/requirements/testing.txt
+++ b/requirements/testing.txt
@ -1,4 +1,4 @@
-# SHA1:97a1d3094738243d47824948a0d6fdcb4a784524
+# SHA1:a36e63b551290f1060a819fe4f1f50bc6200403c
 #
 # This file is autogenerated by pip-compile-multi
 # To update, run:
@ -8,13 +8,17 @@
 -r development.txt
 -r integration.txt
 -e file:.
-    # via -r requirements/base.in
+    # via
+    #   -r requirements/base.in
+    #   -r requirements/testing.in
 appnope==0.1.2
    # via ipython
 astroid==2.6.6
    # via pylint
 backcall==0.2.0
    # via ipython
+cachetools==4.2.4
+    # via google-auth
 coverage==5.5
    # via pytest-cov
 decorator==5.0.9
@ -27,6 +31,43 @@ flask-testing==0.8.1
    # via -r requirements/testing.in
 freezegun==1.1.0
    # via -r requirements/testing.in
+google-api-core[grpc]==2.1.0
+    # via
+    #   google-cloud-bigquery
+    #   google-cloud-bigquery-storage
+    #   google-cloud-core
+    #   pybigquery
+google-auth==2.2.1
+    # via
+    #   google-api-core
+    #   google-auth-oauthlib
+    #   google-cloud-core
+    #   pandas-gbq
+    #   pybigquery
+    #   pydata-google-auth
+google-auth-oauthlib==0.4.6
+    # via
+    #   pandas-gbq
+    #   pydata-google-auth
+google-cloud-bigquery[bqstorage,pandas]==2.28.0
+    # via
+    #   apache-superset
+    #   pandas-gbq
+    #   pybigquery
+google-cloud-bigquery-storage==2.9.1
+    # via google-cloud-bigquery
+google-cloud-core==2.1.0
+    # via google-cloud-bigquery
+google-crc32c==1.3.0
+    # via google-resumable-media
+google-resumable-media==2.0.3
+    # via google-cloud-bigquery
+googleapis-common-protos==1.53.0
+    # via google-api-core
+grpcio==1.41.0
+    # via
+    #   google-api-core
+    #   google-cloud-bigquery
 iniconfig==1.1.1
    # via pytest
 ipdb==0.13.9
@ -43,14 +84,22 @@ jedi==0.18.0
    # via ipython
 lazy-object-proxy==1.6.0
    # via astroid
+libcst==0.3.21
+    # via google-cloud-bigquery-storage
 matplotlib-inline==0.1.2
    # via ipython
 mccabe==0.6.1
    # via pylint
+mypy-extensions==0.4.3
+    # via typing-inspect
+oauthlib==3.1.1
+    # via requests-oauthlib
 openapi-schema-validator==0.1.5
    # via openapi-spec-validator
 openapi-spec-validator==0.3.1
    # via -r requirements/testing.in
+pandas-gbq==0.15.0
+    # via apache-superset
 parameterized==0.8.1
    # via -r requirements/testing.in
 parso==0.8.2
@ -61,8 +110,28 @@ pickleshare==0.7.5
    # via ipython
 prompt-toolkit==3.0.19
    # via ipython
+proto-plus==1.19.2
+    # via
+    #   google-cloud-bigquery
+    #   google-cloud-bigquery-storage
+protobuf==3.18.1
+    # via
+    #   google-api-core
+    #   google-cloud-bigquery
+    #   googleapis-common-protos
+    #   proto-plus
 ptyprocess==0.7.0
    # via pexpect
+pyasn1==0.4.8
+    # via
+    #   pyasn1-modules
+    #   rsa
+pyasn1-modules==0.2.8
+    # via google-auth
+pybigquery==0.10.2
+    # via apache-superset
+pydata-google-auth==1.2.0
+    # via pandas-gbq
 pyfakefs==4.5.0
    # via -r requirements/testing.in
 pygments==2.9.0
@ -82,12 +151,18 @@ pytest-cov==2.12.1
    # via -r requirements/testing.in
 pytest-mock==3.6.1
    # via -r requirements/testing.in
+requests-oauthlib==1.3.0
+    # via google-auth-oauthlib
+rsa==4.7.2
+    # via google-auth
 statsd==3.3.0
    # via -r requirements/testing.in
 traitlets==5.0.5
    # via
    #   ipython
    #   matplotlib-inline
+typing-inspect==0.7.1
+    # via libcst
 wcwidth==0.2.5
    # via prompt-toolkit
 websocket-client==1.2.0
--- a/setup.cfg
+++ b/setup.cfg
@ -30,7 +30,7 @@ combine_as_imports = true
 include_trailing_comma = true
 line_length = 88
 known_first_party = superset
-known_third_party =alembic,apispec,backoff,bleach,cachelib,celery,click,colorama,cron_descriptor,croniter,cryptography,dateutil,deprecation,flask,flask_appbuilder,flask_babel,flask_caching,flask_compress,flask_jwt_extended,flask_login,flask_migrate,flask_sqlalchemy,flask_talisman,flask_testing,flask_wtf,freezegun,geohash,geopy,graphlib,holidays,humanize,isodate,jinja2,jwt,markdown,markupsafe,marshmallow,marshmallow_enum,msgpack,numpy,pandas,parameterized,parsedatetime,pgsanity,pkg_resources,polyline,prison,progress,pyarrow,pyhive,pyparsing,pytest,pytest_mock,pytz,redis,requests,selenium,setuptools,simplejson,slack,sqlalchemy,sqlalchemy_utils,sqlparse,typing_extensions,urllib3,werkzeug,wtforms,wtforms_json,yaml
+known_third_party =alembic,apispec,backoff,bleach,cachelib,celery,click,colorama,cron_descriptor,croniter,cryptography,dateutil,deprecation,flask,flask_appbuilder,flask_babel,flask_caching,flask_compress,flask_jwt_extended,flask_login,flask_migrate,flask_sqlalchemy,flask_talisman,flask_testing,flask_wtf,freezegun,geohash,geopy,graphlib,holidays,humanize,isodate,jinja2,jwt,markdown,markupsafe,marshmallow,marshmallow_enum,msgpack,numpy,pandas,parameterized,parsedatetime,pgsanity,pkg_resources,polyline,prison,progress,pyarrow,pybigquery,pyhive,pyparsing,pytest,pytest_mock,pytz,redis,requests,selenium,setuptools,simplejson,slack,sqlalchemy,sqlalchemy_utils,sqlparse,typing_extensions,urllib3,werkzeug,wtforms,wtforms_json,yaml
 multi_line_output = 3
 order_by_type = false

--- a/tests/integration_tests/db_engine_specs/bigquery_tests.py
+++ b/tests/integration_tests/db_engine_specs/bigquery_tests.py
@ -176,15 +176,13 @@ class TestBigQueryDbEngineSpec(TestDbEngineSpec):
        sys.modules["pandas_gbq"] = mock.MagicMock()
        df = DataFrame()
        database = mock.MagicMock()
-        self.assertRaisesRegexp(
-            Exception,
-            "Could not import libraries",
-            BigQueryEngineSpec.df_to_sql,
-            database=database,
-            table=Table(table="name", schema="schema"),
-            df=df,
-            to_sql_kwargs={},
-        )
+        with self.assertRaises(Exception):
+            BigQueryEngineSpec.df_to_sql(
+                database=database,
+                table=Table(table="name", schema="schema"),
+                df=df,
+                to_sql_kwargs={},
+            )

        invalid_kwargs = [
            {"name": "some_name"},
--- a/tests/unit_tests/db_engine_specs/test_bigquery.py
+++ b/tests/unit_tests/db_engine_specs/test_bigquery.py
@ -17,9 +17,9 @@
 # pylint: disable=unused-argument, import-outside-toplevel, protected-access

 from flask.ctx import AppContext
+from pybigquery.sqlalchemy_bigquery import BigQueryDialect
 from pytest_mock import MockFixture
 from sqlalchemy import select
-from sqlalchemy.dialects.sqlite.base import SQLiteDialect
 from sqlalchemy.sql import sqltypes


@ -59,19 +59,7 @@ def test_get_fields(app_context: AppContext) -> None:
    columns = [{"name": "limit"}, {"name": "name"}, {"name": "project.name"}]
    fields = BigQueryEngineSpec._get_fields(columns)

-    # generic SQL
    query = select(fields)
-    assert (
-        str(query)
-        == 'SELECT "limit" AS "limit", name AS name, "project.name" AS project__name'
-    )
-
-    # BigQuery-specific SQL
-    try:
-        from pybigquery.sqlalchemy_bigquery import BigQueryDialect
-    except ModuleNotFoundError:
-        return
-
    assert str(query.compile(dialect=BigQueryDialect())) == (
        "SELECT `limit` AS `limit`, `name` AS `name`, "
        "`project`.`name` AS `project__name`"
@ -133,42 +121,11 @@ def test_select_star(mocker: MockFixture, app_context: AppContext) -> None:

    # mock the database so we can compile the query
    database = mocker.MagicMock()
-    database.compile_sqla_query = lambda query: str(
-        query.compile(dialect=SQLiteDialect())
-    )
-
-    # use SQLite dialect so we don't need the BQ dependency
-    engine = mocker.MagicMock()
-    engine.dialect = SQLiteDialect()
-
-    sql = BigQueryEngineSpec.select_star(
-        database=database,
-        table_name="my_table",
-        engine=engine,
-        schema=None,
-        limit=100,
-        show_cols=True,
-        indent=True,
-        latest_partition=False,
-        cols=cols,
-    )
-    assert (
-        sql
-        == """SELECT trailer AS trailer
-FROM my_table
-LIMIT ?
-OFFSET ?"""
-    )
-
-    # BigQuery-specific SQL
-    try:
-        from pybigquery.sqlalchemy_bigquery import BigQueryDialect
-    except ModuleNotFoundError:
-        return
-
    database.compile_sqla_query = lambda query: str(
        query.compile(dialect=BigQueryDialect())
    )
+
+    engine = mocker.MagicMock()
    engine.dialect = BigQueryDialect()

    sql = BigQueryEngineSpec.select_star(