From f5ef25233dc3cdf97a77f7234016da1f5079af97 Mon Sep 17 00:00:00 2001
From: Georges Chaudy <chaudyg@gmail.com>
Date: Tue, 18 Nov 2025 12:12:05 +0100
Subject: [PATCH] Enhance SQL KV store and backend diagnostics

- Updated sqlKV struct to maintain a reference to the dbProvider, preventing garbage collection of the database connection.
- Implemented Ping and checkDB methods in sqlKV for verifying database connection health.
- Enhanced kvStorageBackend to implement the DiagnosticsServer interface, adding IsHealthy method for health checks.
- Updated resource server to utilize the new diagnostics capabilities of kvStorageBackend.

These changes improve the reliability and maintainability of the SQL-based KV store and its integration with the resource server.
---
 pkg/storage/unified/resource/sqlkv.go         | 32 ++++++++++++++++---
 .../unified/resource/storage_backend.go       | 20 ++++++++++++
 pkg/storage/unified/sql/server.go             |  4 +--
 3 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/pkg/storage/unified/resource/sqlkv.go b/pkg/storage/unified/resource/sqlkv.go
index 31eec8373f2..a795cf98b5a 100644
--- a/pkg/storage/unified/resource/sqlkv.go
+++ b/pkg/storage/unified/resource/sqlkv.go
@@ -24,8 +24,9 @@ const (
 
 // sqlKV implements the KV interface using SQL storage
 type sqlKV struct {
-	db      db.DB
-	dialect sqltemplate.Dialect
+	dbProvider db.DBProvider // Keep reference to prevent GC
+	db         db.DB
+	dialect    sqltemplate.Dialect
 }
 
 // NewSQLKV creates a new SQL-based KV store
@@ -55,8 +56,9 @@ func NewSQLKV(dbProvider db.DBProvider) (KV, error) {
 	}
 
 	return &sqlKV{
-		db:      dbConn,
-		dialect: dialect,
+		dbProvider: dbProvider, // Keep reference to prevent GC from closing the database
+		db:         dbConn,
+		dialect:    dialect,
 	}, nil
 }
 
@@ -774,3 +776,25 @@ func (k *sqlKV) BatchDelete(ctx context.Context, section string, keys []string)
 func (k *sqlKV) UnixTimestamp(ctx context.Context) (int64, error) {
 	return time.Now().Unix(), nil
 }
+
+// Ping checks if the database connection is alive
+func (k *sqlKV) Ping(ctx context.Context) error {
+	if k.db == nil {
+		return fmt.Errorf("database connection is nil")
+	}
+	return k.db.PingContext(ctx)
+}
+
+// checkDB verifies the database connection is still valid before operations
+func (k *sqlKV) checkDB() error {
+	if k.db == nil {
+		return fmt.Errorf("database connection is nil")
+	}
+	// Quick ping to verify connection is alive
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if err := k.db.PingContext(ctx); err != nil {
+		return fmt.Errorf("database connection is not healthy: %w", err)
+	}
+	return nil
+}
diff --git a/pkg/storage/unified/resource/storage_backend.go b/pkg/storage/unified/resource/storage_backend.go
index 6110c6aacd4..a0fb03ec9bd 100644
--- a/pkg/storage/unified/resource/storage_backend.go
+++ b/pkg/storage/unified/resource/storage_backend.go
@@ -74,6 +74,7 @@ type kvStorageBackend struct {
 
 var _ StorageBackend = &kvStorageBackend{}
 var _ LifecycleHooks = &kvStorageBackend{}
+var _ resourcepb.DiagnosticsServer = &kvStorageBackend{}
 
 type KVBackendOptions struct {
 	KvStore                      KV
@@ -151,6 +152,25 @@ func (k *kvStorageBackend) Stop(ctx context.Context) error {
 	return nil
 }
 
+// IsHealthy implements DiagnosticsServer
+func (k *kvStorageBackend) IsHealthy(ctx context.Context, _ *resourcepb.HealthCheckRequest) (*resourcepb.HealthCheckResponse, error) {
+	// Check if the underlying KV store supports Ping (e.g., sqlKV)
+	type pinger interface {
+		Ping(context.Context) error
+	}
+	if p, ok := k.kv.(pinger); ok {
+		if err := p.Ping(ctx); err != nil {
+			return nil, fmt.Errorf("KV store health check failed: %w", err)
+		}
+	}
+	return &resourcepb.HealthCheckResponse{Status: resourcepb.HealthCheckResponse_SERVING}, nil
+}
+
+// Read implements DiagnosticsServer
+func (k *kvStorageBackend) Read(ctx context.Context, req *resourcepb.ReadRequest) (*resourcepb.ReadResponse, error) {
+	return nil, ErrNotImplementedYet
+}
+
 // runCleanupOldEvents starts a background goroutine that periodically cleans up old events
 func (k *kvStorageBackend) runCleanupOldEvents(ctx context.Context) {
 	// Run cleanup every hour
diff --git a/pkg/storage/unified/sql/server.go b/pkg/storage/unified/sql/server.go
index d6410f14a6b..910a2bf0f48 100644
--- a/pkg/storage/unified/sql/server.go
+++ b/pkg/storage/unified/sql/server.go
@@ -19,6 +19,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/storage/unified/resource"
+	"github.com/grafana/grafana/pkg/storage/unified/resourcepb"
 	"github.com/grafana/grafana/pkg/storage/unified/sql/db/dbimpl"
 )
 
@@ -122,8 +123,7 @@ func NewResourceServer(opts ServerOptions) (resource.ResourceServer, error) {
 			}
 			serverOptions.Backend = kvBackend
 			serverOptions.Lifecycle = kvBackend.(resource.LifecycleHooks)
-			// Note: kvStorageBackend doesn't implement Diagnostics yet
-			// The server will use noopService for diagnostics
+			serverOptions.Diagnostics = kvBackend.(resourcepb.DiagnosticsServer)
 		} else {
 			// Use existing SQL backend
 			backend, err := NewBackend(BackendOptions{