Set custom cookies received from the server

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse Whitehouse

src/databricks/sql/auth/thrift_http_client.py

@@ -167,7 +167,11 @@ def flush(self):
         # Get reply to flush the request
         self.code = self.__resp.status
         self.message = self.__resp.reason
-        self.headers = self.__resp.msg
+        self.headers = self.__resp.headers
+
+        # Saves the cookie sent by the server response
+        if "Set-Cookie" in self.headers:
+            self.setCustomHeaders(dict("Cookie", self.headers["Set-Cookie"]))
 
     @staticmethod
     def basic_proxy_auth_header(proxy):