Get ready for test automation

formatting

Make sql client change backwards compatible, fix broken tests

fix the integration workflow and linting

fix broken tests

fix path

lint + install missing sqlalchemy library

skip cache for one run

fix integration so that it doesn't uninstall sqlalchemy

Signed-off-by: Ben Cassell <[email protected]>

Author: benc-db

.github/workflows/code-quality-checks.yml

@@ -1,5 +1,5 @@
 name: Code Quality Checks
-on: 
+on:
   push:
     branches:
       - main
@@ -157,7 +157,7 @@ jobs:
       - name: Install library
         run: poetry install --no-interaction
       #----------------------------------------------
-      #              black the code
+      #              mypy the code
       #----------------------------------------------
       - name: Mypy
         run: poetry run mypy --install-types --non-interactive src

.github/workflows/integration.yml

@@ -0,0 +1,59 @@
+name: Integration Tests
+on:
+  push:
+    paths-ignore:
+      - "**.MD"
+      - "**.md"
+
+jobs:
+  run-e2e-tests:
+    runs-on: ubuntu-latest
+    environment: azure-prod
+    env:
+      DATABRICKS_SERVER_HOSTNAME: ${{ secrets.DATABRICKS_HOST }}
+      DATABRICKS_HTTP_PATH: ${{ secrets.TEST_PECO_WAREHOUSE_HTTP_PATH }}
+      DATABRICKS_TOKEN: ${{ secrets.DATABRICKS_TOKEN }}
+      DATABRICKS_CATALOG: peco
+      DATABRICKS_USER: ${{ secrets.TEST_PECO_SP_ID }}
+    steps:
+      #----------------------------------------------
+      #       check-out repo and set-up python
+      #----------------------------------------------
+      - name: Check out repository
+        uses: actions/checkout@v3
+      - name: Set up python
+        id: setup-python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+      #----------------------------------------------
+      #  -----  install & configure poetry  -----
+      #----------------------------------------------
+      - name: Install Poetry
+        uses: snok/install-poetry@v1
+        with:
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+          installer-parallel: true
+
+      #----------------------------------------------
+      #       load cached venv if cache exists
+      #----------------------------------------------
+      - name: Load cached venv
+        id: cached-poetry-dependencies
+        uses: actions/cache@v2
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ github.event.repository.name }}-${{ hashFiles('**/poetry.lock') }}
+      #----------------------------------------------
+      # install dependencies if cache does not exist
+      #----------------------------------------------
+      - name: Install dependencies
+        run: poetry install --no-interaction --all-extras
+      #----------------------------------------------
+      #              run test suite
+      #----------------------------------------------
+      - name: Run e2e tests
+        run: poetry run python -m pytest tests/e2e
+      - name: Run SQL Alchemy tests
+        run: poetry run python -m pytest src/databricks/sqlalchemy/test_local

conftest.py

@@ -0,0 +1,44 @@
+import os
+import pytest
+
+
+@pytest.fixture(scope="session")
+def host():
+    return os.getenv("DATABRICKS_SERVER_HOSTNAME")
+
+
+@pytest.fixture(scope="session")
+def http_path():
+    return os.getenv("DATABRICKS_HTTP_PATH")
+
+
+@pytest.fixture(scope="session")
+def access_token():
+    return os.getenv("DATABRICKS_TOKEN")
+
+
+@pytest.fixture(scope="session")
+def ingestion_user():
+    return os.getenv("DATABRICKS_USER")
+
+
+@pytest.fixture(scope="session")
+def catalog():
+    return os.getenv("DATABRICKS_CATALOG")
+
+
+@pytest.fixture(scope="session")
+def schema():
+    return os.getenv("DATABRICKS_SCHEMA", "default")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def connection_details(host, http_path, access_token, ingestion_user, catalog, schema):
+    return {
+        "host": host,
+        "http_path": http_path,
+        "access_token": access_token,
+        "ingestion_user": ingestion_user,
+        "catalog": catalog,
+        "schema": schema,
+    }

pyproject.toml

@@ -58,6 +58,7 @@ exclude = ['ttypes\.py$', 'TCLIService\.py$']
 exclude = '/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist|thrift_api)/'
 
 [tool.pytest.ini_options]
+markers = {"reviewed" = "Test case has been reviewed by Databricks"}
 minversion = "6.0"
 log_cli = "false"
 log_cli_level = "INFO"

src/databricks/sql/client.py

@@ -605,12 +605,15 @@ def _handle_staging_operation(
                     "Local file operations are restricted to paths within the configured staging_allowed_local_path"
                 )
 
-        # TODO: Experiment with DBR sending real headers.
-        # The specification says headers will be in JSON format but the current null value is actually an empty list []
+        # May be real headers, or could be json string
+        headers = (
+            json.loads(row.headers) if isinstance(row.headers, str) else row.headers
+        )
+
         handler_args = {
             "presigned_url": row.presignedUrl,
             "local_file": abs_localFile,
-            "headers": json.loads(row.headers or "{}"),
+            "headers": dict(headers) or {},
         }
 
         logger.debug(

src/databricks/sqlalchemy/pytest.ini

@@ -0,0 +1,44 @@
+import os
+import pytest
+
+
+@pytest.fixture(scope="session")
+def host():
+    return os.getenv("DATABRICKS_SERVER_HOSTNAME")
+
+
+@pytest.fixture(scope="session")
+def http_path():
+    return os.getenv("DATABRICKS_HTTP_PATH")
+
+
+@pytest.fixture(scope="session")
+def access_token():
+    return os.getenv("DATABRICKS_TOKEN")
+
+
+@pytest.fixture(scope="session")
+def ingestion_user():
+    return os.getenv("DATABRICKS_USER")
+
+
+@pytest.fixture(scope="session")
+def catalog():
+    return os.getenv("DATABRICKS_CATALOG")
+
+
+@pytest.fixture(scope="session")
+def schema():
+    return os.getenv("DATABRICKS_SCHEMA", "default")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def connection_details(host, http_path, access_token, ingestion_user, catalog, schema):
+    return {
+        "host": host,
+        "http_path": http_path,
+        "access_token": access_token,
+        "ingestion_user": ingestion_user,
+        "catalog": catalog,
+        "schema": schema,
+    }

src/databricks/sqlalchemy/test_local/conftest.py

@@ -1,6 +1,5 @@
 import datetime
 import decimal
-import os
 from typing import Tuple, Union, List
 from unittest import skipIf
 
@@ -19,7 +18,7 @@
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
 from sqlalchemy.schema import DropColumnComment, SetColumnComment
-from sqlalchemy.types import BOOLEAN, DECIMAL, Date, DateTime, Integer, String
+from sqlalchemy.types import BOOLEAN, DECIMAL, Date, Integer, String
 
 try:
     from sqlalchemy.orm import declarative_base
@@ -49,12 +48,12 @@ def version_agnostic_select(object_to_select, *args, **kwargs):
         return select(object_to_select, *args, **kwargs)
 
 
-def version_agnostic_connect_arguments(catalog=None, schema=None) -> Tuple[str, dict]:
-    HOST = os.environ.get("host")
-    HTTP_PATH = os.environ.get("http_path")
-    ACCESS_TOKEN = os.environ.get("access_token")
-    CATALOG = catalog or os.environ.get("catalog")
-    SCHEMA = schema or os.environ.get("schema")
+def version_agnostic_connect_arguments(connection_details) -> Tuple[str, dict]:
+    HOST = connection_details["host"]
+    HTTP_PATH = connection_details["http_path"]
+    ACCESS_TOKEN = connection_details["access_token"]
+    CATALOG = connection_details["catalog"]
+    SCHEMA = connection_details["schema"]
 
     ua_connect_args = {"_user_agent_entry": USER_AGENT_TOKEN}
 
@@ -77,8 +76,8 @@ def version_agnostic_connect_arguments(catalog=None, schema=None) -> Tuple[str,
 
 
 @pytest.fixture
-def db_engine() -> Engine:
-    conn_string, connect_args = version_agnostic_connect_arguments()
+def db_engine(connection_details) -> Engine:
+    conn_string, connect_args = version_agnostic_connect_arguments(connection_details)
     return create_engine(conn_string, connect_args=connect_args)
 
 
@@ -92,10 +91,11 @@ def run_query(db_engine: Engine, query: Union[str, Text]):
 
 
 @pytest.fixture
-def samples_engine() -> Engine:
-    conn_string, connect_args = version_agnostic_connect_arguments(
-        catalog="samples", schema="nyctaxi"
-    )
+def samples_engine(connection_details) -> Engine:
+    details = connection_details.copy()
+    details["catalog"] = "samples"
+    details["schema"] = "nyctaxi"
+    conn_string, connect_args = version_agnostic_connect_arguments(details)
     return create_engine(conn_string, connect_args=connect_args)
 
 
@@ -141,7 +141,7 @@ def test_connect_args(db_engine):
 def test_pandas_upload(db_engine, metadata_obj):
     import pandas as pd
 
-    SCHEMA = os.environ.get("schema")
+    SCHEMA = "default"
     try:
         df = pd.read_excel(
             "src/databricks/sqlalchemy/test_local/e2e/demo_data/MOCK_DATA.xlsx"
@@ -409,7 +409,9 @@ def test_get_table_names_smoke_test(samples_engine: Engine):
         _names is not None, "get_table_names did not succeed"
 
 
-def test_has_table_across_schemas(db_engine: Engine, samples_engine: Engine):
+def test_has_table_across_schemas(
+    db_engine: Engine, samples_engine: Engine, catalog: str, schema: str
+):
     """For this test to pass these conditions must be met:
     - Table samples.nyctaxi.trips must exist
     - Table samples.tpch.customer must exist
@@ -426,9 +428,6 @@ def test_has_table_across_schemas(db_engine: Engine, samples_engine: Engine):
         )
 
         # 3) Check for a table within a different catalog
-        other_catalog = os.environ.get("catalog")
-        other_schema = os.environ.get("schema")
-
         # Create a table in a different catalog
         with db_engine.connect() as conn:
             conn.execute(text("CREATE TABLE test_has_table (numbers_are_cool INT);"))
@@ -442,8 +441,8 @@ def test_has_table_across_schemas(db_engine: Engine, samples_engine: Engine):
                 assert samples_engine.dialect.has_table(
                     connection=conn,
                     table_name="test_has_table",
-                    schema=other_schema,
-                    catalog=other_catalog,
+                    schema=schema,
+                    catalog=catalog,
                 )
             finally:
                 conn.execute(text("DROP TABLE test_has_table;"))
@@ -503,12 +502,12 @@ def test_get_columns(db_engine, sample_table: str):
 
 class TestCommentReflection:
     @pytest.fixture(scope="class")
-    def engine(self):
-        HOST = os.environ.get("host")
-        HTTP_PATH = os.environ.get("http_path")
-        ACCESS_TOKEN = os.environ.get("access_token")
-        CATALOG = os.environ.get("catalog")
-        SCHEMA = os.environ.get("schema")
+    def engine(self, connection_details: dict):
+        HOST = connection_details["host"]
+        HTTP_PATH = connection_details["http_path"]
+        ACCESS_TOKEN = connection_details["access_token"]
+        CATALOG = connection_details["catalog"]
+        SCHEMA = connection_details["schema"]
 
         connection_string = f"databricks://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}&catalog={CATALOG}&schema={SCHEMA}"
         connect_args = {"_user_agent_entry": USER_AGENT_TOKEN}

src/databricks/sqlalchemy/test_local/e2e/test_basic.py

@@ -64,16 +64,16 @@ def test_extract_3l_namespace_from_bad_constraint_string():
         extract_three_level_identifier_from_constraint_string(input)
 
 
-@pytest.mark.parametrize("schema", [None, "some_schema"])
-def test_build_fk_dict(schema):
+@pytest.mark.parametrize("tschema", [None, "some_schema"])
+def test_build_fk_dict(tschema):
     fk_constraint_string = "FOREIGN KEY (`parent_user_id`) REFERENCES `main`.`some_schema`.`users` (`user_id`)"
 
-    result = build_fk_dict("some_fk_name", fk_constraint_string, schema_name=schema)
+    result = build_fk_dict("some_fk_name", fk_constraint_string, schema_name=tschema)
 
     assert result == {
         "name": "some_fk_name",
         "constrained_columns": ["parent_user_id"],
-        "referred_schema": schema,
+        "referred_schema": tschema,
         "referred_table": "users",
         "referred_columns": ["user_id"],
     }

src/databricks/sqlalchemy/test_local/test_parsing.py

@@ -1,11 +1,11 @@
 # Authentication details for running e2e tests 
-host=""
-http_path=""
-access_token=""
+DATABRICKS_SERVER_HOSTNAME=
+DATABRICKS_HTTP_PATH=
+DATABRICKS_TOKEN=
 
 # Only required to run the PySQLStagingIngestionTestSuite
-staging_ingestion_user=""
+DATABRICKS_USER=
 
 # Only required to run SQLAlchemy tests
-catalog=""
-schema=""
+DATABRICKS_CATALOG=
+DATABRICKS_SCHEMA=