diff --git a/docs/examples/transformer_rules.yaml b/docs/examples/transformer_rules.yaml
index 468931d15..8b3bee194 100644
--- a/docs/examples/transformer_rules.yaml
+++ b/docs/examples/transformer_rules.yaml
@@ -15,6 +15,13 @@ transformations:
# example key only — generate your own with `openssl rand -hex 64`
key_hex: "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132333435363738393a3b3c3d3e3f"
associated_data: "public.test.document_path"
+ country_id:
+ name: lookup_choice
+ parameters:
+ lookup_table: public.countries
+ lookup_column: id
+ generator: deterministic
+ ignore_values: [0, -1]
phone_number:
name: fpe_ff1
parameters:
diff --git a/docs/transformers.md b/docs/transformers.md
index ba01bace7..50f1c5faa 100644
--- a/docs/transformers.md
+++ b/docs/transformers.md
@@ -43,6 +43,7 @@ Each transformer declares how it behaves with respect to uniqueness:
| `greenmask_boolean` | `lossy` |
| `greenmask_choice` | `lossy` |
| `greenmask_firstname` | `lossy` |
+| `lookup_choice` | `lossy` |
| `literal_string` | `lossy` |
| `masking` | `lossy` |
| `neosync_firstname` | `lossy` |
@@ -1262,6 +1263,73 @@ transformations:
Every run with the same key and parameters produces the same output. Tokens can be decrypted with any RFC 5297 AES-SIV implementation, for example Tink's `daead/subtle` package in Go.
+
+
+
+ lookup_choice
+
+**Description:** Replaces a value with one taken from a column of another table, for example a foreign key column pointing at a lookup table. The values are read from the source database once, when the pipeline starts, so the configuration does not have to be regenerated when the lookup table's contents change. It is the live-table counterpart of [`greenmask_choice`](#supported-transformers), which chooses from a list written into the configuration.
+
+**Uniqueness:** `lossy`. Any table with more rows than the lookup column has values produces duplicates, in both generator modes. Cannot be used on a column covered by a unique index. See [Uniqueness and unique indexes](#uniqueness-and-unique-indexes).
+
+| Supported PostgreSQL types |
+| ------------------------------------------------------------------------------------------------------------------------ |
+| `text`, `varchar`, `char`, `bpchar`, `citext`, `bytea`, `boolean`, `int2`, `int4`, `int8`, `float4`, `float8`, `uuid`, `date`, `timestamp`, `timestamptz` |
+
+The type comes from the **lookup column**: the transformer asks PostgreSQL what it is and reports the column types its values can be written to, so a rule pointing a column at a lookup column of an incompatible type is rejected on startup. A narrower integer or float is accepted for a wider column (an `int4` lookup key can fill an `int8` foreign key). A lookup column of any other type is rejected rather than silently skipping the check.
+
+| Parameter | Type | Default | Required | Values |
+| ------------- | -------- | ------- | -------- | --------------------- |
+| lookup_table | string | N/A | Yes | N/A |
+| lookup_column | string | N/A | Yes | N/A |
+| generator | string | random | No | random, deterministic |
+| ignore_values | any[] | [] | No | N/A |
+| postgres_url | string | N/A | Yes | N/A |
+
+`lookup_table` is schema qualified, e.g. `public.countries`; an unqualified name is read from the `public` schema. Both names are quoted as written, so `Public.Countries` looks for a case-sensitive `"Countries"`.
+
+`postgres_url` is required, but the PostgreSQL parser fills it in with the URL of the source database being read, so it only has to be written out when the source is not PostgreSQL.
+
+`ignore_values` removes values from the list after it is read, for placeholder rows such as an "unknown" id. An entry that matches nothing is an error, so a typo or a value written in a form the column never produces is reported rather than quietly leaving the row in the choice set. If it excludes every value, or the lookup column is empty, the pipeline fails to start rather than writing the same value into every row.
+
+`generator: deterministic` picks the value from a hash of the incoming value, so every row that pointed at the same original value still points at one single new value. `generator: random` picks independently for each row and destroys that grouping. Deterministic mode is rejected for `timestamp` and `timestamptz` lookup columns, because a snapshot and a replication event deliver a timestamp in forms that cannot be reduced to the same hash input, so the same row would be mapped differently either side of the cutover.
+
+**Security note:** the deterministic mapping is an unsalted hash over a value set that is usually small and often public. Anyone holding the transformed data and a guess at the lookup table can compute the same hashes and recover much of the original mapping. Deterministic mode preserves structure; it does not hide the values it maps from. The same is true of `greenmask_choice`.
+
+⚠️ **A reference can be left dangling, and the rows are dropped rather than reported.** The values are read from the lookup table in the **source** database. If that table's own key column is itself transformed, or the lookup row is deleted after the pipeline started, the value chosen here will not exist in the target and the foreign key constraint rejects it. Nothing checks this for you. During a snapshot the constraints are restored after the data, so the failure arrives at the end of the load as a constraint violation that names the constraint rather than this rule. Under replication a constraint violation is not retried: in the default mode the row is **dropped** with a `DATALOSS` log line and the checkpoint advances, so the pipeline keeps running without it, and with `deterministic` every row sharing that original value is dropped too. Under `strict_mode` the pipeline stops instead. Leave the lookup table's key column untransformed, with `noop` if the validation mode requires a rule for it.
+
+⚠️ **The mapping only holds while the lookup column does.** The value is chosen by position in the loaded list, so inserting or deleting a single row in the lookup table — or editing `ignore_values` — remaps almost every input the next time the pipeline starts. Deterministic mode is reproducible across restarts for a **fixed** lookup set; it is not stable across changes to it. New rows in the lookup table are not picked up until a restart.
+
+⚠️ **The whole column is loaded into memory, once per rule.** There is no limit on how many values are read, and no paging, and each column rule runs its own query: three columns reading the same lookup table load it three times. Point this at a lookup table, not at a large one. The read is given 30 seconds, so a locked or unreachable lookup table fails startup instead of hanging it.
+
+**Example Configuration:**
+
+```yaml
+transformations:
+ table_transformers:
+ - schema: public
+ table: addresses
+ column_transformers:
+ country_id:
+ name: lookup_choice
+ parameters:
+ lookup_table: public.countries
+ lookup_column: id
+ generator: deterministic
+ ignore_values: [0, -1]
+```
+
+**Input-Output Examples:**
+
+Given a `public.countries` table whose `id` column holds `1, 2, 3`:
+
+| Input Value | Configuration Parameters | Output Value |
+| ----------- | -------------------------- | --------------------- |
+| `7` | `generator: random` | `3` (random) |
+| `7` | `generator: deterministic` | `2` |
+| `7` | `generator: deterministic` | `2` (again, next run) |
+| `8` | `generator: deterministic` | `1` |
+
diff --git a/pkg/transformers/builder/main_test.go b/pkg/transformers/builder/main_test.go
new file mode 100644
index 000000000..4534a71ed
--- /dev/null
+++ b/pkg/transformers/builder/main_test.go
@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package builder
+
+import (
+ "os"
+ "testing"
+
+ "github.com/jackc/pgx/v5/pgtype"
+ "github.com/xataio/pgstream/pkg/transformers/internal/lookup"
+)
+
+// lookupTestValues is what lookup_choice reads in this package's tests
+var lookupTestValues = []any{int64(1), int64(2)}
+
+// lookup_choice reads its values while it is being built, and the tables in
+// this package build every registered transformer. The seam is installed here
+// rather than inside a test because those tests run in parallel and would
+// otherwise race each other writing it.
+func TestMain(m *testing.M) {
+ lookup.NewQuerier = lookup.StubQuerier(pgtype.Int8OID, lookupTestValues, lookup.StubOptions{})
+ os.Exit(m.Run())
+}
diff --git a/pkg/transformers/builder/transformer_builder.go b/pkg/transformers/builder/transformer_builder.go
index 690a511be..3ab466222 100644
--- a/pkg/transformers/builder/transformer_builder.go
+++ b/pkg/transformers/builder/transformer_builder.go
@@ -90,6 +90,12 @@ var TransformersMap = map[transformers.TransformerType]struct {
return transformers.NewEncryptedAESSIVTransformer(cfg.Parameters)
},
},
+ transformers.LookupChoice: {
+ Definition: transformers.LookupChoiceTransformerDefinition(),
+ BuildFn: func(cfg *transformers.Config) (transformers.Transformer, error) {
+ return transformers.NewLookupChoiceTransformer(cfg.Parameters)
+ },
+ },
transformers.FPEFF1: {
Definition: transformers.FPEFF1TransformerDefinition(),
BuildFn: func(cfg *transformers.Config) (transformers.Transformer, error) {
diff --git a/pkg/transformers/builder/transformer_concurrency_test.go b/pkg/transformers/builder/transformer_concurrency_test.go
index 92a2f5129..743880748 100644
--- a/pkg/transformers/builder/transformer_concurrency_test.go
+++ b/pkg/transformers/builder/transformer_concurrency_test.go
@@ -186,6 +186,18 @@ func TestTransformers_ConcurrentTransform(t *testing.T) {
transformers.PGAnonymizer: {
skip: "requires a live PostgreSQL connection",
},
+ transformers.LookupChoice: {
+ params: transformers.ParameterValues{
+ "lookup_table": "public.countries",
+ "lookup_column": "id",
+ "postgres_url": "postgres://user:pass@localhost:5432/db",
+ },
+ input: "hello",
+ supportsGenerators: true,
+ validate: func(t *testing.T, got any) {
+ require.Contains(t, lookupTestValues, got)
+ },
+ },
}
// every registered transformer must have a concurrency test, so that new
diff --git a/pkg/transformers/builder/transformer_uniqueness_test.go b/pkg/transformers/builder/transformer_uniqueness_test.go
index 45c6f8a28..902877ef8 100644
--- a/pkg/transformers/builder/transformer_uniqueness_test.go
+++ b/pkg/transformers/builder/transformer_uniqueness_test.go
@@ -41,6 +41,7 @@ func TestTransformers_UniquenessMatchesDefinition(t *testing.T) {
transformers.NeosyncFullName: {},
transformers.NeosyncEmail: {},
transformers.PGAnonymizer: {"anon_function": "anon.fake_email()", "postgres_url": "postgres://user:pass@localhost:5432/db"},
+ transformers.LookupChoice: {"lookup_table": "public.countries", "lookup_column": "id", "postgres_url": "postgres://user:pass@localhost:5432/db"},
}
for transformerType := range TransformersMap {
diff --git a/pkg/transformers/internal/lookup/lookup.go b/pkg/transformers/internal/lookup/lookup.go
new file mode 100644
index 000000000..a678529fb
--- /dev/null
+++ b/pkg/transformers/internal/lookup/lookup.go
@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: Apache-2.0
+
+// Package lookup holds the connection seam used by transformers that read
+// their values from the database when they are built. It exists as its own
+// package because pkg/transformers/builder's table driven tests construct
+// every registered transformer and have no database, and they cannot reach an
+// unexported variable in pkg/transformers.
+package lookup
+
+import (
+ "context"
+
+ pglib "github.com/xataio/pgstream/internal/postgres"
+)
+
+// NewQuerier opens the connection a lookup load runs on. It is a variable so
+// that tests can build lookup based transformers without a database; replace
+// it with StubQuerier.
+var NewQuerier = func(ctx context.Context, url string) (pglib.Querier, error) {
+ return pglib.NewConnPool(ctx, url)
+}
diff --git a/pkg/transformers/internal/lookup/stub.go b/pkg/transformers/internal/lookup/stub.go
new file mode 100644
index 000000000..efd6e1c3d
--- /dev/null
+++ b/pkg/transformers/internal/lookup/stub.go
@@ -0,0 +1,71 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package lookup
+
+import (
+ "context"
+ "errors"
+
+ "github.com/jackc/pgx/v5/pgconn"
+ pglib "github.com/xataio/pgstream/internal/postgres"
+ pglibmocks "github.com/xataio/pgstream/internal/postgres/mocks"
+)
+
+// StubOptions configures the failure a StubQuerier injects. The zero value
+// serves the values successfully.
+type StubOptions struct {
+ NewQuerierErr error
+ QueryErr error
+ ScanErr error
+ RowsErr error
+ // Closed reports whether the querier and its rows were closed
+ Closed *bool
+ // Query captures the SQL the loader built
+ Query *string
+}
+
+// StubQuerier returns a NewQuerier replacement serving the given values as a
+// single column of the given pg type OID. It lives beside the seam so that
+// every package testing a lookup based transformer shares one test double.
+func StubQuerier(oid uint32, values []any, opts StubOptions) func(context.Context, string) (pglib.Querier, error) {
+ return func(context.Context, string) (pglib.Querier, error) {
+ if opts.NewQuerierErr != nil {
+ return nil, opts.NewQuerierErr
+ }
+ return &pglibmocks.Querier{
+ QueryFn: func(_ context.Context, _ uint, query string, _ ...any) (pglib.Rows, error) {
+ if opts.Query != nil {
+ *opts.Query = query
+ }
+ if opts.QueryErr != nil {
+ return nil, opts.QueryErr
+ }
+ return &pglibmocks.Rows{
+ FieldDescriptionsFn: func() []pgconn.FieldDescription {
+ return []pgconn.FieldDescription{{Name: "lookup", DataTypeOID: oid}}
+ },
+ NextFn: func(i uint) bool { return i <= uint(len(values)) },
+ ScanFn: func(i uint, dest ...any) error {
+ if opts.ScanErr != nil {
+ return opts.ScanErr
+ }
+ value, ok := dest[0].(*any)
+ if !ok {
+ return errors.New("unexpected scan destination")
+ }
+ *value = values[i-1]
+ return nil
+ },
+ ErrFn: func() error { return opts.RowsErr },
+ CloseFn: func() {},
+ }, nil
+ },
+ CloseFn: func(context.Context) error {
+ if opts.Closed != nil {
+ *opts.Closed = true
+ }
+ return nil
+ },
+ }, nil
+ }
+}
diff --git a/pkg/transformers/lookup_choice_transformer.go b/pkg/transformers/lookup_choice_transformer.go
new file mode 100644
index 000000000..8a872558f
--- /dev/null
+++ b/pkg/transformers/lookup_choice_transformer.go
@@ -0,0 +1,432 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package transformers
+
+import (
+ "context"
+ "encoding/binary"
+ "encoding/hex"
+ "errors"
+ "fmt"
+ "time"
+
+ "github.com/jackc/pgx/v5/pgtype"
+ pglib "github.com/xataio/pgstream/internal/postgres"
+ "github.com/xataio/pgstream/pkg/transformers/generators"
+ "github.com/xataio/pgstream/pkg/transformers/internal/lookup"
+)
+
+// LookupChoiceTransformer replaces a value with one taken from a column of
+// another table. The values are read once, when the transformer is built, so
+// the configuration does not have to be regenerated when the lookup table
+// changes.
+type LookupChoiceTransformer struct {
+ values []any
+ generator generators.Generator
+ compatibleTypes []SupportedDataType
+ // dateColumn renders time.Time values the way the replication path
+ // delivers a date, so both ingestion paths hash to the same key
+ dateColumn bool
+}
+
+const (
+ // the index is read from the first 8 bytes the generator produces
+ lookupChoiceIndexSize = 8
+ // transformer constructors receive no context, so the load carries its own
+ // deadline rather than blocking startup indefinitely on a locked or
+ // unreachable lookup table
+ lookupChoiceLoadTimeout = 30 * time.Second
+
+ randomGenerator = "random"
+ deterministicGenerator = "deterministic"
+)
+
+var (
+ errLookupTableNotFound = errors.New("lookup_choice: lookup_table must be provided")
+ errLookupColumnNotFound = errors.New("lookup_choice: lookup_column must be provided")
+ errLookupURLNotFound = errors.New("lookup_choice: postgres_url must be provided")
+ errLookupNoValues = errors.New("lookup_choice: no values loaded")
+
+ lookupChoiceCompatibleTypes = []SupportedDataType{
+ StringDataType,
+ CitextDataType,
+ ByteArrayDataType,
+ BooleanDataType,
+ Integer16DataType,
+ Integer32DataType,
+ Integer64DataType,
+ Float32DataType,
+ Float64DataType,
+ UInt8ArrayOf16DataType,
+ DateDataType,
+ DatetimeDataType,
+ }
+
+ lookupChoiceParams = []Parameter{
+ {
+ Name: "lookup_table",
+ SupportedType: "string",
+ Default: nil,
+ Dynamic: false,
+ Required: true,
+ },
+ {
+ Name: "lookup_column",
+ SupportedType: "string",
+ Default: nil,
+ Dynamic: false,
+ Required: true,
+ },
+ {
+ Name: "postgres_url",
+ SupportedType: "string",
+ Default: nil,
+ Dynamic: false,
+ Required: true,
+ },
+ {
+ Name: "generator",
+ SupportedType: "string",
+ Default: "random",
+ Dynamic: false,
+ Required: false,
+ Values: []any{randomGenerator, deterministicGenerator},
+ },
+ {
+ Name: "ignore_values",
+ SupportedType: "array",
+ Default: nil,
+ Dynamic: false,
+ Required: false,
+ },
+ }
+)
+
+// NewLookupChoiceTransformer reads the values of the configured lookup column
+// and returns a transformer that chooses from them. The connection is only
+// held for the duration of the read.
+func NewLookupChoiceTransformer(params ParameterValues) (*LookupChoiceTransformer, error) {
+ table, found, err := FindParameter[string](params, "lookup_table")
+ if err != nil {
+ return nil, fmt.Errorf("lookup_choice: lookup_table must be a string: %w", err)
+ }
+ if !found || table == "" {
+ return nil, errLookupTableNotFound
+ }
+
+ column, found, err := FindParameter[string](params, "lookup_column")
+ if err != nil {
+ return nil, fmt.Errorf("lookup_choice: lookup_column must be a string: %w", err)
+ }
+ if !found || column == "" {
+ return nil, errLookupColumnNotFound
+ }
+
+ url, found, err := FindParameter[string](params, "postgres_url")
+ if err != nil {
+ return nil, fmt.Errorf("lookup_choice: postgres_url must be a string: %w", err)
+ }
+ // an empty URL would otherwise reach pgx, which falls back to the libpq
+ // environment defaults and reads whichever database those point at
+ if !found || url == "" {
+ return nil, errLookupURLNotFound
+ }
+
+ ignoreValues, _, err := FindParameterArray[any](params, "ignore_values")
+ if err != nil {
+ return nil, fmt.Errorf("lookup_choice: ignore_values must be an array: %w", err)
+ }
+
+ generatorType, err := FindParameterWithDefault(params, "generator", randomGenerator)
+ if err != nil {
+ return nil, fmt.Errorf("lookup_choice: generator must be a string: %w", err)
+ }
+ // validated before the load so that a typo costs a config error rather
+ // than a full table scan
+ if generatorType != randomGenerator && generatorType != deterministicGenerator {
+ return nil, fmt.Errorf("lookup_choice: generator must be one of 'random' or 'deterministic': %w", ErrInvalidParameters)
+ }
+
+ values, columnOID, err := loadLookupValues(url, table, column)
+ if err != nil {
+ return nil, err
+ }
+
+ compatibleTypes, err := lookupColumnTypes(columnOID, values[0])
+ if err != nil {
+ return nil, fmt.Errorf("lookup_choice: column %s of table %s: %w", column, table, err)
+ }
+
+ // the replication path delivers a timestamp as the text wal2json emits
+ // while a snapshot delivers pgx's time.Time, and the two cannot be
+ // rendered to a common form reliably, so the same row would hash to
+ // different values either side of the cutover
+ if generatorType == deterministicGenerator && (columnOID == pgtype.TimestampOID || columnOID == pgtype.TimestamptzOID) {
+ return nil, fmt.Errorf("lookup_choice: column %s of table %s is a timestamp, which the deterministic generator cannot key on consistently across snapshot and replication: %w",
+ column, table, ErrInvalidParameters)
+ }
+
+ values, err = removeIgnoredValues(values, ignoreValues, columnOID == pgtype.DateOID)
+ if err != nil {
+ return nil, err
+ }
+ if len(values) == 0 {
+ return nil, fmt.Errorf("lookup_choice: every value in column %s of table %s is excluded by ignore_values", column, table)
+ }
+
+ return newLookupChoiceTransformer(values, generatorType, compatibleTypes, columnOID == pgtype.DateOID)
+}
+
+// newLookupChoiceTransformer builds the transformer around an already loaded
+// list of values, so that the choosing logic can be tested without a database.
+func newLookupChoiceTransformer(values []any, generatorType string, compatibleTypes []SupportedDataType, dateColumn bool) (*LookupChoiceTransformer, error) {
+ if len(values) == 0 {
+ return nil, errLookupNoValues
+ }
+
+ var generator generators.Generator
+ var err error
+ switch generatorType {
+ case deterministicGenerator:
+ generator, err = generators.NewDeterministicBytesGenerator(lookupChoiceIndexSize)
+ if err != nil {
+ return nil, fmt.Errorf("lookup_choice: error creating deterministic generator: %w", err)
+ }
+ case randomGenerator:
+ generator = generators.NewRandomBytesGenerator(lookupChoiceIndexSize)
+ default:
+ return nil, fmt.Errorf("lookup_choice: generator must be one of 'random' or 'deterministic': %w", ErrInvalidParameters)
+ }
+
+ return &LookupChoiceTransformer{
+ values: values,
+ generator: generator,
+ compatibleTypes: compatibleTypes,
+ dateColumn: dateColumn,
+ }, nil
+}
+
+func (t *LookupChoiceTransformer) Transform(_ context.Context, value Value) (any, error) {
+ index, err := t.generator.Generate([]byte(lookupValueKey(value.TransformValue, t.dateColumn)))
+ if err != nil {
+ return nil, fmt.Errorf("lookup_choice: generating value index: %w", err)
+ }
+ if len(index) < lookupChoiceIndexSize {
+ return nil, fmt.Errorf("lookup_choice: generated index is %d bytes, expected %d", len(index), lookupChoiceIndexSize)
+ }
+
+ return t.values[binary.BigEndian.Uint64(index[:lookupChoiceIndexSize])%uint64(len(t.values))], nil
+}
+
+// CompatibleTypes reports the pg types the lookup column's values can be
+// written to, so that a rule pointing a column at a lookup column of an
+// incompatible type fails on startup rather than mid load.
+func (t *LookupChoiceTransformer) CompatibleTypes() []SupportedDataType {
+ return t.compatibleTypes
+}
+
+func (t *LookupChoiceTransformer) Type() TransformerType {
+ return LookupChoice
+}
+
+func (t *LookupChoiceTransformer) IsDynamic() bool {
+ return false
+}
+
+// Uniqueness is lossy in both generator modes: the values come from a set that
+// is normally much smaller than the number of rows being transformed.
+func (t *LookupChoiceTransformer) Uniqueness() Uniqueness {
+ return UniquenessLossy
+}
+
+// Close is a no-op: the connection used to read the values is released as soon
+// as the read completes.
+func (t *LookupChoiceTransformer) Close() error {
+ return nil
+}
+
+func LookupChoiceTransformerDefinition() *Definition {
+ return &Definition{
+ SupportedTypes: lookupChoiceCompatibleTypes,
+ Parameters: lookupChoiceParams,
+ Uniqueness: UniquenessLossy,
+ }
+}
+
+// loadLookupValues reads the lookup column, returning its values and the pg
+// type OID Postgres reported for it.
+func loadLookupValues(url, table, column string) ([]any, uint32, error) {
+ qualifiedName, err := pglib.NewQualifiedName(table)
+ if err != nil {
+ return nil, 0, fmt.Errorf("lookup_choice: invalid lookup_table %q: %w", table, err)
+ }
+ schema := qualifiedName.Schema()
+ if schema == "" {
+ schema = "public"
+ }
+
+ // BuildFn receives no context, so the caller's cancellation cannot reach
+ // this load; the deadline is what keeps a locked lookup table from
+ // blocking startup with no way out but SIGKILL
+ ctx, cancel := context.WithTimeout(context.Background(), lookupChoiceLoadTimeout)
+ defer cancel()
+
+ querier, err := lookup.NewQuerier(ctx, url)
+ if err != nil {
+ return nil, 0, fmt.Errorf("lookup_choice: creating connection pool: %w", err)
+ }
+ defer querier.Close(ctx)
+
+ // the order is explicit because the deterministic generator picks an index
+ // into this slice, and an unordered scan can return the rows differently on
+ // every run
+ quotedColumn := pglib.QuoteIdentifier(column)
+ query := fmt.Sprintf("SELECT %s FROM %s WHERE %s IS NOT NULL ORDER BY %s",
+ quotedColumn, pglib.QuoteQualifiedIdentifier(schema, qualifiedName.Name()), quotedColumn, quotedColumn)
+
+ rows, err := querier.Query(ctx, query)
+ if err != nil {
+ return nil, 0, fmt.Errorf("lookup_choice: querying column %s of table %s: %w", column, table, err)
+ }
+ defer rows.Close()
+
+ var columnOID uint32
+ if fields := rows.FieldDescriptions(); len(fields) > 0 {
+ columnOID = fields[0].DataTypeOID
+ }
+
+ var values []any
+ for rows.Next() {
+ var value any
+ if err := rows.Scan(&value); err != nil {
+ return nil, 0, fmt.Errorf("lookup_choice: scanning column %s of table %s: %w", column, table, err)
+ }
+ values = append(values, value)
+ }
+ if err := rows.Err(); err != nil {
+ return nil, 0, fmt.Errorf("lookup_choice: reading column %s of table %s: %w", column, table, err)
+ }
+ if len(values) == 0 {
+ return nil, 0, fmt.Errorf("lookup_choice: no values found in column %s of table %s", column, table)
+ }
+
+ return values, columnOID, nil
+}
+
+// removeIgnoredValues drops the configured values from the loaded list. An
+// ignore value that matches nothing is an error: it is either a typo or a
+// value written in a form the column never produces, and silently ignoring it
+// leaves the excluded row in the choice set.
+func removeIgnoredValues(values, ignoreValues []any, dateColumn bool) ([]any, error) {
+ if len(ignoreValues) == 0 {
+ return values, nil
+ }
+
+ // the ignored values come from the configuration and the lookup values
+ // from the database, so they can represent the same value with different
+ // Go types
+ matched := make(map[string]bool, len(ignoreValues))
+ for _, value := range ignoreValues {
+ matched[lookupValueKey(value, dateColumn)] = false
+ }
+
+ kept := make([]any, 0, len(values))
+ for _, value := range values {
+ key := lookupValueKey(value, dateColumn)
+ if _, found := matched[key]; found {
+ matched[key] = true
+ continue
+ }
+ kept = append(kept, value)
+ }
+
+ for key, found := range matched {
+ if !found {
+ return nil, fmt.Errorf("lookup_choice: ignore_values entry %q matches no value in the lookup column: %w", key, ErrInvalidParameters)
+ }
+ }
+
+ return kept, nil
+}
+
+// lookupValueKey gives a value a canonical representation, used both to match
+// ignore_values and to feed the generator. A snapshot delivers the values pgx
+// decodes while replication delivers the text wal2json emits, so the two have
+// to be rendered to the same string or the same row would be mapped
+// differently either side of the cutover.
+func lookupValueKey(value any, dateColumn bool) string {
+ switch v := value.(type) {
+ case string:
+ return v
+ case []byte:
+ return string(v)
+ case [16]byte:
+ return encodeUUID(v)
+ case time.Time:
+ if dateColumn {
+ return v.Format(time.DateOnly)
+ }
+ return v.Format(time.RFC3339)
+ default:
+ return fmt.Sprintf("%v", v)
+ }
+}
+
+// encodeUUID renders the [16]byte pgx decodes a uuid into as the canonical
+// text form replication delivers.
+func encodeUUID(value [16]byte) string {
+ buf := make([]byte, 36)
+ hex.Encode(buf[0:8], value[0:4])
+ buf[8] = '-'
+ hex.Encode(buf[9:13], value[4:6])
+ buf[13] = '-'
+ hex.Encode(buf[14:18], value[6:8])
+ buf[18] = '-'
+ hex.Encode(buf[19:23], value[8:10])
+ buf[23] = '-'
+ hex.Encode(buf[24:36], value[10:16])
+ return string(buf)
+}
+
+// lookupColumnTypes maps the pg type of the lookup column to the column types
+// its values can be written to. Postgres reports the OID, which is exact;
+// extension types have no fixed OID, so those fall back to the Go type pgx
+// decoded. A type that matches neither is rejected, because silently
+// accepting it would disable the compatibility check the parser relies on.
+func lookupColumnTypes(columnOID uint32, sample any) ([]SupportedDataType, error) {
+ switch columnOID {
+ case pgtype.TextOID, pgtype.VarcharOID, pgtype.BPCharOID:
+ return []SupportedDataType{StringDataType, CitextDataType}, nil
+ case pgtype.BoolOID:
+ return []SupportedDataType{BooleanDataType}, nil
+ // a narrower integer or float is assignable to a wider column, which is
+ // the ordinary shape of a foreign key referencing a serial primary key
+ case pgtype.Int2OID:
+ return []SupportedDataType{Integer16DataType, Integer32DataType, Integer64DataType}, nil
+ case pgtype.Int4OID:
+ return []SupportedDataType{Integer32DataType, Integer64DataType}, nil
+ case pgtype.Int8OID:
+ return []SupportedDataType{Integer64DataType}, nil
+ case pgtype.Float4OID:
+ return []SupportedDataType{Float32DataType, Float64DataType}, nil
+ case pgtype.Float8OID:
+ return []SupportedDataType{Float64DataType}, nil
+ case pgtype.UUIDOID:
+ return []SupportedDataType{UInt8ArrayOf16DataType}, nil
+ case pgtype.ByteaOID:
+ return []SupportedDataType{ByteArrayDataType}, nil
+ case pgtype.DateOID:
+ return []SupportedDataType{DateDataType}, nil
+ case pgtype.TimestampOID, pgtype.TimestamptzOID:
+ return []SupportedDataType{DatetimeDataType}, nil
+ }
+
+ switch sample.(type) {
+ case string:
+ return []SupportedDataType{StringDataType, CitextDataType}, nil
+ case []byte:
+ return []SupportedDataType{ByteArrayDataType}, nil
+ default:
+ return nil, fmt.Errorf("unsupported lookup column type with OID %d: %w", columnOID, ErrInvalidParameters)
+ }
+}
diff --git a/pkg/transformers/lookup_choice_transformer_test.go b/pkg/transformers/lookup_choice_transformer_test.go
new file mode 100644
index 000000000..07efb247d
--- /dev/null
+++ b/pkg/transformers/lookup_choice_transformer_test.go
@@ -0,0 +1,446 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package transformers
+
+import (
+ "context"
+ "errors"
+ "testing"
+ "time"
+
+ "github.com/jackc/pgx/v5/pgtype"
+ "github.com/stretchr/testify/require"
+ "github.com/xataio/pgstream/pkg/transformers/internal/lookup"
+)
+
+// setLookupQuerier points the load at a stub serving the given values. The
+// tests that call it must not run in parallel: the seam they replace is
+// package level.
+func setLookupQuerier(t *testing.T, oid uint32, values []any, opts lookup.StubOptions) {
+ t.Helper()
+ original := lookup.NewQuerier
+ t.Cleanup(func() { lookup.NewQuerier = original })
+ lookup.NewQuerier = lookup.StubQuerier(oid, values, opts)
+}
+
+func testUUID(last byte) [16]byte {
+ return [16]byte{0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff, 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, last}
+}
+
+func TestNewLookupChoiceTransformer(t *testing.T) {
+ errTest := errors.New("oh noes")
+
+ validParams := func() ParameterValues {
+ return ParameterValues{
+ "lookup_table": "reference.countries",
+ "lookup_column": "id",
+ "postgres_url": "postgres://user:pass@localhost:5432/db",
+ }
+ }
+ withParam := func(name string, value any) ParameterValues {
+ params := validParams()
+ params[name] = value
+ return params
+ }
+
+ tests := []struct {
+ name string
+ params ParameterValues
+ oid uint32
+ values []any
+ opts lookup.StubOptions
+ wantValues []any
+ wantTypes []SupportedDataType
+ wantErr error
+ }{
+ {
+ name: "ok - values loaded",
+ params: validParams(),
+ oid: pgtype.Int8OID,
+ values: []any{int64(1), int64(2), int64(3)},
+ wantValues: []any{int64(1), int64(2), int64(3)},
+ wantTypes: []SupportedDataType{Integer64DataType},
+ },
+ {
+ name: "ok - a narrower integer column is assignable to a wider one",
+ params: validParams(),
+ oid: pgtype.Int4OID,
+ values: []any{int32(1)},
+ wantValues: []any{int32(1)},
+ wantTypes: []SupportedDataType{Integer32DataType, Integer64DataType},
+ },
+ {
+ name: "ok - an extension type falls back to the decoded Go type",
+ params: validParams(),
+ oid: 16385, // citext has no fixed OID
+ values: []any{"one"},
+ wantValues: []any{"one"},
+ wantTypes: []SupportedDataType{StringDataType, CitextDataType},
+ },
+ {
+ name: "ok - ignore_values removes matching values",
+ // the configuration parses ids as int, while the database
+ // returns them as int64
+ params: withParam("ignore_values", []any{1, 3}),
+ oid: pgtype.Int8OID,
+ values: []any{int64(1), int64(2), int64(3)},
+ wantValues: []any{int64(2)},
+ wantTypes: []SupportedDataType{Integer64DataType},
+ },
+ {
+ name: "ok - ignore_values matches a uuid written as text",
+ params: withParam("ignore_values", []any{"aabbccdd-eeff-0011-2233-445566778801"}),
+ oid: pgtype.UUIDOID,
+ values: []any{testUUID(0x01), testUUID(0x02)},
+ wantValues: []any{testUUID(0x02)},
+ wantTypes: []SupportedDataType{UInt8ArrayOf16DataType},
+ },
+ {
+ name: "ok - ignore_values matches a date written as text",
+ params: withParam("ignore_values", []any{"2024-01-01"}),
+ oid: pgtype.DateOID,
+ values: []any{time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC), time.Date(2024, 1, 2, 0, 0, 0, 0, time.UTC)},
+ wantValues: []any{time.Date(2024, 1, 2, 0, 0, 0, 0, time.UTC)},
+ wantTypes: []SupportedDataType{DateDataType},
+ },
+ {
+ name: "error - ignore_values entry matches nothing",
+ params: withParam("ignore_values", []any{4}),
+ oid: pgtype.Int8OID,
+ values: []any{int64(1), int64(2)},
+ wantErr: ErrInvalidParameters,
+ },
+ {
+ name: "error - lookup table is empty",
+ params: validParams(),
+ oid: pgtype.Int8OID,
+ values: []any{},
+ wantErr: errors.New("lookup_choice: no values found in column id of table reference.countries"),
+ },
+ {
+ name: "error - all values ignored",
+ params: withParam("ignore_values", []any{1, 2}),
+ oid: pgtype.Int8OID,
+ values: []any{int64(1), int64(2)},
+ wantErr: errors.New("lookup_choice: every value in column id of table reference.countries is excluded by ignore_values"),
+ },
+ {
+ name: "error - unsupported lookup column type",
+ params: validParams(),
+ oid: pgtype.NumericOID,
+ values: []any{pgtype.Numeric{}},
+ wantErr: ErrInvalidParameters,
+ },
+ {
+ name: "error - deterministic generator on a timestamp column",
+ params: withParam("generator", "deterministic"),
+ oid: pgtype.TimestamptzOID,
+ values: []any{time.Now()},
+ wantErr: ErrInvalidParameters,
+ },
+ {
+ name: "ok - random generator on a timestamp column",
+ params: withParam("generator", "random"),
+ oid: pgtype.TimestampOID,
+ values: []any{time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)},
+ wantValues: []any{time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)},
+ wantTypes: []SupportedDataType{DatetimeDataType},
+ },
+ {
+ name: "error - lookup_table missing",
+ params: ParameterValues{
+ "lookup_column": "id",
+ "postgres_url": "postgres://user:pass@localhost:5432/db",
+ },
+ wantErr: errLookupTableNotFound,
+ },
+ {
+ name: "error - lookup_column missing",
+ params: ParameterValues{
+ "lookup_table": "reference.countries",
+ "postgres_url": "postgres://user:pass@localhost:5432/db",
+ },
+ wantErr: errLookupColumnNotFound,
+ },
+ {
+ name: "error - postgres_url missing",
+ params: ParameterValues{
+ "lookup_table": "reference.countries",
+ "lookup_column": "id",
+ },
+ wantErr: errLookupURLNotFound,
+ },
+ {
+ name: "error - postgres_url empty",
+ params: withParam("postgres_url", ""),
+ wantErr: errLookupURLNotFound,
+ },
+ {
+ name: "error - unknown generator",
+ params: withParam("generator", "sequential"),
+ oid: pgtype.Int8OID,
+ values: []any{int64(1)},
+ wantErr: ErrInvalidParameters,
+ },
+ {
+ name: "error - ignore_values is not a list",
+ params: withParam("ignore_values", 1),
+ oid: pgtype.Int8OID,
+ values: []any{int64(1)},
+ wantErr: ErrInvalidParameters,
+ },
+ {
+ name: "error - lookup_table is not a valid qualified name",
+ params: withParam("lookup_table", "a.b.c"),
+ oid: pgtype.Int8OID,
+ values: []any{int64(1)},
+ wantErr: errors.New(`lookup_choice: invalid lookup_table "a.b.c": unexpected qualified name format`),
+ },
+ {
+ name: "error - opening the connection fails",
+ params: validParams(),
+ opts: lookup.StubOptions{NewQuerierErr: errTest},
+ wantErr: errTest,
+ },
+ {
+ name: "error - query fails",
+ params: validParams(),
+ opts: lookup.StubOptions{QueryErr: errTest},
+ wantErr: errTest,
+ },
+ {
+ name: "error - scanning a row fails",
+ params: validParams(),
+ oid: pgtype.Int8OID,
+ values: []any{int64(1)},
+ opts: lookup.StubOptions{ScanErr: errTest},
+ wantErr: errTest,
+ },
+ {
+ name: "error - iterating the rows fails",
+ params: validParams(),
+ oid: pgtype.Int8OID,
+ values: []any{int64(1)},
+ opts: lookup.StubOptions{RowsErr: errTest},
+ wantErr: errTest,
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ setLookupQuerier(t, tc.oid, tc.values, tc.opts)
+
+ transformer, err := NewLookupChoiceTransformer(tc.params)
+ if tc.wantErr != nil {
+ require.Error(t, err)
+ if !errors.Is(err, tc.wantErr) {
+ require.EqualError(t, err, tc.wantErr.Error())
+ }
+ return
+ }
+ require.NoError(t, err)
+ require.Equal(t, tc.wantValues, transformer.values)
+ require.Equal(t, tc.wantTypes, transformer.CompatibleTypes())
+ })
+ }
+}
+
+// the load holds a connection only while it reads, so a leak would keep a
+// pool open for the lifetime of the process
+func TestNewLookupChoiceTransformer_closesTheConnection(t *testing.T) {
+ errTest := errors.New("oh noes")
+
+ tests := map[string]lookup.StubOptions{
+ "after a successful load": {},
+ "after a failed read": {RowsErr: errTest},
+ }
+
+ for name, opts := range tests {
+ t.Run(name, func(t *testing.T) {
+ closed := false
+ opts.Closed = &closed
+ setLookupQuerier(t, pgtype.Int8OID, []any{int64(1)}, opts)
+
+ _, err := NewLookupChoiceTransformer(ParameterValues{
+ "lookup_table": "reference.countries",
+ "lookup_column": "id",
+ "postgres_url": "postgres://user:pass@localhost:5432/db",
+ })
+ if opts.RowsErr != nil {
+ require.Error(t, err)
+ } else {
+ require.NoError(t, err)
+ }
+ require.True(t, closed, "the lookup connection was not closed")
+ })
+ }
+}
+
+func TestNewLookupChoiceTransformer_query(t *testing.T) {
+ var gotQuery string
+ setLookupQuerier(t, pgtype.Int8OID, []any{int64(1)}, lookup.StubOptions{Query: &gotQuery})
+
+ _, err := NewLookupChoiceTransformer(ParameterValues{
+ "lookup_table": "reference.countries",
+ "lookup_column": "id",
+ "postgres_url": "postgres://user:pass@localhost:5432/db",
+ })
+ require.NoError(t, err)
+ require.Equal(t, `SELECT "id" FROM "reference"."countries" WHERE "id" IS NOT NULL ORDER BY "id"`, gotQuery)
+
+ // an unqualified table name defaults to the public schema
+ _, err = NewLookupChoiceTransformer(ParameterValues{
+ "lookup_table": "countries",
+ "lookup_column": "id",
+ "postgres_url": "postgres://user:pass@localhost:5432/db",
+ })
+ require.NoError(t, err)
+ require.Equal(t, `SELECT "id" FROM "public"."countries" WHERE "id" IS NOT NULL ORDER BY "id"`, gotQuery)
+}
+
+func TestLookupChoiceTransformer_Transform(t *testing.T) {
+ t.Parallel()
+
+ values := []any{int64(10), int64(20), int64(30), int64(40)}
+ inputs := []string{"a", "b", "c", "d", "e", "f", "g", "h"}
+ intTypes := []SupportedDataType{Integer64DataType}
+
+ newTransformer := func(t *testing.T, values []any, generatorType string) *LookupChoiceTransformer {
+ t.Helper()
+ transformer, err := newLookupChoiceTransformer(values, generatorType, intTypes, false)
+ require.NoError(t, err)
+ return transformer
+ }
+
+ t.Run("random - every output comes from the lookup values", func(t *testing.T) {
+ t.Parallel()
+
+ transformer := newTransformer(t, values, "random")
+ seen := map[any]bool{}
+ for range 100 {
+ got, err := transformer.Transform(context.Background(), NewValue("a", "int8", nil))
+ require.NoError(t, err)
+ require.Contains(t, values, got)
+ seen[got] = true
+ }
+ // the choice must not collapse onto a single value
+ require.Greater(t, len(seen), 1)
+ })
+
+ t.Run("deterministic - the same input always gives the same value", func(t *testing.T) {
+ t.Parallel()
+
+ transformer := newTransformer(t, values, "deterministic")
+ // a separately built transformer over the same values must agree, so
+ // that the mapping survives a restart
+ other := newTransformer(t, values, "deterministic")
+
+ seen := map[any]bool{}
+ for _, input := range inputs {
+ want, err := transformer.Transform(context.Background(), NewValue(input, "text", nil))
+ require.NoError(t, err)
+ require.Contains(t, values, want)
+ seen[want] = true
+
+ for range 10 {
+ got, err := transformer.Transform(context.Background(), NewValue(input, "text", nil))
+ require.NoError(t, err)
+ require.Equal(t, want, got)
+ }
+
+ got, err := other.Transform(context.Background(), NewValue(input, "text", nil))
+ require.NoError(t, err)
+ require.Equal(t, want, got)
+ }
+ require.Greater(t, len(seen), 1)
+ })
+
+ t.Run("deterministic - the mapping is pinned to these values", func(t *testing.T) {
+ t.Parallel()
+
+ // the worked example in docs/transformers.md quotes these outputs; a
+ // change to the hash or the index arithmetic must not pass silently
+ transformer := newTransformer(t, []any{int64(1), int64(2), int64(3)}, "deterministic")
+
+ got, err := transformer.Transform(context.Background(), NewValue(int64(7), "int8", nil))
+ require.NoError(t, err)
+ require.Equal(t, int64(2), got)
+
+ got, err = transformer.Transform(context.Background(), NewValue(int64(8), "int8", nil))
+ require.NoError(t, err)
+ require.Equal(t, int64(1), got)
+ })
+
+ t.Run("deterministic - equal values map to equal values regardless of type", func(t *testing.T) {
+ t.Parallel()
+
+ transformer := newTransformer(t, values, "deterministic")
+
+ fromString, err := transformer.Transform(context.Background(), NewValue("42", "text", nil))
+ require.NoError(t, err)
+ fromBytes, err := transformer.Transform(context.Background(), NewValue([]byte("42"), "bytea", nil))
+ require.NoError(t, err)
+ fromInt, err := transformer.Transform(context.Background(), NewValue(int64(42), "int8", nil))
+ require.NoError(t, err)
+
+ require.Equal(t, fromString, fromBytes)
+ require.Equal(t, fromString, fromInt)
+ })
+
+ t.Run("deterministic - a uuid maps the same from both ingestion paths", func(t *testing.T) {
+ t.Parallel()
+
+ transformer := newTransformer(t, values, "deterministic")
+
+ // a snapshot delivers what pgx decoded, replication delivers the text
+ // wal2json emitted
+ fromSnapshot, err := transformer.Transform(context.Background(), NewValue(testUUID(0x01), "uuid", nil))
+ require.NoError(t, err)
+ fromReplication, err := transformer.Transform(context.Background(), NewValue("aabbccdd-eeff-0011-2233-445566778801", "uuid", nil))
+ require.NoError(t, err)
+
+ require.Equal(t, fromSnapshot, fromReplication)
+ })
+
+ t.Run("deterministic - a date maps the same from both ingestion paths", func(t *testing.T) {
+ t.Parallel()
+
+ transformer, err := newLookupChoiceTransformer(values, "deterministic", intTypes, true)
+ require.NoError(t, err)
+
+ fromSnapshot, err := transformer.Transform(context.Background(), NewValue(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC), "date", nil))
+ require.NoError(t, err)
+ fromReplication, err := transformer.Transform(context.Background(), NewValue("2024-01-01", "date", nil))
+ require.NoError(t, err)
+
+ require.Equal(t, fromSnapshot, fromReplication)
+ })
+
+ t.Run("a single value is always chosen", func(t *testing.T) {
+ t.Parallel()
+
+ transformer := newTransformer(t, []any{int64(7)}, "random")
+ got, err := transformer.Transform(context.Background(), NewValue("a", "int8", nil))
+ require.NoError(t, err)
+ require.Equal(t, int64(7), got)
+ })
+
+ t.Run("error - no values to choose from", func(t *testing.T) {
+ t.Parallel()
+
+ _, err := newLookupChoiceTransformer(nil, "random", intTypes, false)
+ require.ErrorIs(t, err, errLookupNoValues)
+ })
+}
+
+func TestLookupChoiceTransformer_interface(t *testing.T) {
+ t.Parallel()
+
+ transformer, err := newLookupChoiceTransformer([]any{int64(1)}, "random", []SupportedDataType{Integer64DataType}, false)
+ require.NoError(t, err)
+
+ require.Equal(t, LookupChoice, transformer.Type())
+ require.False(t, transformer.IsDynamic())
+ require.Equal(t, UniquenessLossy, transformer.Uniqueness())
+ require.NoError(t, transformer.Close())
+}
diff --git a/pkg/transformers/transformer.go b/pkg/transformers/transformer.go
index 727a99ddb..37104cc04 100644
--- a/pkg/transformers/transformer.go
+++ b/pkg/transformers/transformer.go
@@ -81,6 +81,7 @@ const (
Hstore TransformerType = "hstore"
PGAnonymizer TransformerType = "pg_anonymizer"
EncryptedAESSIV TransformerType = "encrypted_aes_siv"
+ LookupChoice TransformerType = "lookup_choice"
FPEFF1 TransformerType = "fpe_ff1"
)
diff --git a/pkg/wal/processor/transformer/wal_postgres_transformer_parser.go b/pkg/wal/processor/transformer/wal_postgres_transformer_parser.go
index 94cb33d1b..4f93d536d 100644
--- a/pkg/wal/processor/transformer/wal_postgres_transformer_parser.go
+++ b/pkg/wal/processor/transformer/wal_postgres_transformer_parser.go
@@ -137,8 +137,8 @@ func (v *PostgresTransformerParser) ParseAndValidate(ctx context.Context, rules
case "", "noop":
transformerMap.AddNoopTransformer(table.Schema, table.Table, colName)
continue
- case transformers.PGAnonymizer:
- // pg_anonymizer transformer requires a connection pool, set
+ case transformers.PGAnonymizer, transformers.LookupChoice:
+ // these transformers require a connection pool, set
// the source PG URL if not provided
if cfg.Parameters["postgres_url"] == nil {
cfg.Parameters["postgres_url"] = v.connURL
diff --git a/pkg/wal/processor/transformer/wal_postgres_transformer_parser_test.go b/pkg/wal/processor/transformer/wal_postgres_transformer_parser_test.go
index 74a534bb8..7f6fbb3b4 100644
--- a/pkg/wal/processor/transformer/wal_postgres_transformer_parser_test.go
+++ b/pkg/wal/processor/transformer/wal_postgres_transformer_parser_test.go
@@ -16,6 +16,7 @@ import (
pgmocks "github.com/xataio/pgstream/internal/postgres/mocks"
"github.com/xataio/pgstream/pkg/transformers"
"github.com/xataio/pgstream/pkg/transformers/builder"
+ transformermocks "github.com/xataio/pgstream/pkg/transformers/mocks"
)
func TestPostgresTransformerParser_ParseAndValidate(t *testing.T) {
@@ -762,3 +763,108 @@ func Test_validateNumericRange(t *testing.T) {
})
}
}
+
+// transformers that read from the database are given the source URL when the
+// rules do not name one; without it the documented configuration, which omits
+// postgres_url, fails to start
+func TestPostgresTransformerParser_connectionInjection(t *testing.T) {
+ t.Parallel()
+
+ const sourceURL = "postgres://user:pass@source:5432/db"
+
+ querier := func() *pgmocks.Querier {
+ return &pgmocks.Querier{
+ QueryFn: func(_ context.Context, _ uint, query string, _ ...any) (pglib.Rows, error) {
+ switch query {
+ case "SELECT * FROM \"public\".\"test\" LIMIT 0":
+ return &pgmocks.Rows{
+ FieldDescriptionsFn: func() []pgconn.FieldDescription {
+ return []pgconn.FieldDescription{{Name: "id", DataTypeOID: pgtype.Int8OID}}
+ },
+ CloseFn: func() {},
+ ErrFn: func() error { return nil },
+ }, nil
+ case uniqueIndexQuery:
+ return &pgmocks.Rows{
+ CloseFn: func() {},
+ NextFn: func(uint) bool { return false },
+ ErrFn: func() error { return nil },
+ }, nil
+ default:
+ return nil, fmt.Errorf("unexpected query: %s", query)
+ }
+ },
+ }
+ }
+
+ tests := []struct {
+ name string
+ transformer string
+ parameters map[string]any
+ wantPostgresURL any
+ }{
+ {
+ name: "lookup_choice without a url",
+ transformer: "lookup_choice",
+ parameters: map[string]any{"lookup_table": "public.countries", "lookup_column": "id"},
+ wantPostgresURL: sourceURL,
+ },
+ {
+ name: "lookup_choice with its own url",
+ transformer: "lookup_choice",
+ parameters: map[string]any{"lookup_table": "public.countries", "lookup_column": "id", "postgres_url": "postgres://elsewhere"},
+ wantPostgresURL: "postgres://elsewhere",
+ },
+ {
+ name: "pg_anonymizer without a url",
+ transformer: "pg_anonymizer",
+ parameters: map[string]any{"anon_function": "anon.fake_email()"},
+ wantPostgresURL: sourceURL,
+ },
+ {
+ name: "a transformer that needs no connection is left alone",
+ transformer: "string",
+ parameters: nil,
+ wantPostgresURL: nil,
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ t.Parallel()
+
+ var gotURL any
+ parser := PostgresTransformerParser{
+ conn: querier(),
+ connURL: sourceURL,
+ builder: &transformermocks.TransformerBuilder{
+ NewFn: func(cfg *transformers.Config) (transformers.Transformer, error) {
+ gotURL = cfg.Parameters["postgres_url"]
+ return &transformermocks.Transformer{
+ CompatibleTypesFn: func() []transformers.SupportedDataType {
+ return []transformers.SupportedDataType{transformers.AllDataTypes}
+ },
+ }, nil
+ },
+ },
+ pgtypeMap: pglib.NewMapper(querier()),
+ requiredTables: []string{"public.test"},
+ }
+
+ _, err := parser.ParseAndValidate(context.Background(), Rules{
+ ValidationMode: "relaxed",
+ Transformers: []TableRules{
+ {
+ Schema: "public",
+ Table: "test",
+ ColumnRules: map[string]TransformerRules{
+ "id": {Name: tc.transformer, Parameters: tc.parameters},
+ },
+ },
+ },
+ })
+ require.NoError(t, err)
+ require.Equal(t, tc.wantPostgresURL, gotURL)
+ })
+ }
+}
diff --git a/transformers-definition.json b/transformers-definition.json
index a557f81d9..622b88ab6 100644
--- a/transformers-definition.json
+++ b/transformers-definition.json
@@ -554,6 +554,65 @@
}
]
},
+ {
+ "name": "lookup_choice",
+ "supported_types": [
+ "string",
+ "citext",
+ "byte_array",
+ "boolean",
+ "integer16",
+ "integer32",
+ "integer64",
+ "float32",
+ "float64",
+ "uint8_array_of_16",
+ "date",
+ "datetime"
+ ],
+ "uniqueness": "lossy",
+ "parameters": [
+ {
+ "name": "lookup_table",
+ "supported_type": "string",
+ "default": null,
+ "dynamic": false,
+ "required": true
+ },
+ {
+ "name": "lookup_column",
+ "supported_type": "string",
+ "default": null,
+ "dynamic": false,
+ "required": true
+ },
+ {
+ "name": "postgres_url",
+ "supported_type": "string",
+ "default": null,
+ "dynamic": false,
+ "required": true
+ },
+ {
+ "name": "generator",
+ "supported_type": "string",
+ "default": "random",
+ "dynamic": false,
+ "required": false,
+ "values": [
+ "random",
+ "deterministic"
+ ]
+ },
+ {
+ "name": "ignore_values",
+ "supported_type": "array",
+ "default": null,
+ "dynamic": false,
+ "required": false
+ }
+ ]
+ },
{
"name": "masking",
"supported_types": [