@mastra/pg 1.0.0-beta.10 → 1.0.0-beta.11

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.js CHANGED
@@ -1,12 +1,12 @@
1
1
  import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
2
- import { createVectorErrorId, AgentsStorage, TABLE_AGENTS, TABLE_SCHEMAS, createStorageErrorId, normalizePerPage, calculatePagination, MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, ObservabilityStorage, TABLE_SPANS, listTracesArgsSchema, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraStorage, TraceStatus, getSqlType, getDefaultValue, transformScoreRow as transformScoreRow$1 } from '@mastra/core/storage';
2
+ import { createVectorErrorId, TABLE_SCHEMAS, AgentsStorage, TABLE_AGENTS, createStorageErrorId, normalizePerPage, calculatePagination, MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, ObservabilityStorage, TABLE_SPANS, listTracesArgsSchema, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraStorage, TraceStatus, getDefaultValue, transformScoreRow as transformScoreRow$1, getSqlType } from '@mastra/core/storage';
3
3
  import { parseSqlIdentifier, parseFieldKey } from '@mastra/core/utils';
4
4
  import { MastraVector } from '@mastra/core/vector';
5
5
  import { Mutex } from 'async-mutex';
6
6
  import * as pg from 'pg';
7
+ import { Pool } from 'pg';
7
8
  import xxhash from 'xxhash-wasm';
8
9
  import { BaseFilterTranslator } from '@mastra/core/vector/filter';
9
- import pgPromise from 'pg-promise';
10
10
  import { MastraBase } from '@mastra/core/base';
11
11
  import { MessageList } from '@mastra/core/agent';
12
12
  import { saveScorePayloadSchema } from '@mastra/core/evals';
@@ -14,8 +14,11 @@ import { saveScorePayloadSchema } from '@mastra/core/evals';
14
14
  // src/vector/index.ts
15
15
 
16
16
  // src/shared/config.ts
17
+ var isPoolConfig = (cfg) => {
18
+ return "pool" in cfg;
19
+ };
17
20
  var isConnectionStringConfig = (cfg) => {
18
- return "connectionString" in cfg;
21
+ return "connectionString" in cfg && typeof cfg.connectionString === "string";
19
22
  };
20
23
  var isHostConfig = (cfg) => {
21
24
  return "host" in cfg && "database" in cfg && "user" in cfg && "password" in cfg;
@@ -23,16 +26,13 @@ var isHostConfig = (cfg) => {
23
26
  var isCloudSqlConfig = (cfg) => {
24
27
  return "stream" in cfg || "password" in cfg && typeof cfg.password === "function";
25
28
  };
26
- var isClientConfig = (cfg) => {
27
- return "client" in cfg;
28
- };
29
29
  var validateConfig = (name, config) => {
30
30
  if (!config.id || typeof config.id !== "string" || config.id.trim() === "") {
31
31
  throw new Error(`${name}: id must be provided and cannot be empty.`);
32
32
  }
33
- if (isClientConfig(config)) {
34
- if (!config.client) {
35
- throw new Error(`${name}: client must be provided when using client config.`);
33
+ if (isPoolConfig(config)) {
34
+ if (!config.pool) {
35
+ throw new Error(`${name}: pool must be provided when using pool config.`);
36
36
  }
37
37
  return;
38
38
  }
@@ -53,7 +53,7 @@ var validateConfig = (name, config) => {
53
53
  }
54
54
  } else {
55
55
  throw new Error(
56
- `${name}: invalid config. Provide either {client}, {connectionString}, {host,port,database,user,password}, or a pg ClientConfig (e.g., Cloud SQL connector with \`stream\`).`
56
+ `${name}: invalid config. Provide either {pool}, {connectionString}, {host,port,database,user,password}, or a pg ClientConfig (e.g., Cloud SQL connector with \`stream\`).`
57
57
  );
58
58
  }
59
59
  };
@@ -518,8 +518,8 @@ var PgVector = class extends MastraVector {
518
518
  } else if (isCloudSqlConfig(config)) {
519
519
  poolConfig = {
520
520
  ...config,
521
- max: config.max ?? 20,
522
- idleTimeoutMillis: config.idleTimeoutMillis ?? 3e4,
521
+ max: config.pgPoolOptions?.max ?? 20,
522
+ idleTimeoutMillis: config.pgPoolOptions?.idleTimeoutMillis ?? 3e4,
523
523
  connectionTimeoutMillis: 2e3,
524
524
  ...config.pgPoolOptions
525
525
  };
@@ -1661,6 +1661,132 @@ var PgVector = class extends MastraVector {
1661
1661
  }
1662
1662
  }
1663
1663
  };
1664
+
1665
+ // src/storage/client.ts
1666
+ function truncateQuery(query, maxLength = 100) {
1667
+ const normalized = query.replace(/\s+/g, " ").trim();
1668
+ if (normalized.length <= maxLength) {
1669
+ return normalized;
1670
+ }
1671
+ return normalized.slice(0, maxLength) + "...";
1672
+ }
1673
+ var PoolAdapter = class {
1674
+ constructor($pool) {
1675
+ this.$pool = $pool;
1676
+ }
1677
+ connect() {
1678
+ return this.$pool.connect();
1679
+ }
1680
+ async none(query, values) {
1681
+ await this.$pool.query(query, values);
1682
+ return null;
1683
+ }
1684
+ async one(query, values) {
1685
+ const result = await this.$pool.query(query, values);
1686
+ if (result.rows.length === 0) {
1687
+ throw new Error(`No data returned from query: ${truncateQuery(query)}`);
1688
+ }
1689
+ if (result.rows.length > 1) {
1690
+ throw new Error(`Multiple rows returned when one was expected: ${truncateQuery(query)}`);
1691
+ }
1692
+ return result.rows[0];
1693
+ }
1694
+ async oneOrNone(query, values) {
1695
+ const result = await this.$pool.query(query, values);
1696
+ if (result.rows.length === 0) {
1697
+ return null;
1698
+ }
1699
+ if (result.rows.length > 1) {
1700
+ throw new Error(`Multiple rows returned when one or none was expected: ${truncateQuery(query)}`);
1701
+ }
1702
+ return result.rows[0];
1703
+ }
1704
+ async any(query, values) {
1705
+ const result = await this.$pool.query(query, values);
1706
+ return result.rows;
1707
+ }
1708
+ async manyOrNone(query, values) {
1709
+ return this.any(query, values);
1710
+ }
1711
+ async many(query, values) {
1712
+ const result = await this.$pool.query(query, values);
1713
+ if (result.rows.length === 0) {
1714
+ throw new Error(`No data returned from query: ${truncateQuery(query)}`);
1715
+ }
1716
+ return result.rows;
1717
+ }
1718
+ async query(query, values) {
1719
+ return this.$pool.query(query, values);
1720
+ }
1721
+ async tx(callback) {
1722
+ const client = await this.$pool.connect();
1723
+ try {
1724
+ await client.query("BEGIN");
1725
+ const txClient = new TransactionClient(client);
1726
+ const result = await callback(txClient);
1727
+ await client.query("COMMIT");
1728
+ return result;
1729
+ } catch (error) {
1730
+ try {
1731
+ await client.query("ROLLBACK");
1732
+ } catch (rollbackError) {
1733
+ console.error("Transaction rollback failed:", rollbackError);
1734
+ }
1735
+ throw error;
1736
+ } finally {
1737
+ client.release();
1738
+ }
1739
+ }
1740
+ };
1741
+ var TransactionClient = class {
1742
+ constructor(client) {
1743
+ this.client = client;
1744
+ }
1745
+ async none(query, values) {
1746
+ await this.client.query(query, values);
1747
+ return null;
1748
+ }
1749
+ async one(query, values) {
1750
+ const result = await this.client.query(query, values);
1751
+ if (result.rows.length === 0) {
1752
+ throw new Error(`No data returned from query: ${truncateQuery(query)}`);
1753
+ }
1754
+ if (result.rows.length > 1) {
1755
+ throw new Error(`Multiple rows returned when one was expected: ${truncateQuery(query)}`);
1756
+ }
1757
+ return result.rows[0];
1758
+ }
1759
+ async oneOrNone(query, values) {
1760
+ const result = await this.client.query(query, values);
1761
+ if (result.rows.length === 0) {
1762
+ return null;
1763
+ }
1764
+ if (result.rows.length > 1) {
1765
+ throw new Error(`Multiple rows returned when one or none was expected: ${truncateQuery(query)}`);
1766
+ }
1767
+ return result.rows[0];
1768
+ }
1769
+ async any(query, values) {
1770
+ const result = await this.client.query(query, values);
1771
+ return result.rows;
1772
+ }
1773
+ async manyOrNone(query, values) {
1774
+ return this.any(query, values);
1775
+ }
1776
+ async many(query, values) {
1777
+ const result = await this.client.query(query, values);
1778
+ if (result.rows.length === 0) {
1779
+ throw new Error(`No data returned from query: ${truncateQuery(query)}`);
1780
+ }
1781
+ return result.rows;
1782
+ }
1783
+ async query(query, values) {
1784
+ return this.client.query(query, values);
1785
+ }
1786
+ async batch(promises) {
1787
+ return Promise.all(promises);
1788
+ }
1789
+ };
1664
1790
  function resolvePgConfig(config) {
1665
1791
  if ("client" in config) {
1666
1792
  return {
@@ -1670,10 +1796,32 @@ function resolvePgConfig(config) {
1670
1796
  indexes: config.indexes
1671
1797
  };
1672
1798
  }
1673
- const pgp = pgPromise();
1674
- const client = pgp(config);
1799
+ if ("pool" in config) {
1800
+ return {
1801
+ client: new PoolAdapter(config.pool),
1802
+ schemaName: config.schemaName,
1803
+ skipDefaultIndexes: config.skipDefaultIndexes,
1804
+ indexes: config.indexes
1805
+ };
1806
+ }
1807
+ let pool;
1808
+ if ("connectionString" in config) {
1809
+ pool = new Pool({
1810
+ connectionString: config.connectionString,
1811
+ ssl: config.ssl
1812
+ });
1813
+ } else {
1814
+ pool = new Pool({
1815
+ host: config.host,
1816
+ port: config.port,
1817
+ database: config.database,
1818
+ user: config.user,
1819
+ password: config.password,
1820
+ ssl: config.ssl
1821
+ });
1822
+ }
1675
1823
  return {
1676
- client,
1824
+ client: new PoolAdapter(pool),
1677
1825
  schemaName: config.schemaName,
1678
1826
  skipDefaultIndexes: config.skipDefaultIndexes,
1679
1827
  indexes: config.indexes
@@ -1688,6 +1836,87 @@ function getTableName({ indexName, schemaName }) {
1688
1836
  const quotedSchemaName = schemaName;
1689
1837
  return quotedSchemaName ? `${quotedSchemaName}.${quotedIndexName}` : quotedIndexName;
1690
1838
  }
1839
+ function mapToSqlType(type) {
1840
+ switch (type) {
1841
+ case "uuid":
1842
+ return "UUID";
1843
+ case "boolean":
1844
+ return "BOOLEAN";
1845
+ default:
1846
+ return getSqlType(type);
1847
+ }
1848
+ }
1849
+ function generateTableSQL({
1850
+ tableName,
1851
+ schema,
1852
+ schemaName
1853
+ }) {
1854
+ const timeZColumns = Object.entries(schema).filter(([_, def]) => def.type === "timestamp").map(([name]) => {
1855
+ const parsedName = parseSqlIdentifier(name, "column name");
1856
+ return `"${parsedName}Z" TIMESTAMPTZ DEFAULT NOW()`;
1857
+ });
1858
+ const columns = Object.entries(schema).map(([name, def]) => {
1859
+ const parsedName = parseSqlIdentifier(name, "column name");
1860
+ const constraints = [];
1861
+ if (def.primaryKey) constraints.push("PRIMARY KEY");
1862
+ if (!def.nullable) constraints.push("NOT NULL");
1863
+ return `"${parsedName}" ${mapToSqlType(def.type)} ${constraints.join(" ")}`;
1864
+ });
1865
+ const finalColumns = [...columns, ...timeZColumns].join(",\n");
1866
+ const parsedSchemaName = schemaName ? parseSqlIdentifier(schemaName, "schema name") : "";
1867
+ const constraintPrefix = parsedSchemaName ? `${parsedSchemaName}_` : "";
1868
+ const quotedSchemaName = getSchemaName(schemaName);
1869
+ const sql = `
1870
+ CREATE TABLE IF NOT EXISTS ${getTableName({ indexName: tableName, schemaName: quotedSchemaName })} (
1871
+ ${finalColumns}
1872
+ );
1873
+ ${tableName === TABLE_WORKFLOW_SNAPSHOT ? `
1874
+ DO $$ BEGIN
1875
+ IF NOT EXISTS (
1876
+ SELECT 1 FROM pg_constraint WHERE conname = '${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key'
1877
+ ) AND NOT EXISTS (
1878
+ SELECT 1 FROM pg_indexes WHERE indexname = '${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key'
1879
+ ) THEN
1880
+ ALTER TABLE ${getTableName({ indexName: tableName, schemaName: quotedSchemaName })}
1881
+ ADD CONSTRAINT ${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key
1882
+ UNIQUE (workflow_name, run_id);
1883
+ END IF;
1884
+ END $$;
1885
+ ` : ""}
1886
+ ${tableName === TABLE_SPANS ? `
1887
+ DO $$ BEGIN
1888
+ IF NOT EXISTS (
1889
+ SELECT 1 FROM pg_constraint WHERE conname = '${constraintPrefix}mastra_ai_spans_traceid_spanid_pk'
1890
+ ) THEN
1891
+ ALTER TABLE ${getTableName({ indexName: tableName, schemaName: quotedSchemaName })}
1892
+ ADD CONSTRAINT ${constraintPrefix}mastra_ai_spans_traceid_spanid_pk
1893
+ PRIMARY KEY ("traceId", "spanId");
1894
+ END IF;
1895
+ END $$;
1896
+ ` : ""}
1897
+ `;
1898
+ return sql;
1899
+ }
1900
+ function exportSchemas(schemaName) {
1901
+ const statements = [];
1902
+ if (schemaName) {
1903
+ const quotedSchemaName = getSchemaName(schemaName);
1904
+ statements.push(`-- Create schema if it doesn't exist`);
1905
+ statements.push(`CREATE SCHEMA IF NOT EXISTS ${quotedSchemaName};`);
1906
+ statements.push("");
1907
+ }
1908
+ for (const [tableName, schema] of Object.entries(TABLE_SCHEMAS)) {
1909
+ statements.push(`-- Table: ${tableName}`);
1910
+ const sql = generateTableSQL({
1911
+ tableName,
1912
+ schema,
1913
+ schemaName
1914
+ });
1915
+ statements.push(sql.trim());
1916
+ statements.push("");
1917
+ }
1918
+ return statements.join("\n");
1919
+ }
1691
1920
  var schemaSetupRegistry = /* @__PURE__ */ new Map();
1692
1921
  var PgDB = class extends MastraBase {
1693
1922
  client;
@@ -1805,16 +2034,6 @@ var PgDB = class extends MastraBase {
1805
2034
  }
1806
2035
  await registryEntry.promise;
1807
2036
  }
1808
- getSqlType(type) {
1809
- switch (type) {
1810
- case "uuid":
1811
- return "UUID";
1812
- case "boolean":
1813
- return "BOOLEAN";
1814
- default:
1815
- return getSqlType(type);
1816
- }
1817
- }
1818
2037
  getDefaultValue(type) {
1819
2038
  switch (type) {
1820
2039
  case "timestamp":
@@ -1884,52 +2103,10 @@ var PgDB = class extends MastraBase {
1884
2103
  }) {
1885
2104
  try {
1886
2105
  const timeZColumnNames = Object.entries(schema).filter(([_, def]) => def.type === "timestamp").map(([name]) => name);
1887
- const timeZColumns = Object.entries(schema).filter(([_, def]) => def.type === "timestamp").map(([name]) => {
1888
- const parsedName = parseSqlIdentifier(name, "column name");
1889
- return `"${parsedName}Z" TIMESTAMPTZ DEFAULT NOW()`;
1890
- });
1891
- const columns = Object.entries(schema).map(([name, def]) => {
1892
- const parsedName = parseSqlIdentifier(name, "column name");
1893
- const constraints = [];
1894
- if (def.primaryKey) constraints.push("PRIMARY KEY");
1895
- if (!def.nullable) constraints.push("NOT NULL");
1896
- return `"${parsedName}" ${this.getSqlType(def.type)} ${constraints.join(" ")}`;
1897
- });
1898
2106
  if (this.schemaName) {
1899
2107
  await this.setupSchema();
1900
2108
  }
1901
- const finalColumns = [...columns, ...timeZColumns].join(",\n");
1902
- const constraintPrefix = this.schemaName ? `${this.schemaName}_` : "";
1903
- const schemaName = getSchemaName(this.schemaName);
1904
- const sql = `
1905
- CREATE TABLE IF NOT EXISTS ${getTableName({ indexName: tableName, schemaName })} (
1906
- ${finalColumns}
1907
- );
1908
- ${tableName === TABLE_WORKFLOW_SNAPSHOT ? `
1909
- DO $$ BEGIN
1910
- IF NOT EXISTS (
1911
- SELECT 1 FROM pg_constraint WHERE conname = '${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key'
1912
- ) AND NOT EXISTS (
1913
- SELECT 1 FROM pg_indexes WHERE indexname = '${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key'
1914
- ) THEN
1915
- ALTER TABLE ${getTableName({ indexName: tableName, schemaName })}
1916
- ADD CONSTRAINT ${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key
1917
- UNIQUE (workflow_name, run_id);
1918
- END IF;
1919
- END $$;
1920
- ` : ""}
1921
- ${tableName === TABLE_SPANS ? `
1922
- DO $$ BEGIN
1923
- IF NOT EXISTS (
1924
- SELECT 1 FROM pg_constraint WHERE conname = '${constraintPrefix}mastra_ai_spans_traceid_spanid_pk'
1925
- ) THEN
1926
- ALTER TABLE ${getTableName({ indexName: tableName, schemaName: getSchemaName(this.schemaName) })}
1927
- ADD CONSTRAINT ${constraintPrefix}mastra_ai_spans_traceid_spanid_pk
1928
- PRIMARY KEY ("traceId", "spanId");
1929
- END IF;
1930
- END $$;
1931
- ` : ""}
1932
- `;
2109
+ const sql = generateTableSQL({ tableName, schema, schemaName: this.schemaName });
1933
2110
  await this.client.none(sql);
1934
2111
  await this.alterTable({
1935
2112
  tableName,
@@ -2003,7 +2180,7 @@ var PgDB = class extends MastraBase {
2003
2180
  const columnExists = await this.hasColumn(TABLE_SPANS, columnName);
2004
2181
  if (!columnExists) {
2005
2182
  const parsedColumnName = parseSqlIdentifier(columnName, "column name");
2006
- const sqlType = this.getSqlType(columnDef.type);
2183
+ const sqlType = mapToSqlType(columnDef.type);
2007
2184
  const nullable = columnDef.nullable ? "" : "NOT NULL";
2008
2185
  const defaultValue = !columnDef.nullable ? this.getDefaultValue(columnDef.type) : "";
2009
2186
  const alterSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}" ${sqlType} ${nullable} ${defaultValue}`.trim();
@@ -2050,7 +2227,7 @@ var PgDB = class extends MastraBase {
2050
2227
  if (schema[columnName]) {
2051
2228
  const columnDef = schema[columnName];
2052
2229
  const parsedColumnName = parseSqlIdentifier(columnName, "column name");
2053
- const sqlType = this.getSqlType(columnDef.type);
2230
+ const sqlType = mapToSqlType(columnDef.type);
2054
2231
  const nullable = columnDef.nullable ? "" : "NOT NULL";
2055
2232
  const defaultValue = !columnDef.nullable ? this.getDefaultValue(columnDef.type) : "";
2056
2233
  const alterSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}" ${sqlType} ${nullable} ${defaultValue}`.trim();
@@ -2370,9 +2547,9 @@ var PgDB = class extends MastraBase {
2370
2547
  size: result.size || "0",
2371
2548
  definition: result.definition || "",
2372
2549
  method: result.method || "btree",
2373
- scans: parseInt(result.scans) || 0,
2374
- tuples_read: parseInt(result.tuples_read) || 0,
2375
- tuples_fetched: parseInt(result.tuples_fetched) || 0
2550
+ scans: parseInt(String(result.scans)) || 0,
2551
+ tuples_read: parseInt(String(result.tuples_read)) || 0,
2552
+ tuples_fetched: parseInt(String(result.tuples_fetched)) || 0
2376
2553
  };
2377
2554
  } catch (error) {
2378
2555
  throw new MastraError(
@@ -2889,6 +3066,9 @@ function getTableName3({ indexName, schemaName }) {
2889
3066
  const quotedIndexName = `"${indexName}"`;
2890
3067
  return schemaName ? `${schemaName}.${quotedIndexName}` : quotedIndexName;
2891
3068
  }
3069
+ function inPlaceholders(count, startIndex = 1) {
3070
+ return Array.from({ length: count }, (_, i) => `$${i + startIndex}`).join(", ");
3071
+ }
2892
3072
  var MemoryPG = class _MemoryPG extends MemoryStorage {
2893
3073
  #db;
2894
3074
  #schema;
@@ -3160,17 +3340,18 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
3160
3340
  ...metadata
3161
3341
  };
3162
3342
  try {
3343
+ const now = (/* @__PURE__ */ new Date()).toISOString();
3163
3344
  const thread = await this.#db.client.one(
3164
3345
  `UPDATE ${threadTableName}
3165
3346
  SET
3166
3347
  title = $1,
3167
3348
  metadata = $2,
3168
3349
  "updatedAt" = $3,
3169
- "updatedAtZ" = $3
3170
- WHERE id = $4
3350
+ "updatedAtZ" = $4
3351
+ WHERE id = $5
3171
3352
  RETURNING *
3172
3353
  `,
3173
- [title, mergedMetadata, (/* @__PURE__ */ new Date()).toISOString(), id]
3354
+ [title, mergedMetadata, now, now, id]
3174
3355
  );
3175
3356
  return {
3176
3357
  id: thread.id,
@@ -3313,7 +3494,7 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
3313
3494
  const tableName = getTableName3({ indexName: TABLE_MESSAGES, schemaName: getSchemaName3(this.#schema) });
3314
3495
  const query = `
3315
3496
  ${selectStatement} FROM ${tableName}
3316
- WHERE id IN (${messageIds.map((_, i) => `$${i + 1}`).join(", ")})
3497
+ WHERE id IN (${inPlaceholders(messageIds.length)})
3317
3498
  ORDER BY "createdAt" DESC
3318
3499
  `;
3319
3500
  const resultRows = await this.#db.client.manyOrNone(query, messageIds);
@@ -3374,8 +3555,7 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
3374
3555
  const orderByStatement = `ORDER BY "${field}" ${direction}`;
3375
3556
  const selectStatement = `SELECT id, content, role, type, "createdAt", "createdAtZ", thread_id AS "threadId", "resourceId"`;
3376
3557
  const tableName = getTableName3({ indexName: TABLE_MESSAGES, schemaName: getSchemaName3(this.#schema) });
3377
- const threadPlaceholders = threadIds.map((_, i) => `$${i + 1}`).join(", ");
3378
- const conditions = [`thread_id IN (${threadPlaceholders})`];
3558
+ const conditions = [`thread_id IN (${inPlaceholders(threadIds.length)})`];
3379
3559
  const queryParams = [...threadIds];
3380
3560
  let paramIndex = threadIds.length + 1;
3381
3561
  if (resourceId) {
@@ -3532,14 +3712,15 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
3532
3712
  );
3533
3713
  });
3534
3714
  const threadTableName = getTableName3({ indexName: TABLE_THREADS, schemaName: getSchemaName3(this.#schema) });
3715
+ const nowStr = (/* @__PURE__ */ new Date()).toISOString();
3535
3716
  const threadUpdate = t.none(
3536
3717
  `UPDATE ${threadTableName}
3537
3718
  SET
3538
3719
  "updatedAt" = $1,
3539
- "updatedAtZ" = $1
3540
- WHERE id = $2
3720
+ "updatedAtZ" = $2
3721
+ WHERE id = $3
3541
3722
  `,
3542
- [(/* @__PURE__ */ new Date()).toISOString(), threadId]
3723
+ [nowStr, nowStr, threadId]
3543
3724
  );
3544
3725
  await Promise.all([...messageInserts, threadUpdate]);
3545
3726
  });
@@ -3576,8 +3757,8 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
3576
3757
  return [];
3577
3758
  }
3578
3759
  const messageIds = messages.map((m) => m.id);
3579
- const selectQuery = `SELECT id, content, role, type, "createdAt", "createdAtZ", thread_id AS "threadId", "resourceId" FROM ${getTableName3({ indexName: TABLE_MESSAGES, schemaName: getSchemaName3(this.#schema) })} WHERE id IN ($1:list)`;
3580
- const existingMessagesDb = await this.#db.client.manyOrNone(selectQuery, [messageIds]);
3760
+ const selectQuery = `SELECT id, content, role, type, "createdAt", "createdAtZ", thread_id AS "threadId", "resourceId" FROM ${getTableName3({ indexName: TABLE_MESSAGES, schemaName: getSchemaName3(this.#schema) })} WHERE id IN (${inPlaceholders(messageIds.length)})`;
3761
+ const existingMessagesDb = await this.#db.client.manyOrNone(selectQuery, messageIds);
3581
3762
  if (existingMessagesDb.length === 0) {
3582
3763
  return [];
3583
3764
  }
@@ -3638,10 +3819,11 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
3638
3819
  }
3639
3820
  }
3640
3821
  if (threadIdsToUpdate.size > 0) {
3822
+ const threadIds = Array.from(threadIdsToUpdate);
3641
3823
  queries.push(
3642
3824
  t.none(
3643
- `UPDATE ${getTableName3({ indexName: TABLE_THREADS, schemaName: getSchemaName3(this.#schema) })} SET "updatedAt" = NOW(), "updatedAtZ" = NOW() WHERE id IN ($1:list)`,
3644
- [Array.from(threadIdsToUpdate)]
3825
+ `UPDATE ${getTableName3({ indexName: TABLE_THREADS, schemaName: getSchemaName3(this.#schema) })} SET "updatedAt" = NOW(), "updatedAtZ" = NOW() WHERE id IN (${inPlaceholders(threadIds.length)})`,
3826
+ threadIds
3645
3827
  )
3646
3828
  );
3647
3829
  }
@@ -3649,7 +3831,7 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
3649
3831
  await t.batch(queries);
3650
3832
  }
3651
3833
  });
3652
- const updatedMessages = await this.#db.client.manyOrNone(selectQuery, [messageIds]);
3834
+ const updatedMessages = await this.#db.client.manyOrNone(selectQuery, messageIds);
3653
3835
  return (updatedMessages || []).map((row) => {
3654
3836
  const message = this.normalizeMessageRow(row);
3655
3837
  if (typeof message.content === "string") {
@@ -3761,11 +3943,11 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
3761
3943
  values.push(JSON.stringify(updatedResource.metadata));
3762
3944
  paramIndex++;
3763
3945
  }
3764
- updates.push(`"updatedAt" = $${paramIndex}`);
3765
- values.push(updatedResource.updatedAt.toISOString());
3946
+ const updatedAtStr = updatedResource.updatedAt.toISOString();
3947
+ updates.push(`"updatedAt" = $${paramIndex++}`);
3948
+ values.push(updatedAtStr);
3766
3949
  updates.push(`"updatedAtZ" = $${paramIndex++}`);
3767
- values.push(updatedResource.updatedAt.toISOString());
3768
- paramIndex++;
3950
+ values.push(updatedAtStr);
3769
3951
  values.push(resourceId);
3770
3952
  await this.#db.client.none(`UPDATE ${tableName} SET ${updates.join(", ")} WHERE id = $${paramIndex}`, values);
3771
3953
  return updatedResource;
@@ -5044,9 +5226,12 @@ var WorkflowsPG = class _WorkflowsPG extends WorkflowsStorage {
5044
5226
  };
5045
5227
 
5046
5228
  // src/storage/index.ts
5229
+ var DEFAULT_MAX_CONNECTIONS = 20;
5230
+ var DEFAULT_IDLE_TIMEOUT_MS = 3e4;
5047
5231
  var PostgresStore = class extends MastraStorage {
5232
+ #pool;
5048
5233
  #db;
5049
- #pgp;
5234
+ #ownsPool;
5050
5235
  schema;
5051
5236
  isInitialized = false;
5052
5237
  stores;
@@ -5055,59 +5240,26 @@ var PostgresStore = class extends MastraStorage {
5055
5240
  validateConfig("PostgresStore", config);
5056
5241
  super({ id: config.id, name: "PostgresStore", disableInit: config.disableInit });
5057
5242
  this.schema = config.schemaName || "public";
5058
- this.#pgp = pgPromise();
5059
- if (isClientConfig(config)) {
5060
- this.#db = config.client;
5243
+ if (isPoolConfig(config)) {
5244
+ this.#pool = config.pool;
5245
+ this.#ownsPool = false;
5061
5246
  } else {
5062
- let pgConfig;
5063
- if (isConnectionStringConfig(config)) {
5064
- pgConfig = {
5065
- id: config.id,
5066
- connectionString: config.connectionString,
5067
- max: config.max,
5068
- idleTimeoutMillis: config.idleTimeoutMillis,
5069
- ssl: config.ssl
5070
- };
5071
- } else if (isCloudSqlConfig(config)) {
5072
- pgConfig = {
5073
- ...config,
5074
- id: config.id,
5075
- max: config.max,
5076
- idleTimeoutMillis: config.idleTimeoutMillis
5077
- };
5078
- } else if (isHostConfig(config)) {
5079
- pgConfig = {
5080
- id: config.id,
5081
- host: config.host,
5082
- port: config.port,
5083
- database: config.database,
5084
- user: config.user,
5085
- password: config.password,
5086
- ssl: config.ssl,
5087
- max: config.max,
5088
- idleTimeoutMillis: config.idleTimeoutMillis
5089
- };
5090
- } else {
5091
- throw new Error(
5092
- "PostgresStore: invalid config. Provide either {client}, {connectionString}, {host,port,database,user,password}, or a pg ClientConfig (e.g., Cloud SQL connector with `stream`)."
5093
- );
5094
- }
5095
- this.#db = this.#pgp(pgConfig);
5096
- }
5097
- const skipDefaultIndexes = config.skipDefaultIndexes;
5098
- const indexes = config.indexes;
5099
- const domainConfig = { client: this.#db, schemaName: this.schema, skipDefaultIndexes, indexes };
5100
- const scores = new ScoresPG(domainConfig);
5101
- const workflows = new WorkflowsPG(domainConfig);
5102
- const memory = new MemoryPG(domainConfig);
5103
- const observability = new ObservabilityPG(domainConfig);
5104
- const agents = new AgentsPG(domainConfig);
5247
+ this.#pool = this.createPool(config);
5248
+ this.#ownsPool = true;
5249
+ }
5250
+ this.#db = new PoolAdapter(this.#pool);
5251
+ const domainConfig = {
5252
+ client: this.#db,
5253
+ schemaName: this.schema,
5254
+ skipDefaultIndexes: config.skipDefaultIndexes,
5255
+ indexes: config.indexes
5256
+ };
5105
5257
  this.stores = {
5106
- scores,
5107
- workflows,
5108
- memory,
5109
- observability,
5110
- agents
5258
+ scores: new ScoresPG(domainConfig),
5259
+ workflows: new WorkflowsPG(domainConfig),
5260
+ memory: new MemoryPG(domainConfig),
5261
+ observability: new ObservabilityPG(domainConfig),
5262
+ agents: new AgentsPG(domainConfig)
5111
5263
  };
5112
5264
  } catch (e) {
5113
5265
  throw new MastraError(
@@ -5120,6 +5272,32 @@ var PostgresStore = class extends MastraStorage {
5120
5272
  );
5121
5273
  }
5122
5274
  }
5275
+ createPool(config) {
5276
+ if (isConnectionStringConfig(config)) {
5277
+ return new Pool({
5278
+ connectionString: config.connectionString,
5279
+ ssl: config.ssl,
5280
+ max: config.max ?? DEFAULT_MAX_CONNECTIONS,
5281
+ idleTimeoutMillis: config.idleTimeoutMillis ?? DEFAULT_IDLE_TIMEOUT_MS
5282
+ });
5283
+ }
5284
+ if (isHostConfig(config)) {
5285
+ return new Pool({
5286
+ host: config.host,
5287
+ port: config.port,
5288
+ database: config.database,
5289
+ user: config.user,
5290
+ password: config.password,
5291
+ ssl: config.ssl,
5292
+ max: config.max ?? DEFAULT_MAX_CONNECTIONS,
5293
+ idleTimeoutMillis: config.idleTimeoutMillis ?? DEFAULT_IDLE_TIMEOUT_MS
5294
+ });
5295
+ }
5296
+ if (isCloudSqlConfig(config)) {
5297
+ return new Pool(config);
5298
+ }
5299
+ throw new Error("PostgresStore: invalid config");
5300
+ }
5123
5301
  async init() {
5124
5302
  if (this.isInitialized) {
5125
5303
  return;
@@ -5139,19 +5317,32 @@ var PostgresStore = class extends MastraStorage {
5139
5317
  );
5140
5318
  }
5141
5319
  }
5320
+ /**
5321
+ * Database client for executing queries.
5322
+ *
5323
+ * @example
5324
+ * ```typescript
5325
+ * const rows = await store.db.any('SELECT * FROM users WHERE active = $1', [true]);
5326
+ * const user = await store.db.one('SELECT * FROM users WHERE id = $1', [userId]);
5327
+ * ```
5328
+ */
5142
5329
  get db() {
5143
5330
  return this.#db;
5144
5331
  }
5145
- get pgp() {
5146
- return this.#pgp;
5332
+ /**
5333
+ * The underlying pg.Pool for direct database access or ORM integration.
5334
+ */
5335
+ get pool() {
5336
+ return this.#pool;
5147
5337
  }
5148
5338
  /**
5149
- * Closes the pg-promise connection pool.
5150
- *
5151
- * This will close ALL connections in the pool, including pre-configured clients.
5339
+ * Closes the connection pool if it was created by this store.
5340
+ * If a pool was passed in via config, it will not be closed.
5152
5341
  */
5153
5342
  async close() {
5154
- this.pgp.end();
5343
+ if (this.#ownsPool) {
5344
+ await this.#pool.end();
5345
+ }
5155
5346
  }
5156
5347
  };
5157
5348
 
@@ -5254,6 +5445,6 @@ Example Complex Query:
5254
5445
  ]
5255
5446
  }`;
5256
5447
 
5257
- export { AgentsPG, MemoryPG, ObservabilityPG, PGVECTOR_PROMPT, PgVector, PostgresStore, ScoresPG, WorkflowsPG };
5448
+ export { AgentsPG, MemoryPG, ObservabilityPG, PGVECTOR_PROMPT, PgVector, PoolAdapter, PostgresStore, ScoresPG, WorkflowsPG, exportSchemas };
5258
5449
  //# sourceMappingURL=index.js.map
5259
5450
  //# sourceMappingURL=index.js.map