Repository: xataio/pgstream
Pull request #374 (2025-06-02 15:04:29)

[Snapshot] Do not copy generated columns
This PR updates the postgres bulk ingest writer logic to prevent copying generated columns (excluding identities). Generated columns are computed when written, and cannot be written to directly. The COPY command does not support them. 
https://www.postgresql.org/docs/current/ddl-generated-columns.htm 

Instead of modifying the snapshotting process to remove the generated columns from the snapshot, the changes are done on the postgres target side, since the snapshot can be used with any supported target, and this is only an issue for the postgres usecase. For any other target, we'll need to have the generated columns along with the rest, since they won't have the postgres logic embedded to generate it.

Fixes #373 


Patch:
diff --git a/docs/README.md b/docs/README.md
index 87d4b16..c96abc0 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -412,7 +412,7 @@ One of exponential/constant backoff policies can be provided for the search stor
 
 <details>
   <summary>Traces</summary>
-   
+
 | Environment Variable                 | Default | Required | Description                                                                                                                     |
 | ------------------------------------ | ------- | -------- | ------------------------------------------------------------------------------------------------------------------------------- |
 | PGSTREAM_TRACES_ENDPOINT             | N/A     | No       | Endpoint where the pgstream traces will be exported to.                                                                         |
diff --git a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go
index c4c6b28..eca0ed6 100644
--- a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go
+++ b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go
@@ -25,6 +25,9 @@ type BulkIngestWriter struct {
 	batchSenderMapMutex *sync.RWMutex
 	batchSenderMap      map[string]queryBatchSender
 	batchSenderBuilder  func(ctx context.Context) (queryBatchSender, error)
+
+	tableColumns      map[string][]string
+	tableColumnsMutex *sync.RWMutex
 }
 
 const bulkIngestWriter = "postgres_bulk_ingest_writer"
@@ -53,6 +56,8 @@ func NewBulkIngestWriter(ctx context.Context, config *Config, opts ...WriterOpti
 		Writer:              w,
 		batchSenderMapMutex: &sync.RWMutex{},
 		batchSenderMap:      make(map[string]queryBatchSender),
+		tableColumns:        map[string][]string{},
+		tableColumnsMutex:   &sync.RWMutex{},
 	}
 
 	biw.batchSenderBuilder = func(ctx context.Context) (queryBatchSender, error) {
@@ -170,20 +175,23 @@ func (w *BulkIngestWriter) copyFromInsertQueries(ctx context.Context, inserts []
 	if len(inserts) == 0 {
 		return nil
 	}
-	// Get the column names from the first insert query, since all the inserts
+
+	// Get the table name from the first insert query, since all the inserts
 	// in the batch will be for the same schema.table
 	query := inserts[0]
-	srcRows := [][]any{}
-	for _, q := range inserts {
-		srcRows = append(srcRows, q.args)
+	// get column names for the table, removing generated columns, since the
+	// COPY command does not support them.
+	columnNames, err := w.getTableColumnNames(ctx, query.schema, query.table)
+	if err != nil {
+		return err
 	}
 
-	err := w.pgConn.ExecInTx(ctx, func(tx pglib.Tx) error {
+	err = w.pgConn.ExecInTx(ctx, func(tx pglib.Tx) error {
 		if err := w.setReplicationRoleToReplica(ctx, tx); err != nil {
 			return err
 		}
 
-		rowsCopied, err := tx.CopyFrom(ctx, pglib.QuoteQualifiedIdentifier(query.schema, query.table), query.columnNames, srcRows)
+		rowsCopied, err := tx.CopyFrom(ctx, pglib.QuoteQualifiedIdentifier(query.schema, query.table), columnNames, w.getSrcRows(inserts, columnNames))
 		if err != nil {
 			return err
 		}
@@ -200,3 +208,76 @@ func (w *BulkIngestWriter) copyFromInsertQueries(ctx context.Context, inserts []
 	}
 	return nil
 }
+
+func (w *BulkIngestWriter) getTableColumnNames(ctx context.Context, schema, table string) ([]string, error) {
+	key := pglib.QuoteQualifiedIdentifier(schema, table)
+	w.tableColumnsMutex.RLock()
+	columnNames, found := w.tableColumns[key]
+	w.tableColumnsMutex.RUnlock()
+	if found {
+		return columnNames, nil
+	}
+
+	colNames, err := w.queryTableColumnNames(ctx, schema, table)
+	if err != nil {
+		return nil, err
+	}
+
+	w.tableColumnsMutex.Lock()
+	w.tableColumns[key] = colNames
+	w.tableColumnsMutex.Unlock()
+
+	return colNames, nil
+}
+
+func (w *BulkIngestWriter) getSrcRows(inserts []*query, columnNames []string) [][]any {
+	srcRows := [][]any{}
+	columnMap := make(map[string]struct{}, len(columnNames))
+	for _, col := range columnNames {
+		columnMap[pglib.QuoteIdentifier(col)] = struct{}{}
+	}
+	for _, q := range inserts {
+		// we need to make sure we only add the arguments for the
+		// relevant column names (this removes any generated columns row values)
+		rows := []any{}
+		for i, colName := range q.columnNames {
+			if _, found := columnMap[colName]; !found {
+				continue
+			}
+			rows = append(rows, q.args[i])
+		}
+		srcRows = append(srcRows, rows)
+	}
+
+	return srcRows
+}
+
+const tableColumnsQuery = `SELECT attname FROM pg_attribute
+		WHERE attnum > 0
+		AND attrelid = (SELECT c.oid FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2)
+		AND attgenerated = ''`
+
+func (w *BulkIngestWriter) queryTableColumnNames(ctx context.Context, schemaName, tableName string) ([]string, error) {
+	columnNames := []string{}
+	// filter out generated columns (excluding identities) since they will
+	// be generated automatically, and they can't be overwriten.
+	rows, err := w.pgConn.Query(ctx, tableColumnsQuery, tableName, schemaName)
+	if err != nil {
+		return nil, fmt.Errorf("getting table column names for table %s.%s: %w", schemaName, tableName, err)
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var columnName string
+		if err := rows.Scan(&columnName); err != nil {
+			return nil, fmt.Errorf("scanning table column name: %w", err)
+		}
+		columnNames = append(columnNames, columnName)
+	}
+
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	return columnNames, nil
+}


Test patch:
diff --git a/pkg/stream/integration/helper_test.go b/pkg/stream/integration/helper_test.go
index 166f462..6a50e00 100644
--- a/pkg/stream/integration/helper_test.go
+++ b/pkg/stream/integration/helper_test.go
@@ -43,6 +43,14 @@ var (
 	elasticsearchURL string
 )
 
+const (
+	withBulkIngestion    = true
+	withoutBulkIngestion = false
+
+	withGeneratedColumn    = true
+	withoutGeneratedColumn = false
+)
+
 type mockProcessor struct {
 	eventChan chan *wal.Event
 }
@@ -230,17 +238,19 @@ func testWebhookProcessorCfg() stream.ProcessorConfig {
 	}
 }
 
-func testPostgresProcessorCfg(sourcePGURL string) stream.ProcessorConfig {
+func testPostgresProcessorCfg(sourcePGURL string, bulkIngestion bool) stream.ProcessorConfig {
 	return stream.ProcessorConfig{
 		Postgres: &stream.PostgresProcessorConfig{
 			BatchWriter: postgres.Config{
 				URL: targetPGURL,
 				BatchConfig: batch.Config{
-					BatchTimeout: 50 * time.Millisecond,
+					MaxBatchSize: 1,
+					// BatchTimeout: 50 * time.Millisecond,
 				},
 				SchemaLogStore: schemalogpg.Config{
 					URL: sourcePGURL,
 				},
+				BulkIngestEnabled: bulkIngestion,
 			},
 		},
 		Injector: &injector.Config{
diff --git a/pkg/stream/integration/pg_pg_integration_test.go b/pkg/stream/integration/pg_pg_integration_test.go
index a439ead..f37c14f 100644
--- a/pkg/stream/integration/pg_pg_integration_test.go
+++ b/pkg/stream/integration/pg_pg_integration_test.go
@@ -22,8 +22,9 @@ type informationSchemaColumn struct {
 }
 
 type testTableColumn struct {
-	id   int
-	name string
+	id       int
+	name     string
+	username string
 }
 
 func Test_PostgresToPostgres(t *testing.T) {
@@ -33,7 +34,7 @@ func Test_PostgresToPostgres(t *testing.T) {
 
 	cfg := &stream.Config{
 		Listener:  testPostgresListenerCfg(),
-		Processor: testPostgresProcessorCfg(pgurl),
+		Processor: testPostgresProcessorCfg(pgurl, withoutBulkIngestion),
 	}
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -240,7 +241,7 @@ func Test_PostgresToPostgres_Snapshot(t *testing.T) {
 
 	cfg := &stream.Config{
 		Listener:  testPostgresListenerCfgWithSnapshot(snapshotPGURL, targetPGURL, []string{testTable}),
-		Processor: testPostgresProcessorCfg(snapshotPGURL),
+		Processor: testPostgresProcessorCfg(snapshotPGURL, withoutBulkIngestion),
 	}
 	initStream(t, ctx, snapshotPGURL)
 	runStream(t, ctx, cfg)
@@ -310,7 +311,7 @@ func getInformationSchemaColumns(t *testing.T, ctx context.Context, conn pglib.Q
 	return columns
 }
 
-func getTestTableColumns(t *testing.T, ctx context.Context, conn pglib.Querier, query string) []*testTableColumn {
+func getTestTableColumns(t *testing.T, ctx context.Context, conn pglib.Querier, query string, withGeneratedColumn ...bool) []*testTableColumn {
 	rows, err := conn.Query(ctx, query)
 	require.NoError(t, err)
 	defer rows.Close()
@@ -318,7 +319,11 @@ func getTestTableColumns(t *testing.T, ctx context.Context, conn pglib.Querier,
 	columns := []*testTableColumn{}
 	for rows.Next() {
 		column := &testTableColumn{}
-		err := rows.Scan(&column.id, &column.name)
+		if len(withGeneratedColumn) > 0 && withGeneratedColumn[0] {
+			err = rows.Scan(&column.id, &column.name, &column.username)
+		} else {
+			err = rows.Scan(&column.id, &column.name)
+		}
 		require.NoError(t, err)
 		columns = append(columns, column)
 	}
diff --git a/pkg/stream/integration/snapshot_pg_integration_test.go b/pkg/stream/integration/snapshot_pg_integration_test.go
index 04dbea9..58d93bc 100644
--- a/pkg/stream/integration/snapshot_pg_integration_test.go
+++ b/pkg/stream/integration/snapshot_pg_integration_test.go
@@ -33,12 +33,12 @@ func Test_SnapshotToPostgres(t *testing.T) {
 	testTable := "snapshot2pg_integration_test"
 	// create table and populate it before initialising and running pgstream to
 	// ensure the snapshot captures pre-existing schema and data properly
-	execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf("create table %s(id serial primary key, name text)", testTable))
+	execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf(`CREATE TABLE %s(id INTEGER GENERATED ALWAYS AS IDENTITY PRIMARY KEY, name TEXT, username TEXT GENERATED ALWAYS AS ('user_' || name ) STORED)`, testTable))
 	execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf("insert into %s(name) values('a'),('b')", testTable))
 
 	cfg := &stream.Config{
 		Listener:  testSnapshotListenerCfg(snapshotPGURL, targetPGURL, []string{testTable}),
-		Processor: testPostgresProcessorCfg(snapshotPGURL),
+		Processor: testPostgresProcessorCfg(snapshotPGURL, withBulkIngestion),
 	}
 	initStream(t, ctx, snapshotPGURL)
 	runSnapshot(t, ctx, cfg)
@@ -53,24 +53,25 @@ func Test_SnapshotToPostgres(t *testing.T) {
 
 	validation := func() bool {
 		schemaColumns := getInformationSchemaColumns(t, ctx, targetConn, testTable)
-		if len(schemaColumns) != 2 {
+		if len(schemaColumns) != 3 {
 			return false
 		}
 
 		wantSchemaCols := []*informationSchemaColumn{
 			{name: "id", dataType: "integer", isNullable: "NO"},
 			{name: "name", dataType: "text", isNullable: "YES"},
+			{name: "username", dataType: "text", isNullable: "YES"},
 		}
 		require.ElementsMatch(t, wantSchemaCols, schemaColumns)
 
-		columns := getTestTableColumns(t, ctx, targetConn, fmt.Sprintf("select id,name from %s", testTable))
+		columns := getTestTableColumns(t, ctx, targetConn, fmt.Sprintf("select id,name,username from %s", testTable), withGeneratedColumn)
 		if len(columns) != 2 {
 			return false
 		}
 
 		wantCols := []*testTableColumn{
-			{id: 1, name: "a"},
-			{id: 2, name: "b"},
+			{id: 1, name: "a", username: "user_a"},
+			{id: 2, name: "b", username: "user_b"},
 		}
 		require.ElementsMatch(t, wantCols, columns)
 
diff --git a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go
index 12f9ada..fe4469d 100644
--- a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go
+++ b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go
@@ -234,11 +234,31 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 	testQuery := &query{
 		schema:      testSchema,
 		table:       testTable,
-		columnNames: []string{"id", "name"},
+		columnNames: []string{`"id"`, `"name"`},
 		sql:         "INSERT INTO users(id, name) VALUES($1, $2)",
 		args:        []any{1, "alice"},
 	}
 
+	filteredColumnNames := []string{"id"}
+	filteredValues := []any{1}
+
+	validQueryFn := func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+		require.Equal(t, tableColumnsQuery, query)
+		require.Equal(t, []any{testTable, testSchema}, args)
+		return &pgmocks.Rows{
+			CloseFn: func() {},
+			NextFn:  func(i uint) bool { return i == 1 },
+			ScanFn: func(args ...any) error {
+				require.Len(t, args, 1)
+				colName, ok := args[0].(*string)
+				require.True(t, ok, fmt.Sprintf("column name, expected *string, got %T", args[0]))
+				*colName = "id"
+				return nil
+			},
+			ErrFn: func() error { return nil },
+		}, nil
+	}
+
 	tests := []struct {
 		name            string
 		batch           *batch.Batch[*query]
@@ -259,6 +279,9 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 					}
 					return f(tx)
 				},
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					return nil, errors.New("unexpected call to QueryFn with empty batch")
+				},
 			},
 
 			wantErr: nil,
@@ -267,12 +290,13 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 			name:  "ok",
 			batch: batch.NewBatch([]*query{testQuery}, nil),
 			pgConn: &pgmocks.Querier{
+				QueryFn: validQueryFn,
 				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
 					tx := &pgmocks.Tx{
 						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
 							require.Equal(t, pglib.QuoteQualifiedIdentifier(testSchema, testTable), tableName)
-							require.Equal(t, testQuery.columnNames, columnNames)
-							require.Equal(t, testQuery.args, rows[0])
+							require.Equal(t, filteredColumnNames, columnNames)
+							require.Equal(t, filteredValues, rows[0])
 							return 1, nil
 						},
 					}
@@ -286,6 +310,7 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 			name:  "ok - disable triggers",
 			batch: batch.NewBatch([]*query{testQuery}, nil),
 			pgConn: &pgmocks.Querier{
+				QueryFn: validQueryFn,
 				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
 					tx := &pgmocks.Tx{
 						ExecFn: func(ctx context.Context, i uint, s string, a ...any) (pglib.CommandTag, error) {
@@ -301,8 +326,8 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 						},
 						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
 							require.Equal(t, pglib.QuoteQualifiedIdentifier(testSchema, testTable), tableName)
-							require.Equal(t, testQuery.columnNames, columnNames)
-							require.Equal(t, testQuery.args, rows[0])
+							require.Equal(t, filteredColumnNames, columnNames)
+							require.Equal(t, filteredValues, rows[0])
 							return 1, nil
 						},
 					}
@@ -313,10 +338,25 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 
 			wantErr: nil,
 		},
+		{
+			name:  "error - getting column names",
+			batch: batch.NewBatch([]*query{testQuery}, nil),
+			pgConn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					return nil, errTest
+				},
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					return errors.New("unexpected call to CopyFrom")
+				},
+			},
+
+			wantErr: errTest,
+		},
 		{
 			name:  "error - copying from",
 			batch: batch.NewBatch([]*query{testQuery}, nil),
 			pgConn: &pgmocks.Querier{
+				QueryFn: validQueryFn,
 				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
 					tx := &pgmocks.Tx{
 						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
@@ -333,6 +373,7 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 			name:  "error - rows copied mismatch",
 			batch: batch.NewBatch([]*query{testQuery}, nil),
 			pgConn: &pgmocks.Querier{
+				QueryFn: validQueryFn,
 				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
 					tx := &pgmocks.Tx{
 						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
@@ -349,6 +390,7 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 			name:  "error - setting replication role to replica",
 			batch: batch.NewBatch([]*query{testQuery}, nil),
 			pgConn: &pgmocks.Querier{
+				QueryFn: validQueryFn,
 				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
 					tx := &pgmocks.Tx{
 						ExecFn: func(ctx context.Context, i uint, s string, a ...any) (pglib.CommandTag, error) {
@@ -375,6 +417,7 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 			name:  "error - resetting replication role",
 			batch: batch.NewBatch([]*query{testQuery}, nil),
 			pgConn: &pgmocks.Querier{
+				QueryFn: validQueryFn,
 				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
 					tx := &pgmocks.Tx{
 						ExecFn: func(ctx context.Context, i uint, s string, a ...any) (pglib.CommandTag, error) {
@@ -412,10 +455,206 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 					pgConn:          tc.pgConn,
 					disableTriggers: tc.disableTriggers,
 				},
+				tableColumns:      map[string][]string{},
+				tableColumnsMutex: &sync.RWMutex{},
 			}
 
 			err := writer.sendBatch(context.Background(), tc.batch)
+			if !errors.Is(err, tc.wantErr) {
+				require.Equal(t, tc.wantErr, err)
+			}
+		})
+	}
+}
+
+func TestBulkIngestWriter_getTableColumnNames(t *testing.T) {
+	t.Parallel()
+
+	quotedQualifiedTableName := `"test_schema"."test_table"`
+	idColumn := "id"
+
+	tests := []struct {
+		name         string
+		tableColumns map[string][]string
+		pgConn       pglib.Querier
+
+		wantColumns      []string
+		wantTableColumns map[string][]string
+		wantErr          error
+	}{
+		{
+			name:         "ok - empty map",
+			tableColumns: map[string][]string{},
+			pgConn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, tableColumnsQuery, query)
+					require.Equal(t, []any{testTable, testSchema}, args)
+					return &pgmocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(args ...any) error {
+							require.Len(t, args, 1)
+							colName, ok := args[0].(*string)
+							require.True(t, ok, fmt.Sprintf("column name, expected *string, got %T", args[0]))
+							*colName = "id"
+							return nil
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				},
+			},
+
+			wantColumns: []string{idColumn},
+			wantTableColumns: map[string][]string{
+				quotedQualifiedTableName: {idColumn},
+			},
+			wantErr: nil,
+		},
+		{
+			name: "ok - existing table",
+			tableColumns: map[string][]string{
+				quotedQualifiedTableName: {idColumn},
+			},
+			pgConn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					return nil, errors.New("unexpected call to QueryFn")
+				},
+			},
+
+			wantColumns: []string{idColumn},
+			wantTableColumns: map[string][]string{
+				quotedQualifiedTableName: {idColumn},
+			},
+			wantErr: nil,
+		},
+		{
+			name:         "error - quering table columns",
+			tableColumns: map[string][]string{},
+			pgConn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					return nil, errTest
+				},
+			},
+
+			wantColumns:      nil,
+			wantTableColumns: map[string][]string{},
+			wantErr:          errTest,
+		},
+		{
+			name:         "error - scanning table column",
+			tableColumns: map[string][]string{},
+			pgConn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, tableColumnsQuery, query)
+					require.Equal(t, []any{testTable, testSchema}, args)
+					return &pgmocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(args ...any) error {
+							return errTest
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				},
+			},
+
+			wantColumns:      nil,
+			wantTableColumns: map[string][]string{},
+			wantErr:          errTest,
+		},
+		{
+			name:         "error - rows error",
+			tableColumns: map[string][]string{},
+			pgConn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, tableColumnsQuery, query)
+					require.Equal(t, []any{testTable, testSchema}, args)
+					return &pgmocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(args ...any) error {
+							return nil
+						},
+						ErrFn: func() error { return errTest },
+					}, nil
+				},
+			},
+
+			wantColumns:      nil,
+			wantTableColumns: map[string][]string{},
+			wantErr:          errTest,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			w := &BulkIngestWriter{
+				Writer: &Writer{
+					logger: loglib.NewNoopLogger(),
+					pgConn: tc.pgConn,
+				},
+				tableColumns:      tc.tableColumns,
+				tableColumnsMutex: &sync.RWMutex{},
+			}
+
+			colNames, err := w.getTableColumnNames(context.TODO(), "test_schema", "test_table")
 			require.ErrorIs(t, err, tc.wantErr)
+			require.Equal(t, tc.wantColumns, colNames)
+			require.Equal(t, tc.wantTableColumns, w.tableColumns)
+		})
+	}
+}
+
+func TestBulkIngestWriter_getSrcRows(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name        string
+		inserts     []*query
+		columnNames []string
+
+		wantRows [][]any
+	}{
+		{
+			name: "ok - no generated columns",
+			inserts: []*query{
+				{
+					columnNames: []string{`"a"`, `"b"`, `"c"`},
+					args:        []any{"1", "2", "3"},
+				},
+			},
+			columnNames: []string{"a", "b", "c"},
+
+			wantRows: [][]any{
+				{"1", "2", "3"},
+			},
+		},
+		{
+			name: "ok - with generated columns",
+			inserts: []*query{
+				{
+					columnNames: []string{`"a"`, `"b"`, `"c"`, `"d"`, `"e"`},
+					args:        []any{"1", "2", "3", "4", "5"},
+				},
+			},
+			columnNames: []string{"a", "c", "d"},
+
+			wantRows: [][]any{
+				{"1", "3", "4"},
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			w := &BulkIngestWriter{}
+
+			rows := w.getSrcRows(tc.inserts, tc.columnNames)
+			require.Equal(t, tc.wantRows, rows)
 		})
 	}
 }