allow specifying extension directory

ocadaruma · Nov 26, 2023 · ddef340 · ddef340
1 parent 2fde68c
commit ddef340
Showing 1 changed file with 39 additions and 40 deletions.
diff --git a/jdbc/src/main/java/com/mayreh/duckdb_jfr/DuckDBJfrConnection.java b/jdbc/src/main/java/com/mayreh/duckdb_jfr/DuckDBJfrConnection.java
@@ -5,6 +5,7 @@
 import java.io.UncheckedIOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.PreparedStatement;
@@ -16,60 +17,58 @@
  * Connection wrapper for DuckDB with JFR support.
  */
 public class DuckDBJfrConnection implements AutoCloseable {
+    private static final String DUCKDB_JFR_EXTENSION_DIR_PROP = "duckdb.jfr.extension.dir";
+
     private final Connection connection;
     private static final Path extensionPath;
 
     static {
         try {
-            Class.forName("org.duckdb.DuckDBDriver");
-
-            String arch = System.getProperty("os.arch").toLowerCase().trim();
-            String os = System.getProperty("os.name").toLowerCase().trim();
-
-            if (os.startsWith("mac")) {
-                os = "osx";
-                arch = "universal";
-            } else if (os.startsWith("linux")) {
-                os = "linux";
-                switch (arch) {
-                    case "x86_64":
-                    case "amd64":
-                        arch = "amd64";
-                        break;
-                    case "aarch64":
-                    case "arm64":
-                        arch = "arm64";
-                        break;
-                    default:
-                        throw new RuntimeException("Unsupported architecture: " + arch);
+            String extensionDir = System.getProperty(DUCKDB_JFR_EXTENSION_DIR_PROP);
+            if (extensionDir == null) {
+                Path tempDir = Files.createTempDirectory("duckdb_jfr");
+                extensionDir = tempDir.toString();
+            }
+            // we need to use this filename since duckdb identifies init function by filename
+            extensionPath = Paths.get(extensionDir).resolve("libduckdb_jfr_extension.so");
+            if (!Files.exists(extensionPath)) {
+                Class.forName("org.duckdb.DuckDBDriver");
+
+                String arch = System.getProperty("os.arch").toLowerCase().trim();
+                String os = System.getProperty("os.name").toLowerCase().trim();
+
+                if (os.startsWith("mac")) {
+                    os = "osx";
+                    arch = "universal";
+                } else if (os.startsWith("linux")) {
+                    os = "linux";
+                    switch (arch) {
+                        case "x86_64":
+                        case "amd64":
+                            arch = "amd64";
+                            break;
+                        case "aarch64":
+                        case "arm64":
+                            arch = "arm64";
+                            break;
+                        default:
+                            throw new RuntimeException("Unsupported architecture: " + arch);
+                    }
+                } else {
+                    throw new RuntimeException("Unsupported OS: " + os);
+                }
+                String resourceName = String.format("%s-%s/libduckdb_jfr_extension.so", os, arch);
+                try (InputStream s = DuckDBJfrConnection.class.getClassLoader().getResourceAsStream(resourceName)) {
+                    Files.copy(s, extensionPath);
                 }
-            } else {
-                throw new RuntimeException("Unsupported OS: " + os);
             }
-            String resourceName = String.format("%s-%s/libduckdb_jfr_extension.so", os, arch);
-            extensionPath = extractExtension(resourceName);
         } catch (ClassNotFoundException e) {
             throw new RuntimeException(e);
         } catch (IOException e) {
             throw new UncheckedIOException(e);
         }
     }
 
-    private static Path extractExtension(String resourceName) throws IOException {
-        Path tempDir = Files.createTempDirectory("duckdb_jfr");
-        tempDir.toFile().deleteOnExit();
-
-        // we need to use this filename since duckdb identifies init function by filename
-        Path extensionPath = tempDir.resolve("libduckdb_jfr_extension.so");
-        extensionPath.toFile().deleteOnExit();
-
-        try (InputStream s = DuckDBJfrConnection.class.getClassLoader().getResourceAsStream(resourceName)) {
-            Files.copy(s, extensionPath);
-        }
-
-        return extensionPath;
-    }
-
     /**
      * Create a new in-memory DuckDB connection with JFR-extension preloaded.
      */