Drizzle Team•10mo ago

overriding DrizzleClient so it will add a string to all queries

I'm creating multi tenant application uses 1 database. So I need to create row based secuiry. I have DI system, I need to add "slug": "tenentSlug" to all insert, select, etc. queries. Is there a example for overriding default client? Which I can check and learn? I'm using postgresjs.

1 Reply

minester16OP•10mo ago

import { Param, SQL, Table } from 'drizzle-orm';
import {
  PgDialect,
  PgInsertBase,
  PgInsertBuilder,
  PgInsertValue,
  PgSession,
  PgTable,
  QueryResultHKT,
  SelectedFieldsOrdered,
} from 'drizzle-orm/pg-core';
import { is } from 'drizzle-orm';

export interface PgInsertConfig<TTable extends PgTable = PgTable> {
  table: TTable;
  values: Record<string, Param | SQL>[];
  onConflict?: SQL;
  returning?: SelectedFieldsOrdered;
}

export class PgInsertBuilderCustom<
  TTable extends PgTable,
  TQueryResult extends QueryResultHKT
> extends PgInsertBuilder<TTable, TQueryResult> {
  table: TTable;
  session: PgSession;
  dialect: PgDialect;

  constructor(table: TTable, session: PgSession, dialect: PgDialect) {
    super(table, session, dialect);
  }

  override values(
    values: PgInsertValue<TTable> | PgInsertValue<TTable>[]
  ): PgInsertBase<TTable, TQueryResult> {
    values = Array.isArray(values) ? values : [values];
    if (values.length === 0) {
      throw new Error('values() must be called with at least one value');
    }

    const mappedValues = values.map((entry) => {
      const result = {};
      const cols = this.table[Table.Symbol.Columns];
      for (const colKey of Object.keys(entry)) {
        const colValue = entry[colKey];
        result[colKey] = is(colValue, SQL)
          ? colValue
          : new Param(colValue, cols[colKey]);
      }
      return result;
    });
  }
}

import { Param, SQL, Table } from 'drizzle-orm';
import {
  PgDialect,
  PgInsertBase,
  PgInsertBuilder,
  PgInsertValue,
  PgSession,
  PgTable,
  QueryResultHKT,
  SelectedFieldsOrdered,
} from 'drizzle-orm/pg-core';
import { is } from 'drizzle-orm';

export interface PgInsertConfig<TTable extends PgTable = PgTable> {
  table: TTable;
  values: Record<string, Param | SQL>[];
  onConflict?: SQL;
  returning?: SelectedFieldsOrdered;
}

export class PgInsertBuilderCustom<
  TTable extends PgTable,
  TQueryResult extends QueryResultHKT
> extends PgInsertBuilder<TTable, TQueryResult> {
  table: TTable;
  session: PgSession;
  dialect: PgDialect;

  constructor(table: TTable, session: PgSession, dialect: PgDialect) {
    super(table, session, dialect);
  }

  override values(
    values: PgInsertValue<TTable> | PgInsertValue<TTable>[]
  ): PgInsertBase<TTable, TQueryResult> {
    values = Array.isArray(values) ? values : [values];
    if (values.length === 0) {
      throw new Error('values() must be called with at least one value');
    }

    const mappedValues = values.map((entry) => {
      const result = {};
      const cols = this.table[Table.Symbol.Columns];
      for (const colKey of Object.keys(entry)) {
        const colValue = entry[colKey];
        result[colKey] = is(colValue, SQL)
          ? colValue
          : new Param(colValue, cols[colKey]);
      }
      return result;
    });
  }
}

Example. I'm trying to override PgInsertBuilder. What I am doing wrong?

import {
  PgDialect,
  PgInsertBase,
  PgInsertBuilder,
  PgInsertValue,
  PgSession,
  PgTable,
  QueryResultHKT,
} from 'drizzle-orm/pg-core';
import { TenantService } from '../../shared/tenantService';

type HasSlug<T> = 'slug' extends keyof T ? Omit<T, 'slug'> : T;

export class PgInsertBuilderSlug<
  TTable extends PgTable,
  TQueryResult extends QueryResultHKT
> extends PgInsertBuilder<TTable, TQueryResult> {
  private insertBuilder: PgInsertBuilder<TTable, TQueryResult>;
  schemaT: TTable;

  constructor(
    table: TTable,
    session: PgSession,
    dialect: PgDialect,
    private readonly tenantService: TenantService
  ) {
    super(table, session, dialect);
    this.insertBuilder = new PgInsertBuilder(table, session, dialect);

    this.schemaT = table;
    // this.table = table;
    // this.session = session;
    // this.dialect = dialect;
  }

  // Example of a custom method that utilizes the PgInsertBuilder instance
  override values(
    values: HasSlug<PgInsertValue<TTable>> | HasSlug<PgInsertValue<TTable>>[]
  ): PgInsertBase<TTable, TQueryResult> {
    values = Array.isArray(values) ? values : [values];
    if (values.length === 0) {
      throw new Error('values() must be called with at least one value');
    }
    for (const value of values) {
      const hasSlugColumn = 'slug' in this.schemaT;
      if (hasSlugColumn) {
        const tenentSlug = this.tenantService.getTenantSlug();
        if (typeof tenentSlug !== 'string') {
          throw new Error('Tenant slug is not a string');
        }

        (value as { slug: string })['slug'] = tenentSlug;
      }
    }

    return this.insertBuilder.values(values);
  }
}

import {
  PgDialect,
  PgInsertBase,
  PgInsertBuilder,
  PgInsertValue,
  PgSession,
  PgTable,
  QueryResultHKT,
} from 'drizzle-orm/pg-core';
import { TenantService } from '../../shared/tenantService';

type HasSlug<T> = 'slug' extends keyof T ? Omit<T, 'slug'> : T;

export class PgInsertBuilderSlug<
  TTable extends PgTable,
  TQueryResult extends QueryResultHKT
> extends PgInsertBuilder<TTable, TQueryResult> {
  private insertBuilder: PgInsertBuilder<TTable, TQueryResult>;
  schemaT: TTable;

  constructor(
    table: TTable,
    session: PgSession,
    dialect: PgDialect,
    private readonly tenantService: TenantService
  ) {
    super(table, session, dialect);
    this.insertBuilder = new PgInsertBuilder(table, session, dialect);

    this.schemaT = table;
    // this.table = table;
    // this.session = session;
    // this.dialect = dialect;
  }

  // Example of a custom method that utilizes the PgInsertBuilder instance
  override values(
    values: HasSlug<PgInsertValue<TTable>> | HasSlug<PgInsertValue<TTable>>[]
  ): PgInsertBase<TTable, TQueryResult> {
    values = Array.isArray(values) ? values : [values];
    if (values.length === 0) {
      throw new Error('values() must be called with at least one value');
    }
    for (const value of values) {
      const hasSlugColumn = 'slug' in this.schemaT;
      if (hasSlugColumn) {
        const tenentSlug = this.tenantService.getTenantSlug();
        if (typeof tenentSlug !== 'string') {
          throw new Error('Tenant slug is not a string');
        }

        (value as { slug: string })['slug'] = tenentSlug;
      }
    }

    return this.insertBuilder.values(values);
  }
}

Okay managed to override insert like this. Now I'm trying to override where method in select feature. That's quite harder because in source PgSelectBase class is written as abstract class. So I'm not able to create new. What should I do?

Gaming

Programming

overriding DrizzleClient so it will add a string to all queries