feat: allow users to run a preprocessing SQL queries during the promotion stage when PostgreSQL is running and ready to receive commands (#181)

Author: agneum

configs/config.example.logical_generic.yml

@@ -215,6 +215,18 @@ retrieval:
 
     logicalSnapshot:
       options:
+        # Define pre-precessing SQL queries for data patching. For example, "/tmp/scripts/sql".
+        dataPatching:
+          # The Docker image to run data patching container.
+          dockerImage: "sync-instance:12"
+
+          queryPreprocessing:
+            # Path to SQL pre-processing queries. Default: empty string (no pre-processing defined).
+            queryPath: ""
+
+            # Worker limit for parallel queries.
+            maxParallelWorkers: 2
+
         # It is possible to define a pre-precessing script. For example, "/tmp/scripts/custom.sh".
         # Default: empty string (no pre-processing defined).
         # This can be used for scrubbing eliminating PII data, to define data masking, etc.

configs/config.example.logical_rds_iam.yml

@@ -216,6 +216,18 @@ retrieval:
 
     logicalSnapshot:
       options:
+        # Define pre-precessing SQL queries for data patching. For example, "/tmp/scripts/sql".
+        dataPatching:
+          # The Docker image to run data patching container.
+          dockerImage: "sync-instance:12"
+
+          queryPreprocessing:
+            # Path to SQL pre-processing queries. Default: empty string (no pre-processing defined).
+            queryPath: ""
+
+            # Worker limit for parallel queries.
+            maxParallelWorkers: 2
+
         # It is possible to define a pre-precessing script. For example, "/tmp/scripts/custom.sh".
         # Default: empty string (no pre-processing defined).
         # This can be used for scrubbing eliminating PII data, to define data masking, etc.

configs/config.example.physical_generic.yml

@@ -201,6 +201,15 @@ retrieval:
             # Maximum number of health check retries.
             maxRetries: 200
 
+          # It is possible to define pre-precessing SQL queries. For example, "/tmp/scripts/sql".
+          # Default: empty string (no pre-processing defined).
+          queryPreprocessing:
+            # Path to SQL pre-processing queries.
+            queryPath: ""
+
+            # Worker limit for parallel queries.
+            maxParallelWorkers: 2
+
         # It is possible to define a pre-precessing script. For example, "/tmp/scripts/custom.sh".
         # Default: empty string (no pre-processing defined).
         # This can be used for scrubbing eliminating PII data, to define data masking, etc.

configs/config.example.physical_walg.yml

@@ -191,6 +191,15 @@ retrieval:
             # Maximum number of health check retries.
             maxRetries: 200
 
+          # It is possible to define pre-precessing SQL queries. For example, "/tmp/scripts/sql".
+          # Default: empty string (no pre-processing defined).
+          queryPreprocessing:
+            # Path to SQL pre-processing queries.
+            queryPath: ""
+
+            # Worker limit for parallel queries.
+            maxParallelWorkers: 2
+
         # It is possible to define a pre-precessing script. For example, "/tmp/scripts/custom.sh".
         # Default: empty string (no pre-processing defined).
         # This can be used for scrubbing eliminating PII data, to define data masking, etc.

pkg/retrieval/engine/postgres/job_builder.go

@@ -55,7 +55,7 @@ func (s *JobBuilder) BuildJob(jobCfg config.JobConfig) (components.JobRunner, er
 		return physical.NewJob(jobCfg, s.dockerClient, s.globalCfg, s.dbMarker)
 
 	case snapshot.LogicalInitialType:
-		return snapshot.NewLogicalInitialJob(jobCfg, s.cloneManager, s.globalCfg, s.dbMarker)
+		return snapshot.NewLogicalInitialJob(jobCfg, s.dockerClient, s.cloneManager, s.globalCfg, s.dbMarker)
 
 	case snapshot.PhysicalInitialType:
 		return snapshot.NewPhysicalInitialJob(jobCfg, s.dockerClient, s.cloneManager, s.globalCfg, s.dbMarker)

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

@@ -7,42 +7,61 @@ package snapshot
 
 import (
 	"context"
+	"fmt"
 	"path"
 
+	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/api/types/network"
+	"github.com/docker/docker/client"
 	"github.com/pkg/errors"
 
 	dblabCfg "gitlab.com/postgres-ai/database-lab/pkg/config"
+	"gitlab.com/postgres-ai/database-lab/pkg/log"
 	"gitlab.com/postgres-ai/database-lab/pkg/retrieval/config"
 	"gitlab.com/postgres-ai/database-lab/pkg/retrieval/dbmarker"
 	"gitlab.com/postgres-ai/database-lab/pkg/retrieval/engine/postgres/tools"
+	"gitlab.com/postgres-ai/database-lab/pkg/retrieval/engine/postgres/tools/cont"
+	"gitlab.com/postgres-ai/database-lab/pkg/retrieval/engine/postgres/tools/health"
 	"gitlab.com/postgres-ai/database-lab/pkg/retrieval/options"
 	"gitlab.com/postgres-ai/database-lab/pkg/services/provision/databases/postgres/configuration"
 	"gitlab.com/postgres-ai/database-lab/pkg/services/provision/thinclones"
 )
 
 // LogicalInitial describes a job for preparing a logical initial snapshot.
 type LogicalInitial struct {
-	name         string
-	cloneManager thinclones.Manager
-	options      LogicalOptions
-	globalCfg    *dblabCfg.Global
-	dbMarker     *dbmarker.Marker
+	name           string
+	cloneManager   thinclones.Manager
+	dockerClient   *client.Client
+	options        LogicalOptions
+	globalCfg      *dblabCfg.Global
+	dbMarker       *dbmarker.Marker
+	queryProcessor *queryProcessor
 }
 
 // LogicalOptions describes options for a logical initialization job.
 type LogicalOptions struct {
+	DataPatching        DataPatching      `yaml:"dataPatching"`
 	PreprocessingScript string            `yaml:"preprocessingScript"`
 	Configs             map[string]string `yaml:"configs"`
 	Schedule            Scheduler         `yaml:"schedule"`
 }
 
+// DataPatching allows executing queries to transform data before snapshot taking.
+type DataPatching struct {
+	DockerImage        string             `yaml:"dockerImage"`
+	QueryPreprocessing QueryPreprocessing `yaml:"queryPreprocessing"`
+}
+
 const (
 	// LogicalInitialType declares a job type for preparing a logical initial snapshot.
 	LogicalInitialType = "logicalSnapshot"
+
+	patchContainerPrefix = "dblab_patch_"
 )
 
 // NewLogicalInitialJob creates a new logical initial job.
-func NewLogicalInitialJob(cfg config.JobConfig, cloneManager thinclones.Manager,
+func NewLogicalInitialJob(cfg config.JobConfig, dockerClient *client.Client, cloneManager thinclones.Manager,
 	global *dblabCfg.Global, marker *dbmarker.Marker) (*LogicalInitial, error) {
 	li := &LogicalInitial{
 		name:         cfg.Name,
@@ -55,6 +74,12 @@ func NewLogicalInitialJob(cfg config.JobConfig, cloneManager thinclones.Manager,
 		return nil, errors.Wrap(err, "failed to unmarshal configuration options")
 	}
 
+	if li.options.DataPatching.QueryPreprocessing.QueryPath != "" {
+		li.queryProcessor = newQueryProcessor(dockerClient, global.Database.Name(), global.Database.User(),
+			li.options.DataPatching.QueryPreprocessing.QueryPath,
+			li.options.DataPatching.QueryPreprocessing.MaxParallelWorkers)
+	}
+
 	return li, nil
 }
 
@@ -63,13 +88,18 @@ func (s *LogicalInitial) Name() string {
 	return s.name
 }
 
+// patchContainerName returns container name.
+func (s *LogicalInitial) patchContainerName() string {
+	return patchContainerPrefix + s.globalCfg.InstanceID
+}
+
 // Reload reloads job configuration.
 func (s *LogicalInitial) Reload(cfg map[string]interface{}) (err error) {
 	return options.Unmarshal(cfg, &s.options)
 }
 
 // Run starts the job.
-func (s *LogicalInitial) Run(_ context.Context) error {
+func (s *LogicalInitial) Run(ctx context.Context) error {
 	if s.options.PreprocessingScript != "" {
 		if err := runPreprocessingScript(s.options.PreprocessingScript); err != nil {
 			return err
@@ -80,16 +110,24 @@ func (s *LogicalInitial) Run(_ context.Context) error {
 		return errors.Wrap(err, "failed to create PostgreSQL configuration files")
 	}
 
+	dataDir := s.globalCfg.DataDir()
+
 	// Run basic PostgreSQL configuration.
-	if err := configuration.NewCorrector().Run(s.globalCfg.DataDir()); err != nil {
+	if err := configuration.NewCorrector().Run(dataDir); err != nil {
 		return errors.Wrap(err, "failed to adjust PostgreSQL configs")
 	}
 
 	// Apply user defined configs.
-	if err := applyUsersConfigs(s.options.Configs, path.Join(s.globalCfg.DataDir(), "postgresql.conf")); err != nil {
+	if err := applyUsersConfigs(s.options.Configs, path.Join(dataDir, "postgresql.conf")); err != nil {
 		return errors.Wrap(err, "failed to apply user-defined configs")
 	}
 
+	if s.queryProcessor != nil {
+		if err := s.runPreprocessingQueries(ctx, dataDir); err != nil {
+			return errors.Wrap(err, "failed to run preprocessing queries")
+		}
+	}
+
 	dataStateAt := extractDataStateAt(s.dbMarker)
 
 	if _, err := s.cloneManager.CreateSnapshot("", dataStateAt); err != nil {
@@ -100,9 +138,115 @@ func (s *LogicalInitial) Run(_ context.Context) error {
 }
 
 func (s *LogicalInitial) touchConfigFiles() error {
-	if err := tools.TouchFile(path.Join(s.globalCfg.DataDir(), "postgresql.conf")); err != nil {
+	dataDir := s.globalCfg.DataDir()
+
+	if err := tools.TouchFile(path.Join(dataDir, "postgresql.conf")); err != nil {
 		return err
 	}
 
-	return tools.TouchFile(path.Join(s.globalCfg.DataDir(), "pg_hba.conf"))
+	return tools.TouchFile(path.Join(dataDir, "pg_hba.conf"))
+}
+
+func (s *LogicalInitial) runPreprocessingQueries(ctx context.Context, dataDir string) error {
+	pgVersion, err := tools.DetectPGVersion(dataDir)
+	if err != nil {
+		return errors.Wrap(err, "failed to detect the Postgres version")
+	}
+
+	patchImage := s.options.DataPatching.DockerImage
+	if patchImage == "" {
+		patchImage = fmt.Sprintf("postgresai/sync-instance:%s", pgVersion)
+	}
+
+	if err := tools.PullImage(ctx, s.dockerClient, patchImage); err != nil {
+		return errors.Wrap(err, "failed to scan image pulling response")
+	}
+
+	pwd, err := tools.GeneratePassword()
+	if err != nil {
+		return errors.Wrap(err, "failed to generate PostgreSQL password")
+	}
+
+	hostConfig, err := s.buildHostConfig(ctx)
+	if err != nil {
+		return errors.Wrap(err, "failed to build container host config")
+	}
+
+	// Run patch container.
+	patchCont, err := s.dockerClient.ContainerCreate(ctx,
+		s.buildContainerConfig(dataDir, patchImage, pwd),
+		hostConfig,
+		&network.NetworkingConfig{},
+		s.patchContainerName(),
+	)
+	if err != nil {
+		return errors.Wrap(err, "failed to create container")
+	}
+
+	defer tools.RemoveContainer(ctx, s.dockerClient, patchCont.ID, cont.StopPhysicalTimeout)
+
+	defer func() {
+		if err != nil {
+			tools.PrintContainerLogs(ctx, s.dockerClient, s.patchContainerName())
+		}
+	}()
+
+	log.Msg(fmt.Sprintf("Running container: %s. ID: %v", s.patchContainerName(), patchCont.ID))
+
+	if err := s.dockerClient.ContainerStart(ctx, patchCont.ID, types.ContainerStartOptions{}); err != nil {
+		return errors.Wrap(err, "failed to start container")
+	}
+
+	log.Msg("Starting PostgreSQL")
+	log.Msg(fmt.Sprintf("View logs using the command: %s %s", tools.ViewLogsCmd, s.patchContainerName()))
+
+	// Start PostgreSQL instance.
+	if err := tools.RunPostgres(ctx, s.dockerClient, patchCont.ID, dataDir); err != nil {
+		return errors.Wrap(err, "failed to start PostgreSQL instance")
+	}
+
+	log.Msg("Waiting for PostgreSQL is ready")
+
+	if err := tools.CheckContainerReadiness(ctx, s.dockerClient, patchCont.ID); err != nil {
+		return errors.Wrap(err, "failed to readiness check")
+	}
+
+	if err := s.queryProcessor.applyPreprocessingQueries(ctx, patchCont.ID); err != nil {
+		return errors.Wrap(err, "failed to run preprocessing queries")
+	}
+
+	return nil
+}
+
+func (s *LogicalInitial) buildHostConfig(ctx context.Context) (*container.HostConfig, error) {
+	hostConfig := &container.HostConfig{}
+
+	if err := tools.AddVolumesToHostConfig(ctx, s.dockerClient, hostConfig, s.globalCfg.DataDir()); err != nil {
+		return nil, err
+	}
+
+	return hostConfig, nil
+}
+
+func (s *LogicalInitial) buildContainerConfig(clonePath, patchImage, password string) *container.Config {
+	hcInterval := health.DefaultRestoreInterval
+	hcRetries := health.DefaultRestoreRetries
+
+	return &container.Config{
+		Labels: map[string]string{
+			cont.DBLabControlLabel:    cont.DBLabPatchLabel,
+			cont.DBLabInstanceIDLabel: s.globalCfg.InstanceID,
+		},
+		Env: []string{
+			"PGDATA=" + clonePath,
+			"POSTGRES_PASSWORD=" + password,
+		},
+		Image: patchImage,
+		Healthcheck: health.GetConfig(
+			s.globalCfg.Database.User(),
+			s.globalCfg.Database.Name(),
+			health.OptionInterval(hcInterval),
+			health.OptionRetries(hcRetries),
+		),
+	}
 }

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

@@ -71,6 +71,7 @@ type PhysicalInitial struct {
 	scheduler      *cron.Cron
 	schedulerCtx   context.Context
 	promotionMutex sync.Mutex
+	queryProcessor *queryProcessor
 }
 
 // PhysicalOptions describes options for a physical initialization job.
@@ -86,9 +87,10 @@ type PhysicalOptions struct {
 
 // Promotion describes promotion options.
 type Promotion struct {
-	Enabled     bool        `yaml:"enabled"`
-	DockerImage string      `yaml:"dockerImage"`
-	HealthCheck HealthCheck `yaml:"healthCheck"`
+	Enabled            bool               `yaml:"enabled"`
+	DockerImage        string             `yaml:"dockerImage"`
+	HealthCheck        HealthCheck        `yaml:"healthCheck"`
+	QueryPreprocessing QueryPreprocessing `yaml:"queryPreprocessing"`
 }
 
 // HealthCheck describes health check options of a promotion.
@@ -109,6 +111,12 @@ type ScheduleSpec struct {
 	Limit     int    `yaml:"limit"`
 }
 
+// QueryPreprocessing defines query preprocessing options.
+type QueryPreprocessing struct {
+	QueryPath          string `yaml:"queryPath"`
+	MaxParallelWorkers int    `yaml:"maxParallelWorkers"`
+}
+
 // NewPhysicalInitialJob creates a new physical initial job.
 func NewPhysicalInitialJob(cfg config.JobConfig, docker *client.Client, cloneManager thinclones.Manager,
 	global *dblabCfg.Global, marker *dbmarker.Marker) (*PhysicalInitial, error) {
@@ -129,6 +137,12 @@ func NewPhysicalInitialJob(cfg config.JobConfig, docker *client.Client, cloneMan
 		return nil, errors.Wrap(err, "invalid physicalSnapshot configuration")
 	}
 
+	if p.options.Promotion.QueryPreprocessing.QueryPath != "" {
+		p.queryProcessor = newQueryProcessor(docker, global.Database.Name(), global.Database.User(),
+			p.options.Promotion.QueryPreprocessing.QueryPath,
+			p.options.Promotion.QueryPreprocessing.MaxParallelWorkers)
+	}
+
 	p.setupScheduler()
 
 	return p, nil
@@ -261,7 +275,7 @@ func (p *PhysicalInitial) run(ctx context.Context) (err error) {
 	cloneName := fmt.Sprintf("clone%s_%s", pre, preDataStateAt)
 
 	defer func() {
-		if _, ok := err.(*skipSnapshotErr); ok {
+		if _, ok := errors.Cause(err).(*skipSnapshotErr); ok {
 			log.Msg(err.Error())
 			err = nil
 		}
@@ -492,6 +506,12 @@ func (p *PhysicalInitial) promoteInstance(ctx context.Context, clonePath string)
 		}
 	}
 
+	if p.queryProcessor != nil {
+		if err := p.queryProcessor.applyPreprocessingQueries(ctx, promoteCont.ID); err != nil {
+			return errors.Wrap(err, "failed to run preprocessing queries")
+		}
+	}
+
 	// Checkpoint.
 	if err := p.checkpoint(ctx, promoteCont.ID); err != nil {
 		return err