@mastra/spanner 1.6.1 → 1.6.2-alpha.1

package/dist/index.js

@@ -1,6 +1,6 @@
 import { Spanner } from "@google-cloud/spanner";
 import { ErrorCategory, ErrorDomain, MastraError } from "@mastra/core/error";
-import { AGENTS_SCHEMA, AGENT_VERSIONS_SCHEMA, AgentsStorage, BackgroundTasksStorage, BlobStore, ChannelsStorage, DatasetsStorage, ExperimentsStorage, FavoritesStorage, MCPClientsStorage, MCPServersStorage, MCP_CLIENTS_SCHEMA, MCP_CLIENT_VERSIONS_SCHEMA, MCP_SERVERS_SCHEMA, MCP_SERVER_VERSIONS_SCHEMA, METRIC_DISTINCT_COLUMNS, MastraCompositeStore, MemoryStorage, ObservabilityStorage, PROMPT_BLOCKS_SCHEMA, PROMPT_BLOCK_VERSIONS_SCHEMA, PromptBlocksStorage, SCORER_DEFINITIONS_SCHEMA, SCORER_DEFINITION_VERSIONS_SCHEMA, SKILLS_SCHEMA, SKILL_BLOBS_SCHEMA, SKILL_VERSIONS_SCHEMA, SchedulesStorage, ScorerDefinitionsStorage, ScoresStorage, SkillsStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, TABLE_BACKGROUND_TASKS, TABLE_CHANNEL_CONFIG, TABLE_CHANNEL_INSTALLATIONS, TABLE_CONFIGS, TABLE_DATASETS, TABLE_DATASET_ITEMS, TABLE_DATASET_VERSIONS, TABLE_EXPERIMENTS, TABLE_EXPERIMENT_RESULTS, TABLE_FAVORITES, TABLE_MCP_CLIENTS, TABLE_MCP_CLIENT_VERSIONS, TABLE_MCP_SERVERS, TABLE_MCP_SERVER_VERSIONS, TABLE_MESSAGES, TABLE_PROMPT_BLOCKS, TABLE_PROMPT_BLOCK_VERSIONS, TABLE_RESOURCES, TABLE_SCHEDULES, TABLE_SCHEDULE_TRIGGERS, TABLE_SCHEMAS, TABLE_SCORERS, TABLE_SCORER_DEFINITIONS, TABLE_SCORER_DEFINITION_VERSIONS, TABLE_SKILLS, TABLE_SKILL_BLOBS, TABLE_SKILL_VERSIONS, TABLE_SPANS, TABLE_THREADS, TABLE_WORKFLOW_DEFINITIONS, TABLE_WORKFLOW_SNAPSHOT, TABLE_WORKSPACES, TABLE_WORKSPACE_VERSIONS, TraceStatus, WORKFLOW_DEFINITIONS_SCHEMA, WORKSPACES_SCHEMA, WORKSPACE_VERSIONS_SCHEMA, WorkflowDefinitionsStorage, WorkflowsStorage, WorkspacesStorage, calculatePagination, createStorageErrorId, getDefaultValue, hasErrorCode, listMetricsArgsSchema, listTracesArgsSchema, normalizePerPage, normalizeScheduleTarget, toTraceSpans, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
+import { AGENTS_SCHEMA, AGENT_VERSIONS_SCHEMA, AgentsStorage, BackgroundTasksStorage, BlobStore, ChannelsStorage, DatasetsStorage, ExperimentsStorage, FavoritesStorage, MCPClientsStorage, MCPServersStorage, MCP_CLIENTS_SCHEMA, MCP_CLIENT_VERSIONS_SCHEMA, MCP_SERVERS_SCHEMA, MCP_SERVER_VERSIONS_SCHEMA, METRIC_DISTINCT_COLUMNS, MastraCompositeStore, MemoryStorage, ObservabilityStorage, PROMPT_BLOCKS_SCHEMA, PROMPT_BLOCK_VERSIONS_SCHEMA, PromptBlocksStorage, SCORER_DEFINITIONS_SCHEMA, SCORER_DEFINITION_VERSIONS_SCHEMA, SKILLS_SCHEMA, SKILL_BLOBS_SCHEMA, SKILL_VERSIONS_SCHEMA, SchedulesStorage, ScorerDefinitionsStorage, ScoresStorage, SkillsStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, TABLE_BACKGROUND_TASKS, TABLE_CHANNEL_CONFIG, TABLE_CHANNEL_INSTALLATIONS, TABLE_CONFIGS, TABLE_DATASETS, TABLE_DATASET_ITEMS, TABLE_DATASET_VERSIONS, TABLE_EXPERIMENTS, TABLE_EXPERIMENT_RESULTS, TABLE_FAVORITES, TABLE_MCP_CLIENTS, TABLE_MCP_CLIENT_VERSIONS, TABLE_MCP_SERVERS, TABLE_MCP_SERVER_VERSIONS, TABLE_MESSAGES, TABLE_PROMPT_BLOCKS, TABLE_PROMPT_BLOCK_VERSIONS, TABLE_RESOURCES, TABLE_SCHEDULES, TABLE_SCHEDULE_TRIGGERS, TABLE_SCHEMAS, TABLE_SCORERS, TABLE_SCORER_DEFINITIONS, TABLE_SCORER_DEFINITION_VERSIONS, TABLE_SKILLS, TABLE_SKILL_BLOBS, TABLE_SKILL_VERSIONS, TABLE_SPANS, TABLE_THREADS, TABLE_WORKFLOW_DEFINITIONS, TABLE_WORKFLOW_SNAPSHOT, TABLE_WORKSPACES, TABLE_WORKSPACE_VERSIONS, TraceStatus, WORKFLOW_DEFINITIONS_SCHEMA, WORKSPACES_SCHEMA, WORKSPACE_VERSIONS_SCHEMA, WorkflowDefinitionsStorage, WorkflowsStorage, WorkspacesStorage, calculatePagination, createStorageErrorId, getDefaultValue, hasErrorCode, listMetricsArgsSchema, listTracesArgsSchema, matchesExpectedWorkflowStatus, normalizePerPage, normalizeScheduleTarget, toTraceSpans, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
 import { MastraBase } from "@mastra/core/base";
 import { parseSqlIdentifier } from "@mastra/core/utils";
 import { randomUUID } from "crypto";
@@ -3650,8 +3650,11 @@ var DatasetsSpanner = class DatasetsSpanner extends DatasetsStorage {
 			if (args.datasetVersion !== void 0) {
 				sql = `SELECT * FROM ${tableName}
                WHERE ${quoteIdent("id", "column name")} = @id
-                 AND ${quoteIdent("datasetVersion", "column name")} = @datasetVersion
-                 AND ${quoteIdent("isDeleted", "column name")} = FALSE LIMIT 1`;
+                 AND ${quoteIdent("datasetVersion", "column name")} <= @datasetVersion
+                 AND (${quoteIdent("validTo", "column name")} IS NULL OR ${quoteIdent("validTo", "column name")} > @datasetVersion)
+                 AND ${quoteIdent("isDeleted", "column name")} = FALSE
+               ORDER BY ${quoteIdent("datasetVersion", "column name")} DESC
+               LIMIT 1`;
 				params.datasetVersion = args.datasetVersion;
 			} else sql = `SELECT * FROM ${tableName}
                WHERE ${quoteIdent("id", "column name")} = @id
@@ -3982,8 +3985,9 @@ function rowToExperiment(row) {
 		datasetVersion: t.datasetVersion == null ? null : Number(t.datasetVersion),
 		organizationId: t.organizationId ?? null,
 		projectId: t.projectId ?? null,
-		targetType: t.targetType,
-		targetId: String(t.targetId),
+		targetType: t.targetType ?? null,
+		targetId: t.targetId == null ? null : String(t.targetId),
+		scorerIds: t.scorerIds ?? null,
 		status: t.status,
 		totalItems: Number(t.totalItems ?? 0),
 		succeededCount: Number(t.succeededCount ?? 0),
@@ -4015,6 +4019,7 @@ function rowToExperimentResult(row) {
 		startedAt: toDate(t.startedAt),
 		completedAt: toDate(t.completedAt),
 		retryCount: Number(t.retryCount ?? 0),
+		attempt: t.attempt == null ? 0 : Number(t.attempt),
 		traceId: t.traceId ?? null,
 		status: t.status ?? null,
 		tags: t.tags ?? null,
@@ -4067,7 +4072,8 @@ var ExperimentsSpanner = class ExperimentsSpanner extends ExperimentsStorage {
 				"experimentSetId",
 				"comparisonId",
 				"variantId",
-				"trialIndex"
+				"trialIndex",
+				"scorerIds"
 			]
 		});
 		await this.db.alterTable({
@@ -4076,7 +4082,8 @@ var ExperimentsSpanner = class ExperimentsSpanner extends ExperimentsStorage {
 			ifNotExists: [
 				"comment",
 				"organizationId",
-				"projectId"
+				"projectId",
+				"attempt"
 			]
 		});
 		await this.createDefaultIndexes();
@@ -4105,9 +4112,13 @@ var ExperimentsSpanner = class ExperimentsSpanner extends ExperimentsStorage {
 				columns: ["experimentId", "startedAt"]
 			},
 			{
-				name: "mastra_experiment_results_exp_item_idx",
+				name: "mastra_experiment_results_exp_item_attempt_idx",
 				table: TABLE_EXPERIMENT_RESULTS,
-				columns: ["experimentId", "itemId"],
+				columns: [
+					"experimentId",
+					"itemId",
+					"attempt"
+				],
 				unique: true
 			},
 			{
@@ -4124,6 +4135,7 @@ var ExperimentsSpanner = class ExperimentsSpanner extends ExperimentsStorage {
 	}
 	async createDefaultIndexes() {
 		if (this.skipDefaultIndexes) return;
+		await this.db.dropIndex("mastra_experiment_results_exp_item_idx").catch(() => {});
 		await this.db.createIndexes(this.getDefaultIndexDefinitions());
 	}
 	async createCustomIndexes() {
@@ -4153,8 +4165,9 @@ var ExperimentsSpanner = class ExperimentsSpanner extends ExperimentsStorage {
 				datasetVersion: input.datasetVersion ?? null,
 				organizationId: input.organizationId ?? null,
 				projectId: input.projectId ?? null,
-				targetType: input.targetType,
-				targetId: input.targetId,
+				targetType: input.targetType ?? null,
+				targetId: input.targetId ?? null,
+				scorerIds: input.scorerIds ?? null,
 				status: "pending",
 				totalItems: input.totalItems,
 				succeededCount: 0,
@@ -4185,6 +4198,7 @@ var ExperimentsSpanner = class ExperimentsSpanner extends ExperimentsStorage {
 					projectId: experiment.projectId,
 					targetType: experiment.targetType,
 					targetId: experiment.targetId,
+					scorerIds: experiment.scorerIds,
 					status: experiment.status,
 					totalItems: experiment.totalItems,
 					succeededCount: 0,
@@ -4464,6 +4478,7 @@ var ExperimentsSpanner = class ExperimentsSpanner extends ExperimentsStorage {
 				startedAt: input.startedAt,
 				completedAt: input.completedAt,
 				retryCount: input.retryCount,
+				attempt: input.attempt ?? 0,
 				traceId: input.traceId ?? null,
 				status: input.status ?? null,
 				tags: input.tags ?? null,
@@ -4486,6 +4501,7 @@ var ExperimentsSpanner = class ExperimentsSpanner extends ExperimentsStorage {
 					startedAt: result.startedAt,
 					completedAt: result.completedAt,
 					retryCount: result.retryCount,
+					attempt: result.attempt,
 					traceId: result.traceId,
 					status: result.status,
 					tags: result.tags,
@@ -4506,6 +4522,116 @@ var ExperimentsSpanner = class ExperimentsSpanner extends ExperimentsStorage {
 			}, error);
 		}
 	}
+	async upsertExperimentResult(input) {
+		try {
+			const attempt = input.attempt ?? 0;
+			const [rows] = await this.database.run({
+				sql: `SELECT ${quoteIdent("id", "column name")} FROM ${quoteIdent(TABLE_EXPERIMENT_RESULTS, "table name")}
+              WHERE ${quoteIdent("experimentId", "column name")} = @experimentId
+                AND ${quoteIdent("itemId", "column name")} = @itemId
+                AND COALESCE(${quoteIdent("attempt", "column name")}, 0) = @attempt`,
+				params: {
+					experimentId: input.experimentId,
+					itemId: input.itemId,
+					attempt
+				},
+				json: true
+			});
+			let existingId = rows[0]?.id;
+			if (!existingId) try {
+				return await this.addExperimentResult({
+					...input,
+					attempt
+				});
+			} catch (insertError) {
+				if (!hasErrorCode(insertError, /* @__PURE__ */ new Set([6]))) throw insertError;
+				const [winnerRows] = await this.database.run({
+					sql: `SELECT ${quoteIdent("id", "column name")} FROM ${quoteIdent(TABLE_EXPERIMENT_RESULTS, "table name")}
+                  WHERE ${quoteIdent("experimentId", "column name")} = @experimentId
+                    AND ${quoteIdent("itemId", "column name")} = @itemId
+                    AND COALESCE(${quoteIdent("attempt", "column name")}, 0) = @attempt`,
+					params: {
+						experimentId: input.experimentId,
+						itemId: input.itemId,
+						attempt
+					},
+					json: true
+				});
+				existingId = winnerRows[0]?.id;
+				if (!existingId) throw insertError;
+			}
+			await this.db.runDml({
+				sql: `UPDATE ${quoteIdent(TABLE_EXPERIMENT_RESULTS, "table name")} SET
+                ${quoteIdent("itemDatasetVersion", "column name")} = @itemDatasetVersion,
+                ${quoteIdent("organizationId", "column name")} = @organizationId,
+                ${quoteIdent("projectId", "column name")} = @projectId,
+                ${quoteIdent("input", "column name")} = @input,
+                ${quoteIdent("output", "column name")} = @output,
+                ${quoteIdent("groundTruth", "column name")} = @groundTruth,
+                ${quoteIdent("error", "column name")} = @error,
+                ${quoteIdent("startedAt", "column name")} = @startedAt,
+                ${quoteIdent("completedAt", "column name")} = @completedAt,
+                ${quoteIdent("retryCount", "column name")} = @retryCount,
+                ${quoteIdent("attempt", "column name")} = @attempt,
+                ${quoteIdent("traceId", "column name")} = @traceId,
+                ${quoteIdent("status", "column name")} = @status,
+                ${quoteIdent("tags", "column name")} = @tags,
+                ${quoteIdent("toolMockReport", "column name")} = @toolMockReport
+              WHERE ${quoteIdent("id", "column name")} = @id`,
+				params: {
+					id: existingId,
+					itemDatasetVersion: input.itemDatasetVersion ?? null,
+					organizationId: input.organizationId ?? null,
+					projectId: input.projectId ?? null,
+					input: JSON.stringify(input.input),
+					output: input.output == null ? null : JSON.stringify(input.output),
+					groundTruth: input.groundTruth == null ? null : JSON.stringify(input.groundTruth),
+					error: input.error == null ? null : JSON.stringify(input.error),
+					startedAt: input.startedAt,
+					completedAt: input.completedAt,
+					retryCount: input.retryCount,
+					attempt,
+					traceId: input.traceId ?? null,
+					status: input.status ?? null,
+					tags: input.tags == null ? null : JSON.stringify(input.tags),
+					toolMockReport: input.toolMockReport == null ? null : JSON.stringify(input.toolMockReport)
+				},
+				types: {
+					itemDatasetVersion: "int64",
+					organizationId: "string",
+					projectId: "string",
+					input: "json",
+					output: "json",
+					groundTruth: "json",
+					error: "json",
+					traceId: "string",
+					status: "string",
+					tags: "json",
+					toolMockReport: "json"
+				}
+			});
+			const updated = await this.getExperimentResultById({ id: existingId });
+			if (!updated) throw new MastraError({
+				id: createStorageErrorId("SPANNER", "UPSERT_EXPERIMENT_RESULT", "NOT_FOUND"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.USER,
+				text: `Experiment result ${existingId} not found after upsert`,
+				details: { id: existingId }
+			});
+			return updated;
+		} catch (error) {
+			if (error instanceof MastraError) throw error;
+			throw new MastraError({
+				id: createStorageErrorId("SPANNER", "UPSERT_EXPERIMENT_RESULT", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY,
+				details: {
+					experimentId: input.experimentId,
+					itemId: input.itemId
+				}
+			}, error);
+		}
+	}
 	async updateExperimentResult(input) {
 		try {
 			if (input.status === void 0 && input.tags === void 0 && input.comment === void 0) {
@@ -10799,9 +10925,10 @@ var ScoresSpanner = class ScoresSpanner extends ScoresStorage {
 			}, error);
 		}
 		try {
-			const scoreId = randomUUID();
+			const suppliedId = validatedScore.id;
+			const scoreId = suppliedId ?? randomUUID();
 			const now = /* @__PURE__ */ new Date();
-			const { scorer, preprocessStepResult, analyzeStepResult, metadata, input, output, additionalContext, requestContext, entity, ...rest } = validatedScore;
+			const { id: _suppliedId, scorer, preprocessStepResult, analyzeStepResult, metadata, input, output, additionalContext, requestContext, entity, ...rest } = validatedScore;
 			const insertedRecord = {
 				id: scoreId,
 				...rest,
@@ -10817,7 +10944,11 @@ var ScoresSpanner = class ScoresSpanner extends ScoresStorage {
 				createdAt: now,
 				updatedAt: now
 			};
-			await this.db.insert({
+			if (suppliedId) await this.db.upsert({
+				tableName: TABLE_SCORERS,
+				record: insertedRecord
+			});
+			else await this.db.insert({
 				tableName: TABLE_SCORERS,
 				record: insertedRecord
 			});
@@ -12093,9 +12224,14 @@ var WorkflowsSpanner = class WorkflowsSpanner extends WorkflowsStorage {
 							}
 						}, /* @__PURE__ */ new Error(`Snapshot not found for runId ${runId}`));
 					}
+					const { expectedStatus, ...state } = opts;
+					if (!matchesExpectedWorkflowStatus(snapshot.status, expectedStatus)) {
+						await tx.rollback();
+						return;
+					}
 					updated = {
 						...snapshot,
-						...opts
+						...state
 					};
 					await this.db.update({
 						tableName: TABLE_WORKFLOW_SNAPSHOT,