PRE-2175 | Harshad | refresh sessionId/token from the engine response.

.gitignore

@@ -10,3 +10,4 @@ MANIFEST.in
 venv/
 .DS_Store
 gen-py/
+.mise.toml

e6data_python_connector/e6data_grpc.py

@@ -328,6 +328,13 @@ def get_session_id(self):
                 raise e
         return self._session_id
 
+    def _set_session_id(self, refreshed_session_id):
+        if refreshed_session_id:
+            self._session_id = refreshed_session_id
+
+    def set_session_id_from_response(self, response):
+        self._set_session_id(response.sessionId)
+
     def __enter__(self):
         """
         Enters the runtime context related to this object.
@@ -385,10 +392,11 @@ def clear(self, query_id, engine_ip=None):
             queryId=query_id,
             engineIP=engine_ip
         )
-        self._client.clear(
+        clear_response = self._client.clear(
             clear_request,
             metadata=_get_grpc_header(engine_ip=engine_ip, cluster=self.cluster_uuid)
         )
+        self.set_session_id_from_response(clear_response)
 
     def reopen(self):
         """
@@ -412,10 +420,11 @@ def query_cancel(self, engine_ip, query_id):
             sessionId=self.get_session_id,
             queryId=query_id
         )
-        self._client.cancelQuery(
+        cancel_query_response = self._client.cancelQuery(
             cancel_query_request,
             metadata=_get_grpc_header(engine_ip=engine_ip, cluster=self.cluster_uuid)
         )
+        self.set_session_id_from_response(cancel_query_response)
 
     def dry_run(self, query):
         """
@@ -458,6 +467,7 @@ def get_tables(self, catalog, database):
             get_table_request,
             metadata=_get_grpc_header(cluster=self.cluster_uuid)
         )
+        self.set_session_id_from_response(get_table_response)
         return list(get_table_response.tables)
 
     def get_columns(self, catalog, database, table):
@@ -482,6 +492,7 @@ def get_columns(self, catalog, database, table):
             get_columns_request,
             metadata=_get_grpc_header(cluster=self.cluster_uuid)
         )
+        self.set_session_id_from_response(get_columns_response)
         return [{'fieldName': row.fieldName, 'fieldType': row.fieldType} for row in get_columns_response.fieldInfo]
 
     def get_schema_names(self, catalog):
@@ -502,6 +513,7 @@ def get_schema_names(self, catalog):
             get_schema_request,
             metadata=_get_grpc_header(cluster=self.cluster_uuid)
         )
+        self.set_session_id_from_response(get_schema_response)
         return list(get_schema_response.schemas)
 
     def commit(self):
@@ -729,7 +741,9 @@ def clear(self, query_id=None):
             queryId=query_id,
             engineIP=self._engine_ip
         )
-        return self.connection.client.clearOrCancelQuery(clear_request, metadata=self.metadata)
+        clear_response = self.connection.client.clearOrCancelQuery(clear_request, metadata=self.metadata)
+        self.connection.set_session_id_from_response(clear_response)
+        return clear_response
 
     def cancel(self, query_id):
         """
@@ -755,7 +769,9 @@ def status(self, query_id):
             queryId=query_id,
             engineIP=self._engine_ip
         )
-        return self.connection.client.status(status_request, metadata=self.metadata)
+        status_response = self.connection.client.status(status_request, metadata=self.metadata)
+        self.connection.set_session_id_from_response(status_response)
+        return status_response
 
     @re_auth
     def execute(self, operation, parameters=None, **kwargs):
@@ -792,17 +808,21 @@ def execute(self, operation, parameters=None, **kwargs):
                 metadata=self.metadata
             )
 
+            self.connection.set_session_id_from_response(prepare_statement_response)
             self._query_id = prepare_statement_response.queryId
             self._engine_ip = prepare_statement_response.engineIP
+
             execute_statement_request = e6x_engine_pb2.ExecuteStatementRequest(
                 engineIP=self._engine_ip,
                 sessionId=self.connection.get_session_id,
                 queryId=self._query_id,
             )
-            client.executeStatement(
+            execute_statement_response = client.executeStatement(
                 execute_statement_request,
                 metadata=self.metadata
             )
+
+            self.connection.set_session_id_from_response(execute_statement_response)
         else:
             prepare_statement_request = e6x_engine_pb2.PrepareStatementV2Request(
                 sessionId=self.connection.get_session_id,
@@ -816,6 +836,7 @@ def execute(self, operation, parameters=None, **kwargs):
                 timeout=self.connection.grpc_prepare_timeout
             )
 
+            self.connection.set_session_id_from_response(prepare_statement_response)
             self._query_id = prepare_statement_response.queryId
             self._engine_ip = prepare_statement_response.engineIP
 
@@ -824,10 +845,12 @@ def execute(self, operation, parameters=None, **kwargs):
                 sessionId=self.connection.get_session_id,
                 queryId=self._query_id
             )
-            client.executeStatementV2(
+            execute_statement_response = client.executeStatementV2(
                 execute_statement_request,
                 metadata=self.metadata
             )
+
+            self.connection.set_session_id_from_response(execute_statement_response)
         self.update_mete_data()
         return self._query_id
 
@@ -856,6 +879,8 @@ def update_mete_data(self):
             metadata=self.metadata
         )
         buffer = BytesIO(get_result_metadata_response.resultMetaData)
+
+        self.connection.set_session_id_from_response(get_result_metadata_response)
         self._rowcount, self._query_columns_description = get_query_columns_info(buffer)
         self._is_metadata_updated = True
 
@@ -927,6 +952,9 @@ def fetch_batch(self):
             get_next_result_batch_request,
             metadata=self.metadata
         )
+
+        self.connection.set_session_id_from_response(get_next_result_batch_response)
+
         buffer = get_next_result_batch_response.resultBatch
         if not self._is_metadata_updated:
             self.update_mete_data()
@@ -999,6 +1027,8 @@ def explain(self):
             explain_request,
             metadata=self.metadata
         )
+
+        self.connection.set_session_id_from_response(explain_response)
         return explain_response.explain
 
     def explain_analyse(self):
@@ -1017,6 +1047,9 @@ def explain_analyse(self):
             explain_analyze_request,
             metadata=self.metadata
         )
+
+        self.connection.set_session_id_from_response(explain_analyze_response)
+
         return dict(
             is_cached=explain_analyze_response.isCached,
             parsing_time=explain_analyze_response.parsingTime,
@@ -1056,4 +1089,4 @@ class Error(Exception):
 
 for type_id in PRIMITIVE_TYPES:
     name = TypeId._VALUES_TO_NAMES[type_id]
-    setattr(sys.modules[__name__], name, DBAPITypeObject([name]))
+    setattr(sys.modules[__name__], name, DBAPITypeObject([name]))