NO-SNOW: Fix StreamLoader vector metadata detection error handling

Author: sfc-gh-dheyman

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 # Changelog
 - Upcoming release (TBD)
+    - Fix StreamLoader vector metadata detection to handle empty metadata result sets safely and avoid aborting load jobs on non-critical metadata lookup failures.
     - Fix expired session token renewal when polling results (snowflakedb/snowflake-jdbc#2489)   
     - Fix missing minicore async initialization that was dropped during public API restructuring in v4.0.0
     - Adjust level of logging during Driver initialization

src/main/java/net/snowflake/client/internal/api/implementation/metadata/SnowflakeDatabaseMetaDataImpl.java

@@ -3722,9 +3722,27 @@ private ResultSet executeAndReturnEmptyResultIfNotFound(
     try {
       resultSet = statement.executeQuery(sql);
     } catch (SnowflakeSQLException e) {
+      if (metadataType == GET_COLUMNS) {
+        logger.warn(
+            "DatabaseMetaData.getColumns received SQL exception and may return empty result: "
+                + "sqlState={}, errorCode={}, queryId={}, message={}, sql={}",
+            e.getSQLState(),
+            e.getErrorCode(),
+            e.getQueryId(),
+            e.getMessage(),
+            sql);
+      }
       if (e.getSQLState().equals(SqlState.NO_DATA)
           || e.getSQLState().equals(SqlState.BASE_TABLE_OR_VIEW_NOT_FOUND)
           || e.getMessage().contains("Operation is not supported in reader account")) {
+        if (metadataType == GET_COLUMNS) {
+          logger.warn(
+              "DatabaseMetaData.getColumns normalized SQL exception to empty result set: "
+                  + "sqlState={}, queryId={}, sql={}",
+              e.getSQLState(),
+              e.getQueryId(),
+              sql);
+        }
         return SnowflakeDatabaseMetaDataResultSet.getEmptyResult(
             metadataType, statement, e.getQueryId());
       }

src/main/java/net/snowflake/client/internal/loader/StreamLoader.java

@@ -57,6 +57,9 @@ public class StreamLoader implements Loader, Runnable {
   /** Default batch row size */
   private static final long DEFAULT_BATCH_ROW_SIZE = -1L;
 
+  private static final String FAIL_ON_MISSING_COLUMN_METADATA_KEY =
+      SYSTEM_PARAMETER_PREFIX + "failOnMissingColumnMetadata";
+
   public static DatabaseMetaData metadata;
 
   private BufferStage _stage = null;
@@ -632,26 +635,49 @@ public void setVectorColumnType(String vectorType) {
   }
 
   public void setVectorColumns() {
+    _vectorColumnsNameAndSize.clear();
+    if (_columns == null || _columns.isEmpty()) {
+      return;
+    }
+
     try {
       DatabaseMetaData dbmd = _processConn.getMetaData();
       for (String col : _columns) {
         try (ResultSet rs = dbmd.getColumns(_database, _schema, _table, col)) {
-          rs.next();
+          if (!rs.next()) {
+            String message =
+                String.format(
+                    "No metadata row returned for column %s.%s.%s.%s",
+                    _database, _schema, _table, col);
+            logger.warn(message);
+            if (shouldFailOnMissingColumnMetadata()) {
+              throw new Loader.ConnectionError(message);
+            }
+            continue;
+          }
           if (isColumnTypeVector(rs.getString(6))) {
             _vectorColumnsNameAndSize.put(col, rs.getInt(7));
           }
         }
       }
     } catch (SQLException e) {
-      logger.error(e.getMessage(), e);
-      abort(new Loader.ConnectionError(Utils.getCause(e)));
+      logger.warn(
+          "Failed to detect vector columns via metadata query. "
+              + "Vector column type casting will be skipped. Error: {}",
+          e.getMessage());
+      logger.debug("setVectorColumns metadata query failure details", e);
+      _vectorColumnsNameAndSize.clear();
     }
   }
 
   private boolean isColumnTypeVector(String col) {
     return col != null && col.equalsIgnoreCase("vector");
   }
 
+  private boolean shouldFailOnMissingColumnMetadata() {
+    return Boolean.parseBoolean(systemGetProperty(FAIL_ON_MISSING_COLUMN_METADATA_KEY));
+  }
+
   @Override
   public void run() {
     try {

src/test/java/net/snowflake/client/internal/loader/FlatfileReadMultithreadIT.java

@@ -9,6 +9,7 @@
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.List;
@@ -30,6 +31,8 @@
 @Tag(TestTags.LOADER)
 public class FlatfileReadMultithreadIT {
   private final int NUM_RECORDS = 100000;
+  private static final String FAIL_ON_MISSING_COLUMN_METADATA_KEY =
+      "net.snowflake.client.loader.failOnMissingColumnMetadata";
 
   private static final String TARGET_STAGE = "STAGE_MULTITHREAD_LOADER";
   private static String TARGET_SCHEMA;
@@ -64,6 +67,9 @@ public static void tearDownClass() throws Throwable {
   @Test
   public void testIssueSimpleDateFormat() throws Throwable {
     final String targetTable = "TABLE_ISSUE_SIMPLEDATEFORMAT";
+    String originalFailOnMissingMetadataValue =
+        System.getProperty(FAIL_ON_MISSING_COLUMN_METADATA_KEY);
+    System.setProperty(FAIL_ON_MISSING_COLUMN_METADATA_KEY, "true");
     try (Connection testConnection = AbstractDriverIT.getConnection();
         Statement statement = testConnection.createStatement()) {
       try {
@@ -77,11 +83,15 @@ public void testIssueSimpleDateFormat() throws Throwable {
         Thread t2 =
             new Thread(
                 new FlatfileRead(NUM_RECORDS, TARGET_DB, TARGET_SCHEMA, TARGET_STAGE, targetTable));
+        List<Throwable> threadFailures = Collections.synchronizedList(new ArrayList<>());
+        t1.setUncaughtExceptionHandler((thread, throwable) -> threadFailures.add(throwable));
+        t2.setUncaughtExceptionHandler((thread, throwable) -> threadFailures.add(throwable));
 
         t1.start();
         t2.start();
         t1.join();
         t2.join();
+        assertThat("worker thread failures", threadFailures.isEmpty(), equalTo(true));
         try (ResultSet rs =
             statement.executeQuery(
                 String.format(
@@ -93,6 +103,12 @@ public void testIssueSimpleDateFormat() throws Throwable {
       } finally {
         statement.execute(
             String.format("DROP TABLE IF EXISTS %s.%s.%s", TARGET_DB, TARGET_SCHEMA, targetTable));
+        if (originalFailOnMissingMetadataValue == null) {
+          System.clearProperty(FAIL_ON_MISSING_COLUMN_METADATA_KEY);
+        } else {
+          System.setProperty(
+              FAIL_ON_MISSING_COLUMN_METADATA_KEY, originalFailOnMissingMetadataValue);
+        }
       }
     }
   }