Add spark session connection

.circleci/config.yml

@@ -10,6 +10,24 @@ jobs:
       - checkout
       - run: tox -e flake8,unit
 
+  integration-spark-session:
+    environment:
+      DBT_INVOCATION_ENV: circle
+    docker:
+      - image: godatadriven/pyspark:3.1
+    steps:
+      - checkout
+      - run: apt-get update
+      - run: python3 -m pip install --upgrade pip
+      - run: apt-get install -y git gcc g++ unixodbc-dev libsasl2-dev
+      - run: python3 -m pip install tox
+      - run:
+          name: Run integration tests
+          command: tox -e integration-spark-session
+          no_output_timeout: 1h
+      - store_artifacts:
+          path: ./logs
+
   integration-spark-thrift:
     environment:
       DBT_INVOCATION_ENV: circle
@@ -90,7 +108,7 @@ jobs:
           no_output_timeout: 1h
       - store_artifacts:
           path: ./logs
-          
+
   integration-spark-databricks-odbc-endpoint:
     <<: *databricks-odbc
     steps:
@@ -107,6 +125,9 @@ workflows:
   test-everything:
     jobs:
       - unit
+      - integration-spark-session:
+          requires:
+            - unit
       - integration-spark-thrift:
           requires:
             - unit

CHANGELOG.md

@@ -1,3 +1,9 @@
+### Features
+- Add session connection method ([#272](https://github.com/dbt-labs/dbt-spark/issues/272), [#279](https://github.com/dbt-labs/dbt-spark/pull/279))
+
+### Contributors
+- [@JCZuurmond](https://github.com/dbt-labs/dbt-spark/pull/279) ( [#279](https://github.com/dbt-labs/dbt-spark/pull/279))
+
 ## dbt-spark 1.1.0b1 (March 23, 2022)
 
 ### Features

dbt/adapters/spark/connections.py

@@ -55,6 +55,7 @@ class SparkConnectionMethod(StrEnum):
     THRIFT = 'thrift'
     HTTP = 'http'
     ODBC = 'odbc'
+    SESSION = 'session'
 
 
 @dataclass
@@ -133,6 +134,18 @@ def __post_init__(self):
                 "`pip install dbt-spark[PyHive]`"
             )
 
+        if self.method == SparkConnectionMethod.SESSION:
+            try:
+                import pyspark  # noqa: F401
+            except ImportError as e:
+                raise dbt.exceptions.RuntimeException(
+                    f"{self.method} connection method requires "
+                    "additional dependencies. \n"
+                    "Install the additional required dependencies with "
+                    "`pip install dbt-spark[session]`\n\n"
+                    f"ImportError({e.msg})"
+                ) from e
+
     @property
     def type(self):
         return 'spark'
@@ -443,6 +456,12 @@ def open(cls, connection):
 
                     conn = pyodbc.connect(connection_str, autocommit=True)
                     handle = PyodbcConnectionWrapper(conn)
+                elif creds.method == SparkConnectionMethod.SESSION:
+                    from .session import (  # noqa: F401
+                        Connection,
+                        SessionConnectionWrapper,
+                    )
+                    handle = SessionConnectionWrapper(Connection())
                 else:
                     raise dbt.exceptions.DbtProfileError(
                         f"invalid credential method: {creds.method}"

dbt/adapters/spark/session.py

@@ -0,0 +1,221 @@
+"""Spark session integration."""
+
+from __future__ import annotations
+
+import datetime as dt
+from types import TracebackType
+from typing import Any
+
+from dbt.events import AdapterLogger
+from dbt.utils import DECIMALS
+from pyspark.sql import DataFrame, Row, SparkSession
+
+
+logger = AdapterLogger("Spark")
+NUMBERS = DECIMALS + (int, float)
+
+
+class Cursor:
+    """
+    Mock a pyodbc cursor.
+
+    Source
+    ------
+    https://github.com/mkleehammer/pyodbc/wiki/Cursor
+    """
+
+    def __init__(self) -> None:
+        self._df: DataFrame | None = None
+        self._rows: list[Row] | None = None
+
+    def __enter__(self) -> Cursor:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: Exception | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        self.close()
+        return True
+
+    @property
+    def description(
+        self,
+    ) -> list[tuple[str, str, None, None, None, None, bool]]:
+        """
+        Get the description.
+
+        Returns
+        -------
+        out : list[tuple[str, str, None, None, None, None, bool]]
+            The description.
+
+        Source
+        ------
+        https://github.com/mkleehammer/pyodbc/wiki/Cursor#description
+        """
+        if self._df is None:
+            description = list()
+        else:
+            description = [
+                (
+                    field.name,
+                    field.dataType.simpleString(),
+                    None,
+                    None,
+                    None,
+                    None,
+                    field.nullable,
+                )
+                for field in self._df.schema.fields
+            ]
+        return description
+
+    def close(self) -> None:
+        """
+        Close the connection.
+
+        Source
+        ------
+        https://github.com/mkleehammer/pyodbc/wiki/Cursor#close
+        """
+        self._df = None
+        self._rows = None
+
+    def execute(self, sql: str, *parameters: Any) -> None:
+        """
+        Execute a sql statement.
+
+        Parameters
+        ----------
+        sql : str
+            Execute a sql statement.
+        *parameters : Any
+            The parameters.
+
+        Raises
+        ------
+        NotImplementedError
+            If there are parameters given. We do not format sql statements.
+
+        Source
+        ------
+        https://github.com/mkleehammer/pyodbc/wiki/Cursor#executesql-parameters
+        """
+        if len(parameters) > 0:
+            sql = sql % parameters
+        spark_session = SparkSession.builder.enableHiveSupport().getOrCreate()
+        self._df = spark_session.sql(sql)
+
+    def fetchall(self) -> list[Row] | None:
+        """
+        Fetch all data.
+
+        Returns
+        -------
+        out : list[Row] | None
+            The rows.
+
+        Source
+        ------
+        https://github.com/mkleehammer/pyodbc/wiki/Cursor#fetchall
+        """
+        if self._rows is None and self._df is not None:
+            self._rows = self._df.collect()
+        return self._rows
+
+    def fetchone(self) -> Row | None:
+        """
+        Fetch the first output.
+
+        Returns
+        -------
+        out : Row | None
+            The first row.
+
+        Source
+        ------
+        https://github.com/mkleehammer/pyodbc/wiki/Cursor#fetchone
+        """
+        if self._rows is None and self._df is not None:
+            self._rows = self._df.collect()
+
+        if self._rows is not None and len(self._rows) > 0:
+            row = self._rows.pop(0)
+        else:
+            row = None
+
+        return row
+
+
+class Connection:
+    """
+    Mock a pyodbc connection.
+
+    Source
+    ------
+    https://github.com/mkleehammer/pyodbc/wiki/Connection
+    """
+
+    def cursor(self) -> Cursor:
+        """
+        Get a cursor.
+
+        Returns
+        -------
+        out : Cursor
+            The cursor.
+        """
+        return Cursor()
+
+
+class SessionConnectionWrapper(object):
+    """Connection wrapper for the sessoin connection method."""
+
+    def __init__(self, handle):
+        self.handle = handle
+        self._cursor = None
+
+    def cursor(self):
+        self._cursor = self.handle.cursor()
+        return self
+
+    def cancel(self):
+        logger.debug("NotImplemented: cancel")
+
+    def close(self):
+        if self._cursor:
+            self._cursor.close()
+
+    def rollback(self, *args, **kwargs):
+        logger.debug("NotImplemented: rollback")
+
+    def fetchall(self):
+        return self._cursor.fetchall()
+
+    def execute(self, sql, bindings=None):
+        if sql.strip().endswith(";"):
+            sql = sql.strip()[:-1]
+
+        if bindings is None:
+            self._cursor.execute(sql)
+        else:
+            bindings = [self._fix_binding(binding) for binding in bindings]
+            self._cursor.execute(sql, *bindings)
+
+    @property
+    def description(self):
+        return self._cursor.description
+
+    @classmethod
+    def _fix_binding(cls, value):
+        """Convert complex datatypes to primitives that can be loaded by
+        the Spark driver"""
+        if isinstance(value, NUMBERS):
+            return float(value)
+        elif isinstance(value, dt.datetime):
+            return f"'{value.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}'"
+        else:
+            return f"'{value}'"

dev_requirements.txt

@@ -16,5 +16,5 @@ pytest-csv
 
 # Test requirements
 pytest-dbt-adapter==0.6.0
-sasl==0.2.1
+sasl>=0.2.1
 thrift_sasl==0.4.1

setup.py

@@ -61,7 +61,10 @@ def _get_dbt_core_version():
     'PyHive[hive]>=0.6.0,<0.7.0',
     'thrift>=0.11.0,<0.16.0',
 ]
-all_extras = odbc_extras + pyhive_extras
+session_extras = [
+    "pyspark>=3.0.0,<4.0.0"
+]
+all_extras = odbc_extras + pyhive_extras + session_extras
 
 setup(
     name=package_name,
@@ -83,8 +86,9 @@ def _get_dbt_core_version():
     ],
     extras_require={
         "ODBC": odbc_extras,
-        "PyHive":  pyhive_extras,
-        "all": all_extras
+        "PyHive": pyhive_extras,
+        "session": session_extras,
+        "all": all_extras,
     },
     zip_safe=False,
     classifiers=[

tests/integration/conftest.py

@@ -1,10 +1,4 @@
 def pytest_configure(config):
-    config.addinivalue_line(
-        "markers", "profile_databricks_cluster"
-    )
-    config.addinivalue_line(
-        "markers", "profile_databricks_sql_endpoint"
-    )
-    config.addinivalue_line(
-        "markers", "profile_apache_spark"
-    )
+    config.addinivalue_line("markers", "profile_databricks_cluster")
+    config.addinivalue_line("markers", "profile_databricks_sql_endpoint")
+    config.addinivalue_line("markers", "profile_apache_spark")

tests/specs/spark-session.dbtspec

@@ -0,0 +1,17 @@
+target:
+  type: spark
+  method: session
+  host: localhost
+  schema: "analytics_{{ var('_dbt_random_suffix') }}"
+sequences:
+  test_dbt_empty: empty
+  # requires a metastore for persisting over dbt runs
+  # test_dbt_base: base
+  # test_dbt_ephemeral: ephemeral
+  # test_dbt_incremental: incremental
+  # snapshots require delta format
+  # test_dbt_snapshot_strategy_timestamp: snapshot_strategy_timestamp
+  # test_dbt_snapshot_strategy_check_cols: snapshot_strategy_check_cols
+  test_dbt_data_test: data_test
+  test_dbt_schema_test: schema_test
+  test_dbt_ephemeral_data_tests: data_test_ephemeral_models

tox.ini

@@ -57,3 +57,15 @@ deps =
     -r{toxinidir}/requirements.txt
     -r{toxinidir}/dev_requirements.txt
     -e.
+
+[testenv:integration-spark-session]
+basepython = python3
+commands = /bin/bash -c '{envpython} -m pytest -v tests/specs/spark-session.dbtspec'
+passenv =
+    DBT_*
+    PYTEST_*
+    PIP_CACHE_DIR
+deps =
+    -r{toxinidir}/requirements.txt
+    -r{toxinidir}/dev_requirements.txt
+    -e.[session]