Unified Storage: Fix Entity Server unconditionally starting (#88857)

fix Entity Server unconditionally starting

pkg/services/store/entity/sqlstash/create.go

@@ -17,6 +17,10 @@ func (s *sqlEntityServer) Create(ctx context.Context, r *entity.CreateEntityRequ
 	ctx, span := s.tracer.Start(ctx, "storage_server.Create")
 	defer span.End()
 
+	if err := s.Init(); err != nil {
+		return nil, err
+	}
+
 	key, err := entity.ParseKey(r.Entity.Key)
 	if err != nil {
 		return nil, fmt.Errorf("create entity: parse entity key: %w", err)

pkg/services/store/entity/sqlstash/delete.go

@@ -16,6 +16,10 @@ func (s *sqlEntityServer) Delete(ctx context.Context, r *entity.DeleteEntityRequ
 	ctx, span := s.tracer.Start(ctx, "storage_server.Delete")
 	defer span.End()
 
+	if err := s.Init(); err != nil {
+		return nil, err
+	}
+
 	key, err := entity.ParseKey(r.Key)
 	if err != nil {
 		return nil, fmt.Errorf("delete entity: parse entity key: %w", err)

pkg/services/store/entity/sqlstash/sql_storage_server.go

@@ -57,10 +57,6 @@ func ProvideSQLEntityServer(db db.EntityDBInterface, tracer tracing.Tracer /*, c
 		entityServer.log.Warn("error registering storage server metrics", "error", err)
 	}
 
-	if err := entityServer.Init(); err != nil {
-		return nil, fmt.Errorf("initialize Entity Server: %w", err)
-	}
-
 	return entityServer, nil
 }
 
@@ -80,7 +76,7 @@ type sqlEntityServer struct {
 	ctx         context.Context // TODO: remove
 	cancel      context.CancelFunc
 	stream      chan *entity.EntityWatchResponse
-	tracer      tracing.Tracer
+	tracer      trace.Tracer
 
 	once    sync.Once
 	initErr error
@@ -94,6 +90,10 @@ func (s *sqlEntityServer) Init() error {
 		s.initErr = s.init()
 	})
 
+	if s.initErr != nil {
+		return fmt.Errorf("initialize Entity Server: %w", s.initErr)
+	}
+
 	return s.initErr
 }
 
@@ -155,6 +155,12 @@ func (s *sqlEntityServer) init() error {
 }
 
 func (s *sqlEntityServer) IsHealthy(ctx context.Context, r *entity.HealthCheckRequest) (*entity.HealthCheckResponse, error) {
+	ctxLogger := s.log.FromContext(log.WithContextualAttributes(ctx, []any{"method", "isHealthy"}))
+	if err := s.Init(); err != nil {
+		ctxLogger.Error("init error", "error", err)
+		return nil, err
+	}
+
 	if err := s.sqlDB.PingContext(ctx); err != nil {
 		return nil, err
 	}

pkg/services/store/entity/sqlstash/sql_storage_server_test.go

@@ -3,21 +3,37 @@ package sqlstash
 import (
 	"testing"
 
+	sqlmock "github.com/DATA-DOG/go-sqlmock"
 	"github.com/stretchr/testify/require"
+	traceNoop "go.opentelemetry.io/otel/trace/noop"
 
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/sqlstore/session"
 	"github.com/grafana/grafana/pkg/services/store/entity"
+	"github.com/grafana/grafana/pkg/services/store/entity/sqlstash/sqltemplate"
 	"github.com/grafana/grafana/pkg/util/testutil"
 )
 
+func newTestSQLEntityServer(t *testing.T) (*sqlEntityServer, sqlmock.Sqlmock) {
+	db, mock := newMockDBMatchWords(t)
+
+	return &sqlEntityServer{
+		log:    log.NewNopLogger(),
+		tracer: traceNoop.NewTracerProvider().Tracer("test-tracer"),
+
+		sess: new(session.SessionDB), // FIXME
+
+		sqlDB:      db,
+		sqlDialect: sqltemplate.MySQL,
+	}, mock
+}
+
 func TestIsHealthy(t *testing.T) {
 	t.Parallel()
 
 	// test declarations
 	ctx := testutil.NewDefaultTestContext(t)
-	db, mock := newMockDBNopSQL(t)
-	s := &sqlEntityServer{
-		sqlDB: db,
-	}
+	s, mock := newTestSQLEntityServer(t)
 
 	// setup expectations
 	mock.ExpectPing()

pkg/services/store/entity/sqlstash/update.go

@@ -17,6 +17,10 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 	ctx, span := s.tracer.Start(ctx, "storage_server.Update")
 	defer span.End()
 
+	if err := s.Init(); err != nil {
+		return nil, err
+	}
+
 	key, err := entity.ParseKey(r.Entity.Key)
 	if err != nil {
 		return nil, fmt.Errorf("update entity: parse entity key: %w", err)