lib/src/orm/schema/TableBuilder.ts

import {collect, Maybe, ParameterizedCallback, Pipeline} from '../../util'
import {FieldType} from '../types'
import {EscapeValue, QuerySafeValue, raw} from '../dialect/SQLDialect'

/**
 * Base class with shared logic for the various schema
 * builders (table, column, index).
 */
export abstract class SchemaBuilderBase {
    /**
     * Whether or not the schema item should be dropped.
     *  - `exists` - drop if exists
     * @protected
     */
    protected shouldDrop: 'yes'|'no'|'exists' = 'no'

    /**
     * The name the schema item should have if renaming.
     * @protected
     */
    protected shouldRenameTo?: string

    /**
     * If true, apply IF NOT EXISTS syntax.
     * @protected
     */
    protected shouldSkipIfExists = false

    /** True if the schema has been modified since created/loaded. */
    protected dirty = false

    /** True if this resource exists, in some form, in the schema. */
    protected existsInSchema = false

    /** If the resource exists in the schema, the unaltered values it has. */
    public originalFromSchema?: this

    constructor(
        /** The name of the schema item. */
        public readonly name: string,
    ) { }

    /**
     * Clone the properties of this resource to a different instance.
     * @param newBuilder
     */
    public cloneTo(newBuilder: this): this {
        newBuilder.shouldDrop = this.shouldDrop
        newBuilder.shouldRenameTo = this.shouldRenameTo
        newBuilder.shouldSkipIfExists = this.shouldSkipIfExists
        newBuilder.dirty = this.dirty
        newBuilder.existsInSchema = this.existsInSchema
        return newBuilder
    }

    /** True if this resource should be dropped. */
    public isDropping(): boolean {
        return this.shouldDrop === 'yes'
    }

    /** True if this resource should be dropped with IF EXISTS syntax. */
    public isDroppingIfExists(): boolean {
        return this.shouldDrop === 'exists'
    }

    /** True if this resource should be created with IF NOT EXISTS syntax. */
    public isSkippedIfExisting(): boolean {
        return this.shouldSkipIfExists
    }

    /** True if the resource already exists in some form in the schema. */
    public isExisting(): boolean {
        return this.existsInSchema
    }

    /** True if the resource has been modified since created/loaded. */
    public isDirty(): boolean {
        return this.dirty
    }

    /**
     * Get the name this resource should be renamed to, if it exists.
     */
    public getRename(): Maybe<string> {
        return this.shouldRenameTo
    }

    /** Mark the resource to be removed. */
    public drop(): this {
        this.dirty = true
        this.shouldDrop = 'yes'
        return this
    }

    /** Mark the resource to be removed, if it exists. */
    public dropIfExists(): this {
        this.dirty = true
        this.shouldDrop = 'exists'
        return this
    }

    /**
     * Rename the resource to a different name.
     * @param to
     */
    public rename(to: string): this {
        this.dirty = true
        this.shouldRenameTo = to
        return this
    }

    /**
     * Mark the resource to use IF NOT EXISTS syntax.
     */
    public ifNotExists(): this {
        this.shouldSkipIfExists = true
        return this
    }

    /**
     * Used internally.
     * Mark that the resource exists in the schema in some form,
     * and reset the `dirty` flag.
     */
    public flagAsExistingInSchema(): this {
        this.existsInSchema = true
        this.dirty = false
        this.originalFromSchema = this.cloneTo(this.cloneInstance())
        return this
    }

    /** Build and apply a pipeline. */
    pipe<TOut>(builder: (pipeline: Pipeline<this, this>) => Pipeline<this, TOut>): TOut {
        return builder(Pipeline.id()).apply(this)
    }

    tap(op: (x: this) => unknown): this {
        op(this)
        return this
    }

    /**
     * Get a new instance of the concrete implementation of this class.
     * @protected
     */
    protected abstract cloneInstance(): this
}

/**
 * Builder to specify the schema of a table column.
 */
export class ColumnBuilder extends SchemaBuilderBase {
    /** The data type of the column. */
    protected targetType?: FieldType

    /** True if the column should allow NULL values. */
    protected shouldBeNullable = false

    /** The default value of the column, if one should exist. */
    protected defaultValue?: EscapeValue

    /** The data length of this column, if set */
    protected targetLength?: number

    /** True if this is a primary key constraint. */
    protected shouldBePrimary = false

    /** True if this column should contain distinct values. */
    protected shouldBeUnique = false

    constructor(
        name: string,

        /** The table this column belongs to. */
        public readonly parent: TableBuilder,
    ) {
        super(name)
    }

    public cloneTo(newBuilder: this): this {
        super.cloneTo(newBuilder)
        newBuilder.targetType = this.targetType
        newBuilder.shouldBeNullable = this.shouldBeNullable
        newBuilder.defaultValue = this.defaultValue
        newBuilder.targetLength = this.targetLength
        newBuilder.shouldBePrimary = this.shouldBePrimary
        newBuilder.shouldBeUnique = this.shouldBeUnique
        return newBuilder
    }

    /** Get the field type of the column, if it exists. */
    public getType(): Maybe<FieldType> {
        return this.targetType
    }

    /** Get the data-type length of the column, if it exists. */
    public getLength(): Maybe<number> {
        return this.targetLength
    }

    /** Get the default value of the column, if it exists. */
    public getDefaultValue(): Maybe<EscapeValue> {
        return this.defaultValue
    }

    /** True if the column should allow NULL values. */
    public isNullable(): boolean {
        return this.shouldBeNullable
    }

    /** True if the column is a primary key. */
    public isPrimary(): boolean {
        return this.shouldBePrimary
    }

    /** True if the column should require unique values. */
    public isUnique(): boolean {
        return this.shouldBeUnique
    }

    /**
     * Specify the data type of the column.
     * @param type
     */
    public type(type: FieldType): this {
        if ( this.targetType === type ) {
            return this
        }

        this.dirty = true
        this.targetType = type
        return this
    }

    /**
     * Make the column nullable.
     */
    public nullable(): this {
        if ( this.shouldBeNullable ) {
            return this
        }

        this.dirty = true
        this.shouldBeNullable = true
        return this
    }

    /**
     * Make the column non-nullable.
     */
    public required(): this {
        if ( !this.shouldBeNullable ) {
            return this
        }

        this.dirty = true
        this.shouldBeNullable = false
        return this
    }

    /**
     * Specify the default value of the column.
     * @param value
     */
    public default(value: EscapeValue): this {
        if ( this.defaultValue === value ) {
            return this
        }

        this.dirty = true
        this.defaultValue = value
        return this
    }

    /**
     * Set the length of this column's data type.
     * @param value
     */
    public length(value: number): this {
        if ( this.targetLength === value ) {
            return this
        }

        this.dirty = true
        this.targetLength = value
        return this
    }

    /**
     * Make this a primary-key column.
     */
    primary(): this {
        if ( this.shouldBePrimary ) {
            return this
        }

        this.dirty = true
        this.shouldBePrimary = true
        return this
    }

    /**
     * Make this column require distinct values.
     */
    unique(): this {
        if ( this.shouldBeUnique ) {
            return this
        }

        this.dirty = true
        this.shouldBeUnique = true
        return this
    }

    protected cloneInstance(): this {
        return new ColumnBuilder(this.name, this.parent) as this
    }
}

/** Valid constraint types. */
export enum ConstraintType {
    Unique = 'un',
    Check = 'ck',
}

/**
 * Builder to specify the schema of a table constraint.
 */
export class ConstraintBuilder extends SchemaBuilderBase {

    /** The fields included in this constraint. */
    protected fields: Set<string> = new Set<string>()

    /** The type of this constraint. */
    protected constraintType: ConstraintType = ConstraintType.Unique

    /** The expression defining this constraint, if applicable. */
    protected constraintExpression?: QuerySafeValue

    constructor(
        name: string,

        /** The table this constraint belongs to. */
        public readonly parent: TableBuilder,
    ) {
        super(name)
    }

    /** Get the type of this constraint. */
    public getType(): ConstraintType {
        return this.constraintType
    }

    /** Get the fields included in this constraint. */
    public getFields(): string[] {
        return [...this.fields]
    }

    /** Get the expression used to evaluate this constraint, if it exists. */
    public getExpression(): Maybe<QuerySafeValue> {
        return this.constraintExpression
    }

    public cloneTo(newBuilder: this): this {
        super.cloneTo(newBuilder)
        newBuilder.fields = new Set<string>([...this.fields])
        newBuilder.constraintType = this.constraintType
        return newBuilder
    }

    protected cloneInstance(): this {
        return new ConstraintBuilder(this.name, this.parent) as this
    }

    /** Add a field to this constraint. */
    public field(name: string): this {
        if ( this.fields.has(name) ) {
            return this
        }

        this.dirty = true
        this.fields.add(name)
        return this
    }

    /** Remove a field from this constraint. */
    public removeField(name: string): this {
        if ( !this.fields.has(name) ) {
            return this
        }

        this.dirty = true
        this.fields.delete(name)
        return this
    }

    /** Specify the type of this constraint. */
    public type(type: ConstraintType): this {
        if ( this.constraintType === type ) {
            return this
        }

        this.dirty = true
        this.constraintType = type
        return this
    }

    /** Specify the expression used to evaluate this constraint, if applicable. */
    public expression(sql: string | QuerySafeValue): this {
        if ( String(this.constraintExpression) === String(sql) ) {
            return this
        }

        this.dirty = true

        if ( sql instanceof QuerySafeValue ) {
            this.constraintExpression = sql
        }

        this.constraintExpression = raw(sql)
        return this
    }
}

/**
 * Builder to specify the schema of a table index.
 */
export class IndexBuilder extends SchemaBuilderBase {
    /** The fields included in the index. */
    protected fields: Set<string> = new Set<string>()

    /** Fields to remove from the index. */
    protected removedFields: Set<string> = new Set<string>()

    /** True if this is a unique index. */
    protected shouldBeUnique = false

    /** True if this is a primary key index. */
    protected shouldBePrimary = false

    constructor(
        name: string,

        /** The table this index belongs to. */
        public readonly parent: TableBuilder,
    ) {
        super(name)
    }

    public cloneTo(newBuilder: this): this {
        super.cloneTo(newBuilder)
        newBuilder.fields = new Set<string>([...this.fields])
        newBuilder.removedFields = new Set<string>([...this.removedFields])
        newBuilder.shouldBeUnique = this.shouldBeUnique
        newBuilder.shouldBePrimary = this.shouldBePrimary
        return newBuilder
    }

    /** Get the fields in this index. */
    public getFields(): string[] {
        return [...this.fields]
    }

    /** True if this index is a unique index. */
    public isUnique(): boolean {
        return this.shouldBeUnique
    }

    /** True if this index is the primary key index. */
    public isPrimary(): boolean {
        return this.shouldBePrimary
    }

    /**
     * Add the given field to this index.
     * @param name
     */
    public field(name: string): this {
        if ( this.fields.has(name) ) {
            return this
        }

        this.dirty = true
        this.fields.add(name)
        return this
    }

    /**
     * Remove the given field from this index.
     * @param name
     * @protected
     */
    protected removeField(name: string): this {
        if ( !this.fields.has(name) ) {
            return this
        }

        this.dirty = true
        this.removedFields.add(name)
        this.fields.delete(name)
        return this
    }

    /**
     * Make this a primary-key index.
     */
    primary(): this {
        if ( this.shouldBePrimary ) {
            return this
        }

        this.dirty = true
        this.shouldBePrimary = true
        return this
    }

    /**
     * Make this a unique index.
     */
    unique(): this {
        if ( this.shouldBeUnique ) {
            return this
        }

        this.dirty = true
        this.shouldBeUnique = true
        return this
    }

    protected cloneInstance(): this {
        return new IndexBuilder(this.name, this.parent) as this
    }
}

/**
 * Builder to specify the schema of a table.
 */
export class TableBuilder extends SchemaBuilderBase {

    /**
     * Mapping of column name to column schemata.
     * @protected
     */
    protected columns: {[key: string]: ColumnBuilder} = {}

    /**
     * Mapping of index name to index schemata.
     * @protected
     */
    protected indexes: {[key: string]: IndexBuilder} = {}

    /**
     * Mapping of constraint name to constraint schemata.
     * @protected
     */
    protected constraints: {[key: string]: ConstraintBuilder} = {}

    public cloneTo(newBuilder: this): this {
        super.cloneTo(newBuilder)
        newBuilder.columns = {...this.columns}
        newBuilder.indexes = {...this.indexes}
        newBuilder.constraints = {...this.constraints}
        return newBuilder
    }

    /**
     * Get the columns defined on this table.
     */
    public getColumns(): {[key: string]: ColumnBuilder} {
        return {
            ...this.columns,
        }
    }

    /**
     * Get the indices defined on this table.
     */
    public getIndexes(): {[key: string]: IndexBuilder} {
        return {
            ...this.indexes,
        }
    }

    /**
     * Get the constraints defined on this table.
     */
    public getConstraints(): {[key: string]: ConstraintBuilder} {
        return {
            ...this.constraints,
        }
    }

    /**
     * Mark a column to be dropped.
     * @param name
     */
    public dropColumn(name: string): this {
        this.dirty = true
        this.column(name).drop()
        return this
    }

    /**
     * Mark a column to be renamed.
     * @param from
     * @param to
     */
    public renameColumn(from: string, to: string): this {
        this.column(from).rename(to)
        return this
    }

    /**
     * Mark an index to be dropped.
     * @param name
     */
    public dropIndex(name: string): this {
        this.dirty = true
        this.index(name).drop()
        return this
    }

    /**
     * Mark an index to be renamed.
     * @param from
     * @param to
     */
    public renameIndex(from: string, to: string): this {
        this.index(from).rename(to)
        return this
    }

    /**
     * Add a column to this table.
     * @param name
     * @param callback
     */
    public column(name: string, callback?: ParameterizedCallback<ColumnBuilder>): ColumnBuilder {
        if ( !this.columns[name] ) {
            this.dirty = true
            this.columns[name] = new ColumnBuilder(name, this)
        }

        if ( callback ) {
            callback(this.columns[name])
        }

        return this.columns[name]
    }

    /**
     * Add an index to this table.
     * @param name
     * @param callback
     */
    public index(name: string, callback?: ParameterizedCallback<IndexBuilder>): IndexBuilder {
        if ( !this.indexes[name] ) {
            this.dirty = true
            this.indexes[name] = new IndexBuilder(name, this)
        }

        if ( callback ) {
            callback(this.indexes[name])
        }

        return this.indexes[name]
    }

    /**
     * Add a constraint to this table.
     * @param name
     * @param callback
     */
    public constraint(name: string, callback?: ParameterizedCallback<ConstraintBuilder>): ConstraintBuilder {
        if ( !this.constraints[name] ) {
            this.dirty = true
            this.constraints[name] = new ConstraintBuilder(name, this)
        }

        if ( callback ) {
            callback(this.constraints[name])
        }

        return this.constraints[name]
    }

    /**
     * Generate a programmatically-incrementing constraint name.
     * @param suffix
     * @protected
     */
    protected getNextAvailableConstraintName(suffix: ConstraintType): string {
        let current = 1
        let name = `${this.name}_${current}_${suffix}`

        while ( this.constraints[name] ) {
            current += 1
            name = `${this.name}_${current}_${suffix}`
        }

        return name
    }

    /**
     * Add a new check constraint with the given expression.
     * @param expression
     */
    public check(expression: string | QuerySafeValue): this {
        const name = this.getNextAvailableConstraintName(ConstraintType.Check)
        this.constraint(name)
            .type(ConstraintType.Check)
            .expression(expression)

        return this
    }

    /**
     * Add a new unique constraint for the given fields.
     * @param fields
     */
    public unique(...fields: string[]): this {
        // First, check if an existing constraint exists with these fields
        for ( const key in this.constraints ) {
            if ( !Object.prototype.hasOwnProperty.call(this.constraints, key) ) {
                continue
            }

            if ( this.constraints[key].getType() !== ConstraintType.Unique ) {
                continue
            }

            const existingFields = collect<string>(this.constraints[key].getFields())
            const intersection = existingFields.intersect(fields)

            if ( existingFields.length === fields.length && intersection.length === fields.length ) {
                return this
            }
        }

        // If an existing constraint can't satisfy this, create a new one
        const name = this.getNextAvailableConstraintName(ConstraintType.Unique)
        this.constraint(name)
            .type(ConstraintType.Unique)
            .tap(constraint => {
                fields.forEach(field => constraint.field(field))
            })

        return this
    }

    /**
     * Add a primary key (column & index) to this table.
     * @param name
     * @param type
     */
    public primaryKey(name: string, type: FieldType = FieldType.serial): ColumnBuilder {
        this.dirty = true

        return this.column(name)
            .type(type)
            .primary()
    }

    protected cloneInstance(): this {
        return new TableBuilder(this.name) as this
    }
}