Merge upstream/main into add-discovery-datasets

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)
 * Add support to discover content packages based on the datasets defined in the discovery parameter. [#1338](https://github.com/elastic/package-registry/pull/1338)
 
 ### 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)