(feat) Implement replica identity when creating partitions

It follows the parent table.

Signed-off-by: Daniel Vérité <dverite@gmail.com>

Author: dverite

internal/infra/postgresql/partition.go

@@ -140,3 +140,55 @@ func (p Postgres) GetPartitionSettings(schema, table string) (strategy, key stri
 
 	return strategy, key, nil
 }
+
+func (p Postgres) SetPartitionReplicaIdentity(schema, table, parent string) error {
+	var replIdent string
+
+	// Get the replica identity of the parent
+	query_ri := "SELECT relreplident::text from pg_class c JOIN pg_namespace n ON (n.oid=c.relnamespace) WHERE n.nspname=$1 AND c.relname=$2"
+	err := p.conn.QueryRow(p.ctx, query_ri, schema, parent).Scan(&replIdent)
+	if err != nil {
+		return fmt.Errorf("failed to check pg_class.relreplident for parent table: %w", err)
+	}
+
+	if replIdent == "f" { // replica identity = full
+		query_f := fmt.Sprintf("ALTER TABLE %s.%s REPLICA IDENTITY FULL", schema, table)
+		p.logger.Debug("Set identity full", "query", query_f)
+		_, err = p.conn.Exec(p.ctx, query_f)
+		if err != nil {
+			return fmt.Errorf("failed to set replica identity: %w", err)
+		}
+
+	} else if replIdent == "i" { // replica identity = specific index
+		var indexName string
+		/* This query finds the index that is a child of the (only) index
+		   in the parent table having "indisreplident"=true.
+		   "pg_inherits" holds the (parent-index, child-index) relationship. */
+		query_idx := `
+SELECT c_idx_child.relname
+  FROM pg_index i_parent JOIN pg_inherits inh ON (inh.inhparent=i_parent.indexrelid)
+  JOIN pg_index i_child ON (i_child.indexrelid=inh.inhrelid)
+  JOIN pg_class c_parent ON (c_parent.oid=i_parent.indrelid)
+  JOIN pg_namespace ON (pg_namespace.oid=c_parent.relnamespace)
+  JOIN pg_class c_idx_child ON (c_idx_child.oid=inh.inhrelid)
+  JOIN pg_class c_child ON (c_child.oid=i_child.indrelid)
+ WHERE pg_namespace.nspname= $1
+ AND c_parent.relname = $2
+ AND i_parent.indisreplident=true
+ AND c_child.relname = $3
+`
+		err = p.conn.QueryRow(p.ctx, query_idx, schema, parent, table).Scan(&indexName)
+		if err != nil {
+			return fmt.Errorf("failed to find the child index for the new partition: %w", err)
+		}
+
+		query_alter := fmt.Sprintf("ALTER TABLE %s.%s REPLICA IDENTITY USING INDEX %s", schema, table, indexName)
+		p.logger.Debug("Set replica identity", "schema", schema, "table", table, "index", indexName, "query", query_alter)
+
+		_, err := p.conn.Exec(p.ctx, query_alter)
+		if err != nil {
+			return fmt.Errorf("failed to set replica identity on the new partition: %w", err)
+		}
+	}
+	return nil
+}

internal/infra/postgresql/table.go

@@ -4,7 +4,7 @@ import "fmt"
 
 func (p Postgres) CreateTableLikeTable(schema, table, parent string) error {
 	query := fmt.Sprintf("CREATE TABLE %s.%s (LIKE %s.%s INCLUDING ALL)", schema, table, schema, parent)
-	p.logger.Debug("Create table", "query", schema, "table", table, "query", query)
+	p.logger.Debug("Create table", "schema", schema, "table", table, "query", query)
 
 	_, err := p.conn.Exec(p.ctx, query)
 	if err != nil {

pkg/ppm/mocks/PostgreSQLClient.go

@@ -24,6 +24,7 @@ type PostgreSQLClient interface {
 	DropTable(schema, table string) error
 	DetachPartitionConcurrently(schema, table, parent string) error
 	FinalizePartitionDetach(schema, table, parent string) error
+	SetPartitionReplicaIdentity(schema, table, parent string) error
 }
 
 type PPM struct {

pkg/ppm/ppm.go

@@ -167,6 +167,13 @@ func (p PPM) CreatePartition(partitionConfiguration partition.Configuration, par
 			return fmt.Errorf("fail to attach partition: %w", err)
 		}
 
+		err = p.db.SetPartitionReplicaIdentity(partition.Schema, partition.Name, partition.ParentTable)
+		if err != nil {
+			p.logger.Warn("failed to set replica identity", "error", err, "schema", partition.Schema, "table", partition.Name, "attempt", attempt, "max_retries", maxRetries)
+
+			return fmt.Errorf("fail to set replica identity: %w", err)
+		}
+
 		return nil
 	})
 	if err != nil {

pkg/ppm/provisioning.go

@@ -19,6 +19,7 @@ setup() {
   ppm_setup
 }
 
+
 @test "Test that provisioning succeed on up-to-date partitioning" {
   local TABLE=$(generate_table_name)
   local INTERVAL=daily
@@ -434,5 +435,44 @@ EOF
   assert_output "$expected_2"
 
   rm "$CONFIGURATION_FILE"
+}
+
+
+@test "Test setting Replica Identity on new partitions" {
+  local TABLE=$(generate_table_name)
+  create_partitioned_table ${TABLE}
+
+  local CONFIGURATION=$(basic_configuration "$TABLE" monthly created_at 1 2)
+  local CONFIGURATION_FILE=$(generate_configuration_file "${CONFIGURATION}")
+
+  execute_sql "CREATE UNIQUE INDEX idx_ri_test ON \"${TABLE}\"(created_at, id);"
+  execute_sql "ALTER TABLE \"${TABLE}\" REPLICA IDENTITY USING INDEX idx_ri_test;"
+
+  PPM_WORK_DATE="2026-01-01" run "$PPM_PROG" run provisioning -c ${CONFIGURATION_FILE}
+  assert_success
+  assert_output --partial "All partitions are correctly provisioned"
+
+  # The 4 partitions should all have the replica identity set to "i"
+  local check_query="SELECT array_agg(relreplident ORDER BY relname) FROM pg_class WHERE oid in
+   (SELECT inhrelid from pg_inherits where inhparent='\"$TABLE\"'::regclass);"
+
+  run execute_sql_commands "$check_query"
+  assert_output "{i,i,i,i}"
+
+  # Switch the parent replica identity to "full". Afterwards it must be applied to
+  # newly created partitions
+  execute_sql "ALTER TABLE \"${TABLE}\" REPLICA IDENTITY FULL;"
 
+  # Increment the provisioning by one month and run the test again
+  yq eval ".partitions.unittest.preProvisioned = 3" -i ${CONFIGURATION_FILE}
+  PPM_WORK_DATE="2026-01-01" run "$PPM_PROG" run provisioning -c ${CONFIGURATION_FILE}
+  assert_success
+  assert_output --partial "All partitions are correctly provisioned"
+
+  # Now the most recent partition should have its replica identity to full
+  # and the others to indexes
+  run execute_sql_commands "$check_query"
+  assert_output "{i,i,i,i,f}"
+
+  rm "$CONFIGURATION_FILE"
 }

scripts/bats/30_provisioning.bats

@@ -28,7 +28,7 @@ create_table_from_template() {
 
   read -r -d '' QUERY <<EOQ ||
   CREATE TABLE ${TABLE} (
-    id              BIGSERIAL,
+    id              BIGSERIAL NOT NULL,
     temperature     INT,
     created_at      DATE NOT NULL
   ) PARTITION BY RANGE (created_at);