From f0c89164809abf2087f9457f3c085e48e32553e7 Mon Sep 17 00:00:00 2001 From: Tim Band Date: Wed, 23 Sep 2026 14:43:07 +0100 Subject: [PATCH 1/3] Added CopyProposer --- datafaker/db_utils.py | 29 +-- datafaker/interactive/generators.py | 5 +- datafaker/proposers/__init__.py | 3 +- datafaker/proposers/base.py | 76 +++++- datafaker/proposers/intervals.py | 296 +++++++++++++++++----- datafaker/providers.py | 35 ++- tests/examples/date.sql | 14 +- tests/test_generators_dialect.py | 11 +- tests/test_interactive_copy_generators.py | 168 ++++++++++++ tests/test_interactive_date_generators.py | 19 +- 10 files changed, 529 insertions(+), 127 deletions(-) create mode 100644 tests/test_interactive_copy_generators.py diff --git a/datafaker/db_utils.py b/datafaker/db_utils.py index 7ed5f17c..e7e18a9b 100644 --- a/datafaker/db_utils.py +++ b/datafaker/db_utils.py @@ -12,15 +12,7 @@ # pylint: disable=no-name-in-module from psycopg2.errors import UndefinedObject # ty: ignore[unresolved-import] -from sqlalchemy import ( - Column, - Connection, - Engine, - ForeignKey, - create_engine, - event, - select, -) +from sqlalchemy import Connection, Engine, ForeignKey, create_engine, event, select from sqlalchemy.engine.interfaces import DBAPIConnection from sqlalchemy.exc import ( IntegrityError, @@ -280,25 +272,6 @@ def get_orm_metadata( return metadata -def get_fk_column_between( - from_: Table, to_: Table -) -> tuple[Column, ForeignKey] | tuple[None, None]: - """ - Get a foreign key between two tables. - - :param from_: The table on which the foreign key should be defined. - :param to_: The table to which the foreign key should refer. - :return: A pair (column on from_, foreign key to to_) if such a key exists, - or (None, None) if not. - """ - for c in from_.columns: - fk: ForeignKey - for fk in c.foreign_keys: - if fk.column.table == to_: - return (c, fk) - return (None, None) - - def fk_refers_to_ignored_table(fk: ForeignKey) -> bool: """ Test if this foreign key refers to an ignored table. diff --git a/datafaker/interactive/generators.py b/datafaker/interactive/generators.py index a2ca912f..3789afbb 100644 --- a/datafaker/interactive/generators.py +++ b/datafaker/interactive/generators.py @@ -201,6 +201,9 @@ class GeneratorCmd(DbCmd): PRIMARY_PRIVATE_TEXT = "Primary Private" SECONDARY_PRIVATE_TEXT = "Secondary Private on columns {0}" NOT_PRIVATE_TEXT = "Not private" + REQUIRES_NO_SOURCE_DATA_TEXT = ( + "{0}. {2}{1}{3} requires no data from the source database." + ) REQUIRES_SOURCE_DATA_TEXT = ( "{0}. {2}{1}{3} requires the following data from the source database:" ) @@ -910,7 +913,7 @@ def _print_values_queried(self, table_name: str, n: int, prop: Proposer) -> None theme = get_active_theme() if not prop.select_aggregate_clauses() and not prop.custom_queries(): self.print( - "{0}. {2}{1}{3} requires no data from the source database.", + self.REQUIRES_NO_SOURCE_DATA_TEXT, n, prop.name(), theme.function, diff --git a/datafaker/proposers/__init__.py b/datafaker/proposers/__init__.py index bee4c536..98c969a6 100644 --- a/datafaker/proposers/__init__.py +++ b/datafaker/proposers/__init__.py @@ -17,7 +17,7 @@ MultivariateNormalProposerFactory, ) from datafaker.proposers.extract import DateComponentExtractProposerFactory -from datafaker.proposers.intervals import DateAfterProposerFactory +from datafaker.proposers.intervals import CopyProposerFactory, DateAfterProposerFactory from datafaker.proposers.mimesis import ( MimesisDateProposerFactory, MimesisDateTimeProposerFactory, @@ -57,6 +57,7 @@ def everything_factory(config: Mapping, metadata: MetaData) -> ProposerFactory: NullPartitionedNormalProposerFactory(config, metadata), NullPartitionedLogNormalProposerFactory(config, metadata), DateAfterProposerFactory(config, metadata), + CopyProposerFactory(config, metadata), DateComponentExtractProposerFactory(config, metadata), IncrementProposerFactory(), ) diff --git a/datafaker/proposers/base.py b/datafaker/proposers/base.py index 058b7b18..4da4d090 100644 --- a/datafaker/proposers/base.py +++ b/datafaker/proposers/base.py @@ -2,7 +2,8 @@ import re from abc import ABC, abstractmethod -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass from typing import Any, Union import mimesis @@ -14,7 +15,7 @@ from datafaker.dialects import StdDev from datafaker.providers import DistributionProvider -from datafaker.utils import logger +from datafaker.utils import T, logger NumericType = Union[int, float] @@ -291,6 +292,57 @@ def fit_from_buckets(xs: Sequence[NumericType], ys: Sequence[NumericType]) -> fl return sum_diff_squared / (count * count) +@dataclass +class ForeignKeyRelationship: + """How a table is related.""" + + fk_column: Column # the "from" column + target_column: Column # the "to" column that fk_column matches + + def target_table(self) -> Table: + """Get the related table.""" + return self.target_column.table + + +def _get_one(seq: Iterable[T]) -> T | None: + """Get an arbitrary member of the sequence, or None if it is empty.""" + for elt in seq: + return elt + return None + + +def make_foreign_key_relationship(fk_column: Column) -> ForeignKeyRelationship | None: + """Make a ``ForeignKeyRelationship`` from a foreign key column.""" + if (elt := _get_one(fk_column.foreign_keys)) is None: + return None + return ForeignKeyRelationship(fk_column, elt.column) + + +@dataclass +class RelatedColumn: + """A column in a related table, and how it is related.""" + + target: Column + relationship: ForeignKeyRelationship | None + + def add_from_to_query(self, query: Any) -> Any: + """Take a query and add a ``select_from`` and maybe ``join`` as appropriate.""" + if self.relationship is None: + return query.select_from(self.target.table) + query = query.select_from(self.relationship.fk_column.table) + return self.add_join_to_query(query) + + def add_join_to_query(self, query: Any) -> Any: + """Take a query and maybe add ``join`` as appropriate.""" + rel = self.relationship + if rel is None: + return query + return query.join( + rel.target_column.table, + onclause=rel.fk_column == rel.target_column, + ) + + class Buckets: """ Measured buckets for a real distribution. @@ -308,12 +360,12 @@ def __init__( mean: float, stddev: float, count: int, - join_tables: list[Table] | None = None, + anchor_relation: RelatedColumn | None = None, ): """Initialise a Buckets object.""" self.mean = mean self.stddev = stddev - query = self._get_bucket_query(table, column, join_tables) + query = self._get_bucket_query(table, column, anchor_relation) with engine.connect() as connection: raw_buckets = connection.execute(query.group_by("b")) self.buckets: Sequence[int] = [0] * 10 @@ -335,7 +387,7 @@ def _get_bucket_query( self, table: Table | Join, column: Any, - join_tables: list[Table] | None, + anchor_relation: RelatedColumn | None = None, ) -> Any: """Make the query for half-sd buckets.""" bottom = self.mean - 2 * self.stddev @@ -344,8 +396,8 @@ def _get_bucket_query( func.count(column).label("f"), # pylint: disable=not-callable ((func.floor(column) - bottom) / width).label("b"), ).select_from(table) - for jt in join_tables or []: - query = query.join(jt) + if anchor_relation is not None: + query = anchor_relation.add_join_to_query(query) return query @classmethod @@ -354,7 +406,7 @@ def make_buckets( engine: Engine, table: Table | Join, column: Any, - join_tables: list[Table] | None = None, + anchor_relation: RelatedColumn | None = None, ) -> Self | None: """ Construct a Buckets object. @@ -368,14 +420,16 @@ def make_buckets( :param engine: SQLAlchemy engine. :param table: SQLAlchemy table (or joined tables) to pull data from. :param column: SQLAlchemy column or expression to measure. + :param anchor_relation: A relationship to another column whose table + needs to be joined. """ query = select( func.avg(column).label("mean"), StdDev(column).label("stddev"), func.count(column).label("count"), # pylint: disable=not-callable ).select_from(table) - for jt in join_tables or []: - query = query.join(jt) + if anchor_relation is not None: + query = anchor_relation.add_join_to_query(query) with engine.connect() as connection: result = connection.execute(query).first() if result is None or result.stddev is None or getattr(result, "count") < 2: @@ -388,7 +442,7 @@ def make_buckets( result.mean, result.stddev, getattr(result, "count"), - join_tables, + anchor_relation, ) except DatabaseError as exc: logger.debug("Failed to instantiate Buckets object: %s", exc) diff --git a/datafaker/proposers/intervals.py b/datafaker/proposers/intervals.py index e6bb11fc..689a0c0e 100644 --- a/datafaker/proposers/intervals.py +++ b/datafaker/proposers/intervals.py @@ -3,23 +3,28 @@ from collections.abc import Mapping, Sequence from typing import Any -from sqlalchemy import Column, Engine, ForeignKey, MetaData, func, select +from sqlalchemy import Column, Engine, MetaData, func, select from sqlalchemy.types import Date, DateTime -from datafaker.db_utils import get_fk_column_between from datafaker.dialects import Random, SecondsDifference, StdDev -from datafaker.proposers.base import Buckets, Proposer, ProposerFactory, get_column_type +from datafaker.proposers.base import ( + Buckets, + ForeignKeyRelationship, + Proposer, + ProposerFactory, + RelatedColumn, + get_column_type, + make_foreign_key_relationship, +) from datafaker.providers import AnchoredProvider from datafaker.utils import get_property, procrustes -RelatedColumn = tuple[ForeignKey | None, Column] - def _set_roles_for_column( out: dict[str, list[RelatedColumn]], - fk: ForeignKey | None, + fk: ForeignKeyRelationship | None, column: Column, - column_config: Mapping, + columns_config: Mapping, ) -> None: """ Set new entries in ``out`` based on the roles ``column`` has. @@ -29,19 +34,19 @@ def _set_roles_for_column( that role and ``fk`` is the foreign key in the table ``column`` appears in that points to the table ``fcol`` appears in, or None if ``column`` and ``fcol`` are in the same table. - :param fk: Foreign key to the table ``column`` appears in (appears in - new entries set in ``out``). + :param fk: Foreign key relationship to the table ``column`` appears in + (appears in new entries set in ``out``). :param column: The column to be checked for roles. - :param column_config: The ``tables: : columns:`` section + :param columns_config: The ``tables: : columns:`` section of the ``config.yaml`` file. """ - roles: list[Any] = get_property(column_config, [column.name, "roles"], []) + roles: list[Any] = get_property(columns_config, [column.name, "roles"], []) for role in roles: - pair: RelatedColumn = (fk, column) + rc = RelatedColumn(column, fk) if role not in out: - out[role] = [pair] + out[role] = [rc] else: - out[role].append(pair) + out[role].append(rc) def _get_roles( @@ -60,21 +65,22 @@ def _get_roles( """ table = column.table tables_config: dict[str, Any] = get_property(config, "tables", {}) - table_conf: dict[str, Any] = get_property( + columns_config: dict[str, Any] = get_property( tables_config, [str(table.name), "columns"], {} ) role_to_fk_columns: dict[str, list[RelatedColumn]] = {} for col in table.columns: - _set_roles_for_column(role_to_fk_columns, None, col, table_conf) + _set_roles_for_column(role_to_fk_columns, None, col, columns_config) # look for roles in related tables - if col.foreign_keys: - fk = list(col.foreign_keys)[0] - target_table = fk.column.table + if (fk_relationship := make_foreign_key_relationship(col)) is not None: + target_table = fk_relationship.target_table() ft_conf: dict[str, Any] = get_property( tables_config, [str(target_table.name), "columns"], {} ) for fcol in target_table.columns: - _set_roles_for_column(role_to_fk_columns, fk, fcol, ft_conf) + _set_roles_for_column( + role_to_fk_columns, fk_relationship, fcol, ft_conf + ) return role_to_fk_columns @@ -103,7 +109,7 @@ def __init__( sd: float, mean: float, column: Column, - anchor: Column, + anchor_related: RelatedColumn, engine: Engine, buckets: Buckets | None = None, ): @@ -116,13 +122,13 @@ def __init__( :param column: The column to generate. :param anchor: The anchor column, which must be in the same table as column or a related table. - :param dialect: The dialect of the source database. + :param engine: The source database. :param buckets: The buckets for the intervals in the source database. """ super().__init__() self._sd = sd self._mean = mean - self._anchor = anchor + self._anchor_related = anchor_related self._column = column self._engine = engine self._provider = AnchoredProvider(metadata=metadata) @@ -134,27 +140,44 @@ def __init__( [(b - coerce_to_datetime(a, b)).total_seconds() for (a, b) in intervals] ) + @classmethod + def get_query( + cls, + column: Column, + anchor_related: RelatedColumn, + ) -> Any: + """Create the query for the summary of interval lengths.""" + anchor = anchor_related.target + return anchor_related.add_from_to_query( + select( + func.avg(SecondsDifference(column, anchor)).label("mean"), + StdDev(SecondsDifference(column, anchor)).label("sd"), + ) + ) + def function_name(self) -> str: """Get the name of the generator function to call.""" - if self._column.table == self._anchor.table: + if self._anchor_related.relationship is None: return "generic.anchored_provider.normal_date" return "generic.anchored_provider.normal_date_fk" def name(self) -> str: """Get the name of the generator.""" fname = self.function_name() - aname = self._anchor.name - atable = self._anchor.table - if atable == self._column.table: + anchor = self._anchor_related.target + aname = anchor.name + atable = anchor.table + if self._anchor_related.relationship is None: return f"{fname} [anchored to {aname}]" - return f"{fname} [anchored to {aname} of table {atable.name}]" + atable = anchor.table.name + rel = self._anchor_related.relationship.fk_column.name + return f"{fname} [anchored to {aname} of table {atable} (via {rel})]" def nominal_kwargs(self) -> dict[str, Any]: """Get the arguments to be entered into ``config.yaml``.""" column = self._column - anchor = self._anchor - (fk_col, fk) = get_fk_column_between(column.table, anchor.table) - if fk_col is None or fk is None: + anchor = self._anchor_related.target + if self._anchor_related.relationship is None: return { "mean_seconds": ( f'SRC_STATS["auto__{column.table.name}"]' @@ -167,33 +190,34 @@ def nominal_kwargs(self) -> dict[str, Any]: "anchor": f'GENERATED_ROW["{anchor.name}"]', } key = f"auto__interval__{column.table.name}__{column.name}" + fk_col = self._anchor_related.relationship.fk_column.name + on_col = self._anchor_related.relationship.target_column.name return { "dst_db_conn": "dst_db_conn", "anchor_column": f'"{anchor.name}"', "table": f'"{anchor.table.name}"', "mean_seconds": (f'SRC_STATS["{key}"]["results"][0]["mean"]'), "sd_seconds": (f'SRC_STATS["{key}"]["results"][0]["sd"]'), - "anchor_row": f'GENERATED_ROW["{fk_col.name}"]', - "on_column": f'"{fk.column.name}"', + "anchor_row": f'GENERATED_ROW["{fk_col}"]', + "on_column": f'"{on_col}"', } def actual_kwargs(self) -> dict[str, Any]: """Get the kwargs (summary statistics) this generator was instantiated with.""" - column = self._column - anchor = self._anchor - (_, fk) = get_fk_column_between(column.table, anchor.table) - if fk is None: + if self._anchor_related.relationship is None: return { "mean_seconds": self._sd, "sd_seconds": self._mean, "anchor": "1970-01-01", } + anchor = self._anchor_related.target + on_col = self._anchor_related.relationship.target_column.name return { "anchor_column": anchor.name, "table": anchor.table.name, "mean_seconds": self._sd, "sd_seconds": self._mean, - "on_column": fk.column.name, + "on_column": on_col, } def select_aggregate_clauses(self) -> dict[str, dict[str, str]]: @@ -202,10 +226,10 @@ def select_aggregate_clauses(self) -> dict[str, dict[str, str]]: This is for anchors in the same table. """ - column = self._column - anchor = self._anchor - if column.table != anchor.table: + if self._anchor_related.relationship is not None: return {} + column = self._column + anchor = self._anchor_related.target mean_q = func.avg(SecondsDifference(column, anchor)) sd_q = StdDev(SecondsDifference(column, anchor)) @@ -243,23 +267,17 @@ def custom_queries(self) -> dict[str, dict[str, Any]]: This is for anchors in a related table. """ column = self._column - anchor = self._anchor - if column.table == anchor.table: + rel = self._anchor_related + if rel.relationship is None: return {} - query = ( - select( - func.avg(SecondsDifference(column, anchor)).label("mean"), - StdDev(SecondsDifference(column, anchor)).label("sd"), - ) - .select_from(column.table) - .join(anchor.table) - ) + query = DateAfterProposer.get_query(column, rel) return { f"auto__interval__{column.table.name}__{column.name}": { "comments": [ "Mean and standard deviation of the length of time between" - f" column {anchor.name} of table {anchor.table.name}" - f" and column {column.name} of table {column.table.name}." + f" column {rel.target.name} of table {rel.target.table.name}" + f" and column {column.name} of table {column.table.name}" + f" (joined through {rel.relationship.fk_column.name})" ], "query": str( query.compile( @@ -288,7 +306,8 @@ def generate_intervals( :param count: The number of intervals to generate. :return: Each pair consists of an anchor and the generated datetime. """ - query = select(self._anchor).select_from(self._anchor.table) + anchor = self._anchor_related.target + query = select(anchor).select_from(anchor.table) query = query.order_by(Random()).limit(count) with self._engine.connect() as conn: anchors = conn.execute(query).scalars().all() @@ -312,18 +331,14 @@ def __init__(self, config: Mapping, metadata: MetaData): self._config = config self._metadata = metadata - def make_date_after_proposers( - self, engine: Engine, column: Column, anchor: Column + def _make_date_after_proposers( + self, + engine: Engine, + column: Column, + anchor_related: RelatedColumn, ) -> list[DateAfterProposer]: """Create a ``DateAfterProposer`` object.""" - query = select( - func.avg(SecondsDifference(column, anchor)).label("mean"), - StdDev(SecondsDifference(column, anchor)).label("sd"), - ).select_from(column.table) - join_tables = [] - if anchor.table != column.table: - join_tables = [anchor.table] - query = query.join(anchor.table) + query = DateAfterProposer.get_query(column, anchor_related) with engine.connect() as connection: result = connection.execute(query).first() if result is None or result.sd is None: @@ -331,8 +346,8 @@ def make_date_after_proposers( buckets = Buckets.make_buckets( engine, column.table, - SecondsDifference(column, anchor), - join_tables, + SecondsDifference(column, anchor_related.target), + anchor_related, ) return [ DateAfterProposer( @@ -340,7 +355,7 @@ def make_date_after_proposers( result.sd, result.mean, column, - anchor, + anchor_related, engine=engine, buckets=buckets, ) @@ -362,10 +377,153 @@ def get_proposers( if "start" not in roles: return [] other_start_columns = [ - fk_col for fk_col in roles["start"] if fk_col[1] != column + fk_col for fk_col in roles["start"] if fk_col.target != column ] return [ prop for anchor in other_start_columns - for prop in self.make_date_after_proposers(engine, column, anchor[1]) + for prop in self._make_date_after_proposers(engine, column, anchor) + ] + + +class CopyProposer(Proposer): + """Proposer that proposes copying another column.""" + + # pylint: disable=too-many-arguments too-many-positional-arguments + def __init__( + self, + metadata: MetaData, + column: Column, + anchor_related: RelatedColumn, + engine: Engine, + ): + """ + Initialise a date after proposer. + + :param metadata: The metadata of the source database. + :param column: The column to generate. + :param anchor: The anchor column, which must be in the same table as column + or a related table. + :param engine: The source database. + :param buckets: The buckets for the intervals in the source database. + """ + super().__init__() + self._anchor_related = anchor_related + self._column = column + self._engine = engine + self._provider = AnchoredProvider(metadata=metadata) + + def function_name(self) -> str: + """Get the name of the generator function to call.""" + if self._anchor_related.relationship is None: + return "generic.anchored_provider.copy" + return "generic.anchored_provider.copy_fk" + + def name(self) -> str: + """Get the name of the generator.""" + fname = self.function_name() + aname = self._anchor_related.target.name + if self._anchor_related.relationship is None: + return f"{fname} [from {aname}]" + atable = self._anchor_related.target.table + on_col = self._anchor_related.relationship.fk_column.name + return f"{fname} [from {aname} of table {atable.name} (via {on_col})]" + + def nominal_kwargs(self) -> dict[str, Any]: + """Get the arguments to be entered into ``config.yaml``.""" + anchor = self._anchor_related.target + if self._anchor_related.relationship is None: + return { + "anchor": f'GENERATED_ROW["{anchor.name}"]', + } + fk_col = self._anchor_related.relationship.fk_column + on_col = self._anchor_related.relationship.target_column + return { + "dst_db_conn": "dst_db_conn", + "anchor_column": f'"{anchor.name}"', + "table": f'"{anchor.table.name}"', + "anchor_row": f'GENERATED_ROW["{fk_col.name}"]', + "on_column": f'"{on_col.name}"', + } + + def actual_kwargs(self) -> dict[str, Any]: + """Get the kwargs (summary statistics) this generator was instantiated with.""" + anchor = self._anchor_related.target + if self._anchor_related.relationship is None: + return { + "anchor": "1970-01-01", + } + on_col = self._anchor_related.relationship.target_column + return { + "anchor_column": anchor.name, + "table": anchor.table.name, + "on_column": on_col.name, + } + + def select_aggregate_clauses(self) -> dict[str, dict[str, str]]: + """No source data required.""" + return {} + + def custom_queries(self) -> dict[str, dict[str, Any]]: + """No source data required.""" + return {} + + def fit(self, default: float = -1) -> float: + """Get this generator's fit against the real data.""" + return default + + def generate_data(self, count: int) -> list[Any]: + """Generate ``count`` random data points for this column.""" + anchor = self._anchor_related.target + query = select(anchor).select_from(anchor.table) + query = query.order_by(Random()).limit(count) + with self._engine.connect() as conn: + return list(conn.execute(query).scalars().all()) + + +class CopyProposerFactory(ProposerFactory): + """Proposer that proposes copying another column.""" + + def __init__(self, config: Mapping, metadata: MetaData): + """Initialize ``CopyProposerFactory``.""" + super().__init__() + self._config = config + self._metadata = metadata + + def make_copy_fk_proposers( + self, engine: Engine, column: Column, anchor: RelatedColumn + ) -> list[CopyProposer]: + """Create a ``CopyProposer`` object.""" + return [ + CopyProposer( + self._metadata, + column, + anchor, + engine=engine, + ) + ] + + def get_proposers( + self, + columns: list[Column], + engine: Engine, + ) -> Sequence[Proposer]: + """Get all proposers of dates that might be anchored to another column.""" + if len(columns) != 1: + return [] + column = columns[0] + roles = _get_roles(self._config, column) + if "source" not in roles: + return [] + ct = get_column_type(column) + other_source_columns = [ + fk_col + for fk_col in roles["source"] + if fk_col.target != column + and type(get_column_type(fk_col.target)) is type(ct) + ] + return [ + prop + for anchor in other_source_columns + for prop in self.make_copy_fk_proposers(engine, column, anchor) ] diff --git a/datafaker/providers.py b/datafaker/providers.py index 0263ef56..55faabf6 100644 --- a/datafaker/providers.py +++ b/datafaker/providers.py @@ -693,8 +693,8 @@ def normal_date_fk( :param sd_seconds: Standard deviation of the intervals' lengths in seconds. :param table: The name of the table for the anchor column. :param on_column: The name of the column in the foreign table - that ``row`` must match. - :param row: The value in the ``on_column`` column in table ``table`` + that ``anchor_row`` must match. + :param anchor_row: The value in the ``on_column`` column in table ``table`` for the row providing the start of the interval. :param anchor_column: The name of the column in ``table`` providing the start of the interval. @@ -711,3 +711,34 @@ def normal_date_fk( if not isinstance(out, (str, dt.date, dt.datetime)): return None return self.normal_date(mean_seconds, sd_seconds, out) + + def copy(self, anchor: Any) -> Any: + """Copy another column from the same table.""" + return anchor + + def copy_fk( + self, + dst_db_conn: Connection, + table: str, + on_column: str, + anchor_row: Any, + anchor_column: str, + ) -> Any: + """ + Copy a column from a different table. + + :param dst_db_conn: Connection to the destination database. + :param table: The name of the table for the anchor column. + :param on_column: The name of the column in the foreign table + that ``anchor_row`` must match. + :param anchor_row: The value in the ``on_column`` column in table ``table`` + for the row providing the source to copy. + :param anchor_column: The name of the column in ``table`` to copy. + :return: The value to be copied. + """ + mt = self._metadata.tables[table] + query = select(mt.c[anchor_column].label("out")).where( + mt.c[on_column] == anchor_row, + ) + anchor = dst_db_conn.execute(query).first() + return getattr(anchor, "out", None) diff --git a/tests/examples/date.sql b/tests/examples/date.sql index a0155a9d..f24b87b0 100644 --- a/tests/examples/date.sql +++ b/tests/examples/date.sql @@ -23,16 +23,20 @@ CREATE TABLE public.happening ( name TEXT NOT NULL, person_id INTEGER NOT NULL, at_time TIMESTAMP WITH TIME ZONE NOT NULL, - at_date DATE NOT NULL + at_date DATE NOT NULL, + other_person_id INTEGER, + comment TEXT ); ALTER TABLE ONLY public.happening ADD CONSTRAINT happening_pkey PRIMARY KEY (id); ALTER TABLE ONLY public.happening ADD CONSTRAINT person_id_fkey FOREIGN KEY (person_id) REFERENCES public.person(id); +ALTER TABLE ONLY public.happening + ADD CONSTRAINT other_person_id_fkey FOREIGN KEY (other_person_id) REFERENCES public.person(id); ALTER TABLE public.happening OWNER TO postgres; -INSERT INTO public.happening VALUES (1, 'stepped on a tack', 1, '1997-04-20 04:05:06+00:00', '1997-04-20'); -INSERT INTO public.happening VALUES (2, 'had a dream', 2, '1997-04-20 04:04:16+00:00', '1997-04-20'); -INSERT INTO public.happening VALUES (3, 'kicked a can', 1, '2001-12-23 07:05:06+00:00', '2001-12-23'); -INSERT INTO public.happening VALUES (4, 'steppen in gum', 2, '2003-11-03 04:15:26+00:00', '2003-11-03'); +INSERT INTO public.happening VALUES (1, 'stepped on a tack', 1, '1997-04-20 04:05:06+00:00', '1997-04-20', NULL, NULL); +INSERT INTO public.happening VALUES (2, 'had a dream', 2, '1997-04-20 04:04:16+00:00', '1997-04-20', NULL, NULL); +INSERT INTO public.happening VALUES (3, 'kicked a can', 1, '2001-12-23 07:05:06+00:00', '2001-12-23', NULL, NULL); +INSERT INTO public.happening VALUES (4, 'shared a cake', 2, '2003-11-03 04:15:26+00:00', '2003-11-03', 1, NULL); diff --git a/tests/test_generators_dialect.py b/tests/test_generators_dialect.py index ec5e1744..4ddc21bb 100644 --- a/tests/test_generators_dialect.py +++ b/tests/test_generators_dialect.py @@ -11,7 +11,12 @@ from datafaker.dialects import SecondsDifference from datafaker.interactive.generators import get_aggregate_query from datafaker.interactive.missingness import MissingnessType -from datafaker.proposers.base import Buckets, PredefinedProposer, Proposer +from datafaker.proposers.base import ( + Buckets, + PredefinedProposer, + Proposer, + RelatedColumn, +) from datafaker.proposers.choice import ChoiceProposerFactory, ZipfChoiceProposer from datafaker.proposers.continuous import ( ContinuousLogDistributionProposerFactory, @@ -592,7 +597,7 @@ def test_date_after_proposer_stddev_clause_mssql(self) -> None: sd=1.0, mean=86400.0, column=tbl.c.end_date, - anchor=tbl.c.start_date, + anchor_related=RelatedColumn(tbl.c.start_date, None), engine=engine, ) clauses = proposer.select_aggregate_clauses() @@ -620,7 +625,7 @@ def test_date_after_proposer_stddev_clause_postgresql(self) -> None: sd=1.0, mean=86400.0, column=tbl.c.end_date, - anchor=tbl.c.start_date, + anchor_related=RelatedColumn(tbl.c.start_date, None), engine=engine, ) clauses = proposer.select_aggregate_clauses() diff --git a/tests/test_interactive_copy_generators.py b/tests/test_interactive_copy_generators.py new file mode 100644 index 00000000..d6ea5b28 --- /dev/null +++ b/tests/test_interactive_copy_generators.py @@ -0,0 +1,168 @@ +""" Tests for the configure-generators command. """ +from collections.abc import MutableMapping +from typing import Any + +from sqlalchemy import select +from sqlalchemy.orm import aliased + +from datafaker.interactive.base import DbCmd +from tests.test_interactive_generators import MockGeneratorCmd +from tests.utils import GeneratesDBTestCase, MsSqlTestDb + + +class ConfigureCopyGeneratorsWithDateTests(GeneratesDBTestCase): + """Test `configure-generators` with the `instrument.sql` database.""" + + dump_file_path = "date.sql" + database_name = "date_tables" + schema_name = "public" + use_temporary_cwd = True + + def _get_cmd(self, config: MutableMapping[str, Any]) -> MockGeneratorCmd: + """Get the command we are using for this test case.""" + return MockGeneratorCmd( + DbCmd.Settings(self.dsn, self.schema_name, config, self.metadata, None) + ) + + def test_same_table_copy(self) -> None: + """Test that the cross-table copy is proposed and compared.""" + config = { + "tables": { + "happening": { + "columns": { + "name": { + "roles": ["source"], + } + } + } + } + } + with self._get_cmd(config) as gc: + # set up our copy proposer + gc.do_next("happening.comment") + gc.reset() + gc.do_propose("") + proposals = gc.get_proposals() + provider_name = "generic.anchored_provider.copy [from name]" + self.assertIn(provider_name, proposals.keys()) + prop = proposals[provider_name] + gc.reset() + gc.do_compare(str(prop[0])) + self.assertEqual(gc.messages[0][0], gc.NOT_PRIVATE_TEXT) + self.assertEqual(gc.messages[1][0], gc.REQUIRES_NO_SOURCE_DATA_TEXT) + + def test_cross_table_copy(self) -> None: + """Test that the cross-table copy is proposed and compared.""" + config = { + "tables": { + "person": { + "columns": { + "name": { + "roles": ["source"], + } + } + } + } + } + with self._get_cmd(config) as gc: + # set up our copy proposer + gc.do_next("happening.comment") + gc.reset() + gc.do_propose("") + proposals = gc.get_proposals() + provider_name = ( + "generic.anchored_provider.copy_fk" + " [from name of table person (via person_id)]" + ) + self.assertIn(provider_name, proposals.keys()) + prop = proposals[provider_name] + provider_name2 = ( + "generic.anchored_provider.copy_fk" + " [from name of table person (via other_person_id)]" + ) + self.assertIn(provider_name2, proposals.keys()) + prop2 = proposals[provider_name2] + gc.reset() + gc.do_compare(f"{prop[0]} {prop2[0]}") + self.assertEqual(gc.messages[0][0], gc.NOT_PRIVATE_TEXT) + self.assertEqual(gc.messages[1][0], gc.REQUIRES_NO_SOURCE_DATA_TEXT) + + def test_cross_table_copy_actually_copies(self) -> None: + """Test that the cross-table copy is copies.""" + config = { + "tables": { + "person": { + "columns": { + "name": { + "roles": ["source"], + } + } + } + } + } + with self._get_cmd(config) as gc: + # set up our copy proposer + gc.do_next("happening.comment") + gc.reset() + gc.do_propose("") + proposals = gc.get_proposals() + provider_name = ( + "generic.anchored_provider.copy_fk" + " [from name of table person (via person_id)]" + ) + self.assertIn(provider_name, proposals.keys()) + prop = proposals[provider_name] + gc.do_set(str(prop[0])) + gc.do_next("happening.name") + gc.reset() + gc.do_propose("") + proposals = gc.get_proposals() + provider_name = ( + "generic.anchored_provider.copy_fk" + " [from name of table person (via other_person_id)]" + ) + self.assertIn(provider_name, proposals.keys()) + prop = proposals[provider_name] + gc.do_set(str(prop[0])) + gc.do_quit("") + self.generate_data(gc.config, num_passes=50) + assert self.dst_engine is not None + with self.dst_engine.connect() as conn: + pert = self.dst_metadata.tables["person"] + other_pert = aliased(pert) + hapt = self.dst_metadata.tables["happening"] + query = ( + select( + hapt.columns["name"].label("hn"), + hapt.columns["comment"].label("hc"), + pert.columns["name"].label("pn"), + other_pert.columns["name"].label("opn"), + ) + .select_from(hapt) + .join( + pert, + onclause=pert.columns["id"] == hapt.columns["person_id"], + ) + .join( + other_pert, + onclause=other_pert.columns["id"] + == hapt.columns["other_person_id"], + isouter=True, + ) + ) + dst_result = conn.execute(query).fetchall() + saw_other = False + for row in dst_result: + if row.opn: + saw_other = True + self.assertEqual(row.hn, row.opn) + self.assertEqual(row.hc, row.pn) + self.assertTrue(saw_other) + + +# Note that this test won't work with DuckDB because it needs foreign keys to work +class ConfigureCopyGeneratorsWithDateMsSqlTests(ConfigureCopyGeneratorsWithDateTests): + """Test `configure-generators` with `instrument.sql` with MS SQL.""" + + database_type = MsSqlTestDb + schema_name = None diff --git a/tests/test_interactive_date_generators.py b/tests/test_interactive_date_generators.py index 978347e9..3addac8d 100644 --- a/tests/test_interactive_date_generators.py +++ b/tests/test_interactive_date_generators.py @@ -166,7 +166,8 @@ def test_cross_table_interval_end_to_end(self) -> None: proposals = gc.get_proposals() provider_name = ( "generic.anchored_provider.normal_date_fk" - f" [anchored to {anchor} of table {atable}]" + f" [anchored to {anchor} of table {atable}" + " (via manufacturer_id)]" ) self.assertIn(provider_name, proposals.keys()) prop = proposals[provider_name] @@ -216,7 +217,7 @@ class ConfigureGeneratorsWithInstrumentsMsSqlTests( class ConfigureGeneratorsWithDateTests(GeneratesDBTestCase): - """Test `configure-generators` with the `instrument.sql` database.""" + """Test `configure-generators` with the `date.sql` database.""" dump_file_path = "date.sql" database_name = "date_tables" @@ -274,13 +275,15 @@ def test_cross_table_interval_timestamp_vs_date(self) -> None: proposals = gc.get_proposals() provider_name = ( "generic.anchored_provider.normal_date_fk" - f" [anchored to {anchor} of table {atable}]" + f" [anchored to {anchor} of table {atable}" + " (via person_id)]" ) self.assertIn(provider_name, proposals.keys()) prop = proposals[provider_name] provider_name2 = ( "generic.anchored_provider.normal_date_fk" - f" [anchored to {anchor2} of table {atable}]" + f" [anchored to {anchor2} of table {atable}" + " (via person_id)]" ) self.assertIn(provider_name2, proposals.keys()) prop2 = proposals[provider_name2] @@ -295,13 +298,15 @@ def test_cross_table_interval_timestamp_vs_date(self) -> None: proposals = gc.get_proposals() provider_name = ( "generic.anchored_provider.normal_date_fk" - f" [anchored to {anchor} of table {atable}]" + f" [anchored to {anchor} of table {atable}" + " (via person_id)]" ) self.assertIn(provider_name, proposals.keys()) prop = proposals[provider_name] provider_name2 = ( "generic.anchored_provider.normal_date_fk" - f" [anchored to {anchor2} of table {atable}]" + f" [anchored to {anchor2} of table {atable}" + " (via person_id)]" ) self.assertIn(provider_name2, proposals.keys()) prop2 = proposals[provider_name2] @@ -314,7 +319,7 @@ def test_cross_table_interval_timestamp_vs_date(self) -> None: # Note that this test won't work with DuckDB because it needs foreign keys to work class ConfigureGeneratorsWithDateMsSqlTests(ConfigureGeneratorsWithDateTests): - """Test `configure-generators` with `instrument.sql` with MS SQL.""" + """Test `configure-generators` with `date.sql` with MS SQL.""" database_type = MsSqlTestDb schema_name = None From 9d60a6c22732abf0c83ce6c9ba9982439fd402bd Mon Sep 17 00:00:00 2001 From: Tim Band Date: Fri, 25 Sep 2026 13:15:22 +0100 Subject: [PATCH 2/3] Fixed crash in DateAfterProposerFactory Also: - fixed role change text - fixed remove all roles - allow intervals and copying from same column in aliased table --- datafaker/interactive/generators.py | 10 ++-- datafaker/proposers/base.py | 36 ++++++++++++-- datafaker/proposers/intervals.py | 59 +++++++++++++---------- datafaker/proposers/partitioned.py | 19 +------- datafaker/utils.py | 17 +++++++ docs/source/docker.rst | 2 +- pyproject.toml | 2 +- tests/examples/date.sql | 13 +++-- tests/test_interactive_date_generators.py | 25 ++++++++++ 9 files changed, 126 insertions(+), 57 deletions(-) diff --git a/datafaker/interactive/generators.py b/datafaker/interactive/generators.py index 3789afbb..1c3f85b8 100644 --- a/datafaker/interactive/generators.py +++ b/datafaker/interactive/generators.py @@ -37,6 +37,7 @@ schema_qualified_name, set_property, split_column_full_name, + text_list, ) @@ -160,7 +161,8 @@ def do(self, column_roles: RoleEntry | None) -> RoleEntry | None: """Perform the command.""" if column_roles is None: return None - column_roles.new.remove(self.role) + if self.role in column_roles.new: + column_roles.new.remove(self.role) return column_roles @@ -481,7 +483,7 @@ def _copy_roles(self) -> None: """Set role information in the configuration.""" for table_name, table_roles in self.roles.items(): for column_name, entry in table_roles.items(): - if entry.new: + if entry.new or entry.old: set_property( self.config, ["tables", table_name, "columns", column_name, "roles"], @@ -543,8 +545,8 @@ def _print_role_changes(self) -> int: "Changing role set of column {0} of table {1} from {2} to {3}", column_name, table_name, - ", ".join(str(oe) for oe in entry.old), - ", ".join(str(ne) for ne in entry.new), + text_list((e.value for e in entry.old), "empty"), + text_list((e.value for e in entry.new), "empty"), ) return count diff --git a/datafaker/proposers/base.py b/datafaker/proposers/base.py index 4da4d090..2c152b40 100644 --- a/datafaker/proposers/base.py +++ b/datafaker/proposers/base.py @@ -10,6 +10,7 @@ import mimesis.locales from sqlalchemy import Column, Engine, Join, Table, func, select from sqlalchemy.exc import DatabaseError +from sqlalchemy.orm import aliased from sqlalchemy.types import Integer, Numeric, String, TypeEngine from typing_extensions import Self @@ -292,17 +293,34 @@ def fit_from_buckets(xs: Sequence[NumericType], ys: Sequence[NumericType]) -> fl return sum_diff_squared / (count * count) -@dataclass class ForeignKeyRelationship: """How a table is related.""" - fk_column: Column # the "from" column - target_column: Column # the "to" column that fk_column matches + def __init__(self, fk_column: Column, target_column: Column) -> None: + """ + Initialise a ``ForeignKeyRelationship``. + + :param fk_column: The column containing the foreign key (the source). + :param target_column: The column that ``fk_column`` matches in + the target table. + """ + self.fk_column = fk_column + table = aliased(target_column.table) + self.target_column = table.columns[target_column.name] def target_table(self) -> Table: """Get the related table.""" return self.target_column.table + def target_table_name(self) -> str: + """ + Get the name of the target table. + + Don't use ``target_table().name`` because this will give the + name of the aliased table. + """ + return str(self.target_column.table.original.name) + def _get_one(seq: Iterable[T]) -> T | None: """Get an arbitrary member of the sequence, or None if it is empty.""" @@ -342,6 +360,18 @@ def add_join_to_query(self, query: Any) -> Any: onclause=rel.fk_column == rel.target_column, ) + def target_table(self) -> Table | None: + """Get the related table, or None if it's the same table.""" + if self.relationship is None: + return None + return self.relationship.target_table() + + def target_table_name(self) -> str | None: + """Get the name of the target table, or None if it's the same table.""" + if self.relationship is None: + return None + return self.relationship.target_table_name() + class Buckets: """ diff --git a/datafaker/proposers/intervals.py b/datafaker/proposers/intervals.py index 689a0c0e..3b226111 100644 --- a/datafaker/proposers/intervals.py +++ b/datafaker/proposers/intervals.py @@ -73,11 +73,10 @@ def _get_roles( _set_roles_for_column(role_to_fk_columns, None, col, columns_config) # look for roles in related tables if (fk_relationship := make_foreign_key_relationship(col)) is not None: - target_table = fk_relationship.target_table() ft_conf: dict[str, Any] = get_property( - tables_config, [str(target_table.name), "columns"], {} + tables_config, [fk_relationship.target_table_name(), "columns"], {} ) - for fcol in target_table.columns: + for fcol in fk_relationship.target_table().columns: _set_roles_for_column( role_to_fk_columns, fk_relationship, fcol, ft_conf ) @@ -164,20 +163,20 @@ def function_name(self) -> str: def name(self) -> str: """Get the name of the generator.""" fname = self.function_name() - anchor = self._anchor_related.target - aname = anchor.name - atable = anchor.table - if self._anchor_related.relationship is None: + aname = self._anchor_related.target.name + rel = self._anchor_related.relationship + if rel is None: return f"{fname} [anchored to {aname}]" - atable = anchor.table.name - rel = self._anchor_related.relationship.fk_column.name - return f"{fname} [anchored to {aname} of table {atable} (via {rel})]" + atable = rel.target_table_name() + fk_name = rel.fk_column.name + return f"{fname} [anchored to {aname} of table {atable} (via {fk_name})]" def nominal_kwargs(self) -> dict[str, Any]: """Get the arguments to be entered into ``config.yaml``.""" column = self._column anchor = self._anchor_related.target - if self._anchor_related.relationship is None: + rel = self._anchor_related.relationship + if rel is None: return { "mean_seconds": ( f'SRC_STATS["auto__{column.table.name}"]' @@ -190,12 +189,12 @@ def nominal_kwargs(self) -> dict[str, Any]: "anchor": f'GENERATED_ROW["{anchor.name}"]', } key = f"auto__interval__{column.table.name}__{column.name}" - fk_col = self._anchor_related.relationship.fk_column.name - on_col = self._anchor_related.relationship.target_column.name + fk_col = rel.fk_column.name + on_col = rel.target_column.name return { "dst_db_conn": "dst_db_conn", "anchor_column": f'"{anchor.name}"', - "table": f'"{anchor.table.name}"', + "table": f'"{rel.target_table_name()}"', "mean_seconds": (f'SRC_STATS["{key}"]["results"][0]["mean"]'), "sd_seconds": (f'SRC_STATS["{key}"]["results"][0]["sd"]'), "anchor_row": f'GENERATED_ROW["{fk_col}"]', @@ -377,7 +376,10 @@ def get_proposers( if "start" not in roles: return [] other_start_columns = [ - fk_col for fk_col in roles["start"] if fk_col.target != column + fk_col + for fk_col in roles["start"] + # Don't anchor on ourselves + if not (fk_col.relationship is None and fk_col.target == column) ] return [ prop @@ -423,25 +425,27 @@ def name(self) -> str: """Get the name of the generator.""" fname = self.function_name() aname = self._anchor_related.target.name - if self._anchor_related.relationship is None: + rel = self._anchor_related.relationship + if rel is None: return f"{fname} [from {aname}]" - atable = self._anchor_related.target.table - on_col = self._anchor_related.relationship.fk_column.name - return f"{fname} [from {aname} of table {atable.name} (via {on_col})]" + atable = rel.target_table_name() + on_col = rel.fk_column.name + return f"{fname} [from {aname} of table {atable} (via {on_col})]" def nominal_kwargs(self) -> dict[str, Any]: """Get the arguments to be entered into ``config.yaml``.""" anchor = self._anchor_related.target - if self._anchor_related.relationship is None: + rel = self._anchor_related.relationship + if rel is None: return { "anchor": f'GENERATED_ROW["{anchor.name}"]', } - fk_col = self._anchor_related.relationship.fk_column - on_col = self._anchor_related.relationship.target_column + fk_col = rel.fk_column + on_col = rel.target_column return { "dst_db_conn": "dst_db_conn", "anchor_column": f'"{anchor.name}"', - "table": f'"{anchor.table.name}"', + "table": f'"{rel.target_table_name()}"', "anchor_row": f'GENERATED_ROW["{fk_col.name}"]', "on_column": f'"{on_col.name}"', } @@ -503,6 +507,9 @@ def make_copy_fk_proposers( ) ] + def _can_copy_type(self, copy_from, copy_to) -> bool: + return type(copy_from) is type(copy_to) + def get_proposers( self, columns: list[Column], @@ -519,8 +526,10 @@ def get_proposers( other_source_columns = [ fk_col for fk_col in roles["source"] - if fk_col.target != column - and type(get_column_type(fk_col.target)) is type(ct) + # Don't copy from ourself + if not (fk_col.relationship is None and fk_col.target == column) + # Only copy if the types are compatible + and self._can_copy_type(get_column_type(fk_col.target), ct) ] return [ prop diff --git a/datafaker/proposers/partitioned.py b/datafaker/proposers/partitioned.py index 8db2ee9b..f52dc852 100644 --- a/datafaker/proposers/partitioned.py +++ b/datafaker/proposers/partitioned.py @@ -28,7 +28,7 @@ CovariateQuery, MultivariateNormalProposerFactory, ) -from datafaker.utils import T, get_property, logger +from datafaker.utils import T, get_property, logger, text_list NumericType = Union[int, float] @@ -37,23 +37,6 @@ MAXIMUM_CHOICES = 500 -def text_list(items: Iterable[str]) -> str: - """Concatenate the items with commas and one "and".""" - item_i = iter(items) - try: - last_item = next(item_i) - except StopIteration: - return "" - try: - so_far = next(item_i) - except StopIteration: - return last_item - for item in item_i: - so_far += ", " + last_item - last_item = item - return so_far + " and " + last_item - - @dataclass class RowPartition: """A partition where all the rows have the same pattern of NULLs.""" diff --git a/datafaker/utils.py b/datafaker/utils.py index a6a15321..e7ce255e 100644 --- a/datafaker/utils.py +++ b/datafaker/utils.py @@ -457,6 +457,23 @@ def underline_error(e: SyntaxError) -> str: return "\n" + " " * start + "^" * (end - start) +def text_list(items: Iterable[str], empty="") -> str: + """Concatenate the items with commas and one "and".""" + item_i = iter(items) + try: + last_item = next(item_i) + except StopIteration: + return empty + try: + so_far = next(item_i) + except StopIteration: + return last_item + for item in item_i: + so_far += ", " + last_item + last_item = item + return so_far + " and " + last_item + + def gather_from_ast( errors: MutableSequence[tuple], name: str, diff --git a/docs/source/docker.rst b/docs/source/docker.rst index 62fcfe0a..9c21bb2b 100644 --- a/docs/source/docker.rst +++ b/docs/source/docker.rst @@ -23,7 +23,7 @@ and stores the files produced in a directory called ``output``: .. code-block:: shell mkdir output - docker run --rm --user $(id -u):$(id -g) --network host -e SRC_SCHEMA=public -e SRC_DSN=postgresql://tim:tim@localhost:5432/pagila -itv ./output:data datafaker + docker run --rm --user $(id -u):$(id -g) --network host -e SRC_SCHEMA=public -e SRC_DSN=postgresql://tim:tim@localhost:5432/pagila -itv ./output:/data datafaker You do need to create the output folder first. diff --git a/pyproject.toml b/pyproject.toml index 09096059..aa31415b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "datafaker" -version = "0.5.3" +version = "0.5.4" description = "Generates fake SQL data" authors = ["Tim Band <3266052+tim-band@users.noreply.github.com>"] license = "MIT" diff --git a/tests/examples/date.sql b/tests/examples/date.sql index f24b87b0..037b5649 100644 --- a/tests/examples/date.sql +++ b/tests/examples/date.sql @@ -25,7 +25,8 @@ CREATE TABLE public.happening ( at_time TIMESTAMP WITH TIME ZONE NOT NULL, at_date DATE NOT NULL, other_person_id INTEGER, - comment TEXT + comment TEXT, + previous_happening_id INTEGER ); ALTER TABLE ONLY public.happening ADD CONSTRAINT happening_pkey PRIMARY KEY (id); @@ -33,10 +34,12 @@ ALTER TABLE ONLY public.happening ADD CONSTRAINT person_id_fkey FOREIGN KEY (person_id) REFERENCES public.person(id); ALTER TABLE ONLY public.happening ADD CONSTRAINT other_person_id_fkey FOREIGN KEY (other_person_id) REFERENCES public.person(id); +ALTER TABLE ONLY public.happening + ADD CONSTRAINT previous_happening_id_fkey FOREIGN KEY (previous_happening_id) REFERENCES public.happening(id); ALTER TABLE public.happening OWNER TO postgres; -INSERT INTO public.happening VALUES (1, 'stepped on a tack', 1, '1997-04-20 04:05:06+00:00', '1997-04-20', NULL, NULL); -INSERT INTO public.happening VALUES (2, 'had a dream', 2, '1997-04-20 04:04:16+00:00', '1997-04-20', NULL, NULL); -INSERT INTO public.happening VALUES (3, 'kicked a can', 1, '2001-12-23 07:05:06+00:00', '2001-12-23', NULL, NULL); -INSERT INTO public.happening VALUES (4, 'shared a cake', 2, '2003-11-03 04:15:26+00:00', '2003-11-03', 1, NULL); +INSERT INTO public.happening VALUES (1, 'stepped on a tack', 1, '1997-04-20 04:05:06+00:00', '1997-04-20', NULL, NULL, NULL); +INSERT INTO public.happening VALUES (2, 'had a dream', 2, '1997-04-20 04:04:16+00:00', '1997-04-20', NULL, NULL, 1); +INSERT INTO public.happening VALUES (3, 'kicked a can', 1, '2001-12-23 07:05:06+00:00', '2001-12-23', NULL, NULL, 1); +INSERT INTO public.happening VALUES (4, 'shared a cake', 2, '2003-11-03 04:15:26+00:00', '2003-11-03', 1, NULL, NULL); diff --git a/tests/test_interactive_date_generators.py b/tests/test_interactive_date_generators.py index 3addac8d..9fcd4505 100644 --- a/tests/test_interactive_date_generators.py +++ b/tests/test_interactive_date_generators.py @@ -316,6 +316,31 @@ def test_cross_table_interval_timestamp_vs_date(self) -> None: self.assertEqual(gc.messages[1][0], gc.REQUIRES_SOURCE_DATA_TEXT) self.assertEqual(gc.messages[2][0], gc.PROVIDING_VALUES_TEXT) + def test_regression_date_start_crash(self) -> None: + """Test for regression where ``start`` role on a date leads to crash.""" + config = { + "tables": { + "happening": { + "columns": { + "at_date": { + "roles": ["start"], + }, + }, + }, + } + } + with self._get_cmd(config) as gc: + gc.do_next("happening.at_date") + gc.reset() + gc.do_propose("") + proposals = gc.get_proposals() + provider_name = ( + "generic.anchored_provider.normal_date_fk" + " [anchored to at_date of table happening" + " (via previous_happening_id)]" + ) + self.assertIn(provider_name, proposals.keys()) + # Note that this test won't work with DuckDB because it needs foreign keys to work class ConfigureGeneratorsWithDateMsSqlTests(ConfigureGeneratorsWithDateTests): From 37248af5a0930a0f06d6c28b04c7c40e54162281 Mon Sep 17 00:00:00 2001 From: Tim Band Date: Tue, 29 Sep 2026 11:20:30 +0100 Subject: [PATCH 3/3] Copies are much more restricted You can't copy or anchor from nullable (or via nullable) to non-nullable. You can't copy from an FK to an FK that points to a different table. Bugs fixed. --- datafaker/proposers/base.py | 7 ++ datafaker/proposers/intervals.py | 95 +++++++++++++++---- tests/examples/date.sql | 10 +- tests/test_interactive_copy_generators.py | 106 ++++++++++++++++++++-- tests/test_interactive_date_generators.py | 8 +- 5 files changed, 193 insertions(+), 33 deletions(-) diff --git a/datafaker/proposers/base.py b/datafaker/proposers/base.py index 2c152b40..b51f86fa 100644 --- a/datafaker/proposers/base.py +++ b/datafaker/proposers/base.py @@ -308,6 +308,13 @@ def __init__(self, fk_column: Column, target_column: Column) -> None: table = aliased(target_column.table) self.target_column = table.columns[target_column.name] + def __repr__(self): + """Render the relationship as a JOIN.""" + return ( + f"" + ) + def target_table(self) -> Table: """Get the related table.""" return self.target_column.table diff --git a/datafaker/proposers/intervals.py b/datafaker/proposers/intervals.py index 3b226111..c6142cba 100644 --- a/datafaker/proposers/intervals.py +++ b/datafaker/proposers/intervals.py @@ -49,7 +49,7 @@ def _set_roles_for_column( out[role].append(rc) -def _get_roles( +def _get_all_relative_roles( config: Mapping, column: Column, ) -> dict[str, list[RelatedColumn]]: @@ -83,6 +83,77 @@ def _get_roles( return role_to_fk_columns +def _not_anchored_on_ourselves(column: Column, fk_col: RelatedColumn) -> bool: + """ + Return whether these columns are the exact same thing. + + :param column: The column we are wanting to generate for. + :param fk_col: The column we are considering using as an anchor. + :return: False if ``fk_col`` targets the same column as ``column`` + not through a foreign key, so True if ``fk_col`` is through a FK + or if it doesn't target the same column. + """ + return not (fk_col.relationship is None and fk_col.target == column) + + +def _not_anchoring_nonnullable_to_nullable( + column: Column, fk_col: RelatedColumn +) -> bool: + """ + Return True if ``column`` is nullable or ``fk_col`` is not. + + We are counting ``fk_col`` as nullabe if either the target is nullable + or the foreign key itself. + + This is required if ``fk_col`` can be copied to ``col``. + """ + fk_nullable = ( + fk_col.relationship is not None + and fk_col.relationship.fk_column.nullable is not False + ) + fk_col_nullable = fk_nullable or fk_col.target.nullable + return column.nullable or not fk_col_nullable + + +def _not_copying_incompatible_fk(column: Column, fk_col: RelatedColumn) -> bool: + """Return True if ``column`` and ``fk_col`` are incompatible foreign key types.""" + target = fk_col.target + if not column.foreign_keys: + return not target.foreign_keys + if not target.foreign_keys: + return False + tfks = {c.target_fullname for c in target.foreign_keys} + cfks = {c.target_fullname for c in column.foreign_keys} + return tfks == cfks + + +def _get_roles( + config: Mapping, + column: Column, + role_name: str, +) -> list[RelatedColumn]: + """ + Get all the ``role_name`` roles in this table or a related table. + + :param role_name: The role to return. + :return: The list of all columns in the same table as ``column`` or in + tables directly related to that table that have the named role. This + list has the following columns filtered out: Nullable columns (if + ``column`` is nonnullable), foreign keys if ``column`` is not a + foreign key, non-foreign keys if ``column`` is a foreign key + or foreign keys if ``column`` is a foreign key to a different table. + """ + roles = _get_all_relative_roles(config, column) + if role_name not in roles: + return [] + return [ + fk_col + for fk_col in roles[role_name] + if _not_anchored_on_ourselves(column, fk_col) + and _not_anchoring_nonnullable_to_nullable(column, fk_col) + ] + + def coerce_to_datetime( d: datetime.date | datetime.datetime, tzdonor: datetime.datetime, @@ -372,15 +443,7 @@ def get_proposers( ct = get_column_type(column) if not isinstance(ct, (Date, DateTime)): return [] - roles = _get_roles(self._config, column) - if "start" not in roles: - return [] - other_start_columns = [ - fk_col - for fk_col in roles["start"] - # Don't anchor on ourselves - if not (fk_col.relationship is None and fk_col.target == column) - ] + other_start_columns = _get_roles(self._config, column, "start") return [ prop for anchor in other_start_columns @@ -519,17 +582,9 @@ def get_proposers( if len(columns) != 1: return [] column = columns[0] - roles = _get_roles(self._config, column) - if "source" not in roles: - return [] - ct = get_column_type(column) + fk_cols = _get_roles(self._config, column, "source") other_source_columns = [ - fk_col - for fk_col in roles["source"] - # Don't copy from ourself - if not (fk_col.relationship is None and fk_col.target == column) - # Only copy if the types are compatible - and self._can_copy_type(get_column_type(fk_col.target), ct) + fk_col for fk_col in fk_cols if _not_copying_incompatible_fk(column, fk_col) ] return [ prop diff --git a/tests/examples/date.sql b/tests/examples/date.sql index 037b5649..b8e99be7 100644 --- a/tests/examples/date.sql +++ b/tests/examples/date.sql @@ -8,15 +8,19 @@ CREATE TABLE public.person ( id INTEGER NOT NULL, name TEXT NOT NULL, timestamp_of_birth TIMESTAMP WITH TIME ZONE NOT NULL, - date_of_birth DATE NOT NULL + date_of_birth DATE NOT NULL, + parent_of INTEGER ); ALTER TABLE ONLY public.person ADD CONSTRAINT person_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.person + ADD CONSTRAINT parent_of_fkey FOREIGN KEY (parent_of) REFERENCES public.person(id); ALTER TABLE public.person OWNER TO postgres; -INSERT INTO public.person VALUES (1, 'Bobby', '1991-01-08 11:12:13+00:00', '1991-01-08'); -INSERT INTO public.person VALUES (2, 'Mary', '1989-03-04 20:19:18+00:00', '1989-03-04'); +INSERT INTO public.person VALUES (1, 'Bobby', '1991-01-08 11:12:13+00:00', '1991-01-08', NULL); +INSERT INTO public.person VALUES (2, 'Mary', '1989-03-04 20:19:18+00:00', '1989-03-04', NULL); +INSERT INTO public.person VALUES (3, 'Elizabeth', '1951-10-14 00:12:12+00:00', '1951-10-14', 2); CREATE TABLE public.happening ( id INTEGER NOT NULL, diff --git a/tests/test_interactive_copy_generators.py b/tests/test_interactive_copy_generators.py index d6ea5b28..f620a31f 100644 --- a/tests/test_interactive_copy_generators.py +++ b/tests/test_interactive_copy_generators.py @@ -1,6 +1,7 @@ """ Tests for the configure-generators command. """ +import re from collections.abc import MutableMapping -from typing import Any +from typing import Any, Iterable from sqlalchemy import select from sqlalchemy.orm import aliased @@ -25,7 +26,7 @@ def _get_cmd(self, config: MutableMapping[str, Any]) -> MockGeneratorCmd: ) def test_same_table_copy(self) -> None: - """Test that the cross-table copy is proposed and compared.""" + """Test that the same-table copy is proposed and compared.""" config = { "tables": { "happening": { @@ -102,7 +103,7 @@ def test_cross_table_copy_actually_copies(self) -> None: } with self._get_cmd(config) as gc: # set up our copy proposer - gc.do_next("happening.comment") + gc.do_next("happening.name") gc.reset() gc.do_propose("") proposals = gc.get_proposals() @@ -113,7 +114,7 @@ def test_cross_table_copy_actually_copies(self) -> None: self.assertIn(provider_name, proposals.keys()) prop = proposals[provider_name] gc.do_set(str(prop[0])) - gc.do_next("happening.name") + gc.do_next("happening.comment") gc.reset() gc.do_propose("") proposals = gc.get_proposals() @@ -155,10 +156,103 @@ def test_cross_table_copy_actually_copies(self) -> None: for row in dst_result: if row.opn: saw_other = True - self.assertEqual(row.hn, row.opn) - self.assertEqual(row.hc, row.pn) + self.assertEqual(row.hc, row.opn) + self.assertEqual(row.hn, row.pn) self.assertTrue(saw_other) + def test_cross_table_fk_copy(self) -> None: + """Test that the cross-table foreign key copy is proposed.""" + config = { + "tables": { + "person": { + "columns": { + "parent_of": { + "roles": ["source"], + } + } + } + } + } + with self._get_cmd(config) as gc: + gc.do_next("happening.other_person_id") + gc.reset() + gc.do_propose("") + proposals = gc.get_proposals() + provider_name = ( + "generic.anchored_provider.copy_fk" + " [from parent_of of table person (via person_id)]" + ) + self.assertIn(provider_name, proposals.keys()) + + COPY_FROM_RE = re.compile(r"generic.anchored_provider.copy \[from ([^\]]+)\]") + + def get_same_table_copy_proposals(self, proposals: Iterable[str]) -> set[str]: + """Return a set of columns proposed to copy from.""" + return { + p.group(1) + for p in (self.COPY_FROM_RE.match(prop) for prop in proposals) + if p is not None + } + + def test_different_target_fk_copy_is_not_proposed(self) -> None: + """Test a foreign key copy with a different target table is not proposed.""" + config = { + "tables": { + "happening": { + "columns": { + "previous_happening_id": { + "roles": ["source"], + }, + "person_id": { + "roles": ["source"], + }, + } + } + } + } + with self._get_cmd(config) as gc: + gc.do_next("happening.other_person_id") + gc.reset() + gc.do_propose("") + proposals = gc.get_proposals() + provider_name = "generic.anchored_provider.copy [from person_id]" + self.assertIn(provider_name, proposals.keys()) + all_copies = self.get_same_table_copy_proposals(proposals.keys()) + self.assertIn("person_id", all_copies) + self.assertNotIn("previous_happening_id", all_copies) + + def test_nullable_copy_nonnullable_is_not_proposed(self) -> None: + """Test a copy from a nullable column to a nonnullable column is not proposed.""" + config = { + "tables": { + "happening": { + "columns": { + "other_person_id": { + "roles": ["source"], + }, + "person_id": { + "roles": ["source"], + }, + } + } + } + } + with self._get_cmd(config) as gc: + # person -> other_person is allowed + gc.do_next("happening.other_person_id") + gc.reset() + gc.do_propose("") + proposals = gc.get_proposals() + all_copies = self.get_same_table_copy_proposals(proposals.keys()) + self.assertIn("person_id", all_copies) + # other_person -> person is allowed + gc.do_next("happening.person_id") + gc.reset() + gc.do_propose("") + proposals = gc.get_proposals() + all_copies = self.get_same_table_copy_proposals(proposals.keys()) + self.assertNotIn("other_person_id", all_copies) + # Note that this test won't work with DuckDB because it needs foreign keys to work class ConfigureCopyGeneratorsWithDateMsSqlTests(ConfigureCopyGeneratorsWithDateTests): diff --git a/tests/test_interactive_date_generators.py b/tests/test_interactive_date_generators.py index 9fcd4505..aba9b824 100644 --- a/tests/test_interactive_date_generators.py +++ b/tests/test_interactive_date_generators.py @@ -320,9 +320,9 @@ def test_regression_date_start_crash(self) -> None: """Test for regression where ``start`` role on a date leads to crash.""" config = { "tables": { - "happening": { + "person": { "columns": { - "at_date": { + "date_of_birth": { "roles": ["start"], }, }, @@ -336,8 +336,8 @@ def test_regression_date_start_crash(self) -> None: proposals = gc.get_proposals() provider_name = ( "generic.anchored_provider.normal_date_fk" - " [anchored to at_date of table happening" - " (via previous_happening_id)]" + " [anchored to date_of_birth of table person" + " (via person_id)]" ) self.assertIn(provider_name, proposals.keys())