feat: queryPreprocessing in logicalRestore + support inline SQL in any queryPreprocessing (#407):

* support queryPreprocessing in logicalRestore
* add `inline` SQL processing to snapshot jobs (`logicalSnapshot`, `physicalSnapshot`) as well

Author: agneum

engine/configs/config.example.logical_generic.yml

@@ -288,6 +288,18 @@ retrieval:
         #   database2:
         #   databaseN:
 
+        # It is possible to define pre-processing SQL queries. For example, "/tmp/scripts/sql".
+        # Default: empty string (no pre-processing defined).
+        queryPreprocessing:
+          # Path to SQL pre-processing queries. Default: empty string (no pre-processing defined).
+          queryPath: ""
+
+          # Worker limit for parallel queries. Parallelization doesn't work for inline SQL queries.
+          maxParallelWorkers: 2
+
+          # Inline SQL. Queries run after scripts placed in 'queryPath'.
+          inline: ""
+
     logicalSnapshot:
       options:
         # Adjust PostgreSQL configuration
@@ -308,6 +320,9 @@ retrieval:
             # Worker limit for parallel queries.
             maxParallelWorkers: 2
 
+            # Inline SQL. Queries run after scripts placed in 'queryPath'.
+            inline: ""
+
 cloning:
   # Host that will be specified in database connection info for all clones
   # Use public IP address if database connections are allowed from outside

engine/configs/config.example.logical_rds_iam.yml

@@ -289,6 +289,18 @@ retrieval:
         #   database2:
         #   databaseN:
 
+        # It is possible to define pre-processing SQL queries. For example, "/tmp/scripts/sql".
+        # Default: empty string (no pre-processing defined).
+        queryPreprocessing:
+          # Path to SQL pre-processing queries. Default: empty string (no pre-processing defined).
+          queryPath: ""
+
+          # Worker limit for parallel queries. Parallelization doesn't work for inline SQL queries.
+          maxParallelWorkers: 2
+
+          # Inline SQL. Queries run after scripts placed in 'queryPath'.
+          inline: ""
+
     logicalSnapshot:
       options:
         # Adjust PostgreSQL configuration
@@ -309,6 +321,9 @@ retrieval:
             # Worker limit for parallel queries.
             maxParallelWorkers: 2
 
+            # Inline SQL. Queries run after scripts placed in 'queryPath'.
+            inline: ""
+
 cloning:
   # Host that will be specified in database connection info for all clones
   # Use public IP address if database connections are allowed from outside

engine/configs/config.example.physical_generic.yml

@@ -261,6 +261,9 @@ retrieval:
             # Worker limit for parallel queries.
             maxParallelWorkers: 2
 
+            # Inline SQL. Queries run after scripts placed in 'queryPath'.
+            inline: ""
+
           # Add PostgreSQL configuration parameters to the promotion container.
           configs:
             shared_buffers: 2GB

engine/configs/config.example.physical_pgbackrest.yml

@@ -265,6 +265,9 @@ retrieval:
             # Worker limit for parallel queries.
             maxParallelWorkers: 2
 
+            # Inline SQL. Queries run after scripts placed in 'queryPath'.
+            inline: ""
+
           # Add PostgreSQL configuration parameters to the promotion container.
           configs:
             shared_buffers: 2GB

engine/configs/config.example.physical_walg.yml

@@ -251,6 +251,9 @@ retrieval:
             # Worker limit for parallel queries.
             maxParallelWorkers: 2
 
+            # Inline SQL. Queries run after scripts placed in 'queryPath'.
+            inline: ""
+
           # Add PostgreSQL configuration parameters to the promotion container.
           configs:
             shared_buffers: 2GB

engine/internal/retrieval/engine/postgres/logical/restore.go

@@ -31,6 +31,7 @@ import (
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/cont"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/defaults"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/health"
+	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/query"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/options"
 
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/config/global"
@@ -87,19 +88,21 @@ type RestoreJob struct {
 	engineProps       global.EngineProps
 	dbMarker          *dbmarker.Marker
 	dbMark            *dbmarker.Config
+	queryProcessor    *query.Processor
 	isDumpLocationDir bool
 	RestoreOptions
 }
 
 // RestoreOptions defines a logical restore options.
 type RestoreOptions struct {
-	DumpLocation    string                    `yaml:"dumpLocation"`
-	DockerImage     string                    `yaml:"dockerImage"`
-	ContainerConfig map[string]interface{}    `yaml:"containerConfig"`
-	Databases       map[string]DumpDefinition `yaml:"databases"`
-	ForceInit       bool                      `yaml:"forceInit"`
-	ParallelJobs    int                       `yaml:"parallelJobs"`
-	Configs         map[string]string         `yaml:"configs"`
+	DumpLocation       string                    `yaml:"dumpLocation"`
+	DockerImage        string                    `yaml:"dockerImage"`
+	ContainerConfig    map[string]interface{}    `yaml:"containerConfig"`
+	Databases          map[string]DumpDefinition `yaml:"databases"`
+	ForceInit          bool                      `yaml:"forceInit"`
+	ParallelJobs       int                       `yaml:"parallelJobs"`
+	Configs            map[string]string         `yaml:"configs"`
+	QueryPreprocessing query.PreprocessorCfg     `yaml:"queryPreprocessing"`
 }
 
 // Partial defines tables and rules for a partial logical restore.
@@ -125,6 +128,10 @@ func NewJob(cfg config.JobConfig, global *global.Config, engineProps global.Engi
 
 	restoreJob.setDefaults()
 
+	if qp := restoreJob.RestoreOptions.QueryPreprocessing; qp.QueryPath != "" || qp.Inline != "" {
+		restoreJob.queryProcessor = query.NewQueryProcessor(cfg.Docker, qp, global.Database.Name(), global.Database.User())
+	}
+
 	return restoreJob, nil
 }
 
@@ -233,6 +240,10 @@ func (r *RestoreJob) Run(ctx context.Context) (err error) {
 		}
 	}
 
+	if err := r.queryProcessor.ApplyPreprocessingQueries(ctx, containerID); err != nil {
+		return errors.Wrap(err, "failed to run preprocessing queries")
+	}
+
 	dbList, err := r.getDBList(ctx, containerID)
 	if err != nil {
 		return err

engine/internal/retrieval/engine/postgres/snapshot/logical.go

@@ -29,6 +29,7 @@ import (
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/cont"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/health"
+	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/query"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/options"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/telemetry"
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/config/global"
@@ -54,7 +55,7 @@ type LogicalInitial struct {
 	globalCfg      *global.Config
 	engineProps    global.EngineProps
 	dbMarker       *dbmarker.Marker
-	queryProcessor *queryProcessor
+	queryProcessor *query.Processor
 }
 
 // LogicalOptions describes options for a logical initialization job.
@@ -68,7 +69,7 @@ type LogicalOptions struct {
 // DataPatching allows executing queries to transform data before snapshot taking.
 type DataPatching struct {
 	DockerImage        string                 `yaml:"dockerImage"`
-	QueryPreprocessing QueryPreprocessing     `yaml:"queryPreprocessing"`
+	QueryPreprocessing query.PreprocessorCfg  `yaml:"queryPreprocessing"`
 	ContainerConfig    map[string]interface{} `yaml:"containerConfig"`
 }
 
@@ -90,10 +91,8 @@ func NewLogicalInitialJob(cfg config.JobConfig, global *global.Config, enginePro
 		return nil, errors.Wrap(err, "failed to unmarshal configuration options")
 	}
 
-	if li.options.DataPatching.QueryPreprocessing.QueryPath != "" {
-		li.queryProcessor = newQueryProcessor(cfg.Docker, global.Database.Name(), global.Database.User(),
-			li.options.DataPatching.QueryPreprocessing.QueryPath,
-			li.options.DataPatching.QueryPreprocessing.MaxParallelWorkers)
+	if qp := li.options.DataPatching.QueryPreprocessing; qp.QueryPath != "" || qp.Inline != "" {
+		li.queryProcessor = query.NewQueryProcessor(cfg.Docker, qp, global.Database.Name(), global.Database.User())
 	}
 
 	return li, nil
@@ -254,7 +253,7 @@ func (s *LogicalInitial) runPreprocessingQueries(ctx context.Context, dataDir st
 		return errors.Wrap(err, "failed to readiness check")
 	}
 
-	if err := s.queryProcessor.applyPreprocessingQueries(ctx, containerID); err != nil {
+	if err := s.queryProcessor.ApplyPreprocessingQueries(ctx, containerID); err != nil {
 		return errors.Wrap(err, "failed to run preprocessing queries")
 	}
 

engine/internal/retrieval/engine/postgres/snapshot/physical.go

@@ -38,6 +38,7 @@ import (
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/defaults"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/health"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/pgtool"
+	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/query"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/options"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/telemetry"
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/config/global"
@@ -96,7 +97,7 @@ type PhysicalInitial struct {
 	scheduler      *cron.Cron
 	schedulerCtx   context.Context
 	promotionMutex sync.Mutex
-	queryProcessor *queryProcessor
+	queryProcessor *query.Processor
 	tm             *telemetry.Agent
 }
 
@@ -117,7 +118,7 @@ type Promotion struct {
 	DockerImage        string                 `yaml:"dockerImage"`
 	ContainerConfig    map[string]interface{} `yaml:"containerConfig"`
 	HealthCheck        HealthCheck            `yaml:"healthCheck"`
-	QueryPreprocessing QueryPreprocessing     `yaml:"queryPreprocessing"`
+	QueryPreprocessing query.PreprocessorCfg  `yaml:"queryPreprocessing"`
 	Configs            map[string]string      `yaml:"configs"`
 	Recovery           map[string]string      `yaml:"recovery"`
 }
@@ -177,10 +178,8 @@ func NewPhysicalInitialJob(
 		return nil, errors.Wrap(err, "invalid physicalSnapshot configuration")
 	}
 
-	if p.options.Promotion.QueryPreprocessing.QueryPath != "" {
-		p.queryProcessor = newQueryProcessor(cfg.Docker, global.Database.Name(), global.Database.User(),
-			p.options.Promotion.QueryPreprocessing.QueryPath,
-			p.options.Promotion.QueryPreprocessing.MaxParallelWorkers)
+	if qp := p.options.Promotion.QueryPreprocessing; qp.QueryPath != "" || qp.Inline != "" {
+		p.queryProcessor = query.NewQueryProcessor(cfg.Docker, qp, global.Database.Name(), global.Database.User())
 	}
 
 	p.setupScheduler()
@@ -631,7 +630,7 @@ func (p *PhysicalInitial) promoteInstance(ctx context.Context, clonePath string,
 	}
 
 	if p.queryProcessor != nil {
-		if err := p.queryProcessor.applyPreprocessingQueries(ctx, containerID); err != nil {
+		if err := p.queryProcessor.ApplyPreprocessingQueries(ctx, containerID); err != nil {
 			return errors.Wrap(err, "failed to run preprocessing queries")
 		}
 	}