Convert the QueryRunnerState store into a basestore (#17410)

internal/db/saved_queries.go

@@ -3,13 +3,51 @@ package db
 import (
 	"context"
 	"database/sql"
+	"sync"
 	"time"
 
 	"github.com/pkg/errors"
+
+	"github.com/sourcegraph/sourcegraph/internal/db/basestore"
 	"github.com/sourcegraph/sourcegraph/internal/db/dbconn"
+	"github.com/sourcegraph/sourcegraph/internal/db/dbutil"
 )
 
-type queryRunnerState struct{}
+type QueryRunnerStateStore struct {
+	*basestore.Store
+
+	once sync.Once
+}
+
+// NewQueryRunnerStateStoreWithDB instantiates and returns a new QueryRunnerStateStore with prepared statements.
+func NewQueryRunnerStateStoreWithDB(db dbutil.DB) *QueryRunnerStateStore {
+	return &QueryRunnerStateStore{Store: basestore.NewWithDB(db, sql.TxOptions{})}
+}
+
+// NewQueryRunnerStateStoreWithDB instantiates and returns a new QueryRunnerStateStore using the other store handle.
+func NewQueryRunnerStateStoreWith(other basestore.ShareableStore) *QueryRunnerStateStore {
+	return &QueryRunnerStateStore{Store: basestore.NewWithHandle(other.Handle())}
+}
+
+func (s *QueryRunnerStateStore) With(other basestore.ShareableStore) *QueryRunnerStateStore {
+	return &QueryRunnerStateStore{Store: s.Store.With(other)}
+}
+
+func (s *QueryRunnerStateStore) Transact(ctx context.Context) (*QueryRunnerStateStore, error) {
+	txBase, err := s.Store.Transact(ctx)
+	return &QueryRunnerStateStore{Store: txBase}, err
+}
+
+// ensureStore instantiates a basestore.Store if necessary, using the dbconn.Global handle.
+// This function ensures access to dbconn happens after the rest of the code or tests have
+// initialized it.
+func (s *QueryRunnerStateStore) ensureStore() {
+	s.once.Do(func() {
+		if s.Store == nil {
+			s.Store = basestore.NewWithDB(dbconn.Global, sql.TxOptions{})
+		}
+	})
+}
 
 type SavedQueryInfo struct {
 	Query        string
@@ -20,12 +58,14 @@ type SavedQueryInfo struct {
 
 // Get gets the saved query information for the given query. nil
 // is returned if there is no existing saved query info.
-func (s *queryRunnerState) Get(ctx context.Context, query string) (*SavedQueryInfo, error) {
+func (s *QueryRunnerStateStore) Get(ctx context.Context, query string) (*SavedQueryInfo, error) {
+	s.ensureStore()
+
 	info := &SavedQueryInfo{
 		Query: query,
 	}
 	var execDurationNs int64
-	err := dbconn.Global.QueryRowContext(
+	err := s.Handle().DB().QueryRowContext(
 		ctx,
 		"SELECT last_executed, latest_result, exec_duration_ns FROM query_runner_state WHERE query=$1",
 		query,
@@ -44,8 +84,10 @@ func (s *queryRunnerState) Get(ctx context.Context, query string) (*SavedQueryIn
 //
 // It is not safe to call concurrently for the same info.Query, as it uses a
 // poor man's upsert implementation.
-func (s *queryRunnerState) Set(ctx context.Context, info *SavedQueryInfo) error {
-	res, err := dbconn.Global.ExecContext(
+func (s *QueryRunnerStateStore) Set(ctx context.Context, info *SavedQueryInfo) error {
+	s.ensureStore()
+
+	res, err := s.Handle().DB().ExecContext(
 		ctx,
 		"UPDATE query_runner_state SET last_executed=$1, latest_result=$2, exec_duration_ns=$3 WHERE query=$4",
 		info.LastExecuted,
@@ -62,7 +104,7 @@ func (s *queryRunnerState) Set(ctx context.Context, info *SavedQueryInfo) error
 	}
 	if updated == 0 {
 		// Didn't update any row, so insert a new one.
-		_, err := dbconn.Global.ExecContext(
+		_, err := s.Handle().DB().ExecContext(
 			ctx,
 			"INSERT INTO query_runner_state(query, last_executed, latest_result, exec_duration_ns) VALUES($1, $2, $3, $4)",
 			info.Query,
@@ -77,8 +119,10 @@ func (s *queryRunnerState) Set(ctx context.Context, info *SavedQueryInfo) error
 	return nil
 }
 
-func (s *queryRunnerState) Delete(ctx context.Context, query string) error {
-	_, err := dbconn.Global.ExecContext(
+func (s *QueryRunnerStateStore) Delete(ctx context.Context, query string) error {
+	s.ensureStore()
+
+	_, err := s.Handle().DB().ExecContext(
 		ctx,
 		"DELETE FROM query_runner_state WHERE query=$1",
 		query,

internal/db/stores.go

@@ -6,7 +6,7 @@ var (
 	DefaultRepos     = &DefaultRepoStore{}
 	Repos            = &RepoStore{}
 	Phabricator      = &phabricator{}
-	QueryRunnerState = &queryRunnerState{}
+	QueryRunnerState = &QueryRunnerStateStore{}
 	Namespaces       = &namespaces{}
 	Orgs             = &OrgStore{}
 	OrgMembers       = &orgMembers{}