Skip to content

Commit d304173

Browse files
committed
Use default alias and DB pinning for DuckLake
This PR adds the following changes to DuckLake init: - if `ducklake_alias` option is not specified, then default value `ducklake` is used for it - when default URL `jdbc:duckdb:` or explicit memory URL `jdbc:duckdb:memory:` is used, then this URL is changed to `jdbc:duckdb:memory:ducklakemem` to be able to share this DB across multiple connections - DB pinning is enabled automatically (unless disabled by user) - Result set streaming is enabeld automatically (unless disabled by user) - DuckLake initialization is done only once for each DB (and each DuckLake path), `ATTACH IF NOT EXISTS` call is changed to `ATTACH` Testing: test coverage pending, DuckLake is not yet available in `main` branch.
1 parent 66bf0e7 commit d304173

File tree

2 files changed

+65
-35
lines changed

2 files changed

+65
-35
lines changed

src/main/java/org/duckdb/DuckDBDriver.java

Lines changed: 50 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,15 @@ public class DuckDBDriver implements java.sql.Driver {
1818
public static final String JDBC_PIN_DB = "jdbc_pin_db";
1919

2020
static final String DUCKDB_URL_PREFIX = "jdbc:duckdb:";
21+
static final String MEMORY_DB = ":memory:";
2122

2223
private static final String DUCKLAKE_OPTION = "ducklake";
2324
private static final String DUCKLAKE_ALIAS_OPTION = "ducklake_alias";
2425
private static final Pattern DUCKLAKE_ALIAS_OPTION_PATTERN = Pattern.compile("[a-zA-Z0-9_]+");
2526
private static final String DUCKLAKE_URL_PREFIX = "ducklake:";
26-
private static final ReentrantLock DUCKLAKE_INIT_LOCK = new ReentrantLock();
27+
private static final String DUCKLAKE_DEFAULT_DBNAME = MEMORY_DB + "ducklakemem";
28+
private static final LinkedHashSet<String> ducklakeInstances = new LinkedHashSet<>();
29+
private static final ReentrantLock ducklakeInitLock = new ReentrantLock();
2730

2831
private static final LinkedHashMap<String, ByteBuffer> pinnedDbRefs = new LinkedHashMap<>();
2932
private static final ReentrantLock pinnedDbRefsLock = new ReentrantLock();
@@ -42,39 +45,52 @@ public Connection connect(String url, Properties info) throws SQLException {
4245
if (!acceptsURL(url)) {
4346
return null;
4447
}
48+
final Properties props;
4549
if (info == null) {
46-
info = new Properties();
50+
props = new Properties();
4751
} else { // make a copy because we're removing the read only property below
48-
info = (Properties) info.clone();
52+
props = (Properties) info.clone();
4953
}
5054

5155
ParsedProps pp = parsePropsFromUrl(url);
5256
for (Map.Entry<String, String> en : pp.props.entrySet()) {
53-
info.put(en.getKey(), en.getValue());
57+
props.put(en.getKey(), en.getValue());
5458
}
55-
url = pp.shortUrl;
5659

57-
String readOnlyStr = removeOption(info, DUCKDB_READONLY_PROPERTY);
60+
String readOnlyStr = removeOption(props, DUCKDB_READONLY_PROPERTY);
5861
boolean readOnly = isStringTruish(readOnlyStr, false);
59-
info.put("duckdb_api", "jdbc");
62+
props.put("duckdb_api", "jdbc");
6063

6164
// Apache Spark passes this option when SELECT on a JDBC DataSource
6265
// table is performed. It is the internal Spark option and is likely
6366
// passed by mistake, so we need to ignore it to allow the connection
6467
// to be established.
65-
info.remove("path");
68+
props.remove("path");
6669

67-
String pinDbOptStr = removeOption(info, JDBC_PIN_DB);
68-
boolean pinDBOpt = isStringTruish(pinDbOptStr, false);
70+
String ducklake = removeOption(props, DUCKLAKE_OPTION);
71+
String ducklakeAlias = removeOption(props, DUCKLAKE_ALIAS_OPTION, DUCKLAKE_OPTION);
72+
final String shortUrl;
73+
if (null != ducklake) {
74+
setDefaultOptionValue(props, JDBC_PIN_DB, true);
75+
setDefaultOptionValue(props, JDBC_STREAM_RESULTS, true);
76+
String dbName = dbNameFromUrl(pp.shortUrl);
77+
if (MEMORY_DB.equals(dbName)) {
78+
shortUrl = DUCKDB_URL_PREFIX + DUCKLAKE_DEFAULT_DBNAME;
79+
} else {
80+
shortUrl = pp.shortUrl;
81+
}
82+
} else {
83+
shortUrl = pp.shortUrl;
84+
}
6985

70-
String ducklake = removeOption(info, DUCKLAKE_OPTION);
71-
String ducklakeAlias = removeOption(info, DUCKLAKE_ALIAS_OPTION);
86+
String pinDbOptStr = removeOption(props, JDBC_PIN_DB);
87+
boolean pinDBOpt = isStringTruish(pinDbOptStr, false);
7288

73-
DuckDBConnection conn = DuckDBConnection.newConnection(url, readOnly, info);
89+
DuckDBConnection conn = DuckDBConnection.newConnection(shortUrl, readOnly, props);
7490

75-
pinDB(pinDBOpt, url, conn);
91+
pinDB(pinDBOpt, shortUrl, conn);
7692

77-
initDucklake(conn, ducklake, ducklakeAlias);
93+
initDucklake(conn, shortUrl, ducklake, ducklakeAlias);
7894

7995
return conn;
8096
}
@@ -104,23 +120,29 @@ public Logger getParentLogger() throws SQLFeatureNotSupportedException {
104120
throw new SQLFeatureNotSupportedException("no logger");
105121
}
106122

107-
private static void initDucklake(Connection conn, String ducklake, String ducklakeAlias) throws SQLException {
123+
private static void initDucklake(Connection conn, String url, String ducklake, String ducklakeAlias)
124+
throws SQLException {
108125
if (null == ducklake) {
109126
return;
110127
}
111-
DUCKLAKE_INIT_LOCK.lock();
128+
ducklakeInitLock.lock();
112129
try {
113-
String attachQuery = createAttachQuery(ducklake, ducklakeAlias);
114-
try (Statement stmt = conn.createStatement()) {
115-
stmt.execute("INSTALL ducklake");
116-
stmt.execute("LOAD ducklake");
117-
stmt.execute(attachQuery);
118-
if (null != ducklakeAlias) {
119-
stmt.execute("USE " + ducklakeAlias);
130+
String dbName = dbNameFromUrl(url);
131+
String key = dbName + "#" + ducklake;
132+
if (!ducklakeInstances.contains(key)) {
133+
String attachQuery = createAttachQuery(ducklake, ducklakeAlias);
134+
try (Statement stmt = conn.createStatement()) {
135+
stmt.execute("INSTALL ducklake");
136+
stmt.execute("LOAD ducklake");
137+
stmt.execute(attachQuery);
120138
}
139+
ducklakeInstances.add(key);
121140
}
122141
} finally {
123-
DUCKLAKE_INIT_LOCK.unlock();
142+
ducklakeInitLock.unlock();
143+
}
144+
try (Statement stmt = conn.createStatement()) {
145+
stmt.execute("USE " + ducklakeAlias);
124146
}
125147
}
126148

@@ -129,14 +151,10 @@ private static String createAttachQuery(String ducklake, String ducklakeAlias) t
129151
if (!ducklake.startsWith(DUCKLAKE_URL_PREFIX)) {
130152
ducklake = DUCKLAKE_URL_PREFIX + ducklake;
131153
}
132-
String query = "ATTACH IF NOT EXISTS '" + ducklake + "'";
133-
if (null != ducklakeAlias) {
134-
if (!DUCKLAKE_ALIAS_OPTION_PATTERN.matcher(ducklakeAlias).matches()) {
135-
throw new SQLException("Invalid DuckLake alias specified: " + ducklakeAlias);
136-
}
137-
query += " AS " + ducklakeAlias;
154+
if (!DUCKLAKE_ALIAS_OPTION_PATTERN.matcher(ducklakeAlias).matches()) {
155+
throw new SQLException("Invalid DuckLake alias specified: " + ducklakeAlias);
138156
}
139-
return query;
157+
return "ATTACH '" + ducklake + "' AS " + ducklakeAlias;
140158
}
141159

142160
private static ParsedProps parsePropsFromUrl(String url) throws SQLException {

src/main/java/org/duckdb/JdbcUtils.java

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.duckdb;
22

33
import static org.duckdb.DuckDBDriver.DUCKDB_URL_PREFIX;
4+
import static org.duckdb.DuckDBDriver.MEMORY_DB;
45

56
import java.sql.SQLException;
67
import java.util.Properties;
@@ -19,11 +20,22 @@ static <T> T unwrap(Object obj, Class<T> iface) throws SQLException {
1920
}
2021

2122
static String removeOption(Properties props, String opt) {
23+
return removeOption(props, opt, null);
24+
}
25+
26+
static String removeOption(Properties props, String opt, String defaultVal) {
2227
Object obj = props.remove(opt);
2328
if (null != obj) {
2429
return obj.toString().trim();
2530
}
26-
return null;
31+
return defaultVal;
32+
}
33+
34+
static void setDefaultOptionValue(Properties props, String opt, Object value) {
35+
if (props.contains(opt)) {
36+
return;
37+
}
38+
props.put(opt, value);
2739
}
2840

2941
static boolean isStringTruish(String val, boolean defaultVal) throws SQLException {
@@ -56,9 +68,9 @@ static String dbNameFromUrl(String url) throws SQLException {
5668
}
5769
String dbName = shortUrl.substring(DUCKDB_URL_PREFIX.length()).trim();
5870
if (dbName.length() == 0) {
59-
dbName = ":memory:";
71+
dbName = MEMORY_DB;
6072
}
61-
if (dbName.startsWith("memory:")) {
73+
if (dbName.startsWith(MEMORY_DB.substring(1))) {
6274
dbName = ":" + dbName;
6375
}
6476
return dbName;

0 commit comments

Comments
 (0)