SponsorBlockServer/src/databases/Postgres.ts

import { Logger } from '../utils/logger';
import { IDatabase, QueryType } from './IDatabase';
import { Client, Pool, types } from 'pg';

import fs from "fs";

// return numeric (pg_type oid=1700) as float
types.setTypeParser(1700, function(val) {
    return parseFloat(val);
});

// return int8 (pg_type oid=20) as int
types.setTypeParser(20, function(val) {
  return parseInt(val, 10);
});

export class Postgres implements IDatabase {
    private pool: Pool;

    constructor(private config: Record<string, any>) {}

    async init(): Promise<void> {
        this.pool = new Pool(this.config.postgres);

        if (!this.config.readOnly) {
            if (this.config.createDbIfNotExists) {
                await this.createDB();
            }

            if (this.config.createDbIfNotExists && !this.config.readOnly && fs.existsSync(this.config.dbSchemaFileName)) {
                await this.pool.query(this.processUpgradeQuery(fs.readFileSync(this.config.dbSchemaFileName).toString()));
            }

            // Upgrade database if required
            await this.upgradeDB(this.config.fileNamePrefix, this.config.dbSchemaFolder);

            try {
                await this.applyIndexes(this.config.fileNamePrefix, this.config.dbSchemaFolder);
            } catch (e) {
                Logger.warn("Applying indexes failed. See https://github.com/ajayyy/SponsorBlockServer/wiki/Postgres-Extensions for more information.");
                Logger.warn(e);
            }
        }
    }

    async prepare(type: QueryType, query: string, params?: any[]): Promise<any[]> {
        // Convert query to use numbered parameters
        let count = 1;
        for (let char = 0; char < query.length; char++) {
            if (query.charAt(char) === '?') {
                query = query.slice(0, char) + "$" + count + query.slice(char + 1);
                count++;
            }
        }

        Logger.debug(`prepare (postgres): type: ${type}, query: ${query}, params: ${params}`);

        const queryResult = await this.pool.query({text: query, values: params});

        switch (type) {
            case 'get': {
                const value = queryResult.rows[0];
                Logger.debug(`result (postgres): ${JSON.stringify(value)}`);
                return value;
            }
            case 'all': {
                const values = queryResult.rows;
                Logger.debug(`result (postgres): ${values}`);
                return values;
            }
            case 'run': {
                break;
            }
        }
    }

    private async createDB() {
        const client = new Client({
            ...this.config.postgres,
            database: "postgres"
        });

        await client.connect();
        
        if ((await client.query(`SELECT * FROM pg_database WHERE datname = '${this.config.postgres.database}'`)).rowCount == 0) {
            await client.query(`CREATE DATABASE "${this.config.postgres.database}"
                                WITH 
                                OWNER = ${this.config.postgres.user}
                                ENCODING = 'UTF8'
                                LC_COLLATE = 'en_US.utf8'
                                LC_CTYPE = 'en_US.utf8'
                                TABLESPACE = pg_default
                                CONNECTION LIMIT = -1;`
            );
        }

        client.end();
    }

    private async upgradeDB(fileNamePrefix: string, schemaFolder: string) {
        const versionCodeInfo = await this.pool.query("SELECT value FROM config WHERE key = 'version'");
        let versionCode = versionCodeInfo.rows[0] ? versionCodeInfo.rows[0].value : 0;

        let path = schemaFolder + "/_upgrade_" + fileNamePrefix + "_" + (parseInt(versionCode) + 1) + ".sql";
        Logger.debug('db update: trying ' + path);
        while (fs.existsSync(path)) {
            Logger.debug('db update: updating ' + path);
            await this.pool.query(this.processUpgradeQuery(fs.readFileSync(path).toString()));

            versionCode = (await this.pool.query("SELECT value FROM config WHERE key = 'version'"))?.rows[0]?.value;
            path = schemaFolder + "/_upgrade_" + fileNamePrefix + "_" + (parseInt(versionCode) + 1) + ".sql";
            Logger.debug('db update: trying ' + path);
        }
        Logger.debug('db update: no file ' + path);
    }

    private async applyIndexes(fileNamePrefix: string, schemaFolder: string) {
        const path = schemaFolder + "/_" + fileNamePrefix + "_indexes.sql";
        if (fs.existsSync(path)) {
            await this.pool.query(fs.readFileSync(path).toString());
        } else {
            Logger.debug('failed to apply indexes to ' + fileNamePrefix);
        }
    }

    private processUpgradeQuery(query: string): string {
        let result = query;
        result = result.replace(/sha256\((.*?)\)/gm, "encode(digest($1, 'sha256'), 'hex')");
        result = result.replace(/integer/gmi, "NUMERIC");

        return result;
    }
}