[duckdb] Fixed multi-store bug (#1459)

The view created by DuckDBDVRT was called current_version, so it could
not work with multiple stores. Changed it so that the view inherits the
store name instead.

Miscellaneous build changes:

- Fixed the duckdb exclusion rule in all-modules.
- Bumped up the duckdb snapshot dep to the current latest one.

Author: FelixGV

all-modules/build.gradle

@@ -3,7 +3,7 @@
 dependencies {
   rootProject.subprojects.each { subproject ->
     // Excluding venice-duckdb as it's experimental and causes build issues
-    if (subproject.path != project.path && subproject.subprojects.isEmpty() && subproject.path != ':venice-duckdb') {
+    if (subproject.path != project.path && subproject.subprojects.isEmpty() && subproject.path != ':integrations:venice-duckdb') {
       implementation (project(subproject.path)) {
         exclude group: 'org.duckdb'
         exclude module: 'venice-duckdb'

build.gradle

@@ -82,7 +82,7 @@ ext.libraries = [
     commonsLang: 'commons-lang:commons-lang:2.6',
     conscrypt: 'org.conscrypt:conscrypt-openjdk-uber:2.5.2',
     d2: "com.linkedin.pegasus:d2:${pegasusVersion}",
-    duckdbJdbc: "org.duckdb:duckdb_jdbc:1.2.0-20250117.012118-121", // TODO: Remove SNAPSHOT when the real release is published!
+    duckdbJdbc: "org.duckdb:duckdb_jdbc:1.2.0-20250121.012246-127", // TODO: Remove SNAPSHOT when the real release is published!
     failsafe: 'net.jodah:failsafe:2.4.0',
     fastUtil: 'it.unimi.dsi:fastutil:8.3.0',
     grpcNettyShaded: "io.grpc:grpc-netty-shaded:${grpcVersion}",

gradle/spotbugs/exclude.xml

@@ -504,6 +504,8 @@
       <Class name="com.linkedin.venice.duckdb.DuckDBHelloWorldTest"/>
       <Class name="com.linkedin.venice.duckdb.DuckDBAvroToSQLTest"/>
       <Class name="com.linkedin.venice.duckdb.DuckDBDaVinciRecordTransformer"/>
+      <Class name="com.linkedin.venice.duckdb.DuckDBDaVinciRecordTransformerTest"/>
+      <Class name="com.linkedin.venice.endToEnd.DuckDBDaVinciRecordTransformerIntegrationTest"/>
       <Class name="com.linkedin.venice.sql.SQLHelloWorldTest"/>
       <Class name="com.linkedin.venice.sql.SQLUtilsTest"/>
       <Class name="com.linkedin.venice.sql.SQLUtils"/>

integrations/venice-duckdb/src/integrationTest/java/com/linkedin/venice/endToEnd/DuckDBDaVinciRecordTransformerIntegrationTest.java

@@ -158,18 +158,18 @@ public void testRecordTransformer() throws Exception {
 
       clientWithRecordTransformer.subscribeAll().get();
 
-      assertRowCount(duckDBUrl, "subscribeAll() finishes!");
+      assertRowCount(duckDBUrl, storeName, "subscribeAll() finishes!");
 
       clientWithRecordTransformer.unsubscribeAll();
     }
 
-    assertRowCount(duckDBUrl, "DVC gets closed!");
+    assertRowCount(duckDBUrl, storeName, "DVC gets closed!");
   }
 
-  private void assertRowCount(String duckDBUrl, String assertionErrorMsg) throws SQLException {
+  private void assertRowCount(String duckDBUrl, String storeName, String assertionErrorMsg) throws SQLException {
     try (Connection connection = DriverManager.getConnection(duckDBUrl);
         Statement statement = connection.createStatement();
-        ResultSet rs = statement.executeQuery("SELECT count(*) FROM current_version")) {
+        ResultSet rs = statement.executeQuery("SELECT count(*) FROM " + storeName)) {
       assertTrue(rs.next());
       int rowCount = rs.getInt(1);
       assertEquals(

integrations/venice-duckdb/src/main/java/com/linkedin/venice/duckdb/DuckDBDaVinciRecordTransformer.java

@@ -29,8 +29,7 @@ public class DuckDBDaVinciRecordTransformer
   private static final Logger LOGGER = LogManager.getLogger(DuckDBDaVinciRecordTransformer.class);
   private static final String duckDBFilePath = "my_database.duckdb";
   private static final String deleteStatementTemplate = "DELETE FROM %s WHERE %s = ?;";
-  private static final String createViewStatementTemplate =
-      "CREATE OR REPLACE VIEW current_version AS SELECT * FROM %s;";
+  private static final String createViewStatementTemplate = "CREATE OR REPLACE VIEW %s AS SELECT * FROM %s;";
   private static final String dropTableStatementTemplate = "DROP TABLE %s;";
   private final String storeNameWithoutVersionInfo;
   private final String versionTableName;
@@ -132,7 +131,8 @@ public void onStartVersionIngestion(boolean isCurrentVersion) {
 
       if (isCurrentVersion) {
         // Unable to convert to prepared statement as table and column names can't be parameterized
-        String createViewStatement = String.format(createViewStatementTemplate, versionTableName);
+        String createViewStatement =
+            String.format(createViewStatementTemplate, storeNameWithoutVersionInfo, versionTableName);
         stmt.execute(createViewStatement);
       }
     } catch (SQLException e) {
@@ -146,7 +146,8 @@ public void onEndVersionIngestion(int currentVersion) {
         Statement stmt = connection.createStatement()) {
       // Swap to current version
       String currentVersionTableName = buildStoreNameWithVersion(currentVersion);
-      String createViewStatement = String.format(createViewStatementTemplate, currentVersionTableName);
+      String createViewStatement =
+          String.format(createViewStatementTemplate, storeNameWithoutVersionInfo, currentVersionTableName);
       stmt.execute(createViewStatement);
 
       if (currentVersion != getStoreVersion()) {

integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBDaVinciRecordTransformerTest.java

@@ -1,6 +1,7 @@
 package com.linkedin.venice.duckdb;
 
 import static com.linkedin.venice.utils.TestWriteUtils.NAME_RECORD_V1_SCHEMA;
+import static com.linkedin.venice.utils.TestWriteUtils.SIMPLE_USER_WITH_DEFAULT_SCHEMA;
 import static com.linkedin.venice.utils.TestWriteUtils.SINGLE_FIELD_RECORD_SCHEMA;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
@@ -140,7 +141,7 @@ public void testVersionSwap() throws SQLException {
 
     try (Connection connection = DriverManager.getConnection(duckDBUrl);
         Statement stmt = connection.createStatement()) {
-      try (ResultSet rs = stmt.executeQuery("SELECT * FROM current_version")) {
+      try (ResultSet rs = stmt.executeQuery("SELECT * FROM " + storeName)) {
         assertTrue(rs.next(), "There should be a first row!");
         assertEquals(rs.getString("firstName"), "Duck");
         assertEquals(rs.getString("lastName"), "Goose");
@@ -150,12 +151,112 @@ public void testVersionSwap() throws SQLException {
       // Swap here
       recordTransformer_v1.onEndVersionIngestion(2);
 
-      try (ResultSet rs = stmt.executeQuery("SELECT * FROM current_version")) {
+      try (ResultSet rs = stmt.executeQuery("SELECT * FROM " + storeName)) {
         assertTrue(rs.next(), "There should be a first row!");
         assertEquals(rs.getString("firstName"), "Goose");
         assertEquals(rs.getString("lastName"), "Duck");
         assertFalse(rs.next(), "There should be only one row!");
       }
     }
   }
+
+  @Test
+  public void testTwoTablesConcurrently() throws SQLException {
+    String tempDir = Utils.getTempDataDirectory().getAbsolutePath();
+
+    String store1 = "store1";
+    String store2 = "store2";
+
+    DuckDBDaVinciRecordTransformer recordTransformerForStore1 = new DuckDBDaVinciRecordTransformer(
+        1,
+        SINGLE_FIELD_RECORD_SCHEMA,
+        NAME_RECORD_V1_SCHEMA,
+        NAME_RECORD_V1_SCHEMA,
+        false,
+        tempDir,
+        store1,
+        columnsToProject);
+    DuckDBDaVinciRecordTransformer recordTransformerForStore2 = new DuckDBDaVinciRecordTransformer(
+        1,
+        SIMPLE_USER_WITH_DEFAULT_SCHEMA,
+        NAME_RECORD_V1_SCHEMA,
+        NAME_RECORD_V1_SCHEMA,
+        false,
+        tempDir,
+        store2,
+        columnsToProject);
+
+    String duckDBUrl = recordTransformerForStore1.getDuckDBUrl();
+
+    recordTransformerForStore1.onStartVersionIngestion(true);
+
+    GenericRecord keyRecord = new GenericData.Record(SINGLE_FIELD_RECORD_SCHEMA);
+    keyRecord.put("key", "key");
+    Lazy<GenericRecord> lazyKey = Lazy.of(() -> keyRecord);
+
+    GenericRecord valueRecordForStore1 = new GenericData.Record(NAME_RECORD_V1_SCHEMA);
+    valueRecordForStore1.put("firstName", "Duck");
+    valueRecordForStore1.put("lastName", "Goose");
+    Lazy<GenericRecord> lazyValue = Lazy.of(() -> valueRecordForStore1);
+    recordTransformerForStore1.processPut(lazyKey, lazyValue);
+
+    try (Connection connection = DriverManager.getConnection(duckDBUrl);
+        Statement stmt = connection.createStatement()) {
+      try (ResultSet rs = stmt.executeQuery("SELECT * FROM " + store1)) {
+        assertTrue(rs.next(), "There should be a first row!");
+        assertEquals(rs.getString("key"), "key");
+        assertEquals(rs.getString("firstName"), "Duck");
+        assertEquals(rs.getString("lastName"), "Goose");
+        assertFalse(rs.next(), "There should be only one row!");
+      }
+    }
+
+    recordTransformerForStore2.onStartVersionIngestion(true);
+
+    GenericRecord keyRecordForStore2 = new GenericData.Record(SIMPLE_USER_WITH_DEFAULT_SCHEMA);
+    keyRecordForStore2.put("key", "key");
+    keyRecordForStore2.put("value", "value");
+    Lazy<GenericRecord> lazyKeyForStore2 = Lazy.of(() -> keyRecordForStore2);
+
+    GenericRecord valueRecordForStore2 = new GenericData.Record(NAME_RECORD_V1_SCHEMA);
+    valueRecordForStore2.put("firstName", "Duck2");
+    valueRecordForStore2.put("lastName", "Goose2");
+    Lazy<GenericRecord> lazyValueForStore2 = Lazy.of(() -> valueRecordForStore2);
+    recordTransformerForStore2.processPut(lazyKeyForStore2, lazyValueForStore2);
+
+    try (Connection connection = DriverManager.getConnection(duckDBUrl);
+        Statement stmt = connection.createStatement()) {
+      try (ResultSet rs = stmt.executeQuery("SELECT * FROM " + store1)) {
+        assertTrue(rs.next(), "There should be a first row!");
+        assertEquals(rs.getString("key"), "key");
+        assertEquals(rs.getString("firstName"), "Duck");
+        assertEquals(rs.getString("lastName"), "Goose");
+        assertFalse(rs.next(), "There should be only one row!");
+      }
+
+      try (ResultSet rs = stmt.executeQuery("SELECT * FROM " + store2)) {
+        assertTrue(rs.next(), "There should be a first row!");
+        assertEquals(rs.getString("key"), "key");
+        assertEquals(rs.getString("value"), "value");
+        assertEquals(rs.getString("firstName"), "Duck2");
+        assertEquals(rs.getString("lastName"), "Goose2");
+        assertFalse(rs.next(), "There should be only one row!");
+      }
+
+      try (ResultSet rs = stmt.executeQuery(
+          "SELECT s1.key AS s1key, s1.firstName AS s1FirstName, s1.lastName AS s1LastName, "
+              + "s2.key AS s2key, s2.value AS s2value, s2.firstName AS s2FirstName, s2.lastName AS s2LastName "
+              + "FROM " + store1 + " s1 JOIN " + store2 + " s2 ON s1.key = s2.key")) {
+        assertTrue(rs.next(), "There should be a first row!");
+        assertEquals(rs.getString("s1key"), "key");
+        assertEquals(rs.getString("s1FirstName"), "Duck");
+        assertEquals(rs.getString("s1LastName"), "Goose");
+        assertEquals(rs.getString("s2key"), "key");
+        assertEquals(rs.getString("s2value"), "value");
+        assertEquals(rs.getString("s2FirstName"), "Duck2");
+        assertEquals(rs.getString("s2LastName"), "Goose2");
+        assertFalse(rs.next(), "There should be only one row!");
+      }
+    }
+  }
 }