"""Database schema reverse-engineering for unknown production systems."""
import asyncio
import re
from typing import Dict, List, Optional, Tuple, Set
from dataclasses import dataclass
import pandas as pd
import sqlalchemy as sa
from sqlalchemy import text, inspect
from test_data_workbench.core.models import (
SchemaInfo, Table, Column, Relationship, BusinessRule,
EntityType, ConstraintType, PIIColumn
)
# Column-name substrings that indicate a column likely holds personally
# identifiable information. This is a name-only heuristic - it needs no
# row data, so it works even when metadata_only=True skips every SQL
# statement against table data.
#
# Matching is word-boundary-aware: a keyword must be its own token,
# delimited by string start/end, an underscore, or a camelCase transition
# (handled by _normalize_column_name below). A bare substring like "email"
# appearing mid-word would not match. A tradeoff of this precision-first
# approach: keywords glued together with no separator at all (e.g. a column
# literally named "zipcode" with no underscore) won't match "zip" as a
# bounded token; use snake_case or camelCase column names to be detected.
#
# Deliberate judgment calls (documented, not oversights):
# - bare "name" is NOT flagged - too ambiguous ("product_name" vs. a
# person's name); only compound person-name columns are flagged.
# - "username" is NOT flagged - it's an identifier/handle, not a
# plaintext contact detail.
#
# "address" is deliberately generic (street/city/zip/postal all imply a
# physical mailing address), but English also uses "address" for other
# concepts ("email address", "IP address"). classify_column_pii() breaks
# such ties by preferring the longest matching keyword, so those get their
# own explicit compound keywords ("email_address", "ip_address") to win
# out over the shorter, more generic "address".
PII_NAME_PATTERNS: Dict[str, List[str]] = {
"email": ["email", "email_address"],
"name": ["first_name", "last_name", "full_name"],
"phone": ["phone"],
"ssn": ["ssn", "social"],
"address": ["address", "street", "city", "zip", "postal"],
"dob": ["dob", "birth"],
"financial": ["credit", "card", "cvv", "iban"],
"ip_address": ["ip_address", "ip"],
"passport": ["passport"],
"license": ["license"],
}
def _normalize_column_name(name: str) -> str:
"""Lowercase a column name and insert underscores at camelCase
boundaries, so word-boundary matching works regardless of the
naming convention the source database uses."""
with_boundaries = re.sub(r'(?<=[a-z0-9])(?=[A-Z])', '_', name)
return with_boundaries.lower()
[docs]
def classify_column_pii(column_name: str) -> Optional[str]:
"""Return the PII category for a column name, or None if it doesn't
match any known heuristic. Pure name-based check - no DB access.
When a name matches keywords from more than one category (e.g.
"ip_address" and "email_address" both contain the bounded token
"address", which also belongs to the "address" category), the
longest - i.e. most specific - matching keyword wins, so those are
categorized as ip_address / email respectively rather than the more
generic address.
"""
normalized = _normalize_column_name(column_name)
best_category: Optional[str] = None
best_keyword_length = 0
for category, keywords in PII_NAME_PATTERNS.items():
for keyword in keywords:
pattern = r'(?<![a-z0-9])' + re.escape(keyword) + r'(?![a-z0-9])'
if re.search(pattern, normalized) and len(keyword) > best_keyword_length:
best_category = category
best_keyword_length = len(keyword)
return best_category
[docs]
class SchemaAnalyzer:
"""Reverse-engineer any PostgreSQL e-commerce database schema."""
def __init__(self):
self.entity_patterns = {
EntityType.USER: ['user', 'account', 'member', 'auth', 'profile'],
EntityType.CUSTOMER: ['customer', 'client', 'buyer', 'consumer'],
EntityType.PRODUCT: ['product', 'item', 'goods', 'catalog', 'sku'],
EntityType.ORDER: ['order', 'purchase', 'transaction', 'cart', 'checkout'],
EntityType.REVIEW: ['review', 'rating', 'comment', 'feedback', 'testimonial'],
EntityType.CATEGORY: ['category', 'tag', 'type', 'classification', 'group'],
EntityType.INVENTORY: ['inventory', 'stock', 'warehouse', 'supply'],
EntityType.PAYMENT: ['payment', 'billing', 'invoice', 'charge', 'transaction'],
EntityType.SHIPPING: ['shipping', 'delivery', 'fulfillment', 'address'],
EntityType.ANALYTICS: ['analytics', 'metrics', 'stats', 'log', 'event']
}
[docs]
async def analyze_production_schema(self, connection_string: str,
metadata_only: bool = False) -> SchemaInfo:
"""Auto-discover tables, relationships, constraints, data patterns.
When metadata_only=True, only the SQLAlchemy inspector (schema
reflection: get_table_names, get_columns, get_pk_constraint,
get_foreign_keys) is used to describe structure. No SELECT is
issued against any table's row data: sample values are skipped,
row counts are left as None (unknown), and data-derived business
rule detection is skipped entirely. PII column flagging still runs,
since it is a pure column-name heuristic.
"""
try:
engine = sa.create_engine(connection_string, pool_pre_ping=True)
inspector = inspect(engine)
# Get database name
db_name = engine.url.database or "unknown"
# Analyze tables in parallel for speed
table_names = inspector.get_table_names()
tasks = [
self._analyze_table(engine, inspector, table_name, metadata_only)
for table_name in table_names[:20] # Limit for rapid analysis
]
tables = await asyncio.gather(*tasks)
tables = [t for t in tables if t is not None]
# Detect relationships
relationships = self._detect_relationships(tables, inspector)
# Infer business rules from sample data (skipped in metadata_only mode)
business_rules = await self._detect_business_rules(engine, tables, metadata_only)
# Name-based PII flagging - structural only, safe in any mode
pii_columns = self._detect_pii_columns(tables)
return SchemaInfo(
database_name=db_name,
tables=tables,
relationships=relationships,
business_rules=business_rules,
analysis_metadata={
'total_tables': len(table_names),
'analyzed_tables': len(tables),
'analysis_type': 'rapid_production',
'metadata_only': metadata_only,
},
pii_columns=pii_columns,
)
except Exception as e:
# Fallback to minimal schema for team safety
return SchemaInfo(
database_name="error_fallback",
tables=[],
relationships=[],
business_rules=[],
analysis_metadata={'error': str(e), 'fallback_mode': True}
)
async def _analyze_table(self, engine: sa.Engine, inspector, table_name: str,
metadata_only: bool = False) -> Optional[Table]:
"""Analyze single table structure and classify entity type.
metadata_only=True skips every query against row data: sample
values come back empty and row_count is left as None (unknown).
"""
try:
# Get column information
columns_info = inspector.get_columns(table_name)
columns = []
for col_info in columns_info:
constraints = []
# Check constraints
if col_info.get('nullable') is False:
constraints.append(ConstraintType.NOT_NULL)
# Get sample values for pattern detection (skipped in metadata_only mode)
if metadata_only:
sample_values = []
else:
sample_values = await self._get_sample_values(engine, table_name, col_info['name'])
column = Column(
name=col_info['name'],
data_type=str(col_info['type']),
nullable=col_info.get('nullable', True),
default=col_info.get('default'),
constraints=constraints,
sample_values=sample_values
)
columns.append(column)
# Add primary key constraints
pk_constraint = inspector.get_pk_constraint(table_name)
if pk_constraint and pk_constraint['constrained_columns']:
for col_name in pk_constraint['constrained_columns']:
for column in columns:
if column.name == col_name:
column.constraints.append(ConstraintType.PRIMARY_KEY)
# Add foreign key constraints
fk_constraints = inspector.get_foreign_keys(table_name)
for fk in fk_constraints:
for col_name in fk['constrained_columns']:
for column in columns:
if column.name == col_name:
column.constraints.append(ConstraintType.FOREIGN_KEY)
# Get row count estimate (skipped in metadata_only mode - unknown)
if metadata_only:
row_count = None
else:
row_count = await self._estimate_row_count(engine, table_name)
# Classify entity type
entity_type = self.identify_entity_types([table_name])[table_name]
return Table(
name=table_name,
columns=columns,
row_count=row_count,
entity_type=entity_type
)
except Exception:
return None
async def _get_sample_values(self, engine: sa.Engine, table_name: str,
column_name: str, limit: int = 10) -> List:
"""Get sample column values for pattern detection."""
try:
with engine.connect() as conn:
query = text(f'SELECT DISTINCT "{column_name}" FROM "{table_name}" LIMIT {limit}')
result = conn.execute(query)
return [row[0] for row in result if row[0] is not None]
except Exception:
return []
async def _estimate_row_count(self, engine: sa.Engine, table_name: str) -> int:
"""Fast row count estimate."""
try:
with engine.connect() as conn:
# Use EXPLAIN for fast estimate on large tables
query = text(f'SELECT COUNT(*) FROM "{table_name}" LIMIT 1000')
result = conn.execute(query)
return result.scalar() or 0
except Exception:
return 0
[docs]
def identify_entity_types(self, table_names: List[str]) -> Dict[str, EntityType]:
"""Classify tables as users, products, orders, reviews, etc."""
classifications = {}
for table_name in table_names:
table_lower = table_name.lower()
best_match = EntityType.UNKNOWN
max_score = 0
for entity_type, patterns in self.entity_patterns.items():
score = 0
for pattern in patterns:
if pattern in table_lower:
score += len(pattern) / len(table_lower)
if score > max_score:
max_score = score
best_match = entity_type
classifications[table_name] = best_match
return classifications
def _detect_relationships(self, tables: List[Table], inspector) -> List[Relationship]:
"""Detect table relationships from foreign keys and naming patterns."""
relationships = []
for table in tables:
try:
fk_constraints = inspector.get_foreign_keys(table.name)
for fk in fk_constraints:
if fk['constrained_columns'] and fk['referred_columns']:
relationship = Relationship(
from_table=table.name,
to_table=fk['referred_table'],
from_column=fk['constrained_columns'][0],
to_column=fk['referred_columns'][0],
relationship_type="many_to_one",
strength=0.9 # High confidence for explicit FKs
)
relationships.append(relationship)
except Exception:
continue
return relationships
def _detect_pii_columns(self, tables: List[Table]) -> List[PIIColumn]:
"""Flag likely-PII columns purely from their names.
No row data is inspected, so this runs safely in metadata_only
mode as well as full analysis mode. See PII_NAME_PATTERNS for the
keyword list and documented judgment calls.
"""
found = []
for table in tables:
for column in table.columns:
category = classify_column_pii(column.name)
if category is not None:
found.append(PIIColumn(table=table.name, column=column.name, category=category))
return found
async def _detect_business_rules(self, engine: sa.Engine, tables: List[Table],
metadata_only: bool = False) -> List[BusinessRule]:
"""Infer constraints and patterns from existing data.
metadata_only=True skips this entirely - it is purely data-derived
(SELECT * LIMIT 100), so no SQL is issued against table data.
"""
rules = []
if metadata_only:
return rules
for table in tables[:5]: # Limit for rapid analysis
try:
# Sample data for pattern detection
with engine.connect() as conn:
query = text(f'SELECT * FROM "{table.name}" LIMIT 100')
df = pd.read_sql(query, conn)
if not df.empty:
rules.extend(self._analyze_data_patterns(table, df))
except Exception:
continue
return rules
def _analyze_data_patterns(self, table: Table, df: pd.DataFrame) -> List[BusinessRule]:
"""Analyze DataFrame to detect business rules."""
rules = []
for column in df.columns:
col_data = df[column].dropna()
if len(col_data) < 10:
continue
# Email pattern detection
if any(re.match(r'^[\w\.-]+@[\w\.-]+\.\w+$', str(val)) for val in col_data.head()):
rules.append(BusinessRule.create(
f"Column {column} contains email addresses",
"format",
[column],
confidence=0.9
))
# Phone pattern detection
if any(re.match(r'[\+]?[\d\s\-\(\)]{10,}', str(val)) for val in col_data.head()):
rules.append(BusinessRule.create(
f"Column {column} contains phone numbers",
"format",
[column],
confidence=0.8
))
# Numeric range detection
if df[column].dtype in ['int64', 'float64']:
min_val, max_val = col_data.min(), col_data.max()
if min_val >= 0 and max_val <= 100:
rules.append(BusinessRule.create(
f"Column {column} appears to be percentage/rating (0-100)",
"range",
[column],
confidence=0.7
))
return rules