Explicitly add databricks config to DatabricksSession

nrccua · Mar 16, 2024 · dda28e5 · dda28e5
1 parent d74001d
commit dda28e5
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 3 deletions.
diff --git a/HISTORY.rst b/HISTORY.rst
@@ -3,6 +3,11 @@ History
 =======
 
 
+v0.20.15 (2024-03-16)
+
+* Explicitly add databricks config to DatabricksSession.
+
+
 v0.20.14 (2024-03-15)
 
 * Fix bug with spark not being instantiated.

diff --git a/aioradio/ds_utils.py b/aioradio/ds_utils.py
@@ -50,8 +50,12 @@
 
 try:
     from databricks.connect import DatabricksSession
-    spark = DatabricksSession.builder.getOrCreate()
-except ImportError:
+    spark = DatabricksSession.builder.remote(
+        host=os.environ['DATABRICKS_HOST'],
+        token=os.environ['DATABRICKS_TOKEN'],
+        cluster_id=os.environ['DATABRICKS_CLUSTER_ID']
+    ).getOrCreate()
+except Exception:
     spark = SparkSession.builder.getOrCreate()
 
 

diff --git a/aioradio/requirements.txt b/aioradio/requirements.txt
@@ -4,6 +4,7 @@ backoff==2.2.1
 boto3==1.34.34
 botocore==1.34.34
 cython==3.0.8
+databricks-connect==14.3.1
 ddtrace==2.6.5
 dominodatalab==1.3.0
 fakeredis==2.21.1

diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
     long_description = fileobj.read()
 
 setup(name='aioradio',
-    version='0.20.14',
+    version='0.20.15',
     description='Generic asynchronous i/o python utilities for AWS services (SQS, S3, DynamoDB, Secrets Manager), Redis, MSSQL (pyodbc), JIRA and more',
     long_description=long_description,
     long_description_content_type="text/markdown",