1
/*
2
 *  This file is part of CoCalc: Copyright © 2021 Sagemath, Inc.
3
 *  License: MS-RSL – see LICENSE.md for details
4
 */
5

6
import { Client, Pool, PoolClient } from "pg";
7
import { syncSchema } from "@cocalc/database/postgres/schema";
8
import {
9
  pgdatabase as database,
10
  pghost as host,
11
  pguser as user,
12
  pgssl as ssl,
13
} from "@cocalc/backend/data";
14
import { getLogger } from "@cocalc/backend/logger";
15
import { STATEMENT_TIMEOUT_MS } from "../consts";
16
import getCachedPool, { CacheTime } from "./cached";
17
import dbPassword from "./password";
18

19
const L = getLogger("db:pool");
20

21
export * from "./util";
22

23
let pool: Pool | undefined = undefined;
24

25
export default function getPool(cacheTime?: CacheTime): Pool {
26
  if (cacheTime != null) {
27
    return getCachedPool(cacheTime);
28
  }
29
  if (pool == null) {
30
    L.debug(
31
      `creating a new Pool(host:${host}, database:${database}, user:${user}, ssl:${JSON.stringify(ssl)} statement_timeout:${STATEMENT_TIMEOUT_MS}ms)`,
32
    );
33
    pool = new Pool({
34
      password: dbPassword(),
35
      user,
36
      host,
37
      database,
38
      statement_timeout: STATEMENT_TIMEOUT_MS, // fixes https://github.com/sagemathinc/cocalc/issues/6014
39
      // the test suite assumes small pool, or there will be random failures sometimes (?)
40
      max: process.env.PGDATABASE == TEST ? 2 : undefined,
41
      ssl,
42
    });
43

44
    pool.on("error", (err: Error) => {
45
      L.debug("WARNING: Unexpected error on idle client in PG pool", {
46
        err: err.message,
47
        stack: err.stack,
48
      });
49
    });
50
    const end = pool.end.bind(pool);
51
    pool.end = async () => {
52
      pool = undefined;
53
      end();
54
    };
55
  }
56
  return pool;
57
}
58

59
// CRITICAL -- the caller *must* call client.release on the client
60
// that is returned from getTransactionClient()!  E.g., for unit testing
61
// if you don't do this  you exhaust the limit of 2 on the pool size,
62
// (see above) and everything hangs!
63
export async function getTransactionClient(): Promise<PoolClient> {
64
  const client = await getPoolClient();
65
  try {
66
    await client.query("BEGIN");
67
  } catch (err) {
68
    await client.query("ROLLBACK");
69
    client.release();
70
    throw err;
71
  }
72
  return client;
73
}
74

75
export async function getPoolClient(): Promise<PoolClient> {
76
  const pool = await getPool();
77
  return await pool.connect();
78
}
79

80
export function getClient(): Client {
81
  return new Client({ password: dbPassword(), user, host, database, ssl });
82
}
83

84
// This is used for testing.  It ensures the schema is loaded and
85
// test database is defined.
86

87
// Call this with {reset:true} to reset the ephemeral
88
// database to a clean state with the schema loaded.
89
// You *can't* just initEphemeralDatabase({reset:true}) in the pre-amble
90
// of jest tests though, since all the tests are running in parallel, and
91
// they would mess up each other's state...
92
const TEST = "smc_ephemeral_testing_database";
93
export async function initEphemeralDatabase({
94
  reset,
95
}: { reset?: boolean } = {}) {
96
  if (database != TEST) {
97
    throw Error(
98
      `You can't use initEphemeralDatabase() and test using the database if the env variabe PGDATABASE is not set to ${TEST}!`,
99
    );
100
  }
101
  const db = new Pool({
102
    password: dbPassword(),
103
    user,
104
    host,
105
    database: "smc",
106
    statement_timeout: STATEMENT_TIMEOUT_MS,
107
    ssl,
108
  });
109
  db.on("error", (err: Error) => {
110
    L.debug("WARNING: Unexpected error on idle client in PG pool", {
111
      err: err.message,
112
      stack: err.stack,
113
    });
114
  });
115
  const { rows } = await db.query(
116
    "SELECT COUNT(*) AS count FROM pg_catalog.pg_database WHERE datname = $1",
117
    [TEST],
118
  );
119
  //await db.query(`DROP DATABASE IF EXISTS ${TEST}`);
120
  const databaseExists = rows[0].count > 0;
121
  if (!databaseExists) {
122
    await db.query(`CREATE DATABASE ${TEST}`);
123
  }
124
  await db.end();
125
  // sync the schema
126
  await syncSchema();
127
  if (databaseExists && reset) {
128
    // Drop all data from all tables for a clean slate.
129
    // Unfortunately, this can take a little while.
130
    await dropAllData();
131
  }
132
}
133

134
async function dropAllData() {
135
  const pool = getPool();
136
  pool.on("error", (err: Error) => {
137
    L.debug("WARNING: Unexpected error on idle client in PG pool", {
138
      err: err.message,
139
      stack: err.stack,
140
    });
141
  });
142
  if (pool?.["options"]?.database != TEST) {
143
    // safety check!
144
    throw Error(
145
      `You can't use dropAllData() if the env variabe PGDATABASE is not set to ${TEST}!`,
146
    );
147
  }
148
  const client = await pool.connect();
149

150
  try {
151
    // Get all table names
152
    const result = await client.query(
153
      "SELECT tablename FROM pg_tables WHERE schemaname='public'",
154
    );
155
    const tableNames = result.rows.map((row) => row.tablename);
156
    await client.query(`TRUNCATE ${tableNames.join(",")}`);
157
  } catch (err) {
158
    throw err;
159
  } finally {
160
    client.release();
161
  }
162
}
163

164