From 26ed31e7b60ebd86fc00bd2e3e7e2228a4c338db Mon Sep 17 00:00:00 2001
From: Jay Taylor <outtatime@gmail.com>
Date: Fri, 7 Aug 2015 15:01:35 -0700
Subject: [PATCH] Preserve modelStructs so that multiple invocations of
 `db.SingularTable(true|false)' don't clobber the state.

Includes corresponding unit-test.
---
 main.go                    |  5 +++-
 state_preservation_test.go | 61 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 state_preservation_test.go

diff --git a/main.go b/main.go
index 30802205..89271373 100644
--- a/main.go
+++ b/main.go
@@ -130,7 +130,10 @@ func (s *DB) LogMode(enable bool) *DB {
 }
 
 func (s *DB) SingularTable(enable bool) {
-	modelStructs = map[reflect.Type]*ModelStruct{}
+	// Don't clobber/reset modelStructs if it is already initialized.
+	if modelStructs == nil {
+		modelStructs = map[reflect.Type]*ModelStruct{}
+	}
 	s.parent.singularTable = enable
 }
 
diff --git a/state_preservation_test.go b/state_preservation_test.go
new file mode 100644
index 00000000..bda0facb
--- /dev/null
+++ b/state_preservation_test.go
@@ -0,0 +1,61 @@
+package gorm_test
+
+import (
+	"testing"
+)
+
+type (
+	Organization struct {
+		Id   int64
+		Name string `sql:"type:varchar(255);not null;"`
+	}
+
+	App struct {
+		Id             int64
+		Organization   Organization
+		OrganizationId int64  `sql:"not null;"`
+		Name           string `sql:"type:varchar(255);not null;"`
+	}
+)
+
+func TestMultipleSingularTableInvocations(t *testing.T) {
+	DB.SingularTable(true) // Invocation #1.
+	// DB.LogMode(true)
+
+	entities := []interface{}{
+		&Organization{},
+		&App{},
+	}
+	for i := len(entities) - 1; i >= 0; i-- {
+		if err := DB.DropTableIfExists(entities[i]).Error; err != nil {
+			t.Fatalf("Drop table for entity=%+v: %s", entities[i], err)
+		}
+	}
+	if err := DB.AutoMigrate(entities...).Error; err != nil {
+		t.Fatalf("Auto-migrate failed for entity=%+v: %s", entity, err)
+	}
+	if err := DB.Model(&App{}).AddForeignKey("organization_id", "organization(id)", "RESTRICT", "RESTRICT"); err != nil {
+		t.Fatalf("Problem adding OrganizationId foreign-key to App table: %s", err)
+	}
+
+	createFixtures(t)
+}
+
+func createFixtures(t *testing.T) {
+	DB.SingularTable(true) // Invocation #2.  If this were to clobber internal gorm state it can break things.
+
+	org := &Organization{
+		Name: "Some Organization for Testing",
+	}
+	if err := DB.Save(org).Error; err != nil {
+		t.Fatal(err)
+	}
+
+	app := &App{
+		OrganizationId: org.Id,
+		Name:           "my-app-for-test-purposes",
+	}
+	if err := DB.Save(app).Error; err != nil {
+		t.Fatal(err)
+	}
+}