@mastra/mssql 0.2.3 → 0.3.0-alpha.0

package/src/storage/domains/scores/index.ts

@@ -0,0 +1,315 @@
+import type { PaginationInfo, StoragePagination } from '@mastra/core';
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { ScoreRowData } from '@mastra/core/scores';
+import { ScoresStorage, TABLE_SCORERS } from '@mastra/core/storage';
+import type { ConnectionPool } from 'mssql';
+import type { StoreOperationsMSSQL } from '../operations';
+import { getSchemaName, getTableName } from '../utils';
+
+function parseJSON(jsonString: string): any {
+  try {
+    return JSON.parse(jsonString);
+  } catch {
+    return jsonString;
+  }
+}
+
+function transformScoreRow(row: Record<string, any>): ScoreRowData {
+  let input = undefined;
+  let output = undefined;
+  let preprocessStepResult = undefined;
+  let analyzeStepResult = undefined;
+
+  if (row.input) {
+    input = parseJSON(row.input);
+  }
+
+  if (row.output) {
+    output = parseJSON(row.output);
+  }
+
+  if (row.preprocessStepResult) {
+    preprocessStepResult = parseJSON(row.preprocessStepResult);
+  }
+
+  if (row.analyzeStepResult) {
+    analyzeStepResult = parseJSON(row.analyzeStepResult);
+  }
+
+  return {
+    ...row,
+    input,
+    output,
+    preprocessStepResult,
+    analyzeStepResult,
+    createdAt: row.createdAt,
+    updatedAt: row.updatedAt,
+  } as ScoreRowData;
+}
+
+export class ScoresMSSQL extends ScoresStorage {
+  public pool: ConnectionPool;
+  private operations: StoreOperationsMSSQL;
+  private schema?: string;
+
+  constructor({
+    pool,
+    operations,
+    schema,
+  }: {
+    pool: ConnectionPool;
+    operations: StoreOperationsMSSQL;
+    schema?: string;
+  }) {
+    super();
+    this.pool = pool;
+    this.operations = operations;
+    this.schema = schema;
+  }
+
+  async getScoreById({ id }: { id: string }): Promise<ScoreRowData | null> {
+    try {
+      const request = this.pool.request();
+      request.input('p1', id);
+      const result = await request.query(
+        `SELECT * FROM ${getTableName({ indexName: TABLE_SCORERS, schemaName: getSchemaName(this.schema) })} WHERE id = @p1`,
+      );
+
+      if (result.recordset.length === 0) {
+        return null;
+      }
+
+      return transformScoreRow(result.recordset[0]);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_MSSQL_STORE_GET_SCORE_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { id },
+        },
+        error,
+      );
+    }
+  }
+
+  async saveScore(score: Omit<ScoreRowData, 'id' | 'createdAt' | 'updatedAt'>): Promise<{ score: ScoreRowData }> {
+    try {
+      // Generate ID like other storage implementations
+      const scoreId = crypto.randomUUID();
+
+      const { input, output, preprocessStepResult, analyzeStepResult, ...rest } = score;
+      await this.operations.insert({
+        tableName: TABLE_SCORERS,
+        record: {
+          id: scoreId,
+          ...rest,
+          input: JSON.stringify(input),
+          output: JSON.stringify(output),
+          preprocessStepResult: preprocessStepResult ? JSON.stringify(preprocessStepResult) : null,
+          analyzeStepResult: analyzeStepResult ? JSON.stringify(analyzeStepResult) : null,
+          createdAt: new Date().toISOString(),
+          updatedAt: new Date().toISOString(),
+        },
+      });
+
+      const scoreFromDb = await this.getScoreById({ id: scoreId });
+      return { score: scoreFromDb! };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_MSSQL_STORE_SAVE_SCORE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByScorerId({
+    scorerId,
+    pagination,
+  }: {
+    scorerId: string;
+    pagination: StoragePagination;
+    entityId?: string;
+    entityType?: string;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    try {
+      const request = this.pool.request();
+      request.input('p1', scorerId);
+
+      const totalResult = await request.query(
+        `SELECT COUNT(*) as count FROM ${getTableName({ indexName: TABLE_SCORERS, schemaName: getSchemaName(this.schema) })} WHERE [scorerId] = @p1`,
+      );
+
+      const total = totalResult.recordset[0]?.count || 0;
+
+      if (total === 0) {
+        return {
+          pagination: {
+            total: 0,
+            page: pagination.page,
+            perPage: pagination.perPage,
+            hasMore: false,
+          },
+          scores: [],
+        };
+      }
+
+      const dataRequest = this.pool.request();
+      dataRequest.input('p1', scorerId);
+      dataRequest.input('p2', pagination.perPage);
+      dataRequest.input('p3', pagination.page * pagination.perPage);
+
+      const result = await dataRequest.query(
+        `SELECT * FROM ${getTableName({ indexName: TABLE_SCORERS, schemaName: getSchemaName(this.schema) })} WHERE [scorerId] = @p1 ORDER BY [createdAt] DESC OFFSET @p3 ROWS FETCH NEXT @p2 ROWS ONLY`,
+      );
+
+      return {
+        pagination: {
+          total: Number(total),
+          page: pagination.page,
+          perPage: pagination.perPage,
+          hasMore: Number(total) > (pagination.page + 1) * pagination.perPage,
+        },
+        scores: result.recordset.map(row => transformScoreRow(row)),
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_MSSQL_STORE_GET_SCORES_BY_SCORER_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { scorerId },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByRunId({
+    runId,
+    pagination,
+  }: {
+    runId: string;
+    pagination: StoragePagination;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    try {
+      const request = this.pool.request();
+      request.input('p1', runId);
+
+      const totalResult = await request.query(
+        `SELECT COUNT(*) as count FROM ${getTableName({ indexName: TABLE_SCORERS, schemaName: getSchemaName(this.schema) })} WHERE [runId] = @p1`,
+      );
+
+      const total = totalResult.recordset[0]?.count || 0;
+
+      if (total === 0) {
+        return {
+          pagination: {
+            total: 0,
+            page: pagination.page,
+            perPage: pagination.perPage,
+            hasMore: false,
+          },
+          scores: [],
+        };
+      }
+
+      const dataRequest = this.pool.request();
+      dataRequest.input('p1', runId);
+      dataRequest.input('p2', pagination.perPage);
+      dataRequest.input('p3', pagination.page * pagination.perPage);
+
+      const result = await dataRequest.query(
+        `SELECT * FROM ${getTableName({ indexName: TABLE_SCORERS, schemaName: getSchemaName(this.schema) })} WHERE [runId] = @p1 ORDER BY [createdAt] DESC OFFSET @p3 ROWS FETCH NEXT @p2 ROWS ONLY`,
+      );
+
+      return {
+        pagination: {
+          total: Number(total),
+          page: pagination.page,
+          perPage: pagination.perPage,
+          hasMore: Number(total) > (pagination.page + 1) * pagination.perPage,
+        },
+        scores: result.recordset.map(row => transformScoreRow(row)),
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_MSSQL_STORE_GET_SCORES_BY_RUN_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { runId },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByEntityId({
+    entityId,
+    entityType,
+    pagination,
+  }: {
+    pagination: StoragePagination;
+    entityId: string;
+    entityType: string;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    try {
+      const request = this.pool.request();
+      request.input('p1', entityId);
+      request.input('p2', entityType);
+
+      const totalResult = await request.query(
+        `SELECT COUNT(*) as count FROM ${getTableName({ indexName: TABLE_SCORERS, schemaName: getSchemaName(this.schema) })} WHERE [entityId] = @p1 AND [entityType] = @p2`,
+      );
+
+      const total = totalResult.recordset[0]?.count || 0;
+
+      if (total === 0) {
+        return {
+          pagination: {
+            total: 0,
+            page: pagination.page,
+            perPage: pagination.perPage,
+            hasMore: false,
+          },
+          scores: [],
+        };
+      }
+
+      const dataRequest = this.pool.request();
+      dataRequest.input('p1', entityId);
+      dataRequest.input('p2', entityType);
+      dataRequest.input('p3', pagination.perPage);
+      dataRequest.input('p4', pagination.page * pagination.perPage);
+
+      const result = await dataRequest.query(
+        `SELECT * FROM ${getTableName({ indexName: TABLE_SCORERS, schemaName: getSchemaName(this.schema) })} WHERE [entityId] = @p1 AND [entityType] = @p2 ORDER BY [createdAt] DESC OFFSET @p4 ROWS FETCH NEXT @p3 ROWS ONLY`,
+      );
+
+      return {
+        pagination: {
+          total: Number(total),
+          page: pagination.page,
+          perPage: pagination.perPage,
+          hasMore: Number(total) > (pagination.page + 1) * pagination.perPage,
+        },
+        scores: result.recordset.map(row => transformScoreRow(row)),
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_MSSQL_STORE_GET_SCORES_BY_ENTITY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { entityId, entityType },
+        },
+        error,
+      );
+    }
+  }
+}

package/src/storage/domains/traces/index.ts

@@ -0,0 +1,212 @@
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { PaginationInfo, PaginationArgs } from '@mastra/core/storage';
+import { TABLE_TRACES, TracesStorage } from '@mastra/core/storage';
+import { parseFieldKey } from '@mastra/core/utils';
+import sql from 'mssql';
+import type { StoreOperationsMSSQL } from '../operations';
+import { getSchemaName, getTableName } from '../utils';
+
+export class TracesMSSQL extends TracesStorage {
+  public pool: sql.ConnectionPool;
+  private operations: StoreOperationsMSSQL;
+  private schema?: string;
+
+  constructor({
+    pool,
+    operations,
+    schema,
+  }: {
+    pool: sql.ConnectionPool;
+    operations: StoreOperationsMSSQL;
+    schema?: string;
+  }) {
+    super();
+    this.pool = pool;
+    this.operations = operations;
+    this.schema = schema;
+  }
+
+  /** @deprecated use getTracesPaginated instead*/
+  public async getTraces(args: {
+    name?: string;
+    scope?: string;
+    attributes?: Record<string, string>;
+    filters?: Record<string, any>;
+    page: number;
+    perPage?: number;
+    fromDate?: Date;
+    toDate?: Date;
+  }): Promise<any[]> {
+    if (args.fromDate || args.toDate) {
+      (args as any).dateRange = {
+        start: args.fromDate,
+        end: args.toDate,
+      };
+    }
+    const result = await this.getTracesPaginated(args);
+    return result.traces;
+  }
+
+  public async getTracesPaginated(
+    args: {
+      name?: string;
+      scope?: string;
+      attributes?: Record<string, string>;
+      filters?: Record<string, any>;
+    } & PaginationArgs,
+  ): Promise<
+    PaginationInfo & {
+      traces: any[];
+    }
+  > {
+    const { name, scope, page = 0, perPage: perPageInput, attributes, filters, dateRange } = args;
+    const fromDate = dateRange?.start;
+    const toDate = dateRange?.end;
+
+    const perPage = perPageInput !== undefined ? perPageInput : 100;
+    const currentOffset = page * perPage;
+
+    const paramMap: Record<string, any> = {};
+    const conditions: string[] = [];
+    let paramIndex = 1;
+
+    if (name) {
+      const paramName = `p${paramIndex++}`;
+      conditions.push(`[name] LIKE @${paramName}`);
+      paramMap[paramName] = `${name}%`;
+    }
+    if (scope) {
+      const paramName = `p${paramIndex++}`;
+      conditions.push(`[scope] = @${paramName}`);
+      paramMap[paramName] = scope;
+    }
+    if (attributes) {
+      Object.entries(attributes).forEach(([key, value]) => {
+        const parsedKey = parseFieldKey(key);
+        const paramName = `p${paramIndex++}`;
+        conditions.push(`JSON_VALUE([attributes], '$.${parsedKey}') = @${paramName}`);
+        paramMap[paramName] = value;
+      });
+    }
+    if (filters) {
+      Object.entries(filters).forEach(([key, value]) => {
+        const parsedKey = parseFieldKey(key);
+        const paramName = `p${paramIndex++}`;
+        conditions.push(`[${parsedKey}] = @${paramName}`);
+        paramMap[paramName] = value;
+      });
+    }
+    if (fromDate instanceof Date && !isNaN(fromDate.getTime())) {
+      const paramName = `p${paramIndex++}`;
+      conditions.push(`[createdAt] >= @${paramName}`);
+      paramMap[paramName] = fromDate.toISOString();
+    }
+    if (toDate instanceof Date && !isNaN(toDate.getTime())) {
+      const paramName = `p${paramIndex++}`;
+      conditions.push(`[createdAt] <= @${paramName}`);
+      paramMap[paramName] = toDate.toISOString();
+    }
+
+    const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+
+    const countQuery = `SELECT COUNT(*) as total FROM ${getTableName({ indexName: TABLE_TRACES, schemaName: getSchemaName(this.schema) })} ${whereClause}`;
+    let total = 0;
+    try {
+      const countRequest = this.pool.request();
+      Object.entries(paramMap).forEach(([key, value]) => {
+        if (value instanceof Date) {
+          countRequest.input(key, sql.DateTime, value);
+        } else {
+          countRequest.input(key, value);
+        }
+      });
+      const countResult = await countRequest.query(countQuery);
+      total = parseInt(countResult.recordset[0].total, 10);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_MSSQL_STORE_GET_TRACES_PAGINATED_FAILED_TO_RETRIEVE_TOTAL_COUNT',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            name: args.name ?? '',
+            scope: args.scope ?? '',
+          },
+        },
+        error,
+      );
+    }
+
+    if (total === 0) {
+      return {
+        traces: [],
+        total: 0,
+        page,
+        perPage,
+        hasMore: false,
+      };
+    }
+
+    const dataQuery = `SELECT * FROM ${getTableName({ indexName: TABLE_TRACES, schemaName: getSchemaName(this.schema) })} ${whereClause} ORDER BY [seq_id] DESC OFFSET @offset ROWS FETCH NEXT @limit ROWS ONLY`;
+    const dataRequest = this.pool.request();
+    Object.entries(paramMap).forEach(([key, value]) => {
+      if (value instanceof Date) {
+        dataRequest.input(key, sql.DateTime, value);
+      } else {
+        dataRequest.input(key, value);
+      }
+    });
+    dataRequest.input('offset', currentOffset);
+    dataRequest.input('limit', perPage);
+
+    try {
+      const rowsResult = await dataRequest.query(dataQuery);
+      const rows = rowsResult.recordset;
+      const traces = rows.map(row => ({
+        id: row.id,
+        parentSpanId: row.parentSpanId,
+        traceId: row.traceId,
+        name: row.name,
+        scope: row.scope,
+        kind: row.kind,
+        status: JSON.parse(row.status),
+        events: JSON.parse(row.events),
+        links: JSON.parse(row.links),
+        attributes: JSON.parse(row.attributes),
+        startTime: row.startTime,
+        endTime: row.endTime,
+        other: row.other,
+        createdAt: row.createdAt,
+      }));
+
+      return {
+        traces,
+        total,
+        page,
+        perPage,
+        hasMore: currentOffset + traces.length < total,
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_MSSQL_STORE_GET_TRACES_PAGINATED_FAILED_TO_RETRIEVE_TRACES',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            name: args.name ?? '',
+            scope: args.scope ?? '',
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async batchTraceInsert({ records }: { records: Record<string, any>[] }): Promise<void> {
+    this.logger.debug('Batch inserting traces', { count: records.length });
+    await this.operations.batchInsert({
+      tableName: TABLE_TRACES,
+      records,
+    });
+  }
+}

package/src/storage/domains/utils.ts

@@ -0,0 +1,12 @@
+import { parseSqlIdentifier } from '@mastra/core/utils';
+
+export function getSchemaName(schema?: string) {
+  return schema ? `[${parseSqlIdentifier(schema, 'schema name')}]` : undefined;
+}
+
+export function getTableName({ indexName, schemaName }: { indexName: string; schemaName?: string }) {
+  const parsedIndexName = parseSqlIdentifier(indexName, 'index name');
+  const quotedIndexName = `[${parsedIndexName}]`;
+  const quotedSchemaName = schemaName;
+  return quotedSchemaName ? `${quotedSchemaName}.${quotedIndexName}` : quotedIndexName;
+}