All files / third-parties/drizzle-adapter/src/ydb-core/query-builders query-builder.ts

47.36% Statements 9/19
33.33% Branches 4/12
45.45% Functions 5/11
47.36% Lines 9/19

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103                                26x         20x       2x   2x   2x   2x                                                                                         18x                                     18x       18x      
import { entityKind } from 'drizzle-orm/entity'
import { SelectionProxyHandler } from 'drizzle-orm/selection-proxy'
import type { SQL, SQLWrapper } from 'drizzle-orm/sql/sql'
import { type Subquery, WithSubquery } from 'drizzle-orm/subquery'
import { YdbDialect } from '../../ydb/dialect.js'
import { YdbSelectBuilder } from './select.builder.js'
import type { SelectFields } from './select.types.js'
 
type YdbSelectableQuery =
	| SQL
	| {
			getSQL(): SQL
			getSelectedFields(): SelectFields | undefined
	  }
 
export class YdbQueryBuilder {
	static readonly [entityKind] = 'YdbQueryBuilder'
 
	#dialect: YdbDialect | undefined
 
	constructor(dialect?: YdbDialect) {
		this.#dialect = dialect
	}
 
	$with<TAlias extends string>(alias: TAlias) {
		return {
			as: (query: YdbSelectableQuery | ((qb: YdbQueryBuilder) => YdbSelectableQuery)) => {
				let resolved = typeof query === 'function' ? query(this) : query
				let selectedFields =
					'getSelectedFields' in resolved ? (resolved.getSelectedFields() ?? {}) : {}
 
				return new Proxy(
					new WithSubquery(resolved.getSQL(), selectedFields, alias, true),
					new SelectionProxyHandler({
						alias,
						sqlAliasedBehavior: 'alias',
						sqlBehavior: 'error',
					})
				)
			},
		}
	}
 
	with(...queries: Subquery[]) {
		let select = <TFields extends SelectFields | undefined = undefined>(fields?: TFields) =>
			new YdbSelectBuilder(undefined, this.#getDialect(), fields as any, {}, queries)
 
		let selectDistinct = <TFields extends SelectFields | undefined = undefined>(
			fields?: TFields
		) =>
			new YdbSelectBuilder(
				undefined,
				this.#getDialect(),
				fields as any,
				{ distinct: true },
				queries
			)
 
		let selectDistinctOn = <TFields extends SelectFields | undefined = undefined>(
			on: SQLWrapper | SQLWrapper[],
			fields?: TFields
		) =>
			new YdbSelectBuilder(
				undefined,
				this.#getDialect(),
				fields as any,
				{
					distinctOn: Array.isArray(on) ? on : [on],
				},
				queries
			)
 
		return { select, selectDistinct, selectDistinctOn }
	}
 
	select<TFields extends SelectFields | undefined = undefined>(fields?: TFields) {
		return new YdbSelectBuilder(undefined, this.#getDialect(), fields as any)
	}
 
	selectDistinct<TFields extends SelectFields | undefined = undefined>(fields?: TFields) {
		return new YdbSelectBuilder(undefined, this.#getDialect(), fields as any, {
			distinct: true,
		})
	}
 
	selectDistinctOn<TFields extends SelectFields | undefined = undefined>(
		on: SQLWrapper | SQLWrapper[],
		fields?: TFields
	) {
		return new YdbSelectBuilder(undefined, this.#getDialect(), fields as any, {
			distinctOn: Array.isArray(on) ? on : [on],
		})
	}
 
	#getDialect(): YdbDialect {
		Iif (!this.#dialect) {
			this.#dialect = new YdbDialect()
		}
 
		return this.#dialect
	}
}