BED-4194: PG Migrator Testing

cmd/api/src/api/tools/pg.go

@@ -19,6 +19,9 @@ package tools
 import (
  "context"
  "fmt"
+ "net/http"
+ "sync"
+
  "github.com/neo4j/neo4j-go-driver/v5/neo4j/dbtype"
  "github.com/specterops/bloodhound/dawgs"
  "github.com/specterops/bloodhound/dawgs/drivers/neo4j"
@@ -28,16 +31,14 @@ import (
  "github.com/specterops/bloodhound/log"
  "github.com/specterops/bloodhound/src/api"
  "github.com/specterops/bloodhound/src/config"
- "net/http"
- "sync"
 )
 
 type MigratorState string
 
 const (
- stateIdle MigratorState = "idle"
- stateMigrating MigratorState = "migrating"
- stateCanceling MigratorState = "canceling"
+ StateIdle MigratorState = "idle"
+ StateMigrating MigratorState = "migrating"
+ StateCanceling MigratorState = "canceling"
 )
 
 func migrateTypes(ctx context.Context, neoDB, pgDB graph.Database) error {
@@ -188,21 +189,21 @@ func migrateEdges(ctx context.Context, neoDB, pgDB graph.Database, nodeIDMapping
 type PGMigrator struct {
  graphSchema graph.Schema
  graphDBSwitch *graph.DatabaseSwitch
- serverCtx context.Context
+ ServerCtx context.Context
  migrationCancelFunc func()
- state MigratorState
+ State MigratorState
  lock *sync.Mutex
- cfg config.Configuration
+ Cfg config.Configuration
 }
 
 func NewPGMigrator(serverCtx context.Context, cfg config.Configuration, graphSchema graph.Schema, graphDBSwitch *graph.DatabaseSwitch) *PGMigrator {
  return &PGMigrator{
  graphSchema: graphSchema,
  graphDBSwitch: graphDBSwitch,
- serverCtx: serverCtx,
- state: stateIdle,
+ ServerCtx: serverCtx,
+ State: StateIdle,
  lock: &sync.Mutex{},
- cfg: cfg,
+ Cfg: cfg,
  }
 }
 
@@ -213,29 +214,26 @@ func (s *PGMigrator) advanceState(next MigratorState, validTransitions ...Migrat
  isValid := false
 
  for _, validTransition := range validTransitions {
- if s.state == validTransition {
+ if s.State == validTransition {
  isValid = true
  break
  }
  }
 
  if !isValid {
- return fmt.Errorf("migrator state is %s but expected one of: %v", s.state, validTransitions)
+ return fmt.Errorf("migrator state is %s but expected one of: %v", s.State, validTransitions)
  }
 
- s.state = next
+ s.State = next
  return nil
 }
 
 func (s *PGMigrator) SwitchPostgreSQL(response http.ResponseWriter, request *http.Request) {
- if pgDB, err := dawgs.Open(s.serverCtx, pg.DriverName, dawgs.Config{
- TraversalMemoryLimit: size.Gibibyte,
- DriverCfg: s.cfg.Database.PostgreSQLConnectionString(),
- }); err != nil {
+ if pgDB, err := s.OpenPostgresGraphConnection(); err != nil {
  api.WriteJSONResponse(request.Context(), map[string]any{
  "error": fmt.Errorf("failed connecting to PostgreSQL: %w", err),
  }, http.StatusInternalServerError, response)
- } else if err := SetGraphDriver(request.Context(), s.cfg, pg.DriverName); err != nil {
+ } else if err := SetGraphDriver(request.Context(), s.Cfg, pg.DriverName); err != nil {
  api.WriteJSONResponse(request.Context(), map[string]any{
  "error": fmt.Errorf("failed updating graph database driver preferences: %w", err),
  }, http.StatusInternalServerError, response)
@@ -248,14 +246,11 @@ func (s *PGMigrator) SwitchPostgreSQL(response http.ResponseWriter, request *htt
 }
 
 func (s *PGMigrator) SwitchNeo4j(response http.ResponseWriter, request *http.Request) {
- if neo4jDB, err := dawgs.Open(s.serverCtx, neo4j.DriverName, dawgs.Config{
- TraversalMemoryLimit: size.Gibibyte,
- DriverCfg: s.cfg.Neo4J.Neo4jConnectionString(),
- }); err != nil {
+ if neo4jDB, err := s.OpenNeo4jGraphConnection(); err != nil {
  api.WriteJSONResponse(request.Context(), map[string]any{
  "error": fmt.Errorf("failed connecting to Neo4j: %w", err),
  }, http.StatusInternalServerError, response)
- } else if err := SetGraphDriver(request.Context(), s.cfg, neo4j.DriverName); err != nil {
+ } else if err := SetGraphDriver(request.Context(), s.Cfg, neo4j.DriverName); err != nil {
  api.WriteJSONResponse(request.Context(), map[string]any{
  "error": fmt.Errorf("failed updating graph database driver preferences: %w", err),
  }, http.StatusInternalServerError, response)
@@ -267,23 +262,17 @@ func (s *PGMigrator) SwitchNeo4j(response http.ResponseWriter, request *http.Req
  }
 }
 
-func (s *PGMigrator) startMigration() error {
- if err := s.advanceState(stateMigrating, stateIdle); err != nil {
+func (s *PGMigrator) StartMigration() error {
+ if err := s.advanceState(StateMigrating, StateIdle); err != nil {
  return fmt.Errorf("database migration state error: %w", err)
- } else if neo4jDB, err := dawgs.Open(s.serverCtx, neo4j.DriverName, dawgs.Config{
- TraversalMemoryLimit: size.Gibibyte,
- DriverCfg: s.cfg.Neo4J.Neo4jConnectionString(),
- }); err != nil {
+ } else if neo4jDB, err := s.OpenNeo4jGraphConnection(); err != nil {
  return fmt.Errorf("failed connecting to Neo4j: %w", err)
- } else if pgDB, err := dawgs.Open(s.serverCtx, pg.DriverName, dawgs.Config{
- TraversalMemoryLimit: size.Gibibyte,
- DriverCfg: s.cfg.Database.PostgreSQLConnectionString(),
- }); err != nil {
+ } else if pgDB, err := s.OpenPostgresGraphConnection(); err != nil {
  return fmt.Errorf("failed connecting to PostgreSQL: %w", err)
  } else {
  log.Infof("Dispatching live migration from Neo4j to PostgreSQL")
 
- migrationCtx, migrationCancelFunc := context.WithCancel(s.serverCtx)
+ migrationCtx, migrationCancelFunc := context.WithCancel(s.ServerCtx)
  s.migrationCancelFunc = migrationCancelFunc
 
  go func(ctx context.Context) {
@@ -303,7 +292,7 @@ func (s *PGMigrator) startMigration() error {
  log.Infof("Migration to PostgreSQL completed successfully")
  }
 
- if err := s.advanceState(stateIdle, stateMigrating, stateCanceling); err != nil {
+ if err := s.advanceState(StateIdle, StateMigrating, StateCanceling); err != nil {
  log.Errorf("Database migration state management error: %v", err)
  }
  }(migrationCtx)
@@ -313,7 +302,7 @@ func (s *PGMigrator) startMigration() error {
 }
 
 func (s *PGMigrator) MigrationStart(response http.ResponseWriter, request *http.Request) {
- if err := s.startMigration(); err != nil {
+ if err := s.StartMigration(); err != nil {
  api.WriteJSONResponse(request.Context(), map[string]any{
  "error": err.Error(),
  }, http.StatusInternalServerError, response)
@@ -322,8 +311,8 @@ func (s *PGMigrator) MigrationStart(response http.ResponseWriter, request *http.
  }
 }
 
-func (s *PGMigrator) cancelMigration() error {
- if err := s.advanceState(stateCanceling, stateMigrating); err != nil {
+func (s *PGMigrator) CancelMigration() error {
+ if err := s.advanceState(StateCanceling, StateMigrating); err != nil {
  return err
  }
 
@@ -333,7 +322,7 @@ func (s *PGMigrator) cancelMigration() error {
 }
 
 func (s *PGMigrator) MigrationCancel(response http.ResponseWriter, request *http.Request) {
- if err := s.cancelMigration(); err != nil {
+ if err := s.CancelMigration(); err != nil {
  api.WriteJSONResponse(request.Context(), map[string]any{
  "error": err.Error(),
  }, http.StatusInternalServerError, response)
@@ -344,6 +333,20 @@ func (s *PGMigrator) MigrationCancel(response http.ResponseWriter, request *http
 
 func (s *PGMigrator) MigrationStatus(response http.ResponseWriter, request *http.Request) {
  api.WriteJSONResponse(request.Context(), map[string]any{
- "state": s.state,
+ "state": s.State,
  }, http.StatusOK, response)
 }
+
+func (s *PGMigrator) OpenPostgresGraphConnection() (graph.Database, error) {
+ return dawgs.Open(s.ServerCtx, pg.DriverName, dawgs.Config{
+ TraversalMemoryLimit: size.Gibibyte,
+ DriverCfg: s.Cfg.Database.PostgreSQLConnectionString(),
+ })
+}
+
+func (s *PGMigrator) OpenNeo4jGraphConnection() (graph.Database, error) {
+ return dawgs.Open(s.ServerCtx, neo4j.DriverName, dawgs.Config{
+ TraversalMemoryLimit: size.Gibibyte,
+ DriverCfg: s.Cfg.Neo4J.Neo4jConnectionString(),
+ })
+}