Add lookup_choice transformer

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:

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.
 
+</details>
+
+ <details>
+  <summary>lookup_choice</summary>
+
+**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`                   |
+
 </details>
 
  <details>

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())
+}

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) {

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

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 {

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)
+}

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
+	}
+}