Improvements to the SQL storage indexer for observability and reliability (#1334)

Added some improvements into the SQL storage indexer. All spans related to
updating database are now nested under the same parent for better visibility.
Cleanup database process just runs if cursor has been updated in that iteration.
In case of failure dropping the database, it has been added a new column in each
row of the database to set the cursor where that package belongs. All the
entrypoints whose response is a JSON (root, search, categories and package
endpoints) are unified and they use the same method. And it is allowed to
customize the default batch sizes via environment variables.

Author: mrodm

CHANGELOG.md

@@ -9,10 +9,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Breaking changes
 
 ### Bugfixes
+* Fix spans used in SQL storage indexer (technical preview). [#1334](https://github.com/elastic/package-registry/pull/1334)
 
 ### Added
 * Add new method `MustParsePackage` to create new packages, running more validations. [#1333](https://github.com/elastic/package-registry/pull/1333)
-* Allow to customize SQL storage indexer configurations (technical preview). [#1337](https://github.coim/elastic/package-registry/pull/1337)
+* Allow to customize settings related to SQL storage indexer (technical preview). [#1334](https://github.com/elastic/package-registry/pull/1334) [#1337](https://github.coim/elastic/package-registry/pull/1337)
+* Cleanup SQL storage indexer backup database only when a new index version is downloaded (technical preview). [#1337](https://github.coim/elastic/package-registry/pull/1337)
 
 ### Deprecated
 

categories.go

@@ -92,9 +92,7 @@ func categoriesHandlerWithProxyMode(logger *zap.Logger, indexer Indexer, proxyMo
 			return
 		}
 
-		cacheHeaders(w, cacheTime)
-		jsonHeader(w)
-		w.Write(data)
+		serveJSONResponse(r.Context(), w, cacheTime, data)
 	}
 }
 
@@ -242,3 +240,12 @@ func getCategoriesOutput(ctx context.Context, categories map[string]*packages.Ca
 
 	return util.MarshalJSONPretty(outputCategories)
 }
+
+func serveJSONResponse(ctx context.Context, w http.ResponseWriter, cacheTime time.Duration, data []byte) {
+	span, _ := apm.StartSpan(ctx, "Serve JSON Response", "app")
+	defer span.End()
+
+	cacheHeaders(w, cacheTime)
+	jsonHeader(w)
+	w.Write(data)
+}

dev/docker-compose-epr-gcs.yml

@@ -26,6 +26,8 @@ services:
       - EPR_DISABLE_PACKAGE_VALIDATION=true
       - EPR_ADDRESS=0.0.0.0:8080
       # - EPR_LOG_LEVEL=debug
+      # - EPR_SQL_INDEXER_READ_PACKAGES_BATCH_SIZE=2000
+      # - EPR_SQL_INDEXER_DB_INSERT_BATCH_SIZE=2000
       # - EPR_SQL_INDEXER_DATABASE_FOLDER_PATH=/tmp
       # - EPR_SQL_INDEXER_SEARCH_CACHE_SIZE=100
       # - EPR_SQL_INDEXER_SEARCH_CACHE_TTL=10m

dev/docker-compose-epr.yml

@@ -19,6 +19,8 @@ services:
       - EPR_DISABLE_PACKAGE_VALIDATION=true
       - EPR_ADDRESS=0.0.0.0:8080
       # - EPR_LOG_LEVEL=debug
+      # - EPR_SQL_INDEXER_READ_PACKAGES_BATCH_SIZE=2000
+      # - EPR_SQL_INDEXER_DB_INSERT_BATCH_SIZE=2000
       # - EPR_SQL_INDEXER_DATABASE_FOLDER_PATH=/tmp
       # - EPR_SQL_INDEXER_SEARCH_CACHE_SIZE=100
       # - EPR_SQL_INDEXER_SEARCH_CACHE_TTL=10m

dev/launch_epr_service_storage_indexer.sh

@@ -108,6 +108,8 @@ export EPR_CONFIG="${CONFIG_PATH}"
 # export EPR_SQL_INDEXER_DATABASE_FOLDER_PATH=/tmp
 # export EPR_SQL_INDEXER_SEARCH_CACHE_SIZE=100
 # export EPR_SQL_INDEXER_SEARCH_CACHE_TTL=10m
+# export EPR_SQL_INDEXER_READ_PACKAGES_BATCH_SIZE=2000
+# export EPR_SQL_INDEXER_DB_INSERT_BATCH_SIZE=2000
 
 ./package-registry
 

index.go

@@ -26,8 +26,6 @@ func indexHandler(cacheTime time.Duration) (func(w http.ResponseWriter, r *http.
 		return nil, err
 	}
 	return func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", "application/json")
-		cacheHeaders(w, cacheTime)
-		w.Write(body)
+		serveJSONResponse(r.Context(), w, cacheTime, body)
 	}, nil
 }

internal/database/memoryrepository.go

@@ -12,7 +12,11 @@ import (
 	_ "modernc.org/sqlite" // Import the SQLite driver
 )
 
-func NewMemorySQLDB(path string) (*SQLiteRepository, error) {
+type MemorySQLDBOptions struct {
+	Path string
+}
+
+func NewMemorySQLDB(options MemorySQLDBOptions) (*SQLiteRepository, error) {
 	db, err := sql.Open("sqlite", ":memory:")
 	if err != nil {
 		return nil, fmt.Errorf("failed to open database: %w", err)
@@ -22,14 +26,14 @@ func NewMemorySQLDB(path string) (*SQLiteRepository, error) {
 		return nil, fmt.Errorf("failed to connect to database: %w", err)
 	}
 
-	dbRepo, err := newSQLiteRepository(db)
+	dbRepo, err := newSQLiteRepository(sqlDBOptions{db: db})
 	if err != nil {
 		return nil, fmt.Errorf("failed to create SQLite repository: %w", err)
 	}
 
 	if err := dbRepo.Initialize(context.Background()); err != nil {
 		return nil, fmt.Errorf("failed to create database: %w", err)
 	}
-	dbRepo.path = "memory-" + path
+	dbRepo.path = "memory-" + options.Path
 	return dbRepo, nil
 }

internal/database/model.go

@@ -5,6 +5,7 @@
 package database
 
 type Package struct {
+	Cursor          string
 	Name            string
 	Version         string
 	FormatVersion   string

internal/database/sqliterepository.go

@@ -29,6 +29,7 @@ type keyDefinition struct {
 }
 
 var keys = []keyDefinition{
+	{"cursor", "TEXT NOT NULL"},
 	{"name", "TEXT NOT NULL"},
 	{"version", "TEXT NOT NULL"},
 	{"formatVersion", "TEXT NOT NULL"},
@@ -47,20 +48,25 @@ var keys = []keyDefinition{
 const defaultMaxBulkAddBatch = 2000
 
 type SQLiteRepository struct {
-	db              *sql.DB
-	path            string
-	maxBulkAddBatch int
-	numberFields    int
+	db                  *sql.DB
+	path                string
+	maxBulkAddBatchSize int
+	maxTotalArgs        int
 }
 
 var _ Repository = new(SQLiteRepository)
 
-func NewFileSQLDB(path string) (*SQLiteRepository, error) {
+type FileSQLDBOptions struct {
+	Path             string
+	BatchSizeInserts int
+}
+
+func NewFileSQLDB(options FileSQLDBOptions) (*SQLiteRepository, error) {
 	// NOTE: Even using sqlcache (with Ristretto or Redis), data column needs to be processed (Unmarshalled)
 	// again for all the Get queries performed, so there is no advantage in time of using sqlcache with SQLite
 	// for our use case.
 
-	db, err := sql.Open("sqlite", path)
+	db, err := sql.Open("sqlite", options.Path)
 	if err != nil {
 		return nil, fmt.Errorf("failed to open database: %w", err)
 	}
@@ -69,19 +75,27 @@ func NewFileSQLDB(path string) (*SQLiteRepository, error) {
 		return nil, fmt.Errorf("failed to connect to database: %w", err)
 	}
 
-	dbRepo, err := newSQLiteRepository(db)
+	dbRepo, err := newSQLiteRepository(sqlDBOptions{
+		db:               db,
+		batchSizeInserts: options.BatchSizeInserts,
+	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to create SQLite repository: %w", err)
 	}
 	if err := dbRepo.Initialize(context.Background()); err != nil {
 		return nil, fmt.Errorf("failed to create database: %w", err)
 	}
-	dbRepo.path = path
+	dbRepo.path = options.Path
 
 	return dbRepo, nil
 }
 
-func newSQLiteRepository(db *sql.DB) (*SQLiteRepository, error) {
+type sqlDBOptions struct {
+	db               *sql.DB
+	batchSizeInserts int
+}
+
+func newSQLiteRepository(options sqlDBOptions) (*SQLiteRepository, error) {
 	// https://www.sqlite.org/pragma.html#pragma_journal_mode
 	// Not observed any performance difference with WAL mode, so keeping it as DELETE mode for now.
 	// if _, err = db.Exec("PRAGMA journal_mode = WAL;"); err != nil {
@@ -106,10 +120,14 @@ func newSQLiteRepository(db *sql.DB) (*SQLiteRepository, error) {
 	// if _, err := db.Exec("PRAGMA cache_size = -10000;"); err != nil {
 	// 	return nil, fmt.Errorf("failed to update cache_size: %w", err)
 	// }
+	batchSize := defaultMaxBulkAddBatch
+	if options.batchSizeInserts > 0 {
+		batchSize = options.batchSizeInserts
+	}
 	return &SQLiteRepository{
-		db:              db,
-		maxBulkAddBatch: defaultMaxBulkAddBatch,
-		numberFields:    len(keys),
+		db:                  options.db,
+		maxBulkAddBatchSize: batchSize,
+		maxTotalArgs:        batchSize * len(keys),
 	}, nil
 }
 
@@ -119,6 +137,7 @@ func (r *SQLiteRepository) File(ctx context.Context) string {
 
 func (r *SQLiteRepository) Ping(ctx context.Context) error {
 	span, ctx := apm.StartSpan(ctx, "SQL: Ping", "app")
+	span.Context.SetLabel("database.path", r.File(ctx))
 	defer span.End()
 	if r.db == nil {
 		return errors.New("database is not initialized")
@@ -131,14 +150,15 @@ func (r *SQLiteRepository) Ping(ctx context.Context) error {
 
 func (r *SQLiteRepository) Initialize(ctx context.Context) error {
 	span, ctx := apm.StartSpan(ctx, "SQL: Initialize", "app")
+	span.Context.SetLabel("database.path", r.File(ctx))
 	defer span.End()
 	createQuery := strings.Builder{}
 	createQuery.WriteString("CREATE TABLE IF NOT EXISTS ")
 	createQuery.WriteString("packages (")
 	for _, i := range keys {
 		createQuery.WriteString(fmt.Sprintf("%s %s, ", i.Name, i.SQLType))
 	}
-	createQuery.WriteString("PRIMARY KEY (name, version));")
+	createQuery.WriteString("PRIMARY KEY (cursor, name, version));")
 	if _, err := r.db.ExecContext(ctx, createQuery.String()); err != nil {
 		return err
 	}
@@ -159,19 +179,18 @@ func (r *SQLiteRepository) Initialize(ctx context.Context) error {
 
 func (r *SQLiteRepository) BulkAdd(ctx context.Context, database string, pkgs []*Package) error {
 	span, ctx := apm.StartSpan(ctx, "SQL: Insert batches", "app")
+	span.Context.SetLabel("insert.batch.size", r.maxBulkAddBatchSize)
+	span.Context.SetLabel("database.path", r.File(ctx))
 	defer span.End()
 
 	if len(pkgs) == 0 {
 		return nil
 	}
 
 	totalProcessed := 0
-	args := make([]any, 0, r.maxBulkAddBatch*r.numberFields)
+	args := make([]any, 0, r.maxTotalArgs)
 	for {
 		read := 0
-		// reuse args slice
-		args = args[:0]
-
 		var sb strings.Builder
 		sb.WriteString("INSERT INTO ")
 		sb.WriteString(database)
@@ -184,7 +203,7 @@ func (r *SQLiteRepository) BulkAdd(ctx context.Context, database string, pkgs []
 		}
 		sb.WriteString(") values")
 
-		endBatch := totalProcessed + r.maxBulkAddBatch
+		endBatch := totalProcessed + r.maxBulkAddBatchSize
 		for i := totalProcessed; i < endBatch && i < len(pkgs); i++ {
 			sb.WriteString("(")
 			for j := range keys {
@@ -206,6 +225,7 @@ func (r *SQLiteRepository) BulkAdd(ctx context.Context, database string, pkgs []
 			discoveryFields := addCommasToString(pkgs[i].DiscoveryFields)
 
 			args = append(args,
+				pkgs[i].Cursor,
 				pkgs[i].Name,
 				pkgs[i].Version,
 				pkgs[i].FormatVersion,
@@ -240,6 +260,9 @@ func (r *SQLiteRepository) BulkAdd(ctx context.Context, database string, pkgs []
 		if totalProcessed >= len(pkgs) {
 			break
 		}
+
+		// reuse args slice
+		args = args[:0]
 	}
 
 	return nil
@@ -263,6 +286,7 @@ func addCommasToString(s string) string {
 
 func (r *SQLiteRepository) All(ctx context.Context, database string, whereOptions WhereOptions) ([]*Package, error) {
 	span, ctx := apm.StartSpan(ctx, "SQL: Get All", "app")
+	span.Context.SetLabel("database.path", r.File(ctx))
 	defer span.End()
 
 	var all []*Package
@@ -276,6 +300,7 @@ func (r *SQLiteRepository) All(ctx context.Context, database string, whereOption
 
 func (r *SQLiteRepository) AllFunc(ctx context.Context, database string, whereOptions WhereOptions, process func(ctx context.Context, pkg *Package) error) error {
 	span, ctx := apm.StartSpan(ctx, "SQL: Get All (process each package)", "app")
+	span.Context.SetLabel("database.path", r.File(ctx))
 	defer span.End()
 
 	useBaseData := whereOptions == nil || !whereOptions.UseFullData()
@@ -301,8 +326,6 @@ func (r *SQLiteRepository) AllFunc(ctx context.Context, database string, whereOp
 		clause, whereArgs = whereOptions.Where()
 		query.WriteString(clause)
 	}
-	// TODO: remove debug
-	// fmt.Println("Query:", query.String())
 	rows, err := r.db.QueryContext(ctx, query.String(), whereArgs...)
 	if err != nil {
 		return err
@@ -313,6 +336,7 @@ func (r *SQLiteRepository) AllFunc(ctx context.Context, database string, whereOp
 	var pkg Package
 	for rows.Next() {
 		if err := rows.Scan(
+			&pkg.Cursor,
 			&pkg.Name,
 			&pkg.Version,
 			&pkg.FormatVersion,
@@ -345,6 +369,7 @@ func (r *SQLiteRepository) AllFunc(ctx context.Context, database string, whereOp
 
 func (r *SQLiteRepository) Drop(ctx context.Context, table string) error {
 	span, ctx := apm.StartSpan(ctx, "SQL: Drop", "app")
+	span.Context.SetLabel("database.path", r.File(ctx))
 	defer span.End()
 	query := fmt.Sprintf("DROP TABLE IF EXISTS %s", table)
 	_, err := r.db.ExecContext(ctx, query)
@@ -370,16 +395,34 @@ type FilterOptions struct {
 type SQLOptions struct {
 	Filter *FilterOptions
 
+	CurrentCursor string
+
 	IncludeFullData bool // If true, the query will return the full data field instead of the base data field
 }
 
 func (o *SQLOptions) Where() (string, []any) {
-	if o == nil || o.Filter == nil {
+	if o == nil {
 		return "", nil
 	}
 	var sb strings.Builder
 	var args []any
+	// Always filter by cursor
+	if o.CurrentCursor != "" {
+		sb.WriteString("cursor = ?")
+		args = append(args, o.CurrentCursor)
+	}
+
+	if o.Filter == nil {
+		if sb.Len() == 0 {
+			return "", nil
+		}
+		return fmt.Sprintf(" WHERE %s", sb.String()), args
+	}
+
 	if o.Filter.Type != "" {
+		if sb.Len() > 0 {
+			sb.WriteString(" AND ")
+		}
 		sb.WriteString("type = ?")
 		args = append(args, o.Filter.Type)
 	}

internal/storage/index.go

@@ -116,7 +116,7 @@ func LoadPackagesAndCursorFromIndex(ctx context.Context, logger *zap.Logger, sto
 	return anIndex, storageCursor.Current, nil
 }
 
-func loadSearchIndexAllBatches(ctx context.Context, logger *zap.Logger, storageClient *storage.Client, bucketName, rootStoragePath string, aCursor cursor, batchSize int, process func(packages.Packages) error) error {
+func loadSearchIndexAllBatches(ctx context.Context, logger *zap.Logger, storageClient *storage.Client, bucketName, rootStoragePath string, aCursor cursor, batchSize int, process func(context.Context, packages.Packages, string) error) error {
 	span, ctx := apm.StartSpan(ctx, "LoadSearchIndexAll", "app")
 	span.Context.SetLabel("load.method", "batches")
 	span.Context.SetLabel("load.batch.size", batchSize)
@@ -174,7 +174,7 @@ func loadSearchIndexAllBatches(ctx context.Context, logger *zap.Logger, storageC
 			count++
 
 			if count >= batchSize {
-				err = process(packages)
+				err = process(ctx, packages, aCursor.Current)
 				if err != nil {
 					return fmt.Errorf("error processing batch of packages: %w", err)
 				}
@@ -193,15 +193,15 @@ func loadSearchIndexAllBatches(ctx context.Context, logger *zap.Logger, storageC
 		}
 	}
 	if len(packages) > 0 {
-		err = process(packages)
+		err = process(ctx, packages, aCursor.Current)
 		if err != nil {
 			return fmt.Errorf("error processing final batch of packages: %w", err)
 		}
 	}
 	return nil
 }
 
-func LoadPackagesAndCursorFromIndexBatches(ctx context.Context, logger *zap.Logger, storageClient *storage.Client, storageBucketInternal, currentCursor string, batchSize int, process func(packages.Packages) error) (string, error) {
+func LoadPackagesAndCursorFromIndexBatches(ctx context.Context, logger *zap.Logger, storageClient *storage.Client, storageBucketInternal, currentCursor string, batchSize int, process func(context.Context, packages.Packages, string) error) (string, error) {
 	bucketName, rootStoragePath, err := extractBucketNameFromURL(storageBucketInternal)
 	if err != nil {
 		return "", fmt.Errorf("can't extract bucket name from URL (url: %s): %w", storageBucketInternal, err)