Add DuckDBConfig to handle the system resource configurations (#412)

* Add DuckDBConfig to set related configs

* add an api to get the internal configs

* fix building error

* remove trash file

* address review comment

* throw AccioException instead of sql exception

accio-base/pom.xml

@@ -30,6 +30,11 @@
     </properties>
 
     <dependencies>
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>configuration</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.airlift</groupId>
             <artifactId>log</artifactId>
@@ -86,6 +91,11 @@
             <artifactId>netty-all</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>javax.inject</groupId>
+            <artifactId>javax.inject</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>joda-time</groupId>
             <artifactId>joda-time</artifactId>

accio-base/src/main/java/io/accio/base/client/duckdb/DuckDBConfig.java

@@ -0,0 +1,58 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package io.accio.base.client.duckdb;
+
+import io.airlift.configuration.Config;
+import io.airlift.units.DataSize;
+
+public class DuckDBConfig
+{
+    private DataSize memoryLimit = DataSize.of(Runtime.getRuntime().maxMemory() / 2, DataSize.Unit.BYTE);
+    private String homeDirectory;
+    private String tempDirectory = "/tmp/duck";
+
+    public DataSize getMemoryLimit()
+    {
+        return memoryLimit;
+    }
+
+    @Config("duckdb.memory-limit")
+    public void setMemoryLimit(DataSize memoryLimit)
+    {
+        this.memoryLimit = memoryLimit;
+    }
+
+    public String getHomeDirectory()
+    {
+        return homeDirectory;
+    }
+
+    @Config("duckdb.home-directory")
+    public void setHomeDirectory(String homeDirectory)
+    {
+        this.homeDirectory = homeDirectory;
+    }
+
+    public String getTempDirectory()
+    {
+        return tempDirectory;
+    }
+
+    @Config("duckdb.temp-directory")
+    public void setTempDirectory(String tempDirectory)
+    {
+        this.tempDirectory = tempDirectory;
+    }
+}

accio-base/src/main/java/io/accio/base/client/duckdb/DuckdbClient.java

@@ -22,8 +22,11 @@
 import io.accio.base.client.jdbc.JdbcRecordIterator;
 import io.accio.base.metadata.ColumnMetadata;
 import io.airlift.log.Logger;
+import io.airlift.units.DataSize;
 import org.duckdb.DuckDBConnection;
 
+import javax.inject.Inject;
+
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.PreparedStatement;
@@ -43,20 +46,37 @@ public final class DuckdbClient
 {
     private static final Logger LOG = Logger.get(DuckdbClient.class);
     private final Connection duckDBConnection;
+    private final DuckDBConfig duckDBConfig;
 
-    public DuckdbClient()
+    @Inject
+    public DuckdbClient(DuckDBConfig duckDBConfig)
     {
         try {
             // The instance will be cleared after the process end. We don't need to
             // close this connection
             Class.forName("org.duckdb.DuckDBDriver");
             this.duckDBConnection = DriverManager.getConnection("jdbc:duckdb:");
+            this.duckDBConfig = duckDBConfig;
+            init();
         }
         catch (SQLException | ClassNotFoundException e) {
             throw new RuntimeException(e);
         }
     }
 
+    private void init()
+    {
+        DataSize memoryLimit = duckDBConfig.getMemoryLimit();
+        executeDDL(format("SET memory_limit='%s'", memoryLimit.toBytesValueString()));
+        LOG.info("Set memory limit to %s", memoryLimit.toBytesValueString());
+        executeDDL(format("SET temp_directory='%s'", duckDBConfig.getTempDirectory()));
+        LOG.info("Set temp directory to %s", duckDBConfig.getTempDirectory());
+
+        // TODO: Known issue: https://github.com/duckdb/duckdb/issues/10062
+        // executeDDL(format("SET home_directory='%s'", duckDBConfig.getHomeDirectory()));
+        // LOG.info("Set home directory to %s", duckDBConfig.getHomeDirectory());
+    }
+
     @Override
     public AutoCloseableIterator<Object[]> query(String sql)
     {
@@ -172,4 +192,9 @@ public Connection createConnection()
 
         return ((DuckDBConnection) duckDBConnection).duplicate();
     }
+
+    public DuckDBConfig getDuckDBConfig()
+    {
+        return duckDBConfig;
+    }
 }

accio-cache/src/main/java/io/accio/cache/CacheManager.java

@@ -50,6 +50,7 @@
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.accio.base.CatalogSchemaTableName.catalogSchemaTableName;
+import static io.accio.base.metadata.StandardErrorCode.GENERIC_INTERNAL_ERROR;
 import static io.accio.base.metadata.StandardErrorCode.GENERIC_USER_ERROR;
 import static io.accio.cache.EventLogger.Level.ERROR;
 import static io.accio.cache.EventLogger.Level.INFO;
@@ -337,6 +338,18 @@ public void untilTaskDone(CatalogSchemaTableName name)
         Optional.ofNullable(tasks.get(name)).ifPresent(Task::waitUntilDone);
     }
 
+    public List<Object> getDuckDBSettings()
+    {
+        try {
+            ConnectorRecordIterator iter = query("SELECT * FROM duckdb_settings()", List.of());
+            return ImmutableList.copyOf(iter);
+        }
+        catch (Exception e) {
+            LOG.error(e, "Failed to get duckdb settings");
+            throw new AccioException(GENERIC_INTERNAL_ERROR, e);
+        }
+    }
+
     private class Task
     {
         private final TaskInfo taskInfo;

accio-cache/src/main/java/io/accio/cache/CacheModule.java

@@ -16,6 +16,7 @@
 
 import com.google.inject.Binder;
 import com.google.inject.Scopes;
+import io.accio.base.client.duckdb.DuckDBConfig;
 import io.accio.base.client.duckdb.DuckdbClient;
 import io.airlift.configuration.AbstractConfigurationAwareModule;
 
@@ -28,6 +29,7 @@ public class CacheModule
     protected void setup(Binder binder)
     {
         configBinder(binder).bindConfig(DuckdbS3StyleStorageConfig.class);
+        configBinder(binder).bindConfig(DuckDBConfig.class);
         binder.bind(CacheStorageConfig.class).to(DuckdbS3StyleStorageConfig.class).in(Scopes.SINGLETON);
         binder.bind(CacheManager.class).in(Scopes.SINGLETON);
         binder.bind(EventLogger.class).to(Log4jEventLogger.class).in(Scopes.SINGLETON);

accio-main/src/main/java/io/accio/main/web/CacheResource.java

@@ -90,4 +90,11 @@ public void getTaskInfos(
                 .listTaskInfo(catalogName, schemaName)
                 .whenComplete(bindAsyncResponse(asyncResponse));
     }
+
+    @GET
+    @Path("duckdb/settings")
+    public void getDuckDBSettings(@Suspended AsyncResponse asyncResponse)
+    {
+        asyncResponse.resume(Response.ok(cacheManager.getDuckDBSettings()).build());
+    }
 }

accio-testing/src/main/java/io/accio/testing/AbstractTestFramework.java

@@ -17,6 +17,7 @@
 import com.google.common.collect.ImmutableList;
 import io.accio.base.SessionContext;
 import io.accio.base.client.AutoCloseableIterator;
+import io.accio.base.client.duckdb.DuckDBConfig;
 import io.accio.base.client.duckdb.DuckdbClient;
 import io.accio.base.dto.Column;
 import io.accio.base.dto.Manifest;
@@ -64,7 +65,7 @@ public static Model addColumnsToModel(Model model, Column... columns)
     @BeforeClass
     public void init()
     {
-        duckdbClient = new DuckdbClient();
+        duckdbClient = new DuckdbClient(new DuckDBConfig());
         prepareData();
     }
 

accio-validation/src/test/java/io/accio/TestMetricValidation.java

@@ -17,6 +17,7 @@
 import io.accio.base.AccioMDL;
 import io.accio.base.AccioTypes;
 import io.accio.base.client.Client;
+import io.accio.base.client.duckdb.DuckDBConfig;
 import io.accio.base.client.duckdb.DuckdbClient;
 import io.accio.base.dto.Column;
 import io.accio.base.dto.EnumDefinition;
@@ -50,7 +51,7 @@ public class TestMetricValidation
 
     public TestMetricValidation()
     {
-        client = new DuckdbClient();
+        client = new DuckdbClient(new DuckDBConfig());
         sample = AccioMDL.fromManifest(withDefaultCatalogSchema()
                 .setModels(List.of(
                         Model.model("Flight",