immich/server/src/bin/sync-sql.ts

#!/usr/bin/env node
import { INestApplication } from '@nestjs/common';
import { Reflector } from '@nestjs/core';
import { SchedulerRegistry } from '@nestjs/schedule';
import { Test } from '@nestjs/testing';
import { ClassConstructor } from 'class-transformer';
import { ClsModule } from 'nestjs-cls';
import { KyselyModule } from 'nestjs-kysely';
import { OpenTelemetryModule } from 'nestjs-otel';
import { mkdir, rm, writeFile } from 'node:fs/promises';
import { join } from 'node:path';
import { format } from 'sql-formatter';
import { GENERATE_SQL_KEY, GenerateSqlQueries } from 'src/decorators';
import { repositories } from 'src/repositories';
import { AccessRepository } from 'src/repositories/access.repository';
import { ConfigRepository } from 'src/repositories/config.repository';
import { LoggingRepository } from 'src/repositories/logging.repository';
import { AuthService } from 'src/services/auth.service';
import { getKyselyConfig } from 'src/utils/database';

const handleError = (label: string, error: Error | any) => {
  console.error(`${label} error: ${error}`);
};

export class SqlLogger {
  queries: string[] = [];
  errors: Array<{ error: string | Error; query: string }> = [];

  clear() {
    this.queries = [];
    this.errors = [];
  }

  logQuery(query: string) {
    this.queries.push(format(query, { language: 'postgresql' }));
  }

  logQueryError(error: string | Error, query: string) {
    this.errors.push({ error, query });
  }
}

const reflector = new Reflector();

type Repository = ClassConstructor<any>;
type SqlGeneratorOptions = { targetDir: string };

class SqlGenerator {
  private app: INestApplication | null = null;
  private sqlLogger = new SqlLogger();
  private results: Record<string, string[]> = {};

  constructor(private options: SqlGeneratorOptions) {}

  async run() {
    try {
      await this.setup();
      for (const Repository of repositories) {
        if (Repository === LoggingRepository) {
          continue;
        }
        await this.process(Repository);
      }
      await this.write();
      this.stats();
    } finally {
      await this.close();
    }
  }

  private async setup() {
    await rm(this.options.targetDir, { force: true, recursive: true });
    await mkdir(this.options.targetDir);

    if (!process.env.DB_HOSTNAME) {
      process.env.DB_HOSTNAME = 'localhost';
    }
    const { database, cls, otel } = new ConfigRepository().getEnv();

    const moduleFixture = await Test.createTestingModule({
      imports: [
        KyselyModule.forRoot({
          ...getKyselyConfig(database.config),
          log: (event) => {
            if (event.level === 'query') {
              this.sqlLogger.logQuery(event.query.sql);
            } else if (event.level === 'error') {
              this.sqlLogger.logQueryError(event.error as Error, event.query.sql);
              this.sqlLogger.logQuery(event.query.sql);
            }
          },
        }),
        ClsModule.forRoot(cls.config),
        OpenTelemetryModule.forRoot(otel),
      ],
      providers: [...repositories, AuthService, SchedulerRegistry],
    }).compile();

    this.app = await moduleFixture.createNestApplication().init();
  }

  async process(Repository: Repository) {
    if (!this.app) {
      throw new Error('Not initialized');
    }

    const data: string[] = [`-- NOTE: This file is auto generated by ./sql-generator`];
    const instance = this.app.get<Repository>(Repository);

    // normal repositories
    data.push(...(await this.runTargets(instance, `${Repository.name}`)));

    // nested repositories
    if (Repository.name === AccessRepository.name) {
      for (const key of Object.keys(instance)) {
        const subInstance = (instance as any)[key];
        data.push(...(await this.runTargets(subInstance, `${Repository.name}.${key}`)));
      }
    }

    this.results[Repository.name] = data;
  }

  private async runTargets(instance: any, label: string) {
    const data: string[] = [];

    for (const key of this.getPropertyNames(instance)) {
      const target = instance[key];
      if (!(typeof target === 'function')) {
        continue;
      }

      const queries = reflector.get<GenerateSqlQueries[] | undefined>(GENERATE_SQL_KEY, target);
      if (!queries) {
        continue;
      }

      // empty decorator implies calling with no arguments
      if (queries.length === 0) {
        queries.push({ params: [] });
      }

      for (const { name, params, stream } of queries) {
        let queryLabel = `${label}.${key}`;
        if (name) {
          queryLabel += ` (${name})`;
        }

        this.sqlLogger.clear();

        if (stream) {
          try {
            const result: AsyncIterableIterator<unknown> = target.apply(instance, params);
            for await (const _ of result) {
              break;
            }
          } catch (error) {
            handleError(queryLabel, error);
          }
        } else {
          // errors still generate sql, which is all we care about
          await target.apply(instance, params).catch((error: Error) => handleError(queryLabel, error));
        }

        if (this.sqlLogger.queries.length === 0) {
          console.warn(`No queries recorded for ${queryLabel}`);
          continue;
        }

        data.push([`-- ${queryLabel}`, ...this.sqlLogger.queries].join('\n'));
      }
    }

    return data;
  }

  private async write() {
    for (const [repoName, data] of Object.entries(this.results)) {
      // only contains the header
      if (data.length === 1) {
        continue;
      }
      const filename = repoName.replaceAll(/[A-Z]/g, (letter) => `.${letter.toLowerCase()}`).replace('.', '');
      const file = join(this.options.targetDir, `${filename}.sql`);
      await writeFile(file, data.join('\n\n') + '\n');
    }
  }

  private stats() {
    console.log(`Wrote ${Object.keys(this.results).length} files`);
    console.log(`Generated ${Object.values(this.results).flat().length} queries`);
  }

  private async close() {
    if (this.app) {
      await this.app.close();
    }
  }

  private getPropertyNames(instance: any): string[] {
    return Object.getOwnPropertyNames(Object.getPrototypeOf(instance)) as any[];
  }
}

new SqlGenerator({ targetDir: './src/queries' })
  .run()
  .then(() => {
    console.log('Done');
    process.exit(0);
  })
  .catch((error) => {
    console.error(error);
    console.log('Something went wrong');
    process.exit(1);
  });