From 282b8e7a77e72a889988a48065f807ba818b3a39 Mon Sep 17 00:00:00 2001
From: manelcecs <manuelc.p.g@gmail.com>
Date: Thu, 11 Jul 2024 10:40:50 +0200
Subject: [PATCH] Add 'count' function

---
 src/db/util/raw-model.ts | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/src/db/util/raw-model.ts b/src/db/util/raw-model.ts
index d8b778bf..00c05364 100644
--- a/src/db/util/raw-model.ts
+++ b/src/db/util/raw-model.ts
@@ -82,6 +82,20 @@ export type DestroyFn<F extends FieldDefinition> = (args: {
 
 export type TruncateFn = (trx?: Knex.Transaction<any, any>) => Promise<void>;
 
+export type CountFn<F extends FieldDefinition, AdditionalArgs = {}> = (
+  args?: {
+    where?: WhereCond<F>;
+    trx?: Knex.Transaction<any, any>;
+    /**
+     * Will produce `SELECT COUNT(DISTINCT "columnName")` and if there
+     * is an expression in `COUNT`, it will compute the number of input
+     * rows in which the input value is **not null**. So, besides selecting
+     * unique values of "columnName", it will also filter out `NULL` values
+     */
+    distinct?: Array<keyof InstanceDataOf<F>>;
+  } & AdditionalArgs
+) => Promise<bigint>;
+
 export type ModelInternals<F extends FieldDefinition> = {
   readonly type: 'single-table';
   readonly tableName: string;
@@ -101,6 +115,7 @@ export type Model<F extends FieldDefinition, AdditionalFindArgs = {}> = {
   readonly update: UpdateFn<F>;
   readonly destroy: DestroyFn<F>;
   readonly truncate: TruncateFn;
+  readonly count: CountFn<F, AdditionalFindArgs>;
 };
 
 export type InstanceDataOfModel<M extends Model<any>> = M extends Model<infer F>
@@ -251,6 +266,25 @@ export const defineRawModel =
       await builder.truncate();
     };
 
+    const count: CountFn<F> = async ({ where, trx, distinct } = {}) => {
+      const builder = trx ? masterTable().transacting(trx) : replicaTable();
+      const query = builder.where(prepareCondition(where ?? {}));
+
+      if (distinct !== undefined) {
+        query.countDistinct(distinct);
+      } else {
+        query.count('*');
+      }
+
+      const res = (await query) as unknown as Array<{ count: string }>;
+
+      if (res.length !== 1) {
+        throw new Error('Count result must be an array of single element');
+      }
+
+      return BigInt(res[0].count);
+    };
+
     return {
       _internals: {
         type: 'single-table',
@@ -265,5 +299,6 @@ export const defineRawModel =
       update,
       destroy,
       truncate,
+      count,
     };
   };