sqlseed 0.1.0__py3-none-any.whl

sqlseed/config/loader.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import yaml
+
+from sqlseed._utils.logger import get_logger
+from sqlseed.config.models import GeneratorConfig, TableConfig
+
+logger = get_logger(__name__)
+
+
+def load_config(path: str) -> GeneratorConfig:
+    config_path = Path(path)
+    if not config_path.exists():
+        raise FileNotFoundError(f"Configuration file not found: {path}")
+
+    suffix = config_path.suffix.lower()
+    with open(config_path, encoding="utf-8") as f:
+        if suffix in (".yaml", ".yml"):
+            raw = yaml.safe_load(f)
+        elif suffix == ".json":
+            raw = json.load(f)
+        else:
+            raise ValueError(f"Unsupported configuration file format: {suffix}")
+
+    if not isinstance(raw, dict):
+        raise ValueError("Configuration file must contain a YAML/JSON object")
+
+    return GeneratorConfig(**raw)
+
+
+def save_config(config: GeneratorConfig, path: str) -> None:
+    config_path = Path(path)
+    config_path.parent.mkdir(parents=True, exist_ok=True)
+
+    suffix = config_path.suffix.lower()
+    data = config.model_dump(mode="json")
+
+    with open(config_path, "w", encoding="utf-8") as f:
+        if suffix in (".yaml", ".yml"):
+            yaml.dump(data, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
+        elif suffix == ".json":
+            json.dump(data, f, indent=2, ensure_ascii=False)
+        else:
+            raise ValueError(f"Unsupported configuration file format: {suffix}")
+
+    logger.info("Configuration saved", path=path)
+
+
+def generate_template(db_path: str, table_name: str | None = None) -> GeneratorConfig:
+    tables: list[TableConfig] = []
+    if table_name:
+        tables.append(
+            TableConfig(
+                name=table_name,
+                count=1000,
+                columns=[],
+            )
+        )
+
+    return GeneratorConfig(
+        db_path=db_path,
+        tables=tables,
+    )

sqlseed/config/models.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+from typing_extensions import Self
+
+
+class ProviderType(str, Enum):
+    BASE = "base"
+    FAKER = "faker"
+    MIMESIS = "mimesis"
+    CUSTOM = "custom"
+    AI = "ai"
+
+
+class ColumnConstraintsConfig(BaseModel):
+    """列约束配置"""
+
+    unique: bool = False
+    min_value: int | float | None = None
+    max_value: int | float | None = None
+    regex: str | None = None
+    max_retries: int = Field(default=100, gt=0)
+
+
+class ColumnConfig(BaseModel):
+    """
+    列配置 — 支持源列和派生列两种模式。
+
+    源列模式：指定 generator + params
+    派生列模式：指定 derive_from + expression
+    两者不能同时使用。
+    """
+
+    name: str
+
+    # === 源列模式 ===
+    generator: str | None = None
+    provider: ProviderType | None = None
+    params: dict[str, Any] = Field(default_factory=dict)
+    null_ratio: float = Field(default=0.0, ge=0.0, le=1.0)
+
+    # === 派生列模式 ===
+    derive_from: str | None = None  # 源列名
+    expression: str | None = None  # 派生表达式
+
+    # === 约束 ===
+    constraints: ColumnConstraintsConfig | None = None
+
+    @field_validator("null_ratio")
+    @classmethod
+    def validate_null_ratio(cls, v: float) -> float:
+        if not 0.0 <= v <= 1.0:
+            raise ValueError("null_ratio must be between 0.0 and 1.0")
+        return v
+
+    @model_validator(mode="after")
+    def validate_column_mode(self) -> Self:
+        if self.derive_from and self.generator:
+            raise ValueError(f"Column '{self.name}': cannot use both 'generator' and 'derive_from'")
+        if self.derive_from and not self.expression:
+            raise ValueError(f"Column '{self.name}': 'derive_from' requires 'expression'")
+        return self
+
+
+class TableConfig(BaseModel):
+    """单表生成配置"""
+
+    name: str
+    count: int = Field(default=1000, gt=0)
+    batch_size: int = Field(default=5000, gt=0)
+    columns: list[ColumnConfig] = Field(default_factory=list)
+    clear_before: bool = False  # 默认不清空，保护原始数据
+    seed: int | None = None
+    transform: str | None = None  # [NEW] Python 变换脚本路径
+
+
+class ColumnAssociation(BaseModel):
+    """跨表列关联声明 — 用于隐式关联（同名列跨表引用）"""
+
+    column_name: str
+    source_table: str
+    target_tables: list[str] = Field(default_factory=list)
+    strategy: str = "shared_pool"
+
+
+class GeneratorConfig(BaseModel):
+    """全局生成配置"""
+
+    db_path: str
+    provider: ProviderType = ProviderType.MIMESIS
+    locale: str = "en_US"
+    tables: list[TableConfig] = Field(default_factory=list)
+    associations: list[ColumnAssociation] = Field(default_factory=list)
+    optimize_pragma: bool = True
+    log_level: str = "INFO"
+    snapshot_dir: str | None = None

sqlseed/config/snapshot.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from sqlseed._utils.logger import get_logger
+from sqlseed.config.models import GeneratorConfig
+
+logger = get_logger(__name__)
+
+
+class SnapshotManager:
+    def __init__(self, snapshot_dir: str | None = None) -> None:
+        self._snapshot_dir = Path(snapshot_dir) if snapshot_dir else Path("./snapshots")
+
+    def save(
+        self,
+        config: GeneratorConfig,
+        table_name: str,
+        count: int,
+        seed: int | None = None,
+    ) -> str:
+        self._snapshot_dir.mkdir(parents=True, exist_ok=True)
+
+        timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
+        filename = f"{timestamp}_{table_name}.yaml"
+        filepath = self._snapshot_dir / filename
+
+        snapshot_data = {
+            "timestamp": timestamp,
+            "table_name": table_name,
+            "count": count,
+            "seed": seed,
+            "config": config.model_dump(mode="json"),
+        }
+
+        with open(filepath, "w", encoding="utf-8") as f:
+            yaml.dump(snapshot_data, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
+
+        logger.info("Snapshot saved", filepath=str(filepath))
+        return str(filepath)
+
+    def load(self, snapshot_path: str) -> dict[str, Any]:
+        path = Path(snapshot_path)
+        if not path.exists():
+            raise FileNotFoundError(f"Snapshot not found: {snapshot_path}")
+
+        with open(path, encoding="utf-8") as f:
+            data: dict[str, Any] = yaml.safe_load(f)
+
+        return data
+
+    def replay(self, snapshot_path: str) -> Any:
+        from sqlseed.core.orchestrator import DataOrchestrator
+
+        data = self.load(snapshot_path)
+        config_data = data["config"]
+        config = GeneratorConfig(**config_data)
+
+        table_name = data["table_name"]
+        count = data["count"]
+        seed = data.get("seed")
+
+        table_config = None
+        for tc in config.tables:
+            if tc.name == table_name:
+                table_config = tc
+                break
+
+        with DataOrchestrator(
+            db_path=config.db_path,
+            provider_name=config.provider.value,
+            locale=config.locale,
+            optimize_pragma=config.optimize_pragma,
+        ) as orch:
+            return orch.fill_table(
+                table_name=table_name,
+                count=count,
+                seed=seed,
+                batch_size=table_config.batch_size if table_config else 5000,
+                clear_before=table_config.clear_before if table_config else False,
+                column_configs=table_config.columns if table_config else None,
+            )
+
+    def list_snapshots(self) -> list[str]:
+        if not self._snapshot_dir.exists():
+            return []
+        return sorted(str(p) for p in self._snapshot_dir.glob("*.yaml"))

sqlseed/core/__init__.py

@@ -0,0 +1,14 @@
+from sqlseed.core.mapper import ColumnMapper, GeneratorSpec
+from sqlseed.core.orchestrator import DataOrchestrator
+from sqlseed.core.relation import RelationResolver
+from sqlseed.core.result import GenerationResult
+from sqlseed.core.schema import SchemaInferrer
+
+__all__ = [
+    "ColumnMapper",
+    "DataOrchestrator",
+    "GenerationResult",
+    "GeneratorSpec",
+    "RelationResolver",
+    "SchemaInferrer",
+]

sqlseed/core/column_dag.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from sqlseed.core.mapper import GeneratorSpec
+
+
+@dataclass
+class ColumnConstraints:
+    """列级约束"""
+
+    unique: bool = False
+    min_value: int | float | None = None
+    max_value: int | float | None = None
+    regex: str | None = None
+    max_retries: int = 100
+
+
+@dataclass
+class ColumnNode:
+    """DAG 中的一个节点，代表一个列"""
+
+    name: str
+    generator_spec: GeneratorSpec
+    depends_on: list[str] = field(default_factory=list)  # 依赖的源列名
+    expression: str | None = None  # 派生表达式
+    constraints: ColumnConstraints | None = None  # 约束条件
+    is_derived: bool = False  # 是否为派生列
+
+    @property
+    def is_skip(self) -> bool:
+        return self.generator_spec.generator_name == "skip"
+
+
+class ColumnDAG:
+    """构建并管理列依赖图"""
+
+    def build(
+        self,
+        specs: dict[str, GeneratorSpec],
+        column_configs: list[Any] | None = None,
+    ) -> list[ColumnNode]:
+        nodes: dict[str, ColumnNode] = {}
+        config_map: dict[str, Any] = {}
+
+        if column_configs:
+            for cc in column_configs:
+                if hasattr(cc, "name"):
+                    config_map[cc.name] = cc
+
+        for col_name, spec in specs.items():
+            cc = config_map.get(col_name)
+            constraints = None
+            expression = None
+            depends_on = []
+            is_derived = False
+            final_spec = spec
+
+            if cc:
+                if hasattr(cc, "constraints") and cc.constraints:
+                    constraints = ColumnConstraints(
+                        unique=cc.constraints.unique,
+                        max_retries=cc.constraints.max_retries,
+                    )
+                if hasattr(cc, "derive_from") and cc.derive_from:
+                    depends_on = [cc.derive_from]
+                    expression = cc.expression
+                    is_derived = True
+                    final_spec = GeneratorSpec(generator_name="__derive__")
+
+            nodes[col_name] = ColumnNode(
+                name=col_name,
+                generator_spec=final_spec,
+                depends_on=depends_on,
+                expression=expression,
+                constraints=constraints,
+                is_derived=is_derived,
+            )
+
+        return self._topological_sort(nodes)
+
+    def _topological_sort(self, nodes: dict[str, ColumnNode]) -> list[ColumnNode]:
+        """Kahn 算法拓扑排序"""
+        in_degree: dict[str, int] = {name: 0 for name in nodes}
+        adjacency: dict[str, list[str]] = {name: [] for name in nodes}
+
+        for name, node in nodes.items():
+            for dep in node.depends_on:
+                if dep in adjacency:
+                    adjacency[dep].append(name)
+                    in_degree[name] += 1
+
+        queue = [name for name, deg in in_degree.items() if deg == 0]
+        result: list[ColumnNode] = []
+
+        while queue:
+            current = queue.pop(0)
+            result.append(nodes[current])
+            for neighbor in adjacency.get(current, []):
+                in_degree[neighbor] -= 1
+                if in_degree[neighbor] == 0:
+                    queue.append(neighbor)
+
+        if len(result) != len(nodes):
+            raise ValueError("Circular dependency detected in column definitions")
+
+        return result

sqlseed/core/constraints.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class RegisterResult:
+    registered: bool = True
+    need_backtrack: bool = False
+    backtrack_targets: list[str] = field(default_factory=list)
+
+
+class ConstraintSolver:
+    """约束求解器，支持回溯和复合唯一约束
+
+    For large datasets (>100K rows), set probabilistic=True to use
+    a hash-based probabilistic set that trades a small false-positive
+    rate for significantly reduced memory usage.
+    """
+
+    def __init__(
+        self,
+        *,
+        probabilistic: bool = False,
+        expected_count: int = 10000,
+    ) -> None:
+        self._probabilistic = probabilistic
+        self._expected_count = expected_count
+        self._seen: dict[str, set[Any]] = {}
+        self._composite_seen: dict[str, set[tuple[Any, ...]]] = {}
+        if probabilistic:
+            self._hash_seen: dict[str, set[int]] = {}
+
+    def _is_seen(self, column_name: str, value: Any) -> bool:
+        if self._probabilistic:
+            h = hash(value)
+            if column_name not in self._hash_seen:
+                self._hash_seen[column_name] = set()
+            if h in self._hash_seen[column_name]:
+                return True
+            self._hash_seen[column_name].add(h)
+            return False
+        if column_name not in self._seen:
+            self._seen[column_name] = set()
+        if value in self._seen[column_name]:
+            return True
+        self._seen[column_name].add(value)
+        return False
+
+    def _unregister_value(self, column_name: str, value: Any) -> None:
+        if self._probabilistic:
+            if column_name in self._hash_seen:
+                self._hash_seen[column_name].discard(hash(value))
+        elif column_name in self._seen:
+            self._seen[column_name].discard(value)
+
+    def check_and_register(
+        self,
+        column_name: str,
+        value: Any,
+        unique: bool = False,
+    ) -> bool:
+        if unique:
+            return not self._is_seen(column_name, value)
+        return True
+
+    def try_register(
+        self,
+        column_name: str,
+        value: Any,
+        unique: bool = False,
+        source_columns: list[str] | None = None,
+    ) -> RegisterResult:
+        if not unique:
+            return RegisterResult(registered=True)
+
+        if self._is_seen(column_name, value):
+            return RegisterResult(
+                registered=False,
+                need_backtrack=True,
+                backtrack_targets=source_columns if source_columns else [column_name],
+            )
+        return RegisterResult(registered=True)
+
+    def check_composite(
+        self,
+        key_name: str,
+        values: tuple[Any, ...],
+    ) -> bool:
+        if key_name not in self._composite_seen:
+            self._composite_seen[key_name] = set()
+        if values in self._composite_seen[key_name]:
+            return False
+        self._composite_seen[key_name].add(values)
+        return True
+
+    def unregister_composite(
+        self,
+        key_name: str,
+        values: tuple[Any, ...],
+    ) -> None:
+        if key_name in self._composite_seen:
+            self._composite_seen[key_name].discard(values)
+
+    def reset(self) -> None:
+        self._seen.clear()
+        self._composite_seen.clear()
+        if self._probabilistic:
+            self._hash_seen.clear()
+
+    def reset_column(self, column_name: str) -> None:
+        self._seen.pop(column_name, None)
+
+    def unregister(self, column_name: str, value: Any) -> None:
+        self._unregister_value(column_name, value)

sqlseed/core/expression.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import threading
+from typing import Any, ClassVar
+
+import simpleeval
+
+from sqlseed._utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class ExpressionTimeoutError(TimeoutError):
+    pass
+
+
+class ExpressionEngine:
+    """安全表达式求值器"""
+
+    SAFE_FUNCTIONS: ClassVar[dict[str, Any]] = {
+        "len": len,
+        "int": int,
+        "str": str,
+        "float": float,
+        "hex": hex,
+        "oct": oct,
+        "bin": bin,
+        "abs": abs,
+        "min": min,
+        "max": max,
+        "upper": lambda s: s.upper(),
+        "lower": lambda s: s.lower(),
+        "strip": lambda s: s.strip(),
+        "lstrip": lambda s: s.lstrip(),
+        "rstrip": lambda s: s.rstrip(),
+        "zfill": lambda s, w: str(s).zfill(w),
+        "replace": lambda s, old, new: str(s).replace(old, new),
+        "substr": lambda s, start, end=None: str(s)[start:end],
+        "lpad": lambda s, width, char="0": str(s).rjust(width, char),
+        "rpad": lambda s, width, char="0": str(s).ljust(width, char),
+        "concat": lambda *args: "".join(str(a) for a in args),
+    }
+
+    def __init__(self, timeout_seconds: int = 5) -> None:
+        self._timeout = timeout_seconds
+        self._evaluator = simpleeval.SimpleEval()
+        self._evaluator.functions = dict(self.SAFE_FUNCTIONS)
+
+    def evaluate(self, expression: str, context: dict[str, Any]) -> Any:
+        self._evaluator.names = context
+        result: Any = None
+        error: Exception | None = None
+
+        def _eval() -> None:
+            nonlocal result, error
+            try:
+                result = self._evaluator.eval(expression)
+            except Exception as e:
+                error = e
+
+        thread = threading.Thread(target=_eval)
+        thread.start()
+        thread.join(timeout=self._timeout)
+
+        if thread.is_alive():
+            raise ExpressionTimeoutError(f"Expression evaluation timed out after {self._timeout}s: {expression[:100]}")
+
+        if error is not None:
+            raise error
+
+        return result