Refactoring common mysql

server/datastore/mysql/apple_mdm_test.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/VividCortex/mysqlerr"
 	"github.com/fleetdm/fleet/v4/pkg/optjson"
+	"github.com/fleetdm/fleet/v4/server/datastore/mysql/common_mysql"
 	"github.com/fleetdm/fleet/v4/server/datastore/s3"
 	"github.com/fleetdm/fleet/v4/server/fleet"
 	fleetmdm "github.com/fleetdm/fleet/v4/server/mdm"
@@ -4650,7 +4651,7 @@ func testMDMAppleResetEnrollment(t *testing.T, ds *Datastore) {
 	// host has no boostrap package command yet
 	_, err = ds.GetHostBootstrapPackageCommand(ctx, host.UUID)
 	require.Error(t, err)
-	nfe := &notFoundError{}
+	nfe := &common_mysql.NotFoundError{}
 	require.ErrorAs(t, err, &nfe)
 
 	err = ds.RecordHostBootstrapPackage(ctx, "command-uuid", host.UUID)
@@ -5560,7 +5561,7 @@ func TestRestorePendingDEPHost(t *testing.T) {
 				require.Equal(t, expectedHost.Platform, h.Platform)
 				require.Equal(t, expectedHost.TeamID, h.TeamID)
 			} else {
-				nfe := &notFoundError{}
+				nfe := &common_mysql.NotFoundError{}
 				require.ErrorAs(t, err, &nfe)
 			}
 

server/datastore/mysql/common_mysql/common.go

@@ -0,0 +1,141 @@
+package common_mysql
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"net/url"
+	"time"
+
+	"github.com/fleetdm/fleet/v4/server/config"
+	"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
+	"github.com/go-kit/log"
+	"github.com/go-sql-driver/mysql"
+	"github.com/jmoiron/sqlx"
+	"github.com/ngrok/sqlmw"
+)
+
+type DBOptions struct {
+	// MaxAttempts configures the number of retries to connect to the DB
+	MaxAttempts         int
+	Logger              log.Logger
+	ReplicaConfig       *config.MysqlConfig
+	Interceptor         sqlmw.Interceptor
+	TracingConfig       *config.LoggingConfig
+	MinLastOpenedAtDiff time.Duration
+	SqlMode             string
+	PrivateKey          string
+}
+
+func NewDB(conf *config.MysqlConfig, opts *DBOptions, otelDriverName string) (*sqlx.DB, error) {
+	driverName := "mysql"
+	if opts.TracingConfig != nil && opts.TracingConfig.TracingEnabled {
+		if opts.TracingConfig.TracingType == "opentelemetry" {
+			driverName = otelDriverName
+		} else {
+			driverName = "apm/mysql"
+		}
+	}
+	if opts.Interceptor != nil {
+		driverName = "mysql-mw"
+		sql.Register(driverName, sqlmw.Driver(mysql.MySQLDriver{}, opts.Interceptor))
+	}
+	if opts.SqlMode != "" {
+		conf.SQLMode = opts.SqlMode
+	}
+
+	dsn := generateMysqlConnectionString(*conf)
+	db, err := sqlx.Open(driverName, dsn)
+	if err != nil {
+		return nil, err
+	}
+
+	db.SetMaxIdleConns(conf.MaxIdleConns)
+	db.SetMaxOpenConns(conf.MaxOpenConns)
+	db.SetConnMaxLifetime(time.Second * time.Duration(conf.ConnMaxLifetime))
+
+	var dbError error
+	for attempt := 0; attempt < opts.MaxAttempts; attempt++ {
+		dbError = db.Ping()
+		if dbError == nil {
+			// we're connected!
+			break
+		}
+		interval := time.Duration(attempt) * time.Second
+		opts.Logger.Log("mysql", fmt.Sprintf(
+			"could not connect to db: %v, sleeping %v", dbError, interval))
+		time.Sleep(interval)
+	}
+
+	if dbError != nil {
+		return nil, dbError
+	}
+	return db, nil
+}
+
+// generateMysqlConnectionString returns a MySQL connection string using the
+// provided configuration.
+func generateMysqlConnectionString(conf config.MysqlConfig) string {
+	params := url.Values{
+		// using collation implicitly sets the charset too
+		// and it's the recommended way to do it per the
+		// driver documentation:
+		// https://github.com/go-sql-driver/mysql#charset
+		"collation":            []string{"utf8mb4_unicode_ci"},
+		"parseTime":            []string{"true"},
+		"loc":                  []string{"UTC"},
+		"time_zone":            []string{"'-00:00'"},
+		"clientFoundRows":      []string{"true"},
+		"allowNativePasswords": []string{"true"},
+		"group_concat_max_len": []string{"4194304"},
+		"multiStatements":      []string{"true"},
+	}
+	if conf.TLSConfig != "" {
+		params.Set("tls", conf.TLSConfig)
+	}
+	if conf.SQLMode != "" {
+		params.Set("sql_mode", conf.SQLMode)
+	}
+
+	dsn := fmt.Sprintf(
+		"%s:%s@%s(%s)/%s?%s",
+		conf.Username,
+		conf.Password,
+		conf.Protocol,
+		conf.Address,
+		conf.Database,
+		params.Encode(),
+	)
+
+	return dsn
+}
+
+func WithTxx(ctx context.Context, db *sqlx.DB, fn TxFn, logger log.Logger) error {
+	tx, err := db.BeginTxx(ctx, nil)
+	if err != nil {
+		return ctxerr.Wrap(ctx, err, "create transaction")
+	}
+
+	defer func() {
+		if p := recover(); p != nil {
+			if err := tx.Rollback(); err != nil {
+				logger.Log("err", err, "msg", "error encountered during transaction panic rollback")
+			}
+			panic(p)
+		}
+	}()
+
+	if err := fn(tx); err != nil {
+		rbErr := tx.Rollback()
+		if rbErr != nil && rbErr != sql.ErrTxDone {
+			return ctxerr.Wrapf(ctx, err, "got err '%s' rolling back after err", rbErr.Error())
+		}
+		return err
+	}
+
+	if err := tx.Commit(); err != nil {
+		return ctxerr.Wrap(ctx, err, "commit transaction")
+	}
+
+	return nil
+}

server/mdm/android/mysql/errors.go → server/datastore/mysql/common_mysql/errors.go

@@ -1,26 +1,29 @@
-package mysql
-
-// TODO(26218): Refactor this to common_mysql
+package common_mysql
 
 import (
 	"database/sql"
 	"fmt"
+
+	"github.com/fleetdm/fleet/v4/server/fleet"
 )
 
-type notFoundError struct {
+type NotFoundError struct {
 	ID           uint
 	Name         string
 	Message      string
 	ResourceType string
 }
 
-func notFound(kind string) *notFoundError {
-	return &notFoundError{
+// Compile-time interface check.
+var _ fleet.NotFoundError = &NotFoundError{}
+
+func NotFound(kind string) *NotFoundError {
+	return &NotFoundError{
 		ResourceType: kind,
 	}
 }
 
-func (e *notFoundError) Error() string {
+func (e *NotFoundError) Error() string {
 	if e.ID != 0 {
 		return fmt.Sprintf("%s %d was not found in the datastore", e.ResourceType, e.ID)
 	}
@@ -33,28 +36,28 @@ func (e *notFoundError) Error() string {
 	return fmt.Sprintf("%s was not found in the datastore", e.ResourceType)
 }
 
-func (e *notFoundError) WithID(id uint) error {
+func (e *NotFoundError) WithID(id uint) error {
 	e.ID = id
 	return e
 }
 
-func (e *notFoundError) WithName(name string) error {
+func (e *NotFoundError) WithName(name string) error {
 	e.Name = name
 	return e
 }
 
-func (e *notFoundError) WithMessage(msg string) error {
+func (e *NotFoundError) WithMessage(msg string) error {
 	e.Message = msg
 	return e
 }
 
-func (e *notFoundError) IsNotFound() bool {
+func (e *NotFoundError) IsNotFound() bool {
 	return true
 }
 
 // Is helps so that errors.Is(err, sql.ErrNoRows) returns true for an
-// error of type *notFoundError, without having to wrap sql.ErrNoRows
+// error of type *NotFoundError, without having to wrap sql.ErrNoRows
 // explicitly.
-func (e *notFoundError) Is(other error) bool {
+func (e *NotFoundError) Is(other error) bool {
 	return other == sql.ErrNoRows
 }