[*] move prefix to EtcdClient

Author: pashagolub

cmd/etcd_fdw/main.go

@@ -90,6 +90,19 @@ func SetupLogging(logLevel string) error {
 	return nil
 }
 
+// SetupCloseHandler creates a 'listener' on a new goroutine which will notify the
+// program if it receives an interrupt from the OS. We then handle this by calling
+// our clean up procedure and exiting the program.
+func SetupCloseHandler(cancel context.CancelFunc) {
+	c := make(chan os.Signal, 2)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-c
+		logrus.Debug("SetupCloseHandler received an interrupt from OS. Closing session...")
+		cancel()
+	}()
+}
+
 func main() {
 	// Quick check for version flags before full parsing
 	for _, arg := range os.Args[1:] {
@@ -117,16 +130,7 @@ func main() {
 	// Setup graceful shutdown
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-
-	// Setup signal handling for graceful shutdown
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
-
-	go func() {
-		sig := <-sigChan
-		logrus.WithField("signal", sig).Info("Received shutdown signal, initiating graceful shutdown...")
-		cancel()
-	}()
+	SetupCloseHandler(cancel)
 
 	// Connect to PostgreSQL with retry logic
 	pgPool, err := sync.NewWithRetry(ctx, config.PostgresDSN)
@@ -136,18 +140,11 @@ func main() {
 	defer pgPool.Close()
 
 	// Connect to etcd with retry logic
-	var etcdClient *sync.EtcdClient
-	if config.EtcdDSN != "" {
-		var err error
-		etcdClient, err = sync.NewEtcdClientWithRetry(ctx, config.EtcdDSN)
-		if err != nil {
-			logrus.WithError(err).Fatal("Failed to connect to etcd after retries")
-		}
-		defer etcdClient.Close()
+	etcdClient, err := sync.NewEtcdClientWithRetry(ctx, config.EtcdDSN)
+	if err != nil {
+		logrus.WithError(err).Fatal("Failed to connect to etcd after retries")
 	}
-
-	// Get prefix from etcd DSN
-	prefix := sync.GetPrefix(config.EtcdDSN)
+	defer etcdClient.Close()
 
 	// Parse polling interval
 	pollingInterval, err := time.ParseDuration(config.PollingInterval)
@@ -156,7 +153,7 @@ func main() {
 	}
 
 	// Create and start sync service
-	syncService := sync.NewService(pgPool, etcdClient, prefix, pollingInterval)
+	syncService := sync.NewService(pgPool, etcdClient, pollingInterval)
 	if err := syncService.Start(ctx); err != nil && ctx.Err() == nil {
 		logrus.WithError(err).Fatal("Synchronization failed")
 	}

internal/sync/etcd.go

@@ -16,6 +16,7 @@ import (
 // EtcdClient handles all etcd operations for PostgreSQL synchronization
 type EtcdClient struct {
 	*clientv3.Client
+	prefix string
 }
 
 // NewEtcdClient creates a new etcd client with DSN parsing
@@ -34,6 +35,7 @@ func NewEtcdClient(dsn string) (*EtcdClient, error) {
 
 	return &EtcdClient{
 		Client: client,
+		prefix: getPrefix(dsn),
 	}, nil
 }
 
@@ -46,15 +48,15 @@ func (c *EtcdClient) Close() error {
 }
 
 // WatchPrefix sets up a watch for all keys with the given prefix
-func (c *EtcdClient) WatchPrefix(ctx context.Context, prefix string, startRevision int64) clientv3.WatchChan {
+func (c *EtcdClient) WatchPrefix(ctx context.Context, startRevision int64) clientv3.WatchChan {
 	opts := []clientv3.OpOption{clientv3.WithPrefix()}
 	if startRevision > 0 {
 		opts = append(opts, clientv3.WithRev(startRevision+1))
 	}
 
-	watchChan := c.Client.Watch(ctx, prefix, opts...)
+	watchChan := c.Client.Watch(ctx, c.prefix, opts...)
 	logrus.WithFields(logrus.Fields{
-		"prefix":   prefix,
+		"prefix":   c.prefix,
 		"revision": startRevision,
 	}).Info("Started etcd watch")
 
@@ -173,7 +175,7 @@ func NewEtcdClientWithRetry(ctx context.Context, dsn string) (*EtcdClient, error
 }
 
 // WatchWithRecovery wraps the etcd watch functionality with automatic recovery
-func (c *EtcdClient) WatchWithRecovery(ctx context.Context, prefix string, startRevision int64) <-chan clientv3.WatchResponse {
+func (c *EtcdClient) WatchWithRecovery(ctx context.Context, startRevision int64) <-chan clientv3.WatchResponse {
 	watchChan := make(chan clientv3.WatchResponse)
 
 	go func() {
@@ -187,7 +189,7 @@ func (c *EtcdClient) WatchWithRecovery(ctx context.Context, prefix string, start
 				return
 			default:
 				// Attempt to establish watch
-				innerWatchChan := c.WatchPrefix(ctx, prefix, currentRevision)
+				innerWatchChan := c.WatchPrefix(ctx, currentRevision)
 
 				for {
 					select {
@@ -249,7 +251,7 @@ func RetryEtcdOperation(ctx context.Context, operation func() error, operationNa
 // parseEtcdDSN parses etcd DSN format: etcd://[user:password@]host1:port1[,host2:port2]/[prefix]?param=value
 func parseEtcdDSN(dsn string) (*clientv3.Config, error) {
 	if dsn == "" {
-		return nil, fmt.Errorf("etcd DSN is required")
+		return &clientv3.Config{}, nil // Default config
 	}
 
 	// Parse the DSN if provided
@@ -321,8 +323,8 @@ func parseEtcdDSN(dsn string) (*clientv3.Config, error) {
 	return config, nil
 }
 
-// GetPrefix extracts the prefix from the etcd DSN path
-func GetPrefix(dsn string) string {
+// getPrefix extracts the prefix from the etcd DSN path
+func getPrefix(dsn string) string {
 	if dsn == "" || !strings.HasPrefix(dsn, "etcd://") {
 		return "/"
 	}

internal/sync/sync.go

@@ -21,11 +21,10 @@ type Service struct {
 }
 
 // NewService creates a new synchronization service
-func NewService(pgPool PgxIface, etcdClient *EtcdClient, prefix string, pollingInterval time.Duration) *Service {
+func NewService(pgPool PgxIface, etcdClient *EtcdClient, pollingInterval time.Duration) *Service {
 	return &Service{
 		pgPool:          pgPool,
 		etcdClient:      etcdClient,
-		prefix:          prefix,
 		pollingInterval: pollingInterval,
 	}
 }
@@ -109,7 +108,7 @@ func (s *Service) syncEtcdToPostgreSQL(ctx context.Context) error {
 	}
 
 	// Start watching from the next revision with automatic recovery
-	watchChan := s.etcdClient.WatchWithRecovery(ctx, s.prefix, latestRevision)
+	watchChan := s.etcdClient.WatchWithRecovery(ctx, latestRevision)
 
 	for {
 		select {