From 05d5ec480497a0b0952eba2de5305e611cfb6d73 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 bff97c7189c9..7bd11194f687 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 d1791f34a3cd..14c804b2865b 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 1c69cd6fa93e..6b5ebfbf1808 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 98ef58ed301b..8fa28e9a2afd 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 aa5134fb1f73..73c68f8afaf7 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 9b4bd3a96eab..13271021db11 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 5c935007309e..058433f11d64 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 e367bcc09f55..c227010bf2cd 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 554f3830beca..db7d731e3e10 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 3009dc577d16..28bead690111 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 9cb0bb93d6a9..052e4aca882b 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 0d9a47838fd0..8fff28c05e21 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 36dce655eb40..dd6815ae67a0 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 26a7c76a2ab6..c38ddfa8d8d0 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 0655758a60aa..09ed0202dc30 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 459de23fc4cc..f9720b3e0f50 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 {