@mastra/pg 0.1.0-alpha.11 → 0.1.0-alpha.12

package/CHANGELOG.md

@@ -1,5 +1,12 @@
 # @mastra/pg
 
+## 0.1.0-alpha.12
+
+### Patch Changes
+
+- Updated dependencies [a9345f9]
+  - @mastra/core@0.2.0-alpha.102
+
 ## 0.1.0-alpha.11
 
 ### Patch Changes

package/dist/_tsup-dts-rollup.d.ts

@@ -0,0 +1,141 @@
+import { ArrayOperator } from '@mastra/core/filter';
+import { BaseFilterTranslator } from '@mastra/core/filter';
+import { BasicOperator } from '@mastra/core/filter';
+import { ElementOperator } from '@mastra/core/filter';
+import { Filter } from '@mastra/core/filter';
+import { IndexStats } from '@mastra/core/vector';
+import { LogicalOperator } from '@mastra/core/filter';
+import { MastraStorage } from '@mastra/core/storage';
+import { MastraVector } from '@mastra/core/vector';
+import { MessageType } from '@mastra/core/memory';
+import { NumericOperator } from '@mastra/core/filter';
+import { OperatorSupport } from '@mastra/core/filter';
+import { QueryResult } from '@mastra/core/vector';
+import { RegexOperator } from '@mastra/core/filter';
+import { StorageColumn } from '@mastra/core/storage';
+import { StorageGetMessagesArg } from '@mastra/core/storage';
+import { StorageThreadType } from '@mastra/core/memory';
+import { TABLE_NAMES } from '@mastra/core/storage';
+import { WorkflowRunState } from '@mastra/core/workflows';
+
+export declare function buildFilterQuery(filter: Filter, minScore: number): FilterResult;
+
+export declare const FILTER_OPERATORS: Record<string, OperatorFn>;
+
+declare type FilterOperator = {
+    sql: string;
+    needsValue: boolean;
+    transformValue?: (value: any) => any;
+};
+
+export declare interface FilterResult {
+    sql: string;
+    values: any[];
+}
+
+export declare const handleKey: (key: string) => string;
+
+declare type OperatorFn = (key: string, paramIndex: number, value?: any) => FilterOperator;
+
+export declare type OperatorType = BasicOperator | NumericOperator | ArrayOperator | ElementOperator | LogicalOperator | '$contains' | Exclude<RegexOperator, '$options'>;
+
+/**
+ * Translates MongoDB-style filters to PG compatible filters.
+ *
+ * Key differences from MongoDB:
+ *
+ * Logical Operators ($and, $or, $nor):
+ * - Can be used at the top level or nested within fields
+ * - Can take either a single condition or an array of conditions
+ *
+ */
+export declare class PGFilterTranslator extends BaseFilterTranslator {
+    protected getSupportedOperators(): OperatorSupport;
+    translate(filter: Filter): Filter;
+    private translateNode;
+    private translateRegexPattern;
+}
+
+declare class PgVector extends MastraVector {
+    private pool;
+    constructor(connectionString: string);
+    transformFilter(filter?: Filter): Filter;
+    query(indexName: string, queryVector: number[], topK?: number, filter?: Filter, includeVector?: boolean, minScore?: number): Promise<QueryResult[]>;
+    upsert(indexName: string, vectors: number[][], metadata?: Record<string, any>[], ids?: string[]): Promise<string[]>;
+    createIndex(indexName: string, dimension: number, metric?: 'cosine' | 'euclidean' | 'dotproduct'): Promise<void>;
+    listIndexes(): Promise<string[]>;
+    describeIndex(indexName: string): Promise<IndexStats>;
+    deleteIndex(indexName: string): Promise<void>;
+    truncateIndex(indexName: string): Promise<void>;
+    disconnect(): Promise<void>;
+}
+export { PgVector }
+export { PgVector as PgVector_alias_1 }
+
+declare type PostgresConfig = {
+    host: string;
+    port: number;
+    database: string;
+    user: string;
+    password: string;
+} | {
+    connectionString: string;
+};
+export { PostgresConfig }
+export { PostgresConfig as PostgresConfig_alias_1 }
+
+declare class PostgresStore extends MastraStorage {
+    private db;
+    private pgp;
+    constructor(config: PostgresConfig);
+    createTable({ tableName, schema, }: {
+        tableName: TABLE_NAMES;
+        schema: Record<string, StorageColumn>;
+    }): Promise<void>;
+    clearTable({ tableName }: {
+        tableName: TABLE_NAMES;
+    }): Promise<void>;
+    insert({ tableName, record }: {
+        tableName: TABLE_NAMES;
+        record: Record<string, any>;
+    }): Promise<void>;
+    load<R>({ tableName, keys }: {
+        tableName: TABLE_NAMES;
+        keys: Record<string, string>;
+    }): Promise<R | null>;
+    getThreadById({ threadId }: {
+        threadId: string;
+    }): Promise<StorageThreadType | null>;
+    getThreadsByResourceId({ resourceId }: {
+        resourceId: string;
+    }): Promise<StorageThreadType[]>;
+    saveThread({ thread }: {
+        thread: StorageThreadType;
+    }): Promise<StorageThreadType>;
+    updateThread({ id, title, metadata, }: {
+        id: string;
+        title: string;
+        metadata: Record<string, unknown>;
+    }): Promise<StorageThreadType>;
+    deleteThread({ threadId }: {
+        threadId: string;
+    }): Promise<void>;
+    getMessages<T = unknown>({ threadId, selectBy }: StorageGetMessagesArg): Promise<T>;
+    saveMessages({ messages }: {
+        messages: MessageType[];
+    }): Promise<MessageType[]>;
+    persistWorkflowSnapshot({ workflowName, runId, snapshot, }: {
+        workflowName: string;
+        runId: string;
+        snapshot: WorkflowRunState;
+    }): Promise<void>;
+    loadWorkflowSnapshot({ workflowName, runId, }: {
+        workflowName: string;
+        runId: string;
+    }): Promise<WorkflowRunState | null>;
+    close(): Promise<void>;
+}
+export { PostgresStore }
+export { PostgresStore as PostgresStore_alias_1 }
+
+export { }

package/dist/index.d.ts

@@ -0,0 +1,3 @@
+export { PgVector } from './_tsup-dts-rollup.js';
+export { PostgresConfig } from './_tsup-dts-rollup.js';
+export { PostgresStore } from './_tsup-dts-rollup.js';

package/dist/index.js

@@ -0,0 +1,889 @@
+import { MastraVector } from '@mastra/core/vector';
+import pg from 'pg';
+import { BaseFilterTranslator } from '@mastra/core/filter';
+import '@mastra/core/memory';
+import { MastraStorage } from '@mastra/core/storage';
+import '@mastra/core/workflows';
+import pgPromise from 'pg-promise';
+
+// src/vector/index.ts
+var PGFilterTranslator = class extends BaseFilterTranslator {
+  getSupportedOperators() {
+    return {
+      ...BaseFilterTranslator.DEFAULT_OPERATORS,
+      custom: ["$contains", "$size"]
+    };
+  }
+  translate(filter) {
+    if (this.isEmpty(filter)) {
+      return filter;
+    }
+    this.validateFilter(filter);
+    return this.translateNode(filter);
+  }
+  translateNode(node, currentPath = "") {
+    const withPath = (result2) => currentPath ? { [currentPath]: result2 } : result2;
+    if (this.isPrimitive(node)) {
+      return withPath({ $eq: this.normalizeComparisonValue(node) });
+    }
+    if (Array.isArray(node)) {
+      return withPath({ $in: this.normalizeArrayValues(node) });
+    }
+    if (node instanceof RegExp) {
+      return withPath(this.translateRegexPattern(node.source, node.flags));
+    }
+    const entries = Object.entries(node);
+    const result = {};
+    if ("$options" in node && !("$regex" in node)) {
+      throw new Error("$options is not valid without $regex");
+    }
+    if ("$regex" in node) {
+      const options = node.$options || "";
+      return withPath(this.translateRegexPattern(node.$regex, options));
+    }
+    for (const [key, value] of entries) {
+      if (key === "$options") continue;
+      const newPath = currentPath ? `${currentPath}.${key}` : key;
+      if (this.isLogicalOperator(key)) {
+        result[key] = Array.isArray(value) ? value.map((filter) => this.translateNode(filter)) : this.translateNode(value);
+      } else if (this.isOperator(key)) {
+        if (this.isArrayOperator(key) && !Array.isArray(value) && key !== "$elemMatch") {
+          result[key] = [value];
+        } else if (this.isBasicOperator(key) && Array.isArray(value)) {
+          result[key] = JSON.stringify(value);
+        } else {
+          result[key] = value;
+        }
+      } else if (typeof value === "object" && value !== null) {
+        const hasOperators = Object.keys(value).some((k) => this.isOperator(k));
+        if (hasOperators) {
+          result[newPath] = this.translateNode(value);
+        } else {
+          Object.assign(result, this.translateNode(value, newPath));
+        }
+      } else {
+        result[newPath] = this.translateNode(value);
+      }
+    }
+    return result;
+  }
+  translateRegexPattern(pattern, options = "") {
+    if (!options) return { $regex: pattern };
+    const flags = options.split("").filter((f) => "imsux".includes(f)).join("");
+    return { $regex: flags ? `(?${flags})${pattern}` : pattern };
+  }
+};
+var createBasicOperator = (symbol) => {
+  return (key, paramIndex) => ({
+    sql: `CASE 
+      WHEN $${paramIndex}::text IS NULL THEN metadata#>>'{${handleKey(key)}}' IS ${symbol === "=" ? "" : "NOT"} NULL
+      ELSE metadata#>>'{${handleKey(key)}}' ${symbol} $${paramIndex}::text
+    END`,
+    needsValue: true
+  });
+};
+var createNumericOperator = (symbol) => {
+  return (key, paramIndex) => ({
+    sql: `(metadata#>>'{${handleKey(key)}}')::numeric ${symbol} $${paramIndex}`,
+    needsValue: true
+  });
+};
+function buildElemMatchConditions(value, paramIndex) {
+  if (typeof value !== "object" || Array.isArray(value)) {
+    throw new Error("$elemMatch requires an object with conditions");
+  }
+  const conditions = [];
+  const values = [];
+  Object.entries(value).forEach(([field, val]) => {
+    const nextParamIndex = paramIndex + values.length;
+    let paramOperator;
+    let paramKey;
+    let paramValue;
+    if (field.startsWith("$")) {
+      paramOperator = field;
+      paramKey = "";
+      paramValue = val;
+    } else if (typeof val === "object" && !Array.isArray(val)) {
+      const [op, opValue] = Object.entries(val || {})[0] || [];
+      paramOperator = op;
+      paramKey = field;
+      paramValue = opValue;
+    } else {
+      paramOperator = "$eq";
+      paramKey = field;
+      paramValue = val;
+    }
+    const operatorFn = FILTER_OPERATORS[paramOperator];
+    if (!operatorFn) {
+      throw new Error(`Invalid operator: ${paramOperator}`);
+    }
+    const result = operatorFn(paramKey, nextParamIndex, paramValue);
+    const sql = result.sql.replaceAll("metadata#>>", "elem#>>");
+    conditions.push(sql);
+    if (result.needsValue) {
+      values.push(paramValue);
+    }
+  });
+  return {
+    sql: conditions.join(" AND "),
+    values
+  };
+}
+var FILTER_OPERATORS = {
+  $eq: createBasicOperator("="),
+  $ne: createBasicOperator("!="),
+  $gt: createNumericOperator(">"),
+  $gte: createNumericOperator(">="),
+  $lt: createNumericOperator("<"),
+  $lte: createNumericOperator("<="),
+  // Array Operators
+  $in: (key, paramIndex) => ({
+    sql: `metadata#>>'{${handleKey(key)}}' = ANY($${paramIndex}::text[])`,
+    needsValue: true
+  }),
+  $nin: (key, paramIndex) => ({
+    sql: `metadata#>>'{${handleKey(key)}}' != ALL($${paramIndex}::text[])`,
+    needsValue: true
+  }),
+  $all: (key, paramIndex) => ({
+    sql: `CASE WHEN array_length($${paramIndex}::text[], 1) IS NULL THEN false 
+          ELSE (metadata#>'{${handleKey(key)}}')::jsonb ?& $${paramIndex}::text[] END`,
+    needsValue: true
+  }),
+  $elemMatch: (key, paramIndex, value) => {
+    const { sql, values } = buildElemMatchConditions(value, paramIndex);
+    return {
+      sql: `(
+        CASE
+          WHEN jsonb_typeof(metadata->'${handleKey(key)}') = 'array' THEN
+            EXISTS (
+              SELECT 1 
+              FROM jsonb_array_elements(metadata->'${handleKey(key)}') as elem
+              WHERE ${sql}
+            )
+          ELSE FALSE
+        END
+      )`,
+      needsValue: true,
+      transformValue: () => values
+    };
+  },
+  // Element Operators
+  $exists: (key) => ({
+    sql: `metadata ? '${key}'`,
+    needsValue: false
+  }),
+  // Logical Operators
+  $and: (key) => ({ sql: `(${key})`, needsValue: false }),
+  $or: (key) => ({ sql: `(${key})`, needsValue: false }),
+  $not: (key) => ({ sql: `NOT (${key})`, needsValue: false }),
+  $nor: (key) => ({ sql: `NOT (${key})`, needsValue: false }),
+  // Regex Operators
+  $regex: (key, paramIndex) => ({
+    sql: `metadata#>>'{${handleKey(key)}}' ~ $${paramIndex}`,
+    needsValue: true
+  }),
+  $contains: (key, paramIndex) => ({
+    sql: `metadata @> $${paramIndex}::jsonb`,
+    needsValue: true,
+    transformValue: (value) => {
+      const parts = key.split(".");
+      return JSON.stringify(parts.reduceRight((value2, key2) => ({ [key2]: value2 }), value));
+    }
+  }),
+  $size: (key, paramIndex) => ({
+    sql: `(
+      CASE
+        WHEN jsonb_typeof(metadata#>'{${handleKey(key)}}') = 'array' THEN 
+          jsonb_array_length(metadata#>'{${handleKey(key)}}') = $${paramIndex}
+        ELSE FALSE
+      END
+    )`,
+    needsValue: true
+  })
+};
+var handleKey = (key) => {
+  return key.replace(/\./g, ",");
+};
+function buildFilterQuery(filter, minScore) {
+  const values = [minScore];
+  function buildCondition(key, value, parentPath) {
+    if (["$and", "$or", "$not", "$nor"].includes(key)) {
+      return handleLogicalOperator(key, value);
+    }
+    if (!value || typeof value !== "object") {
+      values.push(value);
+      return `metadata#>>'{${handleKey(key)}}' = $${values.length}`;
+    }
+    const [[operator, operatorValue] = []] = Object.entries(value);
+    if (operator === "$not") {
+      const entries = Object.entries(operatorValue);
+      const conditions2 = entries.map(([nestedOp, nestedValue]) => {
+        if (!FILTER_OPERATORS[nestedOp]) {
+          throw new Error(`Invalid operator in $not condition: ${nestedOp}`);
+        }
+        const operatorFn2 = FILTER_OPERATORS[nestedOp];
+        const operatorResult2 = operatorFn2(key, values.length + 1);
+        if (operatorResult2.needsValue) {
+          values.push(nestedValue);
+        }
+        return operatorResult2.sql;
+      }).join(" AND ");
+      return `NOT (${conditions2})`;
+    }
+    const operatorFn = FILTER_OPERATORS[operator];
+    const operatorResult = operatorFn(key, values.length + 1, operatorValue);
+    if (operatorResult.needsValue) {
+      const transformedValue = operatorResult.transformValue ? operatorResult.transformValue(operatorValue) : operatorValue;
+      if (Array.isArray(transformedValue) && operator === "$elemMatch") {
+        values.push(...transformedValue);
+      } else {
+        values.push(transformedValue);
+      }
+    }
+    return operatorResult.sql;
+  }
+  function handleLogicalOperator(key, value, parentPath) {
+    if (key === "$not") {
+      const entries = Object.entries(value);
+      const conditions3 = entries.map(([fieldKey, fieldValue]) => buildCondition(fieldKey, fieldValue)).join(" AND ");
+      return `NOT (${conditions3})`;
+    }
+    if (!value || value.length === 0) {
+      switch (key) {
+        case "$and":
+        case "$nor":
+          return "true";
+        // Empty $and/$nor match everything
+        case "$or":
+          return "false";
+        // Empty $or matches nothing
+        default:
+          return "true";
+      }
+    }
+    const joinOperator = key === "$or" || key === "$nor" ? "OR" : "AND";
+    const conditions2 = value.map((f) => {
+      const entries = Object.entries(f);
+      if (entries.length === 0) return "";
+      const [firstKey, firstValue] = entries[0] || [];
+      if (["$and", "$or", "$not", "$nor"].includes(firstKey)) {
+        return buildCondition(firstKey, firstValue);
+      }
+      return entries.map(([k, v]) => buildCondition(k, v)).join(` ${joinOperator} `);
+    });
+    const joined = conditions2.join(` ${joinOperator} `);
+    const operatorFn = FILTER_OPERATORS[key];
+    return operatorFn(joined, 0, value).sql;
+  }
+  if (!filter) {
+    return { sql: "", values };
+  }
+  const conditions = Object.entries(filter).map(([key, value]) => buildCondition(key, value)).filter(Boolean).join(" AND ");
+  return { sql: conditions ? `WHERE ${conditions}` : "", values };
+}
+
+// src/vector/index.ts
+var PgVector = class extends MastraVector {
+  pool;
+  constructor(connectionString) {
+    super();
+    const basePool = new pg.Pool({
+      connectionString,
+      max: 20,
+      // Maximum number of clients in the pool
+      idleTimeoutMillis: 3e4,
+      // Close idle connections after 30 seconds
+      connectionTimeoutMillis: 2e3
+      // Fail fast if can't connect
+    });
+    const telemetry = this.__getTelemetry();
+    this.pool = telemetry?.traceClass(basePool, {
+      spanNamePrefix: "pg-vector",
+      attributes: {
+        "vector.type": "postgres"
+      }
+    }) ?? basePool;
+  }
+  transformFilter(filter) {
+    const pgFilter = new PGFilterTranslator();
+    const translatedFilter = pgFilter.translate(filter ?? {});
+    return translatedFilter;
+  }
+  async query(indexName, queryVector, topK = 10, filter, includeVector = false, minScore = 0) {
+    const client = await this.pool.connect();
+    try {
+      const vectorStr = `[${queryVector.join(",")}]`;
+      const translatedFilter = this.transformFilter(filter);
+      const { sql: filterQuery, values: filterValues } = buildFilterQuery(translatedFilter, minScore);
+      const query = `
+        WITH vector_scores AS (
+          SELECT
+            vector_id as id,
+            1 - (embedding <=> '${vectorStr}'::vector) as score,
+            metadata
+            ${includeVector ? ", embedding" : ""}
+          FROM ${indexName}
+          ${filterQuery}
+        )
+        SELECT *
+        FROM vector_scores
+        WHERE score > $1
+        ORDER BY score DESC
+        LIMIT ${topK}`;
+      const result = await client.query(query, filterValues);
+      return result.rows.map(({ id, score, metadata, embedding }) => ({
+        id,
+        score,
+        metadata,
+        ...includeVector && embedding && { vector: JSON.parse(embedding) }
+      }));
+    } finally {
+      client.release();
+    }
+  }
+  async upsert(indexName, vectors, metadata, ids) {
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+      const vectorIds = ids || vectors.map(() => crypto.randomUUID());
+      for (let i = 0; i < vectors.length; i++) {
+        const query = `
+            INSERT INTO ${indexName} (vector_id, embedding, metadata)
+            VALUES ($1, $2::vector, $3::jsonb)
+            ON CONFLICT (vector_id)
+            DO UPDATE SET
+                embedding = $2::vector,
+                metadata = $3::jsonb
+            RETURNING embedding::text
+        `;
+        await client.query(query, [vectorIds[i], `[${vectors[i]?.join(",")}]`, JSON.stringify(metadata?.[i] || {})]);
+      }
+      await client.query("COMMIT");
+      return vectorIds;
+    } catch (error) {
+      await client.query("ROLLBACK");
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+  async createIndex(indexName, dimension, metric = "cosine") {
+    const client = await this.pool.connect();
+    try {
+      if (!indexName.match(/^[a-zA-Z_][a-zA-Z0-9_]*$/)) {
+        throw new Error("Invalid index name format");
+      }
+      if (!Number.isInteger(dimension) || dimension <= 0) {
+        throw new Error("Dimension must be a positive integer");
+      }
+      const extensionCheck = await client.query(`
+        SELECT EXISTS (
+          SELECT 1 FROM pg_available_extensions WHERE name = 'vector'
+        );
+      `);
+      if (!extensionCheck.rows[0].exists) {
+        throw new Error("PostgreSQL vector extension is not available. Please install it first.");
+      }
+      await client.query("CREATE EXTENSION IF NOT EXISTS vector");
+      await client.query(`
+        CREATE TABLE IF NOT EXISTS ${indexName} (
+          id SERIAL PRIMARY KEY,
+          vector_id TEXT UNIQUE NOT NULL,
+          embedding vector(${dimension}),
+          metadata JSONB DEFAULT '{}'::jsonb
+        );
+      `);
+      const indexMethod = metric === "cosine" ? "vector_cosine_ops" : metric === "euclidean" ? "vector_l2_ops" : "vector_ip_ops";
+      await client.query(`
+        CREATE INDEX IF NOT EXISTS ${indexName}_vector_idx
+        ON public.${indexName}
+        USING ivfflat (embedding ${indexMethod})
+        WITH (lists = 100);
+      `);
+    } catch (error) {
+      console.error("Failed to create vector table:", error);
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+  async listIndexes() {
+    const client = await this.pool.connect();
+    try {
+      const vectorTablesQuery = `
+            SELECT DISTINCT table_name
+            FROM information_schema.columns
+            WHERE table_schema = 'public'
+            AND udt_name = 'vector';
+        `;
+      const vectorTables = await client.query(vectorTablesQuery);
+      return vectorTables.rows.map((row) => row.table_name);
+    } finally {
+      client.release();
+    }
+  }
+  async describeIndex(indexName) {
+    const client = await this.pool.connect();
+    try {
+      const dimensionQuery = `
+                SELECT atttypmod as dimension
+                FROM pg_attribute
+                WHERE attrelid = $1::regclass
+                AND attname = 'embedding';
+            `;
+      const countQuery = `
+                SELECT COUNT(*) as count
+                FROM ${indexName};
+            `;
+      const metricQuery = `
+            SELECT
+                am.amname as index_method,
+                opclass.opcname as operator_class
+            FROM pg_index i
+            JOIN pg_class c ON i.indexrelid = c.oid
+            JOIN pg_am am ON c.relam = am.oid
+            JOIN pg_opclass opclass ON i.indclass[0] = opclass.oid
+            WHERE c.relname = '${indexName}_vector_idx';
+            `;
+      const [dimResult, countResult, metricResult] = await Promise.all([
+        client.query(dimensionQuery, [indexName]),
+        client.query(countQuery),
+        client.query(metricQuery)
+      ]);
+      let metric = "cosine";
+      if (metricResult.rows.length > 0) {
+        const operatorClass = metricResult.rows[0].operator_class;
+        if (operatorClass.includes("l2")) {
+          metric = "euclidean";
+        } else if (operatorClass.includes("ip")) {
+          metric = "dotproduct";
+        } else if (operatorClass.includes("cosine")) {
+          metric = "cosine";
+        }
+      }
+      return {
+        dimension: dimResult.rows[0].dimension,
+        count: parseInt(countResult.rows[0].count),
+        metric
+      };
+    } catch (e) {
+      await client.query("ROLLBACK");
+      throw new Error(`Failed to describe vector table: ${e.message}`);
+    } finally {
+      client.release();
+    }
+  }
+  async deleteIndex(indexName) {
+    const client = await this.pool.connect();
+    try {
+      await client.query(`DROP TABLE IF EXISTS ${indexName} CASCADE`);
+    } catch (error) {
+      await client.query("ROLLBACK");
+      throw new Error(`Failed to delete vector table: ${error.message}`);
+    } finally {
+      client.release();
+    }
+  }
+  async truncateIndex(indexName) {
+    const client = await this.pool.connect();
+    try {
+      await client.query(`TRUNCATE ${indexName}`);
+    } catch (e) {
+      await client.query("ROLLBACK");
+      throw new Error(`Failed to truncate vector table: ${e.message}`);
+    } finally {
+      client.release();
+    }
+  }
+  async disconnect() {
+    await this.pool.end();
+  }
+};
+var PostgresStore = class extends MastraStorage {
+  db;
+  pgp;
+  constructor(config) {
+    super({ name: "PostgresStore" });
+    this.pgp = pgPromise();
+    this.db = this.pgp(
+      `connectionString` in config ? { connectionString: config.connectionString } : {
+        host: config.host,
+        port: config.port,
+        database: config.database,
+        user: config.user,
+        password: config.password
+      }
+    );
+  }
+  async createTable({
+    tableName,
+    schema
+  }) {
+    try {
+      const columns = Object.entries(schema).map(([name, def]) => {
+        const constraints = [];
+        if (def.primaryKey) constraints.push("PRIMARY KEY");
+        if (!def.nullable) constraints.push("NOT NULL");
+        return `"${name}" ${def.type.toUpperCase()} ${constraints.join(" ")}`;
+      }).join(",\n");
+      const sql = `
+        CREATE TABLE IF NOT EXISTS ${tableName} (
+          ${columns}
+        );
+        ${tableName === MastraStorage.TABLE_WORKFLOW_SNAPSHOT ? `
+        DO $$ BEGIN
+          IF NOT EXISTS (
+            SELECT 1 FROM pg_constraint WHERE conname = 'mastra_workflow_snapshot_workflow_name_run_id_key'
+          ) THEN
+            ALTER TABLE ${tableName}
+            ADD CONSTRAINT mastra_workflow_snapshot_workflow_name_run_id_key
+            UNIQUE (workflow_name, run_id);
+          END IF;
+        END $$;
+        ` : ""}
+      `;
+      await this.db.none(sql);
+    } catch (error) {
+      console.error(`Error creating table ${tableName}:`, error);
+      throw error;
+    }
+  }
+  async clearTable({ tableName }) {
+    try {
+      await this.db.none(`TRUNCATE TABLE ${tableName} CASCADE`);
+    } catch (error) {
+      console.error(`Error clearing table ${tableName}:`, error);
+      throw error;
+    }
+  }
+  async insert({ tableName, record }) {
+    try {
+      const columns = Object.keys(record);
+      const values = Object.values(record);
+      const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
+      await this.db.none(
+        `INSERT INTO ${tableName} (${columns.map((c) => `"${c}"`).join(", ")}) VALUES (${placeholders})`,
+        values
+      );
+    } catch (error) {
+      console.error(`Error inserting into ${tableName}:`, error);
+      throw error;
+    }
+  }
+  async load({ tableName, keys }) {
+    try {
+      const keyEntries = Object.entries(keys);
+      const conditions = keyEntries.map(([key], index) => `"${key}" = $${index + 1}`).join(" AND ");
+      const values = keyEntries.map(([_, value]) => value);
+      const result = await this.db.oneOrNone(`SELECT * FROM ${tableName} WHERE ${conditions}`, values);
+      if (!result) {
+        return null;
+      }
+      if (tableName === MastraStorage.TABLE_WORKFLOW_SNAPSHOT) {
+        const snapshot = result;
+        if (typeof snapshot.snapshot === "string") {
+          snapshot.snapshot = JSON.parse(snapshot.snapshot);
+        }
+        return snapshot;
+      }
+      return result;
+    } catch (error) {
+      console.error(`Error loading from ${tableName}:`, error);
+      throw error;
+    }
+  }
+  async getThreadById({ threadId }) {
+    try {
+      const thread = await this.db.oneOrNone(
+        `SELECT 
+          id,
+          "resourceId",
+          title,
+          metadata,
+          "createdAt",
+          "updatedAt"
+        FROM "${MastraStorage.TABLE_THREADS}"
+        WHERE id = $1`,
+        [threadId]
+      );
+      if (!thread) {
+        return null;
+      }
+      return {
+        ...thread,
+        metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
+        createdAt: thread.createdAt,
+        updatedAt: thread.updatedAt
+      };
+    } catch (error) {
+      console.error(`Error getting thread ${threadId}:`, error);
+      throw error;
+    }
+  }
+  async getThreadsByResourceId({ resourceId }) {
+    try {
+      const threads = await this.db.manyOrNone(
+        `SELECT 
+          id,
+          "resourceId",
+          title,
+          metadata,
+          "createdAt",
+          "updatedAt"
+        FROM "${MastraStorage.TABLE_THREADS}"
+        WHERE "resourceId" = $1`,
+        [resourceId]
+      );
+      return threads.map((thread) => ({
+        ...thread,
+        metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
+        createdAt: thread.createdAt,
+        updatedAt: thread.updatedAt
+      }));
+    } catch (error) {
+      console.error(`Error getting threads for resource ${resourceId}:`, error);
+      throw error;
+    }
+  }
+  async saveThread({ thread }) {
+    try {
+      await this.db.none(
+        `INSERT INTO "${MastraStorage.TABLE_THREADS}" (
+          id,
+          "resourceId",
+          title,
+          metadata,
+          "createdAt",
+          "updatedAt"
+        ) VALUES ($1, $2, $3, $4, $5, $6)
+        ON CONFLICT (id) DO UPDATE SET
+          "resourceId" = EXCLUDED."resourceId",
+          title = EXCLUDED.title,
+          metadata = EXCLUDED.metadata,
+          "createdAt" = EXCLUDED."createdAt",
+          "updatedAt" = EXCLUDED."updatedAt"`,
+        [
+          thread.id,
+          thread.resourceId,
+          thread.title,
+          thread.metadata ? JSON.stringify(thread.metadata) : null,
+          thread.createdAt,
+          thread.updatedAt
+        ]
+      );
+      return thread;
+    } catch (error) {
+      console.error("Error saving thread:", error);
+      throw error;
+    }
+  }
+  async updateThread({
+    id,
+    title,
+    metadata
+  }) {
+    try {
+      const existingThread = await this.getThreadById({ threadId: id });
+      if (!existingThread) {
+        throw new Error(`Thread ${id} not found`);
+      }
+      const mergedMetadata = {
+        ...existingThread.metadata,
+        ...metadata
+      };
+      const thread = await this.db.one(
+        `UPDATE "${MastraStorage.TABLE_THREADS}"
+        SET title = $1,
+            metadata = $2,
+            "updatedAt" = $3
+        WHERE id = $4
+        RETURNING *`,
+        [title, mergedMetadata, (/* @__PURE__ */ new Date()).toISOString(), id]
+      );
+      return {
+        ...thread,
+        metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
+        createdAt: thread.createdAt,
+        updatedAt: thread.updatedAt
+      };
+    } catch (error) {
+      console.error("Error updating thread:", error);
+      throw error;
+    }
+  }
+  async deleteThread({ threadId }) {
+    try {
+      await this.db.tx(async (t) => {
+        await t.none(`DELETE FROM "${MastraStorage.TABLE_MESSAGES}" WHERE thread_id = $1`, [threadId]);
+        await t.none(`DELETE FROM "${MastraStorage.TABLE_THREADS}" WHERE id = $1`, [threadId]);
+      });
+    } catch (error) {
+      console.error("Error deleting thread:", error);
+      throw error;
+    }
+  }
+  async getMessages({ threadId, selectBy }) {
+    try {
+      const messages = [];
+      const limit = typeof selectBy?.last === `number` ? selectBy.last : 40;
+      const include = selectBy?.include || [];
+      if (include.length) {
+        const includeResult = await this.db.manyOrNone(
+          `
+          WITH ordered_messages AS (
+            SELECT 
+              *,
+              ROW_NUMBER() OVER (ORDER BY "createdAt") as row_num
+            FROM "${MastraStorage.TABLE_MESSAGES}"
+            WHERE thread_id = $1
+          )
+          SELECT DISTINCT ON (m.id)
+            m.id, 
+            m.content, 
+            m.role, 
+            m.type,
+            m."createdAt", 
+            m.thread_id AS "threadId"
+          FROM ordered_messages m
+          WHERE m.id = ANY($2)
+          OR EXISTS (
+            SELECT 1 FROM ordered_messages target
+            WHERE target.id = ANY($2)
+            AND (
+              -- Get previous messages based on the max withPreviousMessages
+              (m.row_num >= target.row_num - $3 AND m.row_num < target.row_num)
+              OR
+              -- Get next messages based on the max withNextMessages
+              (m.row_num <= target.row_num + $4 AND m.row_num > target.row_num)
+            )
+          )
+          ORDER BY m.id, m."createdAt"
+          `,
+          [
+            threadId,
+            include.map((i) => i.id),
+            Math.max(...include.map((i) => i.withPreviousMessages || 0)),
+            Math.max(...include.map((i) => i.withNextMessages || 0))
+          ]
+        );
+        messages.push(...includeResult);
+      }
+      const result = await this.db.manyOrNone(
+        `
+        SELECT 
+            id, 
+            content, 
+            role, 
+            type,
+            "createdAt", 
+            thread_id AS "threadId"
+        FROM "${MastraStorage.TABLE_MESSAGES}"
+        WHERE thread_id = $1
+        AND id != ALL($2)
+        ORDER BY "createdAt" DESC
+        LIMIT $3
+        `,
+        [threadId, messages.map((m) => m.id), limit]
+      );
+      messages.push(...result);
+      messages.sort((a, b) => new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime());
+      messages.forEach((message) => {
+        if (typeof message.content === "string") {
+          try {
+            message.content = JSON.parse(message.content);
+          } catch (e) {
+          }
+        }
+      });
+      return messages;
+    } catch (error) {
+      console.error("Error getting messages:", error);
+      throw error;
+    }
+  }
+  async saveMessages({ messages }) {
+    if (messages.length === 0) return messages;
+    try {
+      const threadId = messages[0]?.threadId;
+      if (!threadId) {
+        throw new Error("Thread ID is required");
+      }
+      const thread = await this.getThreadById({ threadId });
+      if (!thread) {
+        throw new Error(`Thread ${threadId} not found`);
+      }
+      await this.db.tx(async (t) => {
+        for (const message of messages) {
+          await t.none(
+            `INSERT INTO "${MastraStorage.TABLE_MESSAGES}" (id, thread_id, content, "createdAt", role, type) 
+             VALUES ($1, $2, $3, $4, $5, $6)`,
+            [
+              message.id,
+              threadId,
+              typeof message.content === "string" ? message.content : JSON.stringify(message.content),
+              message.createdAt || (/* @__PURE__ */ new Date()).toISOString(),
+              message.role,
+              message.type
+            ]
+          );
+        }
+      });
+      return messages;
+    } catch (error) {
+      console.error("Error saving messages:", error);
+      throw error;
+    }
+  }
+  async persistWorkflowSnapshot({
+    workflowName,
+    runId,
+    snapshot
+  }) {
+    try {
+      const now = (/* @__PURE__ */ new Date()).toISOString();
+      await this.db.none(
+        `INSERT INTO "${MastraStorage.TABLE_WORKFLOW_SNAPSHOT}" (
+          workflow_name,
+          run_id,
+          snapshot,
+          "createdAt",
+          "updatedAt"
+        ) VALUES ($1, $2, $3, $4, $5)
+        ON CONFLICT (workflow_name, run_id) DO UPDATE
+        SET snapshot = EXCLUDED.snapshot,
+            "updatedAt" = EXCLUDED."updatedAt"`,
+        [workflowName, runId, JSON.stringify(snapshot), now, now]
+      );
+    } catch (error) {
+      console.error("Error persisting workflow snapshot:", error);
+      throw error;
+    }
+  }
+  async loadWorkflowSnapshot({
+    workflowName,
+    runId
+  }) {
+    try {
+      const result = await this.load({
+        tableName: MastraStorage.TABLE_WORKFLOW_SNAPSHOT,
+        keys: {
+          workflow_name: workflowName,
+          run_id: runId
+        }
+      });
+      if (!result) {
+        return null;
+      }
+      return result.snapshot;
+    } catch (error) {
+      console.error("Error loading workflow snapshot:", error);
+      throw error;
+    }
+  }
+  async close() {
+    this.pgp.end();
+  }
+};
+
+export { PgVector, PostgresStore };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/pg",
-  "version": "0.1.0-alpha.11",
+  "version": "0.1.0-alpha.12",
   "description": "Postgres provider for Mastra - includes both vector and db storage capabilities",
   "type": "module",
   "main": "dist/index.js",
@@ -17,7 +17,7 @@
   "dependencies": {
     "pg": "^8.13.1",
     "pg-promise": "^11.5.4",
-    "@mastra/core": "^0.2.0-alpha.101"
+    "@mastra/core": "^0.2.0-alpha.102"
   },
   "devDependencies": {
     "@microsoft/api-extractor": "^7.49.2",