Postgres - Composite key value fix (#77)

artie-labs · Feb 2, 2024 · 00524c0 · 00524c0
1 parent 5a6051a
commit 00524c0
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/lib/postgres/queries/queries.go b/lib/postgres/queries/queries.go
@@ -38,14 +38,17 @@ type SelectTableQueryArgs struct {
 }
 
 func SelectTableQuery(args SelectTableQueryArgs) string {
-	orderByFragment := strings.Join(quotedIdentifiers(args.OrderBy), ",")
-	if args.Descending {
-		orderByFragment += " DESC"
+	var fragments []string
+	for _, orderBy := range args.OrderBy {
+		fragment := pgx.Identifier{orderBy}.Sanitize()
+		if args.Descending {
+			fragment += " DESC"
+		}
+
+		fragments = append(fragments, fragment)
 	}
-
-	// TODO: Make sure keys are being escaped properly
-	return fmt.Sprintf(`SELECT %s FROM %s ORDER BY %s LIMIT 1`,
-		strings.Join(args.Keys, ","), pgx.Identifier{args.Schema, args.TableName}.Sanitize(), orderByFragment)
+	return fmt.Sprintf(`SELECT %s FROM %s ORDER BY %s LIMIT 1`, strings.Join(args.Keys, ","),
+		pgx.Identifier{args.Schema, args.TableName}.Sanitize(), strings.Join(fragments, ","))
 }
 
 type RetrievePrimaryKeysArgs struct {

diff --git a/lib/postgres/queries/queries_test.go b/lib/postgres/queries/queries_test.go
@@ -48,7 +48,7 @@ func TestSelectTableQuery(t *testing.T) {
 			OrderBy:    []string{"e", "f", "g"},
 			Descending: true,
 		})
-		assert.Equal(t, `SELECT a,b,c FROM "schema"."table" ORDER BY "e","f","g" DESC LIMIT 1`, query)
+		assert.Equal(t, `SELECT a,b,c FROM "schema"."table" ORDER BY "e" DESC,"f" DESC,"g" DESC LIMIT 1`, query)
 	}
 }