feat(pg-pg): support ingestion to target tables having FK relationships (#4219)

> [!NOTE]
This PR is contained to the PG type system only.

> [!IMPORTANT]
PeerDB's INSERTs will no longer trigger regular BEFORE/AFTER triggers on
top of destination tables after this change for PG type system mirrors.
If you were relying on this behaviour, reach out to @Amogh-Bharadwaj on
[Slack](https://slack.peerdb.io/)

### Context
For Postgres to Postgres apples-to-apples migrations, users would often
pre-create their destination tables matching their source table schemas,
including indexes, primary keys and foreign keys.

### Problem statement
Currently, PeerDB's ingestion to target tables can be blocked by foreign
key violations depending on the order of the tables it iterates though.
Ideally, PeerDB's ingestion should be allowed to these tables for the
purpose of migration.

### This PR
This PR sets
[`session_replication_role`](https://www.postgresql.org/docs/17/runtime-config-client.html#GUC-SESSION-REPLICATION-ROLE)
to `replica` within the scope of the session in which inserts are run.

- **Initial load**: Sets `SET LOCAL session_replication_role =
'replica'` on the dedicated per-partition connection in
`syncQRepRecords` before the transaction begins (PG type system only).

- **CDC**: Sets `SET LOCAL session_replication_role = 'replica'` inside
the normalize transaction in `normalizeBatch` (PG type system only).

### Testing
- An end to end test has been added, and verified that without this PR
it fails.

Author: Amogh-Bharadwaj

flow/connectors/postgres/client.go

@@ -48,6 +48,7 @@ const (
 		sync_batch_id=EXCLUDED.sync_batch_id`
 	checkIfJobMetadataExistsSQL          = "SELECT EXISTS(SELECT * FROM %s.%s WHERE mirror_job_name=$1)"
 	updateMetadataForNormalizeRecordsSQL = "UPDATE %s.%s SET normalize_batch_id=$1 WHERE mirror_job_name=$2"
+	setSessionReplicaRoleSQL             = "SET LOCAL session_replication_role = 'replica'"
 
 	getDistinctDestinationTableNamesSQL = `SELECT DISTINCT _peerdb_destination_table_name FROM %s.%s WHERE
 	_peerdb_batch_id=$1`

flow/connectors/postgres/postgres.go

@@ -786,6 +786,16 @@ func (c *PostgresConnector) normalizeBatch(
 	}
 	defer shared.RollbackTx(tx, c.logger)
 
+	for _, tableName := range destinationTableNames {
+		if schema, ok := normalizeStmtGen.tableSchemaMapping[tableName]; ok && schema.System == protos.TypeSystem_PG {
+			if _, err := tx.Exec(ctx, setSessionReplicaRoleSQL); err != nil {
+				return 0, fmt.Errorf("failed to set session_replication_role to replica: %w", err)
+			}
+			c.logger.Info("set session_replication_role to replica for PG type system normalize")
+			break
+		}
+	}
+
 	batch := &pgx.Batch{}
 	var entries []batchEntry
 	for _, destinationTableName := range destinationTableNames {

flow/connectors/postgres/qrep.go

@@ -620,6 +620,13 @@ func syncQRepRecords(
 	}
 	defer shared.RollbackTx(tx, c.logger)
 
+	if config.System == protos.TypeSystem_PG {
+		if _, err := tx.Exec(ctx, setSessionReplicaRoleSQL); err != nil {
+			return 0, nil, fmt.Errorf("failed to set session_replication_role to replica: %w", err)
+		}
+		c.logger.Info("set session_replication_role to replica on destination for PG type system initial load")
+	}
+
 	// Step 2: Insert records into destination table
 	var numRowsSynced int64
 

flow/e2e/postgres_test.go

@@ -1419,3 +1419,106 @@ func (s PeerFlowE2ETestSuitePG) Test_Table_With_Excluded_PK_And_ReplicaIdentityF
 	env.Cancel(s.t.Context())
 	RequireEnvCanceled(s.t, env)
 }
+
+func (s PeerFlowE2ETestSuitePG) Test_PG_PG_Target_Foreign_Keys() {
+	tc := NewTemporalClient(s.t)
+
+	srcTableName1 := s.attachSchemaSuffix("test_pg_pg_fk_src_1")
+	dstTableName1 := s.attachSchemaSuffix("test_pg_pg_fk_dst_1")
+	srcTableName2 := s.attachSchemaSuffix("test_pg_pg_fk_src_2")
+	dstTableName2 := s.attachSchemaSuffix("test_pg_pg_fk_dst_2")
+
+	// Source parent table
+	_, err := s.Conn().Exec(s.t.Context(), fmt.Sprintf(`
+		CREATE TABLE IF NOT EXISTS %s (
+			id SERIAL PRIMARY KEY,
+			val TEXT NOT NULL
+		)`, srcTableName1))
+	require.NoError(s.t, err)
+
+	// Source child table with FK referencing parent
+	_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(`
+		CREATE TABLE IF NOT EXISTS %s (
+			id SERIAL PRIMARY KEY,
+			parent_id BIGINT NOT NULL REFERENCES %s(id),
+			val TEXT NOT NULL
+		)`, srcTableName2, srcTableName1))
+	require.NoError(s.t, err)
+
+	// Destination parent table
+	_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(`
+		CREATE TABLE IF NOT EXISTS %s (
+			id SERIAL PRIMARY KEY,
+			val TEXT NOT NULL
+		)`, dstTableName1))
+	require.NoError(s.t, err)
+
+	// Destination child table with FK referencing destination parent
+	_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(`
+		CREATE TABLE IF NOT EXISTS %s (
+			id SERIAL PRIMARY KEY,
+			parent_id BIGINT NOT NULL REFERENCES %s(id),
+			val TEXT NOT NULL
+		)`, dstTableName2, dstTableName1))
+	require.NoError(s.t, err)
+
+	// Insert rows only into the child table, bypassing FK checks via session_replication_role
+	rowsTx, err := s.Conn().Begin(s.t.Context())
+	require.NoError(s.t, err)
+	_, err = rowsTx.Exec(s.t.Context(), "SET LOCAL session_replication_role = 'replica'")
+	require.NoError(s.t, err)
+	for i := range 5 {
+		_, err = rowsTx.Exec(s.t.Context(), fmt.Sprintf(
+			`INSERT INTO %s(parent_id, val) VALUES ($1, $2)`, srcTableName2), int64(i+100), fmt.Sprintf("child_val_%d", i))
+		require.NoError(s.t, err)
+	}
+	require.NoError(s.t, rowsTx.Commit(s.t.Context()))
+
+	config := &protos.FlowConnectionConfigs{
+		FlowJobName:     s.attachSuffix("test_pg_pg_fk"),
+		DestinationName: s.Peer().Name,
+		TableMappings: []*protos.TableMapping{
+			{
+				SourceTableIdentifier:      srcTableName1,
+				DestinationTableIdentifier: dstTableName1,
+			},
+			{
+				SourceTableIdentifier:      srcTableName2,
+				DestinationTableIdentifier: dstTableName2,
+			},
+		},
+		SourceName:        GeneratePostgresPeer(s.t).Name,
+		MaxBatchSize:      100,
+		DoInitialSnapshot: true,
+		System:            protos.TypeSystem_PG,
+		SoftDeleteColName: "",
+		SyncedAtColName:   "",
+	}
+
+	env := ExecutePeerflow(s.t, tc, config)
+	SetupCDCFlowStatusQuery(s.t, env, config)
+
+	// Verify initial load replicated the child table rows despite FK constraint on destination
+	EnvWaitFor(s.t, env, 3*time.Minute, "wait for child table initial load", func() bool {
+		return s.comparePGTables(srcTableName2, dstTableName2, "id,parent_id,val") == nil
+	})
+
+	// Insert more rows via CDC path, again only into child table bypassing FK
+	cdcTx, err := s.Conn().Begin(s.t.Context())
+	EnvNoError(s.t, env, err)
+	_, err = cdcTx.Exec(s.t.Context(), "SET LOCAL session_replication_role = 'replica'")
+	EnvNoError(s.t, env, err)
+	for i := range 5 {
+		_, err = cdcTx.Exec(s.t.Context(), fmt.Sprintf(
+			`INSERT INTO %s(parent_id, val) VALUES ($1, $2)`, srcTableName2), int64(i+200), fmt.Sprintf("cdc_child_val_%d", i))
+		EnvNoError(s.t, env, err)
+	}
+	EnvNoError(s.t, env, cdcTx.Commit(s.t.Context()))
+
+	EnvWaitFor(s.t, env, 3*time.Minute, "wait for child table cdc", func() bool {
+		return s.comparePGTables(srcTableName2, dstTableName2, "id,parent_id,val") == nil
+	})
+
+	env.Cancel(s.t.Context())
+	RequireEnvCanceled(s.t, env)
+}