fix(postgrest): avoid duplicated columns and prefer fields (#463)

supabase · Jul 13, 2024 · e4f85f3 · e4f85f3
1 parent 0667b9f
commit e4f85f3
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 9 deletions.
diff --git a/Sources/Helpers/HTTP/HTTPRequest.swift b/Sources/Helpers/HTTP/HTTPRequest.swift
@@ -68,3 +68,13 @@ package enum HTTPMethod: String, Sendable {
   case patch = "PATCH"
   case options = "OPTIONS"
 }
+
+extension [URLQueryItem] {
+  package mutating func appendOrUpdate(_ queryItem: URLQueryItem) {
+    if let index = firstIndex(where: { $0.name == queryItem.name }) {
+      self[index] = queryItem
+    } else {
+      self.append(queryItem)
+    }
+  }
+}
diff --git a/Sources/PostgREST/PostgrestQueryBuilder.swift b/Sources/PostgREST/PostgrestQueryBuilder.swift
@@ -27,7 +27,7 @@ public final class PostgrestQueryBuilder: PostgrestBuilder {
       }
       .joined(separator: "")
 
-      $0.request.query.append(URLQueryItem(name: "select", value: cleanedColumns))
+      $0.request.query.appendOrUpdate(URLQueryItem(name: "select", value: cleanedColumns))
 
       if let count {
         $0.request.headers["Prefer"] = "count=\(count.rawValue)"
@@ -73,7 +73,7 @@ public final class PostgrestQueryBuilder: PostgrestBuilder {
       {
         let allKeys = jsonObject.flatMap(\.keys)
         let uniqueKeys = Set(allKeys).sorted()
-        $0.request.query.append(URLQueryItem(
+        $0.request.query.appendOrUpdate(URLQueryItem(
           name: "columns",
           value: uniqueKeys.joined(separator: ",")
         ))
@@ -108,7 +108,7 @@ public final class PostgrestQueryBuilder: PostgrestBuilder {
         "return=\(returning.rawValue)",
       ]
       if let onConflict {
-        $0.request.query.append(URLQueryItem(name: "on_conflict", value: onConflict))
+        $0.request.query.appendOrUpdate(URLQueryItem(name: "on_conflict", value: onConflict))
       }
       $0.request.body = try configuration.encoder.encode(values)
       if let count {
@@ -126,7 +126,7 @@ public final class PostgrestQueryBuilder: PostgrestBuilder {
       {
         let allKeys = jsonObject.flatMap(\.keys)
         let uniqueKeys = Set(allKeys).sorted()
-        $0.request.query.append(URLQueryItem(
+        $0.request.query.appendOrUpdate(URLQueryItem(
           name: "columns",
           value: uniqueKeys.joined(separator: ",")
         ))

diff --git a/Sources/PostgREST/PostgrestTransformBuilder.swift b/Sources/PostgREST/PostgrestTransformBuilder.swift
@@ -22,13 +22,21 @@ public class PostgrestTransformBuilder: PostgrestBuilder {
     }
     .joined(separator: "")
     mutableState.withValue {
-      $0.request.query.append(URLQueryItem(name: "select", value: cleanedColumns))
+      $0.request.query.appendOrUpdate(URLQueryItem(name: "select", value: cleanedColumns))
 
-      if $0.request.headers["Prefer"] != nil {
-        $0.request.headers["Prefer", default: ""] += ","
-      }
+      if let prefer = $0.request.headers["Prefer"] {
+        var components = prefer.components(separatedBy: ",")
+
+        if let index = components.firstIndex(where: { $0.hasPrefix("return=") }) {
+          components[index] = "return=representation"
+        } else {
+          components.append("return=representation")
+        }
 
-      $0.request.headers["Prefer", default: ""] += "return=representation"
+        $0.request.headers["Prefer"] = components.joined(separator: ",")
+      } else {
+        $0.request.headers["Prefer"] = "return=representation"
+      }
     }
     return self
   }

diff --git a/Tests/PostgRESTTests/BuildURLRequestTests.swift b/Tests/PostgRESTTests/BuildURLRequestTests.swift
@@ -141,6 +141,17 @@ final class BuildURLRequestTests: XCTestCase {
             ]
           )
       },
+      TestCase(name: "select after bulk upsert") { client in
+        try client.from("users")
+          .upsert(
+            [
+              User(email: "[email protected]"),
+              User(email: "[email protected]"),
+            ],
+            onConflict: "username"
+          )
+          .select()
+      },
       TestCase(name: "test upsert ignoring duplicates") { client in
         try client.from("users")
           .upsert(User(email: "[email protected]"), ignoreDuplicates: true)

diff --git a/...ESTTests/__Snapshots__/BuildURLRequestTests/testBuildRequest.select-after-bulk-upsert.txt b/...ESTTests/__Snapshots__/BuildURLRequestTests/testBuildRequest.select-after-bulk-upsert.txt
@@ -0,0 +1,8 @@
+curl \
+	--request POST \
+	--header "Accept: application/json" \
+	--header "Content-Type: application/json" \
+	--header "Prefer: resolution=merge-duplicates,return=representation" \
+	--header "X-Client-Info: postgrest-swift/x.y.z" \
+	--data "[{\"email\":\"[email protected]\"},{\"email\":\"[email protected]\"}]" \
+	"https://example.supabase.co/users?columns=email&on_conflict=username&select=*"