pgcdc: fix table metadata

Author: rockwotj

CHANGELOG.md

@@ -17,6 +17,7 @@ All notable changes to this project will be documented in this file.
 ### Fixed
 
 - Fix a snapshot stream consistency issue with `postgres_cdc` where data could be missed if writes where happening during the snapshot phase. (@rockwotj)
+- Fix an issue where `@table` metadata was quoted for the snapshot phase in `postgres_cdc`. (@rockwotj)
 
 ### Changed
 

internal/impl/postgresql/integration_test.go

@@ -924,3 +924,119 @@ read_until:
 	require.Equal(t, expected, sequenceNumbers)
 	batchMu.Unlock()
 }
+
+func TestIntegrationPostgresMetadata(t *testing.T) {
+	t.Parallel()
+	integration.CheckSkip(t)
+	pool, err := dockertest.NewPool("")
+	require.NoError(t, err)
+
+	var (
+		resource *dockertest.Resource
+		db       *sql.DB
+	)
+
+	resource, db, err = ResourceWithPostgreSQLVersion(t, pool, "16")
+	require.NoError(t, err)
+	require.NoError(t, resource.Expire(120))
+
+	hostAndPort := resource.GetHostPort("5432/tcp")
+	hostAndPortSplited := strings.Split(hostAndPort, ":")
+	password := "l]YLSc|4[i56%{gY"
+
+	require.NoError(t, err)
+
+	_, err = db.Exec(`INSERT INTO "FlightsCompositePK" ("Seq", "Name", "CreatedAt") VALUES ($1, $2, $3);`, 1, "delta", "2006-01-02T15:04:05Z07:00")
+	require.NoError(t, err)
+	_, err = db.Exec(`INSERT INTO flights (name, created_at) VALUES ($1, $2);`, "delta", "2006-01-02T15:04:05Z07:00")
+	require.NoError(t, err)
+
+	databaseURL := fmt.Sprintf("user=user_name password=%s dbname=dbname sslmode=disable host=%s port=%s", password, hostAndPortSplited[0], hostAndPortSplited[1])
+	template := fmt.Sprintf(`
+postgres_cdc:
+    dsn: %s
+    slot_name: test_slot_native_decoder
+    stream_snapshot: true
+    snapshot_batch_size: 5
+    schema: public
+    tables:
+      - '"FlightsCompositePK"'
+      - flights
+`, databaseURL)
+
+	streamOutBuilder := service.NewStreamBuilder()
+	require.NoError(t, streamOutBuilder.SetLoggerYAML(`level: TRACE`))
+	require.NoError(t, streamOutBuilder.AddInputYAML(template))
+	require.NoError(t, streamOutBuilder.AddProcessorYAML(`mapping: 'root = @'`))
+
+	var outBatches []any
+	var outBatchMut sync.Mutex
+	require.NoError(t, streamOutBuilder.AddBatchConsumerFunc(func(c context.Context, batch service.MessageBatch) error {
+		outBatchMut.Lock()
+		defer outBatchMut.Unlock()
+		for _, msg := range batch {
+			data, err := msg.AsStructured()
+			require.NoError(t, err)
+			d := data.(map[string]any)
+			if _, ok := d["lsn"]; ok {
+				d["lsn"] = "XXX/XXX" // Consistent LSN for assertions below
+			}
+			outBatches = append(outBatches, data)
+		}
+		return nil
+	}))
+
+	streamOut, err := streamOutBuilder.Build()
+	require.NoError(t, err)
+
+	license.InjectTestService(streamOut.Resources())
+
+	go func() {
+		_ = streamOut.Run(context.Background())
+	}()
+
+	assert.Eventually(t, func() bool {
+		outBatchMut.Lock()
+		defer outBatchMut.Unlock()
+		return len(outBatches) == 2
+	}, time.Second*25, time.Millisecond*100)
+
+	_, err = db.Exec(`INSERT INTO "FlightsCompositePK" ("Seq", "Name", "CreatedAt") VALUES ($1, $2, $3);`, 2, "bravo", "2006-01-02T15:04:05Z07:00")
+	require.NoError(t, err)
+	_, err = db.Exec(`INSERT INTO flights (name, created_at) VALUES ($1, $2);`, "bravo", "2006-01-02T15:04:05Z07:00")
+	require.NoError(t, err)
+
+	assert.EventuallyWithT(t, func(c *assert.CollectT) {
+		outBatchMut.Lock()
+		defer outBatchMut.Unlock()
+		assert.Len(c, outBatches, 4, "got: %#v", outBatches)
+	}, time.Second*25, time.Millisecond*100)
+
+	require.ElementsMatch(
+		t,
+		outBatches,
+		[]any{
+			map[string]any{
+				"operation": "read",
+				"table":     "FlightsCompositePK",
+			},
+			map[string]any{
+				"operation": "read",
+				"table":     "flights",
+			},
+			map[string]any{
+				"operation": "insert",
+				"table":     "flights",
+				"lsn":       "XXX/XXX",
+			},
+			map[string]any{
+				"operation": "insert",
+				"table":     "FlightsCompositePK",
+				"lsn":       "XXX/XXX",
+			},
+		},
+	)
+
+	require.NoError(t, streamOut.StopWithin(time.Second*10))
+
+}

internal/impl/postgresql/pglogicalstream/logical_stream.go

@@ -503,6 +503,15 @@ func (s *Stream) processSnapshot(ctx context.Context, snapshotter *Snapshotter)
 		wg.Go(func() (err error) {
 			s.logger.Debugf("Processing snapshot for table: %v", table)
 
+			unquotedTable, err := sanitize.UnquotePostgresIdentifier(table.Table)
+			if err != nil {
+				return fmt.Errorf("unexpected failure to unquote table name: %w", err)
+			}
+			unquotedSchema, err := sanitize.UnquotePostgresIdentifier(table.Schema)
+			if err != nil {
+				return fmt.Errorf("unexpected failure to unquote schema name: %w", err)
+			}
+
 			avgRowSizeBytes, numRows, err := snapshotter.tableStats(ctx, table)
 			if err != nil {
 				return fmt.Errorf("failed to calculate average row size for table %v: %w", table, err)
@@ -592,8 +601,8 @@ func (s *Stream) processSnapshot(ctx context.Context, snapshotter *Snapshotter)
 					snapshotChangePacket := StreamMessage{
 						LSN:       nil,
 						Operation: ReadOpType,
-						Table:     table.Table,
-						Schema:    table.Schema,
+						Table:     unquotedTable,
+						Schema:    unquotedSchema,
 						Data:      data,
 					}
 

internal/impl/postgresql/pglogicalstream/sanitize/sanitize.go

@@ -381,6 +381,32 @@ func QuotePostgresIdentifier(name string) string {
 	return quoted.String()
 }
 
+// UnquotePostgresIdentifier returns the valid unescaped identifier.
+func UnquotePostgresIdentifier(quoted string) (string, error) {
+	var output strings.Builder
+	// Default to assume we're just going to add quotes and there won't
+	// be any double quotes inside the string that needs escaped.
+	if !strings.HasPrefix(quoted, `"`) || !strings.HasSuffix(quoted, `"`) || len(quoted) < 2 {
+		return "", errors.New("missing quotes for identifier")
+	}
+	unquoted := quoted[1 : len(quoted)-1]
+	output.Grow(len(unquoted))
+	for i := 0; i < len(unquoted); i++ {
+		_ = output.WriteByte(unquoted[i])
+		if unquoted[i] != '"' {
+			continue
+		}
+		if i+1 >= len(unquoted) {
+			return "", fmt.Errorf("invalid quoted identifier: %s", quoted)
+		}
+		if unquoted[i+1] != '"' {
+			return "", fmt.Errorf("invalid quoted identifier: %s", quoted)
+		}
+		i++ // Skip over the next character to handle triple quotes
+	}
+	return output.String(), nil
+}
+
 // NormalizePostgresIdentifier checks if a string is a valid PostgreSQL identifier
 // This follows PostgreSQL's standard naming rules
 func NormalizePostgresIdentifier(name string) (string, error) {

internal/impl/postgresql/pglogicalstream/sanitize/sanitize_test.go

@@ -256,23 +256,31 @@ func TestQuerySanitize(t *testing.T) {
 }
 
 func TestIdentifierValidation(t *testing.T) {
-	quoted := []string{
-		`"FooBar"`,
-		`"Foo""Bar"`,
-		`"Foo""""Bar"`,
+	tests := []struct {
+		quoted   string
+		unquoted string
+	}{
+		{quoted: `"FooBar"`, unquoted: "FooBar"},
+		{quoted: `"Foo""Bar"`, unquoted: `Foo"Bar`},
+		{quoted: `"Foo""""Bar"`, unquoted: `Foo""Bar`},
 	}
 
-	for _, i := range quoted {
-		i := i
-		t.Run(i, func(t *testing.T) {
-			_, err := sanitize.NormalizePostgresIdentifier(i)
+	for _, testcase := range tests {
+		testcase := testcase
+		t.Run(testcase.unquoted, func(t *testing.T) {
+			q, err := sanitize.NormalizePostgresIdentifier(testcase.quoted)
+			require.NoError(t, err)
+			require.Equal(t, testcase.quoted, q)
+			r, err := sanitize.UnquotePostgresIdentifier(q)
 			require.NoError(t, err)
+			require.Equal(t, testcase.unquoted, r)
 		})
 	}
 
 	unquoted := []string{
 		`_Foobar`,
 		strings.Repeat("a", 63),
+		strings.Repeat("A", 63),
 	}
 
 	for _, i := range unquoted {
@@ -281,6 +289,9 @@ func TestIdentifierValidation(t *testing.T) {
 			normalized, err := sanitize.NormalizePostgresIdentifier(i)
 			require.NoError(t, err)
 			require.Equal(t, strconv.Quote(strings.ToLower(i)), normalized)
+			unquoted, err := sanitize.UnquotePostgresIdentifier(normalized)
+			require.NoError(t, err)
+			require.Equal(t, strings.ToLower(i), unquoted)
 		})
 	}