pg-meta-tables.ts277 lines · main
1import { z } from 'zod'
2
3import { DEFAULT_SYSTEM_SCHEMAS } from './constants'
4import { coalesceRowsToArray, filterByList } from './helpers'
5import {
6 ident,
7 joinSqlFragments,
8 keyword,
9 literal,
10 safeSql,
11 type SafeSqlFragment,
12} from './pg-format'
13import { pgColumnArrayZod } from './pg-meta-columns'
14import { COLUMNS_SQL } from './sql/columns'
15import { TABLES_SQL } from './sql/tables'
16
17const pgTablePrimaryKeyZod = z.object({
18 table_id: z.number(),
19 name: z.string(),
20 schema: z.string(),
21 table_name: z.string(),
22})
23
24const 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
35const 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
52const pgTableArrayZod = z.array(pgTableZod)
53
54export type PGTable = z.infer<typeof pgTableZod>
55export type PGTablePrimaryKey = z.infer<typeof pgTablePrimaryKeyZod>
56export type PGTableRelationship = z.infer<typeof pgTableRelationshipZod>
57
58type TableWithoutColumns = Omit<PGTable, 'columns'>
59type TableWithColumns = PGTable
60
61type TableBasedOnIncludeColumns<T extends boolean | undefined> = T extends true
62 ? TableWithColumns
63 : TableWithoutColumns
64
65type TableIdentifier = Pick<PGTable, 'id'> | Pick<PGTable, 'name' | 'schema'>
66
67function 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
77function 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
118function 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
131function 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
141const 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
149type TableCreateParams = {
150 name: string
151 schema?: string
152 comment?: string | null
153 no_transaction?: boolean
154}
155
156function 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
173type 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
184function 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}
233DO $$
234DECLARE
235 r record;
236BEGIN
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;
244END
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`
264BEGIN;
265 ${enableRls}
266 ${forceRls}
267 ${replicaSql}
268 ${primaryKeysSql}
269 ${commentSql}
270 ${schemaSql}
271 ${nameSql}
272COMMIT;`
273
274 return { sql }
275}
276
277export { create, list, remove, retrieve, update }