add session for dataframe

Author: yashlimbad

e6data_python_connector/e6data_grpc.py

@@ -16,6 +16,7 @@
 from io import BytesIO
 from ssl import CERT_NONE, CERT_OPTIONAL, CERT_REQUIRED
 from typing import overload
+import uuid
 
 import grpc
 from grpc._channel import _InactiveRpcError
@@ -201,6 +202,9 @@ def __init__(
         self.grpc_prepare_timeout = self._grpc_options.get('grpc_prepare_timeout') or 10 * 60  # 10 minutes
         self._create_client()
 
+        # initialize session for dataframe
+        self._dataframe_session = DataFrameSession(self)
+
     @property
     def _get_grpc_options(self):
         """
@@ -351,6 +355,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             exc_val (BaseException): The exception instance raised (if any).
             exc_tb (Traceback): The traceback object of the exception (if any).
         """
+        self._dataframe_session.terminate()
         self.close()
 
     def close(self):
@@ -363,6 +368,7 @@ def close(self):
             self._channel.close()
             self._channel = None
         self._session_id = None
+        self._dataframe_session.terminate()
 
     def check_connection(self):
         """
@@ -527,8 +533,15 @@ def cursor(self, catalog_name=None, db_name=None):
         """
         return Cursor(self, database=db_name, catalog_name=catalog_name)
 
-    def load_parquet(self, parquet_path):
-        return DataFrame(self, file_path=parquet_path)
+    def load_parquet(self, parquet_path) -> "DataFrame":
+        dataframe = DataFrame(
+            self,
+            file_path=parquet_path,
+            user_uuid=self._dataframe_session.get_user_uuid,
+            dataframe_number=self._dataframe_session.get_dataframe_number)
+
+        self._dataframe_session.update_dataframe_map(dataframe=dataframe)
+        return dataframe
 
     def rollback(self):
         """
@@ -1035,8 +1048,10 @@ def explain_analyse(self):
 
 class DataFrame:
 
-    def __init__(self, connection: Connection, file_path):
-        self.connection = connection
+    def __init__(self, connection: Connection, file_path, user_uuid, dataframe_number):
+        self._user_uuid = user_uuid
+        self._dataframe_number = dataframe_number
+        self._connection = connection
         self._file_path = file_path
         self._engine_ip = connection.host
         self._sessionId = connection.get_session_id
@@ -1046,21 +1061,17 @@ def __init__(self, connection: Connection, file_path):
         self._batch = None
         self._create_dataframe()
 
-    def __enter__(self):
-        pass
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
-
     def _create_dataframe(self):
-        client = self.connection.client
+        client = self._connection.client
 
         create_dataframe_request = e6x_engine_pb2.CreateDataFrameRequest(
             parquetFilePath=self._file_path,
-            catalog=self.connection.catalog_name,
-            schema=self.connection.database,
+            catalog=self._connection.catalog_name,
+            schema=self._connection.database,
             sessionId=self._sessionId,
-            engineIP=self._engine_ip
+            engineIP=self._engine_ip,
+            userUUID=self._user_uuid,
+            dataframeNumber=self._dataframe_number
         )
 
         create_dataframe_response = client.createDataFrame(
@@ -1073,9 +1084,11 @@ def select(self, *fields) -> "DataFrame":
         for field in fields:
             projection_fields.append(field)
 
-        client = self.connection.client
+        client = self._connection.client
         projection_on_dataframe_request = e6x_engine_pb2.ProjectionOnDataFrameRequest(
+            userUUID=self._user_uuid,
             queryId=self._query_id,
+            dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId,
             field=projection_fields
         )
@@ -1087,9 +1100,11 @@ def select(self, *fields) -> "DataFrame":
         return self
 
     def where(self, where_clause : str) -> "DataFrame":
-        client = self.connection.client
+        client = self._connection.client
         filter_on_dataframe_request = e6x_engine_pb2.FilterOnDataFrameRequest(
+            userUUID=self._user_uuid,
             queryId=self._query_id,
+            dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId,
             whereClause=where_clause
         )
@@ -1129,7 +1144,9 @@ def order_by(self, field_list : list, sort_direction_list = None, null_direction
         client = self.connection.client
 
         orderby_on_dataframe_request = e6x_engine_pb2.OrderByOnDataFrameRequest(
+            userUUID=self._user_uuid,
             queryId=self._query_id,
+            dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId,
             field=orderby_fields,
             sortDirection=sort_direction_request,
@@ -1148,10 +1165,12 @@ def order_by(self, *field_list) -> "DataFrame":
         for field in field_list:
             orderby_fields.append(field)
 
-        client = self.connection.client
+        client = self._connection.client
 
         orderby_on_dataframe_request = e6x_engine_pb2.OrderByOnDataFrameRequest(
+            userUUID=self._user_uuid,
             queryId=self._query_id,
+            dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId,
             field=orderby_fields,
             sortDirection=sort_direction_request,
@@ -1164,9 +1183,11 @@ def order_by(self, *field_list) -> "DataFrame":
         return self
 
     def limit(self, fetch_limit : int) -> "DataFrame":
-        client = self.connection.client
+        client = self._connection.client
         limit_on_dataframe_request = e6x_engine_pb2.LimitOnDataFrameRequest(
+            userUUID=self._user_uuid,
             queryId=self._query_id,
+            dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId,
             fetchLimit=fetch_limit
         )
@@ -1182,9 +1203,11 @@ def show(self):
         return self.fetchall()
 
     def execute(self):
-        client = self.connection.client
+        client = self._connection.client
         execute_dataframe_request = e6x_engine_pb2.ExecuteDataFrameRequest(
+            userUUID=self._user_uuid,
             queryId=self._query_id,
+            dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId
         )
         execute_dataframe_response = client.executeDataFrame(
@@ -1197,15 +1220,15 @@ def _update_meta_data(self):
             sessionId=self._sessionId,
             queryId=self._query_id
         )
-        get_result_metadata_response = self.connection.client.getResultMetadata(
+        get_result_metadata_response = self._connection.client.getResultMetadata(
             result_meta_data_request,
         )
         buffer = BytesIO(get_result_metadata_response.resultMetaData)
         self._rowcount, self._query_columns_description = get_query_columns_info(buffer)
         self._is_metadata_updated = True
 
     def _fetch_batch(self):
-        client = self.connection.client
+        client = self._connection.client
         get_next_result_batch_request = e6x_engine_pb2.GetNextResultBatchRequest(
             engineIP=self._engine_ip,
             sessionId=self._sessionId,
@@ -1233,6 +1256,44 @@ def fetchall(self):
         self._data = None
         return rows
 
+class DataFrameSession:
+    def __init__(self, connection: Connection):
+        self._user_uuid = str(uuid.uuid4())
+        self._connection = connection
+        self._dataframe_count = 0
+        self._dataframe_map = dict()
+        self._is_terminated = False
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.terminate()
+
+    def update_dataframe_map(self, dataframe : "DataFrame"):
+        self._dataframe_map.update({self._dataframe_count, dataframe})
+        self._dataframe_count = self._dataframe_count + 1
+
+    @property
+    def get_user_uuid(self):
+        return self._user_uuid
+
+    @property
+    def get_dataframe_number(self) -> int:
+        return self._dataframe_count
+
+    @property
+    def is_terminated(self) -> bool:
+        return self._is_terminated
+
+    def terminate(self):
+
+        if not self._is_terminated:
+            drop_user_context_request = e6x_engine_pb2.DropUserContextRequest(
+                userUUID=self.get_user_uuid
+            )
+
+            drop_user_context_response = self._connection.client.dropUserContext(drop_user_context_request)
+            self._is_terminated = True
+
+
 
 def poll(self, get_progress_update=True):
     """Poll for and return the raw status data provided by the Hive Thrift REST API.