Revert "Attempt to completely remove query_secret and send an empty UUID."

This reverts commit 1c58313.

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

Author: Jesse Whitehouse

src/databricks/sql/ae.py

@@ -87,6 +87,7 @@ def __init__(
         thrift_backend: "ThriftBackend",
         connection: "Connection",
         query_id: UUID,
+        query_secret: UUID,
         status: Optional[AsyncExecutionStatus] = AsyncExecutionStatus.UNKNOWN,
         execute_statement_response: Optional[
             Union[FakeExecuteStatementResponse, ttypes.TExecuteStatementResp]
@@ -95,6 +96,7 @@ def __init__(
         self._connection = connection
         self._thrift_backend = thrift_backend
         self.query_id = query_id
+        self.query_secret = query_secret
         self.status = status
 
         if execute_statement_response:
@@ -153,7 +155,7 @@ def _thrift_cancel_operation(self) -> None:
     def _thrift_get_operation_status(self) -> ttypes.TGetOperationStatusResp:
         """Execute TGetOperationStatusReq
 
-        Raises an AsyncExecutionError if the query_id is not found on the server.
+        Raises an AsyncExecutionError if the query_id:query_secret pair is not found on the server.
         """
         try:
             return self._thrift_backend._poll_for_status(self.t_operation_handle)
@@ -164,10 +166,10 @@ def _thrift_get_operation_status(self) -> ttypes.TGetOperationStatusResp:
                 ) from e
 
     def serialize(self) -> str:
-        """Return a hex string representing the query_id of this AsyncExecution.
+        """Return a string representing the query_id and secret of this AsyncExecution.
 
-        Use this to preserve a reference to the query_id"""
-        return f"{self.query_id}"
+        Use this to preserve a reference to the query_id and query_secret."""
+        return f"{self.query_id}:{self.query_secret}"
 
     def sync_status(self) -> None:
         """Synchronise the status of this AsyncExecution with the server query execution state."""
@@ -210,7 +212,7 @@ def t_operation_handle(self) -> ttypes.TOperationHandle:
 
         handle = ttypes.TOperationHandle(
             operationId=ttypes.THandleIdentifier(
-                guid=self.query_id.bytes, secret=UUID(int=0).bytes
+                guid=self.query_id.bytes, secret=self.query_secret.bytes
             ),
             operationType=ttypes.TOperationType.EXECUTE_STATEMENT,
             hasResultSet=True,
@@ -236,6 +238,7 @@ def from_thrift_response(
             connection=connection,
             thrift_backend=thrift_backend,
             query_id=UUID(bytes=resp.operationHandle.operationId.guid),
+            query_secret=UUID(bytes=resp.operationHandle.operationId.secret),
             status=_toperationstate_to_ae_status(
                 resp.directResults.operationStatus.operationState
             ),
@@ -248,17 +251,19 @@ def from_query_id_and_secret(
         connection: "Connection",
         thrift_backend: "ThriftBackend",
         query_id: UUID,
+        query_secret: UUID,
     ) -> "AsyncExecution":
-        """Return a valid AsyncExecution object from a query_id.
+        """Return a valid AsyncExecution object from a query_id and query_secret.
 
-        Raises an AsyncExecutionException if the query_id pair is not found on the server.
+        Raises an AsyncExecutionException if the query_id:query_secret pair is not found on the server.
         """
 
         # build a copy of this execution
         ae = cls(
             connection=connection,
             thrift_backend=thrift_backend,
             query_id=query_id,
+            query_secret=query_secret,
         )
         # check to make sure this is a valid one
         ae.sync_status()

src/databricks/sql/client.py

@@ -405,12 +405,13 @@ def execute_async(
         )
 
     def get_async_execution(
-        self, query_id: Union[str, UUID]
+        self, query_id: Union[str, UUID], query_secret: Union[str, UUID]
     ) -> "AsyncExecution":
         """Get an AsyncExecution object for an existing query.
 
         Args:
             query_id: The query id of the query to retrieve
+            query_secret: The query secret of the query to retrieve
 
         Returns:
             An AsyncExecution object that can be used to poll for status and retrieve results.
@@ -421,10 +422,16 @@ def get_async_execution(
         else:
             _qid = UUID(hex=query_id)
 
+        if isinstance(query_secret, UUID):
+            _qs = query_secret
+        else:
+            _qs = UUID(hex=query_secret)
+
         return AsyncExecution.from_query_id_and_secret(
             connection=self,
             thrift_backend=self.thrift_backend,
             query_id=_qid,
+            query_secret=_qs,
         )
 
 

tests/e2e/test_execute_async.py

@@ -82,24 +82,28 @@ def test_cant_get_results_after_cancel(self, long_running_ae: AsyncExecution):
     def test_get_async_execution_can_check_status(
         self, long_running_ae: AsyncExecution
     ):
-        query_id = long_running_ae.serialize()
+        query_id, query_secret = str(long_running_ae.query_id), str(
+            long_running_ae.query_secret
+        )
 
         with self.connection() as conn:
-            ae = conn.get_async_execution(query_id)
+            ae = conn.get_async_execution(query_id, query_secret)
             assert ae.is_running
 
     def test_get_async_execution_can_cancel_across_threads(
         self, long_running_ae: AsyncExecution
     ):
-        query_id = long_running_ae.serialize()
+        query_id, query_secret = str(long_running_ae.query_id), str(
+            long_running_ae.query_secret
+        )
 
-        def cancel_query_in_separate_thread(query_id):
+        def cancel_query_in_separate_thread(query_id, query_secret):
             with self.connection() as conn:
-                ae = conn.get_async_execution(query_id)
+                ae = conn.get_async_execution(query_id, query_secret)
                 ae.cancel()
 
         threading.Thread(
-            target=cancel_query_in_separate_thread, args=(query_id)
+            target=cancel_query_in_separate_thread, args=(query_id, query_secret)
         ).start()
 
         time.sleep(5)
@@ -150,29 +154,30 @@ def test_get_async_execution_with_badly_formed_query_id(self):
                 ae = conn.get_async_execution("foo", "bar")
 
     def test_serialize(self, long_running_ae: AsyncExecution):
-        query_id = long_running_ae.serialize()
+        serialized = long_running_ae.serialize()
+        query_id, query_secret = serialized.split(":")
 
         with self.connection() as conn:
-            ae = conn.get_async_execution(query_id)
+            ae = conn.get_async_execution(query_id, query_secret)
             assert ae.is_running
 
     def test_get_async_execution_no_results_when_direct_results_were_sent(self):
         """It remains to be seen whether results can be fetched repeatedly from a "picked up" execution."""
 
         with self.connection() as conn:
             ae = conn.execute_async(DIRECT_RESULTS_QUERY, {"param": 1})
-            query_id = ae.serialize()
+            query_id, query_secret = ae.serialize().split(":")
             ae.get_result()
 
         with self.connection() as conn:
             with pytest.raises(AsyncExecutionException, match="Query not found"):
-                ae_late = conn.get_async_execution(query_id)
+                ae_late = conn.get_async_execution(query_id, query_secret)
 
     def test_get_async_execution_and_fetch_results(self, long_ish_ae: AsyncExecution):
-        query_id = long_ish_ae.serialize()
+        query_id, query_secret = long_ish_ae.serialize().split(":")
 
         with self.connection() as conn:
-            ae = conn.get_async_execution(query_id)
+            ae = conn.get_async_execution(query_id, query_secret)
 
             while ae.is_running:
                 time.sleep(1)
@@ -189,8 +194,8 @@ def test_get_async_execution_twice(self):
         with self.connection() as conn_1, self.connection() as conn_2:
             ae_1 = conn_1.execute_async(LONG_ISH_QUERY)
 
-            query_id = ae_1.serialize()
-            ae_2 = conn_2.get_async_execution(query_id)
+            query_id, query_secret = ae_1.serialize().split(":")
+            ae_2 = conn_2.get_async_execution(query_id, query_secret)
 
             while ae_1.is_running:
                 time.sleep(1)