set httpfs config in a connection once time

Canner · Jan 29, 2024 · f62127a · f62127a
1 parent 3410e33
commit f62127a
Show file tree

Hide file tree

Showing 10 changed files with 26 additions and 14 deletions.
diff --git a/...va/io/accio/cache/CacheStorageConfig.java → ...ase/client/duckdb/CacheStorageConfig.java b/...va/io/accio/cache/CacheStorageConfig.java → ...ase/client/duckdb/CacheStorageConfig.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package io.accio.cache;
+package io.accio.base.client.duckdb;
 
 public interface CacheStorageConfig
 {

diff --git a/accio-base/src/main/java/io/accio/base/client/duckdb/DuckDBDataSource.java b/accio-base/src/main/java/io/accio/base/client/duckdb/DuckDBDataSource.java
@@ -27,15 +27,19 @@
 import java.sql.SQLException;
 import java.sql.Statement;
 
+import static java.lang.String.format;
+
 public class DuckDBDataSource
         extends BaseDataSource
         implements DataSource, Serializable
 {
     private final DuckDBConnection duckDBConnection;
+    private final DuckdbS3StyleStorageConfig duckdbS3StyleStorageConfig;
 
-    public DuckDBDataSource(DuckDBConnection duckDBConnection)
+    public DuckDBDataSource(DuckDBConnection duckDBConnection, DuckdbS3StyleStorageConfig duckdbS3StyleStorageConfifg)
     {
         this.duckDBConnection = duckDBConnection;
+        this.duckdbS3StyleStorageConfig = duckdbS3StyleStorageConfifg;
     }
 
     @Override
@@ -51,11 +55,14 @@ public Connection getConnection()
         // Refer to the official doc, if we want to create multiple read-write connections,
         // to the same database in-memory database instance, we can use the custom `duplicate()` method.
         // https://duckdb.org/docs/api/java
-        Connection connection = ((DuckDBConnection) duckDBConnection).duplicate();
+        Connection connection = duckDBConnection.duplicate();
         Statement statement = connection.createStatement();
         statement.execute("set search_path = 'main'");
         // install extensions from stable repository
         statement.execute("SET custom_extension_repository = 'http://extensions.duckdb.org'");
+        // init httpfs settings
+        statement.execute(format("SET s3_endpoint='%s'\n", duckdbS3StyleStorageConfig.getEndpoint()));
+        statement.execute(format("SET s3_url_style='%s'\n", duckdbS3StyleStorageConfig.getUrlStyle()));
         return connection;
     }
 

diff --git a/accio-base/src/main/java/io/accio/base/client/duckdb/DuckdbClient.java b/accio-base/src/main/java/io/accio/base/client/duckdb/DuckdbClient.java
@@ -51,15 +51,15 @@ public final class DuckdbClient
     private final HikariDataSource ds;
 
     @Inject
-    public DuckdbClient(DuckDBConfig duckDBConfig)
+    public DuckdbClient(DuckDBConfig duckDBConfig, DuckdbS3StyleStorageConfig duckdbS3StyleStorageConfig)
     {
         try {
             // The instance will be cleared after the process end. We don't need to
             // close this connection
             Class.forName("org.duckdb.DuckDBDriver");
             DuckDBConnection duckDBConnection = (DuckDBConnection) DriverManager.getConnection("jdbc:duckdb:");
             this.duckDBConfig = duckDBConfig;
-            DuckDBDataSource dataSource = new DuckDBDataSource(duckDBConnection);
+            DuckDBDataSource dataSource = new DuckDBDataSource(duckDBConnection, duckdbS3StyleStorageConfig);
             HikariConfig config = new HikariConfig();
             config.setDataSource(dataSource);
             config.setPoolName("MY_POOL");

diff --git a/...cio/cache/DuckdbS3StyleStorageConfig.java → ...nt/duckdb/DuckdbS3StyleStorageConfig.java b/...cio/cache/DuckdbS3StyleStorageConfig.java → ...nt/duckdb/DuckdbS3StyleStorageConfig.java
@@ -12,7 +12,7 @@
  * limitations under the License.
  */
 
-package io.accio.cache;
+package io.accio.base.client.duckdb;
 
 import io.airlift.configuration.Config;
 import io.airlift.configuration.ConfigDescription;
@@ -103,12 +103,10 @@ public String getUrlStyle()
     public String generateDuckdbParquetStatement(String path, String tableName)
     {
         // ref: https://github.com/duckdb/duckdb/issues/1403
-        StringBuilder sb = new StringBuilder("INSTALL httpfs;\n" +
-                "LOAD httpfs;\n");
-        sb.append(format("SET s3_endpoint='%s';\n", endpoint));
+        StringBuilder sb = new StringBuilder();
+        // TODO: check why can't we set s3 access key and secret key in Data source
         accessKey.ifPresent(accessKey -> sb.append(format("SET s3_access_key_id='%s';\n", accessKey)));
         secretKey.ifPresent(secretKey -> sb.append(format("SET s3_secret_access_key='%s';\n", secretKey)));
-        sb.append(format("SET s3_url_style='%s';\n", urlStyle));
         sb.append("BEGIN TRANSACTION;\n");
         sb.append(format("CREATE TABLE \"%s\" AS SELECT * FROM read_parquet('s3://%s');", tableName, path));
         sb.append("COMMIT;\n");

diff --git a/accio-base/src/test/java/io/accio/base/sqlrewrite/AbstractTestFramework.java b/accio-base/src/test/java/io/accio/base/sqlrewrite/AbstractTestFramework.java
@@ -19,6 +19,7 @@
 import io.accio.base.client.AutoCloseableIterator;
 import io.accio.base.client.duckdb.DuckDBConfig;
 import io.accio.base.client.duckdb.DuckdbClient;
+import io.accio.base.client.duckdb.DuckdbS3StyleStorageConfig;
 import io.accio.base.dto.Column;
 import io.accio.base.dto.Manifest;
 import io.accio.base.dto.Model;
@@ -65,7 +66,7 @@ public static Model addColumnsToModel(Model model, Column... columns)
     @BeforeClass
     public void init()
     {
-        duckdbClient = new DuckdbClient(new DuckDBConfig());
+        duckdbClient = new DuckdbClient(new DuckDBConfig(), new DuckdbS3StyleStorageConfig());
         prepareData();
     }
 

diff --git a/accio-cache/src/main/java/io/accio/cache/CacheManager.java b/accio-cache/src/main/java/io/accio/cache/CacheManager.java
@@ -23,6 +23,7 @@
 import io.accio.base.ConnectorRecordIterator;
 import io.accio.base.Parameter;
 import io.accio.base.SessionContext;
+import io.accio.base.client.duckdb.CacheStorageConfig;
 import io.accio.base.client.duckdb.DuckDBConfig;
 import io.accio.base.client.duckdb.DuckdbClient;
 import io.accio.base.dto.CacheInfo;

diff --git a/accio-cache/src/main/java/io/accio/cache/CacheModule.java b/accio-cache/src/main/java/io/accio/cache/CacheModule.java
@@ -16,8 +16,10 @@
 
 import com.google.inject.Binder;
 import com.google.inject.Scopes;
+import io.accio.base.client.duckdb.CacheStorageConfig;
 import io.accio.base.client.duckdb.DuckDBConfig;
 import io.accio.base.client.duckdb.DuckdbClient;
+import io.accio.base.client.duckdb.DuckdbS3StyleStorageConfig;
 import io.airlift.configuration.AbstractConfigurationAwareModule;
 
 import static io.airlift.configuration.ConfigBinder.configBinder;

diff --git a/accio-testing/src/main/java/io/accio/testing/AbstractTestFramework.java b/accio-testing/src/main/java/io/accio/testing/AbstractTestFramework.java
@@ -19,6 +19,7 @@
 import io.accio.base.client.AutoCloseableIterator;
 import io.accio.base.client.duckdb.DuckDBConfig;
 import io.accio.base.client.duckdb.DuckdbClient;
+import io.accio.base.client.duckdb.DuckdbS3StyleStorageConfig;
 import io.accio.base.dto.Column;
 import io.accio.base.dto.Manifest;
 import io.accio.base.dto.Model;
@@ -65,7 +66,7 @@ public static Model addColumnsToModel(Model model, Column... columns)
     @BeforeClass
     public void init()
     {
-        duckdbClient = new DuckdbClient(new DuckDBConfig());
+        duckdbClient = new DuckdbClient(new DuckDBConfig(), new DuckdbS3StyleStorageConfig());
         prepareData();
     }
 

diff --git a/accio-tests/src/test/java/io/accio/testing/bigquery/TestDuckdbTaskManager.java b/accio-tests/src/test/java/io/accio/testing/bigquery/TestDuckdbTaskManager.java
@@ -7,6 +7,7 @@
 import io.accio.base.CatalogSchemaTableName;
 import io.accio.base.client.duckdb.DuckDBConfig;
 import io.accio.base.client.duckdb.DuckdbClient;
+import io.accio.base.client.duckdb.DuckdbS3StyleStorageConfig;
 import io.accio.base.dto.Column;
 import io.accio.base.dto.Model;
 import io.accio.cache.CacheInfoPair;
@@ -78,7 +79,7 @@ public void testCheckCacheMemoryLimit()
 
         DuckDBConfig duckDBConfig = new DuckDBConfig();
         duckDBConfig.setMaxCacheTableSizeRatio(0);
-        try (DuckdbTaskManager taskManager = new DuckdbTaskManager(duckDBConfig, new DuckdbClient(duckDBConfig))) {
+        try (DuckdbTaskManager taskManager = new DuckdbTaskManager(duckDBConfig, new DuckdbClient(duckDBConfig, new DuckdbS3StyleStorageConfig()))) {
             assertThatCode(taskManager::checkCacheMemoryLimit).hasMessageMatching("Cache memory limit exceeded. Usage: .* bytes, Limit: 0.0 bytes");
         }
     }

diff --git a/accio-validation/src/test/java/io/accio/TestMetricValidation.java b/accio-validation/src/test/java/io/accio/TestMetricValidation.java
@@ -19,6 +19,7 @@
 import io.accio.base.client.Client;
 import io.accio.base.client.duckdb.DuckDBConfig;
 import io.accio.base.client.duckdb.DuckdbClient;
+import io.accio.base.client.duckdb.DuckdbS3StyleStorageConfig;
 import io.accio.base.dto.Column;
 import io.accio.base.dto.EnumDefinition;
 import io.accio.base.dto.JoinType;
@@ -51,7 +52,7 @@ public class TestMetricValidation
 
     public TestMetricValidation()
     {
-        client = new DuckdbClient(new DuckDBConfig());
+        client = new DuckdbClient(new DuckDBConfig(), new DuckdbS3StyleStorageConfig());
         sample = AccioMDL.fromManifest(withDefaultCatalogSchema()
                 .setModels(List.of(
                         Model.model("Flight",