Merge pull request #91 from xataio/allow-to-provide-replication-slot-…

…name-on-cli

Allow providing replication slot on the CLI

README.md

@@ -85,6 +85,12 @@ This will create the `pgstream` schema in the configured Postgres database, alon
 pgstream init --pgurl "postgres://postgres:postgres@localhost?sslmode=disable"
 ```
 
+If you want to provide the name of the replication slot to be created instead of using the default value (`pgstream_<dbname>_slot`), you can use the `--replication-slot` flag or set the environment variable `PGSTREAM_POSTGRES_REPLICATION_SLOT_NAME`.
+
+```
+pgstream init --pgurl "postgres://postgres:postgres@localhost?sslmode=disable" --replication-slot test
+```
+
 If there are any issues or if you want to clean up the pgstream setup, you can run the following.
 
 ```
@@ -121,9 +127,10 @@ Here's a list of all the environment variables that can be used to configure the
 <details>
   <summary>Postgres Listener</summary>
 
-| Environment Variable           | Default | Required | Description                                                          |
-| ------------------------------ | ------- | -------- | -------------------------------------------------------------------- |
-| PGSTREAM_POSTGRES_LISTENER_URL | N/A     | Yes      | URL of the Postgres database to connect to for replication purposes. |
+| Environment Variable                    | Default                    | Required | Description                                                          |
+| --------------------------------------- | -------------------------- | -------- | -------------------------------------------------------------------- |
+| PGSTREAM_POSTGRES_LISTENER_URL          | N/A                        | Yes      | URL of the Postgres database to connect to for replication purposes. |
+| PGSTREAM_POSTGRES_REPLICATION_SLOT_NAME | "pgstream\_<dbname>\_slot" | No       | Name of the Postgres replication slot name.                          |
 
 </details>
 

cmd/config.go

@@ -41,6 +41,14 @@ func pgURL() string {
 	return viper.GetString("PGSTREAM_POSTGRES_LISTENER_URL")
 }
 
+func replicationSlotName() string {
+	replicationslot := viper.GetString("replication-slot")
+	if replicationslot != "" {
+		return replicationslot
+	}
+	return viper.GetString("PGSTREAM_POSTGRES_REPLICATION_SLOT_NAME")
+}
+
 func parseStreamConfig() *stream.Config {
 	return &stream.Config{
 		Listener:  parseListenerConfig(),
@@ -65,7 +73,8 @@ func parsePostgresListenerConfig() *stream.PostgresListenerConfig {
 
 	return &stream.PostgresListenerConfig{
 		Replication: pgreplication.Config{
-			PostgresURL: pgURL,
+			PostgresURL:         pgURL,
+			ReplicationSlotName: replicationSlotName(),
 		},
 	}
 }

cmd/init_cmd.go

@@ -18,7 +18,7 @@ var initCmd = &cobra.Command{
 	RunE: func(cmd *cobra.Command, args []string) error {
 		sp, _ := pterm.DefaultSpinner.WithText("initialising pgstream...").Start()
 
-		if err := stream.Init(context.Background(), pgURL()); err != nil {
+		if err := stream.Init(context.Background(), pgURL(), replicationSlotName()); err != nil {
 			sp.Fail(err.Error())
 			return err
 		}
@@ -34,7 +34,7 @@ var tearDownCmd = &cobra.Command{
 	RunE: func(cmd *cobra.Command, args []string) error {
 		sp, _ := pterm.DefaultSpinner.WithText("tearing down pgstream...").Start()
 
-		if err := stream.TearDown(context.Background(), pgURL()); err != nil {
+		if err := stream.TearDown(context.Background(), pgURL(), replicationSlotName()); err != nil {
 			sp.Fail(err.Error())
 			return err
 		}

cmd/root_cmd.go

@@ -22,10 +22,12 @@ func init() {
 	viper.AutomaticEnv()
 
 	rootCmd.PersistentFlags().String("pgurl", "postgres://postgres:postgres@localhost?sslmode=disable", "Postgres URL")
+	rootCmd.PersistentFlags().String("replication-slot", "", "Name of the postgres replication slot to be created")
 	rootCmd.PersistentFlags().StringP("config", "c", "", ".env config file to use if any")
 	rootCmd.PersistentFlags().String("log-level", "debug", "log level for the application")
 
 	viper.BindPFlag("pgurl", rootCmd.PersistentFlags().Lookup("pgurl"))
+	viper.BindPFlag("replication-slot", rootCmd.PersistentFlags().Lookup("replication-slot"))
 	viper.BindPFlag("config", rootCmd.PersistentFlags().Lookup("config"))
 	viper.BindPFlag("PGSTREAM_LOG_LEVEL", rootCmd.PersistentFlags().Lookup("log-level"))
 }

internal/postgres/pg_replication_conn.go

@@ -120,6 +120,10 @@ func (c *ReplicationConn) Close(ctx context.Context) error {
 	return mapError(c.conn.Close(ctx))
 }
 
+func DefaultReplicationSlotName(dbName string) string {
+	return "pgstream_" + dbName + "_slot"
+}
+
 type Error struct {
 	Severity string
 	Msg      string

pkg/stream/integration/setup_test.go

@@ -30,7 +30,7 @@ func TestMain(m *testing.M) {
 		}
 		defer pgcleanup()
 
-		if err := stream.Init(ctx, pgurl); err != nil {
+		if err := stream.Init(ctx, pgurl, ""); err != nil {
 			log.Fatal(err)
 		}
 

pkg/stream/stream_init.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 
+	pglib "github.com/xataio/pgstream/internal/postgres"
 	pgmigrations "github.com/xataio/pgstream/migrations/postgres"
 
 	"github.com/golang-migrate/migrate/v4"
@@ -23,7 +24,7 @@ const (
 
 // Init initialises the pgstream state in the postgres database provided, along
 // with creating the relevant replication slot.
-func Init(ctx context.Context, pgURL string) error {
+func Init(ctx context.Context, pgURL, replicationSlotName string) error {
 	conn, err := newPGConn(ctx, pgURL)
 	if err != nil {
 		return err
@@ -45,9 +46,11 @@ func Init(ctx context.Context, pgURL string) error {
 		return fmt.Errorf("failed to run internal pgstream migrations: %w", err)
 	}
 
-	replicationSlotName, err := getReplicationSlotName(pgURL)
-	if err != nil {
-		return err
+	if replicationSlotName == "" {
+		replicationSlotName, err = getReplicationSlotName(pgURL)
+		if err != nil {
+			return err
+		}
 	}
 
 	if err := createReplicationSlot(ctx, conn, replicationSlotName); err != nil {
@@ -59,16 +62,18 @@ func Init(ctx context.Context, pgURL string) error {
 
 // TearDown removes the pgstream state from the postgres database provided,
 // as well as removing the replication slot.
-func TearDown(ctx context.Context, pgURL string) error {
+func TearDown(ctx context.Context, pgURL, replicationSlotName string) error {
 	conn, err := newPGConn(ctx, pgURL)
 	if err != nil {
 		return err
 	}
 	defer conn.Close(ctx)
 
-	replicationSlotName, err := getReplicationSlotName(pgURL)
-	if err != nil {
-		return err
+	if replicationSlotName == "" {
+		replicationSlotName, err = getReplicationSlotName(pgURL)
+		if err != nil {
+			return err
+		}
 	}
 
 	if err := dropReplicationSlot(ctx, conn, replicationSlotName); err != nil {
@@ -162,5 +167,5 @@ func getReplicationSlotName(pgURL string) (string, error) {
 	if cfg.Database != "" {
 		dbName = cfg.Database
 	}
-	return fmt.Sprintf("pgstream_%s_slot", dbName), nil
+	return pglib.DefaultReplicationSlotName(dbName), nil
 }

pkg/wal/replication/postgres/pg_replication_handler.go

@@ -101,7 +101,7 @@ func (h *Handler) StartReplication(ctx context.Context) error {
 	}
 
 	if h.pgReplicationSlotName == "" {
-		h.pgReplicationSlotName = fmt.Sprintf("pgstream_%s_slot", sysID.DBName)
+		h.pgReplicationSlotName = pglib.DefaultReplicationSlotName(sysID.DBName)
 	}
 
 	logFields := loglib.Fields{