From b747b27afdaa08edb251ee6d4f093015f386cfe1 Mon Sep 17 00:00:00 2001
From: Daniel Nephin <dnephin@hashicorp.com>
Date: Fri, 29 Jan 2021 19:24:41 -0500
Subject: [PATCH] state: remove the need for registerSchema

registerSchema creates some indirection which is not necessary in this
case. newDBSchema can call each of the tables.

Enterprise tables can be added from the existing withEnterpriseSchema
shim.
---
 agent/consul/state/acl.go              |  8 ----
 agent/consul/state/autopilot.go        |  7 +--
 agent/consul/state/catalog_schema.go   |  8 ----
 agent/consul/state/config_entry.go     |  4 --
 agent/consul/state/connect_ca.go       |  6 ---
 agent/consul/state/coordinate.go       |  7 +--
 agent/consul/state/federation_state.go |  4 --
 agent/consul/state/intention.go        |  4 --
 agent/consul/state/kvs.go              |  8 +---
 agent/consul/state/prepared_query.go   |  7 +--
 agent/consul/state/schema.go           | 59 +++++++++++++++-----------
 agent/consul/state/schema_test.go      |  2 +-
 agent/consul/state/session.go          |  8 +---
 agent/consul/state/state_store.go      |  2 +-
 agent/consul/state/system_metadata.go  |  3 --
 agent/consul/state/usage.go            |  4 --
 16 files changed, 46 insertions(+), 95 deletions(-)

diff --git a/agent/consul/state/acl.go b/agent/consul/state/acl.go
index bff97c718..7bd11194f 100644
--- a/agent/consul/state/acl.go
+++ b/agent/consul/state/acl.go
@@ -212,14 +212,6 @@ func (s *TokenExpirationIndex) FromArgs(args ...interface{}) ([]byte, error) {
 	return buf, nil
 }
 
-func init() {
-	registerSchema(tokensTableSchema)
-	registerSchema(policiesTableSchema)
-	registerSchema(rolesTableSchema)
-	registerSchema(bindingRulesTableSchema)
-	registerSchema(authMethodsTableSchema)
-}
-
 // ACLTokens is used when saving a snapshot
 func (s *Snapshot) ACLTokens() (memdb.ResultIterator, error) {
 	iter, err := s.tx.Get("acl-tokens", "id")
diff --git a/agent/consul/state/autopilot.go b/agent/consul/state/autopilot.go
index d1791f34a..14c804b28 100644
--- a/agent/consul/state/autopilot.go
+++ b/agent/consul/state/autopilot.go
@@ -3,8 +3,9 @@ package state
 import (
 	"fmt"
 
-	"github.com/hashicorp/consul/agent/structs"
 	"github.com/hashicorp/go-memdb"
+
+	"github.com/hashicorp/consul/agent/structs"
 )
 
 // autopilotConfigTableSchema returns a new table schema used for storing
@@ -25,10 +26,6 @@ func autopilotConfigTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(autopilotConfigTableSchema)
-}
-
 // Autopilot is used to pull the autopilot config from the snapshot.
 func (s *Snapshot) Autopilot() (*structs.AutopilotConfig, error) {
 	c, err := s.tx.First("autopilot-config", "id")
diff --git a/agent/consul/state/catalog_schema.go b/agent/consul/state/catalog_schema.go
index 1c69cd6fa..6b5ebfbf1 100644
--- a/agent/consul/state/catalog_schema.go
+++ b/agent/consul/state/catalog_schema.go
@@ -347,11 +347,3 @@ func (index *ServiceNameIndex) PrefixFromArgs(args ...interface{}) ([]byte, erro
 	}
 	return val, nil
 }
-
-func init() {
-	registerSchema(nodesTableSchema)
-	registerSchema(servicesTableSchema)
-	registerSchema(checksTableSchema)
-	registerSchema(gatewayServicesTableSchema)
-	registerSchema(meshTopologyTableSchema)
-}
diff --git a/agent/consul/state/config_entry.go b/agent/consul/state/config_entry.go
index 98ef58ed3..8fa28e9a2 100644
--- a/agent/consul/state/config_entry.go
+++ b/agent/consul/state/config_entry.go
@@ -74,10 +74,6 @@ func (s *ConfigEntryLinkIndex) PrefixFromArgs(args ...interface{}) ([]byte, erro
 	return val, nil
 }
 
-func init() {
-	registerSchema(configTableSchema)
-}
-
 // ConfigEntries is used to pull all the config entries for the snapshot.
 func (s *Snapshot) ConfigEntries() ([]structs.ConfigEntry, error) {
 	entries, err := s.tx.Get(tableConfigEntries, "id")
diff --git a/agent/consul/state/connect_ca.go b/agent/consul/state/connect_ca.go
index aa5134fb1..73c68f8af 100644
--- a/agent/consul/state/connect_ca.go
+++ b/agent/consul/state/connect_ca.go
@@ -73,12 +73,6 @@ func caRootTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(caBuiltinProviderTableSchema)
-	registerSchema(caConfigTableSchema)
-	registerSchema(caRootTableSchema)
-}
-
 // CAConfig is used to pull the CA config from the snapshot.
 func (s *Snapshot) CAConfig() (*structs.CAConfiguration, error) {
 	c, err := s.tx.First(tableConnectCAConfig, "id")
diff --git a/agent/consul/state/coordinate.go b/agent/consul/state/coordinate.go
index 9b4bd3a96..13271021d 100644
--- a/agent/consul/state/coordinate.go
+++ b/agent/consul/state/coordinate.go
@@ -3,9 +3,10 @@ package state
 import (
 	"fmt"
 
+	"github.com/hashicorp/go-memdb"
+
 	"github.com/hashicorp/consul/agent/structs"
 	"github.com/hashicorp/consul/lib"
-	"github.com/hashicorp/go-memdb"
 )
 
 // coordinatesTableSchema returns a new table schema used for storing
@@ -47,10 +48,6 @@ func coordinatesTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(coordinatesTableSchema)
-}
-
 // Coordinates is used to pull all the coordinates from the snapshot.
 func (s *Snapshot) Coordinates() (memdb.ResultIterator, error) {
 	iter, err := s.tx.Get("coordinates", "id")
diff --git a/agent/consul/state/federation_state.go b/agent/consul/state/federation_state.go
index 5c9350073..058433f11 100644
--- a/agent/consul/state/federation_state.go
+++ b/agent/consul/state/federation_state.go
@@ -27,10 +27,6 @@ func federationStateTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(federationStateTableSchema)
-}
-
 // FederationStates is used to pull all the federation states for the snapshot.
 func (s *Snapshot) FederationStates() ([]*structs.FederationState, error) {
 	configs, err := s.tx.Get(tableFederationStates, "id")
diff --git a/agent/consul/state/intention.go b/agent/consul/state/intention.go
index e367bcc09..c227010bf 100644
--- a/agent/consul/state/intention.go
+++ b/agent/consul/state/intention.go
@@ -95,10 +95,6 @@ func intentionsTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(intentionsTableSchema)
-}
-
 // LegacyIntentions is used to pull all the intentions from the snapshot.
 //
 // Deprecated: service-intentions config entries are handled as config entries
diff --git a/agent/consul/state/kvs.go b/agent/consul/state/kvs.go
index 554f3830b..db7d731e3 100644
--- a/agent/consul/state/kvs.go
+++ b/agent/consul/state/kvs.go
@@ -4,8 +4,9 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/hashicorp/consul/agent/structs"
 	"github.com/hashicorp/go-memdb"
+
+	"github.com/hashicorp/consul/agent/structs"
 )
 
 // kvsTableSchema returns a new table schema used for storing key/value data for
@@ -48,11 +49,6 @@ func tombstonesTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(kvsTableSchema)
-	registerSchema(tombstonesTableSchema)
-}
-
 // KVs is used to pull the full list of KVS entries for use during snapshots.
 func (s *Snapshot) KVs() (memdb.ResultIterator, error) {
 	iter, err := s.tx.Get("kvs", "id_prefix")
diff --git a/agent/consul/state/prepared_query.go b/agent/consul/state/prepared_query.go
index 3009dc577..28bead690 100644
--- a/agent/consul/state/prepared_query.go
+++ b/agent/consul/state/prepared_query.go
@@ -4,9 +4,10 @@ import (
 	"fmt"
 	"regexp"
 
+	"github.com/hashicorp/go-memdb"
+
 	"github.com/hashicorp/consul/agent/consul/prepared_query"
 	"github.com/hashicorp/consul/agent/structs"
-	"github.com/hashicorp/go-memdb"
 )
 
 // preparedQueriesTableSchema returns a new table schema used for storing
@@ -50,10 +51,6 @@ func preparedQueriesTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(preparedQueriesTableSchema)
-}
-
 // validUUID is used to check if a given string looks like a UUID
 var validUUID = regexp.MustCompile(`(?i)^[\da-f]{8}-[\da-f]{4}-[\da-f]{4}-[\da-f]{4}-[\da-f]{12}$`)
 
diff --git a/agent/consul/state/schema.go b/agent/consul/state/schema.go
index 9cb0bb93d..052e4aca8 100644
--- a/agent/consul/state/schema.go
+++ b/agent/consul/state/schema.go
@@ -6,28 +6,43 @@ import (
 	"github.com/hashicorp/go-memdb"
 )
 
-// schemaFn is an interface function used to create and return
-// new memdb schema structs for constructing an in-memory db.
-type schemaFn func() *memdb.TableSchema
+// newDBSchema creates and returns the memdb schema for the Store.
+func newDBSchema() *memdb.DBSchema {
+	db := &memdb.DBSchema{Tables: make(map[string]*memdb.TableSchema)}
 
-// schemas is used to register schemas with the state store.
-var schemas []schemaFn
-
-// registerSchema registers a new schema with the state store. This should
-// get called at package init() time.
-func registerSchema(fn schemaFn) {
-	schemas = append(schemas, fn)
+	addTableSchemas(db,
+		authMethodsTableSchema,
+		autopilotConfigTableSchema,
+		bindingRulesTableSchema,
+		caBuiltinProviderTableSchema,
+		caConfigTableSchema,
+		caRootTableSchema,
+		checksTableSchema,
+		configTableSchema,
+		coordinatesTableSchema,
+		federationStateTableSchema,
+		gatewayServicesTableSchema,
+		indexTableSchema,
+		intentionsTableSchema,
+		kvsTableSchema,
+		meshTopologyTableSchema,
+		nodesTableSchema,
+		policiesTableSchema,
+		preparedQueriesTableSchema,
+		rolesTableSchema,
+		servicesTableSchema,
+		sessionChecksTableSchema,
+		sessionsTableSchema,
+		systemMetadataTableSchema,
+		tokensTableSchema,
+		tombstonesTableSchema,
+		usageTableSchema,
+	)
+	withEnterpriseSchema(db)
+	return db
 }
 
-// stateStoreSchema is used to return the combined schema for
-// the state store.
-func stateStoreSchema() *memdb.DBSchema {
-	// Create the root DB schema
-	db := &memdb.DBSchema{
-		Tables: make(map[string]*memdb.TableSchema),
-	}
-
-	// Add the tables to the root schema
+func addTableSchemas(db *memdb.DBSchema, schemas ...func() *memdb.TableSchema) {
 	for _, fn := range schemas {
 		schema := fn()
 		if _, ok := db.Tables[schema.Name]; ok {
@@ -35,8 +50,6 @@ func stateStoreSchema() *memdb.DBSchema {
 		}
 		db.Tables[schema.Name] = schema
 	}
-	withEnterpriseSchema(db)
-	return db
 }
 
 // indexTableSchema returns a new table schema used for tracking various indexes
@@ -57,7 +70,3 @@ func indexTableSchema() *memdb.TableSchema {
 		},
 	}
 }
-
-func init() {
-	registerSchema(indexTableSchema)
-}
diff --git a/agent/consul/state/schema_test.go b/agent/consul/state/schema_test.go
index 0d9a47838..8fff28c05 100644
--- a/agent/consul/state/schema_test.go
+++ b/agent/consul/state/schema_test.go
@@ -15,7 +15,7 @@ import (
 )
 
 func TestStateStoreSchema(t *testing.T) {
-	schema := stateStoreSchema()
+	schema := newDBSchema()
 	require.NoError(t, schema.Validate())
 
 	_, err := memdb.NewMemDB(schema)
diff --git a/agent/consul/state/session.go b/agent/consul/state/session.go
index 36dce655e..dd6815ae6 100644
--- a/agent/consul/state/session.go
+++ b/agent/consul/state/session.go
@@ -6,8 +6,9 @@ import (
 	"strings"
 	"time"
 
-	"github.com/hashicorp/consul/agent/structs"
 	"github.com/hashicorp/go-memdb"
+
+	"github.com/hashicorp/consul/agent/structs"
 )
 
 // sessionsTableSchema returns a new table schema used for storing session
@@ -129,11 +130,6 @@ func (index *CheckIDIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
 	return val, nil
 }
 
-func init() {
-	registerSchema(sessionsTableSchema)
-	registerSchema(sessionChecksTableSchema)
-}
-
 // Sessions is used to pull the full list of sessions for use during snapshots.
 func (s *Snapshot) Sessions() (memdb.ResultIterator, error) {
 	iter, err := s.tx.Get("sessions", "id")
diff --git a/agent/consul/state/state_store.go b/agent/consul/state/state_store.go
index 26a7c76a2..c38ddfa8d 100644
--- a/agent/consul/state/state_store.go
+++ b/agent/consul/state/state_store.go
@@ -155,7 +155,7 @@ type sessionCheck struct {
 // NewStateStore creates a new in-memory state storage layer.
 func NewStateStore(gc *TombstoneGC) *Store {
 	// Create the in-memory DB.
-	schema := stateStoreSchema()
+	schema := newDBSchema()
 	db, err := memdb.NewMemDB(schema)
 	if err != nil {
 		// the only way for NewMemDB to error is if the schema is invalid. The
diff --git a/agent/consul/state/system_metadata.go b/agent/consul/state/system_metadata.go
index 0655758a6..09ed0202d 100644
--- a/agent/consul/state/system_metadata.go
+++ b/agent/consul/state/system_metadata.go
@@ -26,9 +26,6 @@ func systemMetadataTableSchema() *memdb.TableSchema {
 		},
 	}
 }
-func init() {
-	registerSchema(systemMetadataTableSchema)
-}
 
 // SystemMetadataEntries used to pull all the system metadata entries for the snapshot.
 func (s *Snapshot) SystemMetadataEntries() ([]*structs.SystemMetadataEntry, error) {
diff --git a/agent/consul/state/usage.go b/agent/consul/state/usage.go
index 459de23fc..f9720b3e0 100644
--- a/agent/consul/state/usage.go
+++ b/agent/consul/state/usage.go
@@ -31,10 +31,6 @@ func usageTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(usageTableSchema)
-}
-
 // UsageEntry represents a count of some arbitrary identifier within the
 // state store, along with the last seen index.
 type UsageEntry struct {