Add server side parameters to http connection method (#825)

* Pass server side parameters to hive connects * Force server_side_parameters keys and values to strings * Undo changes to thrift * Add changie * Revert enfore server side parameters to be string * Update hive http test to check server_side_parameters * Format test config --------- Co-authored-by: colin-rogers-dbt <[email protected]>
dbt-labs · Jul 28, 2023 · 2680d88 · 2680d88
1 parent 53809c3
commit 2680d88
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/.changes/unreleased/Features-20230707-113337.yaml b/.changes/unreleased/Features-20230707-113337.yaml
@@ -0,0 +1,6 @@
+kind: Features
+body: Add server_side_parameters to HTTP connection method
+time: 2023-07-07T11:33:37.794112+02:00
+custom:
+  Author: Fokko,JCZuurmond
+  Issue: "824"
diff --git a/dbt/adapters/spark/connections.py b/dbt/adapters/spark/connections.py
@@ -381,7 +381,10 @@ def open(cls, connection: Connection) -> Connection:
                     token = base64.standard_b64encode(raw_token).decode()
                     transport.setCustomHeaders({"Authorization": "Basic {}".format(token)})
 
-                    conn = hive.connect(thrift_transport=transport)
+                    conn = hive.connect(
+                        thrift_transport=transport,
+                        configuration=creds.server_side_parameters,
+                    )
                     handle = PyhiveConnectionWrapper(conn)
                 elif creds.method == SparkConnectionMethod.THRIFT:
                     cls.validate_creds(creds, ["host", "port", "user", "schema"])

diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py
@@ -39,6 +39,7 @@ def _get_target_http(self, project):
                         "token": "abc123",
                         "organization": "0123456789",
                         "cluster": "01234-23423-coffeetime",
+                        "server_side_parameters": {"spark.driver.memory": "4g"},
                     }
                 },
                 "target": "test",
@@ -147,13 +148,14 @@ def test_http_connection(self):
         config = self._get_target_http(self.project_cfg)
         adapter = SparkAdapter(config)
 
-        def hive_http_connect(thrift_transport):
+        def hive_http_connect(thrift_transport, configuration):
             self.assertEqual(thrift_transport.scheme, "https")
             self.assertEqual(thrift_transport.port, 443)
             self.assertEqual(thrift_transport.host, "myorg.sparkhost.com")
             self.assertEqual(
                 thrift_transport.path, "/sql/protocolv1/o/0123456789/01234-23423-coffeetime"
             )
+            self.assertEqual(configuration["spark.driver.memory"], "4g")
 
         # with mock.patch.object(hive, 'connect', new=hive_http_connect):
         with mock.patch("dbt.adapters.spark.connections.hive.connect", new=hive_http_connect):