Skip to content

Commit

Permalink
Fix transaction problem in SQL Server Memory DeleteAsync method (#712)
Browse files Browse the repository at this point in the history
## Motivation and Context (Why the change? What's the scenario?)
This PR fixes an issue when trying to delete a document from SQL Server
Memory.

---------

Co-authored-by: Devis Lucato <devis@microsoft.com>
  • Loading branch information
marcominerva and dluc committed Jul 30, 2024
1 parent a08b726 commit 6c092a9
Showing 1 changed file with 25 additions and 23 deletions.
48 changes: 25 additions & 23 deletions extensions/SQLServer/SQLServer/SqlServerMemory.cs
Original file line number Diff line number Diff line change
Expand Up @@ -272,16 +272,11 @@ public async IAsyncEnumerable<MemoryRecord> GetListAsync(
}

string queryColumns = "[key], [payload], [tags]";
if (withEmbeddings) { queryColumns += ", [embedding]"; }

if (withEmbeddings)
{
queryColumns += ", [embedding]";
}
if (limit < 0) { limit = int.MaxValue; }

if (limit < 0)
{
limit = int.MaxValue;
}
var list = new List<MemoryRecord>();

var connection = new SqlConnection(this._config.ConnectionString);
await connection.OpenAsync(cancellationToken).ConfigureAwait(false);
Expand All @@ -291,20 +286,20 @@ public async IAsyncEnumerable<MemoryRecord> GetListAsync(
var tagFilters = new TagCollection();

command.CommandText = $@"
WITH [filters] AS
(
SELECT
cast([filters].[key] AS NVARCHAR(256)) COLLATE SQL_Latin1_General_CP1_CI_AS AS [name],
cast([filters].[value] AS NVARCHAR(256)) COLLATE SQL_Latin1_General_CP1_CI_AS AS [value]
FROM openjson(@filters) [filters]
)
SELECT TOP (@limit)
{queryColumns}
FROM
{this.GetFullTableName(this._config.MemoryTableName)}
WHERE 1=1
AND {this.GetFullTableName(this._config.MemoryTableName)}.[collection] = @index
{this.GenerateFilters(index, command.Parameters, filters)};";
WITH [filters] AS
(
SELECT
cast([filters].[key] AS NVARCHAR(256)) COLLATE SQL_Latin1_General_CP1_CI_AS AS [name],
cast([filters].[value] AS NVARCHAR(256)) COLLATE SQL_Latin1_General_CP1_CI_AS AS [value]
FROM openjson(@filters) [filters]
)
SELECT TOP (@limit)
{queryColumns}
FROM
{this.GetFullTableName(this._config.MemoryTableName)}
WHERE 1=1
AND {this.GetFullTableName(this._config.MemoryTableName)}.[collection] = @index
{this.GenerateFilters(index, command.Parameters, filters)};";

command.Parameters.AddWithValue("@index", index);
command.Parameters.AddWithValue("@limit", limit);
Expand All @@ -313,7 +308,9 @@ SELECT TOP (@limit)
var dataReader = await command.ExecuteReaderAsync(cancellationToken).ConfigureAwait(false);
while (await dataReader.ReadAsync(cancellationToken).ConfigureAwait(false))
{
yield return await this.ReadEntryAsync(dataReader, withEmbeddings, cancellationToken).ConfigureAwait(false);
// Iterates over the entries and saves them in a list so that the connection can be closed before returning the results.
var entry = await this.ReadEntryAsync(dataReader, withEmbeddings, cancellationToken).ConfigureAwait(false);
list.Add(entry);
}

await dataReader.DisposeAsync().ConfigureAwait(false);
Expand All @@ -323,6 +320,11 @@ SELECT TOP (@limit)
command.Dispose();
connection.Dispose();
}

foreach (var item in list)
{
yield return item;
}
}

/// <inheritdoc/>
Expand Down

0 comments on commit 6c092a9

Please sign in to comment.