@mastra/dsql 1.4.0 → 1.5.0

package/dist/index.js

@@ -1,6 +1,6 @@
 import { AuroraDSQLClient } from "@aws/aurora-dsql-node-postgres-connector";
 import { ErrorCategory, ErrorDomain, MastraError } from "@mastra/core/error";
-import { AgentsStorage, MastraStorage, MemoryStorage, ObservabilityStorage, ScoresStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, TABLE_MESSAGES, TABLE_RESOURCES, TABLE_SCHEMAS, TABLE_SCORERS, TABLE_SPANS, TABLE_THREADS, TABLE_WORKFLOW_SNAPSHOT, TraceStatus, WorkflowsStorage, calculatePagination, createStorageErrorId, getDefaultValue, getSqlType, listTracesArgsSchema, matchesExpectedWorkflowStatus, normalizePerPage, storageMessageMatchesMetadataFilter, toTraceSpans, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
+import { AgentsStorage, MastraStorage, MemoryStorage, ObservabilityStorage, ScoresStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, TABLE_MESSAGES, TABLE_RESOURCES, TABLE_SCHEMAS, TABLE_SCORERS, TABLE_SPANS, TABLE_THREADS, TABLE_WORKFLOW_SNAPSHOT, TraceStatus, WorkflowsStorage, calculatePagination, createStorageErrorId, executeRetentionPrune, getDefaultValue, getSqlType, listTracesArgsSchema, matchesExpectedWorkflowStatus, normalizePerPage, resolveRetentionTargets, retentionCutoffMs, storageMessageMatchesMetadataFilter, toTraceSpans, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
 import { Pool } from "pg";
 import { MastraBase } from "@mastra/core/base";
 import { parseSqlIdentifier } from "@mastra/core/utils";
@@ -662,6 +662,16 @@ var DsqlDB = class extends MastraBase {
 			}, error);
 		});
 	}
+	async pruneBatch({ tableName, column, cutoff, limit }) {
+		if (!/^[A-Za-z_][A-Za-z0-9_]*$/.test(column)) throw new Error(`Invalid retention identifier: ${tableName}.${column}`);
+		const table = getTableName$1({
+			indexName: tableName,
+			schemaName: getSchemaName$1(this.schemaName)
+		});
+		return (await this.client.query(`DELETE FROM ${table} WHERE ("traceId", "spanId") IN (
+        SELECT "traceId", "spanId" FROM ${table} WHERE "${column}" < $1 ORDER BY "${column}" LIMIT $2
+      )`, [cutoff, limit])).rowCount ?? 0;
+	}
 	async clearTable({ tableName }) {
 		try {
 			await withRetry(async () => {
@@ -2773,6 +2783,11 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 //#endregion
 //#region src/storage/domains/observability/index.ts
 var ObservabilityDSQL = class ObservabilityDSQL extends ObservabilityStorage {
+	static retentionTables = { spans: {
+		table: TABLE_SPANS,
+		column: "startedAtZ",
+		indexed: true
+	} };
 	#db;
 	#schema;
 	#skipDefaultIndexes;
@@ -2791,6 +2806,24 @@ var ObservabilityDSQL = class ObservabilityDSQL extends ObservabilityStorage {
 		this.#skipDefaultIndexes = skipDefaultIndexes;
 		this.#indexes = indexes?.filter((idx) => ObservabilityDSQL.MANAGED_TABLES.includes(idx.table));
 	}
+	async prune(policies, options) {
+		return executeRetentionPrune({
+			domain: "observability",
+			targets: resolveRetentionTargets({
+				policies,
+				descriptor: ObservabilityDSQL.retentionTables,
+				order: ["spans"]
+			}),
+			options,
+			cutoffFor: (target, now) => new Date(retentionCutoffMs(target.policy, now)),
+			deleteBatch: (target, cutoff, limit) => this.#db.pruneBatch({
+				tableName: target.table,
+				column: target.column,
+				cutoff,
+				limit
+			})
+		});
+	}
 	async init() {
 		await this.#db.createTable({
 			tableName: TABLE_SPANS,
@@ -2805,6 +2838,11 @@ var ObservabilityDSQL = class ObservabilityDSQL extends ObservabilityStorage {
 	getDefaultIndexDefinitions() {
 		const schemaPrefix = this.#schema !== "public" ? `${this.#schema}_` : "";
 		return [
+			{
+				name: `${schemaPrefix}mastra_ai_spans_startedatz_idx`,
+				table: TABLE_SPANS,
+				columns: ["startedAtZ"]
+			},
 			{
 				name: `${schemaPrefix}mastra_ai_spans_traceid_startedat_idx`,
 				table: TABLE_SPANS,
@@ -4053,7 +4091,8 @@ var DSQLStore = class extends MastraStorage {
 			super({
 				id: config.id,
 				name: "DSQLStore",
-				disableInit: config.disableInit
+				disableInit: config.disableInit,
+				retention: config.retention
 			});
 			this.schema = config.schemaName || "public";
 			if (isPoolConfig(config)) {