diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..9fc4411 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,24 @@ +name: CI + +on: + push: + branches: + - master + - main + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Setup Go + uses: actions/setup-go@v5 + with: + go-version-file: go.mod + + - name: Run unit tests + run: go test ./... diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..7bef338 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,79 @@ +name: Release + +on: + push: + tags: + - "v*" + workflow_dispatch: + inputs: + version: + description: "Semver version (e.g. 0.2.0 or v0.2.0)" + required: true + type: string + +permissions: + contents: write + +jobs: + release: + runs-on: ubuntu-latest + + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Setup Go + uses: actions/setup-go@v5 + with: + go-version-file: go.mod + + - name: Run unit tests + run: go test ./... + + - name: Resolve release tag + id: tag + shell: bash + run: | + set -euo pipefail + + if [[ "${GITHUB_EVENT_NAME}" == "workflow_dispatch" ]]; then + raw="${{ inputs.version }}" + raw="${raw#v}" + if [[ ! "$raw" =~ ^[0-9]+\.[0-9]+\.[0-9]+(-[0-9A-Za-z.-]+)?(\+[0-9A-Za-z.-]+)?$ ]]; then + echo "Invalid version: $raw" + exit 1 + fi + tag="v${raw}" + else + tag="${GITHUB_REF_NAME}" + fi + + echo "tag_name=$tag" >> "$GITHUB_OUTPUT" + + - name: Create git tag (manual release) + if: github.event_name == 'workflow_dispatch' + shell: bash + run: | + set -euo pipefail + + tag="${{ steps.tag.outputs.tag_name }}" + + if git rev-parse "$tag" >/dev/null 2>&1; then + echo "Tag $tag already exists locally" + else + git tag "$tag" + fi + + if git ls-remote --tags origin "refs/tags/$tag" | grep -q "$tag"; then + echo "Tag $tag already exists on origin" + else + git push origin "$tag" + fi + + - name: Publish GitHub release + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ steps.tag.outputs.tag_name }} + generate_release_notes: true diff --git a/README.md b/README.md index cde35aa..3174db3 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,8 @@ A lightweight Go vector-store library inspired by `Microsoft.Extensions.VectorDa MVP scope: - Go 1.24.x -- Postgres + `pgvector` only +- Postgres + `pgvector` +- MSSQL connector - Record-based core API with optional typed codec wrapper This library can be used to build retrieval systems such as: @@ -21,12 +22,14 @@ This library can be used to build retrieval systems such as: - Go 1.24.x - PostgreSQL 16+ with `pgvector` extension +- SQL Server 2022+ (for `stores/mssql`) - Docker (optional, for integration tests and sample compose flows) ## Project layout - `vectordata`: backend-agnostic core interfaces, record model, filters, typed wrapper - `stores/postgres`: Postgres implementation with `pgxpool` +- `stores/mssql`: SQL Server implementation with `database/sql` - `samples`: runnable demos (see `samples/README.md`) - `docs`: architecture and implementation notes @@ -193,19 +196,46 @@ go test ./... ``` Notes: -- Integration tests start Postgres + pgvector automatically via Testcontainers +- Integration tests start Postgres/pgvector and SQL Server automatically via Testcontainers - Docker daemon must be available when running integration tests - Optional override: set `PGVECTOR_TEST_DSN` to use an existing Postgres instance instead of starting a container +- Optional override: set `MSSQL_TEST_DSN` to use an existing SQL Server instance instead of starting a container + +Run MSSQL integration tests against root compose service: + +```bash +docker compose up -d mssql +MSSQL_TEST_DSN="sqlserver://sa:YourStrong%21Passw0rd@localhost:14339?database=master&encrypt=disable" \ + go test -tags=integration ./stores/mssql +``` ## Docker Compose (optional) `docker-compose.yml` at the repository root is kept for manual local runs. - Use root `docker-compose.yml` when you want a persistent local Postgres+pgvector instance outside tests +- Root compose also includes SQL Server (`mssql`) for local MSSQL connector validation - Use Testcontainers (`go test -tags=integration ./...`) for integration tests - Sample apps have their own compose files at `samples/semantic-search/docker-compose.yml` and `samples/ragrimosa/docker-compose.yml` - Sample Dockerfiles use `golang:1.24-alpine` to match the repo Go version (`1.24.x`) +## Release Automation + +GitHub Actions workflows are configured for: + +- CI on pushes/PRs (`.github/workflows/ci.yml`) +- release publishing (`.github/workflows/release.yml`) + +Release options: + +1. Manual (recommended): run `Release` workflow via GitHub UI with `version` input (`0.2.0` or `v0.2.0`). +2. Tag-driven: push a semver tag and the workflow publishes release notes automatically: + +```bash +git tag v0.2.0 +git push origin v0.2.0 +``` + ## Samples - Overview and entry points: [`samples/README.md`](samples/README.md) diff --git a/docker-compose.yml b/docker-compose.yml index 56399e9..3d6a6a8 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -13,3 +13,13 @@ services: interval: 2s timeout: 2s retries: 30 + + mssql: + image: mcr.microsoft.com/mssql/server:2022-latest + container_name: go-vectorstore-mssql + environment: + ACCEPT_EULA: "Y" + MSSQL_PID: "Developer" + MSSQL_SA_PASSWORD: "YourStrong!Passw0rd" + ports: + - "14339:1433" diff --git a/docs/architecture.md b/docs/architecture.md index 1225ff0..43128bc 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -8,6 +8,7 @@ The project is split into two layers: - `vectordata`: backend-agnostic contracts and primitives - `stores/postgres`: PostgreSQL + pgvector implementation +- `stores/mssql`: SQL Server implementation This keeps the public API stable while allowing additional storage engines later. @@ -23,7 +24,7 @@ The base record type is `vectordata.Record`: Search returns `vectordata.SearchResult`: - `Record`: matched item -- `Distance`: raw pgvector distance +- `Distance`: backend-computed distance value - `Score`: normalized score derived from distance Score normalization: @@ -57,6 +58,11 @@ All methods require `context.Context`. - `EnsureExtension`: `true` - `StrictByDefault`: `true` +`mssql.StoreOptions` defaults (`mssql.DefaultStoreOptions()`): + +- `Schema`: `dbo` +- `StrictByDefault`: `true` + Collection defaults: - Metric defaults to cosine when omitted @@ -71,7 +77,7 @@ Other operational defaults: ## 5) Request Flow -### Ensure collection +### Postgres: Ensure collection `PostgresVectorStore.EnsureCollection`: @@ -90,7 +96,7 @@ metadata jsonb not null default '{}'::jsonb, content text ``` -### Insert / Upsert +### Postgres: Insert / Upsert Bulk writes are chunked and executed as parameterized SQL: @@ -103,7 +109,7 @@ Each record is validated before sending: - vector dimension match - metadata JSON serialization -### Search +### Postgres: Search `SearchByVector` builds a query plan, then executes it: @@ -121,6 +127,22 @@ Each record is validated before sending: The pgvector operators are documented in the [pgvector README](https://github.com/pgvector/pgvector#querying). +### MSSQL: Ensure / Search + +`MSSQLVectorStore.EnsureCollection`: + +1. Ensures target schema exists +2. Ensures internal collection metadata table exists +3. Creates or validates the collection table +4. Persists and validates dimension/metric metadata + +`MSSQLCollection.SearchByVector`: + +1. Loads records from SQL Server +2. Evaluates filters against records in-process +3. Computes distance in-process (cosine/l2/inner product) +4. Applies threshold, sorts by distance, returns top-k + ## 6) Filter System (AST -> SQL) Filters are represented as an AST in `vectordata`: @@ -147,6 +169,8 @@ Behavior details: JSON path extraction behavior comes from PostgreSQL [JSON/JSONB functions and operators](https://www.postgresql.org/docs/current/functions-json.html). +For MSSQL in this MVP, the same AST is evaluated in-process against loaded records. + ## 7) Schema Safety Modes `CollectionSpec.Mode` controls ensure behavior: @@ -238,7 +262,8 @@ Errors are wrapped with context so callers can use `errors.Is(...)` against base Current MVP scope: -- Postgres + pgvector only +- Postgres + pgvector backend +- MSSQL backend (vectors stored as JSON payloads) - single-vector column per collection - metadata filtering through a focused AST @@ -247,8 +272,10 @@ Future extension points: - additional backends under `stores/` - richer filter operators - optional reranking strategies +- native MSSQL vector indexing/query pushdown ## 13) References - pgvector project docs: https://github.com/pgvector/pgvector - PostgreSQL docs: https://www.postgresql.org/docs/current/index.html +- SQL Server docs: https://learn.microsoft.com/sql/sql-server/ diff --git a/go.mod b/go.mod index 23b6be7..14dbbfe 100644 --- a/go.mod +++ b/go.mod @@ -4,7 +4,11 @@ go 1.24.0 toolchain go1.24.9 -require github.com/jackc/pgx/v5 v5.7.2 +require ( + github.com/jackc/pgx/v5 v5.7.2 + github.com/microsoft/go-mssqldb v1.8.0 + github.com/testcontainers/testcontainers-go v0.33.0 +) require ( dario.cat/mergo v1.0.0 // indirect @@ -45,7 +49,6 @@ require ( github.com/shirou/gopsutil/v3 v3.23.12 // indirect github.com/shoenig/go-m1cpu v0.1.6 // indirect github.com/sirupsen/logrus v1.9.3 // indirect - github.com/testcontainers/testcontainers-go v0.33.0 // indirect github.com/tklauser/go-sysconf v0.3.12 // indirect github.com/tklauser/numcpus v0.6.1 // indirect github.com/yusufpapurcu/wmi v1.2.3 // indirect diff --git a/go.sum b/go.sum index 1c956d3..b400ddd 100644 --- a/go.sum +++ b/go.sum @@ -36,6 +36,10 @@ github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= +github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 h1:C/Fq8+Vm3x8+3xB3SO2/3A8Q+7WmfYI8MViMbum/v8A= +github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9/go.mod h1:9x5EikdFyA8Kj4m2Vj93QShv2N6Yf4H6J5D9Q1QzM4Y= +github.com/golang-sql/sqlexp v0.1.0 h1:DKM0nqnvWFfB6K6x2zI2Qf8F9Pwc7P0W8o2JQ7QY9MM= +github.com/golang-sql/sqlexp v0.1.0/go.mod h1:hxXw1K3Zk7K8Y6nG5N7f6d4rV1E8j9a7q6x5y4z3w2Q= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= @@ -57,6 +61,8 @@ github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY= github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0= +github.com/microsoft/go-mssqldb v1.8.0 h1:7cyZ/AT7ycDsEoWPIXibd+aVKFtteUNhDGf3aobP+tw= +github.com/microsoft/go-mssqldb v1.8.0/go.mod h1:6znkekS3T2vp0waiMhen4GPU1BiAsrP+iXHcE7a7rFo= github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0= github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo= github.com/moby/patternmatcher v0.6.0 h1:GmP9lR19aU5GqSSFko+5pRqHi+Ohk1O69aFiKkVGiPk= diff --git a/stores/mssql/collection.go b/stores/mssql/collection.go new file mode 100644 index 0000000..d0f8b76 --- /dev/null +++ b/stores/mssql/collection.go @@ -0,0 +1,397 @@ +package mssql + +import ( + "context" + "database/sql" + "errors" + "fmt" + "sort" + "strings" + + "github.com/gabisonia/go-vectorstore/vectordata" +) + +type writeMode int + +const ( + writeModeInsert writeMode = iota + writeModeUpsert +) + +// MSSQLCollection is a SQL Server-backed vector collection. +type MSSQLCollection struct { + store *MSSQLVectorStore + name string + dimension int + metric vectordata.DistanceMetric +} + +func (c *MSSQLCollection) Name() string { + return c.name +} + +func (c *MSSQLCollection) Dimension() int { + return c.dimension +} + +func (c *MSSQLCollection) Metric() vectordata.DistanceMetric { + return c.metric +} + +func (c *MSSQLCollection) Insert(ctx context.Context, records []vectordata.Record) error { + return c.writeRecords(ctx, records, writeModeInsert) +} + +func (c *MSSQLCollection) Upsert(ctx context.Context, records []vectordata.Record) error { + return c.writeRecords(ctx, records, writeModeUpsert) +} + +func (c *MSSQLCollection) Get(ctx context.Context, id string) (vectordata.Record, error) { + query := fmt.Sprintf("SELECT %s, %s, %s, %s FROM %s WHERE %s = @p1", + quoteIdent(idColumn), + quoteIdent(vectorColumn), + quoteIdent(metadataColumn), + quoteIdent(contentColumn), + c.tableName(), + quoteIdent(idColumn), + ) + + var out vectordata.Record + var vectorRaw string + var metadataRaw string + var content sql.NullString + + err := c.store.db.QueryRowContext(ctx, query, id).Scan(&out.ID, &vectorRaw, &metadataRaw, &content) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return vectordata.Record{}, vectordata.ErrNotFound + } + return vectordata.Record{}, err + } + + vector, err := parseVectorJSON(vectorRaw) + if err != nil { + return vectordata.Record{}, fmt.Errorf("decode vector: %w", err) + } + metadata, err := parseMetadataJSON(metadataRaw) + if err != nil { + return vectordata.Record{}, fmt.Errorf("decode metadata: %w", err) + } + + out.Vector = vector + out.Metadata = metadata + if content.Valid { + value := content.String + out.Content = &value + } + + return out, nil +} + +func (c *MSSQLCollection) Delete(ctx context.Context, ids []string) (int64, error) { + if len(ids) == 0 { + return 0, nil + } + + args := make([]any, 0, len(ids)) + placeholders := make([]string, 0, len(ids)) + for i, id := range ids { + placeholders = append(placeholders, fmt.Sprintf("@p%d", i+1)) + args = append(args, id) + } + + query := fmt.Sprintf("DELETE FROM %s WHERE %s IN (%s)", + c.tableName(), + quoteIdent(idColumn), + strings.Join(placeholders, ", "), + ) + result, err := c.store.db.ExecContext(ctx, query, args...) + if err != nil { + return 0, err + } + rowsAffected, err := result.RowsAffected() + if err != nil { + return 0, err + } + return rowsAffected, nil +} + +func (c *MSSQLCollection) Count(ctx context.Context, filter vectordata.Filter) (int64, error) { + if filter == nil { + query := fmt.Sprintf("SELECT COUNT(*) FROM %s", c.tableName()) + var count int64 + if err := c.store.db.QueryRowContext(ctx, query).Scan(&count); err != nil { + return 0, err + } + return count, nil + } + + records, err := c.loadRecords(ctx, false) + if err != nil { + return 0, err + } + + count := int64(0) + for _, record := range records { + matches, err := matchesFilter(filter, record) + if err != nil { + return 0, err + } + if matches { + count++ + } + } + return count, nil +} + +func (c *MSSQLCollection) SearchByVector(ctx context.Context, vector []float32, topK int, opts vectordata.SearchOptions) ([]vectordata.SearchResult, error) { + if topK <= 0 { + return nil, fmt.Errorf("topK must be > 0") + } + if err := c.validateVectorDimension(vector); err != nil { + return nil, err + } + + projection := resolveProjection(opts.Projection) + metric := defaultMetric(c.metric) + records, err := c.loadRecords(ctx, true) + if err != nil { + return nil, err + } + + results := make([]vectordata.SearchResult, 0, len(records)) + for _, record := range records { + if err := c.validateVectorDimension(record.Vector); err != nil { + return nil, fmt.Errorf("invalid stored vector for record %q: %w", record.ID, err) + } + + matches, err := matchesFilter(opts.Filter, record) + if err != nil { + return nil, err + } + if !matches { + continue + } + + distance, err := distanceBetween(metric, vector, record.Vector) + if err != nil { + return nil, err + } + if opts.Threshold != nil && distance > *opts.Threshold { + continue + } + + results = append(results, vectordata.SearchResult{ + Record: projectRecord(record, projection), + Distance: distance, + Score: vectordata.ScoreFromDistance(metric, distance), + }) + } + + sort.Slice(results, func(i, j int) bool { + if results[i].Distance == results[j].Distance { + return results[i].Record.ID < results[j].Record.ID + } + return results[i].Distance < results[j].Distance + }) + + if len(results) > topK { + results = results[:topK] + } + + return results, nil +} + +func (c *MSSQLCollection) EnsureIndexes(ctx context.Context, opts vectordata.IndexOptions) error { + // SQL Server backend stores vectors as JSON payloads in this MVP, so index options are no-op. + _ = ctx + _ = opts + return nil +} + +func (c *MSSQLCollection) writeRecords(ctx context.Context, records []vectordata.Record, mode writeMode) error { + if len(records) == 0 { + return nil + } + + insertQuery := fmt.Sprintf("INSERT INTO %s (%s, %s, %s, %s) VALUES (@p1, @p2, @p3, @p4)", + c.tableName(), + quoteIdent(idColumn), + quoteIdent(vectorColumn), + quoteIdent(metadataColumn), + quoteIdent(contentColumn), + ) + updateQuery := fmt.Sprintf("UPDATE %s SET %s = @p2, %s = @p3, %s = @p4 WHERE %s = @p1", + c.tableName(), + quoteIdent(vectorColumn), + quoteIdent(metadataColumn), + quoteIdent(contentColumn), + quoteIdent(idColumn), + ) + + for start := 0; start < len(records); start += maxRowsPerStatement { + end := start + maxRowsPerStatement + if end > len(records) { + end = len(records) + } + + tx, err := c.store.db.BeginTx(ctx, nil) + if err != nil { + return err + } + if err := c.writeBatch(ctx, tx, records[start:end], mode, insertQuery, updateQuery); err != nil { + _ = tx.Rollback() + return err + } + if err := tx.Commit(); err != nil { + _ = tx.Rollback() + return err + } + } + + return nil +} + +func (c *MSSQLCollection) writeBatch( + ctx context.Context, + tx *sql.Tx, + records []vectordata.Record, + mode writeMode, + insertQuery string, + updateQuery string, +) error { + for _, record := range records { + if strings.TrimSpace(record.ID) == "" { + return fmt.Errorf("record id is empty") + } + if err := c.validateVectorDimension(record.Vector); err != nil { + return err + } + + vectorPayload, err := vectorJSON(record.Vector) + if err != nil { + return fmt.Errorf("encode vector for record %q: %w", record.ID, err) + } + metadataPayload, err := metadataJSON(record.Metadata) + if err != nil { + return fmt.Errorf("encode metadata for record %q: %w", record.ID, err) + } + + var contentArg any + if record.Content != nil { + contentArg = *record.Content + } + + switch mode { + case writeModeInsert: + if _, err := tx.ExecContext(ctx, insertQuery, record.ID, vectorPayload, metadataPayload, contentArg); err != nil { + return err + } + case writeModeUpsert: + result, err := tx.ExecContext(ctx, updateQuery, record.ID, vectorPayload, metadataPayload, contentArg) + if err != nil { + return err + } + + rowsAffected, err := result.RowsAffected() + if err != nil { + return err + } + if rowsAffected == 0 { + if _, err := tx.ExecContext(ctx, insertQuery, record.ID, vectorPayload, metadataPayload, contentArg); err != nil { + return err + } + } + default: + return fmt.Errorf("unsupported write mode %d", mode) + } + } + + return nil +} + +func (c *MSSQLCollection) loadRecords(ctx context.Context, includeVector bool) ([]vectordata.Record, error) { + selectColumns := []string{quoteIdent(idColumn)} + if includeVector { + selectColumns = append(selectColumns, quoteIdent(vectorColumn)) + } + selectColumns = append(selectColumns, quoteIdent(metadataColumn), quoteIdent(contentColumn)) + + query := fmt.Sprintf("SELECT %s FROM %s", strings.Join(selectColumns, ", "), c.tableName()) + rows, err := c.store.db.QueryContext(ctx, query) + if err != nil { + return nil, err + } + defer rows.Close() + + records := make([]vectordata.Record, 0) + for rows.Next() { + var record vectordata.Record + var vectorRaw string + var metadataRaw string + var content sql.NullString + + if includeVector { + if err := rows.Scan(&record.ID, &vectorRaw, &metadataRaw, &content); err != nil { + return nil, err + } + parsedVector, err := parseVectorJSON(vectorRaw) + if err != nil { + return nil, fmt.Errorf("decode vector: %w", err) + } + record.Vector = parsedVector + } else { + if err := rows.Scan(&record.ID, &metadataRaw, &content); err != nil { + return nil, err + } + } + + metadata, err := parseMetadataJSON(metadataRaw) + if err != nil { + return nil, fmt.Errorf("decode metadata: %w", err) + } + record.Metadata = metadata + + if content.Valid { + value := content.String + record.Content = &value + } + + records = append(records, record) + } + if err := rows.Err(); err != nil { + return nil, err + } + + return records, nil +} + +func (c *MSSQLCollection) validateVectorDimension(vector []float32) error { + if len(vector) != c.dimension { + return fmt.Errorf("%w: expected %d, got %d", vectordata.ErrDimensionMismatch, c.dimension, len(vector)) + } + return nil +} + +func (c *MSSQLCollection) tableName() string { + return qualifiedTable(c.store.opts.Schema, c.name) +} + +func projectRecord(record vectordata.Record, projection vectordata.Projection) vectordata.Record { + projected := vectordata.Record{ID: record.ID} + if projection.IncludeVector { + projected.Vector = append([]float32(nil), record.Vector...) + } + if projection.IncludeMetadata { + metadataCopy := make(map[string]any, len(record.Metadata)) + for key, value := range record.Metadata { + metadataCopy[key] = value + } + projected.Metadata = metadataCopy + } + if projection.IncludeContent && record.Content != nil { + contentCopy := *record.Content + projected.Content = &contentCopy + } + return projected +} diff --git a/stores/mssql/doc.go b/stores/mssql/doc.go new file mode 100644 index 0000000..f025ff7 --- /dev/null +++ b/stores/mssql/doc.go @@ -0,0 +1,2 @@ +// Package mssql provides a SQL Server-backed vectordata implementation. +package mssql diff --git a/stores/mssql/filter_eval.go b/stores/mssql/filter_eval.go new file mode 100644 index 0000000..4a84fb8 --- /dev/null +++ b/stores/mssql/filter_eval.go @@ -0,0 +1,207 @@ +package mssql + +import ( + "fmt" + "reflect" + "strings" + + "github.com/gabisonia/go-vectorstore/vectordata" +) + +func matchesFilter(filter vectordata.Filter, record vectordata.Record) (bool, error) { + if filter == nil { + return true, nil + } + + switch node := filter.(type) { + case vectordata.EqFilter: + left, exists, err := resolveFieldValue(node.Field, record) + if err != nil { + return false, err + } + if !exists { + return false, nil + } + return valuesEqual(left, node.Value), nil + case vectordata.InFilter: + if len(node.Values) == 0 { + return false, fmt.Errorf("%w: IN requires at least one value", vectordata.ErrInvalidFilter) + } + left, exists, err := resolveFieldValue(node.Field, record) + if err != nil { + return false, err + } + if !exists { + return false, nil + } + for _, value := range node.Values { + if valuesEqual(left, value) { + return true, nil + } + } + return false, nil + case vectordata.GtFilter: + left, exists, err := resolveFieldValue(node.Field, record) + if err != nil { + return false, err + } + if !exists { + return false, nil + } + comparison := compareValues(left, node.Value) + return comparison > 0, nil + case vectordata.LtFilter: + left, exists, err := resolveFieldValue(node.Field, record) + if err != nil { + return false, err + } + if !exists { + return false, nil + } + comparison := compareValues(left, node.Value) + return comparison < 0, nil + case vectordata.ExistsFilter: + _, exists, err := resolveFieldValue(node.Field, record) + if err != nil { + return false, err + } + return exists, nil + case vectordata.AndFilter: + if len(node.Children) == 0 { + return false, fmt.Errorf("%w: AND requires at least one child", vectordata.ErrInvalidFilter) + } + for _, child := range node.Children { + if child == nil { + return false, fmt.Errorf("%w: AND contains nil child", vectordata.ErrInvalidFilter) + } + ok, err := matchesFilter(child, record) + if err != nil { + return false, err + } + if !ok { + return false, nil + } + } + return true, nil + case vectordata.OrFilter: + if len(node.Children) == 0 { + return false, fmt.Errorf("%w: OR requires at least one child", vectordata.ErrInvalidFilter) + } + for _, child := range node.Children { + if child == nil { + return false, fmt.Errorf("%w: OR contains nil child", vectordata.ErrInvalidFilter) + } + ok, err := matchesFilter(child, record) + if err != nil { + return false, err + } + if ok { + return true, nil + } + } + return false, nil + case vectordata.NotFilter: + if node.Child == nil { + return false, fmt.Errorf("%w: NOT requires a child", vectordata.ErrInvalidFilter) + } + ok, err := matchesFilter(node.Child, record) + if err != nil { + return false, err + } + return !ok, nil + default: + return false, fmt.Errorf("%w: unsupported node type %T", vectordata.ErrInvalidFilter, filter) + } +} + +func resolveFieldValue(field vectordata.FieldRef, record vectordata.Record) (value any, exists bool, err error) { + switch field.Kind { + case vectordata.FieldColumn: + name := strings.TrimSpace(field.Name) + if name == "" { + return nil, false, fmt.Errorf("%w: column field name is empty", vectordata.ErrInvalidFilter) + } + + switch name { + case idColumn: + return record.ID, true, nil + case contentColumn: + if record.Content == nil { + return nil, false, nil + } + return *record.Content, true, nil + default: + return nil, false, fmt.Errorf("%w: unknown column %q", vectordata.ErrInvalidFilter, name) + } + case vectordata.FieldMetadata: + if len(field.Path) == 0 { + return nil, false, fmt.Errorf("%w: metadata path is empty", vectordata.ErrInvalidFilter) + } + if record.Metadata == nil { + return nil, false, nil + } + + var current any = record.Metadata + for _, segment := range field.Path { + key := strings.TrimSpace(segment) + if key == "" { + return nil, false, fmt.Errorf("%w: metadata path segment is empty", vectordata.ErrInvalidFilter) + } + + asMap, ok := current.(map[string]any) + if !ok { + return nil, false, nil + } + + next, ok := asMap[key] + if !ok { + return nil, false, nil + } + current = next + } + + return current, true, nil + default: + return nil, false, fmt.Errorf("%w: unsupported field kind %q", vectordata.ErrInvalidFilter, field.Kind) + } +} + +func valuesEqual(left, right any) bool { + if left == nil || right == nil { + return left == right + } + + leftNumeric, leftIsNumeric := toFloat64(left) + rightNumeric, rightIsNumeric := toFloat64(right) + if leftIsNumeric && rightIsNumeric { + return leftNumeric == rightNumeric + } + + return reflect.DeepEqual(left, right) +} + +func compareValues(left, right any) int { + leftNumeric, leftIsNumeric := toFloat64(left) + rightNumeric, rightIsNumeric := toFloat64(right) + if leftIsNumeric && rightIsNumeric { + switch { + case leftNumeric < rightNumeric: + return -1 + case leftNumeric > rightNumeric: + return 1 + default: + return 0 + } + } + + leftText := fmt.Sprint(left) + rightText := fmt.Sprint(right) + switch { + case leftText < rightText: + return -1 + case leftText > rightText: + return 1 + default: + return 0 + } +} diff --git a/stores/mssql/filter_eval_test.go b/stores/mssql/filter_eval_test.go new file mode 100644 index 0000000..89a6ae2 --- /dev/null +++ b/stores/mssql/filter_eval_test.go @@ -0,0 +1,113 @@ +package mssql + +import ( + "testing" + + "github.com/gabisonia/go-vectorstore/vectordata" +) + +func TestMatchesFilterMetadataAndComparison(t *testing.T) { + content := "payload" + record := vectordata.Record{ + ID: "doc-1", + Vector: []float32{1, 0}, + Content: &content, + Metadata: map[string]any{ + "category": "news", + "rank": 2, + "nested": map[string]any{ + "k": "v", + }, + }, + } + + filter := vectordata.And( + vectordata.Eq(vectordata.Metadata("category"), "news"), + vectordata.Gt(vectordata.Metadata("rank"), 1), + vectordata.Eq(vectordata.Metadata("nested", "k"), "v"), + ) + + matched, err := matchesFilter(filter, record) + if err != nil { + t.Fatalf("matchesFilter: %v", err) + } + if !matched { + t.Fatalf("expected record to match filter") + } +} + +func TestMatchesFilterColumnExists(t *testing.T) { + content := "x" + record := vectordata.Record{ID: "doc-1", Content: &content} + + matched, err := matchesFilter(vectordata.Exists(vectordata.Column(contentColumn)), record) + if err != nil { + t.Fatalf("matchesFilter: %v", err) + } + if !matched { + t.Fatalf("expected content exists filter to match") + } +} + +func TestMatchesFilterUnknownColumnReturnsError(t *testing.T) { + record := vectordata.Record{ID: "doc-1"} + _, err := matchesFilter(vectordata.Eq(vectordata.Column("unknown"), "x"), record) + if err == nil { + t.Fatalf("expected error for unknown column") + } +} + +func TestDistanceBetweenMetrics(t *testing.T) { + left := []float32{1, 0} + right := []float32{0.8, 0.2} + + cosineDistance, err := distanceBetween(vectordata.DistanceCosine, left, right) + if err != nil { + t.Fatalf("cosine distance: %v", err) + } + if cosineDistance < 0 { + t.Fatalf("expected non-negative cosine distance, got %f", cosineDistance) + } + + l2Distance, err := distanceBetween(vectordata.DistanceL2, left, right) + if err != nil { + t.Fatalf("l2 distance: %v", err) + } + if l2Distance <= 0 { + t.Fatalf("expected positive l2 distance, got %f", l2Distance) + } + + innerProductDistance, err := distanceBetween(vectordata.DistanceInnerProduct, left, right) + if err != nil { + t.Fatalf("inner product distance: %v", err) + } + if innerProductDistance >= 0 { + t.Fatalf("expected negative inner-product distance, got %f", innerProductDistance) + } +} + +func TestProjectRecordHonorsProjection(t *testing.T) { + content := "hello" + record := vectordata.Record{ + ID: "doc-1", + Vector: []float32{1, 2}, + Content: &content, + Metadata: map[string]any{"category": "news"}, + } + + projected := projectRecord(record, vectordata.Projection{ + IncludeVector: false, + IncludeMetadata: true, + IncludeContent: false, + }) + + if len(projected.Vector) != 0 { + t.Fatalf("expected vector to be omitted") + } + if projected.Metadata["category"] != "news" { + t.Fatalf("expected metadata to be present") + } + if projected.Content != nil { + t.Fatalf("expected content to be omitted") + } +} diff --git a/stores/mssql/helpers.go b/stores/mssql/helpers.go new file mode 100644 index 0000000..deb5bdb --- /dev/null +++ b/stores/mssql/helpers.go @@ -0,0 +1,216 @@ +package mssql + +import ( + "encoding/json" + "fmt" + "math" + "strconv" + "strings" + + "github.com/gabisonia/go-vectorstore/vectordata" +) + +const ( + idColumn = "id" + vectorColumn = "vector" + metadataColumn = "metadata" + contentColumn = "content" + collectionMetaTable = "__vector_collections" + maxRowsPerStatement = 500 +) + +func quoteIdent(ident string) string { + return "[" + strings.ReplaceAll(ident, "]", "]]") + "]" +} + +func qualifiedTable(schema, table string) string { + return quoteIdent(schema) + "." + quoteIdent(table) +} + +func objectIDName(schema, table string) string { + return quoteIdent(schema) + "." + quoteIdent(table) +} + +func escapeSQLString(value string) string { + return strings.ReplaceAll(value, "'", "''") +} + +func defaultMetric(metric vectordata.DistanceMetric) vectordata.DistanceMetric { + if metric == "" { + return vectordata.DistanceCosine + } + return metric +} + +func defaultMode(mode vectordata.EnsureMode, strictByDefault bool) vectordata.EnsureMode { + if mode != "" { + return mode + } + if strictByDefault { + return vectordata.EnsureStrict + } + return vectordata.EnsureAutoMigrate +} + +func vectorJSON(vector []float32) (string, error) { + payload, err := json.Marshal(vector) + if err != nil { + return "", err + } + return string(payload), nil +} + +func parseVectorJSON(raw string) ([]float32, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil, nil + } + + var values []float64 + if err := json.Unmarshal([]byte(raw), &values); err != nil { + return nil, err + } + + out := make([]float32, 0, len(values)) + for _, value := range values { + out = append(out, float32(value)) + } + return out, nil +} + +func normalizeMetadata(metadata map[string]any) map[string]any { + if metadata == nil { + return map[string]any{} + } + return metadata +} + +func metadataJSON(metadata map[string]any) (string, error) { + payload, err := json.Marshal(normalizeMetadata(metadata)) + if err != nil { + return "", err + } + return string(payload), nil +} + +func parseMetadataJSON(raw string) (map[string]any, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return map[string]any{}, nil + } + + var metadata map[string]any + if err := json.Unmarshal([]byte(raw), &metadata); err != nil { + return nil, err + } + if metadata == nil { + return map[string]any{}, nil + } + return metadata, nil +} + +func resolveProjection(projection *vectordata.Projection) vectordata.Projection { + if projection == nil { + return vectordata.DefaultProjection() + } + return *projection +} + +func distanceBetween(metric vectordata.DistanceMetric, query, candidate []float32) (float64, error) { + if len(query) != len(candidate) { + return 0, fmt.Errorf("%w: expected %d, got %d", vectordata.ErrDimensionMismatch, len(query), len(candidate)) + } + + switch metric { + case vectordata.DistanceCosine: + return cosineDistance(query, candidate), nil + case vectordata.DistanceL2: + return l2Distance(query, candidate), nil + case vectordata.DistanceInnerProduct: + return -dot(query, candidate), nil + default: + return 0, fmt.Errorf("%w: unsupported distance metric %q", vectordata.ErrSchemaMismatch, metric) + } +} + +func cosineDistance(left, right []float32) float64 { + leftNorm := norm(left) + rightNorm := norm(right) + if leftNorm == 0 || rightNorm == 0 { + return 1 + } + similarity := dot(left, right) / (leftNorm * rightNorm) + return 1 - similarity +} + +func l2Distance(left, right []float32) float64 { + sum := 0.0 + for i := range left { + delta := float64(left[i] - right[i]) + sum += delta * delta + } + return math.Sqrt(sum) +} + +func dot(left, right []float32) float64 { + sum := 0.0 + for i := range left { + sum += float64(left[i] * right[i]) + } + return sum +} + +func norm(vector []float32) float64 { + sum := 0.0 + for _, value := range vector { + f := float64(value) + sum += f * f + } + return math.Sqrt(sum) +} + +func isStringType(dataType string) bool { + switch strings.ToLower(strings.TrimSpace(dataType)) { + case "varchar", "nvarchar", "char", "nchar", "text", "ntext": + return true + default: + return false + } +} + +func toFloat64(value any) (float64, bool) { + switch n := value.(type) { + case int: + return float64(n), true + case int8: + return float64(n), true + case int16: + return float64(n), true + case int32: + return float64(n), true + case int64: + return float64(n), true + case uint: + return float64(n), true + case uint8: + return float64(n), true + case uint16: + return float64(n), true + case uint32: + return float64(n), true + case uint64: + return float64(n), true + case float32: + return float64(n), true + case float64: + return n, true + case json.Number: + parsed, err := strconv.ParseFloat(string(n), 64) + if err != nil { + return 0, false + } + return parsed, true + default: + return 0, false + } +} diff --git a/stores/mssql/mssql_integration_test.go b/stores/mssql/mssql_integration_test.go new file mode 100644 index 0000000..abfbf26 --- /dev/null +++ b/stores/mssql/mssql_integration_test.go @@ -0,0 +1,460 @@ +//go:build integration + +package mssql + +import ( + "context" + "database/sql" + "fmt" + "net" + "net/url" + "os" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/gabisonia/go-vectorstore/vectordata" + _ "github.com/microsoft/go-mssqldb" + testcontainers "github.com/testcontainers/testcontainers-go" + "github.com/testcontainers/testcontainers-go/wait" +) + +const ( + integrationMSSQLUser = "sa" + integrationMSSQLPassword = "YourStrong!Passw0rd" + integrationMSSQLDatabase = "vectorstore_test" +) + +var ( + schemaSeq atomic.Uint64 + integrationDSN string + integrationContainer testcontainers.Container +) + +func TestMain(m *testing.M) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + dsn := strings.TrimSpace(os.Getenv("MSSQL_TEST_DSN")) + if dsn == "" { + container, generatedDSN, err := startMSSQLContainer(ctx) + if err != nil { + fmt.Fprintf(os.Stderr, "failed to start integration container: %v\n", err) + os.Exit(1) + } + integrationContainer = container + integrationDSN = generatedDSN + } else { + integrationDSN = dsn + } + + exitCode := m.Run() + + if integrationContainer != nil { + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) + defer shutdownCancel() + if err := integrationContainer.Terminate(shutdownCtx); err != nil { + fmt.Fprintf(os.Stderr, "failed to terminate integration container: %v\n", err) + if exitCode == 0 { + exitCode = 1 + } + } + } + + os.Exit(exitCode) +} + +func startMSSQLContainer(ctx context.Context) (testcontainers.Container, string, error) { + request := testcontainers.ContainerRequest{ + Image: "mcr.microsoft.com/mssql/server:2022-latest", + ExposedPorts: []string{"1433/tcp"}, + Env: map[string]string{ + "ACCEPT_EULA": "Y", + "MSSQL_SA_PASSWORD": integrationMSSQLPassword, + "MSSQL_PID": "Developer", + }, + WaitingFor: wait.ForLog("SQL Server is now ready for client connections"). + WithStartupTimeout(4 * time.Minute), + } + + container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ + ContainerRequest: request, + Started: true, + }) + if err != nil { + return nil, "", fmt.Errorf("start sql server container: %w", err) + } + + host, err := container.Host(ctx) + if err != nil { + _ = container.Terminate(context.Background()) + return nil, "", fmt.Errorf("resolve container host: %w", err) + } + + mappedPort, err := container.MappedPort(ctx, "1433/tcp") + if err != nil { + _ = container.Terminate(context.Background()) + return nil, "", fmt.Errorf("resolve container port: %w", err) + } + + masterDSN := buildMSSQLDSN(host, mappedPort.Port(), "master") + if err := waitForDatabase(ctx, masterDSN); err != nil { + _ = container.Terminate(context.Background()) + return nil, "", err + } + if err := ensureDatabase(ctx, masterDSN, integrationMSSQLDatabase); err != nil { + _ = container.Terminate(context.Background()) + return nil, "", err + } + + dsn := buildMSSQLDSN(host, mappedPort.Port(), integrationMSSQLDatabase) + if err := waitForDatabase(ctx, dsn); err != nil { + _ = container.Terminate(context.Background()) + return nil, "", err + } + + return container, dsn, nil +} + +func buildMSSQLDSN(host string, port string, database string) string { + u := &url.URL{ + Scheme: "sqlserver", + User: url.UserPassword(integrationMSSQLUser, integrationMSSQLPassword), + Host: net.JoinHostPort(host, port), + } + + query := u.Query() + query.Set("database", database) + query.Set("encrypt", "disable") + u.RawQuery = query.Encode() + + return u.String() +} + +func waitForDatabase(parent context.Context, dsn string) error { + ctx, cancel := context.WithTimeout(parent, 2*time.Minute) + defer cancel() + + for { + db, err := sql.Open("sqlserver", dsn) + if err == nil { + pingCtx, pingCancel := context.WithTimeout(ctx, 4*time.Second) + pingErr := db.PingContext(pingCtx) + pingCancel() + _ = db.Close() + if pingErr == nil { + return nil + } + } + + select { + case <-ctx.Done(): + if err != nil { + return fmt.Errorf("connect integration database: %w", err) + } + return fmt.Errorf("wait for integration database: %w", ctx.Err()) + case <-time.After(500 * time.Millisecond): + } + } +} + +func ensureDatabase(ctx context.Context, dsn string, database string) error { + db, err := sql.Open("sqlserver", dsn) + if err != nil { + return fmt.Errorf("connect master database: %w", err) + } + defer db.Close() + + query := fmt.Sprintf("IF DB_ID(N'%s') IS NULL CREATE DATABASE %s", escapeSQLString(database), quoteIdent(database)) + if _, err := db.ExecContext(ctx, query); err != nil { + return fmt.Errorf("ensure integration database: %w", err) + } + + return nil +} + +func integrationDB(t *testing.T) *sql.DB { + t.Helper() + + dsn := strings.TrimSpace(integrationDSN) + if dsn == "" { + t.Fatal("integration DSN is not initialized") + } + + db, err := sql.Open("sqlserver", dsn) + if err != nil { + t.Fatalf("open db: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := db.PingContext(ctx); err != nil { + _ = db.Close() + t.Fatalf("ping db: %v", err) + } + + t.Cleanup(func() { + _ = db.Close() + }) + + return db +} + +func newTestStore(t *testing.T, db *sql.DB) *MSSQLVectorStore { + t.Helper() + + seq := schemaSeq.Add(1) + schema := fmt.Sprintf("it_%d_%d", time.Now().UnixNano(), seq) + schema = strings.ReplaceAll(schema, "-", "_") + + store, err := NewVectorStore(db, StoreOptions{ + Schema: schema, + StrictByDefault: true, + }) + if err != nil { + t.Fatalf("NewVectorStore: %v", err) + } + + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + cleanupQuery := fmt.Sprintf(` + DECLARE @schema SYSNAME = N'%s'; + DECLARE @dropSql NVARCHAR(MAX) = N''; + SELECT @dropSql = @dropSql + N'DROP TABLE ' + QUOTENAME(SCHEMA_NAME(schema_id)) + N'.' + QUOTENAME(name) + N';' + FROM sys.tables + WHERE schema_id = SCHEMA_ID(@schema); + + IF LEN(@dropSql) > 0 + BEGIN + EXEC sp_executesql @dropSql; + END + + IF SCHEMA_ID(@schema) IS NOT NULL + BEGIN + EXEC(N'DROP SCHEMA ' + QUOTENAME(@schema)); + END + `, escapeSQLString(schema)) + _, _ = db.ExecContext(ctx, cleanupQuery) + }) + + return store +} + +func TestIntegrationEnsureCollection(t *testing.T) { + db := integrationDB(t) + store := newTestStore(t, db) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + + _, err := store.EnsureCollection(ctx, vectordata.CollectionSpec{ + Name: "docs", + Dimension: 3, + Metric: vectordata.DistanceCosine, + Mode: vectordata.EnsureStrict, + }) + if err != nil { + t.Fatalf("EnsureCollection: %v", err) + } + + var count int + err = db.QueryRowContext(ctx, ` + SELECT COUNT(1) + FROM INFORMATION_SCHEMA.TABLES + WHERE TABLE_SCHEMA = @p1 AND TABLE_NAME = @p2 + `, store.opts.Schema, "docs").Scan(&count) + if err != nil { + t.Fatalf("query table exists: %v", err) + } + if count == 0 { + t.Fatalf("expected collection table to exist") + } + + dimension, metric, found, err := store.readCollectionMetadata(ctx, "docs") + if err != nil { + t.Fatalf("readCollectionMetadata: %v", err) + } + if !found { + t.Fatalf("expected collection metadata row to exist") + } + if dimension != 3 { + t.Fatalf("expected dimension 3, got %d", dimension) + } + if metric != vectordata.DistanceCosine { + t.Fatalf("expected metric %q, got %q", vectordata.DistanceCosine, metric) + } +} + +func TestIntegrationUpsertAndGet(t *testing.T) { + db := integrationDB(t) + store := newTestStore(t, db) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + + collection, err := store.EnsureCollection(ctx, vectordata.CollectionSpec{ + Name: "docs", + Dimension: 2, + Metric: vectordata.DistanceCosine, + Mode: vectordata.EnsureStrict, + }) + if err != nil { + t.Fatalf("EnsureCollection: %v", err) + } + + content := "first" + if err := collection.Upsert(ctx, []vectordata.Record{{ + ID: "r1", + Vector: []float32{1, 0}, + Content: &content, + Metadata: map[string]any{ + "category": "news", + "rank": 1, + }, + }}); err != nil { + t.Fatalf("Upsert: %v", err) + } + + record, err := collection.Get(ctx, "r1") + if err != nil { + t.Fatalf("Get: %v", err) + } + if record.ID != "r1" { + t.Fatalf("expected ID r1, got %q", record.ID) + } + if len(record.Vector) != 2 || record.Vector[0] != 1 || record.Vector[1] != 0 { + t.Fatalf("unexpected vector: %#v", record.Vector) + } + if record.Content == nil || *record.Content != "first" { + t.Fatalf("unexpected content: %#v", record.Content) + } + if record.Metadata["category"] != "news" { + t.Fatalf("unexpected metadata category: %#v", record.Metadata["category"]) + } + + updated := "updated" + if err := collection.Upsert(ctx, []vectordata.Record{{ + ID: "r1", + Vector: []float32{0.5, 0.5}, + Content: &updated, + Metadata: map[string]any{ + "category": "blog", + "rank": 3, + }, + }}); err != nil { + t.Fatalf("Upsert second call: %v", err) + } + + updatedRecord, err := collection.Get(ctx, "r1") + if err != nil { + t.Fatalf("Get after upsert: %v", err) + } + if updatedRecord.Content == nil || *updatedRecord.Content != "updated" { + t.Fatalf("content not updated: %#v", updatedRecord.Content) + } + if updatedRecord.Metadata["category"] != "blog" { + t.Fatalf("metadata not updated: %#v", updatedRecord.Metadata) + } +} + +func TestIntegrationSearchByMetric(t *testing.T) { + metrics := []vectordata.DistanceMetric{ + vectordata.DistanceCosine, + vectordata.DistanceL2, + vectordata.DistanceInnerProduct, + } + + for _, metric := range metrics { + t.Run(string(metric), func(t *testing.T) { + db := integrationDB(t) + store := newTestStore(t, db) + + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) + defer cancel() + + collection, err := store.EnsureCollection(ctx, vectordata.CollectionSpec{ + Name: "search_docs", + Dimension: 2, + Metric: metric, + Mode: vectordata.EnsureStrict, + }) + if err != nil { + t.Fatalf("EnsureCollection: %v", err) + } + + err = collection.Upsert(ctx, []vectordata.Record{ + {ID: "a", Vector: []float32{1, 0}, Metadata: map[string]any{"kind": "a"}}, + {ID: "b", Vector: []float32{0.8, 0.2}, Metadata: map[string]any{"kind": "b"}}, + {ID: "c", Vector: []float32{0, 1}, Metadata: map[string]any{"kind": "c"}}, + }) + if err != nil { + t.Fatalf("Upsert: %v", err) + } + + results, err := collection.SearchByVector(ctx, []float32{1, 0}, 2, vectordata.SearchOptions{}) + if err != nil { + t.Fatalf("SearchByVector: %v", err) + } + if len(results) != 2 { + t.Fatalf("expected 2 results, got %d", len(results)) + } + if results[0].Record.ID != "a" || results[1].Record.ID != "b" { + t.Fatalf("unexpected ordering: [%s, %s]", results[0].Record.ID, results[1].Record.ID) + } + }) + } +} + +func TestIntegrationMetadataFilter(t *testing.T) { + db := integrationDB(t) + store := newTestStore(t, db) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + + collection, err := store.EnsureCollection(ctx, vectordata.CollectionSpec{ + Name: "docs", + Dimension: 2, + Metric: vectordata.DistanceCosine, + Mode: vectordata.EnsureStrict, + }) + if err != nil { + t.Fatalf("EnsureCollection: %v", err) + } + + err = collection.Upsert(ctx, []vectordata.Record{ + {ID: "a", Vector: []float32{1, 0}, Metadata: map[string]any{"category": "news", "rank": 1}}, + {ID: "b", Vector: []float32{0.9, 0.1}, Metadata: map[string]any{"category": "news", "rank": 2}}, + {ID: "c", Vector: []float32{0, 1}, Metadata: map[string]any{"category": "other", "rank": 3}}, + }) + if err != nil { + t.Fatalf("Upsert: %v", err) + } + + filter := vectordata.And( + vectordata.Eq(vectordata.Metadata("category"), "news"), + vectordata.Gt(vectordata.Metadata("rank"), 1), + ) + + results, searchErr := collection.SearchByVector(ctx, []float32{1, 0}, 10, vectordata.SearchOptions{Filter: filter}) + count, countErr := collection.Count(ctx, filter) + + if searchErr != nil { + t.Fatalf("SearchByVector with filter: %v", searchErr) + } + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + if results[0].Record.ID != "b" { + t.Fatalf("expected result b, got %s", results[0].Record.ID) + } + if countErr != nil { + t.Fatalf("Count with filter: %v", countErr) + } + if count != 1 { + t.Fatalf("expected count 1, got %d", count) + } +} diff --git a/stores/mssql/schema.go b/stores/mssql/schema.go new file mode 100644 index 0000000..207f63e --- /dev/null +++ b/stores/mssql/schema.go @@ -0,0 +1,298 @@ +package mssql + +import ( + "context" + "database/sql" + "fmt" + + "github.com/gabisonia/go-vectorstore/vectordata" +) + +func (s *MSSQLVectorStore) ensureBaseSchema(ctx context.Context) error { + schemaLiteral := escapeSQLString(s.opts.Schema) + query := fmt.Sprintf("IF SCHEMA_ID(N'%s') IS NULL EXEC(N'CREATE SCHEMA %s')", schemaLiteral, quoteIdent(s.opts.Schema)) + if _, err := s.db.ExecContext(ctx, query); err != nil { + return fmt.Errorf("ensure schema %q: %w", s.opts.Schema, err) + } + + if err := s.ensureCollectionsMetadataTable(ctx); err != nil { + return err + } + + return nil +} + +func (s *MSSQLVectorStore) ensureCollectionsMetadataTable(ctx context.Context) error { + query := fmt.Sprintf(` + IF OBJECT_ID(N'%s', N'U') IS NULL + BEGIN + CREATE TABLE %s ( + %s NVARCHAR(255) NOT NULL PRIMARY KEY, + %s INT NOT NULL, + %s NVARCHAR(64) NOT NULL + ) + END + `, + escapeSQLString(objectIDName(s.opts.Schema, collectionMetaTable)), + qualifiedTable(s.opts.Schema, collectionMetaTable), + quoteIdent("name"), + quoteIdent("dimension"), + quoteIdent("metric"), + ) + + if _, err := s.db.ExecContext(ctx, query); err != nil { + return fmt.Errorf("ensure collection metadata table: %w", err) + } + return nil +} + +func (s *MSSQLVectorStore) ensureTableWithValidation(ctx context.Context, spec vectordata.CollectionSpec, mode vectordata.EnsureMode) error { + exists, err := s.tableExists(ctx, spec.Name) + if err != nil { + return err + } + + if !exists { + if err := s.createCollectionTable(ctx, spec.Name); err != nil { + return err + } + if err := s.upsertCollectionMetadata(ctx, spec.Name, spec.Dimension, spec.Metric); err != nil { + return err + } + return nil + } + + if err := s.validateCollectionSchema(ctx, spec.Name, mode); err != nil { + return err + } + + dimension, metric, found, err := s.readCollectionMetadata(ctx, spec.Name) + if err != nil { + return err + } + if !found { + if mode == vectordata.EnsureStrict { + return fmt.Errorf("%w: missing collection metadata for %q", vectordata.ErrSchemaMismatch, spec.Name) + } + return s.upsertCollectionMetadata(ctx, spec.Name, spec.Dimension, spec.Metric) + } + + if dimension != spec.Dimension { + return fmt.Errorf("%w: expected vector dimension %d, got %d", vectordata.ErrSchemaMismatch, spec.Dimension, dimension) + } + + if defaultMetric(metric) != defaultMetric(spec.Metric) { + return fmt.Errorf("%w: expected metric %q, got %q", vectordata.ErrSchemaMismatch, spec.Metric, metric) + } + + return nil +} + +func (s *MSSQLVectorStore) tableExists(ctx context.Context, table string) (bool, error) { + var count int + err := s.db.QueryRowContext(ctx, ` + SELECT COUNT(1) + FROM INFORMATION_SCHEMA.TABLES + WHERE TABLE_SCHEMA = @p1 AND TABLE_NAME = @p2 + `, s.opts.Schema, table).Scan(&count) + if err != nil { + return false, fmt.Errorf("check table exists: %w", err) + } + + return count > 0, nil +} + +func (s *MSSQLVectorStore) createCollectionTable(ctx context.Context, table string) error { + query := fmt.Sprintf(` + IF OBJECT_ID(N'%s', N'U') IS NULL + BEGIN + CREATE TABLE %s ( + %s NVARCHAR(255) NOT NULL PRIMARY KEY, + %s NVARCHAR(MAX) NOT NULL, + %s NVARCHAR(MAX) NOT NULL DEFAULT N'{}', + %s NVARCHAR(MAX) NULL + ) + END + `, + escapeSQLString(objectIDName(s.opts.Schema, table)), + qualifiedTable(s.opts.Schema, table), + quoteIdent(idColumn), + quoteIdent(vectorColumn), + quoteIdent(metadataColumn), + quoteIdent(contentColumn), + ) + + if _, err := s.db.ExecContext(ctx, query); err != nil { + return fmt.Errorf("create collection table %q: %w", table, err) + } + return nil +} + +func (s *MSSQLVectorStore) validateCollectionSchema(ctx context.Context, table string, mode vectordata.EnsureMode) error { + rows, err := s.db.QueryContext(ctx, ` + SELECT COLUMN_NAME, DATA_TYPE + FROM INFORMATION_SCHEMA.COLUMNS + WHERE TABLE_SCHEMA = @p1 AND TABLE_NAME = @p2 + `, s.opts.Schema, table) + if err != nil { + return fmt.Errorf("read schema columns: %w", err) + } + defer rows.Close() + + columns := make(map[string]string) + for rows.Next() { + var columnName string + var dataType string + if err := rows.Scan(&columnName, &dataType); err != nil { + return fmt.Errorf("scan schema columns: %w", err) + } + columns[columnName] = dataType + } + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate schema columns: %w", err) + } + + idType, ok := columns[idColumn] + if !ok { + return fmt.Errorf("%w: missing column %q", vectordata.ErrSchemaMismatch, idColumn) + } + if !isStringType(idType) { + return fmt.Errorf("%w: expected %q to be string-compatible type, got %q", vectordata.ErrSchemaMismatch, idColumn, idType) + } + + vectorType, ok := columns[vectorColumn] + if !ok { + return fmt.Errorf("%w: missing column %q", vectordata.ErrSchemaMismatch, vectorColumn) + } + if !isStringType(vectorType) { + return fmt.Errorf("%w: expected %q to be string-compatible type, got %q", vectordata.ErrSchemaMismatch, vectorColumn, vectorType) + } + + if err := s.ensurePrimaryKeyOnID(ctx, table); err != nil { + return err + } + + metadataType, hasMetadata := columns[metadataColumn] + if !hasMetadata { + if mode == vectordata.EnsureStrict { + return fmt.Errorf("%w: missing column %q", vectordata.ErrSchemaMismatch, metadataColumn) + } + if err := s.addMetadataColumn(ctx, table); err != nil { + return err + } + } else if !isStringType(metadataType) { + return fmt.Errorf("%w: expected %q to be string-compatible type, got %q", vectordata.ErrSchemaMismatch, metadataColumn, metadataType) + } + + contentType, hasContent := columns[contentColumn] + if !hasContent { + if mode == vectordata.EnsureStrict { + return fmt.Errorf("%w: missing column %q", vectordata.ErrSchemaMismatch, contentColumn) + } + if err := s.addContentColumn(ctx, table); err != nil { + return err + } + } else if !isStringType(contentType) { + return fmt.Errorf("%w: expected %q to be string-compatible type, got %q", vectordata.ErrSchemaMismatch, contentColumn, contentType) + } + + return nil +} + +func (s *MSSQLVectorStore) ensurePrimaryKeyOnID(ctx context.Context, table string) error { + var count int + err := s.db.QueryRowContext(ctx, ` + SELECT COUNT(1) + FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc + INNER JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE kcu + ON tc.CONSTRAINT_NAME = kcu.CONSTRAINT_NAME + AND tc.TABLE_SCHEMA = kcu.TABLE_SCHEMA + AND tc.TABLE_NAME = kcu.TABLE_NAME + WHERE tc.TABLE_SCHEMA = @p1 + AND tc.TABLE_NAME = @p2 + AND tc.CONSTRAINT_TYPE = 'PRIMARY KEY' + AND kcu.COLUMN_NAME = @p3 + `, s.opts.Schema, table, idColumn).Scan(&count) + if err != nil { + return fmt.Errorf("check primary key: %w", err) + } + if count == 0 { + return fmt.Errorf("%w: primary key on %q is required", vectordata.ErrSchemaMismatch, idColumn) + } + return nil +} + +func (s *MSSQLVectorStore) addMetadataColumn(ctx context.Context, table string) error { + query := fmt.Sprintf("ALTER TABLE %s ADD %s NVARCHAR(MAX) NOT NULL DEFAULT N'{}'", qualifiedTable(s.opts.Schema, table), quoteIdent(metadataColumn)) + if _, err := s.db.ExecContext(ctx, query); err != nil { + return fmt.Errorf("auto-migrate metadata column: %w", err) + } + return nil +} + +func (s *MSSQLVectorStore) addContentColumn(ctx context.Context, table string) error { + query := fmt.Sprintf("ALTER TABLE %s ADD %s NVARCHAR(MAX) NULL", qualifiedTable(s.opts.Schema, table), quoteIdent(contentColumn)) + if _, err := s.db.ExecContext(ctx, query); err != nil { + return fmt.Errorf("auto-migrate content column: %w", err) + } + return nil +} + +func (s *MSSQLVectorStore) readCollectionMetadata(ctx context.Context, collectionName string) (dimension int, metric vectordata.DistanceMetric, found bool, err error) { + query := fmt.Sprintf("SELECT %s, %s FROM %s WHERE %s = @p1", + quoteIdent("dimension"), + quoteIdent("metric"), + qualifiedTable(s.opts.Schema, collectionMetaTable), + quoteIdent("name"), + ) + + var metricRaw string + err = s.db.QueryRowContext(ctx, query, collectionName).Scan(&dimension, &metricRaw) + if err != nil { + if err == sql.ErrNoRows { + return 0, "", false, nil + } + return 0, "", false, fmt.Errorf("read collection metadata: %w", err) + } + + return dimension, vectordata.DistanceMetric(metricRaw), true, nil +} + +func (s *MSSQLVectorStore) upsertCollectionMetadata(ctx context.Context, collectionName string, dimension int, metric vectordata.DistanceMetric) error { + query := fmt.Sprintf(` + MERGE %s AS target + USING ( + SELECT @p1 AS %s, @p2 AS %s, @p3 AS %s + ) AS src + ON target.%s = src.%s + WHEN MATCHED THEN + UPDATE SET + target.%s = src.%s, + target.%s = src.%s + WHEN NOT MATCHED THEN + INSERT (%s, %s, %s) + VALUES (src.%s, src.%s, src.%s); + `, + qualifiedTable(s.opts.Schema, collectionMetaTable), + quoteIdent("name"), + quoteIdent("dimension"), + quoteIdent("metric"), + quoteIdent("name"), + quoteIdent("name"), + quoteIdent("dimension"), + quoteIdent("dimension"), + quoteIdent("metric"), + quoteIdent("metric"), + quoteIdent("name"), + quoteIdent("dimension"), + quoteIdent("metric"), + quoteIdent("name"), + quoteIdent("dimension"), + quoteIdent("metric"), + ) + + if _, err := s.db.ExecContext(ctx, query, collectionName, dimension, string(metric)); err != nil { + return fmt.Errorf("upsert collection metadata: %w", err) + } + return nil +} diff --git a/stores/mssql/store.go b/stores/mssql/store.go new file mode 100644 index 0000000..88526bc --- /dev/null +++ b/stores/mssql/store.go @@ -0,0 +1,111 @@ +package mssql + +import ( + "context" + "database/sql" + "fmt" + "strings" + + "github.com/gabisonia/go-vectorstore/vectordata" +) + +// StoreOptions configures MSSQLVectorStore behavior. +type StoreOptions struct { + Schema string + StrictByDefault bool +} + +// DefaultStoreOptions returns production-safe defaults. +func DefaultStoreOptions() StoreOptions { + return StoreOptions{ + Schema: "dbo", + StrictByDefault: true, + } +} + +// MSSQLVectorStore implements vectordata.VectorStore using database/sql. +type MSSQLVectorStore struct { + db *sql.DB + opts StoreOptions +} + +// NewVectorStore creates a SQL Server-backed vector store. +func NewVectorStore(db *sql.DB, opts StoreOptions) (*MSSQLVectorStore, error) { + if db == nil { + return nil, fmt.Errorf("nil sql db") + } + + normalized := opts.withDefaults() + if err := normalized.validate(); err != nil { + return nil, err + } + + return &MSSQLVectorStore{db: db, opts: normalized}, nil +} + +// Collection returns a handle to a collection without schema checks. +func (s *MSSQLVectorStore) Collection(name string, dimension int, metric vectordata.DistanceMetric) vectordata.Collection { + return s.newCollectionHandle(name, dimension, metric) +} + +// EnsureCollection creates or validates a collection schema and returns its handle. +func (s *MSSQLVectorStore) EnsureCollection(ctx context.Context, spec vectordata.CollectionSpec) (vectordata.Collection, error) { + normalizedSpec, mode, err := s.normalizeCollectionSpec(spec) + if err != nil { + return nil, err + } + + if err := s.ensureBaseSchema(ctx); err != nil { + return nil, err + } + if err := s.ensureTableWithValidation(ctx, normalizedSpec, mode); err != nil { + return nil, err + } + + return s.newCollectionHandle(normalizedSpec.Name, normalizedSpec.Dimension, normalizedSpec.Metric), nil +} + +func (s *MSSQLVectorStore) normalizeCollectionSpec(spec vectordata.CollectionSpec) (vectordata.CollectionSpec, vectordata.EnsureMode, error) { + spec.Name = strings.TrimSpace(spec.Name) + if spec.Name == "" { + return vectordata.CollectionSpec{}, "", fmt.Errorf("%w: collection name is empty", vectordata.ErrSchemaMismatch) + } + if spec.Dimension <= 0 { + return vectordata.CollectionSpec{}, "", fmt.Errorf("%w: dimension must be > 0", vectordata.ErrSchemaMismatch) + } + + spec.Metric = defaultMetric(spec.Metric) + if err := spec.Metric.Validate(); err != nil { + return vectordata.CollectionSpec{}, "", err + } + + mode := defaultMode(spec.Mode, s.opts.StrictByDefault) + if mode != vectordata.EnsureStrict && mode != vectordata.EnsureAutoMigrate { + return vectordata.CollectionSpec{}, "", fmt.Errorf("%w: unsupported ensure mode %q", vectordata.ErrSchemaMismatch, mode) + } + + return spec, mode, nil +} + +func (s *MSSQLVectorStore) newCollectionHandle(name string, dimension int, metric vectordata.DistanceMetric) vectordata.Collection { + return &MSSQLCollection{ + store: s, + name: strings.TrimSpace(name), + dimension: dimension, + metric: defaultMetric(metric), + } +} + +func (s StoreOptions) withDefaults() StoreOptions { + if strings.TrimSpace(s.Schema) == "" { + s.Schema = "dbo" + } + return s +} + +func (s StoreOptions) validate() error { + if strings.TrimSpace(s.Schema) == "" { + return fmt.Errorf("%w: schema is empty", vectordata.ErrSchemaMismatch) + } + return nil +}