From ffaa47a777f234482a4a2afddcd7a88151f0bad7 Mon Sep 17 00:00:00 2001
From: steebchen <contact@luca-steeb.com>
Date: Wed, 24 Apr 2024 17:19:56 +0100
Subject: [PATCH] feat(builder): add omit option

---
 generator/templates/_header.gotpl      |   1 +
 generator/templates/actions/find.gotpl |  19 +++++
 test/features/omit/default_test.go     | 103 +++++++++++++++++++++++++
 test/features/omit/schema.prisma       |  19 +++++
 4 files changed, 142 insertions(+)
 create mode 100644 test/features/omit/default_test.go
 create mode 100644 test/features/omit/schema.prisma

diff --git a/generator/templates/_header.gotpl b/generator/templates/_header.gotpl
index fc69e4665..3c769e9e0 100644
--- a/generator/templates/_header.gotpl
+++ b/generator/templates/_header.gotpl
@@ -8,6 +8,7 @@ package {{.Generator.Config.Package}}
 import (
 	"context"
 	"os"
+	"slices"
 	"testing"
 
 	// no-op import for go modules
diff --git a/generator/templates/actions/find.gotpl b/generator/templates/actions/find.gotpl
index a302078af..e84d0fdf6 100644
--- a/generator/templates/actions/find.gotpl
+++ b/generator/templates/actions/find.gotpl
@@ -107,6 +107,25 @@
 				return r
 			}
 
+			func (r {{ $result }}) Omit(params ...{{ $model.Name.GoLowerCase }}PrismaFields) {{ $result }} {
+				var outputs []builder.Output
+
+				var raw []string
+				for _, param := range params {
+					raw = append(raw, string(param))
+				}
+
+				for _, output := range {{ $model.Name.GoLowerCase }}Output {
+					if !slices.Contains(raw, output.Name) {
+						outputs = append(outputs, output)
+					}
+				}
+
+				r.query.Outputs = outputs
+
+				return r
+			}
+
 			{{ if $v.List }}
 				func (r {{ $result }}) OrderBy(params ...{{ $orderByParam }}) {{ $result }} {
 					var fields []builder.Field
diff --git a/test/features/omit/default_test.go b/test/features/omit/default_test.go
new file mode 100644
index 000000000..cf220ed31
--- /dev/null
+++ b/test/features/omit/default_test.go
@@ -0,0 +1,103 @@
+package db
+
+import (
+	"context"
+	"testing"
+
+	"github.com/steebchen/prisma-client-go/test"
+	"github.com/steebchen/prisma-client-go/test/helpers/massert"
+)
+
+type cx = context.Context
+type Func func(t *testing.T, client *PrismaClient, ctx cx)
+
+func str(v string) *string {
+	return &v
+}
+
+func i(v int) *int {
+	return &v
+}
+
+func TestOmit(t *testing.T) {
+	tests := []struct {
+		name   string
+		before []string
+		run    Func
+	}{{
+		name:   "omit",
+		before: nil,
+		run: func(t *testing.T, client *PrismaClient, ctx cx) {
+			user, err := client.User.CreateOne(
+				User.Name.Set("a"),
+				User.Keep.Set("keep"),
+				User.ID.Set("123"),
+				User.Password.Set("password"),
+				User.Age.Set(20),
+			).Exec(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			massert.Equal(t, &UserModel{
+				InnerUser: InnerUser{
+					ID:       "123",
+					Name:     "a",
+					Keep:     "keep",
+					Password: str("password"),
+					Age:      i(20),
+				},
+			}, user)
+
+			expected := &UserModel{
+				InnerUser: InnerUser{
+					Keep: "keep",
+				},
+			}
+
+			users, err := client.User.FindMany(
+				User.Name.Equals("a"),
+			).Omit(
+				User.ID.Field(),
+				User.Password.Field(),
+				User.Age.Field(),
+				User.Name.Field(),
+			).Exec(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			massert.Equal(t, []UserModel{*expected}, users)
+
+			oneUser, err := client.User.FindUnique(
+				User.ID.Equals("123"),
+			).Omit(
+				User.ID.Field(),
+				User.Password.Field(),
+				User.Age.Field(),
+				User.Name.Field(),
+			).Exec(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			massert.Equal(t, expected, oneUser)
+		},
+	}}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			test.RunSerial(t, test.Databases, func(t *testing.T, db test.Database, ctx context.Context) {
+				if db == test.MongoDB {
+					// TODO
+					t.Skip()
+				}
+				client := NewClient()
+				mockDBName := test.Start(t, db, client.Engine, tt.before)
+				defer test.End(t, db, client.Engine, mockDBName)
+				tt.run(t, client, context.Background())
+			})
+		})
+	}
+}
diff --git a/test/features/omit/schema.prisma b/test/features/omit/schema.prisma
new file mode 100644
index 000000000..a479336d0
--- /dev/null
+++ b/test/features/omit/schema.prisma
@@ -0,0 +1,19 @@
+datasource db {
+  provider = "sqlite"
+  url      = env("__REPLACE__")
+}
+
+generator db {
+  provider          = "go run github.com/steebchen/prisma-client-go"
+  output            = "."
+  disableGoBinaries = true
+  package           = "db"
+}
+
+model User {
+  id       String  @id @default(cuid())
+  name     String
+  keep     String
+  age      Int?
+  password String?
+}