Move JdbcRecordSetProvider construction to a module

trinodb · Jan 10, 2025 · 9d41cc9 · 9d41cc9
1 parent e208752
commit 9d41cc9
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
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
@@ -15,10 +15,14 @@
 
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.inject.Binder;
+import com.google.inject.Inject;
 import com.google.inject.Key;
 import com.google.inject.Provider;
 import com.google.inject.Scopes;
+import com.google.inject.Singleton;
 import com.google.inject.multibindings.Multibinder;
+import com.google.inject.multibindings.ProvidesIntoOptional;
+import dev.failsafe.RetryPolicy;
 import io.airlift.configuration.AbstractConfigurationAwareModule;
 import io.trino.plugin.base.mapping.IdentifierMappingModule;
 import io.trino.plugin.base.session.SessionPropertiesProvider;
@@ -29,6 +33,7 @@
 import io.trino.spi.connector.ConnectorAccessControl;
 import io.trino.spi.connector.ConnectorPageSinkProvider;
 import io.trino.spi.connector.ConnectorPageSourceProvider;
+import io.trino.spi.connector.ConnectorRecordSetProvider;
 import io.trino.spi.connector.ConnectorSplitManager;
 import io.trino.spi.function.table.ConnectorTableFunction;
 import io.trino.spi.procedure.Procedure;
@@ -37,6 +42,7 @@
 
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
+import static com.google.inject.multibindings.ProvidesIntoOptional.Type.DEFAULT;
 import static io.airlift.configuration.ConditionalModule.conditionalModule;
 import static io.airlift.configuration.ConfigBinder.configBinder;
 import static io.trino.plugin.base.ClosingBinder.closingBinder;
@@ -143,4 +149,12 @@ public static void bindTablePropertiesProvider(Binder binder, Class<? extends Ta
     {
         tablePropertiesProviderBinder(binder).addBinding().to(type).in(Scopes.SINGLETON);
     }
+
+    @ProvidesIntoOptional(DEFAULT)
+    @Inject
+    @Singleton
+    ConnectorRecordSetProvider recordSetProvider(JdbcClient jdbcClient, @ForRecordCursor ExecutorService executor, RetryPolicy<Object> policy)
+    {
+        return new JdbcRecordSetProvider(jdbcClient, executor, policy);
+    }
 }
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcPageSourceProvider.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcPageSourceProvider.java
@@ -15,11 +15,11 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
-import dev.failsafe.RetryPolicy;
 import io.trino.plugin.jdbc.MergeJdbcPageSource.ColumnAdaptation;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ConnectorPageSource;
 import io.trino.spi.connector.ConnectorPageSourceProvider;
+import io.trino.spi.connector.ConnectorRecordSetProvider;
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.connector.ConnectorSplit;
 import io.trino.spi.connector.ConnectorTableHandle;
@@ -29,7 +29,6 @@
 
 import java.util.List;
 import java.util.Optional;
-import java.util.concurrent.ExecutorService;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
@@ -43,13 +42,13 @@ public class JdbcPageSourceProvider
         implements ConnectorPageSourceProvider
 {
     private final JdbcClient jdbcClient;
-    private final JdbcRecordSetProvider recordSetProvider;
+    private final ConnectorRecordSetProvider recordSetProvider;
 
     @Inject
-    public JdbcPageSourceProvider(JdbcClient jdbcClient, @ForRecordCursor ExecutorService executor, RetryPolicy<Object> policy)
+    public JdbcPageSourceProvider(JdbcClient jdbcClient, ConnectorRecordSetProvider recordSetProvider)
     {
         this.jdbcClient = requireNonNull(jdbcClient, "jdbcClient is null");
-        this.recordSetProvider = new JdbcRecordSetProvider(jdbcClient, executor, policy);
+        this.recordSetProvider = requireNonNull(recordSetProvider, "recordSetProvider is null");
     }
 
     @Override