fixup! Handle composite keys in SqliteDataReader.GetStream()

Handle more cases
dotnet · Dec 4, 2020 · 6c302aa · 6c302aa
1 parent d4c7fec
commit 6c302aa
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 8 deletions.
diff --git a/src/Microsoft.Data.Sqlite.Core/SqliteDataRecord.cs b/src/Microsoft.Data.Sqlite.Core/SqliteDataRecord.cs
@@ -239,6 +239,7 @@ public virtual Stream GetStream(int ordinal)
             if (!_rowidOrdinal.HasValue)
             {
                 _rowidOrdinal = -1;
+                var pkColumns = -1L;
 
                 for (var i = 0; i < FieldCount; i++)
                 {
@@ -280,21 +281,22 @@ public virtual Stream GetStream(int ordinal)
                     if (string.Equals(dataType, "INTEGER", StringComparison.OrdinalIgnoreCase)
                         && primaryKey != 0)
                     {
-                        long pkColumns;
-                        using (var command = _connection.CreateCommand())
+                        if (pkColumns < 0L)
                         {
-                            command.CommandText = "SELECT COUNT(*) FROM pragma_table_info($table) WHERE pk != 0;";
-                            command.Parameters.AddWithValue("$table", tableName);
+                            using (var command = _connection.CreateCommand())
+                            {
+                                command.CommandText = "SELECT COUNT(*) FROM pragma_table_info($table) WHERE pk != 0;";
+                                command.Parameters.AddWithValue("$table", tableName);
 
-                            pkColumns = (long)command.ExecuteScalar();
+                                pkColumns = (long)command.ExecuteScalar();
+                            }
                         }
 
-                        if (pkColumns == 1)
+                        if (pkColumns == 1L)
                         {
                             _rowidOrdinal = i;
+                            break;
                         }
-
-                        break;
                     }
                 }
 

diff --git a/test/Microsoft.Data.Sqlite.Tests/SqliteDataReaderTest.cs b/test/Microsoft.Data.Sqlite.Tests/SqliteDataReaderTest.cs
@@ -468,6 +468,34 @@ public void GetStream_works_when_composite_pk()
             }
         }
 
+        [Fact]
+        public void GetStream_works_when_composite_pk_and_rowid()
+        {
+            using (var connection = new SqliteConnection("Data Source=:memory:"))
+            {
+                connection.Open();
+
+                connection.ExecuteNonQuery(
+                    @"CREATE TABLE DataTable (Id1 INTEGER, Id2 INTEGER, Data BLOB, PRIMARY KEY (Id1, Id2));
+                    INSERT INTO DataTable VALUES (5, 6, X'01020304');");
+
+                var selectCommand = connection.CreateCommand();
+                selectCommand.CommandText = "SELECT Id1, Id2, rowid, Data FROM DataTable WHERE Id1 = 5 AND Id2 = 6";
+                using (var reader = selectCommand.ExecuteReader())
+                {
+                    Assert.True(reader.Read());
+                    using (var sourceStream = reader.GetStream(3))
+                    {
+                        Assert.IsType<SqliteBlob>(sourceStream);
+                        var buffer = new byte[4];
+                        var bytesRead = sourceStream.Read(buffer, 0, 4);
+                        Assert.Equal(4, bytesRead);
+                        Assert.Equal(new byte[] { 0x01, 0x02, 0x03, 0x04 }, buffer);
+                    }
+                }
+            }
+        }
+
         [Fact]
         public void GetStream_throws_when_closed()
         {