Allow to inject IdentityCacheMapping

trinodb · Sep 2, 2024 · 6b97209 · 6b97209
1 parent 2236dd3
commit 6b97209
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 7 deletions.
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DefaultJdbcMetadataFactory.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DefaultJdbcMetadataFactory.java
@@ -29,28 +29,29 @@ public class DefaultJdbcMetadataFactory
     private final JdbcClient jdbcClient;
     private final TimestampTimeZoneDomain timestampTimeZoneDomain;
     private final Set<JdbcQueryEventListener> jdbcQueryEventListeners;
+    private final IdentityCacheMapping identityCacheMapping;
 
     @Inject
     public DefaultJdbcMetadataFactory(
             JdbcClient jdbcClient,
             TimestampTimeZoneDomain timestampTimeZoneDomain,
-            Set<JdbcQueryEventListener> jdbcQueryEventListeners)
+            Set<JdbcQueryEventListener> jdbcQueryEventListeners,
+            IdentityCacheMapping identityCacheMapping)
     {
         this.jdbcClient = requireNonNull(jdbcClient, "jdbcClient is null");
         this.timestampTimeZoneDomain = requireNonNull(timestampTimeZoneDomain, "timestampTimeZoneDomain is null");
         this.jdbcQueryEventListeners = ImmutableSet.copyOf(requireNonNull(jdbcQueryEventListeners, "queryEventListeners is null"));
+        this.identityCacheMapping = requireNonNull(identityCacheMapping, "identityCacheMapping is null");
     }
 
     @Override
     public JdbcMetadata create(JdbcTransactionHandle transaction)
     {
-        // Session stays the same per transaction, therefore session properties don't need to
-        // be a part of cache keys in CachingJdbcClient.
         return create(new CachingJdbcClient(
                 Ticker.systemTicker(),
                 jdbcClient,
                 Set.of(),
-                new SingletonIdentityCacheMapping(),
+                identityCacheMapping,
                 new Duration(1, DAYS),
                 new Duration(1, DAYS),
                 new Duration(1, DAYS),

diff --git a/...c/src/main/java/io/trino/plugin/jdbc/ExtraCredentialsBasedIdentityCacheMappingModule.java b/...c/src/main/java/io/trino/plugin/jdbc/ExtraCredentialsBasedIdentityCacheMappingModule.java
@@ -17,12 +17,14 @@
 import com.google.inject.Module;
 import com.google.inject.Scopes;
 
+import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
+
 public class ExtraCredentialsBasedIdentityCacheMappingModule
         implements Module
 {
     @Override
     public void configure(Binder binder)
     {
-        binder.bind(IdentityCacheMapping.class).to(ExtraCredentialsBasedIdentityCacheMapping.class).in(Scopes.SINGLETON);
+        newOptionalBinder(binder, IdentityCacheMapping.class).setBinding().to(ExtraCredentialsBasedIdentityCacheMapping.class).in(Scopes.SINGLETON);
     }
 }
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcModule.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcModule.java
@@ -60,6 +60,7 @@ public void setup(Binder binder)
         tablePropertiesProviderBinder(binder);
 
         newOptionalBinder(binder, JdbcMetadataFactory.class).setDefault().to(DefaultJdbcMetadataFactory.class).in(Scopes.SINGLETON);
+        newOptionalBinder(binder, IdentityCacheMapping.class).setDefault().to(SingletonIdentityCacheMapping.class).in(Scopes.SINGLETON);
         newOptionalBinder(binder, TimestampTimeZoneDomain.class).setDefault().toInstance(TimestampTimeZoneDomain.ANY);
         newOptionalBinder(binder, Key.get(ConnectorSplitManager.class, ForJdbcDynamicFiltering.class)).setDefault().to(JdbcSplitManager.class).in(Scopes.SINGLETON);
         newOptionalBinder(binder, ConnectorSplitManager.class).setDefault().to(JdbcDynamicFilteringSplitManager.class).in(Scopes.SINGLETON);

diff --git a/plugin/trino-ignite/src/main/java/io/trino/plugin/ignite/IgniteJdbcMetadataFactory.java b/plugin/trino-ignite/src/main/java/io/trino/plugin/ignite/IgniteJdbcMetadataFactory.java
@@ -16,6 +16,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.inject.Inject;
 import io.trino.plugin.jdbc.DefaultJdbcMetadataFactory;
+import io.trino.plugin.jdbc.IdentityCacheMapping;
 import io.trino.plugin.jdbc.JdbcClient;
 import io.trino.plugin.jdbc.JdbcMetadata;
 import io.trino.plugin.jdbc.JdbcQueryEventListener;
@@ -32,9 +33,9 @@ public class IgniteJdbcMetadataFactory
     private final Set<JdbcQueryEventListener> jdbcQueryEventListeners;
 
     @Inject
-    public IgniteJdbcMetadataFactory(JdbcClient jdbcClient, TimestampTimeZoneDomain timestampTimeZoneDomain, Set<JdbcQueryEventListener> jdbcQueryEventListeners)
+    public IgniteJdbcMetadataFactory(JdbcClient jdbcClient, TimestampTimeZoneDomain timestampTimeZoneDomain, Set<JdbcQueryEventListener> jdbcQueryEventListeners, IdentityCacheMapping identityCacheMapping)
     {
-        super(jdbcClient, timestampTimeZoneDomain, jdbcQueryEventListeners);
+        super(jdbcClient, timestampTimeZoneDomain, jdbcQueryEventListeners, identityCacheMapping);
         this.timestampTimeZoneDomain = requireNonNull(timestampTimeZoneDomain, "timestampTimeZoneDomain is null");
         this.jdbcQueryEventListeners = ImmutableSet.copyOf(requireNonNull(jdbcQueryEventListeners, "jdbcQueryEventListeners is null"));
     }