Source code for test_data_workbench.adaptation.generator_factory

"""Auto-create appropriate generators from discovered schema."""

import re
from typing import Dict, List, Set, Optional, Any
from dataclasses import dataclass
from jinja2 import Template

from test_data_workbench.core.models import (
    SchemaInfo, Table, Column, EntityType, BusinessRule,
    GeneratorConfig, Relationship
)


[docs] @dataclass class GeneratorSet: """Collection of generated Python classes and their metadata.""" generators: Dict[str, str] # entity_name -> python_code dependencies: Dict[str, List[str]] # entity_name -> dependency_list configs: Dict[str, GeneratorConfig] fallback_generators: Dict[str, str]
[docs] @dataclass class RelationshipManager: """Manages foreign keys and constraints across generators.""" relationships: List[Relationship] dependency_order: List[str] reference_cache: Dict[str, List[Any]]
[docs] class GeneratorFactory: """Auto-create appropriate generators from discovered schema.""" def __init__(self): self.faker_mapping = { 'varchar': 'self.fake.text', 'text': 'self.fake.text', 'integer': 'self.fake.random_int', 'bigint': 'self.fake.random_int', 'decimal': 'self.fake.pydecimal', 'numeric': 'self.fake.pydecimal', 'boolean': 'self.fake.boolean', 'date': 'self.fake.date', 'timestamp': 'self.fake.date_time', 'timestamptz': 'self.fake.date_time_this_year', 'uuid': 'self.fake.uuid4' } self.entity_specializations = { EntityType.USER: self._create_user_generator, EntityType.CUSTOMER: self._create_customer_generator, EntityType.PRODUCT: self._create_product_generator, EntityType.ORDER: self._create_order_generator, EntityType.REVIEW: self._create_review_generator }
[docs] def create_generators_from_schema(self, schema: SchemaInfo) -> GeneratorSet: """Generate Python classes for each discovered entity type.""" generators = {} dependencies = {} configs = {} fallbacks = {} # Sort tables by dependency order ordered_tables = self._sort_by_dependencies(schema.tables, schema.relationships) for table in ordered_tables: try: # Create specialized generator if entity type is recognized if table.entity_type in self.entity_specializations: generator_code = self.entity_specializations[table.entity_type](table, schema) else: generator_code = self._create_generic_generator(table, schema) generators[table.name] = generator_code dependencies[table.name] = self._get_table_dependencies(table, schema) configs[table.name] = self._create_generator_config(table) fallbacks[table.name] = self._create_fallback_generator(table) except Exception: # Always provide fallback fallbacks[table.name] = self._create_simple_fallback(table) return GeneratorSet( generators=generators, dependencies=dependencies, configs=configs, fallback_generators=fallbacks )
[docs] def build_relationship_handlers(self, relationships: List[Relationship]) -> RelationshipManager: """Create foreign key and constraint management system.""" dependency_order = self._calculate_dependency_order(relationships) return RelationshipManager( relationships=relationships, dependency_order=dependency_order, reference_cache={} )
def _sort_by_dependencies(self, tables: List[Table], relationships: List[Relationship]) -> List[Table]: """Sort tables to generate independent entities first (with cycle detection).""" table_deps = {} for table in tables: deps = [rel.to_table for rel in relationships if rel.from_table == table.name] table_deps[table.name] = deps # Topological sort with cycle detection sorted_names = [] visited = set() in_stack = set() def visit(name: str): if name in visited or name not in table_deps: return if name in in_stack: return # Cycle detected - break it in_stack.add(name) for dep in table_deps[name]: visit(dep) in_stack.discard(name) visited.add(name) sorted_names.append(name) for table in tables: visit(table.name) # Return tables in sorted order table_map = {t.name: t for t in tables} return [table_map[name] for name in sorted_names if name in table_map] def _create_user_generator(self, table: Table, schema: SchemaInfo) -> str: """Create specialized user/customer generator.""" template = Template(''' from faker import Faker from typing import Dict, Any, Optional import random class {{ class_name }}Generator: """Generated {{ entity_type }} data generator.""" def __init__(self, seed: Optional[int] = None): self.fake = Faker() self._generated_emails = set() {%- if has_sequential_pk %} self._pk_counter = 0 {%- endif %} if seed is not None: Faker.seed(seed) random.seed(seed) def generate(self, count: int = 1) -> list[dict]: """Generate {{ entity_type }} records.""" records = [] for _ in range(count): record = { {%- for column in columns %} '{{ column.name }}': {{ column.generator }}, {%- endfor %} } records.append(record) return records def _unique_email(self) -> str: """Generate unique email address.""" while True: email = self.fake.email() if email not in self._generated_emails: self._generated_emails.add(email) return email {%- if has_sequential_pk %} def _next_id(self) -> int: """Return the next sequential primary key value.""" self._pk_counter += 1 return self._pk_counter {%- endif %} ''') columns = [] for col in table.columns: if self._needs_sequential_pk(col): gen = f"self._next_id()" elif col.is_primary_key and 'id' in col.name.lower(): gen = f"self.fake.uuid4()" elif 'email' in col.name.lower(): gen = f"self._unique_email()" elif 'name' in col.name.lower(): gen = f"self.fake.name()" elif 'phone' in col.name.lower(): gen = f"self.fake.phone_number()" else: gen = self._get_column_generator(col) columns.append({'name': col.name, 'generator': gen}) return template.render( class_name=self._to_class_name(table.name), entity_type=table.entity_type.value, columns=columns, has_sequential_pk=self._table_needs_pk_counter(table) ) def _create_product_generator(self, table: Table, schema: SchemaInfo) -> str: """Create specialized product generator. A discovered FK relationship (e.g. ``products.category_id -> categories.id``) takes precedence over the name-based 'category' heuristic below: that heuristic predates FK-awareness and hands out a display-only category *label* (e.g. "Electronics"), which is not a valid foreign key value. When the schema has resolved a real relationship for the column, use the actual parent id via ``reference_data`` instead - see ``_fk_relationships_for_table``. """ fk_relationships = self._fk_relationships_for_table(table, schema) uses_reference_helper = any( fk_relationships[col.name].to_table != table.name for col in table.columns if col.name in fk_relationships ) template = Template(''' from faker import Faker from typing import Dict, Any, Optional, List import random class {{ class_name }}Generator: """Generated {{ entity_type }} data generator.""" def __init__(self, {% if uses_reference_helper %}reference_data: Dict[str, List[Any]] = None, {% endif %}seed: Optional[int] = None): self.fake = Faker() self.categories = ['Electronics', 'Clothing', 'Home', 'Sports', 'Books'] {%- if uses_reference_helper %} self.reference_data = reference_data or {} {%- endif %} {%- if has_sequential_pk %} self._pk_counter = 0 {%- endif %} if seed is not None: Faker.seed(seed) random.seed(seed) def generate(self, count: int = 1) -> list[dict]: """Generate {{ entity_type }} records.""" records = [] for _ in range(count): record = { {%- for column in columns %} '{{ column.name }}': {{ column.generator }}, {%- endfor %} } records.append(record) return records {%- if uses_reference_helper %} def _get_reference_id(self, table_name: str) -> Any: """Get a random id from a referenced parent table's generated values.""" if table_name in self.reference_data and self.reference_data[table_name]: return random.choice(self.reference_data[table_name]) return None {%- endif %} {%- if has_sequential_pk %} def _next_id(self) -> int: """Return the next sequential primary key value.""" self._pk_counter += 1 return self._pk_counter {%- endif %} ''') columns = [] for col in table.columns: if self._needs_sequential_pk(col): gen = f"self._next_id()" elif col.name in fk_relationships and fk_relationships[col.name].to_table != table.name: gen = f"self._get_reference_id('{fk_relationships[col.name].to_table}')" elif 'name' in col.name.lower() or 'title' in col.name.lower(): gen = f"self.fake.catch_phrase()" elif 'price' in col.name.lower(): gen = f"round(random.uniform(10, 1000), 2)" elif 'category' in col.name.lower(): gen = f"random.choice(self.categories)" elif 'description' in col.name.lower(): gen = f"self.fake.text(max_nb_chars=200)" else: gen = self._get_column_generator(col) columns.append({'name': col.name, 'generator': gen}) return template.render( class_name=self._to_class_name(table.name), entity_type=table.entity_type.value, columns=columns, has_sequential_pk=self._table_needs_pk_counter(table), uses_reference_helper=uses_reference_helper, ) def _create_order_generator(self, table: Table, schema: SchemaInfo) -> str: """Create specialized order generator with dependencies.""" template = Template(''' from faker import Faker from typing import Dict, Any, Optional, List import random from datetime import datetime, timedelta class {{ class_name }}Generator: """Generated {{ entity_type }} data generator with dependencies.""" def __init__(self, reference_data: Dict[str, List[Any]] = None, seed: Optional[int] = None): self.fake = Faker() self.reference_data = reference_data or {} self.statuses = ['pending', 'processing', 'shipped', 'delivered', 'cancelled'] {%- if has_sequential_pk %} self._pk_counter = 0 {%- endif %} if seed is not None: Faker.seed(seed) random.seed(seed) def generate(self, count: int = 1) -> list[dict]: """Generate {{ entity_type }} records.""" records = [] for _ in range(count): record = { {%- for column in columns %} '{{ column.name }}': {{ column.generator }}, {%- endfor %} } records.append(record) return records def _get_reference_id(self, table_name: str) -> Any: """Get random ID from reference table.""" if table_name in self.reference_data and self.reference_data[table_name]: return random.choice(self.reference_data[table_name]) return self.fake.uuid4() # Fallback {%- if has_sequential_pk %} def _next_id(self) -> int: """Return the next sequential primary key value.""" self._pk_counter += 1 return self._pk_counter {%- endif %} ''') columns = [] for col in table.columns: if self._needs_sequential_pk(col): gen = f"self._next_id()" elif col.is_foreign_key: # Try to infer referenced table from column name ref_table = self._infer_reference_table(col.name) gen = f"self._get_reference_id('{ref_table}')" elif 'status' in col.name.lower(): gen = f"random.choice(self.statuses)" elif 'total' in col.name.lower() or 'amount' in col.name.lower(): gen = f"round(random.uniform(20, 500), 2)" else: gen = self._get_column_generator(col) columns.append({'name': col.name, 'generator': gen}) return template.render( class_name=self._to_class_name(table.name), entity_type=table.entity_type.value, columns=columns, has_sequential_pk=self._table_needs_pk_counter(table) ) def _create_customer_generator(self, table: Table, schema: SchemaInfo) -> str: """Alias for user generator.""" return self._create_user_generator(table, schema) def _create_review_generator(self, table: Table, schema: SchemaInfo) -> str: """Create specialized review generator.""" template = Template(''' from faker import Faker from typing import Dict, Any, Optional, List import random class {{ class_name }}Generator: """Generated {{ entity_type }} data generator.""" def __init__(self, reference_data: Dict[str, List[Any]] = None, seed: Optional[int] = None): self.fake = Faker() self.reference_data = reference_data or {} {%- if has_sequential_pk %} self._pk_counter = 0 {%- endif %} if seed is not None: Faker.seed(seed) random.seed(seed) def generate(self, count: int = 1) -> list[dict]: """Generate {{ entity_type }} records.""" records = [] for _ in range(count): rating = random.randint(1, 5) record = { {%- for column in columns %} '{{ column.name }}': {{ column.generator }}, {%- endfor %} } records.append(record) return records def _get_reference_id(self, table_name: str) -> Any: """Get random ID from reference table.""" if table_name in self.reference_data and self.reference_data[table_name]: return random.choice(self.reference_data[table_name]) return self.fake.uuid4() {%- if has_sequential_pk %} def _next_id(self) -> int: """Return the next sequential primary key value.""" self._pk_counter += 1 return self._pk_counter {%- endif %} ''') columns = [] for col in table.columns: if self._needs_sequential_pk(col): gen = f"self._next_id()" elif col.is_foreign_key: ref_table = self._infer_reference_table(col.name) gen = f"self._get_reference_id('{ref_table}')" elif 'rating' in col.name.lower(): gen = f"rating" elif 'comment' in col.name.lower() or 'review' in col.name.lower(): gen = f"self.fake.text(max_nb_chars=300)" else: gen = self._get_column_generator(col) columns.append({'name': col.name, 'generator': gen}) return template.render( class_name=self._to_class_name(table.name), entity_type=table.entity_type.value, columns=columns, has_sequential_pk=self._table_needs_pk_counter(table) ) def _create_generic_generator(self, table: Table, schema: SchemaInfo) -> str: """Create generic generator for unknown entity types. A generic table has none of the hand-written FK-awareness the specialized order/review templates do, so foreign keys are resolved generically here from the schema's declared relationships instead of an unrelated random int: - A column with a relationship to a *different* table draws its value from that parent's generated ids via ``reference_data``, threaded in exactly the way the order/review generators expect it (see ``adaptation.constraint_check.generate_sample_dataset`` and ``rapid_deployment`` for how it gets populated). - A column whose relationship targets THIS SAME table (a self-referential FK, e.g. ``categories.parent_id -> categories.id``) draws from ids already generated earlier in this same batch, or is left null - there is no earlier row to point at yet for the first one(s). - An FK-flagged column with no explicit discovered relationship falls back to the same name-based table guess the order/review templates use, so it still gets a plausible-looking value. """ fk_relationships = self._fk_relationships_for_table(table, schema) pk_column = table.primary_keys[0] if table.primary_keys else None columns = [] uses_reference_helper = False uses_self_reference_helper = False self_ref_targets: List[str] = [] for col in table.columns: if self._needs_sequential_pk(col): gen = "self._next_id()" elif col.name in fk_relationships: rel = fk_relationships[col.name] if rel.to_table == table.name: gen = f"self._self_reference_id('{rel.to_column}', {col.nullable})" uses_self_reference_helper = True if rel.to_column not in self_ref_targets: self_ref_targets.append(rel.to_column) else: gen = f"self._get_reference_id('{rel.to_table}')" uses_reference_helper = True elif col.is_foreign_key: ref_table = self._infer_reference_table(col.name) gen = f"self._get_reference_id('{ref_table}')" uses_reference_helper = True else: gen = self._get_column_generator(col) columns.append({'name': col.name, 'generator': gen}) template = Template(''' from faker import Faker from typing import Dict, Any, Optional, List import random class {{ class_name }}Generator: """Generated generic data generator.""" def __init__(self, {% if has_relations %}reference_data: Dict[str, List[Any]] = None, {% endif %}seed: Optional[int] = None): self.fake = Faker() {%- if has_relations %} self.reference_data = reference_data or {} {%- endif %} {%- if has_sequential_pk %} self._pk_counter = 0 {%- endif %} {%- if uses_self_reference_helper %} self._self_generated: Dict[str, list] = {} {%- endif %} if seed is not None: Faker.seed(seed) random.seed(seed) def generate(self, count: int = 1) -> list[dict]: """Generate generic records.""" records = [] for _ in range(count): record = { {%- for column in columns %} '{{ column.name }}': {{ column.generator }}, {%- endfor %} } records.append(record) {%- if uses_self_reference_helper %} {%- for target in self_ref_targets %} self._self_generated.setdefault('{{ target }}', []).append(record.get('{{ target }}')) {%- endfor %} {%- endif %} return records {%- if uses_reference_helper %} def _get_reference_id(self, table_name: str) -> Any: """Get a random id from a referenced parent table's generated values.""" if table_name in self.reference_data and self.reference_data[table_name]: return random.choice(self.reference_data[table_name]) return None {%- endif %} {%- if uses_self_reference_helper %} def _self_reference_id(self, target_column: str, nullable: bool) -> Any: """Pick a value already generated earlier in this batch for a self-referential foreign key, or None for the first row(s) (roots), when the column is nullable.""" prior = self._self_generated.get(target_column) or [] if prior and (not nullable or random.random() < 0.7): return random.choice(prior) if nullable: return None {%- if has_sequential_pk %} return self._pk_counter if target_column == '{{ pk_column_name }}' else None {%- else %} return None {%- endif %} {%- endif %} {%- if has_sequential_pk %} def _next_id(self) -> int: """Return the next sequential primary key value.""" self._pk_counter += 1 return self._pk_counter {%- endif %} ''') return template.render( class_name=self._to_class_name(table.name), columns=columns, has_sequential_pk=self._table_needs_pk_counter(table), has_relations=uses_reference_helper or uses_self_reference_helper, uses_reference_helper=uses_reference_helper, uses_self_reference_helper=uses_self_reference_helper, self_ref_targets=self_ref_targets, pk_column_name=pk_column.name if pk_column else None, ) def _is_integer_column(self, column: Column) -> bool: """True if the column's declared type is in the integer family (integer, bigint, smallint, int4, ...).""" return 'int' in column.data_type.lower() def _needs_sequential_pk(self, column: Column) -> bool: """True if this column is an integer-typed primary key and should therefore get a contiguous sequence (1, 2, 3, ...) instead of a random or uuid value. Non-integer primary keys (uuid, string, etc.) are left to their existing type/name-based strategy.""" return column.is_primary_key and self._is_integer_column(column) def _table_needs_pk_counter(self, table: Table) -> bool: """True if any column in the table needs the sequential-id counter, used to decide whether to emit the counter machinery at all.""" return any(self._needs_sequential_pk(col) for col in table.columns) def _temporal_kind(self, column_name: str) -> Optional[str]: """Classify a column NAME as temporal, independent of its declared database type. Returns ``'date'`` for date-only names (bare ``date``, or a ``*_date`` suffix like ``birth_date``), ``'datetime'`` for names implying a timestamp (``*_at`` such as ``created_at``/ ``updated_at``/``deleted_at``, or bare ``timestamp``/ ``datetime``), and ``None`` when the name doesn't match a well-known temporal pattern. Deliberately conservative and suffix/token based (not a bare substring check) so that e.g. ``update_flag`` or ``validated`` - which contain "date" as a substring - are not misclassified. """ name = column_name.lower() if name in ('timestamp', 'datetime'): return 'datetime' if 'timestamp' in name or 'datetime' in name: return 'datetime' if name == 'date': return 'date' tokens = [t for t in re.split(r'[_\-]+', name) if t] if not tokens: return None if tokens[-1] == 'at': return 'datetime' if tokens[-1] == 'date': return 'date' return None def _temporal_generator_for(self, column_name: str) -> str: """Faker call for a name-detected temporal column.""" kind = self._temporal_kind(column_name) return 'self.fake.date()' if kind == 'date' else 'self.fake.date_time()' def _get_column_generator(self, column: Column) -> str: """Map column type to appropriate Faker method.""" col_type = column.data_type.lower() # Handle specific patterns in column names col_name_lower = column.name.lower() if 'email' in col_name_lower: return 'self.fake.email()' elif 'phone' in col_name_lower: return 'self.fake.phone_number()' elif 'name' in col_name_lower: return 'self.fake.name()' elif 'address' in col_name_lower: return 'self.fake.address()' elif 'city' in col_name_lower: return 'self.fake.city()' elif 'country' in col_name_lower: return 'self.fake.country()' elif 'url' in col_name_lower: return 'self.fake.url()' # Map by data type for db_type, faker_method in self.faker_mapping.items(): if db_type in col_type: # A generic textual type can't tell us whether this is # really a timestamp - SQLite in particular stores dates # with TEXT affinity. Defer to name-based detection before # accepting the generic varchar/text mapping. if db_type in ('varchar', 'text') and self._temporal_kind(column.name): return self._temporal_generator_for(column.name) return faker_method + '()' # Default fallback - also check name-based temporal detection # before generating lorem-ipsum text for a column the type map # didn't recognize at all. if self._temporal_kind(column.name): return self._temporal_generator_for(column.name) return 'self.fake.text(max_nb_chars=50)' def _create_fallback_generator(self, table: Table) -> str: """Create simple fallback generator that always works.""" template = Template(''' import random import string from datetime import datetime, timedelta from typing import Optional class {{ class_name }}FallbackGenerator: """Simple fallback generator - always works.""" def __init__(self, seed: Optional[int] = None): {%- if has_sequential_pk %} self._pk_counter = 0 {%- endif %} if seed is not None: random.seed(seed) def generate(self, count: int = 1) -> list[dict]: """Generate basic records.""" records = [] for _ in range(count): record = { {%- for column in columns %} '{{ column.name }}': {{ column.generator }}, {%- endfor %} } records.append(record) return records {%- if has_sequential_pk %} def _next_id(self) -> int: """Return the next sequential primary key value.""" self._pk_counter += 1 return self._pk_counter {%- endif %} ''') columns = [] for col in table.columns: if self._needs_sequential_pk(col): gen = f"self._next_id()" elif 'int' in col.data_type.lower(): gen = f"random.randint(1, 1000)" elif 'bool' in col.data_type.lower(): gen = f"random.choice([True, False])" elif 'date' in col.data_type.lower() or 'time' in col.data_type.lower() or self._temporal_kind(col.name): # A random-but-seeded timestamp within 2024, instead of the # wall-clock time - keeps generated records reproducible for # a given seed instead of drifting with every run. Also # catches name-detected temporal columns (created_at, etc.) # declared with a generic string type, which would otherwise # fall through to the random-letters branch below. gen = f"datetime(2024, 1, 1) + timedelta(seconds=random.randint(0, 31536000))" else: gen = f"''.join(random.choices(string.ascii_letters, k=10))" columns.append({'name': col.name, 'generator': gen}) return template.render( class_name=self._to_class_name(table.name), columns=columns, has_sequential_pk=self._table_needs_pk_counter(table) ) def _create_simple_fallback(self, table: Table) -> str: """Ultra-simple fallback for critical failures.""" table_name = table.name return f''' def generate_{table_name}(count=1): return [{{"id": i, "data": f"sample_{table_name}_{{i}}"}} for i in range(count)] ''' def _get_table_dependencies(self, table: Table, schema: SchemaInfo) -> List[str]: """Get list of tables this table depends on.""" dependencies = [] for rel in schema.relationships: if rel.from_table == table.name: dependencies.append(rel.to_table) return dependencies def _create_generator_config(self, table: Table) -> GeneratorConfig: """Create configuration for generator.""" return GeneratorConfig( entity_type=table.entity_type, table_name=table.name, dependencies=[] ) def _calculate_dependency_order(self, relationships: List[Relationship]) -> List[str]: """Calculate order to generate data respecting dependencies (with cycle detection).""" deps = {} for rel in relationships: if rel.from_table not in deps: deps[rel.from_table] = [] deps[rel.from_table].append(rel.to_table) # Topological sort with cycle detection order = [] visited = set() in_stack = set() def visit(table: str): if table in visited: return if table in in_stack: return # Cycle detected - break it in_stack.add(table) for dep in deps.get(table, []): visit(dep) in_stack.discard(table) visited.add(table) order.append(table) for table in deps: visit(table) return order def _to_class_name(self, table_name: str) -> str: """Convert table name to PascalCase class name.""" return ''.join(word.capitalize() for word in table_name.split('_')) def _infer_reference_table(self, column_name: str) -> str: """Infer referenced table from foreign key column name.""" # Common patterns: user_id -> users, product_id -> products if column_name.endswith('_id'): base = column_name[:-3] return base + 's' if not base.endswith('s') else base return 'unknown_table' def _fk_relationships_for_table(self, table: Table, schema: SchemaInfo) -> Dict[str, Relationship]: """Map this table's FK column name -> its declared Relationship, for every column the schema analyzer resolved an explicit relationship for. Used to draw a real ``to_table``/``to_column`` instead of guessing a parent table name from the column's own name (see :meth:`_infer_reference_table`), and to detect the self-referential case (``rel.to_table == table.name``, e.g. ``categories.parent_id -> categories.id``).""" return { rel.from_column: rel for rel in schema.relationships if rel.from_table == table.name }