@mastra/spanner 1.2.4-alpha.1 → 1.3.0

package/dist/index.js

@@ -1,6 +1,6 @@
 import { Spanner } from '@google-cloud/spanner';
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
-import { AgentsStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, AGENTS_SCHEMA, AGENT_VERSIONS_SCHEMA, createStorageErrorId, normalizePerPage, calculatePagination, TABLE_FAVORITES, BackgroundTasksStorage, TABLE_BACKGROUND_TASKS, TABLE_SCHEMAS, BlobStore, TABLE_SKILL_BLOBS, SKILL_BLOBS_SCHEMA, ChannelsStorage, TABLE_CHANNEL_INSTALLATIONS, TABLE_CHANNEL_CONFIG, DatasetsStorage, TABLE_DATASETS, TABLE_DATASET_ITEMS, TABLE_DATASET_VERSIONS, TABLE_EXPERIMENT_RESULTS, TABLE_EXPERIMENTS, ExperimentsStorage, FavoritesStorage, TABLE_SKILLS, MCPClientsStorage, TABLE_MCP_CLIENTS, TABLE_MCP_CLIENT_VERSIONS, MCP_CLIENTS_SCHEMA, MCP_CLIENT_VERSIONS_SCHEMA, MCPServersStorage, TABLE_MCP_SERVERS, TABLE_MCP_SERVER_VERSIONS, MCP_SERVERS_SCHEMA, MCP_SERVER_VERSIONS_SCHEMA, MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, ObservabilityStorage, TABLE_SPANS, listTracesArgsSchema, TraceStatus, toTraceSpans, PromptBlocksStorage, TABLE_PROMPT_BLOCKS, TABLE_PROMPT_BLOCK_VERSIONS, PROMPT_BLOCKS_SCHEMA, PROMPT_BLOCK_VERSIONS_SCHEMA, SchedulesStorage, TABLE_SCHEDULES, TABLE_SCHEDULE_TRIGGERS, ScorerDefinitionsStorage, TABLE_SCORER_DEFINITIONS, TABLE_SCORER_DEFINITION_VERSIONS, SCORER_DEFINITIONS_SCHEMA, SCORER_DEFINITION_VERSIONS_SCHEMA, ScoresStorage, TABLE_SCORERS, SkillsStorage, TABLE_SKILL_VERSIONS, SKILLS_SCHEMA, SKILL_VERSIONS_SCHEMA, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, WorkspacesStorage, TABLE_WORKSPACES, TABLE_WORKSPACE_VERSIONS, WORKSPACES_SCHEMA, WORKSPACE_VERSIONS_SCHEMA, MastraCompositeStore, getDefaultValue, TABLE_CONFIGS, listMetricsArgsSchema, normalizeScheduleTarget, transformScoreRow as transformScoreRow$1, METRIC_DISTINCT_COLUMNS } from '@mastra/core/storage';
+import { AgentsStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, AGENTS_SCHEMA, AGENT_VERSIONS_SCHEMA, createStorageErrorId, normalizePerPage, calculatePagination, TABLE_FAVORITES, BackgroundTasksStorage, TABLE_BACKGROUND_TASKS, TABLE_SCHEMAS, BlobStore, TABLE_SKILL_BLOBS, SKILL_BLOBS_SCHEMA, ChannelsStorage, TABLE_CHANNEL_INSTALLATIONS, TABLE_CHANNEL_CONFIG, DatasetsStorage, TABLE_DATASETS, TABLE_DATASET_ITEMS, TABLE_DATASET_VERSIONS, TABLE_EXPERIMENT_RESULTS, TABLE_EXPERIMENTS, hasErrorCode, ExperimentsStorage, FavoritesStorage, TABLE_SKILLS, MCPClientsStorage, TABLE_MCP_CLIENTS, TABLE_MCP_CLIENT_VERSIONS, MCP_CLIENTS_SCHEMA, MCP_CLIENT_VERSIONS_SCHEMA, MCPServersStorage, TABLE_MCP_SERVERS, TABLE_MCP_SERVER_VERSIONS, MCP_SERVERS_SCHEMA, MCP_SERVER_VERSIONS_SCHEMA, MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, ObservabilityStorage, TABLE_SPANS, listTracesArgsSchema, TraceStatus, toTraceSpans, PromptBlocksStorage, TABLE_PROMPT_BLOCKS, TABLE_PROMPT_BLOCK_VERSIONS, PROMPT_BLOCKS_SCHEMA, PROMPT_BLOCK_VERSIONS_SCHEMA, SchedulesStorage, TABLE_SCHEDULES, TABLE_SCHEDULE_TRIGGERS, ScorerDefinitionsStorage, TABLE_SCORER_DEFINITIONS, TABLE_SCORER_DEFINITION_VERSIONS, SCORER_DEFINITIONS_SCHEMA, SCORER_DEFINITION_VERSIONS_SCHEMA, ScoresStorage, TABLE_SCORERS, SkillsStorage, TABLE_SKILL_VERSIONS, SKILLS_SCHEMA, SKILL_VERSIONS_SCHEMA, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, WorkspacesStorage, TABLE_WORKSPACES, TABLE_WORKSPACE_VERSIONS, WORKSPACES_SCHEMA, WORKSPACE_VERSIONS_SCHEMA, MastraCompositeStore, getDefaultValue, TABLE_CONFIGS, listMetricsArgsSchema, normalizeScheduleTarget, transformScoreRow as transformScoreRow$1, METRIC_DISTINCT_COLUMNS } from '@mastra/core/storage';
 import { MastraBase } from '@mastra/core/base';
 import { parseSqlIdentifier } from '@mastra/core/utils';
 import { randomUUID } from 'crypto';
@@ -2993,6 +2993,7 @@ function rowToItem(row) {
     id: String(t.id),
     datasetId: String(t.datasetId),
     datasetVersion: Number(t.datasetVersion),
+    externalId: t.externalId ?? null,
     organizationId: t.organizationId ?? null,
     projectId: t.projectId ?? null,
     input: t.input,
@@ -3057,7 +3058,7 @@ var DatasetsSpanner = class _DatasetsSpanner extends DatasetsStorage {
     await this.db.alterTable({
       tableName: TABLE_DATASET_ITEMS,
       schema: TABLE_SCHEMAS[TABLE_DATASET_ITEMS],
-      ifNotExists: ["organizationId", "projectId"]
+      ifNotExists: ["organizationId", "projectId", "externalId"]
     });
     await this.createDefaultIndexes();
     await this.createCustomIndexes();
@@ -3076,6 +3077,11 @@ var DatasetsSpanner = class _DatasetsSpanner extends DatasetsStorage {
         table: TABLE_DATASET_ITEMS,
         columns: ["datasetId", "datasetVersion"]
       },
+      {
+        name: "mastra_dataset_items_dataset_externalid_version_idx",
+        table: TABLE_DATASET_ITEMS,
+        columns: ["datasetId", "externalId", "datasetVersion"]
+      },
       {
         // Unique invariant: one snapshot row per (datasetId, version). The DESC
         // ordering also serves listDatasetVersions' newest-first scan.
@@ -3119,7 +3125,8 @@ var DatasetsSpanner = class _DatasetsSpanner extends DatasetsStorage {
   async createDataset(input) {
     try {
       const now = /* @__PURE__ */ new Date();
-      const id = randomUUID();
+      const id = input.id ?? randomUUID();
+      if (input.id !== void 0) this.validateCallerDefinedDatasetId(input.id);
       const record = {
         id,
         name: input.name,
@@ -3165,6 +3172,11 @@ var DatasetsSpanner = class _DatasetsSpanner extends DatasetsStorage {
       });
       return record;
     } catch (error) {
+      if (input.id !== void 0 && hasErrorCode(error, /* @__PURE__ */ new Set([6, "ALREADY_EXISTS"]))) {
+        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("SPANNER", "CREATE_DATASET", "FAILED"),
@@ -3587,6 +3599,7 @@ var DatasetsSpanner = class _DatasetsSpanner extends DatasetsStorage {
                 id: args.id,
                 datasetId: args.datasetId,
                 datasetVersion: newVersion,
+                externalId: existing.externalId ?? null,
                 organizationId,
                 projectId,
                 validTo: null,
@@ -3609,6 +3622,7 @@ var DatasetsSpanner = class _DatasetsSpanner extends DatasetsStorage {
               id: args.id,
               datasetId: args.datasetId,
               datasetVersion: newVersion,
+              externalId: existing.externalId ?? null,
               organizationId,
               projectId,
               input: merged.input,
@@ -3668,6 +3682,7 @@ var DatasetsSpanner = class _DatasetsSpanner extends DatasetsStorage {
                 id: args.id,
                 datasetId: args.datasetId,
                 datasetVersion: newVersion,
+                externalId: existing.externalId ?? null,
                 organizationId,
                 projectId,
                 validTo: null,
@@ -3937,55 +3952,75 @@ var DatasetsSpanner = class _DatasetsSpanner extends DatasetsStorage {
   async _doBatchInsertItems(input) {
     try {
       if (input.items.length === 0) return [];
-      const now = /* @__PURE__ */ new Date();
-      const prepared = input.items.map((item) => ({ id: randomUUID(), item }));
       let result = [];
       await this.db.runWithAbortRetry(
         () => this.database.runTransactionAsync(async (tx) => {
           try {
-            const { version: newVersion, organizationId, projectId } = await this.bumpVersion(tx, input.datasetId, now);
-            for (const { id, item } of prepared) {
-              await this.db.insert({
-                tableName: TABLE_DATASET_ITEMS,
-                record: {
-                  id,
+            const [datasetRows] = await tx.run({
+              sql: `SELECT * FROM ${quoteIdent(TABLE_DATASETS, "table name")} WHERE ${quoteIdent("id", "column name")} = @id LIMIT 1`,
+              params: { id: input.datasetId },
+              json: true
+            });
+            const dataset = datasetRows[0];
+            if (!dataset) {
+              throw new MastraError({
+                id: createStorageErrorId("SPANNER", "BATCH_INSERT_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] : []))];
+            let historyRows = [];
+            if (externalIds.length > 0) {
+              const [rows] = await tx.run({
+                sql: `SELECT * FROM ${quoteIdent(TABLE_DATASET_ITEMS, "table name")} WHERE ${quoteIdent("datasetId", "column name")} = @datasetId AND ${quoteIdent("externalId", "column name")} IN UNNEST(@externalIds) ORDER BY ${quoteIdent("datasetVersion", "column name")}`,
+                params: { datasetId: input.datasetId, externalIds },
+                types: { externalIds: { type: "array", child: "string" } },
+                json: true
+              });
+              historyRows = rows.map(rowToItemRow);
+            }
+            const plan = this.planDatasetItemBatch(input.items, historyRows, randomUUID);
+            const resolved = new Map(
+              [...plan.existingCurrentItems].map(([id, row]) => [id, this.datasetItemFromRow(row)])
+            );
+            if (plan.inserts.length > 0) {
+              const now = /* @__PURE__ */ new Date();
+              const {
+                version: newVersion,
+                organizationId,
+                projectId
+              } = await this.bumpVersion(tx, input.datasetId, now);
+              for (const insert of plan.inserts) {
+                const item = {
+                  id: insert.id,
                   datasetId: input.datasetId,
                   datasetVersion: newVersion,
+                  externalId: insert.item.externalId ?? null,
                   organizationId,
                   projectId,
-                  validTo: null,
-                  isDeleted: false,
-                  input: item.input,
-                  groundTruth: item.groundTruth ?? null,
-                  expectedTrajectory: item.expectedTrajectory ?? null,
-                  toolMocks: item.toolMocks ?? null,
-                  requestContext: item.requestContext ?? null,
-                  metadata: item.metadata ?? null,
-                  source: item.source ?? null,
+                  input: insert.item.input,
+                  groundTruth: insert.item.groundTruth,
+                  expectedTrajectory: insert.item.expectedTrajectory,
+                  toolMocks: insert.item.toolMocks,
+                  requestContext: insert.item.requestContext,
+                  metadata: insert.item.metadata,
+                  source: insert.item.source,
                   createdAt: now,
                   updatedAt: now
-                },
-                transaction: tx
-              });
+                };
+                await this.db.insert({
+                  tableName: TABLE_DATASET_ITEMS,
+                  record: { ...item, validTo: null, isDeleted: false },
+                  transaction: tx
+                });
+                resolved.set(item.id, item);
+              }
+              await this.insertVersionRow(tx, input.datasetId, newVersion, now);
             }
-            await this.insertVersionRow(tx, input.datasetId, newVersion, now);
             await tx.commit();
-            result = prepared.map(({ id, item }) => ({
-              id,
-              datasetId: input.datasetId,
-              datasetVersion: newVersion,
-              organizationId,
-              projectId,
-              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
-            }));
+            result = plan.resolvedIds.map((id) => resolved.get(id));
           } catch (err) {
             await tx.rollback().catch((rollbackErr) => {
               throw new AggregateError([err, rollbackErr], "Transaction and rollback both failed");
@@ -4029,6 +4064,7 @@ var DatasetsSpanner = class _DatasetsSpanner extends DatasetsStorage {
                   id: existing.id,
                   datasetId: input.datasetId,
                   datasetVersion: newVersion,
+                  externalId: existing.externalId ?? null,
                   organizationId,
                   projectId,
                   validTo: null,