@mastra/mssql 1.0.0-beta.9 → 1.0.0

package/dist/index.js

@@ -1,5 +1,5 @@
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
-import { MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, TABLE_SCHEMAS, createStorageErrorId, normalizePerPage, calculatePagination, ObservabilityStorage, TABLE_SPANS, SPAN_SCHEMA, listTracesArgsSchema, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraStorage, TraceStatus, getDefaultValue, transformScoreRow as transformScoreRow$1 } from '@mastra/core/storage';
+import { MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, TABLE_SCHEMAS, createStorageErrorId, normalizePerPage, calculatePagination, ObservabilityStorage, TABLE_SPANS, SPAN_SCHEMA, listTracesArgsSchema, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraCompositeStore, TraceStatus, getDefaultValue, transformScoreRow as transformScoreRow$1 } from '@mastra/core/storage';
 import sql from 'mssql';
 import { MessageList } from '@mastra/core/agent';
 import { MastraBase } from '@mastra/core/base';
@@ -334,15 +334,49 @@ ${columns}
         );
         const pkExists = Array.isArray(pkResult.recordset) && pkResult.recordset.length > 0;
         if (!pkExists) {
-          try {
-            const addPkSql = `ALTER TABLE ${getTableName({ indexName: tableName, schemaName: getSchemaName(this.schemaName) })} ADD CONSTRAINT [${pkConstraintName}] PRIMARY KEY ([traceId], [spanId])`;
-            await this.pool.request().query(addPkSql);
-          } catch (pkError) {
-            this.logger?.warn?.(`Failed to add composite primary key to spans table:`, pkError);
+          const duplicateInfo = await this.checkForDuplicateSpans();
+          if (duplicateInfo.hasDuplicates) {
+            const errorMessage = `
+===========================================================================
+MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}
+===========================================================================
+
+Found ${duplicateInfo.duplicateCount} duplicate (traceId, spanId) combinations.
+
+The spans table requires a unique constraint on (traceId, spanId), but your
+database contains duplicate entries that must be resolved first.
+
+To fix this, run the manual migration command:
+
+  npx mastra migrate
+
+This command will:
+  1. Remove duplicate spans (keeping the most complete/recent version)
+  2. Add the required unique constraint
+
+Note: This migration may take some time for large tables.
+===========================================================================
+`;
+            throw new MastraError({
+              id: createStorageErrorId("MSSQL", "MIGRATION_REQUIRED", "DUPLICATE_SPANS"),
+              domain: ErrorDomain.STORAGE,
+              category: ErrorCategory.USER,
+              text: errorMessage
+            });
+          } else {
+            try {
+              const addPkSql = `ALTER TABLE ${getTableName({ indexName: tableName, schemaName: getSchemaName(this.schemaName) })} ADD CONSTRAINT [${pkConstraintName}] PRIMARY KEY ([traceId], [spanId])`;
+              await this.pool.request().query(addPkSql);
+            } catch (pkError) {
+              this.logger?.warn?.(`Failed to add composite primary key to spans table:`, pkError);
+            }
           }
         }
       }
     } catch (error) {
+      if (error instanceof MastraError) {
+        throw error;
+      }
       throw new MastraError(
         {
           id: createStorageErrorId("MSSQL", "CREATE_TABLE", "FAILED"),
@@ -384,6 +418,154 @@ ${columns}
       this.logger?.warn?.(`Failed to migrate spans table ${fullTableName}:`, error);
     }
   }
+  /**
+   * Deduplicates spans with the same (traceId, spanId) combination.
+   * This is needed for databases that existed before the unique constraint was added.
+   *
+   * Priority for keeping spans:
+   * 1. Completed spans (endedAt IS NOT NULL) over incomplete spans
+   * 2. Most recent updatedAt
+   * 3. Most recent createdAt (as tiebreaker)
+   *
+   * Note: This prioritizes migration completion over perfect data preservation.
+   * Old trace data may be lost, which is acceptable for this use case.
+   */
+  async deduplicateSpans() {
+    const fullTableName = getTableName({ indexName: TABLE_SPANS, schemaName: getSchemaName(this.schemaName) });
+    try {
+      const duplicateCheck = await this.pool.request().query(`
+        SELECT TOP 1 1 as has_duplicates
+        FROM ${fullTableName}
+        GROUP BY [traceId], [spanId]
+        HAVING COUNT(*) > 1
+      `);
+      if (!duplicateCheck.recordset || duplicateCheck.recordset.length === 0) {
+        this.logger?.debug?.(`No duplicate spans found in ${fullTableName}`);
+        return;
+      }
+      this.logger?.info?.(`Duplicate spans detected in ${fullTableName}, starting deduplication...`);
+      const result = await this.pool.request().query(`
+        WITH RankedSpans AS (
+          SELECT *, ROW_NUMBER() OVER (
+            PARTITION BY [traceId], [spanId]
+            ORDER BY
+              CASE WHEN [endedAt] IS NOT NULL THEN 0 ELSE 1 END,
+              [updatedAt] DESC,
+              [createdAt] DESC
+          ) as rn
+          FROM ${fullTableName}
+        )
+        DELETE FROM RankedSpans WHERE rn > 1
+      `);
+      this.logger?.info?.(
+        `Deduplication complete: removed ${result.rowsAffected?.[0] ?? 0} duplicate spans from ${fullTableName}`
+      );
+    } catch (error) {
+      this.logger?.warn?.("Failed to deduplicate spans:", error);
+    }
+  }
+  /**
+   * Checks for duplicate (traceId, spanId) combinations in the spans table.
+   * Returns information about duplicates for logging/CLI purposes.
+   */
+  async checkForDuplicateSpans() {
+    const fullTableName = getTableName({ indexName: TABLE_SPANS, schemaName: getSchemaName(this.schemaName) });
+    try {
+      const result = await this.pool.request().query(`
+        SELECT COUNT(*) as duplicate_count
+        FROM (
+          SELECT [traceId], [spanId]
+          FROM ${fullTableName}
+          GROUP BY [traceId], [spanId]
+          HAVING COUNT(*) > 1
+        ) duplicates
+      `);
+      const duplicateCount = result.recordset?.[0]?.duplicate_count ?? 0;
+      return {
+        hasDuplicates: duplicateCount > 0,
+        duplicateCount,
+        tableName: fullTableName
+      };
+    } catch (error) {
+      this.logger?.debug?.(`Could not check for duplicates: ${error}`);
+      return { hasDuplicates: false, duplicateCount: 0, tableName: fullTableName };
+    }
+  }
+  /**
+   * Checks if the PRIMARY KEY constraint on (traceId, spanId) already exists on the spans table.
+   */
+  async spansPrimaryKeyExists() {
+    const schemaPrefix = this.schemaName ? `${parseSqlIdentifier(this.schemaName, "schema name")}_` : "";
+    const pkConstraintName = `${schemaPrefix}mastra_ai_spans_traceid_spanid_pk`;
+    const checkPkRequest = this.pool.request();
+    checkPkRequest.input("constraintName", pkConstraintName);
+    const pkResult = await checkPkRequest.query(
+      `SELECT 1 AS found FROM sys.key_constraints WHERE name = @constraintName`
+    );
+    return Array.isArray(pkResult.recordset) && pkResult.recordset.length > 0;
+  }
+  /**
+   * Manually run the spans migration to deduplicate and add the unique constraint.
+   * This is intended to be called from the CLI when duplicates are detected.
+   *
+   * @returns Migration result with status and details
+   */
+  async migrateSpans() {
+    const fullTableName = getTableName({ indexName: TABLE_SPANS, schemaName: getSchemaName(this.schemaName) });
+    const pkExists = await this.spansPrimaryKeyExists();
+    if (pkExists) {
+      return {
+        success: true,
+        alreadyMigrated: true,
+        duplicatesRemoved: 0,
+        message: `Migration already complete. PRIMARY KEY constraint exists on ${fullTableName}.`
+      };
+    }
+    const duplicateInfo = await this.checkForDuplicateSpans();
+    if (duplicateInfo.hasDuplicates) {
+      this.logger?.info?.(
+        `Found ${duplicateInfo.duplicateCount} duplicate (traceId, spanId) combinations. Starting deduplication...`
+      );
+      await this.deduplicateSpans();
+    } else {
+      this.logger?.info?.(`No duplicate spans found.`);
+    }
+    const schemaPrefix = this.schemaName ? `${parseSqlIdentifier(this.schemaName, "schema name")}_` : "";
+    const pkConstraintName = `${schemaPrefix}mastra_ai_spans_traceid_spanid_pk`;
+    const addPkSql = `ALTER TABLE ${fullTableName} ADD CONSTRAINT [${pkConstraintName}] PRIMARY KEY ([traceId], [spanId])`;
+    await this.pool.request().query(addPkSql);
+    return {
+      success: true,
+      alreadyMigrated: false,
+      duplicatesRemoved: duplicateInfo.duplicateCount,
+      message: duplicateInfo.hasDuplicates ? `Migration complete. Removed duplicates and added PRIMARY KEY constraint to ${fullTableName}.` : `Migration complete. Added PRIMARY KEY constraint to ${fullTableName}.`
+    };
+  }
+  /**
+   * Check migration status for the spans table.
+   * Returns information about whether migration is needed.
+   */
+  async checkSpansMigrationStatus() {
+    const fullTableName = getTableName({ indexName: TABLE_SPANS, schemaName: getSchemaName(this.schemaName) });
+    const pkExists = await this.spansPrimaryKeyExists();
+    if (pkExists) {
+      return {
+        needsMigration: false,
+        hasDuplicates: false,
+        duplicateCount: 0,
+        constraintExists: true,
+        tableName: fullTableName
+      };
+    }
+    const duplicateInfo = await this.checkForDuplicateSpans();
+    return {
+      needsMigration: true,
+      hasDuplicates: duplicateInfo.hasDuplicates,
+      duplicateCount: duplicateInfo.duplicateCount,
+      constraintExists: false,
+      tableName: fullTableName
+    };
+  }
   /**
    * Alters table schema to add columns if they don't exist
    * @param tableName Name of the table
@@ -1010,10 +1192,12 @@ function getTableName2({ indexName, schemaName }) {
 function buildDateRangeFilter(dateRange, fieldName) {
   const filters = {};
   if (dateRange?.start) {
-    filters[`${fieldName}_gte`] = dateRange.start;
+    const suffix = dateRange.startExclusive ? "_gt" : "_gte";
+    filters[`${fieldName}${suffix}`] = dateRange.start;
   }
   if (dateRange?.end) {
-    filters[`${fieldName}_lte`] = dateRange.end;
+    const suffix = dateRange.endExclusive ? "_lt" : "_lte";
+    filters[`${fieldName}${suffix}`] = dateRange.end;
   }
   return filters;
 }
@@ -1031,11 +1215,21 @@ function prepareWhereClause(filters, _schema) {
       const fieldName = key.slice(0, -4);
       conditions.push(`[${parseSqlIdentifier(fieldName, "field name")}] >= @${paramName}`);
       params[paramName] = value instanceof Date ? value.toISOString() : value;
+    } else if (key.endsWith("_gt")) {
+      const paramName = `p${paramIndex++}`;
+      const fieldName = key.slice(0, -3);
+      conditions.push(`[${parseSqlIdentifier(fieldName, "field name")}] > @${paramName}`);
+      params[paramName] = value instanceof Date ? value.toISOString() : value;
     } else if (key.endsWith("_lte")) {
       const paramName = `p${paramIndex++}`;
       const fieldName = key.slice(0, -4);
       conditions.push(`[${parseSqlIdentifier(fieldName, "field name")}] <= @${paramName}`);
       params[paramName] = value instanceof Date ? value.toISOString() : value;
+    } else if (key.endsWith("_lt")) {
+      const paramName = `p${paramIndex++}`;
+      const fieldName = key.slice(0, -3);
+      conditions.push(`[${parseSqlIdentifier(fieldName, "field name")}] < @${paramName}`);
+      params[paramName] = value instanceof Date ? value.toISOString() : value;
     } else if (value === null) {
       conditions.push(`[${parseSqlIdentifier(key, "field name")}] IS NULL`);
     } else if (isInOperator(value)) {
@@ -1247,28 +1441,76 @@ var MemoryMSSQL = class _MemoryMSSQL extends MemoryStorage {
       );
     }
   }
-  async listThreadsByResourceId(args) {
-    const { resourceId, page = 0, perPage: perPageInput, orderBy } = args;
-    if (page < 0) {
+  async listThreads(args) {
+    const { page = 0, perPage: perPageInput, orderBy, filter } = args;
+    try {
+      this.validatePaginationInput(page, perPageInput ?? 100);
+    } catch (error) {
       throw new MastraError({
-        id: createStorageErrorId("MSSQL", "LIST_THREADS_BY_RESOURCE_ID", "INVALID_PAGE"),
+        id: createStorageErrorId("MSSQL", "LIST_THREADS", "INVALID_PAGE"),
         domain: ErrorDomain.STORAGE,
         category: ErrorCategory.USER,
-        text: "Page number must be non-negative",
-        details: {
-          resourceId,
-          page
-        }
+        text: error instanceof Error ? error.message : "Invalid pagination parameters",
+        details: { page, ...perPageInput !== void 0 && { perPage: perPageInput } }
       });
     }
     const perPage = normalizePerPage(perPageInput, 100);
+    try {
+      this.validateMetadataKeys(filter?.metadata);
+    } catch (error) {
+      throw new MastraError({
+        id: createStorageErrorId("MSSQL", "LIST_THREADS", "INVALID_METADATA_KEY"),
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.USER,
+        text: error instanceof Error ? error.message : "Invalid metadata key",
+        details: { metadataKeys: filter?.metadata ? Object.keys(filter.metadata).join(", ") : "" }
+      });
+    }
     const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
     const { field, direction } = this.parseOrderBy(orderBy);
     try {
-      const baseQuery = `FROM ${getTableName2({ indexName: TABLE_THREADS, schemaName: getSchemaName2(this.schema) })} WHERE [resourceId] = @resourceId`;
+      const tableName = getTableName2({ indexName: TABLE_THREADS, schemaName: getSchemaName2(this.schema) });
+      const whereClauses = [];
+      const params = {};
+      if (filter?.resourceId) {
+        whereClauses.push("[resourceId] = @resourceId");
+        params.resourceId = filter.resourceId;
+      }
+      if (filter?.metadata && Object.keys(filter.metadata).length > 0) {
+        let metadataIndex = 0;
+        for (const [key, value] of Object.entries(filter.metadata)) {
+          if (value !== null && typeof value === "object") {
+            throw new MastraError({
+              id: createStorageErrorId("MSSQL", "LIST_THREADS", "INVALID_METADATA_VALUE"),
+              domain: ErrorDomain.STORAGE,
+              category: ErrorCategory.USER,
+              text: `Metadata filter value for key "${key}" must be a scalar type (string, number, boolean, or null), got ${Array.isArray(value) ? "array" : "object"}`,
+              details: { key, valueType: Array.isArray(value) ? "array" : "object" }
+            });
+          }
+          if (value === null) {
+            whereClauses.push(`JSON_VALUE(metadata, '$.${key}') IS NULL`);
+          } else {
+            const paramName = `metadata${metadataIndex}`;
+            whereClauses.push(`JSON_VALUE(metadata, '$.${key}') = @${paramName}`);
+            if (typeof value === "string") {
+              params[paramName] = value;
+            } else if (typeof value === "boolean") {
+              params[paramName] = value ? "true" : "false";
+            } else {
+              params[paramName] = String(value);
+            }
+          }
+          metadataIndex++;
+        }
+      }
+      const whereClause = whereClauses.length > 0 ? `WHERE ${whereClauses.join(" AND ")}` : "";
+      const baseQuery = `FROM ${tableName} ${whereClause}`;
       const countQuery = `SELECT COUNT(*) as count ${baseQuery}`;
       const countRequest = this.pool.request();
-      countRequest.input("resourceId", resourceId);
+      for (const [key, value] of Object.entries(params)) {
+        countRequest.input(key, value);
+      }
       const countResult = await countRequest.query(countQuery);
       const total = parseInt(countResult.recordset[0]?.count ?? "0", 10);
       if (total === 0) {
@@ -1285,7 +1527,9 @@ var MemoryMSSQL = class _MemoryMSSQL extends MemoryStorage {
       const limitValue = perPageInput === false ? total : perPage;
       const dataQuery = `SELECT id, [resourceId], title, metadata, [createdAt], [updatedAt] ${baseQuery} ORDER BY ${orderByField} ${dir} OFFSET @offset ROWS FETCH NEXT @perPage ROWS ONLY`;
       const dataRequest = this.pool.request();
-      dataRequest.input("resourceId", resourceId);
+      for (const [key, value] of Object.entries(params)) {
+        dataRequest.input(key, value);
+      }
       dataRequest.input("offset", offset);
       if (limitValue > 2147483647) {
         dataRequest.input("perPage", sql.BigInt, limitValue);
@@ -1308,13 +1552,17 @@ var MemoryMSSQL = class _MemoryMSSQL extends MemoryStorage {
         hasMore: perPageInput === false ? false : offset + perPage < total
       };
     } catch (error) {
+      if (error instanceof MastraError && error.category === ErrorCategory.USER) {
+        throw error;
+      }
       const mastraError = new MastraError(
         {
-          id: createStorageErrorId("MSSQL", "LIST_THREADS_BY_RESOURCE_ID", "FAILED"),
+          id: createStorageErrorId("MSSQL", "LIST_THREADS", "FAILED"),
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.THIRD_PARTY,
           details: {
-            resourceId,
+            ...filter?.resourceId && { resourceId: filter.resourceId },
+            hasMetadataFilter: !!filter?.metadata,
             page
           }
         },
@@ -2196,6 +2444,22 @@ var ObservabilityMSSQL = class _ObservabilityMSSQL extends ObservabilityStorage
   async dangerouslyClearAll() {
     await this.db.clearTable({ tableName: TABLE_SPANS });
   }
+  /**
+   * Manually run the spans migration to deduplicate and add the unique constraint.
+   * This is intended to be called from the CLI when duplicates are detected.
+   *
+   * @returns Migration result with status and details
+   */
+  async migrateSpans() {
+    return this.db.migrateSpans();
+  }
+  /**
+   * Check migration status for the spans table.
+   * Returns information about whether migration is needed.
+   */
+  async checkSpansMigrationStatus() {
+    return this.db.checkSpansMigrationStatus();
+  }
   get tracingStrategy() {
     return {
       preferred: "batch-with-updates",
@@ -3595,7 +3859,7 @@ var WorkflowsMSSQL = class _WorkflowsMSSQL extends WorkflowsStorage {
 var isPoolConfig = (config) => {
   return "pool" in config;
 };
-var MSSQLStore = class extends MastraStorage {
+var MSSQLStore = class extends MastraCompositeStore {
   pool;
   schema;
   isConnected = null;
@@ -3666,6 +3930,9 @@ var MSSQLStore = class extends MastraStorage {
       await super.init();
     } catch (error) {
       this.isConnected = null;
+      if (error instanceof MastraError) {
+        throw error;
+      }
       throw new MastraError(
         {
           id: createStorageErrorId("MSSQL", "INIT", "FAILED"),