pg-meta-tables.ts277 lines · main
| 1 | import { z } from 'zod' |
| 2 | |
| 3 | import { DEFAULT_SYSTEM_SCHEMAS } from './constants' |
| 4 | import { coalesceRowsToArray, filterByList } from './helpers' |
| 5 | import { |
| 6 | ident, |
| 7 | joinSqlFragments, |
| 8 | keyword, |
| 9 | literal, |
| 10 | safeSql, |
| 11 | type SafeSqlFragment, |
| 12 | } from './pg-format' |
| 13 | import { pgColumnArrayZod } from './pg-meta-columns' |
| 14 | import { COLUMNS_SQL } from './sql/columns' |
| 15 | import { TABLES_SQL } from './sql/tables' |
| 16 | |
| 17 | const pgTablePrimaryKeyZod = z.object({ |
| 18 | table_id: z.number(), |
| 19 | name: z.string(), |
| 20 | schema: z.string(), |
| 21 | table_name: z.string(), |
| 22 | }) |
| 23 | |
| 24 | const pgTableRelationshipZod = z.object({ |
| 25 | id: z.number(), |
| 26 | constraint_name: z.string(), |
| 27 | source_schema: z.string(), |
| 28 | source_table_name: z.string(), |
| 29 | source_column_name: z.string(), |
| 30 | target_table_schema: z.string(), |
| 31 | target_table_name: z.string(), |
| 32 | target_column_name: z.string(), |
| 33 | }) |
| 34 | |
| 35 | const pgTableZod = z.object({ |
| 36 | id: z.number(), |
| 37 | schema: z.string(), |
| 38 | name: z.string(), |
| 39 | rls_enabled: z.boolean(), |
| 40 | rls_forced: z.boolean(), |
| 41 | replica_identity: z.enum(['DEFAULT', 'INDEX', 'FULL', 'NOTHING']), |
| 42 | bytes: z.number(), |
| 43 | size: z.string(), |
| 44 | live_rows_estimate: z.number(), |
| 45 | dead_rows_estimate: z.number(), |
| 46 | comment: z.string().nullable(), |
| 47 | primary_keys: z.array(pgTablePrimaryKeyZod), |
| 48 | relationships: z.array(pgTableRelationshipZod), |
| 49 | columns: pgColumnArrayZod.optional(), |
| 50 | }) |
| 51 | |
| 52 | const pgTableArrayZod = z.array(pgTableZod) |
| 53 | |
| 54 | export type PGTable = z.infer<typeof pgTableZod> |
| 55 | export type PGTablePrimaryKey = z.infer<typeof pgTablePrimaryKeyZod> |
| 56 | export type PGTableRelationship = z.infer<typeof pgTableRelationshipZod> |
| 57 | |
| 58 | type TableWithoutColumns = Omit<PGTable, 'columns'> |
| 59 | type TableWithColumns = PGTable |
| 60 | |
| 61 | type TableBasedOnIncludeColumns<T extends boolean | undefined> = T extends true |
| 62 | ? TableWithColumns |
| 63 | : TableWithoutColumns |
| 64 | |
| 65 | type TableIdentifier = Pick<PGTable, 'id'> | Pick<PGTable, 'name' | 'schema'> |
| 66 | |
| 67 | function getIdentifierWhereClause(identifier: TableIdentifier): SafeSqlFragment { |
| 68 | if ('id' in identifier && identifier.id) { |
| 69 | return safeSql`${ident('id')} = ${literal(identifier.id)}` |
| 70 | } |
| 71 | if ('name' in identifier && identifier.name && identifier.schema) { |
| 72 | return safeSql`${ident('name')} = ${literal(identifier.name)} and ${ident('schema')} = ${literal(identifier.schema)}` |
| 73 | } |
| 74 | throw new Error('Must provide either id or name and schema') |
| 75 | } |
| 76 | |
| 77 | function list<T extends boolean | undefined = true>( |
| 78 | { |
| 79 | includeSystemSchemas = false, |
| 80 | includedSchemas, |
| 81 | excludedSchemas, |
| 82 | limit, |
| 83 | offset, |
| 84 | includeColumns = true as T, |
| 85 | }: { |
| 86 | includeSystemSchemas?: boolean |
| 87 | includedSchemas?: string[] |
| 88 | excludedSchemas?: string[] |
| 89 | limit?: number |
| 90 | offset?: number |
| 91 | includeColumns?: T |
| 92 | } = {} as any |
| 93 | ): { |
| 94 | sql: SafeSqlFragment |
| 95 | zod: z.ZodType<TableBasedOnIncludeColumns<T>[]> |
| 96 | } { |
| 97 | let sql = generateEnrichedTablesSql({ includeColumns }) |
| 98 | const filter = filterByList( |
| 99 | includedSchemas, |
| 100 | excludedSchemas, |
| 101 | !includeSystemSchemas ? DEFAULT_SYSTEM_SCHEMAS : undefined |
| 102 | ) |
| 103 | if (filter) { |
| 104 | sql = safeSql`${sql} where schema ${filter}` |
| 105 | } |
| 106 | if (limit) { |
| 107 | sql = safeSql`${sql} limit ${literal(limit)}` |
| 108 | } |
| 109 | if (offset) { |
| 110 | sql = safeSql`${sql} offset ${literal(offset)}` |
| 111 | } |
| 112 | return { |
| 113 | sql, |
| 114 | zod: pgTableArrayZod, |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | function retrieve(identifier: TableIdentifier): { |
| 119 | sql: SafeSqlFragment |
| 120 | zod: z.ZodType<TableWithColumns> |
| 121 | } { |
| 122 | let whereClause = getIdentifierWhereClause(identifier) |
| 123 | |
| 124 | const sql = safeSql`${generateEnrichedTablesSql({ includeColumns: true })} where ${whereClause};` |
| 125 | return { |
| 126 | sql, |
| 127 | zod: pgTableZod, |
| 128 | } |
| 129 | } |
| 130 | |
| 131 | function remove( |
| 132 | table: Pick<PGTable, 'name' | 'schema'>, |
| 133 | { cascade = false } = {} |
| 134 | ): { sql: SafeSqlFragment } { |
| 135 | const sql = safeSql`DROP TABLE ${ident(table.schema)}.${ident(table.name)} ${ |
| 136 | cascade ? safeSql`CASCADE` : safeSql`RESTRICT` |
| 137 | };` |
| 138 | return { sql } |
| 139 | } |
| 140 | |
| 141 | const generateEnrichedTablesSql = ({ includeColumns }: { includeColumns?: boolean }) => safeSql` |
| 142 | with tables as (${TABLES_SQL}) |
| 143 | ${includeColumns ? safeSql`, columns as (${COLUMNS_SQL})` : safeSql``} |
| 144 | select |
| 145 | * |
| 146 | ${includeColumns ? safeSql`, ${coalesceRowsToArray('columns', safeSql`columns.table_id = tables.id`)}` : safeSql``} |
| 147 | from tables` |
| 148 | |
| 149 | type TableCreateParams = { |
| 150 | name: string |
| 151 | schema?: string |
| 152 | comment?: string | null |
| 153 | no_transaction?: boolean |
| 154 | } |
| 155 | |
| 156 | function create({ name, schema = 'public', comment, no_transaction = false }: TableCreateParams): { |
| 157 | sql: SafeSqlFragment |
| 158 | } { |
| 159 | const tableSql = safeSql`CREATE TABLE ${ident(schema)}.${ident(name)} ();` |
| 160 | const commentSql = |
| 161 | comment != undefined |
| 162 | ? safeSql`COMMENT ON TABLE ${ident(schema)}.${ident(name)} IS ${literal(comment)};` |
| 163 | : safeSql`` |
| 164 | |
| 165 | if (no_transaction) { |
| 166 | const sql = safeSql`${tableSql} ${commentSql}` |
| 167 | return { sql } |
| 168 | } |
| 169 | const sql = safeSql`BEGIN; ${tableSql} ${commentSql} COMMIT;` |
| 170 | return { sql } |
| 171 | } |
| 172 | |
| 173 | type TableUpdateParams = { |
| 174 | name?: string |
| 175 | schema?: string |
| 176 | rls_enabled?: boolean |
| 177 | rls_forced?: boolean |
| 178 | replica_identity?: 'DEFAULT' | 'INDEX' | 'FULL' | 'NOTHING' |
| 179 | replica_identity_index?: string |
| 180 | primary_keys?: Array<{ name: string }> |
| 181 | comment?: string | null |
| 182 | } |
| 183 | |
| 184 | function update( |
| 185 | old: Pick<PGTable, 'id' | 'name' | 'schema'>, |
| 186 | { |
| 187 | name, |
| 188 | schema, |
| 189 | rls_enabled, |
| 190 | rls_forced, |
| 191 | replica_identity, |
| 192 | replica_identity_index, |
| 193 | primary_keys, |
| 194 | comment, |
| 195 | }: TableUpdateParams |
| 196 | ): { sql: SafeSqlFragment } { |
| 197 | const alter = safeSql`ALTER TABLE ${ident(old.schema)}.${ident(old.name)}` |
| 198 | const schemaSql = |
| 199 | schema === undefined ? safeSql`` : safeSql`${alter} SET SCHEMA ${ident(schema)};` |
| 200 | let nameSql = safeSql`` |
| 201 | if (name !== undefined && name !== old.name) { |
| 202 | const currentSchema = schema === undefined ? old.schema : schema |
| 203 | nameSql = safeSql`ALTER TABLE ${ident(currentSchema)}.${ident(old.name)} RENAME TO ${ident(name)};` |
| 204 | } |
| 205 | let enableRls = safeSql`` |
| 206 | if (rls_enabled !== undefined) { |
| 207 | const enable = safeSql`${alter} ENABLE ROW LEVEL SECURITY;` |
| 208 | const disable = safeSql`${alter} DISABLE ROW LEVEL SECURITY;` |
| 209 | enableRls = rls_enabled ? enable : disable |
| 210 | } |
| 211 | let forceRls = safeSql`` |
| 212 | if (rls_forced !== undefined) { |
| 213 | const enable = safeSql`${alter} FORCE ROW LEVEL SECURITY;` |
| 214 | const disable = safeSql`${alter} NO FORCE ROW LEVEL SECURITY;` |
| 215 | forceRls = rls_forced ? enable : disable |
| 216 | } |
| 217 | let replicaSql = safeSql`` |
| 218 | if (replica_identity === undefined) { |
| 219 | // skip |
| 220 | } else if (replica_identity === 'INDEX') { |
| 221 | if (!replica_identity_index) { |
| 222 | throw new Error('replica_identity_index is required when replica_identity is INDEX') |
| 223 | } |
| 224 | replicaSql = safeSql`${alter} REPLICA IDENTITY USING INDEX ${ident(replica_identity_index)};` |
| 225 | } else { |
| 226 | replicaSql = safeSql`${alter} REPLICA IDENTITY ${keyword(replica_identity)};` |
| 227 | } |
| 228 | let primaryKeysSql = safeSql`` |
| 229 | if (primary_keys === undefined) { |
| 230 | // skip |
| 231 | } else { |
| 232 | primaryKeysSql = safeSql`${primaryKeysSql} |
| 233 | DO $$ |
| 234 | DECLARE |
| 235 | r record; |
| 236 | BEGIN |
| 237 | SELECT conname |
| 238 | INTO r |
| 239 | FROM pg_constraint |
| 240 | WHERE contype = 'p' AND conrelid = ${literal(old.id)}; |
| 241 | IF r IS NOT NULL THEN |
| 242 | EXECUTE ${literal(`${alter} DROP CONSTRAINT `)} || quote_ident(r.conname); |
| 243 | END IF; |
| 244 | END |
| 245 | $$; |
| 246 | ` |
| 247 | |
| 248 | if (primary_keys.length === 0) { |
| 249 | // skip |
| 250 | } else { |
| 251 | primaryKeysSql = safeSql`${primaryKeysSql} ${alter} ADD PRIMARY KEY (${joinSqlFragments( |
| 252 | primary_keys.map((x) => ident(x.name)), |
| 253 | ',' |
| 254 | )});` |
| 255 | } |
| 256 | } |
| 257 | const commentSql = |
| 258 | comment == undefined |
| 259 | ? safeSql`` |
| 260 | : safeSql`COMMENT ON TABLE ${ident(old.schema)}.${ident(old.name)} IS ${literal(comment)};` |
| 261 | |
| 262 | // nameSql must be last, right below schemaSql |
| 263 | const sql = safeSql` |
| 264 | BEGIN; |
| 265 | ${enableRls} |
| 266 | ${forceRls} |
| 267 | ${replicaSql} |
| 268 | ${primaryKeysSql} |
| 269 | ${commentSql} |
| 270 | ${schemaSql} |
| 271 | ${nameSql} |
| 272 | COMMIT;` |
| 273 | |
| 274 | return { sql } |
| 275 | } |
| 276 | |
| 277 | export { create, list, remove, retrieve, update } |