Declare unique constraints

CHANGELOG.md

@@ -10,6 +10,13 @@ what shipped rather than what was written down at the time.
 ## Unreleased
 
 ### Added
+- A `Table` can declare **unique constraints** — `unique=["email", ("tenant", "email")]`, any
+  number of them — which psql's `\d` prints as the `UNIQUE CONSTRAINT` lines of the same
+  `Indexes:` footer, named as Postgres names them (`users_email_key`). Unlike a primary
+  key they do not make their columns `not null`. Names longer than Postgres' 63-byte
+  identifier limit are truncated and, if that collides, given a counter, as Postgres
+  does. (#128)
+
 - A `Table` can declare a **primary key** — `Table("commits", {...}, primary_key="sha")`, or a
   tuple for a composite one — and psql's `\d` prints the footer it prints for a real
   table: `"commits_pkey" PRIMARY KEY, btree (sha)`. Its columns are reported `not null`

docs/session-api.md

@@ -68,8 +68,19 @@ Table("commit_files", {"sha": "text", "path": "text"}, primary_key=("sha", "path
 ```
 
 Its columns are then reported `not null`, as they are in Postgres, by both `attnotnull` and
-`information_schema.columns.is_nullable`. Nothing is enforced — pg_mimic stores no rows of its own,
-so a declared key is what the catalog reports rather than a uniqueness check anything applies.
+`information_schema.columns.is_nullable`.
+
+`unique` declares any number of further keys, spelled the same way:
+
+```python
+Table("users", {"id": "integer", "email": "text", "tenant": "text"}, primary_key="id", unique=["email", ("tenant", "email")])
+```
+
+Those appear in the same footer as `UNIQUE CONSTRAINT` lines, and — unlike a primary key — do
+**not** make their columns `not null`.
+
+Nothing is enforced either way: pg_mimic stores no rows of its own, so a declared key is what the
+catalog reports rather than a uniqueness check anything applies.
 
 The nested `{table: {column: type_name}}` dict this returned before is still accepted, so an
 existing session needs no change. A `Schema` is what pg_mimic normalises either into, and it is what

docs/whats-handled.md

@@ -59,9 +59,9 @@ before it ever reaches your `Session`, so real clients/ORMs/`psql` work without
   `column_default`, `udt_name` or `numeric_precision` answers about the columns that exist instead of
   returning nothing. Plus the slice of `pg_catalog` psql's `\dt`, `\d <table>`, `\di` and `\dn` read —
   all built from your `Session.schema()` declaration — see
-  [the session API](./session-api.md#declaring-a-schema). A declared primary key gets the
-  `Indexes:` footer `\d` prints for one, reports its columns as `not null`, and appears in
-  `pg_index`, `pg_constraint` and `\di`.
+  [the session API](./session-api.md#declaring-a-schema). A declared primary key or unique
+  constraint gets the `Indexes:` footer `\d` prints for one, and appears in `pg_index`,
+  `pg_constraint` and `\di`. A primary key additionally reports its columns as `not null`.
 - Multi-statement simple-query batches (`"BEGIN; INSERT ...; COMMIT;"` sent as one `'Q'` message, e.g. by
   `psql -f script.sql`) are split into individual statements (via sqlglot), each getting its own
   `RowDescription`/`DataRow*`/`CommandComplete` — not silently merged or truncated.

pg_mimic/catalog.py

@@ -20,6 +20,7 @@
 
 from __future__ import annotations
 
+import itertools
 import logging
 import re
 from typing import TYPE_CHECKING, NamedTuple
@@ -325,10 +326,50 @@ def _rendered_identifier(name: str) -> str:
     return name if _PLAIN_IDENTIFIER.match(name) else '"' + name.replace('"', '""') + '"'
 
 
-def _key_positions(table: Table) -> list[int]:
+class _Key(NamedTuple):
+    """A declared key, with the one bit that decides how it is named and rendered."""
+
+    table: str
+    columns: tuple[str, ...]
+    is_primary: bool
+
+
+def _declared_keys(table: Table) -> list[_Key]:
+    """Every key a table declares, primary first.
+
+    Order matters only for OID assignment being stable; psql sorts the footer itself
+    (`ORDER BY i.indisprimary DESC, c2.relname`).
+    """
+    keys = [_Key(table.name, table.primary_key, True)] if table.primary_key else []
+    return keys + [_Key(table.name, columns, False) for columns in table.unique]
+
+
+def _key_name(key: _Key, taken: set[str]) -> str:
+    """The name Postgres would give this key's index: `t_pkey`, or `t_a_b_key`.
+
+    The column names go in unquoted whatever they are -- Postgres builds `u1_oddName_key`
+    from `"oddName"`, then quotes the whole name when it prints it.
+
+    Truncated to 63 bytes, which is Postgres' NAMEDATALEN limit and reachable with an
+    ordinary long table name and two or three columns. A truncation that collides gets a
+    counter, as Postgres' own ChooseIndexName does -- two identical *keys* are refused at
+    declaration, so this only fires when two different keys truncate to one name.
+    """
+    stem = f"{key.table}_pkey" if key.is_primary else "_".join([key.table, *key.columns, "key"])
+    name = stem.encode()[:63].decode(errors="ignore")
+    if name not in taken:
+        return name
+    for counter in itertools.count(1):
+        suffixed = f"{name.encode()[: 63 - len(str(counter))].decode(errors='ignore')}{counter}"
+        if suffixed not in taken:
+            return suffixed
+    raise AssertionError("unreachable: itertools.count does not end")
+
+
+def _key_positions(table: Table, key: _Key) -> list[int]:
     """A key's columns as 1-based attnums, in key order -- Postgres' conkey and indkey."""
     positions = {column: position for position, column in enumerate(table.columns, start=1)}
-    return [positions[column] for column in table.primary_key]
+    return [positions[column] for column in key.columns]
 
 
 def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogInput:
@@ -337,6 +378,9 @@ def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogI
     attribute_rows = []
     index_rows: list[dict] = []
     constraint_rows: list[dict] = []
+    # Index names are relation names, so they share one namespace with each other and
+    # have to be unique across the whole schema, not just within a table.
+    key_names: set[str] = set()
     for index, (table_name, table) in enumerate(declared.tables.items()):
         table_oid = _FIRST_USER_OID + index
         class_rows.append(
@@ -349,9 +393,9 @@ def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogI
                 "relam": _HEAP_AM_OID,
                 "reltuples": -1,  # never analysed
                 "relpages": 0,
-                # psql gates the whole `Indexes:` footer on this, so a declared primary
-                # key that does not set it builds rows nothing ever asks for.
-                "relhasindex": bool(table.primary_key),
+                # psql gates the whole `Indexes:` footer on this, so a declared key that
+                # does not set it builds rows nothing ever asks for.
+                "relhasindex": bool(table.primary_key or table.unique),
                 "relpersistence": "p",
                 "reltablespace": 0,
                 "relispartition": False,
@@ -391,11 +435,13 @@ def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogI
                 }
             )
 
-        if table.primary_key:
-            index_oid = _FIRST_INDEX_OID + index
-            constraint_oid = _FIRST_CONSTRAINT_OID + index
-            key_name = f"{table_name}_pkey"
-            key_columns = ", ".join(_rendered_identifier(column) for column in table.primary_key)
+        for key in _declared_keys(table):
+            key_name = _key_name(key, taken=key_names)
+            key_names.add(key_name)
+            index_oid = _FIRST_INDEX_OID + len(index_rows)
+            constraint_oid = _FIRST_CONSTRAINT_OID + len(constraint_rows)
+            key_columns = ", ".join(_rendered_identifier(column) for column in key.columns)
+            positions = _key_positions(table, key)
             # An index is a relation, so it needs its own pg_class row: psql reads the
             # index's *name* as c2.relname off a second pg_class in the same query, and
             # without this row the join drops every index it was told about.
@@ -428,24 +474,30 @@ def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogI
                     # psql echoes everything after " USING " verbatim into the footer,
                     # which is where `btree (sha)` comes from -- so the access method and
                     # the column list have to be in this string, not derived from relam.
-                    "indexdef": f"CREATE UNIQUE INDEX {key_name} ON {_rendered_identifier(table_name)} USING btree ({key_columns})",
+                    "indexdef": (
+                        f"CREATE UNIQUE INDEX {_rendered_identifier(key_name)} "
+                        f"ON {_rendered_identifier(table_name)} USING btree ({key_columns})"
+                    ),
                     "indexrelid": index_oid,
                     "indrelid": table_oid,
-                    "indisprimary": True,
+                    # Both kinds are unique indexes; only one of them is *the* key, and
+                    # psql reads these two flags to choose between "PRIMARY KEY," and
+                    # "UNIQUE CONSTRAINT," in the footer.
+                    "indisprimary": key.is_primary,
                     "indisunique": True,
                     "indisclustered": False,
                     "indisvalid": True,
                     "indisreplident": False,
-                    "indkey": " ".join(str(position) for position in _key_positions(table)),
+                    "indkey": " ".join(str(position) for position in positions),
                 }
             )
             constraint_rows.append(
                 {
-                    "condef": f"PRIMARY KEY ({key_columns})",
+                    "condef": f"{'PRIMARY KEY' if key.is_primary else 'UNIQUE'} ({key_columns})",
                     "oid": constraint_oid,
                     "conname": key_name,
                     "conrelid": table_oid,
-                    "contype": "p",
+                    "contype": "p" if key.is_primary else "u",
                     "conparentid": 0,
                     "conindid": index_oid,
                     "confrelid": 0,
@@ -457,7 +509,7 @@ def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogI
                     "coninhcount": 0,
                     # Postgres' int2[] of attnums, as its array literal -- the column is
                     # declared TEXT here, and a Python list has no wire encoding.
-                    "conkey": "{" + ",".join(str(position) for position in _key_positions(table)) + "}",
+                    "conkey": "{" + ",".join(str(position) for position in positions) + "}",
                 }
             )
 

pg_mimic/declared.py

@@ -59,11 +59,21 @@ class Table:
     what `pg_constraint.conkey` records and what psql prints. pg_mimic enforces
     nothing: nothing here stores rows, so a declared key is what the catalog reports,
     not a uniqueness check anything applies.
+
+    `unique` is any number of further keys, each spelled the same way::
+
+        Table("commits", {...}, primary_key="sha", unique=["author_email", ("a", "b")])
+
+    A primary key and a unique constraint differ in three ways that matter here: there
+    is at most one primary key and any number of unique ones, only a primary key's
+    columns are implicitly NOT NULL, and they are named and rendered differently
+    (`t_pkey` / `PRIMARY KEY (...)` against `t_a_b_key` / `UNIQUE (...)`).
     """
 
     name: str
     columns: Mapping[str, str]
     primary_key: str | Sequence[str] = ()
+    unique: Sequence[str | Sequence[str]] = ()
 
     def __post_init__(self) -> None:
         if not isinstance(self.name, str) or not self.name:
@@ -86,6 +96,16 @@ def __post_init__(self) -> None:
         object.__setattr__(self, "columns", MappingProxyType(columns))
         object.__setattr__(self, "primary_key", _key_columns(self.name, "primary key", self.primary_key, columns))
 
+        declared_unique = [self.unique] if isinstance(self.unique, str) else list(self.unique)
+        keys = tuple(_key_columns(self.name, "unique constraint", key, columns) for key in declared_unique)
+        # Two identical unique keys would name two constraints the same thing, which
+        # `\\d` renders as a duplicated line. Postgres tolerates it and disambiguates with
+        # a counter; here it is almost certainly a copy-paste, so it is refused.
+        for key in keys:
+            if keys.count(key) > 1:
+                raise ValueError(f"{self.name!r} declares the same unique constraint twice: {key}")
+        object.__setattr__(self, "unique", keys)
+
 
 def _key_columns(table: str, kind: str, declared: str | Sequence[str], columns: Mapping[str, str]) -> tuple[str, ...]:
     """A key spec as a tuple of this table's own columns, in the order it was given.

tests/test_declared.py

@@ -164,3 +164,42 @@ def test_a_primary_key_must_name_the_tables_own_columns():
 def test_a_primary_key_cannot_name_a_column_twice():
     with pytest.raises(ValueError, match="twice in its primary key"):
         Table("t", {"a": "text"}, primary_key=("a", "a"))
+
+
+# --- unique constraints --------------------------------------------------------------
+
+
+def test_unique_keys_are_normalised_the_same_way_a_primary_key_is():
+    table = Table("t", {"a": "text", "b": "integer", "c": "text"}, unique=["a", ("b", "c")])
+    assert table.unique == (("a",), ("b", "c"))
+
+
+def test_a_single_unique_column_may_be_spelled_as_a_bare_string():
+    assert Table("t", {"a": "text"}, unique="a").unique == (("a",),)
+
+
+def test_a_table_declares_no_unique_constraints_by_default():
+    assert Table("t", {"a": "text"}).unique == ()
+
+
+def test_a_unique_constraint_must_name_the_tables_own_columns():
+    with pytest.raises(ValueError, match="not one of its columns"):
+        Table("t", {"a": "text"}, unique="b")
+
+
+def test_a_unique_constraint_cannot_name_a_column_twice():
+    with pytest.raises(ValueError, match="twice in its unique constraint"):
+        Table("t", {"a": "text"}, unique=[("a", "a")])
+
+
+def test_the_same_unique_constraint_cannot_be_declared_twice():
+    """Two identical keys would name two constraints the same thing, which `\\d` renders
+    as a duplicated line. Postgres tolerates it with a counter; here it is a copy-paste."""
+    with pytest.raises(ValueError, match="same unique constraint twice"):
+        Table("t", {"a": "text", "b": "text"}, unique=["a", "a"])
+
+
+def test_a_unique_constraint_may_repeat_the_primary_key():
+    """Redundant, and legal in Postgres too -- the names differ, so nothing collides."""
+    table = Table("t", {"a": "text"}, primary_key="a", unique="a")
+    assert (table.primary_key, table.unique) == (("a",), (("a",),))

tests/test_pg_catalog.py

@@ -15,7 +15,7 @@
 import pytest
 import sqlglot
 
-from pg_mimic.catalog import _build_pg_catalog
+from pg_mimic.catalog import _build_pg_catalog, _Key, _key_name
 from pg_mimic.catalog_rewrite import rewrite_for_executor
 from pg_mimic.declared import Schema, Table, resolve
 
@@ -236,3 +236,94 @@ def test_a_primary_keys_columns_are_not_null_and_the_others_are():
 def test_conkey_is_the_key_columns_attnums_in_key_order():
     rows = _run("SELECT conname, conkey FROM pg_catalog.pg_constraint ORDER BY conname", _PK_SCHEMA)
     assert dict(rows) == {"commit_files_pkey": "{1,2}", "commits_pkey": "{1}", "odd_pkey": "{1}"}
+
+
+# --- unique constraints --------------------------------------------------------------
+
+_UNIQUE_SCHEMA = Schema(
+    [
+        Table(
+            "u1",
+            {"id": "integer", "email": "text", "a": "text", "b": "integer", "oddName": "text"},
+            primary_key="id",
+            unique=["email", ("a", "b"), "oddName"],
+        ),
+        Table("only_unique", {"code": "text", "n": "integer"}, unique="code"),
+        Table("plain", {"x": "text"}),
+    ]
+)
+
+
+def _footer_rows(table_name, schema):
+    rows = _run(f"SELECT oid FROM pg_catalog.pg_class WHERE relname = '{table_name}' AND relkind = 'r'", schema)
+    return _run(_INDEX_FOOTER.format(oid=rows[0][0]), schema)
+
+
+def test_psqls_footer_lists_the_primary_key_first_then_unique_by_name():
+    """psql orders it `indisprimary DESC, c2.relname`, so this is the order the lines
+    come out in -- matched against a real PostgreSQL 18 declaring the same table."""
+    rows = _footer_rows("u1", _UNIQUE_SCHEMA)
+    assert [(row[0], row[1], row[7]) for row in rows] == [
+        ("u1_pkey", True, "p"),
+        ("u1_a_b_key", False, "u"),
+        ("u1_email_key", False, "u"),
+        ("u1_oddName_key", False, "u"),
+    ]
+
+
+def test_a_unique_constraint_renders_as_unique_not_primary_key():
+    definitions = {row[0]: (row[5], row[6]) for row in _footer_rows("u1", _UNIQUE_SCHEMA)}
+    assert definitions["u1_a_b_key"][0].endswith(" USING btree (a, b)")
+    assert definitions["u1_a_b_key"][1] == "UNIQUE (a, b)"
+    assert definitions["u1_oddName_key"][1] == 'UNIQUE ("oddName")'
+    assert definitions["u1_pkey"][1] == "PRIMARY KEY (id)"
+
+
+def test_a_unique_index_is_unique_but_not_primary():
+    rows = _footer_rows("only_unique", _UNIQUE_SCHEMA)
+    assert [(row[0], row[1], row[2]) for row in rows] == [("only_unique_code_key", False, True)]
+
+
+def test_a_unique_constraints_columns_are_still_nullable():
+    """The one substantive difference from a primary key: only a primary key's columns
+    are implicitly NOT NULL."""
+    oid = _run("SELECT oid FROM pg_catalog.pg_class WHERE relname = 'u1' AND relkind = 'r'", _UNIQUE_SCHEMA)[0][0]
+    rows = _run(
+        f"SELECT attname, attnotnull FROM pg_catalog.pg_attribute WHERE attrelid = '{oid}' ORDER BY attnum",
+        _UNIQUE_SCHEMA,
+    )
+    assert rows == [("id", True), ("email", False), ("a", False), ("b", False), ("oddName", False)]
+
+
+def test_a_table_with_only_unique_constraints_still_gets_the_footer():
+    """relhasindex gates the whole footer, and a table with no primary key has one."""
+    rows = _run("SELECT relname, relhasindex FROM pg_catalog.pg_class WHERE relkind = 'r'", _UNIQUE_SCHEMA)
+    assert dict(rows) == {"u1": True, "only_unique": True, "plain": False}
+
+
+def test_every_key_gets_its_own_index_and_constraint_oid():
+    indexes = _run("SELECT indexrelid, indrelid FROM pg_catalog.pg_index", _UNIQUE_SCHEMA)
+    constraints = _run("SELECT oid FROM pg_catalog.pg_constraint", _UNIQUE_SCHEMA)
+    tables = _run("SELECT oid FROM pg_catalog.pg_class WHERE relkind = 'r'", _UNIQUE_SCHEMA)
+    index_oids = {oid for oid, _ in indexes}
+    assert len(index_oids) == len(indexes) == 5
+    assert len({oid for (oid,) in constraints}) == 5
+    assert not index_oids & {oid for (oid,) in tables}
+
+
+def test_a_key_name_is_truncated_to_postgres_identifier_length():
+    """63 bytes is NAMEDATALEN - 1, and reachable with an ordinary long table name."""
+    long_table = "t" * 60
+    name = _key_name(_Key(long_table, ("column_one", "column_two"), False), taken=set())
+    assert len(name) == 63
+    assert name.startswith(long_table)
+
+
+def test_two_keys_that_truncate_to_one_name_are_disambiguated():
+    """As Postgres' own ChooseIndexName does. Identical *keys* are refused at
+    declaration, so this only fires for different keys with a shared prefix."""
+    long_table = "t" * 60
+    first = _key_name(_Key(long_table, ("column_one",), False), taken=set())
+    second = _key_name(_Key(long_table, ("column_two",), False), taken={first})
+    assert first != second
+    assert len(second) <= 63