Kysely Typegen (TS version)

kysely-typegen.ts

import camelCase from 'camelcase';
import decamelize from 'decamelize';
import type { MaterializedViewColumn, TableColumn, ViewColumn } from 'extract-pg-schema';
import extractPGSchema from 'extract-pg-schema';
import { readFile, writeFile } from 'node:fs/promises';
import { join } from 'node:path';
import type { ConnectionConfig } from 'pg';
import pg from 'pg';

const dbRootDir = join(import.meta.dirname, '../');

const { extractSchemas } = extractPGSchema;

type MappedTableViewOrMViewColumn = {
    name: string;
    type: string;
    isPrimaryKey?: boolean;
    isNullable?: boolean;
    isUpdatable: boolean;
    isArray: boolean;
    defaultValue?: string | null;
};

type MappedTableViewOrMView = {
    entityName: string;
    propName: string;
    kind: string;
    columns: MappedTableViewOrMViewColumn[];
};

type MappedEnum = {
    name: string;
    values: string[];
};

type Connection = Pick<ConnectionConfig, 'host' | 'port' | 'user' | 'password' | 'database'>;

export class DbTypegen {
    private preRunSqlFiles: string[];
    private prefix: string;
    private outFile: string;
    private connection: Connection;

    constructor({ connection, preRunSqlFiles, prefix = 'DB' }: { connection: Connection, preRunSqlFiles?: string[], prefix?: string }) {
        this.preRunSqlFiles = preRunSqlFiles ?? [];
        this.prefix = prefix;
        this.outFile = join(dbRootDir, 'src', 'db-types.ts');
        this.connection = connection;
    }

    private toPascalCase(name: string) {
        return camelCase(name, { pascalCase: true, preserveConsecutiveUppercase: true });
    }

    private toCamelCase(name: string) {
        return camelCase(name, { pascalCase: false, preserveConsecutiveUppercase: true });
    }

    private toSnakeCase(name: string) {
        return decamelize(name, { separator: '_' }).toUpperCase();
    }

    private unionize(values: string[]) {
        return values.map(v => `    '${v}'`).join(' |\n');
    }

    private async getPublicSchema() {
        const schemas = await extractSchemas(this.connection);

        if (!('public' in schemas)) {
            throw new Error('No public schema found');
        }

        return schemas.public;
    }

    private async mapSchema() {
        const schema = await this.getPublicSchema();

        return ({
            enums: schema.enums.map(e => ({
                name: this.getEntityName(e.name),
                values: e.values
            })),
            tables: schema.tables.map(t => ({
                entityName: this.getEntityName(t.name),
                propName: this.toCamelCase(t.name),
                kind: 'table',
                columns: t.columns.map(c => this.mapColumn(c))
            })),
            views: schema.views.map(v => ({
                entityName: this.getEntityName(v.name),
                propName: this.toCamelCase(v.name),
                kind: 'view',
                columns: v.columns.map(c => this.mapColumn(c))
            })),
            materializedViews: schema.materializedViews.map(mv => ({
                entityName: this.getEntityName(mv.name),
                propName: this.toCamelCase(mv.name),
                kind: 'mView',
                columns: mv.columns.map(c => this.mapColumn(c))
            }))
        });
    }

    private mapColumn(column: TableColumn | ViewColumn | MaterializedViewColumn) {
        if (!['base', 'enum', 'range'].includes(column.type.kind)) {
            throw new Error(`Unsupported type kind: ${column.type.kind}`);
        }

        if (Array.isArray(column.references) && column.references.length > 1) {
            throw new Error('Multiple references not supported');
        }

        return ({
            name: this.toCamelCase(column.name),
            type: this.getColumnType(column),
            isPrimaryKey: column.isPrimaryKey,
            isNullable: column.isNullable,
            isUpdatable: column.isUpdatable,
            isArray: column.isArray,
            defaultValue: column.defaultValue
        });
    }

    private getColumnType(column: TableColumn | ViewColumn | MaterializedViewColumn) {
        if (column?.references?.length === 1) {
            return this.getEntityName(column.references[0].tableName, column.references[0].columnName);
        }
        if (column.type.kind === 'enum') {
            return this.getEnumNameFromType(column);
        }
        if (column.type.kind === 'base') {
            return this.getBaseName(column);
        }
        if (column.type.kind === 'range') {
            return this.getBaseName(column);
        }
        throw new Error(`Unsupported type kind: ${column.type.kind}`);
    }

    private getEntityName(tableName: string, columnName?: string) {
        return columnName
            ? `${this.prefix}${this.toPascalCase(`${tableName}_${columnName}`)}`
            : `${this.prefix}${this.toPascalCase(tableName)}`;
    }

    private getEnumNameFromType(column: TableColumn | ViewColumn | MaterializedViewColumn) {
        const name = column.type.fullName.split('.').at(1);

        if (!name) {
            throw new Error(`Invalid enum type: ${column.type.fullName}`);
        }

        return column.isArray
            ? `${this.prefix}${this.toPascalCase(name)}[]`
            : `${this.prefix}${this.toPascalCase(name)}`;
    }

    private getBaseName(column: TableColumn | ViewColumn | MaterializedViewColumn) {
        switch (column.type.fullName) {
            case 'pg_catalog.int2': return `number${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.int4': return `number${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.int8': return `number${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.numeric': return `string${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.uuid': return `string${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.text': return `string${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.char': return `string${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.varchar': return `string${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.bool': return `boolean${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.timestamp': return `Date${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.timestamptz': return `Date${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.date': return `Date${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.timetz': return `string${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.json': return `object${column.isArray ? '[]' : ''}`;
            case 'pg_catalog.jsonb': return `object${column.isArray ? '[]' : ''}`;
            default:
                throw new Error(`Unsupported base type: ${column.type.fullName}`);
        }
    }

    // REMAPPED
    private getEnumLine(enumColumn: MappedEnum) {
        const values = this.unionize(enumColumn.values);
        return `export type ${enumColumn.name} = \n${values};`;
    }

    // REMAPPED
    private getIdentifierLine(table: MappedTableViewOrMView) {
        const primaryKeyColumns = table.columns.filter(c => c.isPrimaryKey);

        if (primaryKeyColumns.length === 0) {
            return '';
        }

        return primaryKeyColumns.map(column => {
            const name = this.toPascalCase(`${table.entityName}_${column.name}`);
            return `export type ${name} = ${column.type};`;
        }).join('\n');
    }

    // REMAPPED
    private getColumnTypeLine(table: MappedTableViewOrMView, column: MappedTableViewOrMViewColumn) {
        if (table.kind === 'table') {
            const selectable = column.isNullable ? `${column.type} | null` : column.type;
            const insertable = (column.isNullable || column.defaultValue) ? `${column.type} | null` : column.type;
            const updateable = column.isUpdatable ? `${column.type} | null` : 'never';

            return `ColumnType<${selectable}, ${insertable}, ${updateable}>`;
        }

        const selectable = column.isNullable ? `${column.type} | null` : column.type;
        return `ColumnType<${selectable}, never, never>`;
    }

    // REMAPPED
    private getTableLine(tableLike: MappedTableViewOrMView) {
        const props = [];

        for (const column of tableLike.columns) {
            props.push(`    ${column.name}: ${this.getColumnTypeLine(tableLike, column)};`);
        }

        return `export interface ${tableLike.entityName} {\n${props.join('\n')}\n};`;
    }

    // REMAPPED
    private getKyselyExportsLine(tableLike: MappedTableViewOrMView) {
        if (tableLike.kind === 'table') {
            return [
                `export type ${tableLike.entityName}Selectable = Selectable<${tableLike.entityName}>;`,
                `export type ${tableLike.entityName}Insertable = Insertable<${tableLike.entityName}>;`,
                `export type ${tableLike.entityName}Updateable = Updateable<${tableLike.entityName}>;`,
            ].join('\n');
        }

        return `export type ${tableLike.entityName}Selectable = Selectable<${tableLike.entityName}>;`;
    }

    private getDatabaseLine(tables: MappedTableViewOrMView[], views: MappedTableViewOrMView[], materializedViews: MappedTableViewOrMView[]) {
        const props = [];

        for (const table of tables) {
            props.push(`    ${table.propName}: ${table.entityName};`);
        }

        for (const view of views) {
            props.push(`    ${view.propName}: ${view.entityName};`);
        }

        for (const mView of materializedViews) {
            props.push(`    ${mView.propName}: ${mView.entityName};`);
        }

        return `export interface DB {\n${props.join('\n')}\n};`;
    }

    private getDbColumnNamesLine(tables: MappedTableViewOrMView[], views: MappedTableViewOrMView[], materializedViews: MappedTableViewOrMView[]) {
        /**
         * here it probably doesn't make sense to export views and
         * materielized views, since they are not updatable
         */
        const props = [];

        for (const table of tables) {
            props.push(`    ${table.propName}: [${table.columns.map(c => `'${c.name}'`).join(', ')}],`);
        }
        for (const view of views) {
            props.push(`    ${view.propName}: [${view.columns.map(c => `'${c.name}'`).join(', ')}],`);
        }
        for (const mView of materializedViews) {
            props.push(`    ${mView.propName}: [${mView.columns.map(c => `'${c.name}'`).join(', ')}],`);
        }

        return `export const dbColumnNames: DBColumnNames = {\n${props.join('\n')}\n} as const;`;
    }

    private getEnumValuesLine(enums: MappedEnum[]) {
        const props = [];

        for (const pgEnum of enums) {
            const name = this.toSnakeCase(pgEnum.name).replace(/DB_/g, 'DB_TYPE_');
            props.push(`export const ${name} = [${pgEnum.values.map(v => `'${v}'`).join(', ')}] as const;`);
        }

        return props.join('\n');
    }


    public async generate() {
        if (this.preRunSqlFiles.length > 0) {
            const { Client } = pg;
            const client = new Client({
                ...this.connection,
            });
            await client.connect();
            const files = await Promise.all(this.preRunSqlFiles.map(file => readFile(file, 'utf-8')));
            const sqlStatement = files.join('\n\n');
            await client.query(sqlStatement);
            await client.end();
        }

        const { enums, tables, views, materializedViews } = await this.mapSchema();
        const header = `/* File generated automatically, do not edit. */`;
        const imports = `import type { ColumnType, Selectable, Insertable, Updateable } from 'kysely';`;
        const exports = `export type DBColumnNames = { readonly [K in keyof DB]: ReadonlyArray<keyof DB[K]>; };`;

        const result: {
            enums: string[];
            tables: string[];
            views: string[];
            mViews: string[];
            database: string;
            columnNames: string;
            enumValues: string;
        } = {
            enums: [],
            tables: [],
            views: [],
            mViews: [],
            database: '',
            columnNames: '',
            enumValues: ''
        };

        for (const pgEnum of enums) {
            result.enums.push(this.getEnumLine(pgEnum));
        }

        for (const table of tables) {
            result.tables.push(this.getIdentifierLine(table));
            result.tables.push(this.getTableLine(table));
            result.tables.push(this.getKyselyExportsLine(table));
        }

        for (const view of views) {
            result.views.push(this.getTableLine(view));
            result.views.push(this.getKyselyExportsLine(view));
        }

        for (const mView of materializedViews) {
            result.mViews.push(this.getTableLine(mView));
            result.mViews.push(this.getKyselyExportsLine(mView));
        }

        result.database = this.getDatabaseLine(tables, views, materializedViews);
        result.columnNames = this.getDbColumnNamesLine(tables, views, materializedViews);
        result.enumValues = this.getEnumValuesLine(enums);

        const output = [
            header,
            imports,
            ...result.enums,
            ...result.tables,
            ...result.views,
            ...result.mViews,
            exports,
            result.database,
            result.columnNames,
            result.enumValues
        ].filter(Boolean).join('\n\n');

        await writeFile(this.outFile, output, { encoding: 'utf-8' });
    }
}