@mastra/pg 1.0.0-beta.14 → 1.0.0-beta.15

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.js CHANGED
@@ -1,7 +1,7 @@
1
1
  import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
2
- import { createVectorErrorId, TABLE_SCHEMAS, AgentsStorage, TABLE_AGENTS, createStorageErrorId, normalizePerPage, calculatePagination, MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, ObservabilityStorage, TABLE_SPANS, listTracesArgsSchema, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraStorage, TraceStatus, getDefaultValue, transformScoreRow as transformScoreRow$1, getSqlType } from '@mastra/core/storage';
2
+ import { createVectorErrorId, TABLE_SCHEMAS, AgentsStorage, TABLE_AGENTS, createStorageErrorId, normalizePerPage, calculatePagination, MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, ObservabilityStorage, TABLE_SPANS, listTracesArgsSchema, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraCompositeStore, TraceStatus, getDefaultValue, transformScoreRow as transformScoreRow$1, getSqlType } from '@mastra/core/storage';
3
3
  import { parseSqlIdentifier, parseFieldKey } from '@mastra/core/utils';
4
- import { MastraVector } from '@mastra/core/vector';
4
+ import { MastraVector, validateTopK, validateUpsertInput } from '@mastra/core/vector';
5
5
  import { Mutex } from 'async-mutex';
6
6
  import * as pg from 'pg';
7
7
  import { Pool } from 'pg';
@@ -261,8 +261,14 @@ var FILTER_OPERATORS = {
261
261
  };
262
262
  },
263
263
  // Element Operators
264
- $exists: (key) => {
264
+ $exists: (key, paramIndex, value) => {
265
265
  const jsonPathKey = parseJsonPathKey(key);
266
+ if (value === false) {
267
+ return {
268
+ sql: `NOT (metadata ? '${jsonPathKey}')`,
269
+ needsValue: false
270
+ };
271
+ }
266
272
  return {
267
273
  sql: `metadata ? '${jsonPathKey}'`,
268
274
  needsValue: false
@@ -340,17 +346,62 @@ function buildDeleteFilterQuery(filter) {
340
346
  values.push(value);
341
347
  return `metadata#>>'{${parseJsonPathKey(key)}}' = $${values.length}`;
342
348
  }
343
- const [[operator, operatorValue] = []] = Object.entries(value);
349
+ const entries = Object.entries(value);
350
+ if (entries.length > 1) {
351
+ const conditions2 = entries.map(([operator2, operatorValue2]) => {
352
+ if (operator2 === "$not") {
353
+ const nestedEntries = Object.entries(operatorValue2);
354
+ const nestedConditions = nestedEntries.map(([nestedOp, nestedValue]) => {
355
+ if (!FILTER_OPERATORS[nestedOp]) {
356
+ throw new Error(`Invalid operator in $not condition: ${nestedOp}`);
357
+ }
358
+ const operatorFn3 = FILTER_OPERATORS[nestedOp];
359
+ const operatorResult3 = operatorFn3(key, values.length + 1, nestedValue);
360
+ if (operatorResult3.needsValue) {
361
+ const transformedValue = operatorResult3.transformValue ? operatorResult3.transformValue() : nestedValue;
362
+ if (Array.isArray(transformedValue) && nestedOp === "$elemMatch") {
363
+ values.push(...transformedValue);
364
+ } else {
365
+ values.push(transformedValue);
366
+ }
367
+ }
368
+ return operatorResult3.sql;
369
+ }).join(" AND ");
370
+ return `NOT (${nestedConditions})`;
371
+ }
372
+ if (!FILTER_OPERATORS[operator2]) {
373
+ throw new Error(`Invalid operator: ${operator2}`);
374
+ }
375
+ const operatorFn2 = FILTER_OPERATORS[operator2];
376
+ const operatorResult2 = operatorFn2(key, values.length + 1, operatorValue2);
377
+ if (operatorResult2.needsValue) {
378
+ const transformedValue = operatorResult2.transformValue ? operatorResult2.transformValue() : operatorValue2;
379
+ if (Array.isArray(transformedValue) && operator2 === "$elemMatch") {
380
+ values.push(...transformedValue);
381
+ } else {
382
+ values.push(transformedValue);
383
+ }
384
+ }
385
+ return operatorResult2.sql;
386
+ });
387
+ return conditions2.join(" AND ");
388
+ }
389
+ const [[operator, operatorValue] = []] = entries;
344
390
  if (operator === "$not") {
345
- const entries = Object.entries(operatorValue);
346
- const conditions2 = entries.map(([nestedOp, nestedValue]) => {
391
+ const nestedEntries = Object.entries(operatorValue);
392
+ const conditions2 = nestedEntries.map(([nestedOp, nestedValue]) => {
347
393
  if (!FILTER_OPERATORS[nestedOp]) {
348
394
  throw new Error(`Invalid operator in $not condition: ${nestedOp}`);
349
395
  }
350
396
  const operatorFn2 = FILTER_OPERATORS[nestedOp];
351
397
  const operatorResult2 = operatorFn2(key, values.length + 1, nestedValue);
352
398
  if (operatorResult2.needsValue) {
353
- values.push(nestedValue);
399
+ const transformedValue = operatorResult2.transformValue ? operatorResult2.transformValue() : nestedValue;
400
+ if (Array.isArray(transformedValue) && nestedOp === "$elemMatch") {
401
+ values.push(...transformedValue);
402
+ } else {
403
+ values.push(transformedValue);
404
+ }
354
405
  }
355
406
  return operatorResult2.sql;
356
407
  }).join(" AND ");
@@ -417,17 +468,62 @@ function buildFilterQuery(filter, minScore, topK) {
417
468
  values.push(value);
418
469
  return `metadata#>>'{${parseJsonPathKey(key)}}' = $${values.length}`;
419
470
  }
420
- const [[operator, operatorValue] = []] = Object.entries(value);
471
+ const entries = Object.entries(value);
472
+ if (entries.length > 1) {
473
+ const conditions2 = entries.map(([operator2, operatorValue2]) => {
474
+ if (operator2 === "$not") {
475
+ const nestedEntries = Object.entries(operatorValue2);
476
+ const nestedConditions = nestedEntries.map(([nestedOp, nestedValue]) => {
477
+ if (!FILTER_OPERATORS[nestedOp]) {
478
+ throw new Error(`Invalid operator in $not condition: ${nestedOp}`);
479
+ }
480
+ const operatorFn3 = FILTER_OPERATORS[nestedOp];
481
+ const operatorResult3 = operatorFn3(key, values.length + 1, nestedValue);
482
+ if (operatorResult3.needsValue) {
483
+ const transformedValue = operatorResult3.transformValue ? operatorResult3.transformValue() : nestedValue;
484
+ if (Array.isArray(transformedValue) && nestedOp === "$elemMatch") {
485
+ values.push(...transformedValue);
486
+ } else {
487
+ values.push(transformedValue);
488
+ }
489
+ }
490
+ return operatorResult3.sql;
491
+ }).join(" AND ");
492
+ return `NOT (${nestedConditions})`;
493
+ }
494
+ if (!FILTER_OPERATORS[operator2]) {
495
+ throw new Error(`Invalid operator: ${operator2}`);
496
+ }
497
+ const operatorFn2 = FILTER_OPERATORS[operator2];
498
+ const operatorResult2 = operatorFn2(key, values.length + 1, operatorValue2);
499
+ if (operatorResult2.needsValue) {
500
+ const transformedValue = operatorResult2.transformValue ? operatorResult2.transformValue() : operatorValue2;
501
+ if (Array.isArray(transformedValue) && operator2 === "$elemMatch") {
502
+ values.push(...transformedValue);
503
+ } else {
504
+ values.push(transformedValue);
505
+ }
506
+ }
507
+ return operatorResult2.sql;
508
+ });
509
+ return conditions2.join(" AND ");
510
+ }
511
+ const [[operator, operatorValue] = []] = entries;
421
512
  if (operator === "$not") {
422
- const entries = Object.entries(operatorValue);
423
- const conditions2 = entries.map(([nestedOp, nestedValue]) => {
513
+ const nestedEntries = Object.entries(operatorValue);
514
+ const conditions2 = nestedEntries.map(([nestedOp, nestedValue]) => {
424
515
  if (!FILTER_OPERATORS[nestedOp]) {
425
516
  throw new Error(`Invalid operator in $not condition: ${nestedOp}`);
426
517
  }
427
518
  const operatorFn2 = FILTER_OPERATORS[nestedOp];
428
519
  const operatorResult2 = operatorFn2(key, values.length + 1, nestedValue);
429
520
  if (operatorResult2.needsValue) {
430
- values.push(nestedValue);
521
+ const transformedValue = operatorResult2.transformValue ? operatorResult2.transformValue() : nestedValue;
522
+ if (Array.isArray(transformedValue) && nestedOp === "$elemMatch") {
523
+ values.push(...transformedValue);
524
+ } else {
525
+ values.push(transformedValue);
526
+ }
431
527
  }
432
528
  return operatorResult2.sql;
433
529
  }).join(" AND ");
@@ -687,9 +783,7 @@ var PgVector = class extends MastraVector {
687
783
  probes
688
784
  }) {
689
785
  try {
690
- if (!Number.isInteger(topK) || topK <= 0) {
691
- throw new Error("topK must be a positive integer");
692
- }
786
+ validateTopK("PG", topK);
693
787
  if (!Array.isArray(queryVector) || !queryVector.every((x) => typeof x === "number" && Number.isFinite(x))) {
694
788
  throw new Error("queryVector must be an array of finite numbers");
695
789
  }
@@ -774,6 +868,7 @@ var PgVector = class extends MastraVector {
774
868
  ids,
775
869
  deleteFilter
776
870
  }) {
871
+ validateUpsertInput("PG", vectors, metadata, ids);
777
872
  const { tableName } = this.getTableName(indexName);
778
873
  const client = await this.pool.connect();
779
874
  try {
@@ -1849,7 +1944,8 @@ function mapToSqlType(type) {
1849
1944
  function generateTableSQL({
1850
1945
  tableName,
1851
1946
  schema,
1852
- schemaName
1947
+ schemaName,
1948
+ includeAllConstraints = false
1853
1949
  }) {
1854
1950
  const timeZColumns = Object.entries(schema).filter(([_, def]) => def.type === "timestamp").map(([name]) => {
1855
1951
  const parsedName = parseSqlIdentifier(name, "column name");
@@ -1873,9 +1969,9 @@ function generateTableSQL({
1873
1969
  ${tableName === TABLE_WORKFLOW_SNAPSHOT ? `
1874
1970
  DO $$ BEGIN
1875
1971
  IF NOT EXISTS (
1876
- SELECT 1 FROM pg_constraint WHERE conname = '${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key'
1972
+ SELECT 1 FROM pg_constraint WHERE conname = lower('${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key')
1877
1973
  ) AND NOT EXISTS (
1878
- SELECT 1 FROM pg_indexes WHERE indexname = '${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key'
1974
+ SELECT 1 FROM pg_indexes WHERE indexname = lower('${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key')
1879
1975
  ) THEN
1880
1976
  ALTER TABLE ${getTableName({ indexName: tableName, schemaName: quotedSchemaName })}
1881
1977
  ADD CONSTRAINT ${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key
@@ -1883,10 +1979,11 @@ function generateTableSQL({
1883
1979
  END IF;
1884
1980
  END $$;
1885
1981
  ` : ""}
1886
- ${tableName === TABLE_SPANS ? `
1982
+ ${// For spans table: Include PRIMARY KEY in exports, but not in runtime (handled after deduplication)
1983
+ tableName === TABLE_SPANS && includeAllConstraints ? `
1887
1984
  DO $$ BEGIN
1888
1985
  IF NOT EXISTS (
1889
- SELECT 1 FROM pg_constraint WHERE conname = '${constraintPrefix}mastra_ai_spans_traceid_spanid_pk'
1986
+ SELECT 1 FROM pg_constraint WHERE conname = lower('${constraintPrefix}mastra_ai_spans_traceid_spanid_pk')
1890
1987
  ) THEN
1891
1988
  ALTER TABLE ${getTableName({ indexName: tableName, schemaName: quotedSchemaName })}
1892
1989
  ADD CONSTRAINT ${constraintPrefix}mastra_ai_spans_traceid_spanid_pk
@@ -1910,7 +2007,9 @@ function exportSchemas(schemaName) {
1910
2007
  const sql = generateTableSQL({
1911
2008
  tableName,
1912
2009
  schema,
1913
- schemaName
2010
+ schemaName,
2011
+ includeAllConstraints: true
2012
+ // Include all constraints for exports/documentation
1914
2013
  });
1915
2014
  statements.push(sql.trim());
1916
2015
  statements.push("");
@@ -2051,10 +2150,27 @@ var PgDB = class extends MastraBase {
2051
2150
  const columns = Object.keys(record).map((col) => parseSqlIdentifier(col, "column name"));
2052
2151
  const values = this.prepareValuesForInsert(record, tableName);
2053
2152
  const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
2054
- await this.client.none(
2055
- `INSERT INTO ${getTableName({ indexName: tableName, schemaName })} (${columns.map((c) => `"${c}"`).join(", ")}) VALUES (${placeholders})`,
2056
- values
2057
- );
2153
+ const fullTableName = getTableName({ indexName: tableName, schemaName });
2154
+ const columnList = columns.map((c) => `"${c}"`).join(", ");
2155
+ if (tableName === TABLE_SPANS) {
2156
+ const updateColumns = columns.filter((c) => c !== "traceId" && c !== "spanId");
2157
+ if (updateColumns.length > 0) {
2158
+ const updateClause = updateColumns.map((c) => `"${c}" = EXCLUDED."${c}"`).join(", ");
2159
+ await this.client.none(
2160
+ `INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})
2161
+ ON CONFLICT ("traceId", "spanId") DO UPDATE SET ${updateClause}`,
2162
+ values
2163
+ );
2164
+ } else {
2165
+ await this.client.none(
2166
+ `INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})
2167
+ ON CONFLICT ("traceId", "spanId") DO NOTHING`,
2168
+ values
2169
+ );
2170
+ }
2171
+ } else {
2172
+ await this.client.none(`INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})`, values);
2173
+ }
2058
2174
  } catch (error) {
2059
2175
  throw new MastraError(
2060
2176
  {
@@ -2116,8 +2232,46 @@ var PgDB = class extends MastraBase {
2116
2232
  if (tableName === TABLE_SPANS) {
2117
2233
  await this.setupTimestampTriggers(tableName);
2118
2234
  await this.migrateSpansTable();
2235
+ const pkExists = await this.spansPrimaryKeyExists();
2236
+ if (!pkExists) {
2237
+ const duplicateInfo = await this.checkForDuplicateSpans();
2238
+ if (duplicateInfo.hasDuplicates) {
2239
+ const errorMessage = `
2240
+ ===========================================================================
2241
+ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}
2242
+ ===========================================================================
2243
+
2244
+ Found ${duplicateInfo.duplicateCount} duplicate (traceId, spanId) combinations.
2245
+
2246
+ The spans table requires a unique constraint on (traceId, spanId), but your
2247
+ database contains duplicate entries that must be resolved first.
2248
+
2249
+ To fix this, run the manual migration command:
2250
+
2251
+ npx mastra migrate
2252
+
2253
+ This command will:
2254
+ 1. Remove duplicate spans (keeping the most complete/recent version)
2255
+ 2. Add the required unique constraint
2256
+
2257
+ Note: This migration may take some time for large tables.
2258
+ ===========================================================================
2259
+ `;
2260
+ throw new MastraError({
2261
+ id: createStorageErrorId("PG", "MIGRATION_REQUIRED", "DUPLICATE_SPANS"),
2262
+ domain: ErrorDomain.STORAGE,
2263
+ category: ErrorCategory.USER,
2264
+ text: errorMessage
2265
+ });
2266
+ } else {
2267
+ await this.addSpansPrimaryKey();
2268
+ }
2269
+ }
2119
2270
  }
2120
2271
  } catch (error) {
2272
+ if (error instanceof MastraError) {
2273
+ throw error;
2274
+ }
2121
2275
  throw new MastraError(
2122
2276
  {
2123
2277
  id: createStorageErrorId("PG", "CREATE_TABLE", "FAILED"),
@@ -2210,6 +2364,224 @@ var PgDB = class extends MastraBase {
2210
2364
  this.logger?.warn?.(`Failed to migrate spans table ${fullTableName}:`, error);
2211
2365
  }
2212
2366
  }
2367
+ /**
2368
+ * Deduplicates spans in the mastra_ai_spans table before adding the PRIMARY KEY constraint.
2369
+ * Keeps spans based on priority: completed (endedAt NOT NULL) > most recent updatedAt > most recent createdAt.
2370
+ *
2371
+ * Note: This prioritizes migration completion over perfect data preservation.
2372
+ * Old trace data may be lost, which is acceptable for this use case.
2373
+ */
2374
+ async deduplicateSpans() {
2375
+ const fullTableName = getTableName({ indexName: TABLE_SPANS, schemaName: getSchemaName(this.schemaName) });
2376
+ try {
2377
+ const duplicateCheck = await this.client.oneOrNone(`
2378
+ SELECT EXISTS (
2379
+ SELECT 1
2380
+ FROM ${fullTableName}
2381
+ GROUP BY "traceId", "spanId"
2382
+ HAVING COUNT(*) > 1
2383
+ LIMIT 1
2384
+ ) as has_duplicates
2385
+ `);
2386
+ if (!duplicateCheck?.has_duplicates) {
2387
+ this.logger?.debug?.(`No duplicate spans found in ${fullTableName}`);
2388
+ return;
2389
+ }
2390
+ this.logger?.info?.(`Duplicate spans detected in ${fullTableName}, starting deduplication...`);
2391
+ const result = await this.client.query(`
2392
+ DELETE FROM ${fullTableName} t1
2393
+ USING ${fullTableName} t2
2394
+ WHERE t1."traceId" = t2."traceId"
2395
+ AND t1."spanId" = t2."spanId"
2396
+ AND (
2397
+ -- Keep completed spans over incomplete
2398
+ (t1."endedAt" IS NULL AND t2."endedAt" IS NOT NULL)
2399
+ OR
2400
+ -- If both have same completion status, keep more recent updatedAt
2401
+ (
2402
+ (t1."endedAt" IS NULL) = (t2."endedAt" IS NULL)
2403
+ AND (
2404
+ (t1."updatedAt" < t2."updatedAt")
2405
+ OR (t1."updatedAt" IS NULL AND t2."updatedAt" IS NOT NULL)
2406
+ OR
2407
+ -- If updatedAt is the same, keep more recent createdAt
2408
+ (
2409
+ (t1."updatedAt" = t2."updatedAt" OR (t1."updatedAt" IS NULL AND t2."updatedAt" IS NULL))
2410
+ AND (
2411
+ (t1."createdAt" < t2."createdAt")
2412
+ OR (t1."createdAt" IS NULL AND t2."createdAt" IS NOT NULL)
2413
+ OR
2414
+ -- If all else equal, use ctid as tiebreaker
2415
+ (
2416
+ (t1."createdAt" = t2."createdAt" OR (t1."createdAt" IS NULL AND t2."createdAt" IS NULL))
2417
+ AND t1.ctid < t2.ctid
2418
+ )
2419
+ )
2420
+ )
2421
+ )
2422
+ )
2423
+ )
2424
+ `);
2425
+ this.logger?.info?.(
2426
+ `Deduplication complete: removed ${result.rowCount ?? 0} duplicate spans from ${fullTableName}`
2427
+ );
2428
+ } catch (error) {
2429
+ throw new MastraError(
2430
+ {
2431
+ id: createStorageErrorId("PG", "DEDUPLICATE_SPANS", "FAILED"),
2432
+ domain: ErrorDomain.STORAGE,
2433
+ category: ErrorCategory.THIRD_PARTY,
2434
+ details: {
2435
+ tableName: TABLE_SPANS
2436
+ }
2437
+ },
2438
+ error
2439
+ );
2440
+ }
2441
+ }
2442
+ /**
2443
+ * Checks for duplicate (traceId, spanId) combinations in the spans table.
2444
+ * Returns information about duplicates for logging/CLI purposes.
2445
+ */
2446
+ async checkForDuplicateSpans() {
2447
+ const fullTableName = getTableName({ indexName: TABLE_SPANS, schemaName: getSchemaName(this.schemaName) });
2448
+ try {
2449
+ const result = await this.client.oneOrNone(`
2450
+ SELECT COUNT(*) as duplicate_count
2451
+ FROM (
2452
+ SELECT "traceId", "spanId"
2453
+ FROM ${fullTableName}
2454
+ GROUP BY "traceId", "spanId"
2455
+ HAVING COUNT(*) > 1
2456
+ ) duplicates
2457
+ `);
2458
+ const duplicateCount = parseInt(result?.duplicate_count ?? "0", 10);
2459
+ return {
2460
+ hasDuplicates: duplicateCount > 0,
2461
+ duplicateCount,
2462
+ tableName: fullTableName
2463
+ };
2464
+ } catch (error) {
2465
+ this.logger?.debug?.(`Could not check for duplicates: ${error}`);
2466
+ return { hasDuplicates: false, duplicateCount: 0, tableName: fullTableName };
2467
+ }
2468
+ }
2469
+ /**
2470
+ * Checks if the PRIMARY KEY constraint on (traceId, spanId) already exists on the spans table.
2471
+ * Used to skip deduplication when the constraint already exists (migration already complete).
2472
+ */
2473
+ async spansPrimaryKeyExists() {
2474
+ const parsedSchemaName = this.schemaName ? parseSqlIdentifier(this.schemaName, "schema name") : "";
2475
+ const constraintPrefix = parsedSchemaName ? `${parsedSchemaName}_` : "";
2476
+ const constraintName = `${constraintPrefix}mastra_ai_spans_traceid_spanid_pk`;
2477
+ const result = await this.client.oneOrNone(
2478
+ `SELECT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = $1) as exists`,
2479
+ [constraintName]
2480
+ );
2481
+ return result?.exists ?? false;
2482
+ }
2483
+ /**
2484
+ * Adds the PRIMARY KEY constraint on (traceId, spanId) to the spans table.
2485
+ * Should be called AFTER deduplication to ensure no duplicate key violations.
2486
+ */
2487
+ async addSpansPrimaryKey() {
2488
+ const fullTableName = getTableName({ indexName: TABLE_SPANS, schemaName: getSchemaName(this.schemaName) });
2489
+ const parsedSchemaName = this.schemaName ? parseSqlIdentifier(this.schemaName, "schema name") : "";
2490
+ const constraintPrefix = parsedSchemaName ? `${parsedSchemaName}_` : "";
2491
+ const constraintName = `${constraintPrefix}mastra_ai_spans_traceid_spanid_pk`;
2492
+ try {
2493
+ const constraintExists = await this.client.oneOrNone(
2494
+ `
2495
+ SELECT EXISTS (
2496
+ SELECT 1 FROM pg_constraint WHERE conname = $1
2497
+ ) as exists
2498
+ `,
2499
+ [constraintName]
2500
+ );
2501
+ if (constraintExists?.exists) {
2502
+ this.logger?.debug?.(`PRIMARY KEY constraint ${constraintName} already exists on ${fullTableName}`);
2503
+ return;
2504
+ }
2505
+ await this.client.none(`
2506
+ ALTER TABLE ${fullTableName}
2507
+ ADD CONSTRAINT ${constraintName}
2508
+ PRIMARY KEY ("traceId", "spanId")
2509
+ `);
2510
+ this.logger?.info?.(`Added PRIMARY KEY constraint ${constraintName} to ${fullTableName}`);
2511
+ } catch (error) {
2512
+ throw new MastraError(
2513
+ {
2514
+ id: createStorageErrorId("PG", "ADD_SPANS_PRIMARY_KEY", "FAILED"),
2515
+ domain: ErrorDomain.STORAGE,
2516
+ category: ErrorCategory.THIRD_PARTY,
2517
+ details: {
2518
+ tableName: TABLE_SPANS,
2519
+ constraintName
2520
+ }
2521
+ },
2522
+ error
2523
+ );
2524
+ }
2525
+ }
2526
+ /**
2527
+ * Manually run the spans migration to deduplicate and add the unique constraint.
2528
+ * This is intended to be called from the CLI when duplicates are detected.
2529
+ *
2530
+ * @returns Migration result with status and details
2531
+ */
2532
+ async migrateSpans() {
2533
+ const fullTableName = getTableName({ indexName: TABLE_SPANS, schemaName: getSchemaName(this.schemaName) });
2534
+ const pkExists = await this.spansPrimaryKeyExists();
2535
+ if (pkExists) {
2536
+ return {
2537
+ success: true,
2538
+ alreadyMigrated: true,
2539
+ duplicatesRemoved: 0,
2540
+ message: `Migration already complete. PRIMARY KEY constraint exists on ${fullTableName}.`
2541
+ };
2542
+ }
2543
+ const duplicateInfo = await this.checkForDuplicateSpans();
2544
+ if (duplicateInfo.hasDuplicates) {
2545
+ this.logger?.info?.(
2546
+ `Found ${duplicateInfo.duplicateCount} duplicate (traceId, spanId) combinations. Starting deduplication...`
2547
+ );
2548
+ await this.deduplicateSpans();
2549
+ } else {
2550
+ this.logger?.info?.(`No duplicate spans found.`);
2551
+ }
2552
+ await this.addSpansPrimaryKey();
2553
+ return {
2554
+ success: true,
2555
+ alreadyMigrated: false,
2556
+ duplicatesRemoved: duplicateInfo.duplicateCount,
2557
+ message: duplicateInfo.hasDuplicates ? `Migration complete. Removed duplicates and added PRIMARY KEY constraint to ${fullTableName}.` : `Migration complete. Added PRIMARY KEY constraint to ${fullTableName}.`
2558
+ };
2559
+ }
2560
+ /**
2561
+ * Check migration status for the spans table.
2562
+ * Returns information about whether migration is needed.
2563
+ */
2564
+ async checkSpansMigrationStatus() {
2565
+ const fullTableName = getTableName({ indexName: TABLE_SPANS, schemaName: getSchemaName(this.schemaName) });
2566
+ const pkExists = await this.spansPrimaryKeyExists();
2567
+ if (pkExists) {
2568
+ return {
2569
+ needsMigration: false,
2570
+ hasDuplicates: false,
2571
+ duplicateCount: 0,
2572
+ constraintExists: true,
2573
+ tableName: fullTableName
2574
+ };
2575
+ }
2576
+ const duplicateInfo = await this.checkForDuplicateSpans();
2577
+ return {
2578
+ needsMigration: true,
2579
+ hasDuplicates: duplicateInfo.hasDuplicates,
2580
+ duplicateCount: duplicateInfo.duplicateCount,
2581
+ constraintExists: false,
2582
+ tableName: fullTableName
2583
+ };
2584
+ }
2213
2585
  /**
2214
2586
  * Alters table schema to add columns if they don't exist
2215
2587
  * @param tableName Name of the table
@@ -4247,6 +4619,22 @@ var ObservabilityPG = class _ObservabilityPG extends ObservabilityStorage {
4247
4619
  }
4248
4620
  }
4249
4621
  }
4622
+ /**
4623
+ * Manually run the spans migration to deduplicate and add the unique constraint.
4624
+ * This is intended to be called from the CLI when duplicates are detected.
4625
+ *
4626
+ * @returns Migration result with status and details
4627
+ */
4628
+ async migrateSpans() {
4629
+ return this.#db.migrateSpans();
4630
+ }
4631
+ /**
4632
+ * Check migration status for the spans table.
4633
+ * Returns information about whether migration is needed.
4634
+ */
4635
+ async checkSpansMigrationStatus() {
4636
+ return this.#db.checkSpansMigrationStatus();
4637
+ }
4250
4638
  async dangerouslyClearAll() {
4251
4639
  await this.#db.clearTable({ tableName: TABLE_SPANS });
4252
4640
  }
@@ -5417,7 +5805,7 @@ var WorkflowsPG = class _WorkflowsPG extends WorkflowsStorage {
5417
5805
  // src/storage/index.ts
5418
5806
  var DEFAULT_MAX_CONNECTIONS = 20;
5419
5807
  var DEFAULT_IDLE_TIMEOUT_MS = 3e4;
5420
- var PostgresStore = class extends MastraStorage {
5808
+ var PostgresStore = class extends MastraCompositeStore {
5421
5809
  #pool;
5422
5810
  #db;
5423
5811
  #ownsPool;
@@ -5496,6 +5884,9 @@ var PostgresStore = class extends MastraStorage {
5496
5884
  await super.init();
5497
5885
  } catch (error) {
5498
5886
  this.isInitialized = false;
5887
+ if (error instanceof MastraError) {
5888
+ throw error;
5889
+ }
5499
5890
  throw new MastraError(
5500
5891
  {
5501
5892
  id: createStorageErrorId("PG", "INIT", "FAILED"),