add tests, add two-hex-chars on client instead of server, add refresh…

…TableInformation call right on registerTable
snowflakedb · Oct 25, 2024 · 967d5ae · 967d5ae
1 parent ca4f534
commit 967d5ae
Show file tree

Hide file tree

Showing 13 changed files with 301 additions and 812 deletions.
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/IStorageManager.java b/src/main/java/net/snowflake/ingest/streaming/internal/IStorageManager.java
@@ -26,7 +26,7 @@ interface IStorageManager {
   IStorage getStorage(String fullyQualifiedTableName);
 
   /** Informs the storage manager about a new table that's being ingested into by the client. */
-  void registerTable(TableRef tableRef, FileLocationInfo locationInfo);
+  void registerTable(TableRef tableRef);
 
   /**
    * Generate a unique blob path and increment the blob sequencer

diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/InternalStage.java b/src/main/java/net/snowflake/ingest/streaming/internal/InternalStage.java
@@ -97,8 +97,7 @@ class InternalStage implements IStorage {
         (SnowflakeFileTransferMetadataWithAge) null,
         maxUploadRetries);
     Utils.assertStringNotNullOrEmpty("client prefix", clientPrefix);
-    this.fileLocationInfo = fileLocationInfo;
-    this.fileTransferMetadataWithAge = createFileTransferMetadataWithAge(fileLocationInfo);
+    setFileLocationInfo(fileLocationInfo);
   }
 
   /**
@@ -214,8 +213,6 @@ private void putRemote(String fullFilePath, byte[] data, int retryCount)
    */
   synchronized SnowflakeFileTransferMetadataWithAge refreshSnowflakeMetadata(boolean force)
       throws SnowflakeSQLException, IOException {
-    logger.logInfo("Refresh Snowflake metadata, client={} force={}", clientName, force);
-
     if (!force
         && fileTransferMetadataWithAge != null
         && fileTransferMetadataWithAge.timestamp.isPresent()
@@ -224,12 +221,19 @@ synchronized SnowflakeFileTransferMetadataWithAge refreshSnowflakeMetadata(boole
       return fileTransferMetadataWithAge;
     }
 
+    logger.logInfo(
+        "Refresh Snowflake metadata, client={} force={} tableRef={}", clientName, force, tableRef);
+
     FileLocationInfo location =
         this.owningManager.getRefreshedLocation(this.tableRef, Optional.empty());
-    SnowflakeFileTransferMetadataWithAge metadata = createFileTransferMetadataWithAge(location);
-    this.fileLocationInfo = location;
-    this.fileTransferMetadataWithAge = metadata;
-    return metadata;
+    setFileLocationInfo(location);
+    return this.fileTransferMetadataWithAge;
+  }
+
+  private synchronized void setFileLocationInfo(FileLocationInfo fileLocationInfo)
+      throws SnowflakeSQLException, IOException {
+    this.fileTransferMetadataWithAge = createFileTransferMetadataWithAge(fileLocationInfo);
+    this.fileLocationInfo = fileLocationInfo;
   }
 
   static SnowflakeFileTransferMetadataWithAge createFileTransferMetadataWithAge(

diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java b/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java
@@ -116,7 +116,7 @@ public InternalStage getStorage(String fullyQualifiedTableName) {
    * nothing as there's no per-table state yet for FDN tables (that use internal stages).
    */
   @Override
-  public void registerTable(TableRef tableRef, FileLocationInfo locationInfo) {}
+  public void registerTable(TableRef tableRef) {}
 
   /**
    * Gets the latest file location info (with a renewed short-lived access token) for the specified