diff --git a/server/internal/db/query_test.go b/server/internal/db/query_test.go index 6c29924..4910328 100644 --- a/server/internal/db/query_test.go +++ b/server/internal/db/query_test.go @@ -658,6 +658,158 @@ func TestHavingClause(t *testing.T) { } } +func TestQueryReturning(t *testing.T) { + testCases := []struct { + name string + dbType db.DBType + buildQuery func(q *db.Query) *db.Query + expectedSQL string + expectedArgs []any + }{ + { + name: "SQLite INSERT with RETURNING single column", + dbType: db.DBTypeSQLite, + buildQuery: func(q *db.Query) *db.Query { + return q.Insert("users", "name", "email"). + Values(2). + AddArgs("John Doe", "john@example.com"). + Returning("id") + }, + expectedSQL: "INSERT INTO users (name, email) VALUES (?, ?) RETURNING id", + expectedArgs: []any{"John Doe", "john@example.com"}, + }, + { + name: "PostgreSQL INSERT with RETURNING single column", + dbType: db.DBTypePostgres, + buildQuery: func(q *db.Query) *db.Query { + return q.Insert("users", "name", "email"). + Values(2). + AddArgs("John Doe", "john@example.com"). + Returning("id") + }, + expectedSQL: "INSERT INTO users (name, email) VALUES ($1, $2) RETURNING id", + expectedArgs: []any{"John Doe", "john@example.com"}, + }, + { + name: "SQLite INSERT with RETURNING multiple columns", + dbType: db.DBTypeSQLite, + buildQuery: func(q *db.Query) *db.Query { + return q.Insert("users", "name", "email"). + Values(2). + AddArgs("John Doe", "john@example.com"). + Returning("id", "created_at") + }, + expectedSQL: "INSERT INTO users (name, email) VALUES (?, ?) RETURNING id, created_at", + expectedArgs: []any{"John Doe", "john@example.com"}, + }, + { + name: "PostgreSQL INSERT with RETURNING multiple columns", + dbType: db.DBTypePostgres, + buildQuery: func(q *db.Query) *db.Query { + return q.Insert("users", "name", "email"). + Values(2). + AddArgs("John Doe", "john@example.com"). + Returning("id", "created_at") + }, + expectedSQL: "INSERT INTO users (name, email) VALUES ($1, $2) RETURNING id, created_at", + expectedArgs: []any{"John Doe", "john@example.com"}, + }, + { + name: "SQLite UPDATE with RETURNING", + dbType: db.DBTypeSQLite, + buildQuery: func(q *db.Query) *db.Query { + return q.Update("users"). + Set("name").Placeholder("Jane Doe"). + Where("id = ").Placeholder(1). + Returning("name", "updated_at") + }, + expectedSQL: "UPDATE users SET name = ? WHERE id = ? RETURNING name, updated_at", + expectedArgs: []any{"Jane Doe", 1}, + }, + { + name: "PostgreSQL UPDATE with RETURNING", + dbType: db.DBTypePostgres, + buildQuery: func(q *db.Query) *db.Query { + return q.Update("users"). + Set("name").Placeholder("Jane Doe"). + Where("id = ").Placeholder(1). + Returning("name", "updated_at") + }, + expectedSQL: "UPDATE users SET name = $1 WHERE id = $2 RETURNING name, updated_at", + expectedArgs: []any{"Jane Doe", 1}, + }, + { + name: "SQLite DELETE with RETURNING", + dbType: db.DBTypeSQLite, + buildQuery: func(q *db.Query) *db.Query { + return q.Delete(). + From("users"). + Where("id = ").Placeholder(1). + Returning("id", "name", "email") + }, + expectedSQL: "DELETE FROM users WHERE id = ? RETURNING id, name, email", + expectedArgs: []any{1}, + }, + { + name: "PostgreSQL DELETE with RETURNING", + dbType: db.DBTypePostgres, + buildQuery: func(q *db.Query) *db.Query { + return q.Delete(). + From("users"). + Where("id = ").Placeholder(1). + Returning("id", "name", "email") + }, + expectedSQL: "DELETE FROM users WHERE id = $1 RETURNING id, name, email", + expectedArgs: []any{1}, + }, + { + name: "SQLite INSERT with RETURNING *", + dbType: db.DBTypeSQLite, + buildQuery: func(q *db.Query) *db.Query { + return q.Insert("users", "name", "email"). + Values(2). + AddArgs("John Doe", "john@example.com"). + Returning("*") + }, + expectedSQL: "INSERT INTO users (name, email) VALUES (?, ?) RETURNING *", + expectedArgs: []any{"John Doe", "john@example.com"}, + }, + { + name: "PostgreSQL INSERT with RETURNING *", + dbType: db.DBTypePostgres, + buildQuery: func(q *db.Query) *db.Query { + return q.Insert("users", "name", "email"). + Values(2). + AddArgs("John Doe", "john@example.com"). + Returning("*") + }, + expectedSQL: "INSERT INTO users (name, email) VALUES ($1, $2) RETURNING *", + expectedArgs: []any{"John Doe", "john@example.com"}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + query := db.NewQuery(tc.dbType) + result := tc.buildQuery(query) + + if result.String() != tc.expectedSQL { + t.Errorf("Expected SQL: %s, got: %s", tc.expectedSQL, result.String()) + } + + if len(result.Args()) != len(tc.expectedArgs) { + t.Errorf("Expected %d args, got %d", len(tc.expectedArgs), len(result.Args())) + } + + for i, arg := range result.Args() { + if arg != tc.expectedArgs[i] { + t.Errorf("Expected arg %d to be %v, got %v", i, tc.expectedArgs[i], arg) + } + } + }) + } +} + func TestComplexQueries(t *testing.T) { tests := []struct { name string