Add option to set identity when creating TestingConnectorSession

prestodb · Jun 4, 2024 · 3492b98 · 3492b98
1 parent 39e11a5
commit 3492b98
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 5 deletions.
diff --git a/presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraConnector.java b/presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraConnector.java
@@ -36,6 +36,7 @@
 import com.facebook.presto.spi.connector.ConnectorSplitManager;
 import com.facebook.presto.spi.connector.ConnectorSplitManager.SplitSchedulingContext;
 import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
+import com.facebook.presto.spi.security.ConnectorIdentity;
 import com.facebook.presto.testing.TestingConnectorContext;
 import com.facebook.presto.testing.TestingConnectorSession;
 import com.google.common.collect.ImmutableList;
@@ -80,6 +81,7 @@ public class TestCassandraConnector
     private static final Date DATE = new Date();
     private static final ConnectorSession SESSION = new TestingConnectorSession(
             "user",
+            new ConnectorIdentity("user", Optional.empty(), Optional.empty()),
             Optional.of("test"),
             Optional.empty(),
             UTC_KEY,

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
@@ -65,26 +65,32 @@ public class TestingConnectorSession
 
     public TestingConnectorSession(List<PropertyMetadata<?>> properties)
     {
-        this("user", Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, ImmutableMap.of(), new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), ImmutableSet.of(), Optional.empty(), ImmutableMap.of());
+        this("user", new ConnectorIdentity("user", Optional.empty(), Optional.empty()), Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, ImmutableMap.of(), new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), ImmutableSet.of(), Optional.empty(), ImmutableMap.of());
+    }
+
+    public TestingConnectorSession(ConnectorIdentity identity, List<PropertyMetadata<?>> properties)
+    {
+        this(identity.getUser(), identity, Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, ImmutableMap.of(), new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), ImmutableSet.of(), Optional.empty(), ImmutableMap.of());
     }
 
     public TestingConnectorSession(List<PropertyMetadata<?>> properties, Set<String> clientTags)
     {
-        this("user", Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, ImmutableMap.of(), new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), clientTags, Optional.empty(), ImmutableMap.of());
+        this("user", new ConnectorIdentity("user", Optional.empty(), Optional.empty()), Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, ImmutableMap.of(), new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), clientTags, Optional.empty(), ImmutableMap.of());
     }
 
     public TestingConnectorSession(List<PropertyMetadata<?>> properties, Map<String, Object> propertyValues)
     {
-        this("user", Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, propertyValues, new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), ImmutableSet.of(), Optional.empty(), ImmutableMap.of());
+        this("user", new ConnectorIdentity("user", Optional.empty(), Optional.empty()), Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, propertyValues, new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), ImmutableSet.of(), Optional.empty(), ImmutableMap.of());
     }
 
     public TestingConnectorSession(List<PropertyMetadata<?>> properties, Optional<String> schema)
     {
-        this("user", Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, ImmutableMap.of(), new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), ImmutableSet.of(), schema, ImmutableMap.of());
+        this("user", new ConnectorIdentity("user", Optional.empty(), Optional.empty()), Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, ImmutableMap.of(), new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), ImmutableSet.of(), schema, ImmutableMap.of());
     }
 
     public TestingConnectorSession(
             String user,
+            ConnectorIdentity identity,
             Optional<String> source,
             Optional<String> traceToken,
             TimeZoneKey timeZoneKey,
@@ -99,7 +105,7 @@ public TestingConnectorSession(
             Map<SqlFunctionId, SqlInvokedFunction> sessionFunctions)
     {
         this.queryId = queryIdGenerator.createNextQueryId().toString();
-        this.identity = new ConnectorIdentity(requireNonNull(user, "user is null"), Optional.empty(), Optional.empty());
+        this.identity = requireNonNull(identity, "identity is null");
         this.source = requireNonNull(source, "source is null");
         this.traceToken = requireNonNull(traceToken, "traceToken is null");
         this.timeZoneKey = requireNonNull(timeZoneKey, "timeZoneKey is null");

diff --git a/presto-main/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java b/presto-main/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java
@@ -26,6 +26,7 @@
 import com.facebook.presto.common.type.TimestampType;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.spi.StandardErrorCode;
+import com.facebook.presto.spi.security.ConnectorIdentity;
 import com.facebook.presto.testing.TestingConnectorSession;
 import com.facebook.presto.testing.TestingSession;
 import com.facebook.presto.type.SqlIntervalDayTime;
@@ -165,6 +166,7 @@ private void assertCurrentDateAtInstant(TimeZoneKey timeZoneKey, long instant)
         long dateTimeCalculation = currentDate(
                 new TestingConnectorSession(
                         "test",
+                        new ConnectorIdentity("test", Optional.empty(), Optional.empty()),
                         Optional.empty(),
                         Optional.empty(),
                         timeZoneKey,

diff --git a/presto-pinot-toolkit/src/test/java/com/facebook/presto/pinot/TestPinotSplitManager.java b/presto-pinot-toolkit/src/test/java/com/facebook/presto/pinot/TestPinotSplitManager.java
@@ -18,6 +18,7 @@
 import com.facebook.presto.spi.ConnectorSplitSource;
 import com.facebook.presto.spi.plan.AggregationNode;
 import com.facebook.presto.spi.plan.PlanNode;
+import com.facebook.presto.spi.security.ConnectorIdentity;
 import com.facebook.presto.sql.analyzer.FeaturesConfig;
 import com.facebook.presto.sql.planner.iterative.rule.test.PlanBuilder;
 import com.facebook.presto.testing.TestingConnectorSession;
@@ -224,6 +225,7 @@ public static ConnectorSession createSessionWithNumSplits(int numSegmentsPerSpli
     {
         return new TestingConnectorSession(
                 "user",
+                new ConnectorIdentity("user", Optional.empty(), Optional.empty()),
                 Optional.of("test"),
                 Optional.empty(),
                 UTC_KEY,
@@ -246,6 +248,7 @@ public static ConnectorSession createSessionWithLimitLarge(int limitLarge, Pinot
     {
         return new TestingConnectorSession(
                 "user",
+                new ConnectorIdentity("user", Optional.empty(), Optional.empty()),
                 Optional.of("test"),
                 Optional.empty(),
                 UTC_KEY,
@@ -266,6 +269,7 @@ public static ConnectorSession createSessionWithTopNLarge(int topNLarge, PinotCo
     {
         return new TestingConnectorSession(
                 "user",
+                new ConnectorIdentity("user", Optional.empty(), Optional.empty()),
                 Optional.of("test"),
                 Optional.empty(),
                 UTC_KEY,