Add session to getFlightDescriptorForSchema and getFlightDescriptorForTableScan

Thanzeel-Hassan-IBM · tdcmeehan · commit b53d94dfe185 · 2025-06-02T09:27:39.000-04:00
diff --git a/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java b/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java
@@ -131,9 +131,9 @@ public Schema getSchema(ConnectorSession connectorSession, FlightDescriptor flig
 
     public abstract List<SchemaTableName> listTables(ConnectorSession session, Optional<String> schemaName);
 
-    protected abstract FlightDescriptor getFlightDescriptorForSchema(String schemaName, String tableName);
+    protected abstract FlightDescriptor getFlightDescriptorForSchema(ConnectorSession session, String schemaName, String tableName);
 
-    protected abstract FlightDescriptor getFlightDescriptorForTableScan(ArrowTableLayoutHandle tableLayoutHandle);
+    protected abstract FlightDescriptor getFlightDescriptorForTableScan(ConnectorSession session, ArrowTableLayoutHandle tableLayoutHandle);
 
     public Schema getSchemaForTable(ConnectorSession connectorSession, String schemaName, String tableName)
     {
diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/testingConnector/TestingArrowFlightClientHandler.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/testingConnector/TestingArrowFlightClientHandler.java
@@ -74,7 +74,7 @@ public CallOption[] getCallOptions(ConnectorSession connectorSession)
     }
 
     @Override
-    public FlightDescriptor getFlightDescriptorForSchema(String schemaName, String tableName)
+    public FlightDescriptor getFlightDescriptorForSchema(ConnectorSession session, String schemaName, String tableName)
     {
         TestingArrowFlightRequest request = TestingArrowFlightRequest.createDescribeTableRequest(schemaName, tableName);
         return FlightDescriptor.command(requestCodec.toBytes(request));
@@ -139,7 +139,7 @@ public List<SchemaTableName> listTables(ConnectorSession session, Optional<Strin
     }
 
     @Override
-    public FlightDescriptor getFlightDescriptorForTableScan(ArrowTableLayoutHandle tableLayoutHandle)
+    public FlightDescriptor getFlightDescriptorForTableScan(ConnectorSession session, ArrowTableLayoutHandle tableLayoutHandle)
     {
         ArrowTableHandle tableHandle = tableLayoutHandle.getTable();
         String query = new TestingArrowQueryBuilder().buildSql(

Original file line number	Diff line number	Diff line change
`@@ -131,9 +131,9 @@ public Schema getSchema(ConnectorSession connectorSession, FlightDescriptor flig`
`131`	`131`
`132`	`132`	`public abstract List<SchemaTableName> listTables(ConnectorSession session, Optional<String> schemaName);`
`133`	`133`
`134`		`- protected abstract FlightDescriptor getFlightDescriptorForSchema(String schemaName, String tableName);`
	`134`	`+ protected abstract FlightDescriptor getFlightDescriptorForSchema(ConnectorSession session, String schemaName, String tableName);`
`135`	`135`
`136`		`- protected abstract FlightDescriptor getFlightDescriptorForTableScan(ArrowTableLayoutHandle tableLayoutHandle);`
	`136`	`+ protected abstract FlightDescriptor getFlightDescriptorForTableScan(ConnectorSession session, ArrowTableLayoutHandle tableLayoutHandle);`
`137`	`137`
`138`	`138`	`public Schema getSchemaForTable(ConnectorSession connectorSession, String schemaName, String tableName)`
`139`	`139`	`{`
Original file line number	Diff line number	Diff line change
`@@ -74,7 +74,7 @@ public CallOption[] getCallOptions(ConnectorSession connectorSession)`
`74`	`74`	`}`
`75`	`75`
`76`	`76`	`@Override`
`77`		`- public FlightDescriptor getFlightDescriptorForSchema(String schemaName, String tableName)`
	`77`	`+ public FlightDescriptor getFlightDescriptorForSchema(ConnectorSession session, String schemaName, String tableName)`
`78`	`78`	`{`
`79`	`79`	`TestingArrowFlightRequest request = TestingArrowFlightRequest.createDescribeTableRequest(schemaName, tableName);`
`80`	`80`	`return FlightDescriptor.command(requestCodec.toBytes(request));`
`@@ -139,7 +139,7 @@ public List<SchemaTableName> listTables(ConnectorSession session, Optional<Strin`
`139`	`139`	`}`
`140`	`140`
`141`	`141`	`@Override`
`142`		`- public FlightDescriptor getFlightDescriptorForTableScan(ArrowTableLayoutHandle tableLayoutHandle)`
	`142`	`+ public FlightDescriptor getFlightDescriptorForTableScan(ConnectorSession session, ArrowTableLayoutHandle tableLayoutHandle)`
`143`	`143`	`{`
`144`	`144`	`ArrowTableHandle tableHandle = tableLayoutHandle.getTable();`
`145`	`145`	`String query = new TestingArrowQueryBuilder().buildSql(`