@mastra/pg 1.15.1 → 1.16.0-alpha.0

package/dist/index.js

@@ -1,5 +1,5 @@
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
-import { BRANCH_SPAN_TYPES, createVectorErrorId, AgentsStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, TABLE_SCHEMAS, createStorageErrorId, normalizePerPage, calculatePagination, BackgroundTasksStorage, TABLE_BACKGROUND_TASKS, BlobStore, TABLE_SKILL_BLOBS, ChannelsStorage, TABLE_CHANNEL_INSTALLATIONS, TABLE_CHANNEL_CONFIG, DatasetsStorage, TABLE_DATASETS, TABLE_DATASET_ITEMS, TABLE_DATASET_VERSIONS, DATASETS_SCHEMA, TABLE_CONFIGS, DATASET_ITEMS_SCHEMA, DATASET_VERSIONS_SCHEMA, ensureDate, safelyParseJSON, ExperimentsStorage, TABLE_EXPERIMENTS, TABLE_EXPERIMENT_RESULTS, EXPERIMENTS_SCHEMA, EXPERIMENT_RESULTS_SCHEMA, FavoritesStorage, TABLE_FAVORITES, MCPClientsStorage, TABLE_MCP_CLIENTS, TABLE_MCP_CLIENT_VERSIONS, MCPServersStorage, TABLE_MCP_SERVERS, TABLE_MCP_SERVER_VERSIONS, MemoryStorage, TABLE_THREADS, TABLE_RESOURCES, TABLE_MESSAGES, OBSERVATIONAL_MEMORY_TABLE_SCHEMA, NotificationsStorage, TABLE_NOTIFICATIONS, ObservabilityStorage, TABLE_SPANS, listTracesArgsSchema, toTraceSpans, PromptBlocksStorage, TABLE_PROMPT_BLOCKS, TABLE_PROMPT_BLOCK_VERSIONS, SchedulesStorage, TABLE_SCHEDULES, TABLE_SCHEDULE_TRIGGERS, ScorerDefinitionsStorage, TABLE_SCORER_DEFINITIONS, TABLE_SCORER_DEFINITION_VERSIONS, ScoresStorage, TABLE_SCORERS, SkillsStorage, TABLE_SKILLS, TABLE_SKILL_VERSIONS, ToolProviderConnectionsStorage, TABLE_TOOL_PROVIDER_CONNECTIONS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, mergeWorkflowStepResult, WorkspacesStorage, TABLE_WORKSPACES, TABLE_WORKSPACE_VERSIONS, MastraCompositeStore, TraceStatus, getDefaultValue, parseDuration, listBranchesArgsSchema, listLogsArgsSchema, listMetricsArgsSchema, listScoresArgsSchema, listFeedbackArgsSchema, normalizeScheduleTarget, transformScoreRow as transformScoreRow$1, getSqlType, EntityType, METRIC_DISTINCT_COLUMNS } from '@mastra/core/storage';
+import { BRANCH_SPAN_TYPES, createVectorErrorId, AgentsStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, TABLE_SCHEMAS, createStorageErrorId, normalizePerPage, calculatePagination, BackgroundTasksStorage, TABLE_BACKGROUND_TASKS, BlobStore, TABLE_SKILL_BLOBS, ChannelsStorage, TABLE_CHANNEL_INSTALLATIONS, TABLE_CHANNEL_CONFIG, DatasetsStorage, TABLE_DATASETS, TABLE_DATASET_ITEMS, TABLE_DATASET_VERSIONS, DATASETS_SCHEMA, TABLE_CONFIGS, DATASET_ITEMS_SCHEMA, DATASET_VERSIONS_SCHEMA, ensureDate, safelyParseJSON, hasErrorCode, ExperimentsStorage, TABLE_EXPERIMENTS, TABLE_EXPERIMENT_RESULTS, EXPERIMENTS_SCHEMA, EXPERIMENT_RESULTS_SCHEMA, FavoritesStorage, TABLE_FAVORITES, MCPClientsStorage, TABLE_MCP_CLIENTS, TABLE_MCP_CLIENT_VERSIONS, MCPServersStorage, TABLE_MCP_SERVERS, TABLE_MCP_SERVER_VERSIONS, MemoryStorage, TABLE_THREADS, TABLE_RESOURCES, TABLE_MESSAGES, OBSERVATIONAL_MEMORY_TABLE_SCHEMA, NotificationsStorage, TABLE_NOTIFICATIONS, ObservabilityStorage, TABLE_SPANS, listTracesArgsSchema, toTraceSpans, PromptBlocksStorage, TABLE_PROMPT_BLOCKS, TABLE_PROMPT_BLOCK_VERSIONS, SchedulesStorage, TABLE_SCHEDULES, TABLE_SCHEDULE_TRIGGERS, ScorerDefinitionsStorage, TABLE_SCORER_DEFINITIONS, TABLE_SCORER_DEFINITION_VERSIONS, ScoresStorage, TABLE_SCORERS, SkillsStorage, TABLE_SKILLS, TABLE_SKILL_VERSIONS, ToolProviderConnectionsStorage, TABLE_TOOL_PROVIDER_CONNECTIONS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, mergeWorkflowStepResult, WorkspacesStorage, TABLE_WORKSPACES, TABLE_WORKSPACE_VERSIONS, MastraCompositeStore, TraceStatus, getDefaultValue, parseDuration, listBranchesArgsSchema, listLogsArgsSchema, listMetricsArgsSchema, listScoresArgsSchema, listFeedbackArgsSchema, normalizeScheduleTarget, transformScoreRow as transformScoreRow$1, getSqlType, EntityType, METRIC_DISTINCT_COLUMNS } from '@mastra/core/storage';
 import { parseSqlIdentifier, parseFieldKey } from '@mastra/core/utils';
 import { MastraVector, validateTopK, validateUpsertInput } from '@mastra/core/vector';
 import { Mutex } from 'async-mutex';
@@ -5428,6 +5428,7 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
     await this.#addColumnIfNotExists(TABLE_DATASET_ITEMS, "organizationId", "TEXT");
     await this.#addColumnIfNotExists(TABLE_DATASET_ITEMS, "projectId", "TEXT");
     await this.#addColumnIfNotExists(TABLE_DATASET_ITEMS, "toolMocks", "JSONB");
+    await this.#addColumnIfNotExists(TABLE_DATASET_ITEMS, "externalId", "TEXT");
     await this.createDefaultIndexes();
     await this.createCustomIndexes();
   }
@@ -5446,6 +5447,11 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
         table: TABLE_DATASET_ITEMS,
         columns: ["datasetId", "datasetVersion"]
       },
+      {
+        name: "idx_dataset_items_external_id_history",
+        table: TABLE_DATASET_ITEMS,
+        columns: ["datasetId", "externalId", "datasetVersion"]
+      },
       {
         name: "idx_dataset_items_dataset_validto_deleted",
         table: TABLE_DATASET_ITEMS,
@@ -5528,6 +5534,7 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
       id: row.id,
       datasetId: row.datasetId,
       datasetVersion: row.datasetVersion,
+      externalId: row.externalId ?? null,
       organizationId: row.organizationId ?? null,
       projectId: row.projectId ?? null,
       input: safelyParseJSON(row.input),
@@ -5546,6 +5553,7 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
       id: row.id,
       datasetId: row.datasetId,
       datasetVersion: row.datasetVersion,
+      externalId: row.externalId ?? null,
       organizationId: row.organizationId ?? null,
       projectId: row.projectId ?? null,
       validTo: row.validTo,
@@ -5572,7 +5580,8 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
   // --- Dataset CRUD ---
   async createDataset(input) {
     try {
-      const id = crypto.randomUUID();
+      const id = input.id ?? crypto.randomUUID();
+      if (input.id !== void 0) this.validateCallerDefinedDatasetId(input.id);
       const now = /* @__PURE__ */ new Date();
       const nowIso = now.toISOString();
       await this.#db.insert({
@@ -5617,6 +5626,11 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
         updatedAt: now
       };
     } catch (error) {
+      if (input.id !== void 0 && hasErrorCode(error, /* @__PURE__ */ new Set(["23505"]))) {
+        const existing = await this.getDatasetById({ id: input.id });
+        if (existing) return this.resolveExistingDataset(existing, { ...input, id: input.id });
+      }
+      if (error instanceof MastraError) throw error;
       throw new MastraError(
         {
           id: createStorageErrorId("PG", "CREATE_DATASET", "FAILED"),
@@ -5889,11 +5903,12 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
         parentOrganizationId = row.organizationId ?? null;
         parentProjectId = row.projectId ?? null;
         await t.none(
-          `INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,NULL,false,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16)`,
+          `INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,false,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)`,
           [
             id,
             args.datasetId,
             newVersion,
+            args.externalId ?? null,
             parentOrganizationId,
             parentProjectId,
             JSON.stringify(args.input),
@@ -5993,11 +6008,12 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
           [newVersion, args.id]
         );
         await t.none(
-          `INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,NULL,false,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16)`,
+          `INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,false,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)`,
           [
             args.id,
             args.datasetId,
             newVersion,
+            existing.externalId ?? null,
             parentOrganizationId,
             parentProjectId,
             JSON.stringify(mergedInput),
@@ -6077,11 +6093,12 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
           [newVersion, id]
         );
         await t.none(
-          `INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,NULL,true,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16)`,
+          `INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,true,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)`,
           [
             id,
             datasetId,
             newVersion,
+            existing.externalId ?? null,
             parentOrganizationId,
             parentProjectId,
             JSON.stringify(existing.input),
@@ -6116,78 +6133,93 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
   }
   async _doBatchInsertItems(input) {
     try {
-      const dataset = await this.getDatasetById({ id: input.datasetId });
-      if (!dataset) {
-        throw new MastraError({
-          id: createStorageErrorId("PG", "BULK_ADD_ITEMS", "DATASET_NOT_FOUND"),
-          domain: ErrorDomain.STORAGE,
-          category: ErrorCategory.USER,
-          details: { datasetId: input.datasetId }
-        });
-      }
+      if (input.items.length === 0) return [];
       const datasetsTable = getTableName2({ indexName: TABLE_DATASETS, schemaName: getSchemaName2(this.#schema) });
       const itemsTable = getTableName2({ indexName: TABLE_DATASET_ITEMS, schemaName: getSchemaName2(this.#schema) });
       const versionsTable = getTableName2({
         indexName: TABLE_DATASET_VERSIONS,
         schemaName: getSchemaName2(this.#schema)
       });
-      const now = /* @__PURE__ */ new Date();
-      const nowIso = now.toISOString();
-      const versionId = crypto.randomUUID();
-      const itemsWithIds = input.items.map((itemInput) => ({ id: crypto.randomUUID(), input: itemInput }));
-      const parentOrganizationId = dataset.organizationId ?? null;
-      const parentProjectId = dataset.projectId ?? null;
-      let newVersion;
+      let results = [];
       await this.#db.client.tx(async (t) => {
-        const row = await t.one(
-          `UPDATE ${datasetsTable} SET "version" = "version" + 1 WHERE "id" = $1 RETURNING "version"`,
+        const dataset = await t.oneOrNone(
+          `SELECT "version", "organizationId", "projectId" FROM ${datasetsTable} WHERE "id" = $1 FOR UPDATE`,
           [input.datasetId]
         );
-        newVersion = row.version;
-        for (const { id, input: itemInput } of itemsWithIds) {
-          await t.none(
-            `INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,NULL,false,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16)`,
-            [
+        if (!dataset)
+          throw new MastraError({
+            id: createStorageErrorId("PG", "BULK_ADD_ITEMS", "DATASET_NOT_FOUND"),
+            domain: ErrorDomain.STORAGE,
+            category: ErrorCategory.USER,
+            details: { datasetId: input.datasetId }
+          });
+        const externalIds = [...new Set(input.items.flatMap((item) => item.externalId ? [item.externalId] : []))];
+        const historyRows = externalIds.length ? await t.manyOrNone(
+          `SELECT * FROM ${itemsTable} WHERE "datasetId" = $1 AND "externalId" = ANY($2::text[]) ORDER BY "datasetVersion"`,
+          [input.datasetId, externalIds]
+        ) : [];
+        const plan = this.planDatasetItemBatch(
+          input.items,
+          historyRows.map((row) => this.transformItemRowFull(row)),
+          () => crypto.randomUUID()
+        );
+        const resolved = new Map(
+          [...plan.existingCurrentItems].map(([id, row]) => [id, this.datasetItemFromRow(row)])
+        );
+        if (plan.inserts.length > 0) {
+          const newVersion = Number(dataset.version) + 1;
+          const now = /* @__PURE__ */ new Date();
+          const nowIso = now.toISOString();
+          await t.none(`UPDATE ${datasetsTable} SET "version" = $2 WHERE "id" = $1`, [input.datasetId, newVersion]);
+          for (const { id, item } of plan.inserts) {
+            await t.none(
+              `INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,false,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)`,
+              [
+                id,
+                input.datasetId,
+                newVersion,
+                item.externalId ?? null,
+                dataset.organizationId ?? null,
+                dataset.projectId ?? null,
+                JSON.stringify(item.input),
+                jsonbArg(item.groundTruth),
+                jsonbArg(item.expectedTrajectory),
+                jsonbArg(item.toolMocks),
+                jsonbArg(item.requestContext),
+                jsonbArg(item.metadata),
+                jsonbArg(item.source),
+                nowIso,
+                nowIso,
+                nowIso,
+                nowIso
+              ]
+            );
+            resolved.set(id, {
               id,
-              input.datasetId,
-              newVersion,
-              parentOrganizationId,
-              parentProjectId,
-              JSON.stringify(itemInput.input),
-              jsonbArg(itemInput.groundTruth),
-              jsonbArg(itemInput.expectedTrajectory),
-              jsonbArg(itemInput.toolMocks),
-              jsonbArg(itemInput.requestContext),
-              jsonbArg(itemInput.metadata),
-              jsonbArg(itemInput.source),
-              nowIso,
-              nowIso,
-              nowIso,
-              nowIso
-            ]
+              datasetId: input.datasetId,
+              datasetVersion: newVersion,
+              externalId: item.externalId ?? null,
+              organizationId: dataset.organizationId ?? null,
+              projectId: dataset.projectId ?? null,
+              input: item.input,
+              groundTruth: item.groundTruth,
+              expectedTrajectory: item.expectedTrajectory,
+              toolMocks: item.toolMocks,
+              requestContext: item.requestContext,
+              metadata: item.metadata,
+              source: item.source,
+              createdAt: now,
+              updatedAt: now
+            });
+          }
+          await t.none(
+            `INSERT INTO ${versionsTable} ("id","datasetId","version","createdAt","createdAtZ") VALUES ($1,$2,$3,$4,$5)`,
+            [crypto.randomUUID(), input.datasetId, newVersion, nowIso, nowIso]
           );
         }
-        await t.none(
-          `INSERT INTO ${versionsTable} ("id","datasetId","version","createdAt","createdAtZ") VALUES ($1,$2,$3,$4,$5)`,
-          [versionId, input.datasetId, newVersion, nowIso, nowIso]
-        );
+        results = plan.resolvedIds.map((id) => resolved.get(id));
       });
-      return itemsWithIds.map(({ id, input: itemInput }) => ({
-        id,
-        datasetId: input.datasetId,
-        datasetVersion: newVersion,
-        organizationId: parentOrganizationId,
-        projectId: parentProjectId,
-        input: itemInput.input,
-        groundTruth: itemInput.groundTruth,
-        expectedTrajectory: itemInput.expectedTrajectory,
-        toolMocks: itemInput.toolMocks,
-        requestContext: itemInput.requestContext,
-        metadata: itemInput.metadata,
-        source: itemInput.source,
-        createdAt: now,
-        updatedAt: now
-      }));
+      return results;
     } catch (error) {
       if (error instanceof MastraError) throw error;
       throw new MastraError(
@@ -6241,11 +6273,12 @@ var DatasetsPG = class _DatasetsPG extends DatasetsStorage {
             [newVersion, item.id]
           );
           await t.none(
-            `INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,NULL,true,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16)`,
+            `INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,true,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)`,
             [
               item.id,
               input.datasetId,
               newVersion,
+              item.externalId ?? null,
               parentOrganizationId,
               parentProjectId,
               JSON.stringify(item.input),