revert change to connection_uuid

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/thrift_backend.py

@@ -72,9 +72,6 @@
     "_retry_delay_default": (float, 5, 1, 60),
 }
 
-# Add thread local storage
-_connection_uuid = threading.local()
-
 
 class ThriftBackend:
     CLOSED_OP_STATE = ttypes.TOperationState.CLOSED_STATE
@@ -226,7 +223,7 @@ def __init__(
             raise
 
         self._request_lock = threading.RLock()
-        _connection_uuid.value = None
+        self._connection_uuid = None
 
     # TODO: Move this bounding logic into DatabricksRetryPolicy for v3 (PECO-918)
     def _initialize_retry_args(self, kwargs):
@@ -259,14 +256,14 @@ def _initialize_retry_args(self, kwargs):
             )
 
     @staticmethod
-    def _check_response_for_error(response):
+    def _check_response_for_error(response, connection_uuid=None):
         if response.status and response.status.statusCode in [
             ttypes.TStatusCode.ERROR_STATUS,
             ttypes.TStatusCode.INVALID_HANDLE_STATUS,
         ]:
             raise DatabaseError(
                 response.status.errorMessage,
-                connection_uuid=getattr(_connection_uuid, "value", None),
+                connection_uuid=connection_uuid,
             )
 
     @staticmethod
@@ -320,7 +317,7 @@ def _handle_request_error(self, error_info, attempt, elapsed):
             network_request_error = RequestError(
                 user_friendly_error_message,
                 full_error_info_context,
-                getattr(_connection_uuid, "value", None),
+                self._connection_uuid,
                 error_info.error,
             )
             logger.info(network_request_error.message_with_context())
@@ -493,7 +490,7 @@ def attempt_request(attempt):
             if not isinstance(response_or_error_info, RequestErrorInfo):
                 # log nothing here, presume that main request logging covers
                 response = response_or_error_info
-                ThriftBackend._check_response_for_error(response)
+                ThriftBackend._check_response_for_error(response, self._connection_uuid)
                 return response
 
             error_info = response_or_error_info
@@ -508,7 +505,7 @@ def _check_protocol_version(self, t_open_session_resp):
                 "Error: expected server to use a protocol version >= "
                 "SPARK_CLI_SERVICE_PROTOCOL_V2, "
                 "instead got: {}".format(protocol_version),
-                connection_uuid=getattr(_connection_uuid, "value", None),
+                connection_uuid=self._connection_uuid,
             )
 
     def _check_initial_namespace(self, catalog, schema, response):
@@ -522,15 +519,15 @@ def _check_initial_namespace(self, catalog, schema, response):
             raise InvalidServerResponseError(
                 "Setting initial namespace not supported by the DBR version, "
                 "Please use a Databricks SQL endpoint or a cluster with DBR >= 9.0.",
-                connection_uuid=getattr(_connection_uuid, "value", None),
+                connection_uuid=self._connection_uuid,
             )
 
         if catalog:
             if not response.canUseMultipleCatalogs:
                 raise InvalidServerResponseError(
                     "Unexpected response from server: Trying to set initial catalog to {}, "
                     + "but server does not support multiple catalogs.".format(catalog),  # type: ignore
-                    connection_uuid=getattr(_connection_uuid, "value", None),
+                    connection_uuid=self._connection_uuid,
                 )
 
     def _check_session_configuration(self, session_configuration):
@@ -545,7 +542,7 @@ def _check_session_configuration(self, session_configuration):
                     TIMESTAMP_AS_STRING_CONFIG,
                     session_configuration[TIMESTAMP_AS_STRING_CONFIG],
                 ),
-                connection_uuid=getattr(_connection_uuid, "value", None),
+                connection_uuid=self._connection_uuid,
             )
 
     def open_session(self, session_configuration, catalog, schema):
@@ -576,7 +573,7 @@ def open_session(self, session_configuration, catalog, schema):
             response = self.make_request(self._client.OpenSession, open_session_req)
             self._check_initial_namespace(catalog, schema, response)
             self._check_protocol_version(response)
-            _connection_uuid.value = (
+            self._connection_uuid = (
                 self.handle_to_hex_id(response.sessionHandle)
                 if response.sessionHandle
                 else None
@@ -605,7 +602,7 @@ def _check_command_not_in_error_or_closed_state(
                         and self.guid_to_hex_id(op_handle.operationId.guid),
                         "diagnostic-info": get_operations_resp.diagnosticInfo,
                     },
-                    connection_uuid=getattr(_connection_uuid, "value", None),
+                    connection_uuid=self._connection_uuid,
                 )
             else:
                 raise ServerOperationError(
@@ -615,7 +612,7 @@ def _check_command_not_in_error_or_closed_state(
                         and self.guid_to_hex_id(op_handle.operationId.guid),
                         "diagnostic-info": None,
                     },
-                    connection_uuid=getattr(_connection_uuid, "value", None),
+                    connection_uuid=self._connection_uuid,
                 )
         elif get_operations_resp.operationState == ttypes.TOperationState.CLOSED_STATE:
             raise DatabaseError(
@@ -626,7 +623,7 @@ def _check_command_not_in_error_or_closed_state(
                     "operation-id": op_handle
                     and self.guid_to_hex_id(op_handle.operationId.guid)
                 },
-                connection_uuid=getattr(_connection_uuid, "value", None),
+                connection_uuid=self._connection_uuid,
             )
 
     def _poll_for_status(self, op_handle):
@@ -649,7 +646,7 @@ def _create_arrow_table(self, t_row_set, lz4_compressed, schema_bytes, descripti
         else:
             raise OperationalError(
                 "Unsupported TRowSet instance {}".format(t_row_set),
-                connection_uuid=getattr(_connection_uuid, "value", None),
+                connection_uuid=self._connection_uuid,
             )
         return convert_decimals_in_arrow_table(arrow_table, description), num_rows
 
@@ -658,7 +655,7 @@ def _get_metadata_resp(self, op_handle):
         return self.make_request(self._client.GetResultSetMetadata, req)
 
     @staticmethod
-    def _hive_schema_to_arrow_schema(t_table_schema):
+    def _hive_schema_to_arrow_schema(t_table_schema, connection_uuid=None):
         def map_type(t_type_entry):
             if t_type_entry.primitiveEntry:
                 return {
@@ -690,7 +687,7 @@ def map_type(t_type_entry):
                 # even for complex types
                 raise OperationalError(
                     "Thrift protocol error: t_type_entry not a primitiveEntry",
-                    connection_uuid=getattr(_connection_uuid, "value", None),
+                    connection_uuid=connection_uuid,
                 )
 
         def convert_col(t_column_desc):
@@ -701,7 +698,7 @@ def convert_col(t_column_desc):
         return pyarrow.schema([convert_col(col) for col in t_table_schema.columns])
 
     @staticmethod
-    def _col_to_description(col):
+    def _col_to_description(col, connection_uuid=None):
         type_entry = col.typeDesc.types[0]
 
         if type_entry.primitiveEntry:
@@ -711,7 +708,7 @@ def _col_to_description(col):
         else:
             raise OperationalError(
                 "Thrift protocol error: t_type_entry not a primitiveEntry",
-                connection_uuid=getattr(_connection_uuid, "value", None),
+                connection_uuid=connection_uuid,
             )
 
         if type_entry.primitiveEntry.type == ttypes.TTypeId.DECIMAL_TYPE:
@@ -725,17 +722,18 @@ def _col_to_description(col):
                 raise OperationalError(
                     "Decimal type did not provide typeQualifier precision, scale in "
                     "primitiveEntry {}".format(type_entry.primitiveEntry),
-                    connection_uuid=getattr(_connection_uuid, "value", None),
+                    connection_uuid=connection_uuid,
                 )
         else:
             precision, scale = None, None
 
         return col.columnName, cleaned_type, None, None, precision, scale, None
 
     @staticmethod
-    def _hive_schema_to_description(t_table_schema):
+    def _hive_schema_to_description(t_table_schema, connection_uuid=None):
         return [
-            ThriftBackend._col_to_description(col) for col in t_table_schema.columns
+            ThriftBackend._col_to_description(col, connection_uuid)
+            for col in t_table_schema.columns
         ]
 
     def _results_message_to_execute_response(self, resp, operation_state):
@@ -756,7 +754,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
                         t_result_set_metadata_resp.resultFormat
                     ]
                 ),
-                connection_uuid=getattr(_connection_uuid, "value", None),
+                connection_uuid=self._connection_uuid,
             )
         direct_results = resp.directResults
         has_been_closed_server_side = direct_results and direct_results.closeOperation
@@ -766,13 +764,16 @@ def _results_message_to_execute_response(self, resp, operation_state):
             or direct_results.resultSet.hasMoreRows
         )
         description = self._hive_schema_to_description(
-            t_result_set_metadata_resp.schema
+            t_result_set_metadata_resp.schema,
+            self._connection_uuid,
         )
 
         if pyarrow:
             schema_bytes = (
                 t_result_set_metadata_resp.arrowSchema
-                or self._hive_schema_to_arrow_schema(t_result_set_metadata_resp.schema)
+                or self._hive_schema_to_arrow_schema(
+                    t_result_set_metadata_resp.schema, self._connection_uuid
+                )
                 .serialize()
                 .to_pybytes()
             )
@@ -833,13 +834,16 @@ def get_execution_result(self, op_handle, cursor):
         is_staging_operation = t_result_set_metadata_resp.isStagingOperation
         has_more_rows = resp.hasMoreRows
         description = self._hive_schema_to_description(
-            t_result_set_metadata_resp.schema
+            t_result_set_metadata_resp.schema,
+            self._connection_uuid,
         )
 
         if pyarrow:
             schema_bytes = (
                 t_result_set_metadata_resp.arrowSchema
-                or self._hive_schema_to_arrow_schema(t_result_set_metadata_resp.schema)
+                or self._hive_schema_to_arrow_schema(
+                    t_result_set_metadata_resp.schema, self._connection_uuid
+                )
                 .serialize()
                 .to_pybytes()
             )
@@ -893,23 +897,27 @@ def get_query_state(self, op_handle) -> "TOperationState":
         return operation_state
 
     @staticmethod
-    def _check_direct_results_for_error(t_spark_direct_results):
+    def _check_direct_results_for_error(t_spark_direct_results, connection_uuid=None):
         if t_spark_direct_results:
             if t_spark_direct_results.operationStatus:
                 ThriftBackend._check_response_for_error(
-                    t_spark_direct_results.operationStatus
+                    t_spark_direct_results.operationStatus,
+                    connection_uuid,
                 )
             if t_spark_direct_results.resultSetMetadata:
                 ThriftBackend._check_response_for_error(
-                    t_spark_direct_results.resultSetMetadata
+                    t_spark_direct_results.resultSetMetadata,
+                    connection_uuid,
                 )
             if t_spark_direct_results.resultSet:
                 ThriftBackend._check_response_for_error(
-                    t_spark_direct_results.resultSet
+                    t_spark_direct_results.resultSet,
+                    connection_uuid,
                 )
             if t_spark_direct_results.closeOperation:
                 ThriftBackend._check_response_for_error(
-                    t_spark_direct_results.closeOperation
+                    t_spark_direct_results.closeOperation,
+                    connection_uuid,
                 )
 
     def execute_command(
@@ -1058,7 +1066,7 @@ def get_columns(
 
     def _handle_execute_response(self, resp, cursor):
         cursor.active_op_handle = resp.operationHandle
-        self._check_direct_results_for_error(resp.directResults)
+        self._check_direct_results_for_error(resp.directResults, self._connection_uuid)
 
         final_operation_state = self._wait_until_command_done(
             resp.operationHandle,
@@ -1069,7 +1077,7 @@ def _handle_execute_response(self, resp, cursor):
 
     def _handle_execute_response_async(self, resp, cursor):
         cursor.active_op_handle = resp.operationHandle
-        self._check_direct_results_for_error(resp.directResults)
+        self._check_direct_results_for_error(resp.directResults, self._connection_uuid)
 
     def fetch_results(
         self,
@@ -1104,7 +1112,7 @@ def fetch_results(
                 "fetch_results failed due to inconsistency in the state between the client and the server. Expected results to start from {} but they instead start at {}, some result batches must have been skipped".format(
                     expected_row_start_offset, resp.results.startRowOffset
                 ),
-                connection_uuid=getattr(_connection_uuid, "value", None),
+                connection_uuid=self._connection_uuid,
             )
 
         queue = ResultSetQueueFactory.build_queue(