@mastra/mssql 1.0.0-beta.12 → 1.0.0-beta.13

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
@@ -2262,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",
@@ -3661,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;
@@ -3732,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"),