feat: Adding head and count option on select(), insert(), update(), d…

…elete(), and rpc() (#147) * select method now accepts home and count option * removed head and count option from PostgrestTransformBuilder and added home and count option to select in PostgrestQueryBuilder * select now returns count value when count option is selected * Enabled head for select * Updated the test snapshots to include count field * Set default values for head and count Co-authored-by: Bobbie Soedirgo <bobbie@soedirgo.dev> * Added basic test cases for select with head and count parameters * Added count and head option to rpc, and added count option to insert, update, delete * Added count tests for insert, update, delete rpc * Added head and count test for rpc * Updated test so that count value for 'planned' and 'estimated' would match any number Co-authored-by: Bobbie Soedirgo <bobbie@soedirgo.dev>
supabase · Jan 15, 2021 · 8cf6e00 · 8cf6e00
1 parent d56d69a
commit 8cf6e00
Show file tree

Hide file tree

Showing 7 changed files with 986 additions and 20 deletions.
diff --git a/src/PostgrestClient.ts b/src/PostgrestClient.ts
@@ -1,6 +1,5 @@
 import PostgrestQueryBuilder from './lib/PostgrestQueryBuilder'
 import PostgrestTransformBuilder from './lib/PostgrestTransformBuilder'
-import { PostgrestBuilder } from './lib/types'
 
 export default class PostgrestClient {
   url: string

diff --git a/src/lib/PostgrestQueryBuilder.ts b/src/lib/PostgrestQueryBuilder.ts
@@ -17,8 +17,19 @@ export default class PostgrestQueryBuilder<T> extends PostgrestBuilder<T> {
    * Performs vertical filtering with SELECT.
    *
    * @param columns  The columns to retrieve, separated by commas.
+   * @param head  When set to true, select will void data.
+   * @param count  Count algorithm to use to count rows in a table.
    */
-  select(columns = '*'): PostgrestFilterBuilder<T> {
+  select(
+    columns = '*',
+    {
+      head = false,
+      count = null,
+    }: {
+      head?: boolean
+      count?: null | 'exact' | 'planned' | 'estimated'
+    } = {}
+  ): PostgrestFilterBuilder<T> {
     this.method = 'GET'
     // Remove whitespaces except when quoted
     let quoted = false
@@ -35,6 +46,12 @@ export default class PostgrestQueryBuilder<T> extends PostgrestBuilder<T> {
       })
       .join('')
     this.url.searchParams.set('select', cleanedColumns)
+    if (count) {
+      this.headers['Prefer'] = `count=${count}`
+    }
+    if (head) {
+      this.method = 'HEAD'
+    }
     return new PostgrestFilterBuilder(this)
   }
 
@@ -52,10 +69,12 @@ export default class PostgrestQueryBuilder<T> extends PostgrestBuilder<T> {
       upsert = false,
       onConflict,
       returning = 'representation',
+      count = null,
     }: {
       upsert?: boolean
       onConflict?: string
       returning?: 'minimal' | 'representation'
+      count?: null | 'exact' | 'planned' | 'estimated'
     } = {}
   ): PostgrestFilterBuilder<T> {
     this.method = 'POST'
@@ -67,6 +86,10 @@ export default class PostgrestQueryBuilder<T> extends PostgrestBuilder<T> {
 
     if (upsert && onConflict !== undefined) this.url.searchParams.set('on_conflict', onConflict)
     this.body = values
+    if (count) {
+      prefersHeaders.push(`count=${count}`)
+      this.headers['Prefer'] = prefersHeaders.join(',')
+    }
     return new PostgrestFilterBuilder(this)
   }
 
@@ -78,11 +101,23 @@ export default class PostgrestQueryBuilder<T> extends PostgrestBuilder<T> {
    */
   update(
     values: Partial<T>,
-    { returning = 'representation' }: { returning?: 'minimal' | 'representation' } = {}
+    {
+      returning = 'representation',
+      count = null,
+    }: {
+      returning?: 'minimal' | 'representation'
+      count?: null | 'exact' | 'planned' | 'estimated'
+    } = {}
   ): PostgrestFilterBuilder<T> {
     this.method = 'PATCH'
-    this.headers['Prefer'] = `return=${returning}`
+    let prefersHeaders = []
+    prefersHeaders.push(`return=${returning}`)
+    this.headers['Prefer'] = prefersHeaders.join(',')
     this.body = values
+    if (count) {
+      prefersHeaders.push(`count=${count}`)
+      this.headers['Prefer'] = prefersHeaders.join(',')
+    }
     return new PostgrestFilterBuilder(this)
   }
 
@@ -93,16 +128,41 @@ export default class PostgrestQueryBuilder<T> extends PostgrestBuilder<T> {
    */
   delete({
     returning = 'representation',
-  }: { returning?: 'minimal' | 'representation' } = {}): PostgrestFilterBuilder<T> {
+    count = null,
+  }: {
+    returning?: 'minimal' | 'representation'
+    count?: null | 'exact' | 'planned' | 'estimated'
+  } = {}): PostgrestFilterBuilder<T> {
     this.method = 'DELETE'
-    this.headers['Prefer'] = `return=${returning}`
+    let prefersHeaders = []
+    prefersHeaders.push(`return=${returning}`)
+    this.headers['Prefer'] = prefersHeaders.join(',')
+    if (count) {
+      prefersHeaders.push(`count=${count}`)
+      this.headers['Prefer'] = prefersHeaders.join(',')
+    }
     return new PostgrestFilterBuilder(this)
   }
 
   /** @internal */
-  rpc(params?: object): PostgrestTransformBuilder<T> {
+  rpc(
+    params?: object,
+    {
+      head = false,
+      count = null,
+    }: {
+      head?: boolean
+      count?: null | 'exact' | 'planned' | 'estimated'
+    } = {}
+  ): PostgrestTransformBuilder<T> {
     this.method = 'POST'
     this.body = params
+    if (count) {
+      this.headers['Prefer'] = `count=${count}`
+    }
+    if (head) {
+      this.method = 'HEAD'
+    }
     return new PostgrestTransformBuilder(this)
   }
 }
diff --git a/src/lib/PostgrestTransformBuilder.ts b/src/lib/PostgrestTransformBuilder.ts
@@ -59,10 +59,7 @@ export default class PostgrestTransformBuilder<T> extends PostgrestBuilder<T> {
    * @param count  The maximum no. of rows to limit to.
    * @param foreignTable  The foreign table to use (for foreign columns).
    */
-  limit(
-    count: number,
-    { foreignTable }: { foreignTable?: string } = {}
-  ): this {
+  limit(count: number, { foreignTable }: { foreignTable?: string } = {}): this {
     const key = typeof foreignTable === 'undefined' ? 'limit' : `"${foreignTable}".limit`
     this.url.searchParams.set(key, `${count}`)
     return this
@@ -75,11 +72,7 @@ export default class PostgrestTransformBuilder<T> extends PostgrestBuilder<T> {
    * @param to  The last index to which to limit the result, inclusive.
    * @param foreignTable  The foreign table to use (for foreign columns).
    */
-  range(
-    from: number,
-    to: number,
-    { foreignTable }: { foreignTable?: string } = {}
-  ): this {
+  range(from: number, to: number, { foreignTable }: { foreignTable?: string } = {}): this {
     const keyOffset = typeof foreignTable === 'undefined' ? 'offset' : `"${foreignTable}".offset`
     const keyLimit = typeof foreignTable === 'undefined' ? 'limit' : `"${foreignTable}".limit`
     this.url.searchParams.set(keyOffset, `${from}`)

diff --git a/src/lib/types.ts b/src/lib/types.ts
@@ -26,12 +26,14 @@ interface PostgrestResponseSuccess<T> extends PostgrestResponseBase {
   error: null
   data: T[]
   body: T[]
+  count: number | null
 }
 interface PostgrestResponseFailure extends PostgrestResponseBase {
   error: PostgrestError
   data: null
   // For backward compatibility: body === data
   body: null
+  count: null
 }
 export type PostgrestResponse<T> = PostgrestResponseSuccess<T> | PostgrestResponseFailure
 
@@ -80,18 +82,36 @@ export abstract class PostgrestBuilder<T> implements PromiseLike<PostgrestRespon
       body: JSON.stringify(this.body),
     })
       .then(async (res) => {
-        let error, data
+        let error, data, count
         if (res.ok) {
           error = null
-          const isReturnMinimal = this.headers['Prefer']?.split(',').includes('return=minimal')
-          data = isReturnMinimal ? null : await res.json()
+          if (this.method !== 'HEAD') {
+            const isReturnMinimal = this.headers['Prefer']?.split(',').includes('return=minimal')
+            data = isReturnMinimal ? null : await res.json()
+          } else {
+            data = null
+          }
+
+          const countHeader = this.headers['Prefer']?.match(/count=(exact|planned|estimated)/)
+          if (countHeader) {
+            const contentRange = res.headers.get('content-range')?.split('/')
+            if (contentRange && contentRange.length > 1) {
+              count = parseInt(contentRange[1])
+            } else {
+              count = null
+            }
+          } else {
+            count = null
+          }
         } else {
           error = await res.json()
           data = null
+          count = null
         }
         const postgrestResponse: PostgrestResponse<T> = {
           error,
           data,
+          count,
           status: res.status,
           statusText: res.statusText,
           body: data,