import { Knex } from "knex" ; import { whereFilterHash } from "../dataloaders/findDataLoader" ; import { afterTransaction , beforeTransaction , buildQuery , deTagId , Entity , EntityConstructor , entityLimit , EntityManager , EntityMetadata , EnumArrayFieldSerde , FilterAndSettings , getMetadata , hasSerde , keyToNumber , maybeResolveReferenceToId , OneToManyCollection , tagIds , } from "../index" ; import { ManyToManyCollection , OneToOneReferenceImpl } from "../relations" ; import { JoinRow } from "../relations/ManyToManyCollection" ; import { JoinRowTodo , Todo } from "../Todo" ; import { getOrSet , partition , zeroTo } from "../utils" ; import { Driver } from "./driver" ; import { IdAssigner , SequenceIdAssigner } from "./IdAssigner" ; export interface PostgresDriverOpts { idAssigner ? : IdAssigner ; } /**
 * Implements the `Driver` interface for Postgres.
 *
 * This is the canonical driver implementation and leverages several aspects of
 * Postgres for the best performance, i.e.:
 *
 * - Deferred foreign key constraints are used to allow bulk inserting/updating
 * all entities without any topographic sorting/ordering issues.
 *
 * - Sequences are used to bulk-assign + then bulk-insert all entities, to again
 * avoid issues with ordering issues (cannot bulk insert A w/o knowing the id of B).
 *
 * - We use a pg-specific bulk update syntax.
 */ export class PostgresDriver implements Driver { private readonly idAssigner : IdAssigner ; constructor ( private readonly knex , opts ? ) { this . idAssigner = opts ?. idAssigner ?? new SequenceIdAssigner ( ) ; } load < T extends Entity > ( em , meta , untaggedIds ) { const knex = this . getMaybeInTxnKnex ( em ) ; return knex . select ( "*" ) . from ( meta . tableName ) . whereIn ( "id" , untaggedIds ) ; } loadManyToMany < T extends Entity , U extends Entity > ( em , collection , keys ) { const knex = this . getMaybeInTxnKnex ( em ) ; const columns = { } ; keys . forEach ( ( key ) => { const [ columnId , id ] = key . split ( "=" ) ; getOrSet ( columns , columnId , [ ] ) . push ( id ) ; } ) ; let query = knex . select ( "*" ) . from ( collection . joinTableName ) ; Object . entries ( columns ) . forEach ( ( [ columnId , values ] ) => { const meta = collection . columnName == columnId ? getMetadata ( collection . entity ) : collection . otherMeta ; query = query . orWhereIn ( columnId , values . map ( ( id ) => keyToNumber ( meta , id ) ! ) ) ; } ) ; return query . orderBy ( "id" ) ; } loadOneToMany < T extends Entity , U extends Entity > ( em , collection , untaggedIds ) { const knex = this . getMaybeInTxnKnex ( em ) ; return knex . select ( "*" ) . from ( collection . otherMeta . tableName ) . whereIn ( collection . otherColumnName , untaggedIds ) . orderBy ( "id" ) ; } loadOneToOne < T extends Entity , U extends Entity > ( em , reference , untaggedIds ) { const knex = this . getMaybeInTxnKnex ( em ) ; return knex . select ( "*" ) . from ( reference . otherMeta . tableName ) . whereIn ( reference . otherColumnName , untaggedIds ) . orderBy ( "id" ) ; } async find < T extends Entity > ( em , type , queries ) { const knex = this . getMaybeInTxnKnex ( em ) ; if ( queries . length === 1 ) { return [ ensureUnderLimit ( await buildQuery ( knex , type , queries [ 0 ] ) ) ] ; } const uniqueQueries = [ ] ; const queryToUnique = { } ; queries . forEach ( ( q , i ) => { let j = uniqueQueries . findIndex ( ( uq ) => whereFilterHash ( uq ) === whereFilterHash ( q ) ) ; if ( j === - 1 ) { uniqueQueries . push ( q ) ; j = uniqueQueries . length - 1 ; } queryToUnique [ i ] = j ; } ) ; if ( uniqueQueries . length === 1 ) { const rows = ensureUnderLimit ( await buildQuery ( knex , type , queries [ 0 ] ) ) ; return queries . map ( ( ) => rows ) ; } const tagged = uniqueQueries . map ( ( queryAndSettings , i ) => { const query = buildQuery ( knex , type , queryAndSettings ) as Knex . QueryBuilder ; return query . select ( knex . raw ( `${ i } as __tag` ) , knex . raw ( "row_number() over () as __row" ) ) ; } ) ; const meta = getMetadata ( type ) ; let query = knex . select ( "*" , knex . raw ( "-1 as __tag" ) , knex . raw ( "-1 as __row" ) ) . from ( meta . tableName ) . orderBy ( "__tag" , "__row" ) . where ( { id : - 1 } ) ; tagged . forEach ( ( add ) => { query = query . unionAll ( add , true ) ; } ) ; const rows = ensureUnderLimit ( await query ) ; const resultForUniques = [ ] ; uniqueQueries . forEach ( ( q , i ) => { resultForUniques [ i ] = [ ] ; } ) ; rows . forEach ( ( row ) => { resultForUniques [ row [ "__tag" ] ] . push ( row ) ; } ) ; const result = [ ] ; queries . forEach ( ( q , i ) => { result [ i ] = resultForUniques [ queryToUnique [ i ] ] ; } ) ; return result ; } async transaction < T > ( em , fn , isolationLevel ? ) { const knex = this . getMaybeInTxnKnex ( em ) ; const alreadyInTxn = "commit" in knex ; if ( alreadyInTxn ) { return fn ( knex as Knex . Transaction ) ; } return await knex . transaction ( async ( txn ) => { em . currentTxnKnex = txn ; try { if ( isolationLevel ) { await txn . raw ( "set transaction isolation level serializable;" ) ; } await beforeTransaction ( em , txn ) ; const result = await fn ( txn ) ; await afterTransaction ( em , txn ) ; return result ; } finally { em . currentTxnKnex = undefined ; } } ) ; } async flushEntities ( em , todos ) { const knex = this . getMaybeInTxnKnex ( em ) ; const updatedAt = new Date ( ) ; await this . idAssigner . assignNewIds ( knex , todos ) ; for await ( const todo of Object . values ( todos ) ) { if ( todo ) { const meta = todo . metadata ; if ( todo . inserts . length > 0 ) { await batchInsert ( knex , meta , todo . inserts ) ; } if ( todo . updates . length > 0 ) { todo . updates . forEach ( ( e ) => { e . __orm . originalData [ "updatedAt" ] = e . __orm . data [ "updatedAt" ] ; e . __orm . data [ "updatedAt" ] = updatedAt ; } ) ; await batchUpdate ( knex , meta , todo . updates ) ; } if ( todo . deletes . length > 0 ) { await batchDelete ( knex , meta , todo . deletes ) ; } } } } async flushJoinTables ( em , joinRows ) { const knex = this . getMaybeInTxnKnex ( em ) ; for await ( const [ joinTableName , { m2m , newRows , deletedRows } ] of Object . entries ( joinRows ) ) { if ( newRows . length > 0 ) { const sql = cleanSql ( `
          INSERT INTO ${ joinTableName } (${ m2m . columnName }, ${ m2m . otherColumnName })
          VALUES ${ zeroTo ( newRows . length ) . map ( ( ) => "(?, ?) " ) . join ( ", " ) }
          ON CONFLICT (${ m2m . columnName }, ${ m2m . otherColumnName }) DO UPDATE SET id = ${ joinTableName }.id
          RETURNING id;
        ` ) ; const meta1 = getMetadata ( m2m . entity ) ; const meta2 = m2m . otherMeta ; const bindings = newRows . flatMap ( ( row ) => { return [ keyToNumber ( meta1 , maybeResolveReferenceToId ( row [ m2m . columnName ] ) ) ! , keyToNumber ( meta2 , maybeResolveReferenceToId ( row [ m2m . otherColumnName ] ) ) ! , ] ; } ) ; const { rows } = await knex . raw ( sql , bindings ) ; for ( let i = 0 ; i < rows . length ; i ++ ) { newRows [ i ] . id = rows [ i ] . id ; } } if ( deletedRows . length > 0 ) { const [ haveIds , noIds ] = partition ( deletedRows , ( r ) => r . id !== - 1 ) ; if ( haveIds . length > 0 ) { await knex ( joinTableName ) . del ( ) . whereIn ( "id" , haveIds . map ( ( e ) => e . id ! ) ) ; } if ( noIds . length > 0 ) { const data = noIds . map ( ( e ) => [ deTagId ( m2m . meta , maybeResolveReferenceToId ( e [ m2m . columnName ] ) ! ) , deTagId ( m2m . otherMeta , maybeResolveReferenceToId ( e [ m2m . otherColumnName ] ) ! ) , ] as any ) ; await knex ( joinTableName ) . del ( ) . whereIn ( [ m2m . columnName , m2m . otherColumnName ] , data ) ; } } } } private getMaybeInTxnKnex ( em ) { return em . currentTxnKnex || this . knex ; } } async function batchInsert ( knex , meta , entities ) { const fields = Object . values ( meta . fields ) . filter ( hasSerde ) ; const columns = fields . flatMap ( ( f ) => f . serde . columns ) ; const bindings = entities . flatMap ( ( entity ) => columns . map ( ( c ) => c . dbValue ( entity . __orm . data ) ?? null ) ) ; const sql = `
    INSERT INTO ${ meta . tableName } (${ columns . map ( ( c ) => `"${ c . columnName }"` ) . join ( ", " ) })
    VALUES ${ entities . map ( ( ) => `(${ columns . map ( ( ) => `?` ) . join ( ", " ) })` ) . join ( "," ) }
  ` ; await knex . raw ( cleanSql ( sql ) , bindings ) ; } async function batchUpdate ( knex , meta , entities ) { const changedFields = new Set < string > ( ) ; changedFields . add ( "id" ) ; changedFields . add ( "updatedAt" ) ; entities . forEach ( ( entity ) => { Object . keys ( entity . __orm . originalData ) . forEach ( ( key ) => changedFields . add ( key ) ) ; } ) ; if ( changedFields . size === 2 && ! entities . some ( ( e ) => e . __orm . isTouched ) ) { return ; } const fields = Object . values ( meta . fields ) . filter ( ( f ) => changedFields . has ( f . fieldName ) ) . filter ( hasSerde ) ; const columns = fields . flatMap ( ( f ) => f . serde . columns ) ; let sql ; let bindings ; if ( fields . some ( ( f ) => f . serde instanceof EnumArrayFieldSerde ) ) { bindings = entities . flatMap ( ( entity ) => [ ... columns . map ( ( c ) => c . dbValue ( entity . __orm . data ) ?? null ) , entity . __orm . originalData . updatedAt , ] ) ; sql = `
      UPDATE ${ meta . tableName }
      SET ${ columns . filter ( ( c ) => c . columnName !== "id" ) . map ( ( c ) => `"${ c . columnName }" = data."${ c . columnName }"` ) . join ( ", " ) }
      FROM (
        VALUES ${ entities . map ( ( ) => `(${ columns . map ( ( c ) => `?::${ c . dbType }` ) . join ( ", " ) }, ?::timestamp)` ) . join ( "," ) }
      ) AS data(${ columns . map ( ( c ) => c . columnName ) . join ( "," ) },original_updated_at)
      WHERE
        ${ meta . tableName }.id = data.id
        AND date_trunc('milliseconds', ${ meta . tableName }.updated_at) = data.original_updated_at
      RETURNING ${ meta . tableName }.id
   ` ; } else { bindings = columns . map ( ( ) => [ ] ) ; const originalUpdatedAt = [ ] ; bindings . push ( originalUpdatedAt ) ; for ( const entity of entities ) { columns . forEach ( ( c , i ) => { bindings [ i ] . push ( c . dbValue ( entity . __orm . data ) ?? null ) ; } ) ; originalUpdatedAt . push ( entity . __orm . originalData . updatedAt ) ; } sql = `
      UPDATE ${ meta . tableName }
      SET ${ columns . filter ( ( c ) => c . columnName !== "id" ) . map ( ( c ) => `"${ c . columnName }" = data."${ c . columnName }"` ) . join ( ", " ) }
      FROM (
        SELECT
          ${ columns . map ( ( c ) => `unnest(?::${ c . dbType }[]) as "${ c . columnName }"` ) . join ( ", " ) },
          unnest(?::timestamptz[]) as original_updated_at
      ) as data
      WHERE
        ${ meta . tableName }.id = data.id
        AND date_trunc('milliseconds', ${ meta . tableName }.updated_at) = data.original_updated_at
      RETURNING ${ meta . tableName }.id
     ` ; } const result = await knex . raw ( cleanSql ( sql ) , bindings ) ; if ( result . rows . length !== entities . length ) { const updated = new Set ( result . rows . map ( ( r ) => r . id ) ) ; const missing = entities . map ( ( e ) => e . idOrFail ) . filter ( ( id ) => ! updated . has ( id ) ) ; throw new Error ( `Oplock failure for ${ tagIds ( meta , missing ) . join ( ", " ) }` ) ; } } async function batchDelete ( knex , meta , entities ) { await knex ( meta . tableName ) . del ( ) . whereIn ( "id" , entities . map ( ( e ) => keyToNumber ( meta , e . id ! ) . toString ( ) ) ) ; entities . forEach ( ( entity ) => ( entity . __orm . deleted = "deleted" ) ) ; } /** Strips new lines/indentation from our `UPDATE` string; doesn't do any actual SQL param escaping/etc. */ function cleanSql ( sql ) { return sql . trim ( ) . replace ( /\n/g , "" ) . replace ( /  +/g , " " ) ; } function ensureUnderLimit ( rows ) { if ( rows . length >= entityLimit ) { throw new Error ( `Query returned more than ${ entityLimit } rows` ) ; } return rows ; }