Merge remote-tracking branch 'origin/backend-interface' into fetch-interface

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -1166,6 +1166,11 @@ def _handle_execute_response_async(self, resp, cursor):
         cursor.active_command_id = command_id
         self._check_direct_results_for_error(resp.directResults)
 
+    def _handle_execute_response_async(self, resp, cursor):
+        command_id = CommandId.from_thrift_handle(resp.operationHandle)
+        cursor.active_command_id = command_id
+        self._check_direct_results_for_error(resp.directResults)
+
     def fetch_results(
         self,
         command_id: CommandId,

src/databricks/sql/backend/types.py

@@ -47,7 +47,7 @@ def guid_to_hex_id(guid: bytes) -> str:
     try:
         this_uuid = uuid.UUID(bytes=guid)
     except Exception as e:
-        logger.debug(f"Unable to convert bytes to UUID: {guid} -- {str(e)}")
+        logger.debug(f"Unable to convert bytes to UUID: {guid!r} -- {str(e)}")
         return str(guid)
     return str(this_uuid)
 

src/databricks/sql/client.py

@@ -232,7 +232,6 @@ def read(self) -> Optional[OAuthToken]:
         self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True)
         self._cursors = []  # type: List[Cursor]
 
-        # Create the session
         self.session = Session(
             server_hostname,
             http_path,
@@ -245,11 +244,6 @@ def read(self) -> Optional[OAuthToken]:
         )
         self.session.open()
 
-        logger.info(
-            "Successfully opened connection with session "
-            + str(self.get_session_id_hex())
-        )
-
         self.use_inline_params = self._set_use_inline_params_with_warning(
             kwargs.get("use_inline_params", False)
         )
@@ -788,6 +782,14 @@ def execute(
             async_op=False,
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
         )
+        self.active_result_set = ResultSet(
+            self.connection,
+            execute_response,
+            self.backend,
+            self.buffer_size_bytes,
+            self.arraysize,
+            self.connection.use_cloud_fetch,
+        )
 
         if self.active_result_set.is_staging_operation:
             self._handle_staging_operation(
@@ -853,6 +855,8 @@ def get_query_state(self) -> CommandState:
         :return:
         """
         self._check_not_closed()
+        if self.active_command_id is None:
+            raise Error("No active command to get state for")
         return self.backend.get_query_state(self.active_command_id)
 
     def is_query_pending(self):

src/databricks/sql/session.py

@@ -84,6 +84,8 @@ def __init__(
             **kwargs,
         )
 
+        self.protocol_version = None
+
     def open(self):
         self._session_id = self.backend.open_session(
             session_configuration=self.session_configuration,

tests/e2e/test_complex_types.py

@@ -14,7 +14,7 @@ def table_fixture(self, connection_details):
             # Create the table
             cursor.execute(
                 """
-                CREATE TABLE IF NOT EXISTS pysql_test_complex_types_table (
+                CREATE TABLE IF NOT EXISTS pysql_e2e_test_complex_types_table (
                     array_col ARRAY<STRING>,
                     map_col MAP<STRING, INTEGER>,
                     struct_col STRUCT<field1: STRING, field2: INTEGER>
@@ -24,7 +24,7 @@ def table_fixture(self, connection_details):
             # Insert a record
             cursor.execute(
                 """
-                INSERT INTO pysql_test_complex_types_table
+                INSERT INTO pysql_e2e_test_complex_types_table
                 VALUES (
                     ARRAY('a', 'b', 'c'),
                     MAP('a', 1, 'b', 2, 'c', 3),
@@ -34,7 +34,7 @@ def table_fixture(self, connection_details):
             )
             yield
             # Clean up the table after the test
-            cursor.execute("DROP TABLE IF EXISTS pysql_test_complex_types_table")
+            cursor.execute("DROP TABLE IF EXISTS pysql_e2e_test_complex_types_table")
 
     @pytest.mark.parametrize(
         "field,expected_type",
@@ -45,7 +45,7 @@ def test_read_complex_types_as_arrow(self, field, expected_type, table_fixture):
 
         with self.cursor() as cursor:
             result = cursor.execute(
-                "SELECT * FROM pysql_test_complex_types_table LIMIT 1"
+                "SELECT * FROM pysql_e2e_test_complex_types_table LIMIT 1"
             ).fetchone()
 
         assert isinstance(result[field], expected_type)
@@ -57,7 +57,7 @@ def test_read_complex_types_as_string(self, field, table_fixture):
             extra_params={"_use_arrow_native_complex_types": False}
         ) as cursor:
             result = cursor.execute(
-                "SELECT * FROM pysql_test_complex_types_table LIMIT 1"
+                "SELECT * FROM pysql_e2e_test_complex_types_table LIMIT 1"
             ).fetchone()
 
         assert isinstance(result[field], str)

tests/unit/test_client.py

@@ -102,8 +102,10 @@ def test_closing_connection_closes_commands(self, mock_result_set_class):
                 # Close the connection
                 connection.close()
 
-                # After closing the connection, the close method should have been called on the result set
-                mock_result_set.close.assert_called_once_with()
+                self.assertTrue(
+                    mock_result_set_class.return_value.has_been_closed_server_side
+                )
+                mock_result_set_class.return_value.close.assert_called_once_with()
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_cant_open_cursor_on_closed_connection(self, mock_client_class):

tests/unit/test_thrift_backend.py

@@ -178,7 +178,9 @@ def test_okay_protocol_versions_succeed(self, tcli_service_client_cass):
 
         for protocol_version in good_protocol_versions:
             t_http_client_instance.OpenSession.return_value = ttypes.TOpenSessionResp(
-                status=self.okay_status, serverProtocolVersion=protocol_version
+                status=self.okay_status,
+                serverProtocolVersion=protocol_version,
+                sessionHandle=self.session_handle,
             )
 
             thrift_backend = self._make_fake_thrift_backend()
@@ -2079,6 +2081,7 @@ def _construct_open_session_with_namespace(self, can_use_multiple_cats, cat, sch
             serverProtocolVersion=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4,
             canUseMultipleCatalogs=can_use_multiple_cats,
             initialNamespace=ttypes.TNamespace(catalogName=cat, schemaName=schem),
+            sessionHandle=self.session_handle,
         )
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
@@ -2178,6 +2181,7 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
             serverProtocolVersion=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V3,
             canUseMultipleCatalogs=True,
             initialNamespace=ttypes.TNamespace(catalogName="cat", schemaName="schem"),
+            sessionHandle=self.session_handle,
         )
 
         backend = ThriftDatabricksClient(