Added CopyProposer

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.

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
 
 
@@ -201,6 +203,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:"
     )
@@ -478,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"],
@@ -540,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
 
@@ -910,7 +915,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,

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(),
     )

datafaker/proposers/base.py

@@ -2,19 +2,21 @@
 
 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
 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
 
 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 +293,93 @@ def fit_from_buckets(xs: Sequence[NumericType], ys: Sequence[NumericType]) -> fl
     return sum_diff_squared / (count * count)
 
 
+class ForeignKeyRelationship:
+    """How a table is related."""
+
+    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 __repr__(self):
+        """Render the relationship as a JOIN."""
+        return (
+            f"<JOIN {self.target_column.table.original.name} ON"
+            f" {self.fk_column} = {self.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."""
+    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,
+        )
+
+    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:
     """
     Measured buckets for a real distribution.
@@ -308,12 +397,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 +424,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 +433,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 +443,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 +457,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 +479,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)