SqliteVec: support per-connection initialization callbacks

MEVD/src/SqliteVec/SqliteCollection.cs

@@ -31,6 +31,9 @@ public class SqliteCollection<TKey, TRecord> : VectorStoreCollection<TKey, TReco
     /// <summary>The connection string for the SQLite database represented by this <see cref="SqliteVectorStore"/>.</summary>
     private readonly string _connectionString;
 
+    /// <summary>Optional callback that runs on each connection after it is opened and the sqlite-vec extension is loaded.</summary>
+    private readonly Action<SqliteConnection>? _connectionInitializer;
+
     /// <summary>The mapper to use when mapping between the consumer data model and the SQLite record.</summary>
     private readonly SqliteMapper<TRecord> _mapper;
 
@@ -102,6 +105,7 @@ internal SqliteCollection(string connectionString, string name, Func<SqliteColle
 
         // Assign.
         _connectionString = connectionString;
+        _connectionInitializer = options.ConnectionInitializer;
         Name = name;
         _model = modelFactory(options);
 
@@ -389,13 +393,12 @@ public override async Task DeleteAsync(IEnumerable<TKey> keys, CancellationToken
 
     #region private
 
-    private async ValueTask<SqliteConnection> GetConnectionAsync(CancellationToken cancellationToken = default)
-    {
-        var connection = new SqliteConnection(_connectionString);
-        await connection.OpenAsync(cancellationToken).ConfigureAwait(false);
-        connection.LoadVector();
-        return connection;
-    }
+    private ValueTask<SqliteConnection> GetConnectionAsync(CancellationToken cancellationToken = default)
+        => SqliteConnectionFactory.OpenAsync(
+            _connectionString,
+            _connectionInitializer,
+            loadVector: true,
+            cancellationToken: cancellationToken);
 
     private async IAsyncEnumerable<VectorSearchResult<TRecord>> EnumerateAndMapSearchResultsAsync(
         List<SqliteWhereCondition> conditions,

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;
     }
 
     /// <summary>
@@ -31,4 +33,28 @@ internal SqliteCollectionOptions(SqliteCollectionOptions? source) : base(source)
     /// If not provided, collection name with prefix will be used as virtual table name.
     /// </remarks>
     public string? VectorVirtualTableName { get; set; }
+
+    /// <summary>
+    /// Gets or sets a callback that runs on every connection this collection opens.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// If defined, the callback runs after the connection is opened and the sqlite-vec extension is loaded.
+    /// </para>
+    /// <para>
+    /// Use it for per-connection settings that have no connection string equivalent, such as
+    /// <c>PRAGMA synchronous=NORMAL</c>, <c>PRAGMA cache_size</c> or <c>PRAGMA temp_store=MEMORY</c>, or for loading additional extensions.
+    /// </para>
+    /// <para>
+    /// When the collection is created from a <see cref="SqliteVectorStore"/>, the callback configured on the vector store options
+    /// is used; this property applies to collections constructed directly.
+    /// </para>
+    /// <para>
+    /// 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.
+    /// </para>
+    /// </remarks>
+    public Action<SqliteConnection>? ConnectionInitializer { get; set; }
 }

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;
+
+/// <summary>
+/// Creates the connections used by the vector store and its collections.
+/// </summary>
+internal static class SqliteConnectionFactory
+{
+    /// <summary>
+    /// Opens a connection, runs the shared setup steps in order, and hands it to the caller.
+    /// </summary>
+    /// <remarks>
+    /// The caller owns the returned connection. If any setup step fails, the connection is disposed here
+    /// because the caller never receives it.
+    /// </remarks>
+    public static async ValueTask<SqliteConnection> OpenAsync(
+        string connectionString,
+        Action<SqliteConnection>? 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;
+        }
+    }
+}

MEVD/src/SqliteVec/SqliteVec.csproj

@@ -1,7 +1,7 @@
 <Project Sdk="Microsoft.NET.Sdk">
 
   <PropertyGroup>
-    <Version>1.0.3-preview</Version>
+    <Version>1.0.4-preview</Version>
     <AssemblyName>CommunityToolkit.VectorData.SqliteVec</AssemblyName>
     <RootNamespace>$(AssemblyName)</RootNamespace>
     <TargetFrameworks>net10.0;net8.0</TargetFrameworks>

MEVD/src/SqliteVec/SqliteVectorStore.cs

@@ -33,6 +33,9 @@ public sealed class SqliteVectorStore : VectorStore
 
     private readonly IEmbeddingGenerator? _embeddingGenerator;
 
+    /// <summary>Optional callback that runs on each connection after it is opened.</summary>
+    private readonly Action<SqliteConnection>? _connectionInitializer;
+
     /// <summary>
     /// Initializes a new instance of the <see cref="SqliteVectorStore"/> class.
     /// </summary>
@@ -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<TKey, TRecord> GetCollection<TKey, TRecord>(str
                 {
                     Definition = definition,
                     VectorVirtualTableName = _vectorVirtualTableName,
-                    EmbeddingGenerator = _embeddingGenerator
+                    EmbeddingGenerator = _embeddingGenerator,
+                    ConnectionInitializer = _connectionInitializer
                 });
 
     /// <inheritdoc />
@@ -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<string> 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;

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;
     }
 
     /// <summary>
@@ -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.
     /// </summary>
     public IEmbeddingGenerator? EmbeddingGenerator { get; set; }
+
+    /// <summary>
+    /// Gets or sets a callback that runs on every connection this vector store and the collections created from it open.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// If defined, the callback runs after the connection is opened and, on connections used for collection operations, after the sqlite-vec
+    /// extension is loaded.
+    /// </para>
+    /// <para>
+    /// Use it for per-connection settings that have no connection string equivalent, such as
+    /// <c>PRAGMA synchronous=NORMAL</c>, <c>PRAGMA cache_size</c> or <c>PRAGMA temp_store=MEMORY</c>, or for loading additional extensions.
+    /// </para>
+    /// <para>
+    /// 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.
+    /// </para>
+    /// </remarks>
+    public Action<SqliteConnection>? ConnectionInitializer { get; set; }
 }