@mastra/pg 0.12.3 → 0.12.4-alpha.0

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

@@ -0,0 +1,900 @@
+import { MessageList } from '@mastra/core/agent';
+import type { MastraMessageContentV2 } from '@mastra/core/agent';
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { MastraMessageV1, MastraMessageV2, StorageThreadType } from '@mastra/core/memory';
+import {
+  MemoryStorage,
+  resolveMessageLimit,
+  TABLE_MESSAGES,
+  TABLE_RESOURCES,
+  TABLE_THREADS,
+} from '@mastra/core/storage';
+import type { StorageGetMessagesArg, PaginationInfo, StorageResourceType } from '@mastra/core/storage';
+import type { IDatabase } from 'pg-promise';
+import type { StoreOperationsPG } from '../operations';
+import { getTableName, getSchemaName } from '../utils';
+
+export class MemoryPG extends MemoryStorage {
+  private client: IDatabase<{}>;
+  private schema: string;
+  private operations: StoreOperationsPG;
+
+  constructor({
+    client,
+    schema,
+    operations,
+  }: {
+    client: IDatabase<{}>;
+    schema: string;
+    operations: StoreOperationsPG;
+  }) {
+    super();
+    this.client = client;
+    this.schema = schema;
+    this.operations = operations;
+  }
+
+  async getThreadById({ threadId }: { threadId: string }): Promise<StorageThreadType | null> {
+    try {
+      const tableName = getTableName({ indexName: TABLE_THREADS, schemaName: getSchemaName(this.schema) });
+
+      const thread = await this.client.oneOrNone<StorageThreadType & { createdAtZ: Date; updatedAtZ: Date }>(
+        `SELECT * FROM ${tableName} WHERE id = $1`,
+        [threadId],
+      );
+
+      if (!thread) {
+        return null;
+      }
+
+      return {
+        id: thread.id,
+        resourceId: thread.resourceId,
+        title: thread.title,
+        metadata: typeof thread.metadata === 'string' ? JSON.parse(thread.metadata) : thread.metadata,
+        createdAt: thread.createdAtZ || thread.createdAt,
+        updatedAt: thread.updatedAtZ || thread.updatedAt,
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_GET_THREAD_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            threadId,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  /**
+   * @deprecated use getThreadsByResourceIdPaginated instead
+   */
+  public async getThreadsByResourceId(args: { resourceId: string }): Promise<StorageThreadType[]> {
+    const { resourceId } = args;
+
+    try {
+      const tableName = getTableName({ indexName: TABLE_THREADS, schemaName: getSchemaName(this.schema) });
+      const baseQuery = `FROM ${tableName} WHERE "resourceId" = $1`;
+      const queryParams: any[] = [resourceId];
+
+      const dataQuery = `SELECT id, "resourceId", title, metadata, "createdAt", "updatedAt" ${baseQuery} ORDER BY "createdAt" DESC`;
+      const rows = await this.client.manyOrNone(dataQuery, queryParams);
+      return (rows || []).map(thread => ({
+        ...thread,
+        metadata: typeof thread.metadata === 'string' ? JSON.parse(thread.metadata) : thread.metadata,
+        createdAt: thread.createdAt,
+        updatedAt: thread.updatedAt,
+      }));
+    } catch (error) {
+      this.logger.error(`Error getting threads for resource ${resourceId}:`, error);
+      return [];
+    }
+  }
+
+  public async getThreadsByResourceIdPaginated(args: {
+    resourceId: string;
+    page: number;
+    perPage: number;
+  }): Promise<PaginationInfo & { threads: StorageThreadType[] }> {
+    const { resourceId, page = 0, perPage: perPageInput } = args;
+    try {
+      const tableName = getTableName({ indexName: TABLE_THREADS, schemaName: getSchemaName(this.schema) });
+      const baseQuery = `FROM ${tableName} WHERE "resourceId" = $1`;
+      const queryParams: any[] = [resourceId];
+      const perPage = perPageInput !== undefined ? perPageInput : 100;
+      const currentOffset = page * perPage;
+
+      const countQuery = `SELECT COUNT(*) ${baseQuery}`;
+      const countResult = await this.client.one(countQuery, queryParams);
+      const total = parseInt(countResult.count, 10);
+
+      if (total === 0) {
+        return {
+          threads: [],
+          total: 0,
+          page,
+          perPage,
+          hasMore: false,
+        };
+      }
+
+      const dataQuery = `SELECT id, "resourceId", title, metadata, "createdAt", "updatedAt" ${baseQuery} ORDER BY "createdAt" DESC LIMIT $2 OFFSET $3`;
+      const rows = await this.client.manyOrNone(dataQuery, [...queryParams, perPage, currentOffset]);
+
+      const threads = (rows || []).map(thread => ({
+        ...thread,
+        metadata: typeof thread.metadata === 'string' ? JSON.parse(thread.metadata) : thread.metadata,
+        createdAt: thread.createdAt, // Assuming already Date objects or ISO strings
+        updatedAt: thread.updatedAt,
+      }));
+
+      return {
+        threads,
+        total,
+        page,
+        perPage,
+        hasMore: currentOffset + threads.length < total,
+      };
+    } catch (error) {
+      const mastraError = new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_GET_THREADS_BY_RESOURCE_ID_PAGINATED_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            resourceId,
+            page,
+          },
+        },
+        error,
+      );
+      this.logger?.error?.(mastraError.toString());
+      this.logger?.trackException(mastraError);
+      return { threads: [], total: 0, page, perPage: perPageInput || 100, hasMore: false };
+    }
+  }
+
+  async saveThread({ thread }: { thread: StorageThreadType }): Promise<StorageThreadType> {
+    try {
+      const tableName = getTableName({ indexName: TABLE_THREADS, schemaName: getSchemaName(this.schema) });
+      await this.client.none(
+        `INSERT INTO ${tableName} (
+          id,
+          "resourceId",
+          title,
+          metadata,
+          "createdAt",
+          "createdAtZ",
+          "updatedAt",
+          "updatedAtZ"
+        ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+        ON CONFLICT (id) DO UPDATE SET
+          "resourceId" = EXCLUDED."resourceId",
+          title = EXCLUDED.title,
+          metadata = EXCLUDED.metadata,
+          "createdAt" = EXCLUDED."createdAt",
+          "createdAtZ" = EXCLUDED."createdAtZ",
+          "updatedAt" = EXCLUDED."updatedAt",
+          "updatedAtZ" = EXCLUDED."updatedAtZ"`,
+        [
+          thread.id,
+          thread.resourceId,
+          thread.title,
+          thread.metadata ? JSON.stringify(thread.metadata) : null,
+          thread.createdAt,
+          thread.createdAt,
+          thread.updatedAt,
+          thread.updatedAt,
+        ],
+      );
+
+      return thread;
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_SAVE_THREAD_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            threadId: thread.id,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async updateThread({
+    id,
+    title,
+    metadata,
+  }: {
+    id: string;
+    title: string;
+    metadata: Record<string, unknown>;
+  }): Promise<StorageThreadType> {
+    const threadTableName = getTableName({ indexName: TABLE_THREADS, schemaName: getSchemaName(this.schema) });
+    // First get the existing thread to merge metadata
+    const existingThread = await this.getThreadById({ threadId: id });
+    if (!existingThread) {
+      throw new MastraError({
+        id: 'MASTRA_STORAGE_PG_STORE_UPDATE_THREAD_FAILED',
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.USER,
+        text: `Thread ${id} not found`,
+        details: {
+          threadId: id,
+          title,
+        },
+      });
+    }
+
+    // Merge the existing metadata with the new metadata
+    const mergedMetadata = {
+      ...existingThread.metadata,
+      ...metadata,
+    };
+
+    try {
+      const thread = await this.client.one<StorageThreadType & { createdAtZ: Date; updatedAtZ: Date }>(
+        `UPDATE ${threadTableName}
+                    SET 
+                        title = $1,
+                        metadata = $2,
+                        "updatedAt" = $3,
+                        "updatedAtZ" = $3
+                    WHERE id = $4
+                    RETURNING *
+                `,
+        [title, mergedMetadata, new Date().toISOString(), id],
+      );
+
+      return {
+        id: thread.id,
+        resourceId: thread.resourceId,
+        title: thread.title,
+        metadata: typeof thread.metadata === 'string' ? JSON.parse(thread.metadata) : thread.metadata,
+        createdAt: thread.createdAtZ || thread.createdAt,
+        updatedAt: thread.updatedAtZ || thread.updatedAt,
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_UPDATE_THREAD_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            threadId: id,
+            title,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async deleteThread({ threadId }: { threadId: string }): Promise<void> {
+    try {
+      const tableName = getTableName({ indexName: TABLE_MESSAGES, schemaName: getSchemaName(this.schema) });
+      const threadTableName = getTableName({ indexName: TABLE_THREADS, schemaName: getSchemaName(this.schema) });
+      await this.client.tx(async t => {
+        // First delete all messages associated with this thread
+        await t.none(`DELETE FROM ${tableName} WHERE thread_id = $1`, [threadId]);
+
+        // Then delete the thread
+        await t.none(`DELETE FROM ${threadTableName} WHERE id = $1`, [threadId]);
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_DELETE_THREAD_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            threadId,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  private async _getIncludedMessages({
+    threadId,
+    selectBy,
+    orderByStatement,
+  }: {
+    threadId: string;
+    selectBy: StorageGetMessagesArg['selectBy'];
+    orderByStatement: string;
+  }) {
+    const include = selectBy?.include;
+    if (!include) return null;
+
+    const unionQueries: string[] = [];
+    const params: any[] = [];
+    let paramIdx = 1;
+    const tableName = getTableName({ indexName: TABLE_MESSAGES, schemaName: getSchemaName(this.schema) });
+
+    for (const inc of include) {
+      const { id, withPreviousMessages = 0, withNextMessages = 0 } = inc;
+      // if threadId is provided, use it, otherwise use threadId from args
+      const searchId = inc.threadId || threadId;
+      unionQueries.push(
+        `
+            SELECT * FROM (
+              WITH ordered_messages AS (
+                SELECT 
+                  *,
+                  ROW_NUMBER() OVER (${orderByStatement}) as row_num
+                FROM ${tableName}
+                WHERE thread_id = $${paramIdx}
+              )
+              SELECT
+                m.id, 
+                m.content, 
+                m.role, 
+                m.type,
+                m."createdAt", 
+                m.thread_id AS "threadId",
+                m."resourceId"
+              FROM ordered_messages m
+              WHERE m.id = $${paramIdx + 1}
+              OR EXISTS (
+                SELECT 1 FROM ordered_messages target
+                WHERE target.id = $${paramIdx + 1}
+                AND (
+                  -- Get previous messages based on the max withPreviousMessages
+                  (m.row_num <= target.row_num + $${paramIdx + 2} AND m.row_num > target.row_num)
+                  OR
+                  -- Get next messages based on the max withNextMessages
+                  (m.row_num >= target.row_num - $${paramIdx + 3} AND m.row_num < target.row_num)
+                )
+              )
+            ) AS query_${paramIdx}
+            `, // Keep ASC for final sorting after fetching context
+      );
+      params.push(searchId, id, withPreviousMessages, withNextMessages);
+      paramIdx += 4;
+    }
+    const finalQuery = unionQueries.join(' UNION ALL ') + ' ORDER BY "createdAt" ASC';
+    const includedRows = await this.client.manyOrNone(finalQuery, params);
+    const seen = new Set<string>();
+    const dedupedRows = includedRows.filter(row => {
+      if (seen.has(row.id)) return false;
+      seen.add(row.id);
+      return true;
+    });
+    return dedupedRows;
+  }
+
+  /**
+   * @deprecated use getMessagesPaginated instead
+   */
+  public async getMessages(args: StorageGetMessagesArg & { format?: 'v1' }): Promise<MastraMessageV1[]>;
+  public async getMessages(args: StorageGetMessagesArg & { format: 'v2' }): Promise<MastraMessageV2[]>;
+  public async getMessages(
+    args: StorageGetMessagesArg & {
+      format?: 'v1' | 'v2';
+    },
+  ): Promise<MastraMessageV1[] | MastraMessageV2[]> {
+    const { threadId, format, selectBy } = args;
+
+    const selectStatement = `SELECT id, content, role, type, "createdAt", thread_id AS "threadId", "resourceId"`;
+    const orderByStatement = `ORDER BY "createdAt" DESC`;
+    const limit = resolveMessageLimit({ last: selectBy?.last, defaultLimit: 40 });
+
+    try {
+      let rows: any[] = [];
+      const include = selectBy?.include || [];
+
+      if (include?.length) {
+        const includeMessages = await this._getIncludedMessages({ threadId, selectBy, orderByStatement });
+        if (includeMessages) {
+          rows.push(...includeMessages);
+        }
+      }
+
+      const excludeIds = rows.map(m => m.id);
+      const tableName = getTableName({ indexName: TABLE_MESSAGES, schemaName: getSchemaName(this.schema) });
+      const excludeIdsParam = excludeIds.map((_, idx) => `$${idx + 2}`).join(', ');
+      let query = `${selectStatement} FROM ${tableName} WHERE thread_id = $1 
+        ${excludeIds.length ? `AND id NOT IN (${excludeIdsParam})` : ''}
+        ${orderByStatement}
+        LIMIT $${excludeIds.length + 2}
+        `;
+      const queryParams: any[] = [threadId, ...excludeIds, limit];
+      const remainingRows = await this.client.manyOrNone(query, queryParams);
+      rows.push(...remainingRows);
+
+      const fetchedMessages = (rows || []).map(message => {
+        if (typeof message.content === 'string') {
+          try {
+            message.content = JSON.parse(message.content);
+          } catch {
+            /* ignore */
+          }
+        }
+        if (message.type === 'v2') delete message.type;
+        return message as MastraMessageV1;
+      });
+
+      // Sort all messages by creation date
+      const sortedMessages = fetchedMessages.sort(
+        (a, b) => new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime(),
+      );
+
+      return format === 'v2'
+        ? sortedMessages.map(
+            m =>
+              ({ ...m, content: m.content || { format: 2, parts: [{ type: 'text', text: '' }] } }) as MastraMessageV2,
+          )
+        : sortedMessages;
+    } catch (error) {
+      const mastraError = new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_GET_MESSAGES_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            threadId,
+          },
+        },
+        error,
+      );
+      this.logger?.error?.(mastraError.toString());
+      this.logger?.trackException(mastraError);
+      return [];
+    }
+  }
+
+  public async getMessagesPaginated(
+    args: StorageGetMessagesArg & {
+      format?: 'v1' | 'v2';
+    },
+  ): Promise<PaginationInfo & { messages: MastraMessageV1[] | MastraMessageV2[] }> {
+    const { threadId, format, selectBy } = args;
+    const { page = 0, perPage: perPageInput, dateRange } = selectBy?.pagination || {};
+    const fromDate = dateRange?.start;
+    const toDate = dateRange?.end;
+
+    const selectStatement = `SELECT id, content, role, type, "createdAt", thread_id AS "threadId", "resourceId"`;
+    const orderByStatement = `ORDER BY "createdAt" DESC`;
+
+    const messages: MastraMessageV2[] = [];
+
+    if (selectBy?.include?.length) {
+      const includeMessages = await this._getIncludedMessages({ threadId, selectBy, orderByStatement });
+      if (includeMessages) {
+        messages.push(...includeMessages);
+      }
+    }
+
+    try {
+      const perPage =
+        perPageInput !== undefined ? perPageInput : resolveMessageLimit({ last: selectBy?.last, defaultLimit: 40 });
+      const currentOffset = page * perPage;
+
+      const conditions: string[] = [`thread_id = $1`];
+      const queryParams: any[] = [threadId];
+      let paramIndex = 2;
+
+      if (fromDate) {
+        conditions.push(`"createdAt" >= $${paramIndex++}`);
+        queryParams.push(fromDate);
+      }
+      if (toDate) {
+        conditions.push(`"createdAt" <= $${paramIndex++}`);
+        queryParams.push(toDate);
+      }
+      const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+
+      const tableName = getTableName({ indexName: TABLE_MESSAGES, schemaName: getSchemaName(this.schema) });
+      const countQuery = `SELECT COUNT(*) FROM ${tableName} ${whereClause}`;
+      const countResult = await this.client.one(countQuery, queryParams);
+      const total = parseInt(countResult.count, 10);
+
+      if (total === 0 && messages.length === 0) {
+        return {
+          messages: [],
+          total: 0,
+          page,
+          perPage,
+          hasMore: false,
+        };
+      }
+
+      const excludeIds = messages.map(m => m.id);
+      const excludeIdsParam = excludeIds.map((_, idx) => `$${idx + paramIndex}`).join(', ');
+      paramIndex += excludeIds.length;
+
+      const dataQuery = `${selectStatement} FROM ${tableName} ${whereClause} ${excludeIds.length ? `AND id NOT IN (${excludeIdsParam})` : ''}${orderByStatement} LIMIT $${paramIndex++} OFFSET $${paramIndex++}`;
+      const rows = await this.client.manyOrNone(dataQuery, [...queryParams, ...excludeIds, perPage, currentOffset]);
+      messages.push(...(rows || []));
+
+      // Parse content back to objects if they were stringified during storage
+      const messagesWithParsedContent = messages.map(message => {
+        if (typeof message.content === 'string') {
+          try {
+            return { ...message, content: JSON.parse(message.content) };
+          } catch {
+            // If parsing fails, leave as string (V1 message)
+            return message;
+          }
+        }
+        return message;
+      });
+
+      const list = new MessageList().add(messagesWithParsedContent, 'memory');
+      const messagesToReturn = format === `v2` ? list.get.all.v2() : list.get.all.v1();
+
+      return {
+        messages: messagesToReturn,
+        total,
+        page,
+        perPage,
+        hasMore: currentOffset + rows.length < total,
+      };
+    } catch (error) {
+      const mastraError = new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_GET_MESSAGES_PAGINATED_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            threadId,
+            page,
+          },
+        },
+        error,
+      );
+      this.logger?.error?.(mastraError.toString());
+      this.logger?.trackException(mastraError);
+      return { messages: [], total: 0, page, perPage: perPageInput || 40, hasMore: false };
+    }
+  }
+
+  async saveMessages(args: { messages: MastraMessageV1[]; format?: undefined | 'v1' }): Promise<MastraMessageV1[]>;
+  async saveMessages(args: { messages: MastraMessageV2[]; format: 'v2' }): Promise<MastraMessageV2[]>;
+  async saveMessages({
+    messages,
+    format,
+  }:
+    | { messages: MastraMessageV1[]; format?: undefined | 'v1' }
+    | { messages: MastraMessageV2[]; format: 'v2' }): Promise<MastraMessageV2[] | MastraMessageV1[]> {
+    if (messages.length === 0) return messages;
+
+    const threadId = messages[0]?.threadId;
+    if (!threadId) {
+      throw new MastraError({
+        id: 'MASTRA_STORAGE_PG_STORE_SAVE_MESSAGES_FAILED',
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.THIRD_PARTY,
+        text: `Thread ID is required`,
+      });
+    }
+
+    // Check if thread exists
+    const thread = await this.getThreadById({ threadId });
+    if (!thread) {
+      throw new MastraError({
+        id: 'MASTRA_STORAGE_PG_STORE_SAVE_MESSAGES_FAILED',
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.THIRD_PARTY,
+        text: `Thread ${threadId} not found`,
+        details: {
+          threadId,
+        },
+      });
+    }
+
+    try {
+      const tableName = getTableName({ indexName: TABLE_MESSAGES, schemaName: getSchemaName(this.schema) });
+      await this.client.tx(async t => {
+        // Execute message inserts and thread update in parallel for better performance
+        const messageInserts = messages.map(message => {
+          if (!message.threadId) {
+            throw new Error(
+              `Expected to find a threadId for message, but couldn't find one. An unexpected error has occurred.`,
+            );
+          }
+          if (!message.resourceId) {
+            throw new Error(
+              `Expected to find a resourceId for message, but couldn't find one. An unexpected error has occurred.`,
+            );
+          }
+          return t.none(
+            `INSERT INTO ${tableName} (id, thread_id, content, "createdAt", "createdAtZ", role, type, "resourceId") 
+             VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+             ON CONFLICT (id) DO UPDATE SET
+              thread_id = EXCLUDED.thread_id,
+              content = EXCLUDED.content,
+              role = EXCLUDED.role,
+              type = EXCLUDED.type,
+              "resourceId" = EXCLUDED."resourceId"`,
+            [
+              message.id,
+              message.threadId,
+              typeof message.content === 'string' ? message.content : JSON.stringify(message.content),
+              message.createdAt || new Date().toISOString(),
+              message.createdAt || new Date().toISOString(),
+              message.role,
+              message.type || 'v2',
+              message.resourceId,
+            ],
+          );
+        });
+
+        const threadTableName = getTableName({ indexName: TABLE_THREADS, schemaName: getSchemaName(this.schema) });
+        const threadUpdate = t.none(
+          `UPDATE ${threadTableName} 
+                        SET 
+                            "updatedAt" = $1,
+                            "updatedAtZ" = $1
+                        WHERE id = $2
+                    `,
+          [new Date().toISOString(), threadId],
+        );
+
+        await Promise.all([...messageInserts, threadUpdate]);
+      });
+
+      // Parse content back to objects if they were stringified during storage
+      const messagesWithParsedContent = messages.map(message => {
+        if (typeof message.content === 'string') {
+          try {
+            return { ...message, content: JSON.parse(message.content) };
+          } catch {
+            // If parsing fails, leave as string (V1 message)
+            return message;
+          }
+        }
+        return message;
+      });
+
+      const list = new MessageList().add(messagesWithParsedContent, 'memory');
+      if (format === `v2`) return list.get.all.v2();
+      return list.get.all.v1();
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_SAVE_MESSAGES_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            threadId,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async updateMessages({
+    messages,
+  }: {
+    messages: (Partial<Omit<MastraMessageV2, 'createdAt'>> & {
+      id: string;
+      content?: {
+        metadata?: MastraMessageContentV2['metadata'];
+        content?: MastraMessageContentV2['content'];
+      };
+    })[];
+  }): Promise<MastraMessageV2[]> {
+    if (messages.length === 0) {
+      return [];
+    }
+
+    const messageIds = messages.map(m => m.id);
+
+    const selectQuery = `SELECT id, content, role, type, "createdAt", thread_id AS "threadId", "resourceId" FROM ${getTableName({ indexName: TABLE_MESSAGES, schemaName: getSchemaName(this.schema) })} WHERE id IN ($1:list)`;
+
+    const existingMessagesDb = await this.client.manyOrNone(selectQuery, [messageIds]);
+
+    if (existingMessagesDb.length === 0) {
+      return [];
+    }
+
+    // Parse content from string to object for merging
+    const existingMessages: MastraMessageV2[] = existingMessagesDb.map(msg => {
+      if (typeof msg.content === 'string') {
+        try {
+          msg.content = JSON.parse(msg.content);
+        } catch {
+          // ignore if not valid json
+        }
+      }
+      return msg as MastraMessageV2;
+    });
+
+    const threadIdsToUpdate = new Set<string>();
+
+    await this.client.tx(async t => {
+      const queries = [];
+      const columnMapping: Record<string, string> = {
+        threadId: 'thread_id',
+      };
+
+      for (const existingMessage of existingMessages) {
+        const updatePayload = messages.find(m => m.id === existingMessage.id);
+        if (!updatePayload) continue;
+
+        const { id, ...fieldsToUpdate } = updatePayload;
+        if (Object.keys(fieldsToUpdate).length === 0) continue;
+
+        threadIdsToUpdate.add(existingMessage.threadId!);
+        if (updatePayload.threadId && updatePayload.threadId !== existingMessage.threadId) {
+          threadIdsToUpdate.add(updatePayload.threadId);
+        }
+
+        const setClauses: string[] = [];
+        const values: any[] = [];
+        let paramIndex = 1;
+
+        const updatableFields = { ...fieldsToUpdate };
+
+        // Special handling for content: merge in code, then update the whole field
+        if (updatableFields.content) {
+          const newContent = {
+            ...existingMessage.content,
+            ...updatableFields.content,
+            // Deep merge metadata if it exists on both
+            ...(existingMessage.content?.metadata && updatableFields.content.metadata
+              ? {
+                  metadata: {
+                    ...existingMessage.content.metadata,
+                    ...updatableFields.content.metadata,
+                  },
+                }
+              : {}),
+          };
+          setClauses.push(`content = $${paramIndex++}`);
+          values.push(newContent);
+          delete updatableFields.content;
+        }
+
+        for (const key in updatableFields) {
+          if (Object.prototype.hasOwnProperty.call(updatableFields, key)) {
+            const dbColumn = columnMapping[key] || key;
+            setClauses.push(`"${dbColumn}" = $${paramIndex++}`);
+            values.push(updatableFields[key as keyof typeof updatableFields]);
+          }
+        }
+
+        if (setClauses.length > 0) {
+          values.push(id);
+          const sql = `UPDATE ${getTableName({ indexName: TABLE_MESSAGES, schemaName: getSchemaName(this.schema) })} SET ${setClauses.join(', ')} WHERE id = $${paramIndex}`;
+          queries.push(t.none(sql, values));
+        }
+      }
+
+      if (threadIdsToUpdate.size > 0) {
+        queries.push(
+          t.none(
+            `UPDATE ${getTableName({ indexName: TABLE_THREADS, schemaName: getSchemaName(this.schema) })} SET "updatedAt" = NOW(), "updatedAtZ" = NOW() WHERE id IN ($1:list)`,
+            [Array.from(threadIdsToUpdate)],
+          ),
+        );
+      }
+
+      if (queries.length > 0) {
+        await t.batch(queries);
+      }
+    });
+
+    // Re-fetch to return the fully updated messages
+    const updatedMessages = await this.client.manyOrNone<MastraMessageV2>(selectQuery, [messageIds]);
+
+    return (updatedMessages || []).map(message => {
+      if (typeof message.content === 'string') {
+        try {
+          message.content = JSON.parse(message.content);
+        } catch {
+          /* ignore */
+        }
+      }
+      return message;
+    });
+  }
+
+  async getResourceById({ resourceId }: { resourceId: string }): Promise<StorageResourceType | null> {
+    const tableName = getTableName({ indexName: TABLE_RESOURCES, schemaName: getSchemaName(this.schema) });
+    const result = await this.client.oneOrNone<StorageResourceType & { createdAtZ: Date; updatedAtZ: Date }>(
+      `SELECT * FROM ${tableName} WHERE id = $1`,
+      [resourceId],
+    );
+
+    if (!result) {
+      return null;
+    }
+
+    return {
+      id: result.id,
+      createdAt: result.createdAtZ || result.createdAt,
+      updatedAt: result.updatedAtZ || result.updatedAt,
+      workingMemory: result.workingMemory,
+      metadata: typeof result.metadata === 'string' ? JSON.parse(result.metadata) : result.metadata,
+    };
+  }
+
+  async saveResource({ resource }: { resource: StorageResourceType }): Promise<StorageResourceType> {
+    await this.operations.insert({
+      tableName: TABLE_RESOURCES,
+      record: {
+        ...resource,
+        metadata: JSON.stringify(resource.metadata),
+      },
+    });
+
+    return resource;
+  }
+
+  async updateResource({
+    resourceId,
+    workingMemory,
+    metadata,
+  }: {
+    resourceId: string;
+    workingMemory?: string;
+    metadata?: Record<string, unknown>;
+  }): Promise<StorageResourceType> {
+    const existingResource = await this.getResourceById({ resourceId });
+
+    if (!existingResource) {
+      // Create new resource if it doesn't exist
+      const newResource: StorageResourceType = {
+        id: resourceId,
+        workingMemory,
+        metadata: metadata || {},
+        createdAt: new Date(),
+        updatedAt: new Date(),
+      };
+      return this.saveResource({ resource: newResource });
+    }
+
+    const updatedResource = {
+      ...existingResource,
+      workingMemory: workingMemory !== undefined ? workingMemory : existingResource.workingMemory,
+      metadata: {
+        ...existingResource.metadata,
+        ...metadata,
+      },
+      updatedAt: new Date(),
+    };
+
+    const tableName = getTableName({ indexName: TABLE_RESOURCES, schemaName: getSchemaName(this.schema) });
+
+    const updates: string[] = [];
+    const values: any[] = [];
+    let paramIndex = 1;
+
+    if (workingMemory !== undefined) {
+      updates.push(`"workingMemory" = $${paramIndex}`);
+      values.push(workingMemory);
+      paramIndex++;
+    }
+
+    if (metadata) {
+      updates.push(`metadata = $${paramIndex}`);
+      values.push(JSON.stringify(updatedResource.metadata));
+      paramIndex++;
+    }
+
+    updates.push(`"updatedAt" = $${paramIndex}`);
+    values.push(updatedResource.updatedAt.toISOString());
+    updates.push(`"updatedAtZ" = $${paramIndex++}`);
+    values.push(updatedResource.updatedAt.toISOString());
+
+    paramIndex++;
+
+    values.push(resourceId);
+
+    await this.client.none(`UPDATE ${tableName} SET ${updates.join(', ')} WHERE id = $${paramIndex}`, values);
+
+    return updatedResource;
+  }
+}