From 8dd2832a2aef579abddbd5bfcde314965de1eba2 Mon Sep 17 00:00:00 2001 From: romeerez Date: Sat, 22 Jun 2024 17:17:01 +0200 Subject: [PATCH] JS computed columns (#270) also, change `returnType` type to be `undefined` by default. (#299) --- .changeset/silly-waves-itch.md | 9 + BREAKING_CHANGES.md | 18 + docs/src/.vitepress/config.ts | 4 + docs/src/guide/computed-columns.md | 212 ++++++++++ docs/src/guide/orm-and-query-builder.md | 5 - packages/core/src/query.ts | 20 +- packages/core/src/raw.ts | 8 +- packages/core/src/utils.ts | 7 - packages/orm/src/baseTable.test.ts | 53 +-- packages/orm/src/orm.test.ts | 16 +- packages/orm/src/orm.ts | 6 +- packages/orm/src/relations/belongsTo.ts | 15 +- .../orm/src/relations/hasAndBelongsToMany.ts | 4 +- packages/orm/src/repo.ts | 2 +- packages/qb/pqb/src/columns/unknown.ts | 4 +- .../pqb/src/common/queryResultProcessing.ts | 66 ++++ packages/qb/pqb/src/common/utils.ts | 1 - packages/qb/pqb/src/modules/computed.test.ts | 240 +++++++++++- packages/qb/pqb/src/modules/computed.ts | 141 ++++++- packages/qb/pqb/src/query/db.ts | 2 +- packages/qb/pqb/src/query/query.ts | 11 +- packages/qb/pqb/src/queryMethods/create.ts | 15 +- packages/qb/pqb/src/queryMethods/from.ts | 2 + packages/qb/pqb/src/queryMethods/get.utils.ts | 16 +- .../qb/pqb/src/queryMethods/join/_join.ts | 37 +- .../qb/pqb/src/queryMethods/json.utils.ts | 2 +- packages/qb/pqb/src/queryMethods/map.ts | 8 +- .../qb/pqb/src/queryMethods/merge.test.ts | 10 +- packages/qb/pqb/src/queryMethods/merge.ts | 3 +- .../pqb/src/queryMethods/queryMethods.test.ts | 7 + .../qb/pqb/src/queryMethods/queryMethods.ts | 13 +- packages/qb/pqb/src/queryMethods/select.ts | 370 +++++++++++++----- packages/qb/pqb/src/queryMethods/sql.ts | 2 +- packages/qb/pqb/src/queryMethods/then.ts | 240 +++++++----- packages/qb/pqb/src/queryMethods/update.ts | 3 +- packages/qb/pqb/src/sql/data.ts | 20 +- packages/qb/pqb/src/sql/rawSql.ts | 11 +- packages/qb/pqb/src/sql/select.ts | 19 +- packages/qb/pqb/src/sql/toSQL.ts | 2 +- 39 files changed, 1309 insertions(+), 315 deletions(-) create mode 100644 .changeset/silly-waves-itch.md create mode 100644 docs/src/guide/computed-columns.md create mode 100644 packages/qb/pqb/src/common/queryResultProcessing.ts diff --git a/.changeset/silly-waves-itch.md b/.changeset/silly-waves-itch.md new file mode 100644 index 000000000..9658d206d --- /dev/null +++ b/.changeset/silly-waves-itch.md @@ -0,0 +1,9 @@ +--- +'pqb': minor +'orchid-core': minor +'orchid-orm': minor +--- + +JS computed columns; + +Change `returnType` type to be `undefined` by default. diff --git a/BREAKING_CHANGES.md b/BREAKING_CHANGES.md index 67f53adac..3f3c1c2bb 100644 --- a/BREAKING_CHANGES.md +++ b/BREAKING_CHANGES.md @@ -1,5 +1,23 @@ # Breaking changes +## orchid-orm 1.31.0 + +Computed columns change, see the [docs](https://orchid-orm.netlify.app/guide/computed-columns.html). + +For SQL computed columns: + +```ts +// before +computed = this.setComputed({ + fullName: (q) => q.sql`...`.type((t) => t.string()), +}); + +// after +computed = this.setComputed((q) => ({ + fullName: q.sql`...`.type((t) => t.string()), +})); +``` + ## orchid-orm 1.30.0 The `text` column type no longer accepts `min` and `max` params. diff --git a/docs/src/.vitepress/config.ts b/docs/src/.vitepress/config.ts index ac8b85020..b9710a987 100644 --- a/docs/src/.vitepress/config.ts +++ b/docs/src/.vitepress/config.ts @@ -78,6 +78,10 @@ export default { text: 'JSON functions', link: '/guide/json', }, + { + text: 'Computed columns', + link: '/guide/computed-columns', + }, { text: 'Window functions', link: '/guide/window', diff --git a/docs/src/guide/computed-columns.md b/docs/src/guide/computed-columns.md new file mode 100644 index 000000000..6ddfaa92d --- /dev/null +++ b/docs/src/guide/computed-columns.md @@ -0,0 +1,212 @@ +# Computed columns + +OrchidORM supports defining columns that are calculated on the fly, +either by injecting SQL into a `SELECT` statement, or by computing values in runtime on JS side. + +Note that unlike regular columns, computed columns are not selected by default. + +Alternatively, you can add a generated column in the migration (see [generated](/guide/migration-column-methods#generated-column)), +such column will persist in the database. + +## SQL computed column + +SQL computed column is going to unwrap into the given SQL when selecting it from the table. + +In the following example, selecting `fullName` will unwrap into `"firstName" || ' ' || "lastName"` SQL: + +```ts +export class UserTable extends BaseTable { + readonly table = 'user'; + columns = this.setColumns((t) => ({ + id: t.identity().primaryKey(), + firstName: t.string(), + lastName: t.string(), + })); + + computed = this.setComputed((q) => ({ + fullName: q.sql`${q.column('firstName')} || ' ' || ${q.column( + 'lastName', + )}`.type((t) => t.string()), + randomizedName: q + .sql(() => q.sql`${Math.random()} ${q.column('firstName')}`) + .type((t) => t.string()), + })); +} +``` + +`randomizedName` in the example is defined with `` q.sql(() => q.sql`...`) `` syntax that makes it dynamic, +so that a new random value will be selected for every query. + +Such can be column can be selected, can be used for filtering and ordering, available in nested sub-queries. + +```ts +// select all columns + the computed +db.user.select('*', 'fullName') + +// use in nested select +db.chat.find(id).select({ + messages: (q) => q.messages.select({ + // select fullName for a single row + sender: (q) => q.sender.select('fullName') + // `pluck` will load a flat array of values + receipients: (q) => + q.receipients + .pluck('fullName') + // works for filtering + .where({ fullName: { startsWith: 'x' } }) + // works for ordering + .order('fullName'), + }) +}) + +// can be selected for a joined table +db.post.join('author').select('author.fullName') + +// can be returned from `insert`, `create`, `update`, `delete`, `upsert` +db.user.select('fullName').insert(data) +``` + +## JS runtime computed + +Define a runtime computed column to compute values after loading results. + +Unlike SQL computed columns, these columns aren't suitable for filtering or ordering records, they only can be used in selects. + +```ts +export class UserTable extends BaseTable { + readonly table = 'user'; + columns = this.setColumns((t) => ({ + id: t.identity().primaryKey(), + firstName: t.string(), + lastName: t.string(), + })); + + computed = this.setComputed((q) => ({ + fullName: q.computeAtRuntime( + // define columns that it depends on + ['firstName', 'lastName'], + // only columns defined above are available in the callback + (record) => `${record.firstName} ${record.lastName}`, + ), + })); +} +``` + +The runtime computed column is available in all kinds of selections. + +It will automatically select dependencies, if they weren't selected, +and will dispose dependencies after computing a value if they weren't selected. + +```ts +const record = await db.user.select('firstName', 'fullName'); +record.firstName; // was selected +record.fullName; // was computed +record.lastName; // TS error: it was selected but then disposed + +db.char.find(id).select({ + messages: (q) => q.messages.select({ + // select fullName for a single row + sender: (q) => q.sender.select('fullName') + // `pluck` will collect a flat array of values + receipients: (q) => q.receipients.pluck('fullName') + }) +}) + +// can be selected for a joined table +db.post.join('author').select('author.fullName') + +// can be returned from `insert`, `create`, `update`, `delete`, `upsert` +db.user.select('fullName').insert(data) +``` + +## Async computed columns + +Asynchronously fetching data for records one-by-one would take a lot of loading time, +it's much better to load data in batches. + +```ts +interface WeatherData { + country: string; + city: string; + weatherInfo: SomeStructure; +} + +export class UserTable extends BaseTable { + readonly table = 'user'; + columns = this.setColumns((t) => ({ + id: t.identity().primaryKey(), + country: t.string(), + city: t.string(), + })); + + computed = this.setComputed((q) => ({ + weather: q.computeBatchAtRuntime( + // define columns that it depends on + ['country', 'city'], + // load weather data for all users using a single fetch + async (users): Promise<(SomeStructure | undefined)[]> => { + // to not query the same location twice + const uniqueLocations = new Set( + users.map((user) => `${user.country} ${user.city}`), + ); + + // fetch data for all locations at once + const weatherData: WeatherData[] = await fetchWeatherData({ + location: [...uniqueLocations], + }); + + // return array with weather data for every user + return users.map( + (user) => + weatherData.find( + (wd) => wd.country === user.country && wd.city === user.city, + )?.weatherInfo, + ); + }, + ), + })); +} +``` + +`computeBatchAtRuntime` can also take a synchronous function. + +From a querying perspective, there is no difference from a [computeAtRuntime](#js-runtime-computed) column, +it works and acts in the same way. + +```ts +db.user.select('*', 'weather'); + +// a city can have millions of people, +// but the weather is loaded just once +db.city.find(id).select({ + users: (q) => q.users.select('name', 'weather'), +}); +``` + +Only a single batch of records is processed even when loading a nested query. + +Let's say we have 10 countries, every country has 10 cities, with 100 users in each. + +The `weather` computed column will be called just once with 10_000 of records. + +```ts +db.country.select({ + cities: (q) => + q.cities.select({ + users: (q) => q.users.select('name', 'weather'), + }), +}); +``` + +A city may have a mayor, but that's not always the case. +Null records are omitted when passing data to a computed column. + +```ts +db.country.select({ + cities: (q) => + q.cities.select({ + // city hasOne mayor, not required + mayor: (q) => q.mayor.select('name', 'weather')., + }), +}); +``` diff --git a/docs/src/guide/orm-and-query-builder.md b/docs/src/guide/orm-and-query-builder.md index 99d2f46ae..ff11987f0 100644 --- a/docs/src/guide/orm-and-query-builder.md +++ b/docs/src/guide/orm-and-query-builder.md @@ -852,11 +852,6 @@ await db.some.unscope('default'); [//]: # 'has JSDoc' -You can add a generated column in the migration (see [generated](/guide/migration-column-methods#generated-column)), -such column will persist in the database, it can be indexed. - -Or you can add a computed column on the ORM level, without adding it to the database, in such a way: - ```ts import { BaseTable } from './baseTable'; diff --git a/packages/core/src/query.ts b/packages/core/src/query.ts index d78824c6d..51cc65531 100644 --- a/packages/core/src/query.ts +++ b/packages/core/src/query.ts @@ -3,6 +3,7 @@ import { TransactionState } from './adapter'; import { EmptyObject, FnUnknownToUnknown, + MaybePromise, RecordKeyTrue, RecordUnknown, } from './utils'; @@ -77,6 +78,7 @@ export type CoreQueryScopes = { }; export type QueryReturnType = + | undefined | 'all' | 'one' | 'oneOrThrow' @@ -186,11 +188,25 @@ export const getValueKey = Symbol('get'); // function to parse a single column after loading the data export type ColumnParser = FnUnknownToUnknown; -// functions to parse columns after loading the data +// To parse all returned rows. Unlike column parser, can return a promise. +export interface BatchParser { + path: string[]; + fn: (path: string[], queryResult: { rows: unknown[] }) => MaybePromise; +} + +// set of value parsers // key is a name of a selected column, // or it can be a `getValueKey` to parse single values requested by the `.get()`, `.count()`, or similar methods export type ColumnsParsers = { [K in string | getValueKey]?: ColumnParser }; +// set of batch parsers +// is only triggered when loading all, +// or when using `hookSelect` or computed columns that convert response to `all` internally. +// key is a name of a selected column, +// or it can be a `getValueKey` to parse single values requested by the `.get()`, `.count()`, or similar methods +export type BatchParsers = BatchParser[]; + +// result transformer: function for `transform`, object for `map` export type QueryDataTransform = | FnUnknownToUnknown | { map: FnUnknownToUnknown }; @@ -246,7 +262,7 @@ export const applyTransforms = ( ): unknown => { for (const fn of fns) { if ('map' in fn) { - if (returnType === 'all') { + if (!returnType || returnType === 'all') { result = (result as unknown[]).map(fn.map); } else { result = fn.map(result); diff --git a/packages/core/src/raw.ts b/packages/core/src/raw.ts index 02b40d0e0..1b9c4b24b 100644 --- a/packages/core/src/raw.ts +++ b/packages/core/src/raw.ts @@ -64,13 +64,13 @@ export const isTemplateLiteralArgs = ( // Argument type for `sql` function. // It can take a template literal, an object `{ raw: string, values?: Record }`, // or a function to build SQL lazily. -export type SQLArgs = StaticSQLArgs | [DynamicSQLArg]; +export type SQLArgs = StaticSQLArgs | [DynamicSQLArg]; // Function for sql method to build SQL lazily (dynamically). // May be used for computed column to build a different SQL in different executions. -export type DynamicSQLArg = ( - sql: (...args: StaticSQLArgs) => Expression, -) => Expression; +export type DynamicSQLArg = ( + sql: (...args: StaticSQLArgs) => Expression, +) => Expression; // SQL arguments for a non-lazy SQL expression. export type StaticSQLArgs = diff --git a/packages/core/src/utils.ts b/packages/core/src/utils.ts index 0e890a68e..a9fc784f7 100644 --- a/packages/core/src/utils.ts +++ b/packages/core/src/utils.ts @@ -353,13 +353,6 @@ export const callWithThis = function (this: T, cb: (arg: T) => R): R { return cb(this); }; -export const cloneInstance = (instance: T): T => { - return Object.assign( - Object.create(Object.getPrototypeOf(instance)), - instance, - ); -}; - export const pick = ( obj: T, keys: Keys[], diff --git a/packages/orm/src/baseTable.test.ts b/packages/orm/src/baseTable.test.ts index ddc057596..193b96c56 100644 --- a/packages/orm/src/baseTable.test.ts +++ b/packages/orm/src/baseTable.test.ts @@ -576,35 +576,36 @@ describe('baseTable', () => { }); describe('select', () => { - it.only('should select record with computed', async () => { - const q = local.profile - .select({ - user: (q) => - q.user.select('sqlComputed', 'runtimeComputed', 'batchComputed'), - }) - .take(); - - expectSql( - q.toSQL(), - ` - SELECT row_to_json("user".*) "user" - FROM "profile" - LEFT JOIN LATERAL ( - SELECT - "user"."name" || ' ' || "user"."userKey" "sqlComputed", - "user"."id" "Id", - "user"."name" "Name" - FROM "user" - WHERE "user"."id" = "profile"."userId" - AND "user"."userKey" = "profile"."profileKey" - ) "user" ON true - LIMIT 1 - `, - ); + it('should select record with computed', async () => { + const q = local.profile.select({ + user: (q) => + q.user.select('sqlComputed', 'runtimeComputed', 'batchComputed'), + }); const res = await q; - console.log(res.user); + assertType< + typeof res, + { + user: + | { + sqlComputed: string; + runtimeComputed: string; + batchComputed: string; + } + | undefined; + }[] + >(); + + expect(res).toEqual([ + { + user: { + sqlComputed: `${userData.Name} ${userData.UserKey}`, + runtimeComputed: `${userId} ${userData.Name}`, + batchComputed: `${userId} ${userData.Name}`, + }, + }, + ]); }); }); }); diff --git a/packages/orm/src/orm.test.ts b/packages/orm/src/orm.test.ts index f1e102fbb..adfa55e19 100644 --- a/packages/orm/src/orm.test.ts +++ b/packages/orm/src/orm.test.ts @@ -149,8 +149,15 @@ describe('orm', () => { describe('$from', () => { it('should have method `$from` with proper handling of type, where operators, parsers', async () => { const ChatId = await db.chat.get('IdOfChat').create(chatData); - const AuthorId = await db.user.get('Id').create(userData); - await db.message.insert({ ...messageData, ChatId, AuthorId }); + const [AuthorId1, AuthorId2] = await db.user + .pluck('Id') + .insertMany([userData, userData]); + + await db.message.createMany([ + { ...messageData, ChatId, AuthorId: AuthorId1 }, + { ...messageData, ChatId, AuthorId: AuthorId2 }, + { ...messageData, ChatId, AuthorId: AuthorId2 }, + ]); const inner = db.user.select('createdAt', { alias: 'Name', @@ -192,6 +199,11 @@ describe('orm', () => { alias: 'name', messagesCount: 1, }, + { + createdAt: expect.any(Date), + alias: 'name', + messagesCount: 2, + }, ]); }); }); diff --git a/packages/orm/src/orm.ts b/packages/orm/src/orm.ts index 46157ab4d..ea508e071 100644 --- a/packages/orm/src/orm.ts +++ b/packages/orm/src/orm.ts @@ -27,6 +27,10 @@ import { TransactionState, } from 'orchid-core'; +interface FromQuery extends Query { + returnType: 'all'; +} + export type OrchidORM = { [K in keyof T]: DbTable>; } & { @@ -98,7 +102,7 @@ export type OrchidORM = { */ $from>>( arg: Arg, - ): FromResult; + ): FromResult; $close(): Promise; }; diff --git a/packages/orm/src/relations/belongsTo.ts b/packages/orm/src/relations/belongsTo.ts index 73fc8b8bb..3757a2996 100644 --- a/packages/orm/src/relations/belongsTo.ts +++ b/packages/orm/src/relations/belongsTo.ts @@ -364,10 +364,10 @@ const nestedInsert = ({ query, primaryKeys }: State) => { .create; } - created = await _queryCreateMany( + created = (await _queryCreateMany( t.select(...primaryKeys), items as never, - ); + )) as never; } else { created = emptyArray; } @@ -442,7 +442,10 @@ const nestedUpdate = ({ query, primaryKeys, foreignKeys, len }: State) => { } else if (params.create) { const q = query.clone(); q.q.select = primaryKeys; - const record = await _queryCreate(q, params.create); + const record = (await _queryCreate( + q, + params.create, + )) as unknown as RecordUnknown; for (let i = 0; i < len; i++) { update[foreignKeys[i]] = record[primaryKeys[i]]; } @@ -489,7 +492,11 @@ const nestedUpdate = ({ query, primaryKeys, foreignKeys, len }: State) => { typeof upsert.create === 'function' ? upsert.create() : upsert.create; - const result = await _queryCreate(query.select(...primaryKeys), data); + + const result = (await _queryCreate( + query.select(...primaryKeys), + data, + )) as unknown as RecordUnknown; const collectData: RecordUnknown = {}; state.collect = { diff --git a/packages/orm/src/relations/hasAndBelongsToMany.ts b/packages/orm/src/relations/hasAndBelongsToMany.ts index 40a1ab74e..90a1cbfd8 100644 --- a/packages/orm/src/relations/hasAndBelongsToMany.ts +++ b/packages/orm/src/relations/hasAndBelongsToMany.ts @@ -583,10 +583,10 @@ const nestedInsert = ({ } } - created = await _queryCreateMany( + created = (await _queryCreateMany( t.select(...throughPrimaryKeys), records, - ); + )) as never; } else { created = []; } diff --git a/packages/orm/src/repo.ts b/packages/orm/src/repo.ts index 32f00e2e8..3a7df828b 100644 --- a/packages/orm/src/repo.ts +++ b/packages/orm/src/repo.ts @@ -9,7 +9,7 @@ type QueryMethods = Record< type QueryOne = { [K in keyof T]: K extends 'returnType' - ? Exclude + ? Exclude : T[K]; }; diff --git a/packages/qb/pqb/src/columns/unknown.ts b/packages/qb/pqb/src/columns/unknown.ts index feeb94b1d..550012928 100644 --- a/packages/qb/pqb/src/columns/unknown.ts +++ b/packages/qb/pqb/src/columns/unknown.ts @@ -7,9 +7,11 @@ import { defaultSchemaConfig } from './defaultSchemaConfig'; export class UnknownColumn< Schema extends ColumnSchemaConfig, > extends VirtualColumn { + static instance = new UnknownColumn(defaultSchemaConfig); + constructor(schema: Schema) { super(schema, schema.unknown() as never); } } -RawSQL.prototype.result = { value: new UnknownColumn(defaultSchemaConfig) }; +RawSQL.prototype.result = { value: UnknownColumn.instance }; diff --git a/packages/qb/pqb/src/common/queryResultProcessing.ts b/packages/qb/pqb/src/common/queryResultProcessing.ts new file mode 100644 index 000000000..0905572db --- /dev/null +++ b/packages/qb/pqb/src/common/queryResultProcessing.ts @@ -0,0 +1,66 @@ +import { QueryData } from '../sql'; +import { applyTransforms, QueryReturnType, RecordString } from 'orchid-core'; +import { QueryBatchResult } from '../queryMethods'; + +export const applyBatchTransforms = ( + query: QueryData, + batches: QueryBatchResult[], +) => { + if (query.transform) { + for (const item of batches) { + item.parent[item.key] = applyTransforms( + query.returnType, + query.transform, + item.data, + ); + } + } +}; + +export const finalizeNestedHookSelect = ( + batches: QueryBatchResult[], + returnType: QueryReturnType, + tempColumns: Set | undefined, + renames: RecordString | undefined, + key: string, +) => { + if (renames) { + for (const { data } of batches) { + for (const record of data) { + if (record) { + for (const a in renames) { + record[a] = record[renames[a]]; + } + } + } + } + } + + if (tempColumns?.size) { + for (const { data } of batches) { + for (const record of data) { + if (record) { + for (const key of tempColumns) { + delete record[key]; + } + } + } + } + } + + if (returnType === 'one' || returnType === 'oneOrThrow') { + for (const batch of batches) { + batch.data = batch.data[0]; + } + } else if (returnType === 'pluck') { + for (const { data } of batches) { + for (let i = 0; i < data.length; i++) { + data[i] = data[i][key]; + } + } + } else if (returnType === 'value' || returnType === 'valueOrThrow') { + for (const item of batches) { + item.parent[item.key] = item.data[0]?.[key]; + } + } +}; diff --git a/packages/qb/pqb/src/common/utils.ts b/packages/qb/pqb/src/common/utils.ts index 5f41601c1..5fa2b9cbf 100644 --- a/packages/qb/pqb/src/common/utils.ts +++ b/packages/qb/pqb/src/common/utils.ts @@ -36,7 +36,6 @@ export type ExpressionOutput< export const getClonedQueryData = (query: QueryData): QueryData => { const cloned = { ...query }; delete cloned[toSQLCacheKey]; - if (cloned.parsers) cloned.parsers = { ...cloned.parsers }; cloneQuery(cloned); return cloned as QueryData; }; diff --git a/packages/qb/pqb/src/modules/computed.test.ts b/packages/qb/pqb/src/modules/computed.test.ts index 5add5fcd8..e76c4e9fd 100644 --- a/packages/qb/pqb/src/modules/computed.test.ts +++ b/packages/qb/pqb/src/modules/computed.test.ts @@ -6,6 +6,7 @@ import { userData as partialUserData, } from '../test-utils/test-utils'; import { Query } from '../query/query'; +import { NotFoundError } from '../errors'; const User = testDb( 'user', @@ -18,9 +19,9 @@ const User = testDb( undefined, { computed: (q) => ({ - nameAndKey: q.sql`${q.column('name')} || ' ' || ${q.column( - 'userKey', - )}`.type((t) => t.string()), + nameAndKey: q + .sql(() => q.sql`${q.column('name')} || ' ' || ${q.column('userKey')}`) + .type((t) => t.string()), runtimeComputed: q.computeAtRuntime( ['id', 'name'], (record) => `${record.id} ${record.name}`, @@ -647,6 +648,239 @@ describe('computed', () => { }); }); + describe('sub-select', () => { + it('should select many', async () => { + const res = await User.select({ + users: () => + User.select({ + users: () => User.select('runtimeComputed', 'batchComputed'), + }), + }); + + assertType< + typeof res, + { + users: { + users: { runtimeComputed: string; batchComputed: string }[]; + }[]; + }[] + >(); + + expect(res).toEqual([ + { + users: [ + { + users: [ + { + runtimeComputed: `${userId} ${userData.name}`, + batchComputed: `${userId} ${userData.name}`, + }, + ], + }, + ], + }, + ]); + }); + + it('should select one optional', async () => { + const res = await User.select({ + user: () => + User.select({ + user: () => + User.select('runtimeComputed', 'batchComputed').takeOptional(), + }).takeOptional(), + }).takeOptional(); + + assertType< + typeof res, + | { + user: + | { + user: + | { runtimeComputed: string; batchComputed: string } + | undefined; + } + | undefined; + } + | undefined + >(); + + expect(res).toEqual({ + user: { + user: { + runtimeComputed: `${userId} ${userData.name}`, + batchComputed: `${userId} ${userData.name}`, + }, + }, + }); + }); + + it('should return undefined when one optional is not found', async () => { + const res = await User.select({ + user: () => + User.select({ + user: () => + User.select('runtimeComputed', 'batchComputed').findOptional(0), + }).takeOptional(), + }).takeOptional(); + + expect(res).toEqual({ + user: { user: null }, + }); + }); + + it('should select one required', async () => { + const res = await User.select({ + user: () => + User.select({ + user: () => + User.select('runtimeComputed', 'batchComputed').take(), + }).take(), + }).take(); + + assertType< + typeof res, + { + user: { + user: { runtimeComputed: string; batchComputed: string }; + }; + } + >(); + + expect(res).toEqual({ + user: { + user: { + runtimeComputed: `${userId} ${userData.name}`, + batchComputed: `${userId} ${userData.name}`, + }, + }, + }); + }); + + it('should throw if one is not found', async () => { + const q = User.select({ + user: () => + User.select({ + user: () => + User.select('runtimeComputed', 'batchComputed').find(0), + }).take(), + }).take(); + + await expect(q).rejects.toThrow(NotFoundError); + }); + + it('should select a pluck', async () => { + const id = await User.get('id').insert(userData); + + const res = await User.select({ + users: () => + User.select({ + runtimeComputed: () => User.pluck('runtimeComputed'), + batchComputed: () => User.pluck('batchComputed'), + }), + }); + + const expected = { + runtimeComputed: [ + `${userId} ${userData.name}`, + `${id} ${userData.name}`, + ], + batchComputed: [ + `${userId} ${userData.name}`, + `${id} ${userData.name}`, + ], + }; + + expect(res).toEqual([ + { + users: [expected, expected], + }, + { + users: [expected, expected], + }, + ]); + }); + + it('should select an optional value', async () => { + const res = await User.select({ + user: () => + User.select({ + runtimeComputed: () => User.getOptional('runtimeComputed'), + batchComputed: () => User.getOptional('batchComputed'), + }).take(), + }).take(); + + assertType< + typeof res, + { + user: { + runtimeComputed: string | undefined; + batchComputed: string | undefined; + }; + } + >(); + + expect(res).toEqual({ + user: { + runtimeComputed: `${userId} ${userData.name}`, + batchComputed: `${userId} ${userData.name}`, + }, + }); + }); + + it('should select undefined for optional value when is not found', async () => { + const res = await User.select({ + user: () => + User.select({ + runtimeComputed: () => + User.find(0).getOptional('runtimeComputed'), + batchComputed: () => User.find(0).getOptional('batchComputed'), + }).take(), + }).take(); + + expect(res).toEqual({ + user: { + runtimeComputed: undefined, + batchComputed: undefined, + }, + }); + }); + + it('should select a required value', async () => { + const res = await User.select({ + user: () => + User.select({ + runtimeComputed: () => User.get('runtimeComputed'), + batchComputed: () => User.get('batchComputed'), + }).take(), + }).take(); + + assertType< + typeof res, + { user: { runtimeComputed: string; batchComputed: string } } + >(); + + expect(res).toEqual({ + user: { + runtimeComputed: `${userId} ${userData.name}`, + batchComputed: `${userId} ${userData.name}`, + }, + }); + }); + + it('should throw when a required value is not found', async () => { + const q = User.select({ + user: () => + User.select({ + runtimeComputed: () => User.find(0).get('runtimeComputed'), + batchComputed: () => User.find(0).get('batchComputed'), + }).take(), + }).take(); + + await expect(q).rejects.toThrow(NotFoundError); + }); + }); + describe('where', () => { it('should not support computed columns', () => { // @ts-expect-error computed column should not be allowed diff --git a/packages/qb/pqb/src/modules/computed.ts b/packages/qb/pqb/src/modules/computed.ts index 0c4af2c9b..43b6b0746 100644 --- a/packages/qb/pqb/src/modules/computed.ts +++ b/packages/qb/pqb/src/modules/computed.ts @@ -6,12 +6,23 @@ import { QueryColumn, QueryColumns, QueryMetaBase, + QueryReturnType, + RecordString, RecordUnknown, } from 'orchid-core'; import { Query, QueryOrExpression } from '../query/query'; -import { ExpressionMethods, SqlMethod } from '../queryMethods'; +import { + ExpressionMethods, + QueryBatchResult, + SqlMethod, +} from '../queryMethods'; import { RelationsBase } from '../relations'; -import { ColumnType } from '../columns'; +import { ColumnType, UnknownColumn } from '../columns'; +import { QueryData } from '../sql'; +import { + applyBatchTransforms, + finalizeNestedHookSelect, +} from '../common/queryResultProcessing'; declare module 'orchid-core' { interface ColumnDataBase { @@ -98,8 +109,8 @@ export const applyComputedColumns = ( (q.q.computeds ??= {})[key] = item; } else { ( - ((q.shape as QueryColumns)[key] = item.result - .value as never) as ColumnType + ((q.shape as QueryColumns)[key] = + item.result.value || UnknownColumn.instance) as ColumnType ).data.computed = item as Expression; } } @@ -108,3 +119,125 @@ export const applyComputedColumns = ( q as unknown as RecordUnknown ).computeBatchAtRuntime = undefined; }; + +export const processComputedResult = (query: QueryData, result: unknown) => { + let promises: Promise[] | undefined; + + for (const key in query.selectedComputeds) { + const computed = query.selectedComputeds[key]; + if (computed.kind === 'one') { + for (const record of result as RecordUnknown[]) { + record[key] = computed.fn(record); + } + } else { + const res = computed.fn(result); + if (Array.isArray(res)) { + saveBatchComputed(key, result, res); + } else { + (promises ??= []).push( + (res as Promise).then((res) => + saveBatchComputed(key, result, res), + ), + ); + } + } + } + + if (!promises) return; + return Promise.all(promises); +}; + +export const processComputedBatches = ( + query: QueryData, + batches: QueryBatchResult[], + originalReturnType: QueryReturnType, + returnType: QueryReturnType, + tempColumns: Set | undefined, + renames: RecordString | undefined, + key: string, +) => { + let promises: Promise[] | undefined; + + for (const key in query.selectedComputeds) { + const computed = query.selectedComputeds[key]; + if (computed.kind === 'one') { + for (const { data } of batches) { + for (const record of data) { + if (record) { + record[key] = computed.fn(record); + } + } + } + } else { + for (const { data } of batches) { + let present; + let blanks: Set | undefined; + if (!returnType || returnType === 'all') { + present = data; + } else { + present = []; + blanks = new Set(); + for (let i = 0; i < data.length; i++) { + if (data[i]) { + present.push(data[i]); + } else { + blanks.add(i); + } + } + } + + const res = computed.fn(present); + if (Array.isArray(res)) { + saveBatchComputed(key, data, res, blanks); + } else { + (promises ??= []).push( + (res as Promise).then((res) => + saveBatchComputed(key, data, res, blanks), + ), + ); + } + } + } + } + + if (!promises) return; + + return Promise.all(promises).then(() => { + finalizeNestedHookSelect( + batches, + originalReturnType, + tempColumns, + renames, + key, + ); + + applyBatchTransforms(query, batches); + }); +}; + +const saveBatchComputed = ( + key: string, + result: unknown, + res: unknown[], + blanks?: Set, +) => { + const len = (result as unknown[]).length; + const actual = res.length + (blanks?.size || 0); + if (len !== actual) { + throw new Error( + `Incorrect length of batch computed result for column ${key}. Expected ${len}, received ${actual}.`, + ); + } + + if (blanks) { + for (let i = 0, r = 0; i < len; i++) { + if (!blanks.has(i)) { + (result as RecordUnknown[])[i][key] = res[r++]; + } + } + } else { + for (let i = 0; i < len; i++) { + (result as RecordUnknown[])[i][key] = res[i]; + } + } +}; diff --git a/packages/qb/pqb/src/query/db.ts b/packages/qb/pqb/src/query/db.ts index 4b6900278..de35f714e 100644 --- a/packages/qb/pqb/src/query/db.ts +++ b/packages/qb/pqb/src/query/db.ts @@ -199,7 +199,7 @@ export interface Db< QueryBase { result: Pick[number]>; // Pick is optimal queryBuilder: Db; - returnType: Query['returnType']; + returnType: undefined; then: QueryThen>; windows: Query['windows']; defaultSelectColumns: DefaultSelectColumns; diff --git a/packages/qb/pqb/src/query/query.ts b/packages/qb/pqb/src/query/query.ts index e4036e72e..09f2f8a13 100644 --- a/packages/qb/pqb/src/query/query.ts +++ b/packages/qb/pqb/src/query/query.ts @@ -23,8 +23,8 @@ import { QueryColumn, QueryColumns, QueryInternalBase, - QueryReturnType, QueryThen, + RecordKeyTrue, RecordUnknown, } from 'orchid-core'; import { QueryBase } from './queryBase'; @@ -245,23 +245,20 @@ export interface QueryWithTable extends Query { table: string; } -export const queryTypeWithLimitOne = { +export const queryTypeWithLimitOne: RecordKeyTrue = { one: true, oneOrThrow: true, value: true, valueOrThrow: true, -} as { [K in QueryReturnType]: true | undefined }; +}; export const isQueryReturnsAll = (q: Query) => !q.q.returnType || q.q.returnType === 'all'; -export type QueryReturnsAll = - QueryReturnType extends T ? true : T extends 'all' ? true : false; - export type GetQueryResult< T extends PickQueryReturnType, Result extends QueryColumns, -> = QueryReturnsAll extends true +> = T['returnType'] extends undefined | 'all' ? ColumnShapeOutput[] : T['returnType'] extends 'one' ? ColumnShapeOutput | undefined diff --git a/packages/qb/pqb/src/queryMethods/create.ts b/packages/qb/pqb/src/queryMethods/create.ts index bb932ce8b..5a58d2769 100644 --- a/packages/qb/pqb/src/queryMethods/create.ts +++ b/packages/qb/pqb/src/queryMethods/create.ts @@ -2,7 +2,6 @@ import { PickQueryMetaResultRelationsWithDataReturnTypeShape, Query, QueryOrExpression, - QueryReturnsAll, queryTypeWithLimitOne, SetQueryKind, SetQueryKindResult, @@ -152,7 +151,7 @@ export type CreateRelationsDataOmittingFKeys< // - if it is a `pluck` query, forces it to return a single value type CreateResult = T extends { isCount: true } ? SetQueryKind - : QueryReturnsAll extends true + : T['returnType'] extends undefined | 'all' ? SetQueryReturnsOneKindResult> : T['returnType'] extends 'pluck' ? SetQueryReturnsColumnKindResult> @@ -160,7 +159,7 @@ type CreateResult = T extends { isCount: true } type CreateRawOrFromResult = T extends { isCount: true } ? SetQueryKind - : QueryReturnsAll extends true + : T['returnType'] extends undefined | 'all' ? SetQueryReturnsOneKind : T['returnType'] extends 'pluck' ? SetQueryReturnsColumnKind @@ -175,7 +174,7 @@ type InsertResult< T extends CreateSelf, BT, > = T['meta']['hasSelect'] extends true - ? QueryReturnsAll extends true + ? T['returnType'] extends undefined | 'all' ? SetQueryReturnsOneKindResult> : T['returnType'] extends 'pluck' ? SetQueryReturnsColumnKindResult> @@ -184,7 +183,7 @@ type InsertResult< type InsertRawOrFromResult = T['meta']['hasSelect'] extends true - ? QueryReturnsAll extends true + ? T['returnType'] extends undefined | 'all' ? SetQueryReturnsOneKind : T['returnType'] extends 'pluck' ? SetQueryReturnsColumnKind @@ -524,7 +523,7 @@ const insert = ( // so that author.books.create(data) will actually perform the `from` kind of create if (!q.kind) q.kind = kind; - const { select, returnType = 'all' } = q; + const { select, returnType } = q; if (!select) { if (returnType !== 'void') q.returnType = 'rowCount'; @@ -534,7 +533,7 @@ const insert = ( } else if (returnType === 'value' || returnType === 'valueOrThrow') { q.returnType = 'pluck'; } - } else if (returnType === 'all') { + } else if (!returnType || returnType === 'all') { q.returnType = 'from' in values ? values.from.q.returnType : 'one'; } else if (returnType === 'pluck') { q.returnType = 'valueOrThrow'; @@ -555,7 +554,7 @@ const getFromSelectColumns = ( obj?: { columns: string[] }, many?: boolean, ) => { - if (!many && !queryTypeWithLimitOne[(from as Query).q.returnType]) { + if (!many && !queryTypeWithLimitOne[(from as Query).q.returnType as string]) { throw new Error( 'Cannot create based on a query which returns multiple records', ); diff --git a/packages/qb/pqb/src/queryMethods/from.ts b/packages/qb/pqb/src/queryMethods/from.ts index 938683d38..89d76176a 100644 --- a/packages/qb/pqb/src/queryMethods/from.ts +++ b/packages/qb/pqb/src/queryMethods/from.ts @@ -132,6 +132,7 @@ export function queryFrom< const { shape } = data; const parsers = (data.parsers ??= {}); const computeds = (data.computeds ??= {}); + // TODO: batchParsers for (const item of arg) { if (typeof item === 'string') { const w = (data.withShapes as WithConfigs)[item]; @@ -154,6 +155,7 @@ export function queryFrom< data.as ||= q.q.as || q.table || 't'; data.shape = getShapeFromSelect(arg as QueryBase, true) as ColumnsShapeBase; data.parsers = q.q.parsers; + data.batchParsers = q.q.batchParsers; } data.from = arg as Query; diff --git a/packages/qb/pqb/src/queryMethods/get.utils.ts b/packages/qb/pqb/src/queryMethods/get.utils.ts index 6b81effc4..24c2019ae 100644 --- a/packages/qb/pqb/src/queryMethods/get.utils.ts +++ b/packages/qb/pqb/src/queryMethods/get.utils.ts @@ -80,23 +80,29 @@ const _get = < (q as SelectQueryData)[getValueKey] = type; - setParserForSelectedString( + const selected = setParserForSelectedString( query as unknown as Query, arg, getQueryAs(query as unknown as Query), getValueKey, ); - q.expr = new SelectItemExpression(query as unknown as Query, arg, type); + q.select = selected + ? [ + (q.expr = new SelectItemExpression( + query as unknown as Query, + selected, + type, + )), + ] + : undefined; } else { type = arg.result.value; (q as SelectQueryData)[getValueKey] = type; addParserForRawExpression(query as unknown as Query, getValueKey, arg); - q.expr = arg; + q.select = [(q.expr = arg)]; } - q.select = [q.expr]; - return setQueryOperators( query as unknown as Query, type?.operators || Operators.any, diff --git a/packages/qb/pqb/src/queryMethods/join/_join.ts b/packages/qb/pqb/src/queryMethods/join/_join.ts index 824f42868..7124911ae 100644 --- a/packages/qb/pqb/src/queryMethods/join/_join.ts +++ b/packages/qb/pqb/src/queryMethods/join/_join.ts @@ -4,6 +4,7 @@ import { Query, } from '../../query/query'; import { + BatchParsers, ColumnsParsers, ColumnsShapeBase, ColumnTypeBase, @@ -56,6 +57,7 @@ export const _join = < let joinKey: string | undefined; let shape: QueryColumns | undefined; let parsers: ColumnsParsers | undefined; + let batchParsers: BatchParsers | undefined; let computeds: ComputedColumns | undefined; let joinSubQuery = false; @@ -82,6 +84,7 @@ export const _join = < if (joinKey) { shape = getShapeFromSelect(q, joinSubQuery); parsers = q.q.parsers; + batchParsers = q.q.batchParsers; computeds = q.q.computeds; if (joinSubQuery) { @@ -97,11 +100,13 @@ export const _join = < shape = getShapeFromSelect(relation.relationConfig.query); const r = relation.relationConfig.query; parsers = r.q.parsers; + batchParsers = r.q.batchParsers; computeds = r.q.computeds; } else { const w = (query as unknown as PickQueryQ).q.withShapes?.[joinKey]; shape = w?.shape; computeds = w?.computeds; + // TODO batchParsers if (shape) { // clone the shape to mutate it below, in other cases the shape is newly created @@ -149,6 +154,12 @@ export const _join = < j.q.parsers, ); + if (j.q.batchParsers) { + ((query as unknown as PickQueryQ).q.joinedBatchParsers ??= {})[ + joinKey + ] = j.q.batchParsers; + } + setQueryObjectValue( query as unknown as PickQueryQ, 'joinedComputeds', @@ -156,12 +167,26 @@ export const _join = < j.q.computeds, ); } else { - addAllShapesAndParsers(query, joinKey, shape, parsers, computeds); + addAllShapesAndParsers( + query, + joinKey, + shape, + parsers, + batchParsers, + computeds, + ); } } else if (require && 'r' in joinArgs && isQueryNone(joinArgs.r)) { return _queryNone(query) as JoinResult; } else { - addAllShapesAndParsers(query, joinKey, shape, parsers, computeds); + addAllShapesAndParsers( + query, + joinKey, + shape, + parsers, + batchParsers, + computeds, + ); } return pushQueryValue(query as unknown as PickQueryQ, 'join', { @@ -175,6 +200,7 @@ const addAllShapesAndParsers = ( joinKey?: string, shape?: QueryColumns, parsers?: ColumnsParsers, + batchParsers?: BatchParsers, computeds?: ComputedColumns, ) => { if (!joinKey) return; @@ -183,6 +209,10 @@ const addAllShapesAndParsers = ( setQueryObjectValue(query as PickQueryQ, 'joinedParsers', joinKey, parsers); + if (batchParsers) { + ((query as PickQueryQ).q.joinedBatchParsers ??= {})[joinKey] = batchParsers; + } + setQueryObjectValue( query as PickQueryQ, 'joinedComputeds', @@ -263,6 +293,9 @@ export const _joinLateral = < const shape = getShapeFromSelect(result, true); setQueryObjectValue(q, 'joinedShapes', joinKey, shape); setQueryObjectValue(q, 'joinedParsers', joinKey, result.q.parsers); + if (result.q.batchParsers) { + (q.q.joinedBatchParsers ??= {})[joinKey] = result.q.batchParsers; + } } as ||= getQueryAs(result); diff --git a/packages/qb/pqb/src/queryMethods/json.utils.ts b/packages/qb/pqb/src/queryMethods/json.utils.ts index d36accdd8..096188818 100644 --- a/packages/qb/pqb/src/queryMethods/json.utils.ts +++ b/packages/qb/pqb/src/queryMethods/json.utils.ts @@ -21,7 +21,7 @@ export function queryJson( _queryGetOptional( q, new RawSQL( - queryTypeWithLimitOne[(self as Query).q.returnType] + queryTypeWithLimitOne[(self as Query).q.returnType as string] ? `row_to_json("t".*)` : coalesce !== false ? `COALESCE(json_agg(row_to_json("t".*)), '[]')` diff --git a/packages/qb/pqb/src/queryMethods/map.ts b/packages/qb/pqb/src/queryMethods/map.ts index 80f29b6d5..4d9ba394b 100644 --- a/packages/qb/pqb/src/queryMethods/map.ts +++ b/packages/qb/pqb/src/queryMethods/map.ts @@ -1,4 +1,4 @@ -import { Query, QueryReturnsAll } from '../query/query'; +import { Query } from '../query/query'; import { QueryColumn, QueryThen, RecordUnknown } from 'orchid-core'; import { pushQueryValue } from '../query/queryUtils'; @@ -46,7 +46,7 @@ export class QueryMap { map( this: T, fn: ( - input: QueryReturnsAll extends true + input: T['returnType'] extends undefined | 'all' ? T['then'] extends QueryThen<(infer Data)[]> ? Data : never @@ -58,9 +58,7 @@ export class QueryMap { [K in keyof T]: K extends 'result' ? { [K in keyof Result]: QueryColumn } : K extends 'then' - ? QueryThen< - QueryReturnsAll extends true ? Result[] : Result - > + ? QueryThen : T[K]; } { return pushQueryValue(this.clone(), 'transform', { map: fn }) as never; diff --git a/packages/qb/pqb/src/queryMethods/merge.test.ts b/packages/qb/pqb/src/queryMethods/merge.test.ts index f17c29f66..3cd93883d 100644 --- a/packages/qb/pqb/src/queryMethods/merge.test.ts +++ b/packages/qb/pqb/src/queryMethods/merge.test.ts @@ -14,13 +14,7 @@ import { testZodColumnTypes as t, testDb, } from 'test-utils'; -import { - emptyObject, - Expression, - getValueKey, - noop, - QueryReturnType, -} from 'orchid-core'; +import { emptyObject, Expression, getValueKey, noop } from 'orchid-core'; import { ComputedColumn } from '../modules/computed'; @@ -47,7 +41,7 @@ describe('merge queries', () => { it('should have default return type if none of the queries have it', () => { const q = User.merge(User); - assertType(); + assertType(); }); it('should use left return type unless right has it', () => { diff --git a/packages/qb/pqb/src/queryMethods/merge.ts b/packages/qb/pqb/src/queryMethods/merge.ts index d3d848d90..45078ff8a 100644 --- a/packages/qb/pqb/src/queryMethods/merge.ts +++ b/packages/qb/pqb/src/queryMethods/merge.ts @@ -29,7 +29,7 @@ export type MergeQuery< : K extends 'result' ? MergeQueryResult : K extends 'returnType' - ? QueryReturnType extends Q['returnType'] + ? Q['returnType'] extends undefined ? T['returnType'] : Q['returnType'] : K extends 'then' @@ -68,6 +68,7 @@ const mergableObjects: RecordBoolean = { defaults: true, joinedShapes: true, joinedParsers: true, + joinedBatchParsers: true, selectedComputeds: true, }; diff --git a/packages/qb/pqb/src/queryMethods/queryMethods.test.ts b/packages/qb/pqb/src/queryMethods/queryMethods.test.ts index f55162740..d2b6dca0d 100644 --- a/packages/qb/pqb/src/queryMethods/queryMethods.test.ts +++ b/packages/qb/pqb/src/queryMethods/queryMethods.test.ts @@ -118,6 +118,13 @@ describe('queryMethods', () => { expect(received).toEqual(expected); }); + + it('should be disabled in a sub-query', () => { + // @ts-expect-error rows is disabled in a sub-query + User.select({ + x: () => User.rows(), + }); + }); }); describe('pluck', () => { diff --git a/packages/qb/pqb/src/queryMethods/queryMethods.ts b/packages/qb/pqb/src/queryMethods/queryMethods.ts index 25eb0baa8..7c1a2a16f 100644 --- a/packages/qb/pqb/src/queryMethods/queryMethods.ts +++ b/packages/qb/pqb/src/queryMethods/queryMethods.ts @@ -21,7 +21,6 @@ import { } from '../common/utils'; import { OrderTsQueryConfig, - SelectItem, SelectQueryData, SortDir, toSQL, @@ -354,8 +353,16 @@ export class QueryMethods { ): SetQueryReturnsPluck { const q = (this as unknown as Query).clone(); q.q.returnType = 'pluck'; - (q.q as SelectQueryData).select = [select as SelectItem]; - addParserForSelectItem(q as never, q.q.as || q.table, 'pluck', select); + + const selected = addParserForSelectItem( + q as never, + q.q.as || q.table, + 'pluck', + select, + ); + (q.q as SelectQueryData).select = selected + ? [selected as never] + : undefined; return q as never; } diff --git a/packages/qb/pqb/src/queryMethods/select.ts b/packages/qb/pqb/src/queryMethods/select.ts index 68d1d507e..5755c3ad1 100644 --- a/packages/qb/pqb/src/queryMethods/select.ts +++ b/packages/qb/pqb/src/queryMethods/select.ts @@ -4,7 +4,6 @@ import { PickQueryQAndInternal, Query, QueryMetaHasSelect, - QueryReturnsAll, WithDataBase, } from '../query/query'; import { @@ -14,25 +13,22 @@ import { ColumnsShapeToPluck, } from '../columns'; import { JSONTextColumn } from '../columns/json'; -import { pushQueryArray, pushQueryValue } from '../query/queryUtils'; +import { pushQueryArray } from '../query/queryUtils'; import { SelectAsValue, SelectItem, SelectQueryData, ToSQLQuery } from '../sql'; -import { QueryResult } from '../adapter'; import { - applyTransforms, ColumnsParsers, ColumnTypeBase, - emptyArray, Expression, getValueKey, HookSelect, isExpression, - MaybeArray, PickQueryMeta, QueryColumn, QueryColumns, QueryMetaBase, QueryReturnType, QueryThen, + RecordString, RecordUnknown, setColumnData, setParserToQuery, @@ -46,11 +42,15 @@ import { import { RawSQL } from '../sql/rawSql'; import { defaultSchemaConfig } from '../columns/defaultSchemaConfig'; import { RelationsBase } from '../relations'; -import { filterResult, parseRecord } from './then'; +import { parseRecord } from './then'; import { _queryNone, isQueryNone } from './none'; import { NotFoundError } from '../errors'; -import { ComputedColumns } from '../modules/computed'; +import { ComputedColumns, processComputedBatches } from '../modules/computed'; +import { + applyBatchTransforms, + finalizeNestedHookSelect, +} from '../common/queryResultProcessing'; interface SelectSelf { shape: QueryColumns; @@ -82,7 +82,9 @@ interface SelectAsArg { [K in keyof T]: K extends keyof T['relations'] ? T['relations'][K]['relationConfig']['methodQuery'] : T[K]; - }) => QueryBase | Expression); + }) => + | (QueryBase & { returnType: Exclude }) + | Expression); } // Result type of select without the ending object argument. @@ -261,17 +263,16 @@ type SelectAsValueResult< // query that returns a single value becomes a column of that value // query that returns 'pluck' becomes a column with array type of specific value type // query that returns a single record becomes an object column, possibly nullable -export type SelectSubQueryResult = QueryReturnsAll< - Arg['returnType'] -> extends true - ? ColumnsShapeToObjectArray - : Arg['returnType'] extends 'value' | 'valueOrThrow' - ? Arg['result']['value'] - : Arg['returnType'] extends 'pluck' - ? ColumnsShapeToPluck - : Arg['returnType'] extends 'one' - ? ColumnsShapeToNullableObject - : ColumnsShapeToObject; +export type SelectSubQueryResult = + Arg['returnType'] extends undefined | 'all' + ? ColumnsShapeToObjectArray + : Arg['returnType'] extends 'value' | 'valueOrThrow' + ? Arg['result']['value'] + : Arg['returnType'] extends 'pluck' + ? ColumnsShapeToPluck + : Arg['returnType'] extends 'one' + ? ColumnsShapeToNullableObject + : ColumnsShapeToObject; // add a parser for a raw expression column // is used by .select and .get methods @@ -284,14 +285,6 @@ export const addParserForRawExpression = ( if (type?.parseFn) setParserToQuery(q.q, key, type.parseFn); }; -// these are used as a wrapper to pass sub query result to `parseRecord` -const subQueryResult: QueryResult = { - // sub query can't return a rowCount, use -1 as for impossible case - rowCount: -1, - rows: emptyArray, - fields: emptyArray, -}; - // add parsers when selecting a full joined table by name or alias const addParsersForSelectJoined = ( q: PickQueryQ, @@ -302,8 +295,26 @@ const addParsersForSelectJoined = ( if (parsers) { setParserToQuery(q.q, as, (row) => parseRecord(parsers, row)); } + + const batchParsers = q.q.joinedBatchParsers?.[arg]; + if (batchParsers) { + (q.q.batchParsers ??= []).push( + ...batchParsers.map((x) => ({ + path: [as as string, ...x.path], + fn: x.fn, + })), + ); + } }; +export interface QueryBatchResult { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + data: any; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + parent: any; + key: PropertyKey; +} + // add parser for a single key-value pair of selected object export const addParserForSelectItem = ( q: T, @@ -316,65 +327,183 @@ export const addParserForSelectItem = ( addParserForRawExpression(q as never, key, arg); } else { const { q: query } = arg; - if (query.hookSelect || query.parsers || query.transform) { - setParserToQuery((q as unknown as Query).q, key, (item) => { - const { hookSelect } = query; - - const returnType = query.returnType || 'all'; - const tempReturnType = hookSelect ? 'all' : returnType; - - subQueryResult.rows = - returnType === 'value' || returnType === 'valueOrThrow' - ? [[item]] - : returnType === 'one' || returnType === 'oneOrThrow' - ? [item] - : (item as unknown[]); - - let result = query.handleResult( - arg, - tempReturnType, - subQueryResult, - true, - ); - - if (hookSelect) { - result = filterResult( - arg, - returnType, - subQueryResult, - result, - new Set(), - ); - } - return query.transform - ? applyTransforms(returnType, query.transform, result) - : result; - }); + if (query.batchParsers) { + const batchParsers = ((q as unknown as Query).q.batchParsers ??= []); + for (const bp of query.batchParsers) { + batchParsers.push({ path: [key, ...bp.path], fn: bp.fn }); + } } - if ( - query.returnType === 'valueOrThrow' || - query.returnType === 'oneOrThrow' - ) { - pushQueryValue( - q as unknown as PickQueryQ, - 'transform', - (data: MaybeArray) => { - if (Array.isArray(data)) { - for (const item of data) { - if (item[key as string] === undefined) { - throw new NotFoundError(q as unknown as Query); + if (query.hookSelect || query.parsers || query.transform) { + const batchParsers = ((q as unknown as Query).q.batchParsers ??= []); + + batchParsers.push({ + path: [key], + fn: (path, queryResult) => { + const { rows } = queryResult; + const originalReturnType = query.returnType || 'all'; + let returnType = originalReturnType; + const { hookSelect } = query; + const batches: QueryBatchResult[] = []; + + let last = path.length; + if (returnType === 'value' || returnType === 'valueOrThrow') { + if (hookSelect) { + batches.push = (item) => { + // if the item has no key, it means value return was implicitly turned into 'one' return, + // happens when getting a computed column + if (!(key in item)) { + returnType = returnType === 'value' ? 'one' : 'oneOrThrow'; + } + batches.push = Array.prototype.push; + return batches.push(item); + }; + } else { + last--; + } + } + + collectNestedSelectBatches(batches, rows, path, last); + + switch (returnType) { + case 'all': { + const { parsers } = query; + if (parsers) { + for (const { data } of batches) { + for (const one of data) { + parseRecord(parsers, one); + } + } } + break; } - } else { - if (data[key as string] === undefined) { - throw new NotFoundError(q as unknown as Query); + case 'one': + case 'oneOrThrow': { + const { parsers } = query; + if (parsers) { + if (returnType === 'one') { + for (const { data } of batches) { + if (data) parseRecord(parsers, data); + } + } else { + for (const { data } of batches) { + if (!data) throw new NotFoundError(arg); + parseRecord(parsers, data); + } + } + } else if (returnType !== 'one') { + for (const { data } of batches) { + if (!data) throw new NotFoundError(arg); + } + } + + if (hookSelect) { + for (const batch of batches) { + batch.data = [batch.data]; + } + } + + break; + } + case 'pluck': { + const parse = query.parsers?.pluck; + if (parse) { + for (const { data } of batches) { + for (let i = 0; i < data.length; i++) { + (data as unknown as RecordUnknown)[i] = parse(data[i]); + } + } + } + + // not transforming data for hookSelect because it's set to load 'all' elsewhere for this case + + break; + } + case 'value': + case 'valueOrThrow': { + const parse = query.parsers?.[getValueKey]; + if (parse) { + if (returnType === 'value') { + for (const { data } of batches) { + data[key] = + data[key] === undefined + ? arg.q.notFoundDefault + : parse(data[key]); + } + } else { + for (const { data } of batches) { + if (data[key] === undefined) throw new NotFoundError(arg); + data[key] = parse(data[key]); + } + } + } else if (returnType !== 'value') { + for (const { data } of batches) { + if (data[key] === undefined) throw new NotFoundError(arg); + } + } + + if (hookSelect) { + for (const batch of batches) { + batch.data = [batch.data]; + } + } + + break; } } - return data; + + if (hookSelect) { + let tempColumns: Set | undefined; + let renames: RecordString | undefined; + for (const column of hookSelect.keys()) { + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + const as = hookSelect!.get(column)!.as; + if (as) (renames ??= {})[column] = as; + + (tempColumns ??= new Set())?.add(as || column); + } + + if (renames) { + for (const { data } of batches) { + for (const record of data) { + if (record) { + for (const a in renames) { + const value = record[renames[a]]; + record[renames[a]] = record[a]; + record[a] = value; + } + } + } + } + } + + if (query.selectedComputeds) { + const maybePromise = processComputedBatches( + query, + batches, + originalReturnType, + returnType, + tempColumns, + renames, + key, + ); + if (maybePromise) return maybePromise; + } + + finalizeNestedHookSelect( + batches, + originalReturnType, + tempColumns, + renames, + key, + ); + } + + applyBatchTransforms(query, batches); + return; }, - ); + }); } } return arg; @@ -383,6 +512,48 @@ export const addParserForSelectItem = ( return setParserForSelectedString(q as never, arg as string, as, key); }; +const collectNestedSelectBatches = ( + batches: QueryBatchResult[], + rows: unknown[], + path: string[], + last: number, +) => { + const stack: { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + data: any; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + parent: any; + key: PropertyKey; + i: number; + }[] = rows.map( + (row) => + ({ + data: row, + i: 0, + } as never), + ); + + while (stack.length > 0) { + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + const item = stack.pop()!; + const { i } = item; + if (i === last) { + batches.push(item); + continue; + } + + const { data } = item; + const key = path[i]; + if (Array.isArray(data)) { + for (let key = 0; key < data.length; key++) { + stack.push({ data: data[key], parent: data, key, i }); + } + } else if (data && typeof data === 'object') { + stack.push({ data: data[key], parent: data, key, i: i + 1 }); + } + } +}; + // reuse SQL for empty array for JSON agg expressions const emptyArrSQL = new RawSQL("'[]'"); @@ -419,25 +590,29 @@ export const processSelectArg = ( query = value.json(false); value.q.coalesceValue = emptyArrSQL; } else if (returnType === 'pluck') { - query = value - .wrap(value.baseQuery.clone()) - .jsonAgg(value.q.select[0]); + // no select in case of plucking a computed + query = value.q.select + ? value.wrap(value.baseQuery.clone()).jsonAgg(value.q.select[0]) + : value.json(false); value.q.coalesceValue = emptyArrSQL; } else { - if ( - (returnType === 'value' || returnType === 'valueOrThrow') && - value.q.select - ) { - // todo: investigate what is this for - if (typeof value.q.select[0] === 'string') { - value.q.select[0] = { - selectAs: { r: value.q.select[0] }, - }; + if (returnType === 'value' || returnType === 'valueOrThrow') { + if (value.q.select) { + // todo: investigate what is this for + if (typeof value.q.select[0] === 'string') { + value.q.select[0] = { + selectAs: { r: value.q.select[0] }, + }; + } + + query = value; + } else { + query = value.json(false); } + } else { + query = value; } - - query = value; } let asOverride = key; @@ -509,6 +684,15 @@ export const setParserForSelectedString = ( const parser = q.q.joinedParsers?.[table]?.[column]; if (parser) setParserToQuery(q.q, columnAs || column, parser); + const batchParsers = q.q.joinedBatchParsers?.[table]; + if (batchParsers) { + for (const bp of batchParsers) { + if (bp.path[0] === column) { + (q.q.batchParsers ??= []).push(bp); + } + } + } + const computeds = q.q.joinedComputeds?.[table]; if (computeds?.[column]) { const computed = computeds[column]; diff --git a/packages/qb/pqb/src/queryMethods/sql.ts b/packages/qb/pqb/src/queryMethods/sql.ts index bf981e7b0..139b90306 100644 --- a/packages/qb/pqb/src/queryMethods/sql.ts +++ b/packages/qb/pqb/src/queryMethods/sql.ts @@ -148,7 +148,7 @@ export class SqlMethod { ): RawSQL, ColumnTypes>; sql( this: PickQueryColumnTypes, - ...args: [DynamicSQLArg] + ...args: [DynamicSQLArg>] ): DynamicRawSQL, ColumnTypes>; sql(this: PickQueryColumnTypes, ...args: unknown[]) { // eslint-disable-next-line @typescript-eslint/no-explicit-any diff --git a/packages/qb/pqb/src/queryMethods/then.ts b/packages/qb/pqb/src/queryMethods/then.ts index 97518d1a3..745a515f0 100644 --- a/packages/qb/pqb/src/queryMethods/then.ts +++ b/packages/qb/pqb/src/queryMethods/then.ts @@ -1,7 +1,7 @@ import { Query } from '../query/query'; import { NotFoundError, QueryError } from '../errors'; -import { QueryArraysResult, QueryResult } from '../adapter'; -import { CommonQueryData, QueryAfterHook, QueryBeforeHook } from '../sql'; +import { QueryResult } from '../adapter'; +import { HandleResult, QueryAfterHook, QueryBeforeHook } from '../sql'; import pg from 'pg'; import { AdapterBase, @@ -11,6 +11,7 @@ import { ColumnsParsers, emptyArray, getValueKey, + MaybePromise, QueryReturnType, RecordString, RecordUnknown, @@ -19,10 +20,12 @@ import { TransactionState, } from 'orchid-core'; import { commitSql } from './transaction'; +import { processComputedResult } from '../modules/computed'; export const queryMethodByReturnType: { - [K in QueryReturnType]: 'query' | 'arrays'; + [K in string]: 'query' | 'arrays'; } = { + undefined: 'query', all: 'query', rows: 'arrays', pluck: 'arrays', @@ -68,7 +71,7 @@ if (process.versions.bun) { if (!trx) return maybeWrappedThen; return (resolve, reject) => { - // Here `transactionStorage.getStore()` returns undefined, + // Here `transactionStorage.getStore()` tempReturnType undefined, // need to set the `trx` value to the store to workaround the bug. return this.internal.transactionStorage.run(trx, () => { return maybeWrappedThen.call(this, resolve, reject); @@ -92,15 +95,6 @@ Object.defineProperty(Then.prototype, 'then', { }, }); -export const handleResult: CommonQueryData['handleResult'] = ( - q, - returnType, - result: QueryResult, - isSubQuery?: true, -) => { - return parseResult(q, q.q.parsers, returnType, result, isSubQuery); -}; - function maybeWrappedThen( this: Query, resolve?: Resolve, @@ -205,7 +199,10 @@ const then = async ( sql = q.toSQL(); const { hookSelect } = sql; const { returnType = 'all' } = query; - const returns = hookSelect ? 'all' : returnType; + const tempReturnType = + hookSelect || (returnType === 'rows' && q.q.batchParsers) + ? 'all' + : returnType; let result: unknown; let queryResult; @@ -222,7 +219,7 @@ const then = async ( } queryResult = (await adapter[ - hookSelect ? 'query' : (queryMethodByReturnType[returnType] as 'query') + queryMethodByReturnType[tempReturnType] as 'query' ](sql)) as QueryResult; if (query.patchResult) { @@ -235,13 +232,11 @@ const then = async ( sql = undefined; } - result = query.handleResult(q, returns, queryResult); + result = query.handleResult(q, tempReturnType, queryResult); } else { // autoPreparedStatements in batch doesn't seem to make sense - const queryMethod = hookSelect - ? 'query' - : (queryMethodByReturnType[returnType] as 'query'); + const queryMethod = queryMethodByReturnType[tempReturnType] as 'query'; if (!trx) { if (query.log) logData = query.log.beforeQuery(beginSql); @@ -284,9 +279,18 @@ const then = async ( } // eslint-disable-next-line @typescript-eslint/no-non-null-assertion - result = query.handleResult(q, returns, queryResult!); + result = query.handleResult(q, tempReturnType, queryResult!); } + if ( + result && + typeof result === 'object' && + typeof (result as RecordUnknown).then === 'function' + ) { + result = await result; + } + + // TODO: move computeds after parsing let tempColumns: Set | undefined; let renames: RecordString | undefined; if (hookSelect) { @@ -307,30 +311,11 @@ const then = async ( } } } - } - if (query.selectedComputeds) { - let promises: Promise[] | undefined; - - for (const key in query.selectedComputeds) { - const computed = query.selectedComputeds[key]; - if (computed.kind === 'one') { - for (const record of result as RecordUnknown[]) { - record[key] = computed.fn(record); - } - } else { - const res = computed.fn(result); - if (Array.isArray(res)) saveBatchComputed(key, result, res); - else - (promises ??= []).push( - (res as Promise).then((res) => - saveBatchComputed(key, result, res), - ), - ); - } + if (query.selectedComputeds) { + const promise = processComputedResult(query, result); + if (promise) await promise; } - - if (promises) await Promise.all(promises); } const hasAfterHook = afterHooks || afterCommitHooks || query.after; @@ -366,7 +351,7 @@ const then = async ( } // can be set by hooks or by computed columns - if (hookSelect) { + if (hookSelect || tempReturnType !== returnType) { if (renames) { for (const record of result as RecordUnknown[]) { for (const a in renames) { @@ -411,23 +396,6 @@ const then = async ( } }; -export const saveBatchComputed = ( - key: string, - result: unknown, - res: unknown[], -) => { - const len = (result as unknown[]).length; - if (len !== res.length) { - throw new Error( - `Incorrect length of batch computed result for column ${key}. Expected ${len}, received ${res.length}.`, - ); - } - - for (let i = 0; i < len; i++) { - (result as RecordUnknown[])[i][key] = res[i]; - } -}; - const assignError = (to: QueryError, from: pg.DatabaseError) => { to.message = from.message; (to as { length?: number }).length = from.length; @@ -452,66 +420,115 @@ const assignError = (to: QueryError, from: pg.DatabaseError) => { return to; }; -export const parseResult = ( - q: Query, - parsers: ColumnsParsers | undefined, - returnType: QueryReturnType | undefined = 'all', - result: QueryResult, - isSubQuery?: boolean, -): unknown => { +export const handleResult: HandleResult = ( + q, + returnType, + result, + isSubQuery, +) => { + const { parsers } = q.q; + switch (returnType) { case 'all': { if (q.q.throwOnNotFound && result.rows.length === 0) throw new NotFoundError(q); + const promise = parseBatch(q, result); + const { rows } = result; + if (parsers) { for (const row of rows) { parseRecord(parsers, row); } } - return rows; + + return promise ? promise.then(() => rows) : rows; } case 'one': { - const row = result.rows[0]; - if (!row) return; + const { rows } = result; + if (!rows.length) return; - return parsers ? parseRecord(parsers, row) : row; + const promise = parseBatch(q, result); + + if (parsers) parseRecord(parsers, rows[0]); + + return promise ? promise.then(() => rows[0]) : rows[0]; } case 'oneOrThrow': { - const row = result.rows[0]; - if (!row) throw new NotFoundError(q); + const { rows } = result; + if (!rows.length) throw new NotFoundError(q); + + const promise = parseBatch(q, result); - return parsers ? parseRecord(parsers, row) : row; + if (parsers) parseRecord(parsers, rows[0]); + + return promise ? promise.then(() => rows[0]) : rows[0]; } case 'rows': { - return parsers - ? parseRows( - parsers, - (result as unknown as QueryArraysResult).fields, - result.rows, - ) - : result.rows; + const { rows } = result; + + const promise = parseBatch(q, result); + if (promise) { + return promise.then(() => { + if (parsers) parseRows(parsers, result.fields, rows); + + return rows; + }); + } else if (parsers) { + parseRows(parsers, result.fields, rows); + } + + return rows; } case 'pluck': { - const pluck = parsers?.pluck; - if (pluck) { - return result.rows.map(isSubQuery ? pluck : (row) => pluck(row[0])); - } else if (isSubQuery) { - return result.rows; + const { rows } = result; + + const promise = parseBatch(q, result); + + if (promise) { + return promise.then(() => { + parsePluck(parsers, isSubQuery, rows); + + return rows; + }); } - return result.rows.map((row) => row[0]); + + parsePluck(parsers, isSubQuery, rows); + + return rows; } case 'value': { - const value = result.rows[0]?.[0]; - return value !== undefined - ? parseValue(value, parsers) + const { rows } = result; + + const promise = parseBatch(q, result); + + if (promise) { + return promise.then(() => { + return rows[0]?.[0] !== undefined + ? parseValue(rows[0][0], parsers) + : q.q.notFoundDefault; + }); + } + + return rows[0]?.[0] !== undefined + ? parseValue(rows[0][0], parsers) : q.q.notFoundDefault; } case 'valueOrThrow': { - const value = result.rows[0]?.[0]; - if (value === undefined) throw new NotFoundError(q); - return parseValue(value, parsers); + const { rows } = result; + + const promise = parseBatch(q, result); + + if (promise) { + return promise.then(() => { + if (rows[0]?.[0] === undefined) throw new NotFoundError(q); + return parseValue(rows[0][0], parsers); + }); + } + + if (rows[0]?.[0] === undefined) throw new NotFoundError(q); + return parseValue(rows[0][0], parsers); } case 'rowCount': { if (q.q.throwOnNotFound && result.rowCount === 0) { @@ -525,8 +542,21 @@ export const parseResult = ( } }; +const parseBatch = (q: Query, queryResult: QueryResult): MaybePromise => { + let promises: Promise[] | undefined; + + if (q.q.batchParsers) { + for (const parser of q.q.batchParsers) { + const res = parser.fn(parser.path, queryResult); + if (res) (promises ??= []).push(res); + } + } + + return promises && (Promise.all(promises) as never); +}; + // eslint-disable-next-line @typescript-eslint/no-explicit-any -export const parseRecord = (parsers: ColumnsParsers, row: any) => { +export const parseRecord = (parsers: ColumnsParsers, row: any): unknown => { for (const key in parsers) { if (key in row) { row[key] = (parsers[key] as ColumnParser)(row[key]); @@ -540,7 +570,7 @@ const parseRows = ( fields: { name: string }[], // eslint-disable-next-line @typescript-eslint/no-explicit-any rows: any[], -) => { +): void => { for (let i = fields.length - 1; i >= 0; i--) { const parser = parsers[fields[i].name]; if (parser) { @@ -549,10 +579,26 @@ const parseRows = ( } } } - return rows; }; -const parseValue = (value: unknown, parsers?: ColumnsParsers) => { +const parsePluck = ( + parsers: ColumnsParsers | undefined, + isSubQuery: true | undefined, + rows: unknown[], +): void => { + const pluck = parsers?.pluck; + if (pluck) { + for (let i = 0; i < rows.length; i++) { + rows[i] = pluck(isSubQuery ? rows[i] : (rows[i] as RecordUnknown)[0]); + } + } else if (!isSubQuery) { + for (let i = 0; i < rows.length; i++) { + rows[i] = (rows[i] as RecordUnknown)[0]; + } + } +}; + +const parseValue = (value: unknown, parsers?: ColumnsParsers): unknown => { const parser = parsers?.[getValueKey]; return parser ? parser(value) : value; }; diff --git a/packages/qb/pqb/src/queryMethods/update.ts b/packages/qb/pqb/src/queryMethods/update.ts index d87cb7018..364c6e1bc 100644 --- a/packages/qb/pqb/src/queryMethods/update.ts +++ b/packages/qb/pqb/src/queryMethods/update.ts @@ -2,7 +2,6 @@ import { PickQueryMetaResultRelationsWithDataReturnTypeShape, Query, QueryOrExpression, - QueryReturnsAll, SetQueryKind, SetQueryReturnsRowCount, } from '../query/query'; @@ -63,7 +62,7 @@ type UpdateColumn = type UpdateRelationData< T extends UpdateSelf, Rel extends RelationConfigBase, -> = QueryReturnsAll extends true +> = T['returnType'] extends undefined | 'all' ? Rel['dataForUpdate'] : Rel['one'] extends true ? Rel['dataForUpdate'] | Rel['dataForUpdateOne'] diff --git a/packages/qb/pqb/src/sql/data.ts b/packages/qb/pqb/src/sql/data.ts index d313bd995..f381882e8 100644 --- a/packages/qb/pqb/src/sql/data.ts +++ b/packages/qb/pqb/src/sql/data.ts @@ -32,6 +32,8 @@ import { ExpressionChain, QueryDataTransform, HookSelect, + BatchParsers, + MaybePromise, } from 'orchid-core'; import { RelationQuery } from '../relations'; @@ -80,16 +82,18 @@ export type QueryDataFromItem = string | Query | Expression; export interface QueryDataJoinTo extends PickQueryTable, PickQueryQ {} +export type HandleResult = ( + q: Query, + returnType: QueryReturnType, + result: QueryResult, + isSubQuery?: true, +) => MaybePromise; + export interface CommonQueryData { adapter: Adapter; shape: ColumnsShapeBase; patchResult?(q: Query, queryResult: QueryResult): Promise; - handleResult( - q: Query, - returnType: QueryReturnType, - result: QueryResult, - isSubQuery?: true, - ): unknown; + handleResult: HandleResult; returnType: QueryReturnType; wrapInTransaction?: boolean; throwOnNotFound?: boolean; @@ -98,6 +102,7 @@ export interface CommonQueryData { joinTo?: QueryDataJoinTo; joinedShapes?: JoinedShapes; joinedParsers?: JoinedParsers; + joinedBatchParsers?: { [K: string]: BatchParsers }; joinedComputeds?: { [K: string]: ComputedColumns }; joinedForSelect?: string; innerJoinLateral?: true; @@ -113,6 +118,7 @@ export interface CommonQueryData { or?: WhereItem[][]; coalesceValue?: unknown | Expression; parsers?: ColumnsParsers; + batchParsers?: BatchParsers; notFoundDefault?: unknown; defaults?: RecordUnknown; // for runtime computed dependencies @@ -300,6 +306,7 @@ export interface PickQueryDataShapeAndJoinedShapes { joinedShapes?: JoinedShapes; } +// TODO: what if destructure when setting instead of when cloning? export const cloneQuery = (q: QueryData) => { if (q.with) q.with = q.with.slice(0); if (q.select) q.select = q.select.slice(0); @@ -310,6 +317,7 @@ export const cloneQuery = (q: QueryData) => { if (q.after) q.after = q.after.slice(0); if (q.joinedShapes) q.joinedShapes = { ...q.joinedShapes }; if (q.scopes) q.scopes = { ...q.scopes }; + if (q.parsers) q.parsers = { ...q.parsers }; // may have data for updating timestamps on any kind of query if ((q as UpdateQueryData).updateData) { diff --git a/packages/qb/pqb/src/sql/rawSql.ts b/packages/qb/pqb/src/sql/rawSql.ts index 33f3c9aab..d18ef27fc 100644 --- a/packages/qb/pqb/src/sql/rawSql.ts +++ b/packages/qb/pqb/src/sql/rawSql.ts @@ -3,7 +3,7 @@ import { DynamicSQLArg, emptyObject, Expression, - ExpressionChain, + ExpressionData, ExpressionTypeMethod, isTemplateLiteralArgs, QueryColumn, @@ -140,15 +140,16 @@ export class DynamicRawSQL< > extends Expression { declare columnTypes: ColumnTypes; result: { value: T } = emptyObject as { value: T }; - q: { chain?: ExpressionChain } = {}; + q: ExpressionData; - constructor(public fn: DynamicSQLArg) { + constructor(public fn: DynamicSQLArg) { super(); + this.q = { expr: this }; } // Calls the given function to get SQL from it. makeSQL(ctx: ToSQLCtx, quotedAs?: string): string { - return this.fn(raw).toSQL(ctx, quotedAs); + return this.fn(raw as never).toSQL(ctx, quotedAs); } } @@ -156,7 +157,7 @@ DynamicRawSQL.prototype.type = ExpressionTypeMethod.prototype.type; export function raw(...args: StaticSQLArgs): RawSQL>; export function raw( - ...args: [DynamicSQLArg] + ...args: [DynamicSQLArg>] ): DynamicRawSQL>; export function raw(...args: SQLArgs) { return isTemplateLiteralArgs(args) diff --git a/packages/qb/pqb/src/sql/select.ts b/packages/qb/pqb/src/sql/select.ts index 4313b3552..5f5338d94 100644 --- a/packages/qb/pqb/src/sql/select.ts +++ b/packages/qb/pqb/src/sql/select.ts @@ -255,6 +255,7 @@ const pushSubQuerySql = ( default: throw new UnhandledTypeError(query as Query, returnType); } + if (sql) list.push(`${coalesce(ctx, query, sql, quotedAs)} "${as}"`); return; } @@ -268,21 +269,27 @@ const pushSubQuerySql = ( case 'pluck': { const { select } = query.q; const first = select?.[0]; - if (!select || !first) { + if (!first && query.q.computeds?.[as]) { + query = queryJson(query) as unknown as typeof query; + } else if (!first) { throw new OrchidOrmInternalError( query as Query, `Nothing was selected for pluck`, ); + } else { + const cloned = query.clone(); + cloned.q.select = [{ selectAs: { c: first } }] as SelectItem[]; + query = queryWrap(cloned, cloned.baseQuery.clone()); + _queryGetOptional(query, new RawSQL(`COALESCE(json_agg("c"), '[]')`)); } - - const cloned = query.clone(); - cloned.q.select = [{ selectAs: { c: first } }] as SelectItem[]; - query = queryWrap(cloned, cloned.baseQuery.clone()); - _queryGetOptional(query, new RawSQL(`COALESCE(json_agg("c"), '[]')`)); break; } case 'value': case 'valueOrThrow': + if (query.q.computeds?.[as]) { + query = queryJson(query) as unknown as typeof query; + } + break; case 'rows': case 'rowCount': case 'void': diff --git a/packages/qb/pqb/src/sql/toSQL.ts b/packages/qb/pqb/src/sql/toSQL.ts index ccb14565f..92a03ddcd 100644 --- a/packages/qb/pqb/src/sql/toSQL.ts +++ b/packages/qb/pqb/src/sql/toSQL.ts @@ -218,7 +218,7 @@ export function pushLimitSQL( q: SelectQueryData, ) { if (!q.returnsOne) { - if (queryTypeWithLimitOne[q.returnType]) { + if (queryTypeWithLimitOne[q.returnType as string]) { sql.push(`LIMIT 1`); } else if (q.limit) { sql.push(`LIMIT ${addValue(values, q.limit)}`);