ts-prorm-orm 1.1.0

package/dist/dialects/yugabytedb/index.js

@@ -0,0 +1,424 @@
+"use strict";
+/**
+ * YugabyteDB dialect implementation for the TypeScript ORM
+ *
+ * YugabyteDB's YSQL API reuses the PostgreSQL query layer (it is literally
+ * forked from the PostgreSQL source), so it speaks the PostgreSQL wire
+ * protocol and is largely SQL-compatible with PostgreSQL. This dialect
+ * therefore `extends PostgresDialect` and reuses the same `pg` driver
+ * (`library === 'pg'`). Notable differences from the Postgres dialect that
+ * this class layers on top:
+ *
+ * - Default port is 5433 (not Postgres's 5432). YugabyteDB's YSQL server
+ *   listens on 5433 by default; 5432 is used by other services in some
+ *   deployments.
+ * - `SPLIT INTO n TABLETS` / `SPLIT AT VALUES (...)`: because YugabyteDB
+ *   shards ("tablets") tables across nodes, `CREATE TABLE`/`CREATE INDEX`
+ *   accepts pre-splitting clauses to control the initial tablet layout.
+ *   `SPLIT INTO n TABLETS` presplits a hash-sharded table/index into `n`
+ *   tablets; `SPLIT AT VALUES ((a), (b), ...)` presplits a range-sharded
+ *   table/index at the given range boundaries. See
+ *   https://docs.yugabyte.com/preview/api/ysql/the-sql-language/statements/ddl_create_table/#split-into
+ * - Colocated tables: small tables can be co-located onto a single tablet
+ *   with `WITH (colocated = true)` (older syntax) / `WITH (colocation = true)`
+ *   to avoid the overhead of a tablet-per-table. See
+ *   https://docs.yugabyte.com/preview/architecture/docdb-sharding/colocated-tables/
+ * - Primary-key column ordering: unlike Postgres, a YugabyteDB primary key
+ *   column carries a sharding/sort directive — `HASH` (hash-sharded, the
+ *   default for the first PK column) or `ASC`/`DESC` (range-sharded). This is
+ *   emitted as `PRIMARY KEY (id HASH, created_at DESC)`. See
+ *   https://docs.yugabyte.com/preview/architecture/docdb-sharding/sharding/
+ * - `CREATE TABLEGROUP`: YugabyteDB groups tables that should share tablets
+ *   into a tablegroup; `createTableGroup()` is a small helper that emits the
+ *   DDL and `buildCreateTableGroupSQL()` builds it synchronously.
+ *
+ * Behavioral notes shared with PostgreSQL (inherited, called out here so
+ * callers know the caveats):
+ * - Because YSQL is the PostgreSQL wire protocol, connection handling,
+ *   parameter binding, transactions, and most query building are inherited
+ *   unchanged from {@link PostgresDialect}.
+ * - `CREATE INDEX ... CONCURRENTLY` behaves differently: in YugabyteDB online
+ *   index backfill is the default mechanism and `CONCURRENTLY` has its own
+ *   semantics/restrictions (and cannot run inside a transaction block), so it
+ *   should not be assumed to behave identically to vanilla PostgreSQL.
+ * - Some PostgreSQL system-catalog (`pg_catalog`) internals differ: YugabyteDB
+ *   emulates the catalogs for compatibility but a handful of columns/values
+ *   reflect its distributed storage rather than a single-node heap.
+ */
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.YugabyteDBDialect = void 0;
+exports.createYugabyteDBDialect = createYugabyteDBDialect;
+const index_1 = require("../postgres/index");
+/**
+ * YugabyteDB dialect. Postgres-wire-compatible: reuses the `pg` driver and the
+ * PostgreSQL query-building layer, layering on YugabyteDB's distributed-table
+ * DDL (tablet splitting, colocation, hash/range primary keys, tablegroups).
+ */
+class YugabyteDBDialect extends index_1.PostgresDialect {
+    constructor(config) {
+        // YugabyteDB's YSQL server defaults to port 5433 (not Postgres's 5432).
+        // An explicit `config.port` still wins because it's spread last.
+        super({ port: 5433, ...config });
+        // Override the inherited `readonly name = 'postgres'`. YugabyteDB reuses the
+        // pg driver, so `library` ('pg') is inherited unchanged.
+        this.name = 'yugabytedb';
+    }
+    /**
+     * Quote a `SPLIT AT VALUES` / `PRIMARY KEY` scalar for inlining into DDL:
+     * numbers pass through, everything else is single-quoted and escaped.
+     */
+    quoteSplitValue(value) {
+        if (typeof value === 'number') {
+            return String(value);
+        }
+        return `'${String(value).replace(/'/g, "''")}'`;
+    }
+    /**
+     * Build the trailing YugabyteDB-specific clauses for a `CREATE TABLE`
+     * statement (colocation, tablegroup, tablet splitting), in the order
+     * YugabyteDB expects them. Returns an empty string when none apply.
+     */
+    buildYugabyteTableClauses(options) {
+        if (!options) {
+            return '';
+        }
+        const parts = [];
+        // Colocation: `WITH (colocated = true)` / `WITH (colocation = true)`.
+        if (options.colocation !== undefined) {
+            parts.push(`WITH (colocation = ${options.colocation ? 'true' : 'false'})`);
+        }
+        else if (options.colocated !== undefined) {
+            parts.push(`WITH (colocated = ${options.colocated ? 'true' : 'false'})`);
+        }
+        // Tablegroup placement.
+        if (options.tablegroup) {
+            parts.push(`TABLEGROUP ${this.escapeId(options.tablegroup)}`);
+        }
+        // Presplit a hash-sharded table into N tablets.
+        if (options.splitIntoTablets !== undefined) {
+            parts.push(`SPLIT INTO ${options.splitIntoTablets} TABLETS`);
+        }
+        // Presplit a range-sharded table at explicit boundaries.
+        if (options.splitAtValues && options.splitAtValues.length > 0) {
+            const points = options.splitAtValues
+                .map((row) => `(${row.map((v) => this.quoteSplitValue(v)).join(', ')})`)
+                .join(', ');
+            parts.push(`SPLIT AT VALUES (${points})`);
+        }
+        return parts.length > 0 ? ` ${parts.join(' ')}` : '';
+    }
+    /**
+     * Format a column `DEFAULT` value for inlining into DDL. Mirrors the base
+     * PostgreSQL dialect's behavior (which lives in a private helper we cannot
+     * call from here): numbers/booleans pass through; strings that look like a
+     * SQL function call (`gen_random_uuid()`), a `CURRENT_*` keyword, or `NULL`
+     * are emitted raw rather than quoted; `DataType` instances that expose
+     * `toDefaultValue()` (e.g. `DataTypes.NOW`) delegate to it; everything else
+     * is single-quoted and escaped.
+     */
+    formatDefaultValue(value) {
+        if (value === null) {
+            return 'NULL';
+        }
+        if (typeof value === 'number' || typeof value === 'boolean') {
+            return String(value);
+        }
+        if (typeof value === 'string') {
+            if (value.toUpperCase().includes('CURRENT_') ||
+                value.toUpperCase() === 'NULL' ||
+                /^[a-zA-Z_][a-zA-Z0-9_.]*\([^)]*\)$/.test(value.trim())) {
+                return value;
+            }
+            return `'${value.replace(/'/g, "''")}'`;
+        }
+        if (typeof value === 'object' &&
+            value !== null &&
+            typeof value.toDefaultValue === 'function') {
+            return value.toDefaultValue();
+        }
+        return `'${String(value).replace(/'/g, "''")}'`;
+    }
+    /**
+     * Emit a column-level `REFERENCES` (foreign key) clause, if the column
+     * carries one. Handles composite reference targets.
+     */
+    buildColumnReferencesSql(definition) {
+        if (!definition.references) {
+            return '';
+        }
+        const refField = definition.references.field;
+        const refFieldSql = Array.isArray(refField)
+            ? `(${refField.map((f) => this.escapeId(f)).join(', ')})`
+            : `(${this.escapeId(refField)})`;
+        let sql = ` REFERENCES ${this.escapeId(definition.references.table)}${refFieldSql}`;
+        if (definition.references.onDelete) {
+            sql += ` ON DELETE ${definition.references.onDelete}`;
+        }
+        if (definition.references.onUpdate) {
+            sql += ` ON UPDATE ${definition.references.onUpdate}`;
+        }
+        return sql;
+    }
+    /**
+     * Build a single column definition for a `CREATE TABLE` statement. Kept
+     * self-contained (rather than reusing the base class's private helper) so
+     * this dialect can synchronously assemble DDL. Column-level `primaryKey` is
+     * intentionally NOT emitted here — YugabyteDB primary keys are emitted as a
+     * table-level `PRIMARY KEY (...)` (see {@link buildPrimaryKeyClause}) so the
+     * per-column `HASH`/`ASC`/`DESC` directive can be attached.
+     */
+    buildColumnDefinitionSql(columnName, definition) {
+        const dt = definition.type;
+        const isSerial = definition.autoIncrement && (dt.key === 'INTEGER' || !dt.key);
+        let sql = `${this.escapeId(columnName)} ${isSerial ? 'SERIAL' : this.getDataTypeSql(definition.type)}`;
+        if (definition.allowNull === false) {
+            sql += ' NOT NULL';
+        }
+        if (definition.defaultValue !== undefined) {
+            sql += ` DEFAULT ${this.formatDefaultValue(definition.defaultValue)}`;
+        }
+        if (definition.unique) {
+            sql += ' UNIQUE';
+        }
+        sql += this.buildColumnReferencesSql(definition);
+        return sql;
+    }
+    /**
+     * Build a table-level constraint clause (`PRIMARY KEY`/`UNIQUE`/
+     * `FOREIGN KEY`/`CHECK`) for inlining into a `CREATE TABLE` statement.
+     * Mirrors the base PostgreSQL dialect's private helper (which we cannot call
+     * from here). Returns null for an unrecognized/empty constraint.
+     */
+    buildYbConstraintSql(constraint) {
+        const name = constraint.name ? `${this.escapeId(constraint.name)} ` : '';
+        const fields = constraint.fields?.map((f) => this.escapeId(f)).join(', ') || '';
+        switch (constraint.type) {
+            case 'PRIMARY KEY':
+                return `${name}PRIMARY KEY (${fields})`;
+            case 'UNIQUE':
+                return `${name}UNIQUE (${fields})`;
+            case 'FOREIGN KEY': {
+                if (!constraint.references)
+                    return null;
+                const refField = constraint.references.field;
+                const refFieldSql = Array.isArray(refField)
+                    ? `(${refField.map((f) => this.escapeId(f)).join(', ')})`
+                    : `(${this.escapeId(refField)})`;
+                let fkSql = `${name}FOREIGN KEY (${fields}) REFERENCES ${this.escapeId(constraint.references.table)}${refFieldSql}`;
+                if (constraint.references.onDelete) {
+                    fkSql += ` ON DELETE ${constraint.references.onDelete}`;
+                }
+                if (constraint.references.onUpdate) {
+                    fkSql += ` ON UPDATE ${constraint.references.onUpdate}`;
+                }
+                return fkSql;
+            }
+            case 'CHECK':
+                if (!constraint.check)
+                    return null;
+                return `${name}CHECK (${constraint.check})`;
+            default:
+                return null;
+        }
+    }
+    /**
+     * Build the table-level `PRIMARY KEY (col HASH, col2 ASC, ...)` clause from
+     * either the explicit {@link YugabyteTableOptions.primaryKey} directive list
+     * or, failing that, any columns flagged with column-level `primaryKey: true`
+     * (which get no explicit sharding directive). Returns null when there is no
+     * primary key to emit.
+     */
+    buildPrimaryKeyClause(columns, options) {
+        if (options?.primaryKey && options.primaryKey.length > 0) {
+            const cols = options.primaryKey
+                .map((pk) => `${this.escapeId(pk.name)}${pk.order ? ` ${pk.order}` : ''}`)
+                .join(', ');
+            return `PRIMARY KEY (${cols})`;
+        }
+        const pkColumns = Object.entries(columns)
+            .filter(([, def]) => def.primaryKey)
+            .map(([name]) => this.escapeId(name));
+        if (pkColumns.length > 0) {
+            return `PRIMARY KEY (${pkColumns.join(', ')})`;
+        }
+        return null;
+    }
+    /**
+     * Synchronously build a YugabyteDB `CREATE TABLE` statement, including the
+     * distributed-table clauses (hash/range primary key, colocation,
+     * tablegroup, tablet splitting). Used by {@link createTable} and directly
+     * testable without a live connection.
+     */
+    buildCreateTableSQL(tableName, columns, options) {
+        const defs = [];
+        const hasTablePk = !!(options?.primaryKey && options.primaryKey.length > 0);
+        for (const [columnName, definition] of Object.entries(columns)) {
+            // When a table-level PK directive is supplied, strip any column-level
+            // primaryKey flag so we don't emit two PRIMARY KEY clauses.
+            const def = hasTablePk ? { ...definition, primaryKey: false } : definition;
+            defs.push(this.buildColumnDefinitionSql(columnName, def));
+        }
+        const pkClause = this.buildPrimaryKeyClause(columns, options);
+        if (pkClause) {
+            defs.push(pkClause);
+        }
+        // Table-level constraints (FOREIGN KEY / UNIQUE / CHECK / PRIMARY KEY).
+        // A PRIMARY KEY passed here is skipped when a YugabyteDB PK directive
+        // (column-level or `options.primaryKey`) already produced a PK clause, to
+        // avoid emitting two PRIMARY KEY definitions.
+        if (options?.constraints) {
+            for (const constraint of options.constraints) {
+                if (constraint.type === 'PRIMARY KEY' && pkClause) {
+                    continue;
+                }
+                const constraintSql = this.buildYbConstraintSql(constraint);
+                if (constraintSql) {
+                    defs.push(constraintSql);
+                }
+            }
+        }
+        let sql = 'CREATE TABLE';
+        if (options?.ifNotExists) {
+            sql += ' IF NOT EXISTS';
+        }
+        sql += ` ${this.escapeId(tableName)} (${defs.join(', ')})`;
+        sql += this.buildYugabyteTableClauses(options);
+        return sql;
+    }
+    /**
+     * Create a new table. Overrides {@link PostgresDialect.createTable} so the
+     * YugabyteDB-specific clauses (tablet splitting, colocation, hash/range
+     * primary keys, tablegroup) are emitted. Falls back to the inherited
+     * Postgres implementation when no YugabyteDB-specific option is present, so
+     * indexes/comments/uniqueKeys handling is preserved for the common case.
+     */
+    async createTable(tableName, columns, options) {
+        const usesYugabyteFeatures = !!options &&
+            (options.splitIntoTablets !== undefined ||
+                (options.splitAtValues && options.splitAtValues.length > 0) ||
+                options.colocated !== undefined ||
+                options.colocation !== undefined ||
+                options.tablegroup !== undefined ||
+                (options.primaryKey && options.primaryKey.length > 0));
+        if (!usesYugabyteFeatures) {
+            return super.createTable(tableName, columns, options);
+        }
+        await this.query(this.buildCreateTableSQL(tableName, columns, options));
+        // Post-creation steps. These mirror PostgresDialect.createTable so that
+        // uniqueKeys/comment/initialAutoIncrement/indexes are still honored when a
+        // YugabyteDB-specific clause forced us down the custom DDL path (they are
+        // emitted as separate statements, exactly as the base dialect does).
+        if (options) {
+            await this.applyPostCreateTableSteps(tableName, columns, options);
+        }
+    }
+    /**
+     * Emit the follow-up statements PostgresDialect runs after `CREATE TABLE`
+     * (unique keys, table comment, sequence start value, secondary indexes).
+     * Kept here because the base dialect performs them inline in its own
+     * `createTable`, which the YugabyteDB custom-DDL path bypasses.
+     */
+    async applyPostCreateTableSteps(tableName, columns, options) {
+        // Unique keys — supports both the array format and the legacy object map.
+        if (options.uniqueKeys) {
+            if (Array.isArray(options.uniqueKeys)) {
+                for (const uk of options.uniqueKeys) {
+                    const constraintName = uk.name || `${tableName}_${uk.fields.join('_')}_key`;
+                    await this.query(`ALTER TABLE ${this.escapeId(tableName)} ADD CONSTRAINT ${this.escapeId(constraintName)} UNIQUE (${uk.fields.map((f) => this.escapeId(f)).join(', ')})`);
+                }
+            }
+            else {
+                for (const [keyName, fields] of Object.entries(options.uniqueKeys)) {
+                    const constraintName = keyName || `${tableName}_${fields.join('_')}_key`;
+                    await this.query(`ALTER TABLE ${this.escapeId(tableName)} ADD CONSTRAINT ${this.escapeId(constraintName)} UNIQUE (${fields.map((f) => this.escapeId(f)).join(', ')})`);
+                }
+            }
+        }
+        // Table comment.
+        if (options.comment) {
+            await this.query(`COMMENT ON TABLE ${this.escapeId(tableName)} IS '${options.comment.replace(/'/g, "''")}'`);
+        }
+        // Sequence start value for the auto-increment (SERIAL) column.
+        if (options.initialAutoIncrement) {
+            for (const [columnName, definition] of Object.entries(columns)) {
+                if (definition.autoIncrement) {
+                    const sequenceName = `${tableName}_${columnName}_seq`;
+                    await this.query(`ALTER SEQUENCE ${this.escapeId(sequenceName)} START WITH ${options.initialAutoIncrement}`);
+                }
+            }
+        }
+        // Secondary indexes.
+        if (options.indexes) {
+            for (const index of options.indexes) {
+                await this.addIndex(tableName, index.name || `idx_${tableName}_${index.fields.join('_')}`, index.fields, {
+                    unique: index.unique,
+                    type: index.type,
+                    using: index.using,
+                    where: index.where,
+                    expression: index.expression,
+                    include: index.include,
+                });
+            }
+        }
+    }
+    /**
+     * Add an index to a table.
+     *
+     * Overrides {@link PostgresDialect.addIndex} only to account for
+     * YugabyteDB's `CREATE INDEX ... CONCURRENTLY` divergence: in YugabyteDB
+     * index backfill is already online by default, and `CONCURRENTLY` carries
+     * different semantics/restrictions (notably it cannot run inside a
+     * transaction block). Like the CockroachDB dialect, a requested
+     * `concurrently` flag is ignored (with a warning) rather than emitted, so
+     * the generated SQL stays safe to run in the common (transactional) case.
+     * All other index-building behavior is inherited unchanged.
+     */
+    async addIndex(tableName, indexName, fields = [], options) {
+        if (options?.concurrently) {
+            console.warn('YugabyteDB: CREATE INDEX CONCURRENTLY is online by default and has ' +
+                'distinct semantics (it cannot run inside a transaction block). The ' +
+                '`concurrently` option will be ignored.');
+        }
+        return super.addIndex(tableName, indexName, fields, options);
+    }
+    /**
+     * Build a `CREATE TABLEGROUP` statement synchronously.
+     */
+    buildCreateTableGroupSQL(name, options) {
+        let sql = 'CREATE TABLEGROUP';
+        if (options?.ifNotExists) {
+            sql += ' IF NOT EXISTS';
+        }
+        sql += ` ${this.escapeId(name)}`;
+        if (options?.owner) {
+            sql += ` OWNER ${this.escapeId(options.owner)}`;
+        }
+        return sql;
+    }
+    /**
+     * Create a tablegroup (`CREATE TABLEGROUP name`). Tables created with the
+     * `tablegroup` option share the tablegroup's tablets.
+     */
+    async createTableGroup(name, options) {
+        await this.query(this.buildCreateTableGroupSQL(name, options));
+    }
+    /**
+     * Drop a tablegroup (`DROP TABLEGROUP [IF EXISTS] name`).
+     */
+    async dropTableGroup(name, options) {
+        let sql = 'DROP TABLEGROUP';
+        if (options?.ifExists) {
+            sql += ' IF EXISTS';
+        }
+        sql += ` ${this.escapeId(name)}`;
+        await this.query(sql);
+    }
+}
+exports.YugabyteDBDialect = YugabyteDBDialect;
+/**
+ * Create a new YugabyteDB dialect instance.
+ */
+function createYugabyteDBDialect(options) {
+    return new YugabyteDBDialect(options || {});
+}
+exports.default = YugabyteDBDialect;