SNOW-1640968 chunk downloader fix (#1022)

snowflakedb · Oct 9, 2024 · 950fa55 · 950fa55
1 parent ef6db7e
commit 950fa55
Show file tree

Hide file tree

Showing 4 changed files with 327 additions and 260 deletions.
diff --git a/Snowflake.Data.Tests/UnitTests/SFReusableChunkTest.cs b/Snowflake.Data.Tests/UnitTests/SFReusableChunkTest.cs
@@ -219,6 +219,33 @@ public void TestResetClearsChunkData()
             Assert.AreEqual(0, chunk.ChunkIndex);
             Assert.AreEqual(chunkInfo.url, chunk.Url);
             Assert.AreEqual(chunkInfo.rowCount, chunk.RowCount);
+            Assert.AreEqual(chunkInfo.uncompressedSize, chunk.UncompressedSize);
+            Assert.Greater(chunk.data.blockCount, 0);
+            Assert.Greater(chunk.data.metaBlockCount, 0);
+        }
+
+        [Test]
+        public void TestClearRemovesAllChunkData()
+        {
+            const int RowCount = 3;
+            string data = "[ [\"1\"],  [\"2\"],  [\"3\"] ]";
+            var chunk = PrepareChunkAsync(data, 1, RowCount).Result;
+
+            ExecResponseChunk chunkInfo = new ExecResponseChunk()
+            {
+                url = "new_url",
+                uncompressedSize = 100,
+                rowCount = 200
+            };
+
+            chunk.Clear();
+
+            Assert.AreEqual(0, chunk.ChunkIndex);
+            Assert.AreEqual(null, chunk.Url);
+            Assert.AreEqual(0, chunk.RowCount);
+            Assert.AreEqual(0, chunk.UncompressedSize);
+            Assert.AreEqual(0, chunk.data.blockCount);
+            Assert.AreEqual(0, chunk.data.metaBlockCount);
         }
 
         private async Task<SFReusableChunk> PrepareChunkAsync(string stringData, int colCount, int rowCount)

diff --git a/Snowflake.Data/Core/BaseResultChunk.cs b/Snowflake.Data/Core/BaseResultChunk.cs
@@ -9,21 +9,21 @@ namespace Snowflake.Data.Core
     public abstract class BaseResultChunk : IResultChunk
     {
         internal abstract ResultFormat ResultFormat { get; }
-        
+
         public int RowCount { get; protected set; }
-        
+
         public int ColumnCount { get; protected set; }
-        
+
         public int ChunkIndex { get; protected set; }
 
         internal int CompressedSize;
-        
+
         internal int UncompressedSize;
 
         internal string Url { get; set; }
 
         internal string[,] RowSet { get; set; }
-        
+
         public int GetRowCount() => RowCount;
 
         public int GetChunkIndex() => ChunkIndex;
@@ -32,11 +32,11 @@ public abstract class BaseResultChunk : IResultChunk
         public abstract UTF8Buffer ExtractCell(int rowIndex, int columnIndex);
 
         public abstract UTF8Buffer ExtractCell(int columnIndex);
-        
+
         internal abstract bool Next();
-        
+
         internal abstract bool Rewind();
-        
+
         internal virtual void Reset(ExecResponseChunk chunkInfo, int chunkIndex)
         {
             RowCount = chunkInfo.rowCount;
@@ -46,6 +46,15 @@ internal virtual void Reset(ExecResponseChunk chunkInfo, int chunkIndex)
             UncompressedSize = chunkInfo.uncompressedSize;
         }
 
+        internal virtual void Clear()
+        {
+            RowCount = 0;
+            Url = null;
+            ChunkIndex = 0;
+            CompressedSize = 0;
+            UncompressedSize = 0;
+        }
+
         internal virtual void ResetForRetry()
         {
         }