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

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
 
 
 
 
 
 
 
 
 
 
 
26x
3x
 
 
 
 
 
3x
3x
3x
3x
3x
 
 
 
3x
 
 
 
3x
 
 
 
 
 
 
2x
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
import { entityKind } from 'drizzle-orm/entity'
import { SQL, type SQLWrapper, sql as yql } from 'drizzle-orm/sql/sql'
import type { YdbSession } from '../session.js'
 
export interface YdbCountBuilderParams {
	source: SQLWrapper
	filters?: SQL | undefined
	session: Pick<YdbSession, 'count'>
}
 
export class YdbCountBuilder extends SQL<number> {
	static override readonly [entityKind] = 'YdbCountBuilder'
	readonly [Symbol.toStringTag] = 'YdbCountBuilder'
 
	readonly #session: Pick<YdbSession, 'count'>
	readonly #countSql: SQL<number>
 
	constructor(params: YdbCountBuilderParams) {
		let embeddedCount = YdbCountBuilder.buildEmbeddedCount(params.source, params.filters)
		super(embeddedCount.queryChunks)
		this.#session = params.session
		this.mapWith(Number)
		this.#countSql = YdbCountBuilder.buildCount(params.source, params.filters)
	}
 
	static buildEmbeddedCount(source: SQLWrapper, filters?: SQL): SQL<number> {
		return yql`(select count(*) from ${source}${yql.raw(' where ').if(filters)}${filters})`
	}
 
	static buildCount(source: SQLWrapper, filters?: SQL): SQL<number> {
		return yql`select count(*) as count from ${source}${yql.raw(' where ').if(filters)}${filters}`
	}
 
	then<TResult1 = number, TResult2 = never>(
		onfulfilled?: ((value: number) => TResult1 | PromiseLike<TResult1>) | null,
		onrejected?: ((reason: unknown) => TResult2 | PromiseLike<TResult2>) | null
	): Promise<TResult1 | TResult2> {
		return Promise.resolve(this.#session.count(this.#countSql)).then(onfulfilled, onrejected)
	}
 
	catch<TResult = never>(
		onrejected?: ((reason: unknown) => TResult | PromiseLike<TResult>) | null
	): Promise<number | TResult> {
		return this.then(undefined, onrejected)
	}
 
	finally(onfinally?: (() => void) | null): Promise<number> {
		return this.then(
			(value) => {
				onfinally?.()
				return value
			},
			(reason) => {
				onfinally?.()
				throw reason
			}
		)
	}
}