fix(ctid-load): fix for mixed case partitioned tables (#4234)

### Why
When the CTID based partitioning initial load setting is enabled,
initial loads for mixed case table names do not sync data to target,
because the input we pass to `to_regclass` in queries was not quoted,
and Postgres assumes lowercase, making those queries not return anything

### What
- Fix `to_regclass` usage in estimated table size function
- Make `classifyTableName` return OID of the table as well
- Replace usage of `to_regclass` with comparing with OID directly
- Adapt E2E test to cover mixed case usage

Fixes: DBI-718

Author: Amogh-Bharadwaj

flow/connectors/postgres/postgres.go

@@ -360,9 +360,13 @@ func (c *PostgresConnector) GetDatabaseVariant(ctx context.Context) (protos.Data
 }
 
 func (c *PostgresConnector) GetTableSizeEstimatedBytes(ctx context.Context, tableIdentifier string) (int64, error) {
+	parsed, err := common.ParseTableIdentifier(tableIdentifier)
+	if err != nil {
+		return 0, fmt.Errorf("failed to parse table identifier %s: %w", tableIdentifier, err)
+	}
 	tableSizeQuery := "SELECT pg_relation_size(to_regclass($1))"
 	var tableSizeBytes pgtype.Int8
-	if err := c.conn.QueryRow(ctx, tableSizeQuery, tableIdentifier).Scan(&tableSizeBytes); err != nil {
+	if err := c.conn.QueryRow(ctx, tableSizeQuery, parsed.String()).Scan(&tableSizeBytes); err != nil {
 		return 0, err
 	}
 	if !tableSizeBytes.Valid {

flow/connectors/postgres/qrep_partition.go

@@ -130,13 +130,13 @@ func CTIDBlockPartitioningFunc(ctx context.Context, pp PartitionParams) ([]*prot
 	}
 	switch {
 	case tc.relkind == "p":
-		blocksPerTable, err := getPartitionedTables(ctx, pp.tx, tc.qualifiedName)
+		blocksPerTable, err := getPartitionedTables(ctx, pp.tx, tc.oid)
 		if err != nil {
 			return nil, fmt.Errorf("failed to get child partitioned tables: %w", err)
 		}
 		return ctidPartitionsForChildTables(pp, blocksPerTable)
 	case tc.relhassubclass:
-		blocksPerTable, err := getInheritedTables(ctx, pp.tx, tc.qualifiedName)
+		blocksPerTable, err := getInheritedTables(ctx, pp.tx, tc.oid, tc.qualifiedName)
 		if err != nil {
 			return nil, fmt.Errorf("failed to get child inherited tables: %w", err)
 		}
@@ -305,19 +305,22 @@ func getTableBlockCount(ctx context.Context, tx pgx.Tx, table string) (int64, er
 type tableClassification struct {
 	qualifiedName  string
 	relkind        string
+	oid            uint32
 	relhassubclass bool
 }
 
 func classifyTable(ctx context.Context, tx pgx.Tx, table string) (tableClassification, error) {
 	var classification tableClassification
 	err := tx.QueryRow(ctx,
 		// We fetch the qualified name from pg_class instead of using the input directly because
-		// the input table name is quoted; and we want the unquoted canonical form for uniformity
-		`SELECT format('%s.%s', n.nspname, c.relname), c.relkind::text, c.relhassubclass
+		// the input table name is quoted; and we want the unquoted canonical form for uniformity.
+		// We also fetch the OID so downstream queries can join on it directly, avoiding
+		// to_regclass which lowercases unquoted identifiers and breaks mixed-case table names.
+		`SELECT c.oid, format('%s.%s', n.nspname, c.relname), c.relkind::text, c.relhassubclass
 		 FROM pg_class c
 		 JOIN pg_namespace n ON c.relnamespace = n.oid
 		 WHERE c.oid = to_regclass($1)`,
-		table).Scan(&classification.qualifiedName, &classification.relkind, &classification.relhassubclass)
+		table).Scan(&classification.oid, &classification.qualifiedName, &classification.relkind, &classification.relhassubclass)
 	if err != nil {
 		if errors.Is(err, pgx.ErrNoRows) {
 			return classification, fmt.Errorf("table %s not found in pg_class", table)
@@ -329,30 +332,32 @@ func classifyTable(ctx context.Context, tx pgx.Tx, table string) (tableClassific
 
 // getPartitionedTables returns a map of partitioned child table names to their block counts,
 // recursively handle multi-level partitioned tables.
-func getPartitionedTables(ctx context.Context, tx pgx.Tx, table string) (map[string]int64, error) {
+// Uses OID-based join to avoid to_regclass which lowercases unquoted mixed-case identifiers.
+func getPartitionedTables(ctx context.Context, tx pgx.Tx, parentOID uint32) (map[string]int64, error) {
 	rows, err := tx.Query(ctx, `
-		SELECT format('%s.%s', n.nspname, c.relname), c.relkind::text,
+		SELECT c.oid, format('%s.%s', n.nspname, c.relname), c.relkind::text,
 			(pg_relation_size(c.oid) / current_setting('block_size')::int)::bigint
 		FROM pg_inherits i
 		JOIN pg_class c ON i.inhrelid = c.oid
 		JOIN pg_namespace n ON c.relnamespace = n.oid
-		WHERE i.inhparent = to_regclass($1)
+		WHERE i.inhparent = $1
 		ORDER BY c.relname
-	`, table)
+	`, parentOID)
 	if err != nil {
-		return nil, fmt.Errorf("failed to query child tables for %s: %w", table, err)
+		return nil, fmt.Errorf("failed to query child tables for oid %d: %w", parentOID, err)
 	}
 	defer rows.Close()
 
 	type tableInfo struct {
 		name   string
 		kind   string
 		blocks int64
+		oid    uint32
 	}
 	var tableInfos []tableInfo
 	for rows.Next() {
 		var info tableInfo
-		if err := rows.Scan(&info.name, &info.kind, &info.blocks); err != nil {
+		if err := rows.Scan(&info.oid, &info.name, &info.kind, &info.blocks); err != nil {
 			return nil, fmt.Errorf("failed to scan child table: %w", err)
 		}
 		tableInfos = append(tableInfos, info)
@@ -364,7 +369,7 @@ func getPartitionedTables(ctx context.Context, tx pgx.Tx, table string) (map[str
 	blocksPerPartitionedTable := make(map[string]int64)
 	for _, info := range tableInfos {
 		if info.kind == "p" {
-			leaveTables, err := getPartitionedTables(ctx, tx, info.name)
+			leaveTables, err := getPartitionedTables(ctx, tx, info.oid)
 			if err != nil {
 				return nil, fmt.Errorf("failed to get partitions of %s: %w", info.name, err)
 			}
@@ -378,40 +383,44 @@ func getPartitionedTables(ctx context.Context, tx pgx.Tx, table string) (map[str
 
 // getInheritedTables returns a map of table names to their block counts for an inherited
 // table hierarchy. Unlike partitioned tables, the parent itself stores data and is included.
-func getInheritedTables(ctx context.Context, tx pgx.Tx, table string) (map[string]int64, error) {
+// Uses OID-based join to avoid to_regclass which lowercases unquoted mixed-case identifiers.
+func getInheritedTables(ctx context.Context, tx pgx.Tx, parentOID uint32, parentName string) (map[string]int64, error) {
 	blocksPerTable := make(map[string]int64)
 
-	numBlocks, err := getTableBlockCount(ctx, tx, table)
-	if err != nil {
-		return nil, err
+	var numBlocks int64
+	if err := tx.QueryRow(ctx,
+		"SELECT (pg_relation_size($1) / current_setting('block_size')::int)::bigint",
+		parentOID).Scan(&numBlocks); err != nil {
+		return nil, fmt.Errorf("failed to get block count for %s: %w", parentName, err)
 	}
 	if numBlocks > 0 {
-		blocksPerTable[table] = numBlocks
+		blocksPerTable[parentName] = numBlocks
 	}
 
 	rows, err := tx.Query(ctx, `
-		SELECT format('%s.%s', n.nspname, c.relname), c.relhassubclass,
+		SELECT c.oid, format('%s.%s', n.nspname, c.relname), c.relhassubclass,
 			(pg_relation_size(c.oid) / current_setting('block_size')::int)::bigint
 		FROM pg_inherits i
 		JOIN pg_class c ON i.inhrelid = c.oid
 		JOIN pg_namespace n ON c.relnamespace = n.oid
-		WHERE i.inhparent = to_regclass($1)
+		WHERE i.inhparent = $1
 		ORDER BY c.relname
-	`, table)
+	`, parentOID)
 	if err != nil {
-		return nil, fmt.Errorf("failed to query child tables for %s: %w", table, err)
+		return nil, fmt.Errorf("failed to query child tables for %s: %w", parentName, err)
 	}
 	defer rows.Close()
 
 	type tableInfo struct {
 		name        string
-		hasSubclass bool
 		blocks      int64
+		oid         uint32
+		hasSubclass bool
 	}
 	var children []tableInfo
 	for rows.Next() {
 		var info tableInfo
-		if err := rows.Scan(&info.name, &info.hasSubclass, &info.blocks); err != nil {
+		if err := rows.Scan(&info.oid, &info.name, &info.hasSubclass, &info.blocks); err != nil {
 			return nil, fmt.Errorf("failed to scan child table: %w", err)
 		}
 		children = append(children, info)
@@ -425,7 +434,7 @@ func getInheritedTables(ctx context.Context, tx pgx.Tx, table string) (map[strin
 			blocksPerTable[child.name] = child.blocks
 		}
 		if child.hasSubclass {
-			childTables, err := getInheritedTables(ctx, tx, child.name)
+			childTables, err := getInheritedTables(ctx, tx, child.oid, child.name)
 			if err != nil {
 				return nil, fmt.Errorf("failed to get inherited tables of %s: %w", child.name, err)
 			}

flow/connectors/postgres/qrep_partition_test.go

@@ -346,19 +346,33 @@ func TestCTIDPartitioningOnPartitionedTable(t *testing.T) {
 
 func TestCTIDPartitioningOnMultiLevelPartitionedTable(t *testing.T) {
 	t.Parallel()
-	schemaName, conn, connStr := setupTestSchema(t)
+	_, conn, connStr := setupTestSchema(t)
 
-	// Root table partitioned by region
-	rootTable := schemaName + ".multi_level"
-	_, err := conn.Exec(t.Context(), fmt.Sprintf(`
+	// Use a mixed-case schema name as well to verify OID-based lookups work
+	// for both schema and table identifiers (to_regclass lowercases unquoted identifiers).
+	schemaName := "Test_" + shared.RandomString(8)
+	schemaNameDDL := `"` + schemaName + `"`
+	_, err := conn.Exec(t.Context(), `CREATE SCHEMA `+schemaNameDDL)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		if _, err := conn.Exec(t.Context(), `DROP SCHEMA `+schemaNameDDL+` CASCADE`); err != nil {
+			t.Logf("Failed to drop schema: %v", err)
+		}
+	})
+
+	// Mixed-case names to verify OID-based lookups work (to_regclass lowercases unquoted identifiers).
+	// rootTable is the unquoted form passed to PeerDB; rootTableDDL is the quoted form for SQL DDL.
+	rootTable := schemaName + ".MultiLevel"
+	rootTableDDL := schemaNameDDL + `."MultiLevel"`
+	_, err = conn.Exec(t.Context(), fmt.Sprintf(`
 		CREATE TABLE %s (
 			id SERIAL,
 			region INT NOT NULL,
 			category INT NOT NULL,
 			value TEXT,
 			PRIMARY KEY (region, category, id)
 		) PARTITION BY RANGE (region)
-	`, rootTable))
+	`, rootTableDDL))
 	require.NoError(t, err)
 
 	rowsPerPartition := 20
@@ -367,15 +381,15 @@ func TestCTIDPartitioningOnMultiLevelPartitionedTable(t *testing.T) {
 
 	// Two mid-level partitions, each sub-partitioned by category
 	for r := range numMidLevelTables {
-		mid := fmt.Sprintf("%s.region_%d", schemaName, r)
+		mid := fmt.Sprintf(`%s."Region_%d"`, schemaNameDDL, r)
 		_, err = conn.Exec(t.Context(), fmt.Sprintf(
 			`CREATE TABLE %s PARTITION OF %s FOR VALUES FROM (%d) TO (%d) PARTITION BY RANGE (category)`,
-			mid, rootTable, r*50, (r+1)*50))
+			mid, rootTableDDL, r*50, (r+1)*50))
 		require.NoError(t, err)
 
 		// Three leaf partitions per mid-level
 		for c := range numLeafChildTablesPerMidLevelTable {
-			leaf := fmt.Sprintf("%s.region_%d_cat_%d", schemaName, r, c)
+			leaf := fmt.Sprintf(`%s."Region_%d_Cat_%d"`, schemaNameDDL, r, c)
 			_, err = conn.Exec(t.Context(), fmt.Sprintf(
 				`CREATE TABLE %s PARTITION OF %s FOR VALUES FROM (%d) TO (%d)`,
 				leaf, mid, c*33, (c+1)*33))
@@ -384,14 +398,15 @@ func TestCTIDPartitioningOnMultiLevelPartitionedTable(t *testing.T) {
 	}
 
 	// 6 leaf tables total, insert 20 rows into each
+	// expectedLeaves uses unquoted names since that's what format('%s.%s', ...) in pg_class returns
 	expectedLeaves := make([]string, 0, numLeafChildTablesPerMidLevelTable*numMidLevelTables)
 	for r := range numMidLevelTables {
 		for c := range numLeafChildTablesPerMidLevelTable {
-			leaf := fmt.Sprintf("%s.region_%d_cat_%d", schemaName, r, c)
+			leaf := fmt.Sprintf("%s.Region_%d_Cat_%d", schemaName, r, c)
 			expectedLeaves = append(expectedLeaves, leaf)
 			for j := range rowsPerPartition {
 				_, err = conn.Exec(t.Context(), fmt.Sprintf(
-					`INSERT INTO %s (region, category, value) VALUES ($1, $2, $3)`, rootTable),
+					`INSERT INTO %s (region, category, value) VALUES ($1, $2, $3)`, rootTableDDL),
 					r*50, c*33+j%33, fmt.Sprintf("v_%d_%d_%d", r, c, j))
 				require.NoError(t, err)
 			}
@@ -404,7 +419,7 @@ func TestCTIDPartitioningOnMultiLevelPartitionedTable(t *testing.T) {
 		conn:    conn,
 		logger:  log.NewStructuredLogger(slog.With(slog.String(string(shared.FlowNameKey), "testMultiLevel"))),
 	}
-	query := fmt.Sprintf(`SELECT * FROM %s WHERE ctid BETWEEN {{.start}} AND {{.end}}`, rootTable)
+	query := fmt.Sprintf(`SELECT * FROM %s WHERE ctid BETWEEN {{.start}} AND {{.end}}`, rootTableDDL)
 	partitions, err := c.GetQRepPartitions(t.Context(), &protos.QRepConfig{
 		FlowJobName:           "test_ctid_multi_level",
 		NumRowsPerPartition:   10,
@@ -431,7 +446,7 @@ func TestCTIDPartitioningOnMultiLevelPartitionedTable(t *testing.T) {
 		require.Positive(t, childTableCounts[leaf])
 	}
 	for r := range numMidLevelTables {
-		mid := fmt.Sprintf("%s.region_%d", schemaName, r)
+		mid := fmt.Sprintf("%s.Region_%d", schemaName, r)
 		require.Zero(t, childTableCounts[mid], mid)
 	}
 }
@@ -597,27 +612,42 @@ func TestCtidPartitionsForChildTablesOffsetNumberBounds(t *testing.T) {
 
 func TestCTIDPartitioningOnInheritedTable(t *testing.T) {
 	t.Parallel()
-	schemaName, conn, connStr := setupTestSchema(t)
+	_, conn, connStr := setupTestSchema(t)
 
-	parentTable := schemaName + ".parent_inh"
-	_, err := conn.Exec(t.Context(), fmt.Sprintf(`CREATE TABLE %s (id SERIAL PRIMARY KEY, value TEXT)`, parentTable))
+	// Use a mixed-case schema name as well to verify OID-based lookups work
+	// for both schema and table identifiers (to_regclass lowercases unquoted identifiers).
+	schemaName := "Test_" + shared.RandomString(8)
+	schemaNameDDL := `"` + schemaName + `"`
+	_, err := conn.Exec(t.Context(), `CREATE SCHEMA `+schemaNameDDL)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		if _, err := conn.Exec(t.Context(), `DROP SCHEMA `+schemaNameDDL+` CASCADE`); err != nil {
+			t.Logf("Failed to drop schema: %v", err)
+		}
+	})
+
+	// Mixed-case names to verify OID-based lookups work (to_regclass lowercases unquoted identifiers).
+	// parentTable is the unquoted form passed to PeerDB; parentTableDDL is the quoted form for SQL DDL.
+	parentTable := schemaName + ".ParentInh"
+	parentTableDDL := schemaNameDDL + `."ParentInh"`
+	_, err = conn.Exec(t.Context(), fmt.Sprintf(`CREATE TABLE %s (id SERIAL PRIMARY KEY, value TEXT)`, parentTableDDL))
 	require.NoError(t, err)
 
 	numChildren := 3
 	rowsPerTable := 20
-	childTables := make([]string, numChildren)
+	childTablesDDL := make([]string, numChildren)
 	for i := range numChildren {
-		childTables[i] = fmt.Sprintf("%s.child_inh_%d", schemaName, i)
-		_, err = conn.Exec(t.Context(), fmt.Sprintf(`CREATE TABLE %s () INHERITS (%s)`, childTables[i], parentTable))
+		childTablesDDL[i] = fmt.Sprintf(`%s."ChildInh_%d"`, schemaNameDDL, i)
+		_, err = conn.Exec(t.Context(), fmt.Sprintf(`CREATE TABLE %s () INHERITS (%s)`, childTablesDDL[i], parentTableDDL))
 		require.NoError(t, err)
 	}
 
 	for j := range rowsPerTable {
 		_, err = conn.Exec(t.Context(), fmt.Sprintf(
-			`INSERT INTO %s (value) VALUES ($1)`, parentTable), fmt.Sprintf("parent_%d", j))
+			`INSERT INTO %s (value) VALUES ($1)`, parentTableDDL), fmt.Sprintf("parent_%d", j))
 		require.NoError(t, err)
 	}
-	for i, child := range childTables {
+	for i, child := range childTablesDDL {
 		for j := range rowsPerTable {
 			_, err = conn.Exec(t.Context(), fmt.Sprintf(
 				`INSERT INTO %s (value) VALUES ($1)`, child), fmt.Sprintf("child_%d_%d", i, j))
@@ -648,10 +678,11 @@ func TestCTIDPartitioningOnInheritedTable(t *testing.T) {
 			childTablesCovered[ctr.Table] = true
 		}
 	}
+	// pg_class returns unquoted names from format('%s.%s', ...) so assertions use the unquoted form
 	require.Len(t, childTablesCovered, 1+numChildren)
-	require.True(t, childTablesCovered[parentTable])
-	for _, child := range childTables {
-		require.True(t, childTablesCovered[child])
+	require.True(t, childTablesCovered[schemaName+".ParentInh"])
+	for i := range numChildren {
+		require.True(t, childTablesCovered[fmt.Sprintf("%s.ChildInh_%d", schemaName, i)])
 	}
 }