Livy connection support for Spark SQL models

dbt/adapters/spark/connections.py

@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+from threading import Lock
 
 from dbt.adapters.contracts.connection import (
     AdapterResponse,
@@ -13,6 +14,7 @@
 
 from dbt_common.utils.encoding import DECIMALS
 from dbt.adapters.spark import __version__
+from dbt.adapters.spark.livysession import LivyConnectionManager, LivySessionConnectionWrapper
 
 try:
     from TCLIService.ttypes import TOperationState as ThriftState
@@ -30,7 +32,7 @@
 import sqlparams
 from dbt_common.dataclass_schema import StrEnum
 from dataclasses import dataclass, field
-from typing import Any, Dict, Optional, Union, Tuple, List, Generator, Iterable, Sequence
+from typing import Any, Dict, Optional, Union, Tuple, List, Generator, Iterable, Sequence, Hashable
 
 from abc import ABC, abstractmethod
 
@@ -47,7 +49,7 @@
 import time
 
 logger = AdapterLogger("Spark")
-
+lock = Lock()
 NUMBERS = DECIMALS + (int, float)
 
 
@@ -60,6 +62,7 @@ class SparkConnectionMethod(StrEnum):
     HTTP = "http"
     ODBC = "odbc"
     SESSION = "session"
+    LIVY = "livy"
 
 
 @dataclass
@@ -83,6 +86,8 @@ class SparkCredentials(Credentials):
     use_ssl: bool = False
     server_side_parameters: Dict[str, str] = field(default_factory=dict)
     retry_all: bool = False
+    livy_session_parameters: Dict[str, Any] = field(default_factory=dict)
+    verify_ssl_certificate: Optional[bool] = True
 
     @classmethod
     def __pre_deserialize__(cls, data: Any) -> Any:
@@ -346,6 +351,7 @@ class SparkConnectionManager(SQLConnectionManager):
     SPARK_CLUSTER_HTTP_PATH = "/sql/protocolv1/o/{organization}/{cluster}"
     SPARK_SQL_ENDPOINT_HTTP_PATH = "/sql/1.0/endpoints/{endpoint}"
     SPARK_CONNECTION_URL = "{host}:{port}" + SPARK_CLUSTER_HTTP_PATH
+    connection_managers: Dict[Hashable, LivyConnectionManager] = {}
 
     @contextmanager
     def exception_handler(self, sql: str) -> Generator[None, None, None]:
@@ -527,6 +533,45 @@ def open(cls, connection: Connection) -> Connection:
                     handle = SessionConnectionWrapper(
                         Connection(server_side_parameters=creds.server_side_parameters)
                     )
+                elif creds.method == SparkConnectionMethod.LIVY:
+                    # connect to livy interactive session
+
+                    lock.acquire()
+                    try:
+                        thread_id = cls.get_thread_identifier()
+
+                        if thread_id not in SparkConnectionManager.connection_managers:
+                            if len(SparkConnectionManager.connection_managers) > 0:
+                                # Return same livy session
+                                livyConnMgr = list(
+                                    SparkConnectionManager.connection_managers.values()
+                                )[0]
+                                SparkConnectionManager.connection_managers[thread_id] = livyConnMgr
+                            else:
+                                SparkConnectionManager.connection_managers[
+                                    thread_id
+                                ] = LivyConnectionManager()
+
+                        handle = LivySessionConnectionWrapper(  # type: ignore
+                            SparkConnectionManager.connection_managers[thread_id].connect(
+                                creds.host,
+                                creds.user,
+                                creds.password,
+                                creds.auth,
+                                creds.livy_session_parameters,
+                                creds.verify_ssl_certificate,
+                            )
+                        )
+                        connection.state = ConnectionState.OPEN
+
+                    except Exception as ex:
+                        logger.debug("Connection error: {}".format(ex))
+                        connection.state = ConnectionState.FAIL
+                        raise ex
+
+                    finally:
+                        lock.release()
+
                 else:
                     raise DbtConfigError(f"invalid credential method: {creds.method}")
                 break

dbt/adapters/spark/impl.py

@@ -516,6 +516,19 @@ def debug_query(self) -> None:
         """Override for DebugTask method"""
         self.execute("select 1 as id")
 
+        self.connections.get_thread_connection().handle.close()
+
+    def cleanup_connections(self) -> None:
+        self.connections.cleanup_all()
+        logger.debug("cleanup_connections")
+
+        # close all sessions
+        for conn_mgr in SparkConnectionManager.connection_managers:
+            SparkConnectionManager.connection_managers[conn_mgr].delete_session()
+
+        # reset connection_manager list
+        SparkConnectionManager.connection_managers = {}
+
 
 # spark does something interesting with joins when both tables have the same
 # static values for the join condition and complains that the join condition is