(arrow) branch main updated: GH-38737: [Java] Fix JDBC caching of SqlInfo values (#38739)

[li...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new 5583dbeca4 GH-38737: [Java] Fix JDBC caching of SqlInfo values (#38739)
5583dbeca4 is described below

commit 5583dbeca4f36d6eadba979c619c4a07dbb2095f
Author: James Duong <ja...@improving.com>
AuthorDate: Thu Nov 16 06:20:12 2023 -0800

    GH-38737: [Java] Fix JDBC caching of SqlInfo values (#38739)
    
    ### Rationale for this change
    The cache of SqlInfo properties that ArrowDatabaseMetaData maintains isn't populated in a thread-safe way. This can cause JDBC applications trying to retrieve several properties from DatabaseMetaData to encounter missing properties when they shouldn't.
    
    ### What changes are included in this PR?
    - Changed the checking for the cache being populated to be based on an AtomicBoolean marking that the cache is fully populated, rather than just checking if the cache is empty.
    - Avoid having multiple threads call getSqlInfo() if they see that the cache is empty concurrently.
    
    ### Are these changes tested?
    Verified existing unit tests.
    
    ### Are there any user-facing changes?
    No.
    * Closes: #38737
    
    Authored-by: James Duong <ja...@improving.com>
    Signed-off-by: David Li <li...@gmail.com>
---
 .../apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java  | 16 +++++++++++-----
 .../arrow/driver/jdbc/utils/MockFlightSqlProducer.java   | 10 +++++-----
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java
index da2b0b00ed..3487e58a64 100644
--- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java
+++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java
@@ -45,11 +45,11 @@ import java.sql.DatabaseMetaData;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.EnumMap;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
@@ -145,8 +145,8 @@ public class ArrowDatabaseMetadata extends AvaticaDatabaseMetaData {
           Field.notNullable("IS_AUTOINCREMENT", Types.MinorType.VARCHAR.getType()),
           Field.notNullable("IS_GENERATEDCOLUMN", Types.MinorType.VARCHAR.getType())
       ));
-  private final Map<SqlInfo, Object> cachedSqlInfo =
-      Collections.synchronizedMap(new EnumMap<>(SqlInfo.class));
+  private final AtomicBoolean isCachePopulated = new AtomicBoolean(false);
+  private final Map<SqlInfo, Object> cachedSqlInfo = new EnumMap<>(SqlInfo.class);
   private static final Map<Integer, Integer> sqlTypesToFlightEnumConvertTypes = new HashMap<>();
 
   static {
@@ -729,10 +729,15 @@ public class ArrowDatabaseMetadata extends AvaticaDatabaseMetaData {
                                                  final Class<T> desiredType)
       throws SQLException {
     final ArrowFlightConnection connection = getConnection();
-    if (cachedSqlInfo.isEmpty()) {
-      final FlightInfo sqlInfo = connection.getClientHandler().getSqlInfo();
+    if (!isCachePopulated.get()) {
+      // Lock-and-populate the cache. Only issue the call to getSqlInfo() once,
+      // populate the cache, then mark it as populated.
+      // Note that multiple callers from separate threads can see that the cache is not populated, but only
+      // one thread will try to populate the cache. Other threads will see the cache is non-empty when acquiring
+      // the lock on the cache and skip population.
       synchronized (cachedSqlInfo) {
         if (cachedSqlInfo.isEmpty()) {
+          final FlightInfo sqlInfo = connection.getClientHandler().getSqlInfo();
           try (final ResultSet resultSet =
                    ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo(
                        connection, sqlInfo, null)) {
@@ -741,6 +746,7 @@ public class ArrowDatabaseMetadata extends AvaticaDatabaseMetaData {
                   resultSet.getObject("value"));
             }
           }
+          isCachePopulated.set(true);
         }
       }
     }
diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
index 75a7396931..2b65f8f5a0 100644
--- a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
+++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
@@ -105,8 +105,8 @@ public final class MockFlightSqlProducer implements FlightSqlProducer {
 
   private final Map<String, Integer> actionTypeCounter = new HashMap<>();
 
-  private static FlightInfo getFightInfoExportedAndImportedKeys(final Message message,
-                                                                final FlightDescriptor descriptor) {
+  private static FlightInfo getFlightInfoExportedAndImportedKeys(final Message message,
+                                                                 final FlightDescriptor descriptor) {
     return getFlightInfo(message, Schemas.GET_IMPORTED_KEYS_SCHEMA, descriptor);
   }
 
@@ -529,14 +529,14 @@ public final class MockFlightSqlProducer implements FlightSqlProducer {
   public FlightInfo getFlightInfoExportedKeys(final CommandGetExportedKeys commandGetExportedKeys,
                                               final CallContext callContext,
                                               final FlightDescriptor flightDescriptor) {
-    return getFightInfoExportedAndImportedKeys(commandGetExportedKeys, flightDescriptor);
+    return getFlightInfoExportedAndImportedKeys(commandGetExportedKeys, flightDescriptor);
   }
 
   @Override
   public FlightInfo getFlightInfoImportedKeys(final CommandGetImportedKeys commandGetImportedKeys,
                                               final CallContext callContext,
                                               final FlightDescriptor flightDescriptor) {
-    return getFightInfoExportedAndImportedKeys(commandGetImportedKeys, flightDescriptor);
+    return getFlightInfoExportedAndImportedKeys(commandGetImportedKeys, flightDescriptor);
   }
 
   @Override
@@ -544,7 +544,7 @@ public final class MockFlightSqlProducer implements FlightSqlProducer {
       final CommandGetCrossReference commandGetCrossReference,
       final CallContext callContext,
       final FlightDescriptor flightDescriptor) {
-    return getFightInfoExportedAndImportedKeys(commandGetCrossReference, flightDescriptor);
+    return getFlightInfoExportedAndImportedKeys(commandGetCrossReference, flightDescriptor);
   }
 
   @Override