Add MetricsReporter interface so that databaseBackend's can share their connection counts with CensusManager

builtin/logical/database/backend.go

@@ -174,8 +174,10 @@ func (b *databaseBackend) collectPluginInstanceGaugeValues(context.Context) ([]m
 type databaseBackend struct {
 	// connections holds configured database connections by config name
 	createConnectionLock sync.Mutex
-	connections          *syncmap.SyncMap[string, *dbPluginInstance]
-	logger               log.Logger
+
+	// Connections are loaded lazily, when a connection to a specific database plugin is needed.
+	connections *syncmap.SyncMap[string, *dbPluginInstance]
+	logger      log.Logger
 
 	*framework.Backend
 	// credRotationQueue is an in-memory priority queue used to track Static Roles
@@ -505,6 +507,26 @@ func (b *databaseBackend) getDatabaseConfigNameFromRotationID(path string) (stri
 	return res[1], nil
 }
 
+// GetConnectionMetrics returns a count of the active Database connections.
+// The returned count depends on the database connections map which is not
+// gauranteed to be an exhaustive list of all configured connections.
+func (b *databaseBackend) GetConnectionMetrics() (map[string]int, error) {
+	// Access the private b.connections field here
+	counts := make(map[string]int)
+	connectionsCopy := b.connections.Values()
+
+	for _, v := range connectionsCopy {
+		dbType, err := v.database.Type()
+		if err != nil {
+			continue
+		}
+
+		counts[dbType]++
+	}
+
+	return counts, nil
+}
+
 const backendHelp = `
 The database backend supports using many different databases
 as secret backends, including but not limited to:

builtin/logical/database/backend_test.go

@@ -23,10 +23,12 @@ import (
 	"github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/vault/helper/builtinplugins"
 	"github.com/hashicorp/vault/helper/namespace"
+	"github.com/hashicorp/vault/helper/pluginconsts"
 	"github.com/hashicorp/vault/helper/testhelpers/certhelpers"
 	"github.com/hashicorp/vault/helper/testhelpers/corehelpers"
 	postgreshelper "github.com/hashicorp/vault/helper/testhelpers/postgresql"
 	vaulthttp "github.com/hashicorp/vault/http"
+	"github.com/hashicorp/vault/plugins/database/cassandra"
 	"github.com/hashicorp/vault/plugins/database/postgresql"
 	v4 "github.com/hashicorp/vault/sdk/database/dbplugin"
 	v5 "github.com/hashicorp/vault/sdk/database/dbplugin/v5"
@@ -39,6 +41,7 @@ import (
 	_ "github.com/jackc/pgx/v4"
 	"github.com/mitchellh/mapstructure"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func getClusterPostgresDBWithFactory(t *testing.T, factory logical.Factory) (*vault.TestCluster, logical.SystemView) {
@@ -92,6 +95,15 @@ func TestBackend_PluginMain_PostgresMultiplexed(t *testing.T) {
 	v5.ServeMultiplex(postgresql.New)
 }
 
+// TestBackend_PluginMain_CassandraMultiplexed tests the Cassandra database plugin
+func TestBackend_PluginMain_CassandraMultiplexed(t *testing.T) {
+	if os.Getenv(pluginutil.PluginVaultVersionEnv) == "" {
+		return
+	}
+
+	v5.ServeMultiplex(cassandra.New)
+}
+
 func TestBackend_RoleUpgrade(t *testing.T) {
 	storage := &logical.InmemStorage{}
 	backend := &databaseBackend{}
@@ -1593,6 +1605,133 @@ func TestBackend_ConnectionURL_redacted(t *testing.T) {
 	}
 }
 
+// TestBackend_GetConnectionMetrics tests the GetConnectionMetrics method
+// of the database backend to ensure it correctly counts the number of connections
+// for each plugin type.
+func TestBackend_GetConnectionMetrics(t *testing.T) {
+	cluster, sys := getCluster(t)
+	defer cluster.Cleanup()
+
+	vault.TestAddTestPlugin(t, cluster.Cores[0].Core, "cassandra-database-plugin", consts.PluginTypeDatabase, "", "TestBackend_PluginMain_CassandraMultiplexed",
+		[]string{fmt.Sprintf("%s=%s", pluginutil.PluginCACertPEMEnv, cluster.CACertPEMFile)})
+	vault.TestAddTestPlugin(t, cluster.Cores[0].Core, "postgresql-database-plugin", consts.PluginTypeDatabase, "", "TestBackend_PluginMain_PostgresMultiplexed",
+		[]string{fmt.Sprintf("%s=%s", pluginutil.PluginCACertPEMEnv, cluster.CACertPEMFile)})
+
+	postgresConfig := map[string]interface{}{
+		"name":              "postgres-plugin-test",
+		"plugin_name":       "postgresql-database-plugin",
+		"connection_url":    "some_postgres_url",
+		"username":          "postgres",
+		"password":          "secret",
+		"verify_connection": false,
+	}
+
+	cassandraConfig := map[string]interface{}{
+		"name":              "cassandra-plugin-test",
+		"plugin_name":       "cassandra-database-plugin",
+		"hosts":             "some_cassandra_url",
+		"username":          "cassandra",
+		"password":          "secret",
+		"verify_connection": false,
+	}
+
+	postgresConfigV2 := map[string]interface{}{
+		"name":              "postgres-plugin-test-second-connection",
+		"plugin_name":       "postgresql-database-plugin",
+		"connection_url":    "some_postgres_url",
+		"username":          "postgres",
+		"password":          "secret",
+		"verify_connection": false,
+	}
+
+	cassandraConfigV2 := map[string]interface{}{
+		"name":              "cassandra-plugin-test-second-connection",
+		"plugin_name":       "cassandra-database-plugin",
+		"hosts":             "some_cassandra_url",
+		"username":          "cassandra",
+		"password":          "secret",
+		"verify_connection": false,
+	}
+
+	tests := []struct {
+		name              string
+		connectionConfigs []map[string]interface{}
+		expectedCounts    map[string]int
+	}{
+		{
+			name: "1 Postgres, 1 Cassandra connection",
+			connectionConfigs: []map[string]interface{}{
+				postgresConfig,
+				cassandraConfig,
+			},
+			expectedCounts: map[string]int{
+				pluginconsts.DbPostgresqlPlugin: 1,
+				pluginconsts.DbCassandraPlugin:  1,
+			},
+		},
+		{
+			name: "1 Postgres connection",
+			connectionConfigs: []map[string]interface{}{
+				postgresConfig,
+			},
+			expectedCounts: map[string]int{
+				pluginconsts.DbPostgresqlPlugin: 1,
+			},
+		},
+		{
+			name:              "No Connections",
+			connectionConfigs: []map[string]interface{}{},
+			expectedCounts:    map[string]int{},
+		},
+		{
+			name: " 2 Postgres, 2 Cassandra connections",
+			connectionConfigs: []map[string]interface{}{
+				postgresConfig,
+				postgresConfigV2,
+				cassandraConfig,
+				cassandraConfigV2,
+			},
+			expectedCounts: map[string]int{
+				pluginconsts.DbPostgresqlPlugin: 2,
+				pluginconsts.DbCassandraPlugin:  2,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config := logical.TestBackendConfig()
+			config.StorageView = &logical.InmemStorage{}
+			config.System = sys
+
+			b, err := Factory(context.Background(), config)
+			require.NoError(t, err)
+			defer b.Cleanup(context.Background())
+
+			// Create connections
+			for _, connData := range tt.connectionConfigs {
+				req := &logical.Request{
+					Operation: logical.UpdateOperation,
+					Path:      fmt.Sprintf("config/%s", connData["name"]),
+					Storage:   config.StorageView,
+					Data:      connData,
+				}
+				resp, err := b.HandleRequest(namespace.RootContext(nil), req)
+				if err != nil || (resp != nil && resp.IsError()) {
+					t.Fatalf("Failed to configure connection %s: err:%s resp:%#v\n", connData["name"], err, resp)
+				}
+			}
+
+			// Get metric counts
+			metricsReporter, ok := b.(logical.MetricsReporter)
+			require.True(t, ok)
+			connectionCount, err := metricsReporter.GetConnectionMetrics()
+			require.NoError(t, err)
+			require.Equal(t, tt.expectedCounts, connectionCount, "Unexpected connection counts")
+		})
+	}
+}
+
 type hangingPlugin struct{}
 
 func (h hangingPlugin) Initialize(_ context.Context, req v5.InitializeRequest) (v5.InitializeResponse, error) {

helper/pluginconsts/plugin_consts.go

@@ -26,6 +26,14 @@ const (
 	AuthTypeSAML             = "saml"
 	AuthTypeApprole          = "approle"
 	AuthTypeJWT              = "jwt"
+	DbCassandraPlugin        = "cassandra"
+	DbHanaPlugin             = "hdb"
+	DbInfluxDBPlugin         = "influxdb"
+	DbMongoDBPlugin          = "mongodb"
+	DbMsSQLPlugin            = "mssql"
+	DbMySQLPlugin            = "mysql"
+	DbRedshiftPlugin         = "redshift"
+	DbPostgresqlPlugin       = "pgx"
 	SecretEngineAD           = "ad"
 	SecretEngineAlicloud     = "alicloud"
 	SecretEngineAWS          = "aws"

sdk/logical/logical.go

@@ -203,4 +203,10 @@ type PluginVersioner interface {
 	PluginVersion() PluginVersion
 }
 
+// MetricsReporter is an optional interface that returns a
+// metric. Currently only implemented by the database backend.
+type MetricsReporter interface {
+	GetConnectionMetrics() (map[string]int, error)
+}
+
 var EmptyPluginVersion = PluginVersion{""}