pqb 0.3.2 → 0.3.4

package/src/queryMethods/queryMethods.test.ts

@@ -1,17 +1,16 @@
 import { raw } from '../common';
-import { SelectQueryData } from '../sql';
 import {
   expectQueryNotMutated,
   adapter,
   User,
   Profile,
-  AssertEqual,
   useTestDatabase,
   db,
   expectSql,
   userData,
   now,
   assertType,
+  UserRecord,
 } from '../test-utils';
 import { NumberColumn } from '../columnSchema';
 import { NotFoundError } from '../errors';
@@ -26,8 +25,7 @@ describe('queryMethods', () => {
       expect(cloned.table).toBe(User.table);
       expect(cloned.shape).toBe(User.shape);
 
-      const eq: AssertEqual<typeof User, typeof cloned> = true;
-      expect(eq).toBe(true);
+      assertType<typeof User, typeof cloned>();
     });
   });
 
@@ -36,19 +34,11 @@ describe('queryMethods', () => {
       const sql = User.toSql();
       expectSql(sql, `SELECT * FROM "user"`);
 
-      const eq: AssertEqual<typeof sql, { text: string; values: unknown[] }> =
-        true;
-      expect(eq).toBe(true);
+      assertType<typeof sql, { text: string; values: unknown[] }>();
     });
   });
 
   describe('.all', () => {
-    it('should remove `take` from query if it is set', () => {
-      const q = User.take();
-      expect((q.query as SelectQueryData)?.take).toBe(true);
-      expect((q.all().query as SelectQueryData)?.take).toBe(undefined);
-    });
-
     it('should produce correct sql', () => {
       expectSql(User.all().toSql(), `SELECT * FROM "user"`);
     });
@@ -67,8 +57,7 @@ describe('queryMethods', () => {
         .then((res) => res.rows[0]);
 
       const user = await q.take();
-      const eq: AssertEqual<typeof user, typeof User.type> = true;
-      expect(eq).toBe(true);
+      assertType<typeof user, UserRecord>();
 
       expect(user).toEqual({
         ...expected,
@@ -95,8 +84,7 @@ describe('queryMethods', () => {
         .then((res) => res.rows[0]);
 
       const user = await q.takeOptional();
-      const eq: AssertEqual<typeof user, typeof User.type | undefined> = true;
-      expect(eq).toBe(true);
+      assertType<typeof user, UserRecord | undefined>();
 
       expect(user).toEqual({
         ...expected,
@@ -107,8 +95,7 @@ describe('queryMethods', () => {
 
     it('should return undefined if not found', async () => {
       const user = await User.takeOptional();
-      const eq: AssertEqual<typeof user, typeof User.type | undefined> = true;
-      expect(eq).toBe(true);
+      assertType<typeof user, UserRecord | undefined>();
 
       expect(user).toBe(undefined);
     });
@@ -122,12 +109,6 @@ describe('queryMethods', () => {
       const received = await User.rows();
       expect(received).toEqual(expected);
     });
-
-    it('removes `take` from query data', () => {
-      expect((User.take().rows().query as SelectQueryData)?.take).toBe(
-        undefined,
-      );
-    });
   });
 
   describe('pluck', () => {
@@ -141,16 +122,14 @@ describe('queryMethods', () => {
       const result = await User.pluck('createdAt');
       expect(result).toEqual([now, now, now]);
 
-      const eq: AssertEqual<typeof result, Date[]> = true;
-      expect(eq).toBe(true);
+      assertType<typeof result, Date[]>();
     });
 
     it('should support raw expression', async () => {
       const result = await User.pluck(raw<NumberColumn>('123'));
       expect(result).toEqual([123, 123, 123]);
 
-      const eq: AssertEqual<typeof result, number[]> = true;
-      expect(eq).toBe(true);
+      assertType<typeof result, number[]>();
     });
   });
 
@@ -159,12 +138,6 @@ describe('queryMethods', () => {
       const received = await User.exec();
       expect(received).toEqual(undefined);
     });
-
-    it('removes `take` from query data', () => {
-      expect((User.take().exec().query as SelectQueryData)?.take).toBe(
-        undefined,
-      );
-    });
   });
 
   describe('distinct', () => {
@@ -247,8 +220,7 @@ describe('queryMethods', () => {
       const q = User.all();
       const query = q.find(1);
 
-      const eq: AssertEqual<Awaited<typeof query>, typeof User.type> = true;
-      expect(eq).toBe(true);
+      assertType<Awaited<typeof query>, UserRecord>();
 
       expectSql(
         query.toSql(),
@@ -266,8 +238,7 @@ describe('queryMethods', () => {
       const q = User.all();
       const query = q.find(raw('$1 + $2', 1, 2));
 
-      const eq: AssertEqual<Awaited<typeof query>, typeof User.type> = true;
-      expect(eq).toBe(true);
+      assertType<Awaited<typeof query>, UserRecord>();
 
       expectSql(
         query.toSql(),
@@ -287,11 +258,7 @@ describe('queryMethods', () => {
       const q = User.all();
       const query = q.findOptional(1);
 
-      const eq: AssertEqual<
-        Awaited<typeof query>,
-        typeof User.type | undefined
-      > = true;
-      expect(eq).toBe(true);
+      assertType<Awaited<typeof query>, UserRecord | undefined>();
 
       expectSql(
         query.toSql(),
@@ -309,11 +276,7 @@ describe('queryMethods', () => {
       const q = User.all();
       const query = q.findOptional(raw('$1 + $2', 1, 2));
 
-      const eq: AssertEqual<
-        Awaited<typeof query>,
-        typeof User.type | undefined
-      > = true;
-      expect(eq).toBe(true);
+      assertType<Awaited<typeof query>, UserRecord | undefined>();
 
       expectSql(
         query.toSql(),
@@ -355,11 +318,7 @@ describe('queryMethods', () => {
       const q = User.all();
       const query = q.findByOptional({ name: 's' });
 
-      const eq: AssertEqual<
-        Awaited<typeof query>,
-        typeof User.type | undefined
-      > = true;
-      expect(eq).toBe(true);
+      assertType<Awaited<typeof query>, UserRecord | undefined>();
 
       expectSql(
         query.toSql(),
@@ -373,11 +332,7 @@ describe('queryMethods', () => {
       const q = User.all();
       const query = q.findByOptional({ name: raw(`'string'`) });
 
-      const eq: AssertEqual<
-        Awaited<typeof query>,
-        typeof User.type | undefined
-      > = true;
-      expect(eq).toBe(true);
+      assertType<Awaited<typeof query>, UserRecord | undefined>();
 
       expectSql(
         query.toSql(),

package/src/queryMethods/queryMethods.ts

@@ -13,12 +13,7 @@ import {
   SetQueryTableAlias,
   SetQueryWindows,
 } from '../query';
-import {
-  applyMixins,
-  getClonedQueryData,
-  GetTypesOrRaw,
-  PropertyKeyUnionToArray,
-} from '../utils';
+import { applyMixins, EmptyObject, getClonedQueryData } from '../utils';
 import {
   SelectItem,
   SelectQueryData,
@@ -28,11 +23,7 @@ import {
   ToSqlOptions,
   TruncateQueryData,
 } from '../sql';
-import {
-  pushQueryArray,
-  pushQueryValue,
-  removeFromQuery,
-} from '../queryDataUtils';
+import { pushQueryArray, pushQueryValue } from '../queryDataUtils';
 import { Then } from './then';
 import { BooleanColumn } from '../columnSchema';
 import { Aggregate } from './aggregate';
@@ -56,6 +47,7 @@ import { QueryLog } from './log';
 import { QueryCallbacks } from './callbacks';
 import { QueryUpsert } from './upsert';
 import { QueryGet } from './get';
+import { MergeQueryMethods } from './merge';
 
 export type WindowArg<T extends Query> = Record<
   string,
@@ -69,7 +61,7 @@ export type WindowArgDeclaration<T extends Query = Query> = {
 
 type WindowResult<T extends Query, W extends WindowArg<T>> = SetQueryWindows<
   T,
-  PropertyKeyUnionToArray<keyof W>
+  Record<keyof W, true>
 >;
 
 export type OrderArg<T extends Query> =
@@ -103,10 +95,11 @@ export interface QueryMethods
     QueryLog,
     QueryCallbacks,
     QueryUpsert,
-    QueryGet {}
+    QueryGet,
+    MergeQueryMethods {}
 
 export class QueryMethods {
-  windows!: PropertyKey[];
+  windows!: EmptyObject;
   __model!: Query;
 
   all<T extends Query>(this: T): SetQueryReturnsAll<T> {
@@ -115,7 +108,6 @@ export class QueryMethods {
 
   _all<T extends Query>(this: T): SetQueryReturnsAll<T> {
     this.query.returnType = 'all';
-    removeFromQuery(this, 'take');
     return this as unknown as SetQueryReturnsAll<T>;
   }
 
@@ -125,7 +117,6 @@ export class QueryMethods {
 
   _take<T extends Query>(this: T): SetQueryReturnsOne<T> {
     this.query.returnType = 'oneOrThrow';
-    this.query.take = true;
     return this as unknown as SetQueryReturnsOne<T>;
   }
 
@@ -135,7 +126,6 @@ export class QueryMethods {
 
   _takeOptional<T extends Query>(this: T): SetQueryReturnsOneOptional<T> {
     this.query.returnType = 'one';
-    this.query.take = true;
     return this as unknown as SetQueryReturnsOneOptional<T>;
   }
 
@@ -145,7 +135,6 @@ export class QueryMethods {
 
   _rows<T extends Query>(this: T): SetQueryReturnsRows<T> {
     this.query.returnType = 'rows';
-    removeFromQuery(this, 'take');
     return this as unknown as SetQueryReturnsRows<T>;
   }
 
@@ -161,7 +150,6 @@ export class QueryMethods {
     select: S,
   ): SetQueryReturnsPluck<T, S> {
     this.query.returnType = 'pluck';
-    removeFromQuery(this, 'take');
     (this.query as SelectQueryData).select = [select as SelectItem];
     addParserForSelectItem(this, this.query.as || this.table, 'pluck', select);
     return this as unknown as SetQueryReturnsPluck<T, S>;
@@ -173,7 +161,6 @@ export class QueryMethods {
 
   _exec<T extends Query>(this: T): SetQueryReturnsVoid<T> {
     this.query.returnType = 'void';
-    removeFromQuery(this, 'take');
     return this as unknown as SetQueryReturnsVoid<T>;
   }
 
@@ -197,35 +184,33 @@ export class QueryMethods {
 
   find<T extends Query>(
     this: T,
-    ...args: GetTypesOrRaw<T['schema']['primaryTypes']>
+    value: T['shape'][T['singlePrimaryKey']]['type'] | RawExpression,
   ): SetQueryReturnsOne<WhereResult<T>> {
-    return this.clone()._find(...args);
+    return this.clone()._find(value);
   }
 
   _find<T extends Query>(
     this: T,
-    ...args: GetTypesOrRaw<T['schema']['primaryTypes']>
+    value: T['shape'][T['singlePrimaryKey']]['type'] | RawExpression,
   ): SetQueryReturnsOne<WhereResult<T>> {
-    const conditions: Partial<T['type']> = {};
-    this.schema.primaryKeys.forEach((key: string, i: number) => {
-      conditions[key as keyof T['type']] = args[i];
-    });
-    return this._where(conditions as WhereArg<T>)._take();
+    return this._where({
+      [this.singlePrimaryKey]: value,
+    } as WhereArg<T>)._take();
   }
 
   findOptional<T extends Query>(
     this: T,
-    ...args: GetTypesOrRaw<T['schema']['primaryTypes']>
+    value: T['shape'][T['singlePrimaryKey']]['type'] | RawExpression,
   ): SetQueryReturnsOneOptional<WhereResult<T>> {
-    return this.clone()._findOptional(...args);
+    return this.clone()._findOptional(value);
   }
 
   _findOptional<T extends Query>(
     this: T,
-    ...args: GetTypesOrRaw<T['schema']['primaryTypes']>
+    value: T['shape'][T['singlePrimaryKey']]['type'] | RawExpression,
   ): SetQueryReturnsOneOptional<WhereResult<T>> {
     return this._find(
-      ...args,
+      value,
     ).takeOptional() as unknown as SetQueryReturnsOneOptional<WhereResult<T>>;
   }
 
@@ -362,7 +347,6 @@ export class QueryMethods {
     const q = this._getOptional(raw<BooleanColumn>('true'));
     q.query.notFoundDefault = false;
     q.query.coalesceValue = false;
-    delete q.query.take;
     return q as unknown as SetQueryReturnsValue<T, BooleanColumn>;
   }
 
@@ -412,4 +396,5 @@ applyMixins(QueryMethods, [
   QueryCallbacks,
   QueryUpsert,
   QueryGet,
+  MergeQueryMethods,
 ]);

package/src/queryMethods/select.test.ts

@@ -29,6 +29,10 @@ const insertUserAndProfile = async () => {
 describe('selectMethods', () => {
   useTestDatabase();
 
+  it('table should have all columns selected if select was not applied', () => {
+    assertType<Awaited<typeof User>, UserRecord[]>();
+  });
+
   describe('select', () => {
     it('should have no effect if no columns provided', () => {
       const q = User.all();

package/src/queryMethods/select.ts

@@ -4,7 +4,9 @@ import {
   ColumnsParsers,
   Query,
   QueryBase,
+  QueryReturnsAll,
   QuerySelectAll,
+  queryTypeWithLimitOne,
 } from '../query';
 import {
   ArrayOfColumnsObjects,
@@ -75,7 +77,7 @@ type SelectResult<
 
 type SelectSubQueryResult<
   Arg extends Query & { [isRequiredRelationKey]?: boolean },
-> = Arg['returnType'] extends 'all'
+> = QueryReturnsAll<Arg['returnType']> extends true
   ? ArrayOfColumnsObjects<Arg['result']>
   : Arg['returnType'] extends 'valueOrThrow'
   ? Arg['result']['value']
@@ -107,7 +109,7 @@ export const addParserForSelectItem = <T extends Query>(
     const rel = arg(q);
     const parsers = getQueryParsers(rel);
     if (parsers) {
-      if (rel.query.take) {
+      if (queryTypeWithLimitOne[rel.query.returnType]) {
         addParserToQuery(q.query, key, (item) => parseRecord(parsers, item));
       } else {
         addParserToQuery(q.query, key, (items) =>

package/src/queryMethods/then.ts

@@ -49,7 +49,7 @@ export const handleResult: CommonQueryData['handleResult'] = async (
   q,
   result: QueryResult,
 ) => {
-  return parseResult(q, q.query.returnType, result);
+  return parseResult(q, q.query.returnType || 'all', result);
 };
 
 const then = async (
@@ -89,7 +89,7 @@ const then = async (
     }
 
     const queryResult = await q.query.adapter[
-      queryMethodByReturnType[q.query.returnType] as 'query'
+      queryMethodByReturnType[q.query.returnType || 'all'] as 'query'
     ](sql);
 
     if (q.query.log) {
@@ -119,7 +119,7 @@ const then = async (
 
 export const parseResult = (
   q: Query,
-  returnType: QueryReturnType,
+  returnType: QueryReturnType | undefined = 'all',
   result: QueryResult,
 ): unknown => {
   switch (returnType) {

package/src/queryMethods/update.ts

@@ -1,4 +1,4 @@
-import { Query, SetQueryReturnsRowCount } from '../query';
+import { Query, QueryReturnsAll, SetQueryReturnsRowCount } from '../query';
 import { pushQueryArray, pushQueryValue } from '../queryDataUtils';
 import { isRaw, RawExpression, StringKey } from '../common';
 import {
@@ -43,7 +43,7 @@ type UpdateBelongsToData<T extends Query, Rel extends BelongsToRelation> =
   | {
       create: InsertData<Rel['nestedCreateQuery']>;
     }
-  | (T['returnType'] extends 'all'
+  | (QueryReturnsAll<T['returnType']> extends true
       ? never
       : {
           upsert: {
@@ -56,7 +56,7 @@ type UpdateHasOneData<T extends Query, Rel extends HasOneRelation> =
   | { disconnect: boolean }
   | { delete: boolean }
   | { update: UpdateData<Rel['model']> }
-  | (T['returnType'] extends 'all'
+  | (QueryReturnsAll<T['returnType']> extends true
       ? never
       :
           | { set: WhereArg<Rel['model']> }
@@ -77,7 +77,7 @@ type UpdateHasManyData<T extends Query, Rel extends HasManyRelation> = {
     where: MaybeArray<WhereArg<Rel['model']>>;
     data: UpdateData<Rel['model']>;
   };
-} & (T['returnType'] extends 'all'
+} & (QueryReturnsAll<T['returnType']> extends true
   ? EmptyObject
   : {
       set?: MaybeArray<WhereArg<Rel['model']>>;
@@ -107,9 +107,9 @@ type UpdateRawArgs<T extends Query, ForceAll extends boolean> = (
   ? [update: RawExpression]
   : [update: RawExpression, forceAll: true];
 
-type UpdateResult<T extends Query> = T['hasSelect'] extends false
-  ? SetQueryReturnsRowCount<T>
-  : T;
+type UpdateResult<T extends Query> = T['hasSelect'] extends true
+  ? T
+  : SetQueryReturnsRowCount<T>;
 
 type ChangeCountArg<T extends Query> =
   | keyof T['shape']
@@ -179,7 +179,7 @@ export class Update {
     const prependRelations: Record<string, Record<string, unknown>> = {};
     const appendRelations: Record<string, Record<string, unknown>> = {};
 
-    const originalReturnType = query.returnType;
+    const originalReturnType = query.returnType || 'all';
 
     for (const key in data) {
       if (relations[key]) {
@@ -232,11 +232,13 @@ export class Update {
       query.returnType = 'all';
 
       if (state?.updateLater) {
-        this.schema.primaryKeys.forEach((key: string) => {
-          if (!query.select?.includes('*') && !query.select?.includes(key)) {
-            this._select(key as StringKey<keyof T['selectable']>);
-          }
-        });
+        if (!query.select?.includes('*')) {
+          this.primaryKeys.forEach((key) => {
+            if (!query.select?.includes(key)) {
+              this._select(key as StringKey<keyof T['selectable']>);
+            }
+          });
+        }
       }
 
       const { handleResult } = query;
@@ -250,7 +252,7 @@ export class Update {
           await Promise.all(state.updateLaterPromises as Promise<void>[]);
 
           const t = this.__model.clone().transacting(q);
-          const keys = this.schema.primaryKeys as string[];
+          const keys = this.primaryKeys;
           (
             t._whereIn as unknown as (
               keys: string[],
@@ -287,6 +289,11 @@ export class Update {
         appendRelationKeys.map((relationName) => {
           return (q: Query, result: Record<string, unknown>[]) => {
             const all = resultOfTypeAll || result;
+
+            if (q.query.returnType !== originalReturnType) {
+              q.query.returnType = originalReturnType;
+            }
+
             return (
               relations[relationName].nestedUpdate as HasOneNestedUpdate
             )?.(q, all, appendRelations[relationName] as NestedUpdateOneItem);

package/src/sql/fromAndAs.ts

@@ -1,7 +1,7 @@
 import { getRaw, isRaw } from '../common';
 import { quoteSchemaAndTable } from './common';
 import { QueryBase } from '../query';
-import { queryKeysOfNotSimpleQuery, SelectQueryData } from './types';
+import { checkIfASimpleQuery, SelectQueryData } from './types';
 import { ToSqlCtx } from './toSql';
 
 export const pushFromAndAs = (
@@ -38,9 +38,8 @@ const getFrom = (
       }
 
       const q = query.from.query;
-      const keys = Object.keys(q) as (keyof SelectQueryData)[];
       // if query contains more than just schema return (SELECT ...)
-      if (keys.some((key) => queryKeysOfNotSimpleQuery.includes(key))) {
+      if (!checkIfASimpleQuery(q)) {
         const sql = query.from.toSql({ values });
         return `(${sql.text})`;
       }

package/src/sql/select.ts

@@ -142,7 +142,7 @@ const pushSubQuerySql = (
   values: unknown[],
   list: string[],
 ) => {
-  const { returnType } = query.query;
+  const { returnType = 'all' } = query.query;
   switch (returnType) {
     case 'all':
     case 'one':
@@ -160,7 +160,6 @@ const pushSubQuerySql = (
       select[0] = { selectAs: { c: first } } as SelectItem;
       query = query._wrap(query.__model.clone()) as unknown as typeof query;
       query._getOptional(raw<StringColumn>(`COALESCE(json_agg("c"), '[]')`));
-      delete query.query.take;
       break;
     }
     case 'rows':

package/src/sql/toSql.ts

@@ -1,5 +1,5 @@
 import { getRaw, isRaw } from '../common';
-import { Query } from '../query';
+import { Query, queryTypeWithLimitOne } from '../query';
 import { addValue, q, qc } from './common';
 import { JoinItem, QueryData, Sql } from './types';
 import { pushDistinctSql } from './distinct';
@@ -156,8 +156,9 @@ const makeSql = (model: Query, { values = [] }: ToSqlOptions = {}): Sql => {
     pushOrderBySql(ctx, quotedAs, query.order);
   }
 
-  if (query.take || query.limit !== undefined) {
-    sql.push(`LIMIT ${addValue(values, query.take ? 1 : query.limit)}`);
+  const limit = queryTypeWithLimitOne[query.returnType] ? 1 : query.limit;
+  if (limit) {
+    sql.push(`LIMIT ${addValue(values, limit)}`);
   }
 
   if (query.offset) {

package/src/sql/types.ts

@@ -21,8 +21,13 @@ export type Sql = {
 };
 
 // used in `from` logic to decide if convert query to sql or just write table name
-export const queryKeysOfNotSimpleQuery: (keyof SelectQueryData)[] = [
-  'take',
+export const checkIfASimpleQuery = (q: QueryData) => {
+  if (q.returnType && q.returnType !== 'all') return false;
+  const keys = Object.keys(q) as (keyof SelectQueryData)[];
+  return !keys.some((key) => queryKeysOfNotSimpleQuery.includes(key));
+};
+
+const queryKeysOfNotSimpleQuery: (keyof SelectQueryData)[] = [
   'with',
   'as',
   'from',
@@ -51,7 +56,6 @@ export type CommonQueryData = {
   inTransaction?: boolean;
   wrapInTransaction?: boolean;
   throwOnNotFound?: boolean;
-  take?: boolean;
   with?: WithItem[];
   withShapes?: Record<string, ColumnsShape>;
   schema?: string;

package/src/sql/window.ts

@@ -5,7 +5,7 @@ import { expressionToSql, q } from './common';
 import { orderByToSql } from './orderBy';
 
 export const windowToSql = <T extends Query>(
-  window: T['windows'][number] | WindowDeclaration | RawExpression,
+  window: keyof T['windows'] | WindowDeclaration | RawExpression,
   values: unknown[],
   quotedAs?: string,
 ) => {

package/src/sql/with.ts

@@ -7,26 +7,31 @@ export const pushWithSql = (
   ctx: ToSqlCtx,
   withData: Exclude<QueryData['with'], undefined>,
 ) => {
-  withData.forEach((withItem) => {
-    const [name, options, query] = withItem;
+  if (!withData.length) return;
 
-    let inner: string;
-    if (isRaw(query)) {
-      inner = getRaw(query, ctx.values);
-    } else {
-      inner = query.toSql({ values: ctx.values }).text;
-    }
+  ctx.sql.push(
+    'WITH',
+    withData
+      .map((withItem) => {
+        const [name, options, query] = withItem;
 
-    ctx.sql.push(
-      `WITH ${options.recursive ? 'RECURSIVE ' : ''}${q(name)}${
-        options.columns ? `(${options.columns.map(q).join(', ')})` : ''
-      } AS ${
-        options.materialized
-          ? 'MATERIALIZED '
-          : options.notMaterialized
-          ? 'NOT MATERIALIZED '
-          : ''
-      }(${inner})`,
-    );
-  });
+        let inner: string;
+        if (isRaw(query)) {
+          inner = getRaw(query, ctx.values);
+        } else {
+          inner = query.toSql({ values: ctx.values }).text;
+        }
+
+        return `${options.recursive ? 'RECURSIVE ' : ''}${q(name)}${
+          options.columns ? `(${options.columns.map(q).join(', ')})` : ''
+        } AS ${
+          options.materialized
+            ? 'MATERIALIZED '
+            : options.notMaterialized
+            ? 'NOT MATERIALIZED '
+            : ''
+        }(${inner})`;
+      })
+      .join(', '),
+  );
 };