Refactor

FerretDB · AlekSi · Jul 27, 2023 · Jul 26, 2023 · Jul 26, 2023 · Jul 26, 2023
commit 6a3251e4bd04dc7db8c4d08b2ff4d852e34b77ea
@@ -31,9 +31,12 @@ import (
 	"go.uber.org/zap"
 	"golang.org/x/exp/maps"
 	"golang.org/x/exp/slices"
+	_ "modernc.org/sqlite" // register database/sql driver
 
+	"github.com/FerretDB/FerretDB/internal/util/fsql"
 	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
 	"github.com/FerretDB/FerretDB/internal/util/resource"
+	"github.com/prometheus/client_golang/prometheus"
 )
 
 // filenameExtension represents SQLite database filename extension.
@@ -47,11 +50,39 @@ type Pool struct {
 	l   *zap.Logger
 
 	rw  sync.RWMutex
-	dbs map[string]*db
+	dbs map[string]*fsql.DB
 
 	token *resource.Token
 }
 
+// openDB opens existing database or creates a new one.
+//
+// All valid FerretDB database names are valid SQLite database names / file names,
+// so no validation is needed.
+// One exception is very long full path names for the filesystem,
+// but we don't check it.
+func openDB(uri string, l *zap.Logger) (*fsql.DB, error) {
+	db, err := sql.Open("sqlite", uri)
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	// TODO https://github.com/FerretDB/FerretDB/issues/2909
+
+	// TODO https://github.com/FerretDB/FerretDB/issues/2755
+	db.SetConnMaxIdleTime(0)
+	db.SetConnMaxLifetime(0)
+	// db.SetMaxIdleConns(5)
+	// db.SetMaxOpenConns(5)
+
+	if err = db.Ping(); err != nil {
+		_ = db.Close()
+		return nil, lazyerrors.Error(err)
+	}
+
+	return fsql.WrapDB(uri, db, l), nil
+}
+
 // New creates a pool for SQLite databases in the directory specified by SQLite URI.
 //
 // All databases are opened on creation.
@@ -69,7 +100,7 @@ func New(u string, l *zap.Logger) (*Pool, error) {
 	p := &Pool{
 		uri:   uri,
 		l:     l,
-		dbs:   make(map[string]*db, len(matches)),
+		dbs:   make(map[string]*fsql.DB, len(matches)),
 		token: resource.NewToken(),
 	}
 
@@ -81,7 +112,7 @@ func New(u string, l *zap.Logger) (*Pool, error) {
 
 		p.l.Debug("Opening existing database.", zap.String("name", name), zap.String("uri", uri))
 
-		db, err := openDB(uri)
+		db, err := openDB(uri, l)
 		if err != nil {
 			p.Close()
 			return nil, lazyerrors.Error(err)
@@ -138,7 +169,7 @@ func (p *Pool) List(ctx context.Context) []string {
 }
 
 // GetExisting returns an existing database by valid name, or nil.
-func (p *Pool) GetExisting(ctx context.Context, name string) *sql.DB {
+func (p *Pool) GetExisting(ctx context.Context, name string) *fsql.DB {
 	p.rw.RLock()
 	defer p.rw.RUnlock()
 
@@ -147,28 +178,28 @@ func (p *Pool) GetExisting(ctx context.Context, name string) *sql.DB {
 		return nil
 	}
 
-	return db.sqlDB
+	return db
 }
 
 // GetOrCreate returns an existing database by valid name, or creates a new one.
 //
 // Returned boolean value indicates whether the database was created.
-func (p *Pool) GetOrCreate(ctx context.Context, name string) (*sql.DB, bool, error) {
-	sqlDB := p.GetExisting(ctx, name)
-	if sqlDB != nil {
-		return sqlDB, false, nil
+func (p *Pool) GetOrCreate(ctx context.Context, name string) (*fsql.DB, bool, error) {
+	db := p.GetExisting(ctx, name)
+	if db != nil {
+		return db, false, nil
 	}
 
 	p.rw.Lock()
 	defer p.rw.Unlock()
 
 	// it might have been created by a concurrent call
 	if db := p.dbs[name]; db != nil {
-		return db.sqlDB, false, nil
+		return db, false, nil
 	}
 
 	uri := p.databaseURI(name)
-	db, err := openDB(uri)
+	db, err := openDB(uri, p.l)
 	if err != nil {
 		return nil, false, lazyerrors.Errorf("%s: %w", uri, err)
 	}
@@ -177,7 +208,7 @@ func (p *Pool) GetOrCreate(ctx context.Context, name string) (*sql.DB, bool, err
 
 	p.dbs[name] = db
 
-	return db.sqlDB, true, nil
+	return db, true, nil
 }
 
 // Drop closes and removes a database by valid name.
@@ -202,3 +233,24 @@ func (p *Pool) Drop(ctx context.Context, name string) bool {
 
 	return true
 }
+
+// Describe implements prometheus.Collector.
+func (p *Pool) Describe(ch chan<- *prometheus.Desc) {
+	// FIXME that's not correct if there no databases
+	prometheus.DescribeByCollect(p, ch)
+}
+
+// Collect implements prometheus.Collector.
+func (p *Pool) Collect(ch chan<- prometheus.Metric) {
+	p.rw.RLock()
+	defer p.rw.RUnlock()
+
+	for _, db := range p.dbs {
+		db.Collect(ch)
+	}
+}
+
+// check interfaces
+var (
+	_ prometheus.Collector = (*Pool)(nil)
+)
@@ -28,6 +28,7 @@ import (
 	sqlitelib "modernc.org/sqlite/lib"
 
 	"github.com/FerretDB/FerretDB/internal/backends/sqlite/metadata/pool"
+	"github.com/FerretDB/FerretDB/internal/util/fsql"
 	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
 	"github.com/FerretDB/FerretDB/internal/util/must"
 )
@@ -74,12 +75,12 @@ func (r *Registry) DatabaseList(ctx context.Context) []string {
 }
 
 // DatabaseGetExisting returns a connection to existing database or nil if it doesn't exist.
-func (r *Registry) DatabaseGetExisting(ctx context.Context, dbName string) *sql.DB {
+func (r *Registry) DatabaseGetExisting(ctx context.Context, dbName string) *fsql.DB {
 	return r.p.GetExisting(ctx, dbName)
 }
 
 // DatabaseGetOrCreate returns a connection to existing database or newly created database.
-func (r *Registry) DatabaseGetOrCreate(ctx context.Context, dbName string) (*sql.DB, error) {
+func (r *Registry) DatabaseGetOrCreate(ctx context.Context, dbName string) (*fsql.DB, error) {
 	db, created, err := r.p.GetOrCreate(ctx, dbName)
 	if err != nil {
 		return nil, lazyerrors.Error(err)

@@ -16,15 +16,110 @@
 package fsql
 
 import (
+	"context"
 	"database/sql"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"go.uber.org/zap"
+
+	"github.com/FerretDB/FerretDB/internal/util/observability"
+	"github.com/FerretDB/FerretDB/internal/util/resource"
 )
 
-// DB is a subset of [*database/sql.DB] methods that we use.
-type DB interface {
+// sqlDB is a subset of [*database/sql.DB] methods that we use.
+//
+// It mainly exist to check interfaces.
+type sqlDB interface {
+	Close() error
+	QueryContext(context.Context, string, ...any) (*sql.Rows, error)
+	QueryRowContext(context.Context, string, ...any) *sql.Row
+	ExecContext(context.Context, string, ...any) (sql.Result, error)
 	Stats() sql.DBStats
 }
 
+type DB struct {
+	*metricsCollector
+
+	sqlDB sqlDB
+	l     *zap.Logger
+	token *resource.Token
+}
+
+// TODO do not pass both name and l?
+func WrapDB(name string, db *sql.DB, l *zap.Logger) *DB {
+	res := &DB{
+		metricsCollector: newMetricsCollector(name, db.Stats),
+		sqlDB:            db,
+		l:                l,
+		token:            resource.NewToken(),
+	}
+
+	resource.Track(res, res.token)
+
+	return res
+}
+
+func (db *DB) Close() error {
+	resource.Untrack(db, db.token)
+	return db.sqlDB.Close()
+}
+
+func (db *DB) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
+	defer observability.FuncCall(ctx)()
+
+	start := time.Now()
+
+	fields := []zap.Field{zap.Any("args", args)}
+	db.l.Sugar().With(fields).Debugf(">>> %s", query)
+
+	rows, err := db.sqlDB.QueryContext(ctx, query, args...)
+
+	fields = append(fields, zap.Duration("duration", time.Since(start)), zap.Error(err))
+	db.l.Sugar().With(fields).Debugf("<<< %s", query)
+
+	return rows, err
+}
+
+func (db *DB) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row {
+	defer observability.FuncCall(ctx)()
+
+	start := time.Now()
+
+	fields := []zap.Field{zap.Any("args", args)}
+	db.l.Sugar().With(fields).Debugf(">>> %s", query)
+
+	row := db.sqlDB.QueryRowContext(ctx, query, args...)
+
+	fields = append(fields, zap.Duration("duration", time.Since(start)), zap.Error(row.Err()))
+	db.l.Sugar().With(fields).Debugf("<<< %s", query)
+
+	return row
+}
+
+func (db *DB) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
+	defer observability.FuncCall(ctx)()
+
+	start := time.Now()
+
+	fields := []zap.Field{zap.Any("args", args)}
+	db.l.Sugar().With(fields).Debugf(">>> %s", query)
+
+	res, err := db.sqlDB.ExecContext(ctx, query, args...)
+
+	fields = append(fields, zap.Duration("duration", time.Since(start)), zap.Error(err))
+	db.l.Sugar().With(fields).Debugf("<<< %s", query)
+
+	return res, err
+}
+
+func (db *DB) Stats() sql.DBStats {
+	return db.sqlDB.Stats()
+}
+
 // check interfaces
 var (
-	_ DB = (*sql.DB)(nil)
+	_ sqlDB                = (*sql.DB)(nil)
+	_ sqlDB                = (*DB)(nil)
+	_ prometheus.Collector = (*DB)(nil)
 )