diff --git a/MEVD/src/SqliteVec/SqliteCollection.cs b/MEVD/src/SqliteVec/SqliteCollection.cs index 475368e..9c6f298 100644 --- a/MEVD/src/SqliteVec/SqliteCollection.cs +++ b/MEVD/src/SqliteVec/SqliteCollection.cs @@ -31,6 +31,9 @@ public class SqliteCollection : VectorStoreCollectionThe connection string for the SQLite database represented by this . private readonly string _connectionString; + /// Optional callback that runs on each connection after it is opened and the sqlite-vec extension is loaded. + private readonly Action? _connectionInitializer; + /// The mapper to use when mapping between the consumer data model and the SQLite record. private readonly SqliteMapper _mapper; @@ -102,6 +105,7 @@ internal SqliteCollection(string connectionString, string name, Func keys, CancellationToken #region private - private async ValueTask GetConnectionAsync(CancellationToken cancellationToken = default) - { - var connection = new SqliteConnection(_connectionString); - await connection.OpenAsync(cancellationToken).ConfigureAwait(false); - connection.LoadVector(); - return connection; - } + private ValueTask GetConnectionAsync(CancellationToken cancellationToken = default) + => SqliteConnectionFactory.OpenAsync( + _connectionString, + _connectionInitializer, + loadVector: true, + cancellationToken: cancellationToken); private async IAsyncEnumerable> EnumerateAndMapSearchResultsAsync( List conditions, diff --git a/MEVD/src/SqliteVec/SqliteCollectionOptions.cs b/MEVD/src/SqliteVec/SqliteCollectionOptions.cs index e613f59..84407e0 100644 --- a/MEVD/src/SqliteVec/SqliteCollectionOptions.cs +++ b/MEVD/src/SqliteVec/SqliteCollectionOptions.cs @@ -1,6 +1,7 @@ // Licensed to the .NET Foundation under one or more agreements. // The .NET Foundation licenses this file to you under the MIT license. +using Microsoft.Data.Sqlite; using Microsoft.Extensions.VectorData; namespace CommunityToolkit.VectorData.SqliteVec; @@ -22,6 +23,7 @@ public SqliteCollectionOptions() internal SqliteCollectionOptions(SqliteCollectionOptions? source) : base(source) { VectorVirtualTableName = source?.VectorVirtualTableName; + ConnectionInitializer = source?.ConnectionInitializer; } /// @@ -31,4 +33,28 @@ internal SqliteCollectionOptions(SqliteCollectionOptions? source) : base(source) /// If not provided, collection name with prefix will be used as virtual table name. /// public string? VectorVirtualTableName { get; set; } + + /// + /// Gets or sets a callback that runs on every connection this collection opens. + /// + /// + /// + /// If defined, the callback runs after the connection is opened and the sqlite-vec extension is loaded. + /// + /// + /// Use it for per-connection settings that have no connection string equivalent, such as + /// PRAGMA synchronous=NORMAL, PRAGMA cache_size or PRAGMA temp_store=MEMORY, or for loading additional extensions. + /// + /// + /// When the collection is created from a , the callback configured on the vector store options + /// is used; this property applies to collections constructed directly. + /// + /// + /// The collection owns the connection and disposes it when the operation completes. The callback must not close or dispose + /// the connection, and must not leave a transaction or a data reader open. If the callback throws, the operation fails and + /// the exception is surfaced to the caller. The callback can run on several connections at the same time, so it must be + /// thread-safe. + /// + /// + public Action? ConnectionInitializer { get; set; } } diff --git a/MEVD/src/SqliteVec/SqliteConnectionFactory.cs b/MEVD/src/SqliteVec/SqliteConnectionFactory.cs new file mode 100644 index 0000000..fec2b95 --- /dev/null +++ b/MEVD/src/SqliteVec/SqliteConnectionFactory.cs @@ -0,0 +1,46 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using Microsoft.Data.Sqlite; + +namespace CommunityToolkit.VectorData.SqliteVec; + +/// +/// Creates the connections used by the vector store and its collections. +/// +internal static class SqliteConnectionFactory +{ + /// + /// Opens a connection, runs the shared setup steps in order, and hands it to the caller. + /// + /// + /// The caller owns the returned connection. If any setup step fails, the connection is disposed here + /// because the caller never receives it. + /// + public static async ValueTask OpenAsync( + string connectionString, + Action? initializer, + bool loadVector, + CancellationToken cancellationToken) + { + var connection = new SqliteConnection(connectionString); + + try + { + await connection.OpenAsync(cancellationToken).ConfigureAwait(false); + + if (loadVector) + { + connection.LoadVector(); + } + + initializer?.Invoke(connection); + return connection; + } + catch + { + await connection.DisposeAsync().ConfigureAwait(false); + throw; + } + } +} diff --git a/MEVD/src/SqliteVec/SqliteVec.csproj b/MEVD/src/SqliteVec/SqliteVec.csproj index fa7643b..a1fbbff 100644 --- a/MEVD/src/SqliteVec/SqliteVec.csproj +++ b/MEVD/src/SqliteVec/SqliteVec.csproj @@ -1,7 +1,7 @@  - 1.0.3-preview + 1.0.4-preview CommunityToolkit.VectorData.SqliteVec $(AssemblyName) net10.0;net8.0 diff --git a/MEVD/src/SqliteVec/SqliteVectorStore.cs b/MEVD/src/SqliteVec/SqliteVectorStore.cs index 0c42ed5..b3161e4 100644 --- a/MEVD/src/SqliteVec/SqliteVectorStore.cs +++ b/MEVD/src/SqliteVec/SqliteVectorStore.cs @@ -33,6 +33,9 @@ public sealed class SqliteVectorStore : VectorStore private readonly IEmbeddingGenerator? _embeddingGenerator; + /// Optional callback that runs on each connection after it is opened. + private readonly Action? _connectionInitializer; + /// /// Initializes a new instance of the class. /// @@ -47,6 +50,7 @@ public SqliteVectorStore(string connectionString, SqliteVectorStoreOptions? opti options ??= SqliteVectorStoreOptions.Default; _vectorVirtualTableName = options.VectorVirtualTableName; _embeddingGenerator = options.EmbeddingGenerator; + _connectionInitializer = options.ConnectionInitializer; var connectionStringBuilder = new SqliteConnectionStringBuilder(connectionString); @@ -71,7 +75,8 @@ public override SqliteCollection GetCollection(str { Definition = definition, VectorVirtualTableName = _vectorVirtualTableName, - EmbeddingGenerator = _embeddingGenerator + EmbeddingGenerator = _embeddingGenerator, + ConnectionInitializer = _connectionInitializer }); /// @@ -83,7 +88,8 @@ public override SqliteDynamicCollection GetDynamicCollection(string name, Vector { Definition = definition, VectorVirtualTableName = _vectorVirtualTableName, - EmbeddingGenerator = _embeddingGenerator + EmbeddingGenerator = _embeddingGenerator, + ConnectionInitializer = _connectionInitializer } ); #pragma warning restore IDE0090 @@ -95,8 +101,11 @@ public override async IAsyncEnumerable ListCollectionNamesAsync([Enumera const string TablePropertyName = "name"; const string Query = $"SELECT {TablePropertyName} FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';"; - using var connection = new SqliteConnection(_connectionString); - await connection.OpenAsync(cancellationToken).ConfigureAwait(false); + using var connection = await SqliteConnectionFactory.OpenAsync( + _connectionString, + _connectionInitializer, + loadVector: false, + cancellationToken: cancellationToken).ConfigureAwait(false); using var command = connection.CreateCommand(); command.CommandText = Query; diff --git a/MEVD/src/SqliteVec/SqliteVectorStoreOptions.cs b/MEVD/src/SqliteVec/SqliteVectorStoreOptions.cs index 4835b67..4d364c3 100644 --- a/MEVD/src/SqliteVec/SqliteVectorStoreOptions.cs +++ b/MEVD/src/SqliteVec/SqliteVectorStoreOptions.cs @@ -1,6 +1,7 @@ // Licensed to the .NET Foundation under one or more agreements. // The .NET Foundation licenses this file to you under the MIT license. +using Microsoft.Data.Sqlite; using Microsoft.Extensions.AI; namespace CommunityToolkit.VectorData.SqliteVec; @@ -23,6 +24,7 @@ internal SqliteVectorStoreOptions(SqliteVectorStoreOptions? source) { VectorVirtualTableName = source?.VectorVirtualTableName; EmbeddingGenerator = source?.EmbeddingGenerator; + ConnectionInitializer = source?.ConnectionInitializer; } /// @@ -37,4 +39,25 @@ internal SqliteVectorStoreOptions(SqliteVectorStoreOptions? source) /// Gets or sets the default embedding generator to use when generating vectors embeddings with this vector store. /// public IEmbeddingGenerator? EmbeddingGenerator { get; set; } + + /// + /// Gets or sets a callback that runs on every connection this vector store and the collections created from it open. + /// + /// + /// + /// If defined, the callback runs after the connection is opened and, on connections used for collection operations, after the sqlite-vec + /// extension is loaded. + /// + /// + /// Use it for per-connection settings that have no connection string equivalent, such as + /// PRAGMA synchronous=NORMAL, PRAGMA cache_size or PRAGMA temp_store=MEMORY, or for loading additional extensions. + /// + /// + /// The vector store owns the connection and disposes it when the operation completes. The callback must not close or dispose + /// the connection, and must not leave a transaction or a data reader open. If the callback throws, the operation fails and + /// the exception is surfaced to the caller. The callback can run on several connections at the same time, so it must be + /// thread-safe. + /// + /// + public Action? ConnectionInitializer { get; set; } } diff --git a/MEVD/test/SqliteVec.ConformanceTests/SqliteConnectionInitializerTests.cs b/MEVD/test/SqliteVec.ConformanceTests/SqliteConnectionInitializerTests.cs new file mode 100644 index 0000000..07ce40a --- /dev/null +++ b/MEVD/test/SqliteVec.ConformanceTests/SqliteConnectionInitializerTests.cs @@ -0,0 +1,168 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System.Data; +using CommunityToolkit.VectorData.SqliteVec; +using Microsoft.Data.Sqlite; +using Microsoft.Extensions.VectorData; +using Xunit; + +namespace SqliteVec.ConformanceTests; + +public sealed class SqliteConnectionInitializerTests : IDisposable +{ + private readonly string _databasePath = Path.Combine(Path.GetTempPath(), $"sqlitevec-init-{Guid.NewGuid():N}.db"); + + private string ConnectionString => $"Data Source={_databasePath};Pooling=false"; + + [Fact] + public async Task CollectionCallbackRunsOncePerOperationAndCanSetPragmas() + { + var invocations = 0; + var synchronousValues = new List(); + + var options = new SqliteCollectionOptions + { + ConnectionInitializer = connection => + { + invocations++; + + using var setCommand = connection.CreateCommand(); + setCommand.CommandText = "PRAGMA synchronous=NORMAL"; + setCommand.ExecuteNonQuery(); + + using var getCommand = connection.CreateCommand(); + getCommand.CommandText = "PRAGMA synchronous"; + using var reader = getCommand.ExecuteReader(); + Assert.True(reader.Read()); + synchronousValues.Add(reader.GetInt64(0)); + } + }; + + var collection = new SqliteCollection(ConnectionString, "InitializerCollection", options); + + await collection.EnsureCollectionExistsAsync(); + await collection.UpsertAsync(new TestRecord + { + Key = "key1", + Text = "hello", + Vector = new([1f, 2f, 3f, 4f]) + }); + Assert.NotNull(await collection.GetAsync("key1")); + Assert.True(await collection.CollectionExistsAsync()); + + // Each operation opens its own connection and runs the callback on it. + Assert.Equal(4, invocations); + // NORMAL is 1, FULL is 2 (default). + Assert.All(synchronousValues, value => Assert.Equal(1, value)); + } + + [Fact] + public async Task StoreCallbackRunsForStoreOperationsAndFlowsToCollections() + { + var invocations = 0; + + var store = new SqliteVectorStore(ConnectionString, new() + { + ConnectionInitializer = _ => invocations++ + }); + + await foreach (var _ in store.ListCollectionNamesAsync()) + { + } + + var collection = store.GetCollection("InitializerStoreCollection"); + await collection.EnsureCollectionExistsAsync(); + await collection.UpsertAsync(new TestRecord + { + Key = "key2", + Text = "hello", + Vector = new([1f, 2f, 3f, 4f]) + }); + + // One connection for ListCollectionNamesAsync, one each for ensure and upsert on the collection. + Assert.Equal(3, invocations); + } + + [Fact] + public async Task StoreCallbackRunsForDynamicCollectionOperations() + { + var invocations = 0; + + var store = new SqliteVectorStore(ConnectionString, new() + { + ConnectionInitializer = _ => invocations++ + }); + + Assert.False(await store.CollectionExistsAsync("MissingCollection")); + await store.EnsureCollectionDeletedAsync("MissingCollection"); + + // CollectionExistsAsync and EnsureCollectionDeletedAsync each build a dynamic collection which must call the initializer + Assert.Equal(2, invocations); + } + + [Fact] + public async Task CallbackRunsAfterVectorExtensionIsLoaded() + { + string? vecVersion = null; + + var options = new SqliteCollectionOptions + { + // vec_version is provided by the sqlite-vec extension, so this only works if the + // extension is loaded before the callback runs. + ConnectionInitializer = connection => + { + using var command = connection.CreateCommand(); + command.CommandText = "SELECT vec_version()"; + vecVersion = command.ExecuteScalar() as string; + } + }; + + var collection = new SqliteCollection(ConnectionString, "InitializerExtensionCollection", options); + await collection.EnsureCollectionExistsAsync(); + + Assert.False(string.IsNullOrEmpty(vecVersion)); + } + + [Fact] + public async Task ThrowingInitializerPropagatesExceptionAndDisposesConnection() + { + SqliteConnection? capturedConnection = null; + + var options = new SqliteCollectionOptions + { + ConnectionInitializer = connection => + { + capturedConnection = connection; + throw new InvalidOperationException("Initializer failed."); + } + }; + + var collection = new SqliteCollection(ConnectionString, "InitializerThrowCollection", options); + + var exception = await Assert.ThrowsAsync( + () => collection.EnsureCollectionExistsAsync()); + Assert.Equal("Initializer failed.", exception.Message); + + // The caller never received the connection, so the connector must dispose it itself. + Assert.NotNull(capturedConnection); + Assert.Equal(ConnectionState.Closed, capturedConnection!.State); + } + + public void Dispose() + { + File.Delete(_databasePath); + } + + private sealed class TestRecord + { + [VectorStoreKey] + public string? Key { get; set; } + + [VectorStoreData] + public string? Text { get; set; } + + [VectorStoreVector(dimensions: 4, DistanceFunction = DistanceFunction.EuclideanDistance)] + public ReadOnlyMemory? Vector { get; set; } + } +}