From 5f40401e6de986abac35b4251105cb252cc5e268 Mon Sep 17 00:00:00 2001
From: Luca Steeb <contact@luca-steeb.com>
Date: Tue, 12 Dec 2023 01:31:23 +0700
Subject: [PATCH] fix(engine-binary): adapt build tag generation (#1125)

---
 binaries/bindata/bindata.go | 13 ++++++------
 binaries/platform/info.go   | 40 +++++++++++++++++++++++++++++++++++++
 generator/run.go            | 10 +++-------
 3 files changed, 50 insertions(+), 13 deletions(-)
 create mode 100644 binaries/platform/info.go

diff --git a/binaries/bindata/bindata.go b/binaries/bindata/bindata.go
index 082fac0c..23fa5d3c 100644
--- a/binaries/bindata/bindata.go
+++ b/binaries/bindata/bindata.go
@@ -7,11 +7,12 @@ import (
 	"os"
 
 	"github.com/steebchen/prisma-client-go/binaries"
+	"github.com/steebchen/prisma-client-go/binaries/platform"
 )
 
 // TODO go fmt files after creation
 
-func WriteFile(name, pkg, platform, from, to string) error {
+func WriteFile(name, pkg, from, to string, info platform.Info) error {
 	f, err := os.Create(to)
 	if err != nil {
 		return fmt.Errorf("generate open go file: %w", err)
@@ -20,7 +21,7 @@ func WriteFile(name, pkg, platform, from, to string) error {
 	//goland:noinspection GoUnhandledErrorResult
 	defer f.Close()
 
-	if err := writeHeader(f, pkg, name, platform); err != nil {
+	if err := writeHeader(f, pkg, name, info); err != nil {
 		return fmt.Errorf("write header: %w", err)
 	}
 
@@ -31,10 +32,10 @@ func WriteFile(name, pkg, platform, from, to string) error {
 	return nil
 }
 
-func writeHeader(w io.Writer, pkg, name, platform string) error {
+func writeHeader(w io.Writer, pkg string, name string, info platform.Info) error {
 	_, err := fmt.Fprintf(w, `// Code generated by Prisma Client Go. DO NOT EDIT.
-//go:build !codeanalysis && !prisma_ignore && %s
-// +build !codeanalysis,!prisma_ignore,%s
+//go:build !codeanalysis && !prisma_ignore && %s && %s
+// +build !codeanalysis,!prisma_ignore,%s,%s
 
 //nolint
 package %s
@@ -46,7 +47,7 @@ import (
 func init() {
 	unpack.Unpack(data, "%s", "%s")
 }
-`, platform, platform, pkg, name, binaries.EngineVersion)
+`, info.Platform, info.Arch, info.Platform, info.Arch, pkg, name, binaries.EngineVersion)
 	return err
 }
 
diff --git a/binaries/platform/info.go b/binaries/platform/info.go
new file mode 100644
index 00000000..d4896512
--- /dev/null
+++ b/binaries/platform/info.go
@@ -0,0 +1,40 @@
+package platform
+
+import "strings"
+
+type Info struct {
+	Platform string
+	Arch     string
+}
+
+func MapBinaryTarget(name string) Info {
+	return Info{
+		Platform: mapBinaryTargetToPlatform(name),
+		Arch:     mapBinaryTargetToArch(name),
+	}
+}
+
+func mapBinaryTargetToPlatform(name string) string {
+	switch {
+	case strings.Contains(name, "linux") ||
+		strings.Contains(name, "debian") ||
+		strings.Contains(name, "rhel") ||
+		strings.Contains(name, "musl"):
+		return "linux"
+	case strings.Contains(name, "darwin"):
+		return "darwin"
+	case strings.Contains(name, "windows"):
+		return "windows"
+	default:
+		return "linux"
+	}
+}
+
+func mapBinaryTargetToArch(name string) string {
+	switch {
+	case strings.Contains(name, "arm64"):
+		return "arm64"
+	default:
+		return "!arm64"
+	}
+}
diff --git a/generator/run.go b/generator/run.go
index 1991ee73..db3e730c 100644
--- a/generator/run.go
+++ b/generator/run.go
@@ -8,7 +8,6 @@ import (
 	"go/format"
 	"os"
 	"path"
-	"runtime"
 	"strings"
 	"text/template"
 
@@ -189,15 +188,12 @@ func generateBinaries(input *Root) error {
 
 func generateQueryEngineFiles(binaryTargets []string, pkg, outputDir string) error {
 	for _, name := range binaryTargets {
-		pt := runtime.GOOS
-		if strings.Contains(name, "debian") || strings.Contains(name, "rhel") || strings.Contains(name, "musl") {
-			pt = "linux"
-		}
-
 		if name == "native" {
 			name = platform.BinaryPlatformNameStatic()
 		}
 
+		info := platform.MapBinaryTarget(name)
+
 		name = TransformBinaryTarget(name)
 
 		enginePath := binaries.GetEnginePath(binaries.GlobalCacheDir(), "query-engine", name)
@@ -206,7 +202,7 @@ func generateQueryEngineFiles(binaryTargets []string, pkg, outputDir string) err
 		to := path.Join(outputDir, filename)
 
 		// TODO check if already exists, but make sure version matches
-		if err := bindata.WriteFile(name, pkg, pt, enginePath, to); err != nil {
+		if err := bindata.WriteFile(name, pkg, enginePath, to, info); err != nil {
 			return fmt.Errorf("generate write go file: %w", err)
 		}