From 8876d31317c965a25b175ac347ceb521cef46a34 Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <7413593+a8m@users.noreply.github.com>
Date: Tue, 26 Jan 2021 22:31:32 +0200
Subject: [PATCH] entc/gen: fix feature-flag cleanups (#1210)

---
 entc/gen/feature.go    | 23 +++++++++++++++++++++--
 entc/gen/graph_test.go | 22 +++++++++++++++++++++-
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/entc/gen/feature.go b/entc/gen/feature.go
index 5339afbda..6626913da 100644
--- a/entc/gen/feature.go
+++ b/entc/gen/feature.go
@@ -5,6 +5,7 @@
 package gen
 
 import (
+	"io/ioutil"
 	"os"
 	"path/filepath"
 )
@@ -39,7 +40,7 @@ var (
 		Default:     false,
 		Description: "Schema snapshot stores a snapshot of ent/schema and auto-solve merge-conflict (issue #852)",
 		cleanup: func(c *Config) error {
-			return os.RemoveAll(filepath.Join(c.Target, "internal"))
+			return remove(filepath.Join(c.Target, "internal"), "schema.go")
 		},
 	}
 
@@ -58,7 +59,7 @@ var (
 			},
 		},
 		cleanup: func(c *Config) error {
-			return os.RemoveAll(filepath.Join(c.Target, "internal"))
+			return remove(filepath.Join(c.Target, "internal"), "schemaconfig.go")
 		},
 	}
 
@@ -115,3 +116,21 @@ type Feature struct {
 	// e.g. delete files from previous codegen runs.
 	cleanup func(*Config) error
 }
+
+// remove file (if exists) and its dir if it's empty.
+func remove(dir, file string) error {
+	if err := os.Remove(filepath.Join(dir, file)); err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return err
+	}
+	infos, err := ioutil.ReadDir(dir)
+	if err != nil {
+		return err
+	}
+	if len(infos) == 0 {
+		return os.Remove(dir)
+	}
+	return nil
+}
diff --git a/entc/gen/graph_test.go b/entc/gen/graph_test.go
index ce29d56a7..fffb3ed69 100644
--- a/entc/gen/graph_test.go
+++ b/entc/gen/graph_test.go
@@ -291,6 +291,7 @@ func TestGraph_Gen(t *testing.T) {
 		Storage:   drivers[0],
 		Templates: []*Template{external},
 		IDType:    &field.TypeInfo{Type: field.TypeInt},
+		Features:  AllFeatures,
 	}, &load.Schema{
 		Name: "T1",
 		Fields: []*load.Field{
@@ -315,8 +316,27 @@ func TestGraph_Gen(t *testing.T) {
 		_, err := os.Stat(fmt.Sprintf(fmt.Sprintf("%s/%s.go", target, format), "t1"))
 		require.NoError(err)
 	}
-	_, err = os.Stat(target + "/external.go")
+	_, err = os.Stat(filepath.Join(target, "external.go"))
 	require.NoError(err)
+
+	// Generated feature templates.
+	_, err = os.Stat(filepath.Join(target, "internal", "schema.go"))
+	require.NoError(err)
+	_, err = os.Stat(filepath.Join(target, "internal", "schemaconfig.go"))
+	require.NoError(err)
+	// Rerun codegen with only one feature-flag.
+	graph.Features = []Feature{FeatureSnapshot}
+	require.NoError(graph.Gen())
+	// Generated feature templates.
+	_, err = os.Stat(filepath.Join(target, "internal", "schema.go"))
+	require.NoError(err)
+	_, err = os.Stat(filepath.Join(target, "internal", "schemaconfig.go"))
+	require.True(os.IsNotExist(err))
+	// Rerun codegen without any feature-flags.
+	graph.Features = nil
+	require.NoError(graph.Gen())
+	_, err = os.Stat(filepath.Join(target, "internal"))
+	require.True(os.IsNotExist(err))
 }
 
 func ensureStructTag(name string) Hook {