lib/src/orm/dialect/PostgreSQLDialect.ts

import {EscapeValue, QuerySafeValue, raw, SQLDialect} from './SQLDialect'
import {Constraint, inverseFieldType, isConstraintGroup, isConstraintItem, QuerySource, SpecifiedField} from '../types'
import {AbstractBuilder} from '../builder/AbstractBuilder'
import {ColumnBuilder, ConstraintBuilder, ConstraintType, IndexBuilder, TableBuilder} from '../schema/TableBuilder'
import {collect, Collectable, Collection, ErrorWithContext, hasOwnProperty, Maybe} from '../../util'

/**
 * An implementation of the SQLDialect specific to PostgreSQL.
 * @todo joins
 * @todo sub-selects
 */
export class PostgreSQLDialect extends SQLDialect {

    public escape(value: EscapeValue): QuerySafeValue {
        if ( value instanceof QuerySafeValue ) {
            return value
        } else if ( Array.isArray(value) || value instanceof Collection ) {
            return new QuerySafeValue(value, `(${value.map(v => this.escape(v)).join(',')})`)
        } else if ( String(value).toLowerCase() === 'true' || value === true ) {
            return new QuerySafeValue(value, 'TRUE')
        } else if ( String(value).toLowerCase() === 'false' || value === false ) {
            return new QuerySafeValue(value, 'FALSE')
        } else if ( typeof value === 'number' ) {
            return new QuerySafeValue(value, `${value}`)
        } else if ( value instanceof Date ) {
            const pad = (val: number) => val < 10 ? `0${val}` : `${val}`
            const [y, m, d, h, i, s] = [
                `${value.getFullYear()}`,
                `${pad(value.getMonth() + 1)}`,
                `${pad(value.getDate())}`,
                `${pad(value.getHours())}`,
                `${pad(value.getMinutes())}`,
                `${pad(value.getSeconds())}`,
            ]

            return new QuerySafeValue(value, `'${y}-${m}-${d} ${h}:${i}:${s}'`)
        } else if ( value === null || typeof value === 'undefined' ) {
            return new QuerySafeValue(value, 'NULL')
        } else if ( !isNaN(Number(value)) ) {
            return new QuerySafeValue(value, String(Number(value)))
        } else {
            const escaped = value.replace(/'/g, '\'\'')  // .replace(/"/g, '\\"').replace(/`/g, '\\`')
            return new QuerySafeValue(value, `'${escaped}'`)
        }
    }

    public renderQuerySource(source: QuerySource): string {
        if ( source instanceof QuerySafeValue ) {
            return String(source)
        } else if ( typeof source === 'string' ) {
            return source.replace(/"/g, '""')
                .split('.')
                .map(x => '"' + x + '"')
                .join('.')
        }

        return `${this.renderQuerySource(source.table)} AS "${source.alias.replace(/"/g, '""')}"`
    }

    public renderCount(query: string): string {
        return [
            'SELECT COUNT(*) AS "extollo_render_count"',
            'FROM (',
            ...query.split('\n').map(x => `  ${x}`),
            ') AS extollo_target_query',
        ].join('\n')
    }

    public renderRangedSelect(query: string, start: number, end: number): string {
        return [
            'SELECT *',
            'FROM (',
            ...query.split('\n').map(x => `  ${x}`),
            ') AS extollo_target_query',
            `OFFSET ${start} LIMIT ${start === end ? ((end - start) + 1) : (end - start)}`,
        ].join('\n')
    }

    /** Render the fields from the builder class to PostgreSQL syntax. */
    protected renderFields(builder: AbstractBuilder<any>): string[] {
        return builder.appliedFields.map((field: SpecifiedField) => {
            let columnString: string
            if ( typeof field === 'string' ) {
                columnString = field.split('.').map(x => `"${x}"`)
                    .join('.')
            } else if ( field instanceof QuerySafeValue ) {
                columnString = field.toString()
            } else if ( typeof field.field === 'string' ) {
                columnString = field.field.split('.').map(x => `"${x}"`)
                    .join('.')
            } else {
                columnString = field.field.toString()
            }

            let aliasString = ''
            if ( typeof field !== 'string' && !(field instanceof QuerySafeValue) ) {
                aliasString = ` AS "${field.alias}"`
            }

            return `${columnString}${aliasString}`
        })
    }

    public renderSelect(builder: AbstractBuilder<any>): string {
        const rawSql = builder.appliedRawSql
        if ( rawSql ) {
            return rawSql
        }

        const indent = (item: string, level = 1) => Array(level + 1).fill('')
            .join('  ') + item
        const queryLines = [
            `SELECT${builder.appliedDistinction ? ' DISTINCT' : ''}`,
        ]

        // Add fields
        // FIXME error if no fields
        const fields = this.renderFields(builder).map(x => indent(x))
            .join(',\n')

        queryLines.push(fields)

        // Add table source
        // FIXME error if no source
        const source = builder.querySource
        if ( source ) {
            queryLines.push('FROM ' + this.renderQuerySource(source))
        }

        // Add constraints
        const wheres = this.renderConstraints(builder.appliedConstraints)
        if ( wheres.trim() ) {
            queryLines.push('WHERE')
            queryLines.push(wheres)
        }

        // Add group by
        if ( builder.appliedGroupings?.length ) {
            const grouping = builder.appliedGroupings.map(group => {
                return indent(group.split('.').map(x => `"${x}"`)
                    .join('.'))
            }).join(',\n')

            queryLines.push('GROUP BY')
            queryLines.push(grouping)
        }

        // Add order by
        if ( builder.appliedOrder?.length ) {
            const ordering = builder.appliedOrder.map(x => indent(`${x.field.split('.').map(y => '"' + y + '"')
                .join('.')} ${x.direction}`)).join(',\n')
            queryLines.push('ORDER BY')
            queryLines.push(ordering)
        }

        // Add limit/offset
        const pagination = builder.appliedPagination
        if ( pagination.take ) {
            queryLines.push(`LIMIT ${pagination.take}${pagination.skip ? ' OFFSET ' + pagination.skip : ''}`)
        } else if ( pagination.skip ) {
            queryLines.push(`OFFSET ${pagination.skip}`)
        }

        return queryLines.join('\n')
    }

    public renderBatchUpdate(builder: AbstractBuilder<any>, primaryKey: string, dataRows: Collectable<{[key: string]: EscapeValue}>): string {
        const rows = Collection.normalize(dataRows)
        const rawSql = builder.appliedRawSql
        if ( rawSql ) {
            return rawSql
        }

        const queryLines: string[] = []

        // Add table source
        let source = builder.querySource
        if ( !source ) {
            throw new ErrorWithContext('No table specified for update query')
        }

        source = (typeof source !== 'string' && !(source instanceof QuerySafeValue)) ? source : {
            table: source,
            alias: 'extollo_update_source',
        }

        const sourceAlias = source.alias
        const sourceTable = source.table

        queryLines.push('UPDATE ' + this.renderQuerySource(source))

        queryLines.push('SET')
        const updateFields = this.getAllFieldsFromUpdateRows(rows)

        const updateTuples = rows.map(row => {
            return updateFields.map(field => {
                if ( hasOwnProperty(row, field) ) {
                    return this.escape(row[field])
                }

                // FIXME: This is fairly inefficient. Probably a better way with a FROM ... SELECT
                // return raw(`"${sourceAlias}"."${field}"`)
                return raw(`(SELECT "${field}" FROM ${sourceTable} WHERE "${primaryKey}" = ${this.escape(row[primaryKey])})`)
            })
        })

        queryLines.push(updateFields.map(field => `  "${field}" = "extollo_update_tuple"."${field}"`).join(',\n'))

        queryLines.push('FROM (VALUES')

        queryLines.push(
            updateTuples.map(tuple => `  (${tuple.implode(', ')})`).join(',\n'),
        )

        queryLines.push(`) as extollo_update_tuple(${updateFields.map(x => `"${x}"`).join(', ')})`)

        queryLines.push(`WHERE "extollo_update_tuple"."${primaryKey}" = "${sourceAlias}"."${primaryKey}" AND (`)

        queryLines.push(this.renderConstraints(builder.appliedConstraints, 2))

        queryLines.push(`)`)

        return queryLines.join('\n')
    }

    private getAllFieldsFromUpdateRows(rows: Collection<{[key: string]: EscapeValue}>): Collection<string> {
        return rows.reduce((fields: Collection<string>, row) => {
            Object.keys(row).forEach(key => {
                if ( !fields.includes(key) ) {
                    fields.push(key)
                }
            })
            return fields
        }, collect<string>())
    }

    // TODO support FROM, RETURNING
    public renderUpdate(builder: AbstractBuilder<any>, data: {[key: string]: EscapeValue}): string {
        const rawSql = builder.appliedRawSql
        if ( rawSql ) {
            return rawSql
        }

        const queryLines: string[] = []

        // Add table source
        const source = builder.querySource
        if ( source ) {
            queryLines.push('UPDATE ' + this.renderQuerySource(source))
        }

        queryLines.push(this.renderUpdateSet(data))

        // Add constraints
        const wheres = this.renderConstraints(builder.appliedConstraints)
        if ( wheres.trim() ) {
            queryLines.push('WHERE')
            queryLines.push(wheres)
        }

        const fields = this.renderFields(builder).map(x => `  ${x}`)
            .join(',\n')

        if ( fields ) {
            queryLines.push('RETURNING')
            queryLines.push(fields)
        }

        return queryLines.join('\n')
    }

    public renderExistential(builder: AbstractBuilder<any>): string {
        const rawSql = builder.appliedRawSql
        if ( rawSql ) {
            return `
                SELECT EXISTS(
                    ${rawSql}
                )
            `
        }

        const query = builder.clone()
            .clearFields()
            .field(raw('TRUE'))
            .limit(1)

        return this.renderSelect(query)
    }

    // FIXME: subquery support here and with select
    public renderInsert(builder: AbstractBuilder<any>, data: {[key: string]: EscapeValue}|{[key: string]: EscapeValue}[] = []): string {
        const rawSql = builder.appliedRawSql
        if ( rawSql ) {
            return rawSql
        }

        const indent = (item: string, level = 1) => Array(level + 1).fill('')
            .join('  ') + item
        const queryLines: string[] = []

        if ( !Array.isArray(data) ) {
            data = [data]
        }

        if ( data.length < 1 ) {
            return ''
        }

        const columns = Object.keys(data[0])

        // Add table source
        const source = builder.querySource
        if ( source ) {
            queryLines.push('INSERT INTO ' + this.renderQuerySource(source)
                + (columns.length ? ` (${columns.map(x => `"${x}"`).join(', ')})` : ''))
        }

        if ( Array.isArray(data) && !data.length ) {
            queryLines.push('DEFAULT VALUES')
        } else {
            queryLines.push('VALUES')

            const valueString = data.map(row => {
                const values = columns.map(x => this.escape(row[x]))
                return indent(`(${values.join(', ')})`)
            })
                .join(',\n')

            queryLines.push(valueString)
        }

        // Add return fields
        if ( builder.appliedFields?.length ) {
            queryLines.push('RETURNING')
            const fields = this.renderFields(builder).map(x => indent(x))
                .join(',\n')

            queryLines.push(fields)
        }

        return queryLines.join('\n')
    }

    public renderDelete(builder: AbstractBuilder<any>): string {
        const rawSql = builder.appliedRawSql
        if ( rawSql ) {
            return rawSql
        }

        const indent = (item: string, level = 1) => Array(level + 1).fill('')
            .join('  ') + item
        const queryLines: string[] = []

        // Add table source
        const source = builder.querySource
        if ( source ) {
            queryLines.push('DELETE FROM ' + this.renderQuerySource(source))
        }

        // Add constraints
        const wheres = this.renderConstraints(builder.appliedConstraints)
        if ( wheres.trim() ) {
            queryLines.push('WHERE')
            queryLines.push(wheres)
        }

        // Add return fields
        if ( builder.appliedFields?.length ) {
            queryLines.push('RETURNING')

            const fields = this.renderFields(builder).map(x => indent(x))
                .join(',\n')

            queryLines.push(fields)
        }

        return queryLines.join('\n')
    }

    public renderConstraints(allConstraints: Constraint[], startingLevel = 1): string {
        const constraintsToSql = (constraints: Constraint[], level = startingLevel): string => {
            const indent = Array(level * 2).fill(' ')
                .join('')
            const statements = []

            for ( const constraint of constraints ) {
                if ( isConstraintGroup(constraint) ) {
                    statements.push(`${indent}${statements.length < 1 ? '' : constraint.preop + ' '}(\n${constraintsToSql(constraint.items, level + 1)}\n${indent})`)
                } else if ( isConstraintItem(constraint) ) {
                    if ( Array.isArray(constraint.operand) && !constraint.operand.length ) {
                        statements.push(`${indent}1 = 0  -- ${constraint.field} ${constraint.operator} empty set`)
                        continue
                    }

                    const field: string = constraint.field.split('.').map(x => `"${x}"`)
                        .join('.')
                    statements.push(`${indent}${statements.length < 1 ? '' : constraint.preop + ' '}${field} ${constraint.operator} ${this.escape(constraint.operand).value}`)
                } else if ( constraint instanceof QuerySafeValue ) {
                    statements.push(`${indent}${statements.length < 1 ? '' : 'AND '}${constraint.toString()}`)
                }
            }

            return statements.filter(Boolean).join('\n')
        }

        return constraintsToSql(allConstraints)
    }

    public renderUpdateSet(data: {[key: string]: EscapeValue}): string {
        const sets = []
        for ( const key in data ) {
            if ( !Object.prototype.hasOwnProperty.call(data, key) ) {
                continue
            }

            sets.push(`  "${key}" = ${this.escape(data[key])}`)
        }

        return `SET\n${sets.join(',\n')}`
    }

    public renderCreateTable(builder: TableBuilder): string {
        const cols = this.renderTableColumns(builder).map(x => `    ${x}`)

        const builderConstraints = builder.getConstraints()
        const constraints: string[] = []
        for ( const constraintName in builderConstraints ) {
            if ( !Object.prototype.hasOwnProperty.call(builderConstraints, constraintName) ) {
                continue
            }

            const constraintBuilder = builderConstraints[constraintName]
            const constraintDefinition = this.renderConstraintDefinition(constraintBuilder)
            if ( constraintDefinition ) {
                constraints.push(`    CONSTRAINT ${constraintDefinition}`)
            }
        }

        const parts = [
            `CREATE TABLE ${builder.isSkippedIfExisting() ? 'IF NOT EXISTS ' : ''}${builder.name} (`,
            [
                ...cols,
                ...constraints,
            ].join(',\n'),
            `)`,
        ]

        return parts.join('\n')
    }

    public renderTableColumns(builder: TableBuilder): string[] {
        const defined = builder.getColumns()
        const rendered: string[] = []

        for ( const columnName in defined ) {
            if ( !Object.prototype.hasOwnProperty.call(defined, columnName) ) {
                continue
            }

            const columnBuilder = defined[columnName]
            rendered.push(this.renderColumnDefinition(columnBuilder))
        }

        return rendered
    }

    /**
     * Given a constraint schema-builder, render the constraint definition.
     * @param builder
     * @protected
     */
    protected renderConstraintDefinition(builder: ConstraintBuilder): Maybe<string> {
        const constraintType = builder.getType()
        if ( constraintType === ConstraintType.Unique ) {
            const fields = builder.getFields()
                .map(x => `"${x}"`)
                .join(',')

            return `${builder.name} UNIQUE(${fields})`
        } else if ( constraintType === ConstraintType.Check ) {
            const expression = builder.getExpression()
            if ( !expression ) {
                throw new ErrorWithContext('Cannot create check constraint without expression.', {
                    constraintName: builder.name,
                    tableName: builder.parent.name,
                })
            }

            return `${builder.name} CHECK(${expression})`
        }
    }

    /**
     * Given a column-builder, render the SQL-definition as used in
     * CREATE TABLE and ALTER TABLE statements.
     * @fixme Type `serial` only exists on CREATE TABLE... queries
     * @param builder
     * @protected
     */
    protected renderColumnDefinition(builder: ColumnBuilder): string {
        const type = builder.getType()
        if ( !type ) {
            throw new ErrorWithContext(`Missing field type for column: ${builder.name}`, {
                columnName: builder.name,
                columnType: type,
            })
        }

        let render = `"${builder.name}"  ${inverseFieldType(type)}`

        if ( builder.getLength() ) {
            render += `(${builder.getLength()})`
        }

        const defaultValue = builder.getDefaultValue()
        if ( typeof defaultValue !== 'undefined' ) {
            render += ` DEFAULT ${this.escape(defaultValue)}`
        }

        if ( builder.isPrimary() ) {
            render += ` CONSTRAINT ${builder.name}_pk PRIMARY KEY`
        }

        if ( builder.isUnique() ) {
            render += ` UNIQUE`
        }

        render += ` ${builder.isNullable() ? 'NULL' : 'NOT NULL'}`
        return render
    }

    public renderDropTable(builder: TableBuilder): string {
        return `DROP TABLE ${builder.isSkippedIfExisting() ? 'IF EXISTS ' : ''}${builder.name}`
    }

    public renderCreateIndex(builder: IndexBuilder): string {
        const cols = builder.getFields().map(x => `"${x}"`)
        const parts = [
            `CREATE ${builder.isUnique() ? 'UNIQUE ' : ''}INDEX ${builder.isSkippedIfExisting() ? 'IF NOT EXISTS ' : ''}${builder.name}`,
            `    ON ${builder.parent.name}`,
            `    (${cols.join(',')})`,
        ]

        return parts.join('\n')
    }

    public renderAlterTable(builder: TableBuilder): string {
        const alters: string[] = []
        const columns = builder.getColumns()

        for ( const columnName in columns ) {
            if ( !Object.prototype.hasOwnProperty.call(columns, columnName) ) {
                continue
            }

            const columnBuilder = columns[columnName]
            if ( !columnBuilder.isExisting() ) {
                // The column doesn't exist on the table, but was added to the schema
                alters.push(`    ADD COLUMN ${this.renderColumnDefinition(columnBuilder)}`)
            } else if ( columnBuilder.isDirty() && columnBuilder.originalFromSchema ) {
                // The column exists in the table, but was modified in the schema
                if ( columnBuilder.isDropping() || columnBuilder.isDroppingIfExists() ) {
                    alters.push(`    DROP COLUMN "${columnBuilder.name}"`)
                    continue
                }

                // Change the data type of the column
                if ( columnBuilder.getType() !== columnBuilder.originalFromSchema.getType() ) {
                    const renderedType = `${columnBuilder.getType()}${columnBuilder.getLength() ? `(${columnBuilder.getLength()})` : ''}`
                    alters.push(`    ALTER COLUMN "${columnBuilder.name}" TYPE ${renderedType}`)
                }

                // Change the default value of the column
                if ( columnBuilder.getDefaultValue() !== columnBuilder.originalFromSchema.getDefaultValue() ) {
                    alters.push(`    ALTER COLUMN "${columnBuilder.name}" SET default ${this.escape(columnBuilder.getDefaultValue())}`)
                }

                // Change the nullable-status of the column
                if ( columnBuilder.isNullable() !== columnBuilder.originalFromSchema.isNullable() ) {
                    if ( columnBuilder.isNullable() ) {
                        alters.push(`    ALTER COLUMN "${columnBuilder.name}" DROP NOT NULL`)
                    } else {
                        alters.push(`    ALTER COLUMN "${columnBuilder.name}" SET NOT NULL`)
                    }
                }

                // Change the name of the column
                if ( columnBuilder.getRename() ) {
                    alters.push(`    RENAME COLUMN "${columnBuilder.name}" TO "${columnBuilder.getRename()}"`)
                }
            }
        }

        const constraints = builder.getConstraints()
        for ( const constraintName in constraints ) {
            if ( !Object.prototype.hasOwnProperty.call(constraints, constraintName) ) {
                continue
            }

            const constraintBuilder = constraints[constraintName]

            // Drop the constraint if specified
            if ( constraintBuilder.isDropping() ) {
                alters.push(`    DROP CONSTRAINT ${constraintBuilder.name}`)
                continue
            }

            // Drop the constraint with IF EXISTS if specified
            if ( constraintBuilder.isDroppingIfExists() ) {
                alters.push(`    DROP CONSTRAINT IF EXISTS ${constraintBuilder.name}`)
                continue
            }

            // Otherwise, drop and recreate the constraint if it was modified
            if ( constraintBuilder.isDirty() ) {
                if ( constraintBuilder.isExisting() ) {
                    alters.push(`    DROP CONSTRAINT IF EXISTS ${constraintBuilder.name}`)
                }

                const constraintDefinition = this.renderConstraintDefinition(constraintBuilder)
                if ( constraintDefinition ) {
                    alters.push(`    ADD CONSTRAINT ${constraintDefinition}`)
                }
            }
        }

        if ( builder.getRename() ) {
            alters.push(`    RENAME TO "${builder.getRename()}"`)
        }

        return 'ALTER TABLE ' + builder.name + '\n' + alters.join(',\n')
    }

    public renderDropIndex(builder: IndexBuilder): string {
        return `DROP INDEX ${builder.isDroppingIfExists() ? 'IF EXISTS ' : ''}${builder.name}`
    }

    public renderTransaction(queries: string[]): string {
        const parts = [
            'BEGIN',
            ...queries,
            'COMMIT',
        ]

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

    public renderRenameIndex(builder: IndexBuilder): string {
        return `ALTER INDEX ${builder.name} RENAME TO ${builder.getRename()}`
    }

    public renderRecreateIndex(builder: IndexBuilder): string {
        return `${this.renderDropIndex(builder)};\n\n${this.renderCreateIndex(builder)}`
    }

    public renderDropColumn(builder: ColumnBuilder): string {
        const parts = [
            `ALTER TABLE ${builder.parent.name} ${builder.parent.isSkippedIfExisting() ? 'IF EXISTS ' : ''}`,
            `    DROP COLUMN ${builder.isSkippedIfExisting() ? 'IF EXISTS ' : ''}${builder.name}`,
        ]

        return parts.join('\n')
    }

    public currentTimestamp(): QuerySafeValue {
        return raw('NOW()')
    }
}