Add forConnectorId method to ConnectorSession

Add a method forConnectorId to ConnectorSession to allow converting ConnectorSession of one connector to a ConnectorSession for another connector.
prestodb · Feb 18, 2025 · 4435387 · 4435387
1 parent d6b494a
commit 4435387
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 0 deletions.
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java b/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java
@@ -1235,6 +1235,12 @@ public RuntimeStats getRuntimeStats()
             {
                 return session.getRuntimeStats();
             }
+
+            @Override
+            public ConnectorSession forConnectorId(ConnectorId connectorId)
+            {
+                return this;
+            }
         };
     }
 

diff --git a/presto-main/src/main/java/com/facebook/presto/FullConnectorSession.java b/presto-main/src/main/java/com/facebook/presto/FullConnectorSession.java
@@ -199,4 +199,10 @@ public RuntimeStats getRuntimeStats()
     {
         return session.getRuntimeStats();
     }
+
+    @Override
+    public ConnectorSession forConnectorId(ConnectorId connectorId)
+    {
+        return new FullConnectorSession(session, identity);
+    }
 }
diff --git a/presto-main/src/main/java/com/facebook/presto/testing/TestingConnectorSession.java b/presto-main/src/main/java/com/facebook/presto/testing/TestingConnectorSession.java
@@ -17,6 +17,7 @@
 import com.facebook.presto.common.function.SqlFunctionProperties;
 import com.facebook.presto.common.type.TimeZoneKey;
 import com.facebook.presto.execution.QueryIdGenerator;
+import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.WarningCollector;
@@ -224,6 +225,26 @@ public RuntimeStats getRuntimeStats()
         return new RuntimeStats();
     }
 
+    @Override
+    public ConnectorSession forConnectorId(ConnectorId connectorId)
+    {
+        return new TestingConnectorSession(
+                sqlFunctionProperties.getSessionUser(),
+                identity,
+                source,
+                traceToken,
+                timeZoneKey,
+                locale,
+                startTime,
+                ImmutableList.copyOf(properties.values()),
+                propertyValues,
+                sqlFunctionProperties.isLegacyRowFieldOrdinalAccessEnabled(),
+                clientInfo,
+                clientTags,
+                schema,
+                sessionFunctions);
+    }
+
     @Override
     public String toString()
     {

diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorSession.java b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorSession.java
@@ -66,4 +66,10 @@ default boolean isReadConstraints()
     WarningCollector getWarningCollector();
 
     RuntimeStats getRuntimeStats();
+
+    /**
+     * returns a ConnectorSession for a specific ConnectorId
+     * @return
+     */
+    ConnectorSession forConnectorId(ConnectorId connectorId);
 }
diff --git a/presto-spi/src/test/java/com/facebook/presto/spi/TestingSession.java b/presto-spi/src/test/java/com/facebook/presto/spi/TestingSession.java
@@ -130,6 +130,12 @@ public RuntimeStats getRuntimeStats()
         {
             return new RuntimeStats();
         }
+
+        @Override
+        public ConnectorSession forConnectorId(ConnectorId connectorId)
+        {
+            return this;
+        }
     };
 
     private TestingSession() {}