changes for planner stickiness

Author: yashlimbad

e6data_python_connector/e6data_grpc.py

@@ -188,6 +188,8 @@ def __init__(
         self.cluster_uuid = cluster_uuid
         self._session_id = None
         self._host = host
+        # engine ip for stickiness
+        self._engine_ip = None
         self._port = port
 
         self._secure_channel = secure
@@ -322,6 +324,7 @@ def get_session_id(self):
                                 metadata=_get_grpc_header(cluster=self.cluster_uuid)
                             )
                             self._session_id = authenticate_response.sessionId
+                            self._engine_ip = authenticate_response.engineIP
                         else:
                             raise e
                     else:
@@ -534,25 +537,23 @@ 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) -> "DataFrame":
-        dataframe = DataFrame(
-            self,
-            file_path=parquet_path,
-            dataframe_number=self._dataframe_session.get_dataframe_number,
-            table_name=None
-        )
+        dataframe = DataFrame(self,
+                              file_path=parquet_path,
+                              dataframe_number=self._dataframe_session.get_dataframe_number,
+                              table_name=None,
+                              engine_ip=self._engine_ip)
 
         self._dataframe_session.update_dataframe_map(dataframe=dataframe)
         return dataframe
 
     def load_table(self, table_name, database = None, catalog = None) -> "DataFrame":
-        dataframe = DataFrame(
-            self,
-            file_path=None,
-            dataframe_number=self._dataframe_session.get_dataframe_number,
-            table_name=table_name,
-            database=database,
-            catalog=catalog
-        )
+        dataframe = DataFrame(self,
+                              file_path=None,
+                              dataframe_number=self._dataframe_session.get_dataframe_number,
+                              table_name=table_name,
+                              engine_ip=self._engine_ip,
+                              database=database,
+                              catalog=catalog)
 
         self._dataframe_session.update_dataframe_map(dataframe=dataframe)
         return dataframe
@@ -1062,15 +1063,15 @@ def explain_analyse(self):
 
 class DataFrame:
 
-    def __init__(self, connection: Connection, file_path, dataframe_number, table_name, database = None, catalog = None):
+    def __init__(self, connection: Connection, file_path, dataframe_number, table_name, engine_ip, database = None, catalog = None):
         self._dataframe_number = dataframe_number
         self._connection = connection
         self._catalog = self._connection.catalog_name if catalog is None else catalog
         self._database = self._connection.database if database is None else database
         self._table_name = table_name
         self._file_path = file_path
-        self._engine_ip = connection.host
         self._sessionId = connection.get_session_id
+        self._engine_ip = engine_ip
         self._is_metadata_updated = False
         self._query_id = None
         self._data = None
@@ -1106,6 +1107,7 @@ def select(self, *fields : str) -> "DataFrame":
             queryId=self._query_id,
             dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId,
+            engineIP=self._engine_ip,
             field=projection_fields
         )
 
@@ -1139,6 +1141,7 @@ def get_agg_enum(function_name : str) -> AggregateFunction | None:
             queryId=self._query_id,
             dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId,
+            engineIP=self._engine_ip,
             aggregateFunctionMap=agg_function_map,
             groupBy=group_by
         )
@@ -1153,6 +1156,7 @@ def where(self, where_clause : str) -> "DataFrame":
             queryId=self._query_id,
             dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId,
+            engineIP=self._engine_ip,
             whereClause=where_clause
         )
 
@@ -1173,6 +1177,7 @@ def order_by(self, *field_list : str) -> "DataFrame":
             queryId=self._query_id,
             dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId,
+            engineIP=self._engine_ip,
             orderByFieldMap=order_by_map
         )
 
@@ -1185,6 +1190,7 @@ def limit(self, fetch_limit : int) -> "DataFrame":
             queryId=self._query_id,
             dataframeNumber=self._dataframe_number,
             sessionId=self._sessionId,
+            engineIP=self._engine_ip,
             fetchLimit=fetch_limit
         )
 
@@ -1201,7 +1207,8 @@ def execute(self):
         execute_dataframe_request = e6x_engine_pb2.ExecuteDataFrameRequest(
             queryId=self._query_id,
             dataframeNumber=self._dataframe_number,
-            sessionId=self._sessionId
+            sessionId=self._sessionId,
+            engineIP=self._engine_ip,
         )
         client.executeDataFrame(execute_dataframe_request)
 
@@ -1248,11 +1255,13 @@ def fetchall(self):
         return rows
 
 class DataFrameSession:
-    def __init__(self, connection: Connection):
+    def __init__(self, connection: Connection, planner_ip):
         self._connection = connection
         self._dataframe_count = 0
         self._dataframe_map = dict()
         self._is_terminated = False
+        self._session_id = connection.get_session_id
+        self._planner_ip = planner_ip
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.terminate()
@@ -1269,10 +1278,15 @@ def get_dataframe_number(self) -> int:
     def is_terminated(self) -> bool:
         return self._is_terminated
 
+    @property
+    def planner_ip(self):
+        return self._planner_ip
+
     def terminate(self):
         if not self._is_terminated:
             drop_user_context_request = e6x_engine_pb2.DropUserContextRequest(
-                sessionId=self._connection.get_session_id
+                sessionId=self._session_id,
+                engineIP=self._planner_ip
             )
 
             self._connection.client.dropUserContext(drop_user_context_request)