From 3dd7bf74a1930d6792641c166caad4373575fb56 Mon Sep 17 00:00:00 2001
From: Ben <ben@zfben.com>
Date: Fri, 14 Feb 2025 01:24:54 +0000
Subject: [PATCH] feat: add orWhere method to QueryBuilder for flexible query
 conditions

---
 src/__tests__/query-builder/query.test.ts | 29 +++++++++
 src/query-builder.ts                      | 78 +++++++++++++++++++++--
 2 files changed, 101 insertions(+), 6 deletions(-)

diff --git a/src/__tests__/query-builder/query.test.ts b/src/__tests__/query-builder/query.test.ts
index ad49258..a658b88 100644
--- a/src/__tests__/query-builder/query.test.ts
+++ b/src/__tests__/query-builder/query.test.ts
@@ -198,6 +198,35 @@ describe('QueryBuilder/query', () => {
     })
   })
 
+  describe('orWhere', () => {
+    it('column and value', async () => {
+      const result = await new QueryBuilder(client, 'query')
+        .where('name', 'Alice')
+        .orWhere('name', 'Bob')
+        .pluck('id')
+
+      expect(result).toEqual([1, 2])
+    })
+
+    it('column, operator, and value', async () => {
+      const result = await new QueryBuilder(client, 'query')
+        .where('name', 'Alice')
+        .orWhere('name', '!=', 'Bob')
+        .pluck('id')
+
+      expect(result).toEqual([1])
+    })
+
+    it('multiple conditions', async () => {
+      const result = await new QueryBuilder(client, 'query')
+        .where('name', 'Alice')
+        .orWhere('id', '>', 1)
+        .pluck('id')
+
+      expect(result).toEqual([1, 2])
+    })
+  })
+
   describe('limit', () => {
     it('limits the number of results', async () => {
       const result = await new QueryBuilder(client, 'query')
diff --git a/src/query-builder.ts b/src/query-builder.ts
index 82e583c..aaa395a 100644
--- a/src/query-builder.ts
+++ b/src/query-builder.ts
@@ -85,6 +85,7 @@ export class QueryBuilder<
   private table: T
   private selectColumns: (ColumnName<T> | JsonSelectField<T>)[] = []
   private whereConditions: {
+    type: 'AND' | 'OR'
     column: ColumnName<T> | string
     operator: Operator
     value: any
@@ -168,6 +169,7 @@ export class QueryBuilder<
       !['IS NULL', 'IS NOT NULL'].includes(operatorOrValue)
     ) {
       this.whereConditions.push({
+        type: 'AND',
         column,
         operator: '=',
         value: operatorOrValue,
@@ -180,6 +182,68 @@ export class QueryBuilder<
       throw new Error(`Invalid operator: ${operatorOrValue}`)
 
     this.whereConditions.push({
+      type: 'AND',
+      column,
+      operator: operatorOrValue,
+      value,
+    })
+
+    return this
+  }
+
+  /**
+   * Applies an OR WHERE condition to the query builder.
+   * @param column - The column to filter on.
+   * @param operator - The operator to use for comparison.
+   * @param value - The value to compare against.
+   * @example
+   * ```ts
+   * await query('users').where('id', 1).orWhere('id', 2) // WHERE id = 1 OR id = 2
+   * ```
+   */
+  orWhere<C extends ColumnName<T>>(
+    column: C,
+    operator: (typeof NormalOperators)[number],
+    value?: ColumnValue<T, C>
+  ): QueryBuilder<T, TResult>
+  orWhere<C extends ColumnName<T>>(
+    column: C,
+    operator: (typeof ArrayOperators)[number],
+    value: ColumnValue<T, C>[]
+  ): QueryBuilder<T, TResult>
+  orWhere<C extends ColumnName<T>>(
+    column: C,
+    operator: (typeof NullOperators)[number]
+  ): QueryBuilder<T, TResult>
+  orWhere<C extends ColumnName<T>>(
+    column: C,
+    operator: (typeof JsonOperators)[number],
+    value: Partial<ColumnValue<T, C>>
+  ): QueryBuilder<T, TResult>
+  orWhere<C extends ColumnName<T>>(
+    column: C,
+    value: ColumnValue<T, C>
+  ): QueryBuilder<T, TResult>
+  orWhere(column: ColumnName<T>, operatorOrValue: Operator | any, value?: any) {
+    if (
+      typeof value === 'undefined' &&
+      !['IS NULL', 'IS NOT NULL'].includes(operatorOrValue)
+    ) {
+      this.whereConditions.push({
+        type: 'OR',
+        column,
+        operator: '=',
+        value: operatorOrValue,
+      })
+
+      return this
+    }
+
+    if (!Operators.includes(operatorOrValue))
+      throw new Error(`Invalid operator: ${operatorOrValue}`)
+
+    this.whereConditions.push({
+      type: 'OR',
       column,
       operator: operatorOrValue,
       value,
@@ -249,23 +313,25 @@ export class QueryBuilder<
       sql.push(
         'WHERE',
         this.whereConditions
-          .map(({ column, operator, value }) => {
+          .map(({ type, column, operator, value }, index) => {
+            const prefix = index > 0 ? `${type} ` : ''
+
             if (operator === 'IS NULL' || operator === 'IS NOT NULL')
-              return `${escapeIdentifier(column)} ${operator}`
+              return `${prefix}${escapeIdentifier(column)} ${operator}`
 
             if (operator === 'IN' || operator === 'NOT IN') {
               params.push(value)
 
               if (mode === 'update')
-                return `${escapeIdentifier(column)} = ANY(?)`
+                return `${prefix}${escapeIdentifier(column)} = ANY(?)`
 
-              return `${escapeIdentifier(column)} ${operator} (${value.map(() => '?').join(',')})`
+              return `${prefix}${escapeIdentifier(column)} ${operator} (${value.map(() => '?').join(',')})`
             }
 
             params.push(value)
-            return `${escapeIdentifier(column)} ${operator} ?`
+            return `${prefix}${escapeIdentifier(column)} ${operator} ?`
           })
-          .join(' AND ')
+          .join(' ')
       )
     }