sqlseed 0.1.0__py3-none-any.whl

sqlseed/generators/mimesis_provider.py

@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlseed._utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class MimesisProvider:
+    """Mimesis-based data generator adapter."""
+
+    def __init__(self) -> None:
+        self._generic: Any = None
+        self._locale: str = "en"
+        self._seed: int | None = None
+        self._init_mimesis()
+
+    def _init_mimesis(self) -> None:
+        try:
+            from mimesis import Generic
+            from mimesis.locales import Locale
+
+            locale_enum = Locale(self._locale)
+            self._generic = Generic(locale_enum)
+        except ImportError:
+            raise ImportError("Mimesis is not installed. Install it with: pip install sqlseed[mimesis]") from None
+
+    @property
+    def name(self) -> str:
+        return "mimesis"
+
+    def set_locale(self, locale: str) -> None:
+        locale_map = {
+            "en_US": "en",
+            "en_GB": "en",
+            "zh_CN": "zh",
+            "zh_TW": "zh",
+            "ja_JP": "ja",
+            "ko_KR": "ko",
+            "de_DE": "de",
+            "fr_FR": "fr",
+            "es_ES": "es",
+            "pt_BR": "pt-br",
+            "ru_RU": "ru",
+        }
+        self._locale = locale_map.get(locale, locale.split("_", maxsplit=1)[0])
+        self._init_mimesis()
+
+    def set_seed(self, seed: int) -> None:
+        from mimesis import Generic
+        from mimesis.locales import Locale
+
+        self._seed = seed
+        locale_enum = Locale(self._locale)
+        self._generic = Generic(locale_enum, seed=seed)
+
+    def generate_string(
+        self,
+        *,
+        min_length: int = 1,
+        max_length: int = 100,
+        charset: str | None = None,
+    ) -> str:
+        import string
+
+        if charset == "alphanumeric":
+            chars = string.ascii_letters + string.digits
+        elif charset == "alpha":
+            chars = string.ascii_letters
+        elif charset == "digits":
+            chars = string.digits
+        elif charset is not None:
+            chars = charset
+        else:
+            chars = string.ascii_letters + string.digits + " _-"
+        length = self._generic.numeric.integer_number(start=min_length, end=max_length)
+        return "".join(self._generic.random.choice(chars) for _ in range(length))
+
+    def generate_integer(self, *, min_value: int = 0, max_value: int = 999999) -> int:
+        return self._generic.numeric.integer_number(start=min_value, end=max_value)
+
+    def generate_float(
+        self,
+        *,
+        min_value: float = 0.0,
+        max_value: float = 999999.0,
+        precision: int = 2,
+    ) -> float:
+        return round(
+            self._generic.numeric.float_number(start=min_value, end=max_value, precision=precision),
+            precision,
+        )
+
+    def generate_boolean(self) -> bool:
+        return self._generic.development.boolean()
+
+    def generate_bytes(self, *, length: int = 16) -> bytes:
+        return self._generic.cryptographic.token_bytes(length)
+
+    def generate_name(self) -> str:
+        return self._generic.person.full_name()
+
+    def generate_first_name(self) -> str:
+        return self._generic.person.first_name()
+
+    def generate_last_name(self) -> str:
+        return self._generic.person.last_name()
+
+    def generate_email(self) -> str:
+        return self._generic.person.email()
+
+    def generate_phone(self) -> str:
+        return self._generic.person.phone_number()
+
+    def generate_address(self) -> str:
+        return self._generic.address.address()
+
+    def generate_company(self) -> str:
+        return self._generic.finance.company()
+
+    def generate_url(self) -> str:
+        return self._generic.internet.url()
+
+    def generate_ipv4(self) -> str:
+        return self._generic.internet.ip_v4()
+
+    def generate_uuid(self) -> str:
+        return str(self._generic.cryptographic.uuid_object())
+
+    def generate_date(self, *, start_year: int = 2000, end_year: int | None = None) -> str:
+        from datetime import datetime
+
+        if end_year is None:
+            end_year = datetime.now().year
+        date = self._generic.datetime.date(start=start_year, end=end_year)
+        return str(date)
+
+    def generate_datetime(self, *, start_year: int = 2000, end_year: int | None = None) -> str:
+        from datetime import datetime
+
+        if end_year is None:
+            end_year = datetime.now().year
+        dt = self._generic.datetime.datetime(start=start_year, end=end_year)
+        return str(dt)
+
+    def generate_timestamp(self) -> int:
+        return self._generic.datetime.timestamp()
+
+    def generate_text(self, *, min_length: int = 50, max_length: int = 200) -> str:
+        text = self._generic.text.text(quantity=1)
+        while len(text) < min_length:
+            text += " " + self._generic.text.text(quantity=1)
+        return text[:max_length]
+
+    def generate_sentence(self) -> str:
+        return self._generic.text.sentence()
+
+    def generate_password(self, *, length: int = 16) -> str:
+        return self._generic.person.password(length=length)
+
+    def generate_choice(self, choices: list[Any]) -> Any:
+        return self._generic.random.choice(choices)
+
+    def generate_json(self, *, schema: dict[str, Any] | None = None) -> str:
+        import json
+
+        if schema is None:
+            data = {
+                "id": self.generate_integer(min_value=1, max_value=999999),
+                "name": self.generate_name(),
+                "active": self.generate_boolean(),
+            }
+        else:
+            data = self._generate_from_schema(schema)
+        return json.dumps(data)
+
+    def _generate_from_schema(self, schema: dict[str, Any]) -> Any:
+        schema_type = schema.get("type", "string")
+        if schema_type == "string":
+            return self.generate_string(min_length=5, max_length=20)
+        if schema_type == "integer":
+            return self.generate_integer()
+        if schema_type == "number":
+            return self.generate_float()
+        if schema_type == "boolean":
+            return self.generate_boolean()
+        if schema_type == "array":
+            items = schema.get("items", {"type": "string"})
+            count = self._generic.numeric.integer_number(start=1, end=5)
+            return [self._generate_from_schema(items) for _ in range(count)]
+        if schema_type == "object":
+            properties = schema.get("properties", {})
+            return {k: self._generate_from_schema(v) for k, v in properties.items()}
+        return self.generate_string()
+
+    def generate_pattern(self, *, regex: str) -> str:
+        import random
+
+        import rstr
+
+        rng = random.Random(self._seed)
+        return rstr.Rstr(rng).xeger(regex)

sqlseed/generators/registry.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from sqlseed._utils.logger import get_logger
+from sqlseed.generators._protocol import DataProvider
+from sqlseed.generators.base_provider import BaseProvider
+
+logger = get_logger(__name__)
+
+
+class ProviderRegistry:
+    def __init__(self) -> None:
+        self._providers: dict[str, DataProvider] = {}
+        self._default_name: str = "base"
+        self._register_builtin()
+
+    def _register_builtin(self) -> None:
+        base = BaseProvider()
+        self._providers["base"] = base
+
+    def register(self, provider: DataProvider) -> None:
+        name = provider.name
+        self._providers[name] = provider
+        logger.debug("Registered provider", name=name)
+
+    def register_from_entry_points(self) -> None:
+        try:
+            from importlib.metadata import entry_points
+
+            eps = entry_points()
+            sqlseed_eps = eps.select(group="sqlseed") if hasattr(eps, "select") else eps.get("sqlseed", [])  # type: ignore[arg-type]
+            for ep in sqlseed_eps:
+                try:
+                    provider_cls = ep.load()
+                    provider = provider_cls()
+                    if isinstance(provider, DataProvider):
+                        self.register(provider)
+                        logger.info("Auto-discovered provider", name=ep.name)
+                except Exception as e:
+                    logger.warning("Failed to load provider", name=ep.name, error=e)
+        except Exception as e:
+            logger.debug("Entry point discovery failed", error=e)
+
+    def get(self, name: str | None = None) -> DataProvider:
+        provider_name = name or self._default_name
+        if provider_name not in self._providers:
+            available = ", ".join(self._providers.keys())
+            raise ValueError(f"Provider '{provider_name}' not found. Available: {available}")
+        return self._providers[provider_name]
+
+    def set_default(self, name: str) -> None:
+        if name not in self._providers:
+            raise ValueError(f"Provider '{name}' not found")
+        self._default_name = name
+
+    @property
+    def default_name(self) -> str:
+        return self._default_name
+
+    @property
+    def available_providers(self) -> list[str]:
+        return list(self._providers.keys())
+
+    def ensure_provider(self, name: str) -> DataProvider:
+        if name in self._providers:
+            return self._providers[name]
+
+        if name == "faker":
+            try:
+                from sqlseed.generators.faker_provider import FakerProvider
+
+                provider: DataProvider = FakerProvider()
+                self.register(provider)
+                return provider
+            except ImportError:
+                raise ImportError("Faker is not installed. Install it with: pip install sqlseed[faker]") from None
+        elif name == "mimesis":
+            try:
+                from sqlseed.generators.mimesis_provider import MimesisProvider
+
+                provider = MimesisProvider()
+                self.register(provider)
+                return provider
+            except ImportError:
+                raise ImportError("Mimesis is not installed. Install it with: pip install sqlseed[mimesis]") from None
+
+        raise ValueError(f"Unknown provider: {name}")

sqlseed/generators/stream.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+import random
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from sqlseed.core.column_dag import ColumnNode
+    from sqlseed.core.constraints import ConstraintSolver
+    from sqlseed.core.expression import ExpressionEngine
+    from sqlseed.core.mapper import GeneratorSpec
+    from sqlseed.core.transform import RowTransformFn
+
+
+class UnknownGeneratorError(Exception):
+    def __init__(self, generator_name: str, column_name: str | None = None) -> None:
+        self.generator_name = generator_name
+        self.column_name = column_name
+        super().__init__(f"Unknown generator '{generator_name}'{f' for column {column_name}' if column_name else ''}")
+
+
+class DataStream:
+    def __init__(
+        self,
+        dag_nodes: list[ColumnNode],
+        provider: Any,
+        expr_engine: ExpressionEngine,
+        constraint_solver: ConstraintSolver,
+        transform_fn: RowTransformFn | None = None,
+        seed: int | None = None,
+    ) -> None:
+        self._nodes = dag_nodes
+        self._provider = provider
+        self._expr_engine = expr_engine
+        self._constraint_solver = constraint_solver
+        self._transform_fn = transform_fn
+
+        self._rng = random.Random(seed)
+        if seed is not None:
+            self._provider.set_seed(seed)
+
+    def generate(
+        self,
+        count: int,
+        batch_size: int = 5000,
+    ) -> Iterator[list[dict[str, Any]]]:
+        generated = 0
+        while generated < count:
+            current_batch_size = min(batch_size, count - generated)
+            batch = [self._generate_row() for _ in range(current_batch_size)]
+            yield batch
+            generated += current_batch_size
+
+    def _generate_row(self) -> dict[str, Any]:
+        max_total_retries = 1000
+        total_retries = 0
+
+        while total_retries < max_total_retries:
+            row: dict[str, Any] = {}
+            generated_values: dict[str, Any] = {}
+            backtrack_to: int | None = None
+
+            for idx, node in enumerate(self._nodes):
+                if node.is_skip:
+                    continue
+
+                col_name = node.name
+                max_retries = node.constraints.max_retries if node.constraints else 100
+                is_unique = node.constraints.unique if node.constraints else False
+                source_columns = node.depends_on if node.is_derived else None
+
+                if backtrack_to is not None and idx < backtrack_to:
+                    continue
+
+                col_success = False
+                for _ in range(max_retries):
+                    if node.is_derived and node.expression:
+                        ctx = {"row": row, "value": row.get(node.depends_on[0]) if node.depends_on else None}
+                        val = self._expr_engine.evaluate(node.expression, ctx)
+                    else:
+                        val = self._apply_generator(node.generator_spec)
+
+                    result = self._constraint_solver.try_register(
+                        col_name,
+                        val,
+                        unique=is_unique,
+                        source_columns=source_columns,
+                    )
+                    if result.registered:
+                        row[col_name] = val
+                        generated_values[col_name] = val
+                        col_success = True
+                        break
+                    if result.need_backtrack and source_columns:
+                        for bt_col in source_columns:
+                            if bt_col in generated_values:
+                                self._constraint_solver.unregister(bt_col, generated_values[bt_col])
+                                del generated_values[bt_col]
+                                row.pop(bt_col, None)
+                        bt_idx = self._find_node_index(source_columns[0])
+                        if bt_idx is not None:
+                            backtrack_to = bt_idx
+                        col_success = False
+                        break
+
+                if not col_success:
+                    if backtrack_to is not None:
+                        break
+                    for col, val in generated_values.items():
+                        self._constraint_solver.unregister(col, val)
+                    generated_values.clear()
+                    row.clear()
+                    break
+
+            if backtrack_to is not None:
+                total_retries += 1
+                continue
+
+            if generated_values:
+                if self._transform_fn:
+                    ctx = {"row_number": total_retries}
+                    row = self._transform_fn(row, ctx)
+                return row
+
+            total_retries += 1
+
+        raise RuntimeError("Failed to generate row satisfying all constraints after maximum retries.")
+
+    def _find_node_index(self, col_name: str) -> int | None:
+        for i, node in enumerate(self._nodes):
+            if node.name == col_name:
+                return i
+        return None
+
+    def _apply_generator(self, spec: GeneratorSpec) -> Any:
+        if spec.null_ratio > 0 and self._rng.random() < spec.null_ratio:
+            return None
+
+        method_name = f"generate_{spec.generator_name}"
+        if hasattr(self._provider, method_name):
+            method = getattr(self._provider, method_name)
+            return method(**spec.params) if spec.params else method()
+
+        if spec.generator_name == "choice" and "choices" in spec.params:
+            return self._rng.choice(spec.params["choices"])
+
+        if spec.generator_name == "foreign_key":
+            return self._handle_foreign_key(spec)
+
+        return self._provider.generate_string(**spec.params) if spec.params else self._provider.generate_string()
+
+    def _handle_foreign_key(self, spec: GeneratorSpec) -> Any:
+        ref_values = spec.params.get("_ref_values", [])
+        if ref_values:
+            return self._rng.choice(ref_values)
+        return self._provider.generate_integer(min_value=1, max_value=spec.params.get("max_ref", 999999))