Fix _use_cert_as_auth_flag

The `authorization_header` wasn't set

Author: Niall Egan <niall.egan@databricks.com>

Author: Niall Egan

cmdexec/clients/python/src/databricks/sql/client.py

@@ -95,13 +95,14 @@ def __init__(self,
         self.host = server_hostname
         self.port = kwargs.get("_port", 443)
 
+        authorization_header = []
         if kwargs.get("_username") and kwargs.get("_password"):
             auth_credentials = "{username}:{password}".format(
                 username=kwargs.get("_username"), password=kwargs.get("_password")).encode("UTF-8")
             auth_credentials_base64 = base64.standard_b64encode(auth_credentials).decode("UTF-8")
-            authorization_header = "Basic {}".format(auth_credentials_base64)
+            authorization_header = [("Authorization", "Basic {}".format(auth_credentials_base64))]
         elif access_token:
-            authorization_header = "Bearer {}".format(access_token)
+            authorization_header = [("Authorization", "Bearer {}".format(access_token))]
         elif not (kwargs.get("_use_cert_as_auth") and kwargs.get("_tls_client_cert_file")):
             raise ValueError("No valid authentication settings. Please provide an access token.")
 
@@ -111,9 +112,8 @@ def __init__(self,
             useragent_header = "{}/{} ({})".format(USER_AGENT_NAME, __version__,
                                                    kwargs.get("_user_agent_entry"))
 
-        base_headers = [("Authorization", authorization_header),
-                        ("X-Databricks-Sqlgateway-CommandService-Mode", "grpc-thrift"),
-                        ("User-Agent", useragent_header)]
+        base_headers = [("X-Databricks-Sqlgateway-CommandService-Mode", "grpc-thrift"),
+                        ("User-Agent", useragent_header)] + authorization_header
 
         if not kwargs.get("_skip_routing_headers"):
             base_headers.append(self._http_path_to_routing_header(http_path))

cmdexec/clients/python/tests/tests.py

@@ -39,6 +39,46 @@ def test_close_uses_the_correct_session_id(self, mock_client_class):
         _, close_session_request = instance.make_request.call_args[0]
         self.assertEqual(close_session_request.session_id, mock_response.session_id)
 
+    @patch("%s.client.CmdExecBaseHttpClient" % PACKAGE_NAME)
+    def test_auth_args(self, mock_client_class):
+        instance = mock_client_class.return_value
+        mock_response = MagicMock()
+        mock_response.session_id = b'\x22'
+        instance.make_request.return_value = mock_response
+
+        # Test that the following auth args work:
+        # token = foo,
+        # token = None, _username = foo, _password = bar
+        # token = None, _tls_client_cert_file = something, _use_cert_as_auth = True
+        connection_args = [
+            {
+                "server_hostname": "foo",
+                "http_path": None,
+                "access_token": "tok",
+                "_skip_routing_headers": True,
+            },
+            {
+                "server_hostname": "foo",
+                "http_path": None,
+                "_username": "foo",
+                "_password": "bar",
+                "access_token": None,
+                "_skip_routing_headers": True,
+            },
+            {
+                "server_hostname": "foo",
+                "http_path": None,
+                "_tls_client_cert_file": "something",
+                "_use_cert_as_auth": True,
+                "access_token": None,
+                "_skip_routing_headers": True,
+            },
+        ]
+
+        for args in connection_args:
+            connection = databricks.sql.connect(**args)
+            connection.close()
+
     @patch("%s.client.CmdExecBaseHttpClient" % PACKAGE_NAME)
     @patch("%s.client.ResultSet" % PACKAGE_NAME)
     def test_closing_connection_closes_commands(self, mock_result_set_class, mock_client_class):