From 4de9ab485f066c8bef769547213a6bf5fc02f40a Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <ariel@mashraki.co.il>
Date: Fri, 10 Nov 2023 10:44:59 +0200
Subject: [PATCH] dialect/entsql: add schema to package annotation

---
 dialect/entsql/annotation.go | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/dialect/entsql/annotation.go b/dialect/entsql/annotation.go
index 73a6b185f8..8ad1950679 100644
--- a/dialect/entsql/annotation.go
+++ b/dialect/entsql/annotation.go
@@ -9,6 +9,16 @@ import "entgo.io/ent/schema"
 // Annotation is a builtin schema annotation for attaching
 // SQL metadata to schema objects for both codegen and runtime.
 type Annotation struct {
+	// The Schema option allows setting the schema which the table belongs to.
+	// Note, this option is no-op for Ent default migration engine. However, schema
+	// extensions (like Atlas) can accept this option and implement it accordingly.
+	//
+	//	entsql.Annotation{
+	//		Schema: "public",
+	//	}
+	//
+	Schema string `json:"schema,omitempty"`
+
 	// The Table option allows overriding the default table
 	// name that is generated by ent. For example:
 	//
@@ -146,6 +156,21 @@ func (Annotation) Name() string {
 	return "EntSQL"
 }
 
+// The Schema option allows setting the schema which the table belongs to.
+// Note, this option is no-op for Ent default migration engine. However, schema
+// extensions (like Atlas) can accept this option and implement it accordingly.
+//
+//	func (T) Annotations() []schema.Annotation {
+//		return []schema.Annotation{
+//			entsql.Schema("public"),
+//		}
+//	}
+func Schema(s string) *Annotation {
+	return &Annotation{
+		Schema: s,
+	}
+}
+
 // The Table option allows overriding the default table
 // name that is generated by ent. For example:
 //
@@ -160,6 +185,14 @@ func Table(t string) *Annotation {
 	}
 }
 
+// SchemaTable allows setting both schema and table name in one annotation.
+func SchemaTable(s, t string) *Annotation {
+	return &Annotation{
+		Schema: s,
+		Table:  t,
+	}
+}
+
 // Check allows injecting custom "DDL" for setting an unnamed "CHECK" clause in "CREATE TABLE".
 //
 //	entsql.Annotation{
@@ -273,6 +306,9 @@ func (a Annotation) Merge(other schema.Annotation) schema.Annotation {
 	default:
 		return a
 	}
+	if s := ant.Schema; s != "" {
+		a.Schema = s
+	}
 	if t := ant.Table; t != "" {
 		a.Table = t
 	}