ns-orm 0.0.0__py3-none-any.whl

ns_orm/migrations/loader.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import re
+from importlib.machinery import SourceFileLoader
+from importlib.util import module_from_spec, spec_from_loader
+from pathlib import Path
+from typing import Any, Optional
+
+from ns_orm.migrations.migration import Migration
+
+_MIGRATION_RE = re.compile(r"^\d{14}_.+\.py$")
+
+
+class MigrationLoader:
+    def __init__(self, *, migrations_dir: Path):
+        self.migrations_dir = migrations_dir
+
+    def list_files(self) -> list[Path]:
+        if not self.migrations_dir.exists():
+            return []
+        files = [
+            p
+            for p in self.migrations_dir.iterdir()
+            if p.is_file() and _MIGRATION_RE.match(p.name)
+        ]
+        return sorted(files, key=lambda p: p.name)
+
+    def load_all(self) -> list[Migration]:
+        migrations: list[Migration] = []
+        for p in self.list_files():
+            migrations.append(self.load_file(p))
+        return migrations
+
+    def load_file(self, path: Path) -> Migration:
+        loader = SourceFileLoader(f"ns_orm_user_migration_{path.stem}", str(path))
+        spec = spec_from_loader(loader.name, loader)
+        if spec is None:
+            raise RuntimeError(f"Failed to load migration: {path}")
+        mod = module_from_spec(spec)
+        assert spec.loader is not None
+        spec.loader.exec_module(mod)
+        mig = getattr(mod, "migration", None)
+        if not isinstance(mig, Migration):
+            raise RuntimeError(
+                f"Migration file must define `migration: Migration`: {path}"
+            )
+        return mig
+
+    def latest_state(self) -> Optional[dict[str, Any]]:
+        migrations = self.load_all()
+        if not migrations:
+            return None
+        return migrations[-1].state

ns_orm/migrations/migration.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from ns_orm.migrations.operations import Operation
+
+
+@dataclass(frozen=True)
+class Migration:
+    name: str
+    operations: list[Operation]
+    state: dict[str, Any]
+    dependencies: list[str] = ()

ns_orm/migrations/operations.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from ns_orm.dialects import Dialect
+from ns_orm.exceptions import QueryError
+
+
+def _fielddef_from_value(value: Any) -> Any:
+    if isinstance(value, dict) and "__class__" in value:
+        cls_name = value.get("__class__", "")
+        kwargs = {k: v for k, v in value.items() if k != "__class__"}
+        from ns_orm import fields as f
+
+        cls = getattr(f, cls_name, None)
+        if cls is None:
+            return f.Int()
+        try:
+            return cls(**kwargs)
+        except Exception:
+            return cls()
+    return value
+
+
+@dataclass(frozen=True)
+class Operation:
+    def sql(self, dialect: Dialect) -> list[str]:
+        raise NotImplementedError
+
+
+@dataclass(frozen=True)
+class CreateTable(Operation):
+    table_name: str
+    columns: list[tuple[str, Any]]
+    pk_name: str
+    fks: list[tuple[str, str, str, str]]
+    unique_together: list[tuple[str, ...]] = ()
+
+    def sql(self, dialect: Dialect) -> list[str]:
+        cols = [(name, _fielddef_from_value(fdef)) for name, fdef in self.columns]
+        ddl = dialect.ddl_create_table(
+            table_name=self.table_name,
+            columns=cols,
+            pk_name=self.pk_name,
+            fks=self.fks,
+            uniques=[],
+            unique_together=self.unique_together,
+        )
+        return [ddl]
+
+
+@dataclass(frozen=True)
+class AddColumn(Operation):
+    table_name: str
+    column_name: str
+    column_def: Any
+
+    def sql(self, dialect: Dialect) -> list[str]:
+        col_def = _fielddef_from_value(self.column_def)
+        col = dialect.quote_ident(self.column_name)
+        ty = dialect.type_sql(col_def)
+        parts = [
+            f"ALTER TABLE {dialect.quote_ident(self.table_name)} ADD COLUMN {col} {ty}"
+        ]
+        nullable = getattr(col_def, "nullable", None)
+        if nullable is False:
+            parts[0] += " NOT NULL"
+        if getattr(col_def, "unique", False):
+            parts[0] += " UNIQUE"
+        return parts
+
+
+@dataclass(frozen=True)
+class DropColumn(Operation):
+    table_name: str
+    column_name: str
+
+    def sql(self, dialect: Dialect) -> list[str]:
+        if dialect.name == "sqlite":
+            raise QueryError("SQLite 不支持直接 DROP COLUMN，请手工重建表或扩展迁移器")
+        return [
+            f"ALTER TABLE {dialect.quote_ident(self.table_name)} "
+            f"DROP COLUMN {dialect.quote_ident(self.column_name)}"
+        ]
+
+
+@dataclass(frozen=True)
+class DropTable(Operation):
+    table_name: str
+
+    def sql(self, dialect: Dialect) -> list[str]:
+        return [f"DROP TABLE IF EXISTS {dialect.quote_ident(self.table_name)}"]

ns_orm/migrations/state.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from dataclasses import asdict, is_dataclass
+from typing import Any
+
+from ns_orm.fields import ForeignKey, ManyToMany
+
+
+def field_to_dict(field: Any) -> dict[str, Any]:
+    if is_dataclass(field):
+        data = asdict(field)
+    else:
+        data = dict(getattr(field, "__dict__", {}))
+    data["__class__"] = field.__class__.__name__
+    return data
+
+
+def model_state(model: type[Any]) -> dict[str, Any]:
+    fields: dict[str, Any] = {}
+    fks: dict[str, Any] = {}
+    m2m: dict[str, Any] = {}
+
+    for name, fdef in model._meta.fields.items():
+        if isinstance(fdef, ManyToMany):
+            m2m[name] = field_to_dict(fdef)
+            continue
+        if isinstance(fdef, ForeignKey):
+            fks[name] = field_to_dict(fdef)
+        fields[name] = field_to_dict(fdef)
+
+    return {
+        "table_name": model.table_name(),
+        "schema": getattr(model._meta, "schema", None),
+        "pk_name": model.pk_name(),
+        "fields": fields,
+        "fks": fks,
+        "m2m": m2m,
+    }
+
+
+def project_state(models: list[type[Any]]) -> dict[str, Any]:
+    return {m.__name__: model_state(m) for m in models}

ns_orm/migrations/writer.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from ns_orm.migrations.migration import Migration
+from ns_orm.migrations.operations import (
+    AddColumn,
+    CreateTable,
+    DropColumn,
+    DropTable,
+    Operation,
+)
+
+
+class MigrationWriter:
+    def __init__(self, *, migrations_dir: Path):
+        self.migrations_dir = migrations_dir
+
+    def write(self, migration: Migration) -> Path:
+        self.migrations_dir.mkdir(parents=True, exist_ok=True)
+        path = self.migrations_dir / f"{migration.name}.py"
+        content = self._render(migration)
+        path.write_text(content, encoding="utf-8")
+        return path
+
+    def _render(self, migration: Migration) -> str:
+        ops_src = ",\n        ".join(self._render_op(op) for op in migration.operations)
+        state_json = json.dumps(
+            migration.state, ensure_ascii=False, indent=2, sort_keys=True
+        )
+        return (
+            "from __future__ import annotations\n\n"
+            "from ns_orm.migrations.migration import Migration\n"
+            "from ns_orm.migrations.operations import (\n"
+            "    AddColumn,\n"
+            "    CreateTable,\n"
+            "    DropColumn,\n"
+            "    DropTable,\n"
+            ")\n\n"
+            f"migration = Migration(\n"
+            f"    name={migration.name!r},\n"
+            f"    dependencies={list(migration.dependencies)!r},\n"
+            f"    operations=[\n"
+            f"        {ops_src}\n"
+            f"    ],\n"
+            f"    state={state_json},\n"
+            f")\n"
+        )
+
+    def _render_op(self, op: Operation) -> str:
+        if isinstance(op, CreateTable):
+            return (
+                "CreateTable(\n"
+                f"            table_name={op.table_name!r},\n"
+                f"            columns={op.columns!r},\n"
+                f"            pk_name={op.pk_name!r},\n"
+                f"            fks={op.fks!r},\n"
+                f"            unique_together={list(op.unique_together)!r},\n"
+                "        )"
+            )
+        if isinstance(op, AddColumn):
+            return (
+                "AddColumn(\n"
+                f"            table_name={op.table_name!r},\n"
+                f"            column_name={op.column_name!r},\n"
+                f"            column_def={op.column_def!r},\n"
+                "        )"
+            )
+        if isinstance(op, DropColumn):
+            return (
+                "DropColumn(\n"
+                f"            table_name={op.table_name!r},\n"
+                f"            column_name={op.column_name!r},\n"
+                "        )"
+            )
+        if isinstance(op, DropTable):
+            return f"DropTable(\n            table_name={op.table_name!r},\n        )"
+        raise TypeError(f"Unsupported operation: {op}")

ns_orm/model.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import re
+from typing import Any, ClassVar, TypeVar
+
+from pydantic import BaseModel
+
+from ns_orm.exceptions import ModelDefinitionError
+from ns_orm.fields import FieldDef, ForeignKey, ManyToMany
+from ns_orm.typing import ModelMetaInfo
+from ns_orm.utils import (
+    get_class_namespace,
+    get_type_hints_with_extras,
+    is_classvar,
+    resolve_ref,
+    unwrap_annotated,
+)
+
+TModel = TypeVar("TModel", bound="Model")
+
+
+def _default_table_name(cls_name: str) -> str:
+    s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", cls_name)
+    return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower()
+
+
+class Model(BaseModel):
+    class Config:
+        extra = "allow"
+        validate_assignment = True
+
+    _meta: ClassVar[ModelMetaInfo]
+    _registry: ClassVar[dict[str, type[Model]]] = {}
+    objects: ClassVar[Any]
+
+    @classmethod
+    def table_name(cls) -> str:
+        return cls._meta.table_name
+
+    @classmethod
+    def pk_name(cls) -> str:
+        return cls._meta.pk_name
+
+    @classmethod
+    def _resolve_model(cls, ref: str) -> type[Model]:
+        if ref in cls._registry:
+            return cls._registry[ref]
+        simple = ref.rsplit(".", 1)[-1]
+        if simple in cls._registry:
+            return cls._registry[simple]
+        raise ModelDefinitionError(f"Unresolved model reference: {ref}")
+
+    def pk_value(self) -> Any:
+        return getattr(self, self.pk_name())
+
+    def to_db_dict(self) -> dict[str, Any]:
+        data = self.dict()
+        cols = {k: v for k, v in data.items() if k in self._meta.fields}
+        return cols
+
+    def save(self, db: Any) -> Any:
+        from ns_orm.query import QuerySet
+
+        qs = QuerySet(self.__class__, db=db)
+        return qs.save_instance(self)
+
+    async def asave(self, db: Any) -> Any:
+        from ns_orm.query import AsyncQuerySet
+
+        qs = AsyncQuerySet(self.__class__, db=db)
+        return await qs.save_instance(self)
+
+    def delete(self, db: Any) -> int:
+        from ns_orm.query import QuerySet
+
+        qs = QuerySet(self.__class__, db=db)
+        return qs.filter(**{self.pk_name(): self.pk_value()}).delete()
+
+    async def adelete(self, db: Any) -> int:
+        from ns_orm.query import AsyncQuerySet
+
+        qs = AsyncQuerySet(self.__class__, db=db)
+        return await qs.filter(**{self.pk_name(): self.pk_value()}).delete()
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+        if cls.__name__ == "Model":
+            return
+
+        hints = get_type_hints_with_extras(cls)
+        namespace = get_class_namespace(cls)
+
+        fields: dict[str, FieldDef] = {}
+        relations: dict[str, tuple[str, ForeignKey]] = {}
+        m2m: dict[str, ManyToMany] = {}
+
+        for name, annotation in hints.items():
+            if is_classvar(annotation):
+                continue
+            base, meta = unwrap_annotated(annotation)
+            resolved_base = resolve_ref(base, namespace)
+            _ = resolved_base
+            for m in meta:
+                if isinstance(m, ForeignKey):
+                    fields[name] = m
+                    rel_name = name[:-3] if name.endswith("_id") else f"{name}_obj"
+                    relations[rel_name] = (name, m)
+                elif isinstance(m, FieldDef):
+                    fields[name] = m
+                elif isinstance(m, ManyToMany):
+                    m2m[name] = m
+
+        if not fields:
+            raise ModelDefinitionError(
+                f"{cls.__name__} has no ORM fields. "
+                "Use typing.Annotated[..., FieldDef()]."
+            )
+
+        table_name = getattr(
+            getattr(cls, "Meta", object), "table_name", None
+        ) or _default_table_name(cls.__name__)
+        schema = getattr(getattr(cls, "Meta", object), "schema", None)
+        connect_name = getattr(getattr(cls, "Meta", object), "connect_name", None)
+
+        pk_candidates = [
+            n for n, f in fields.items() if getattr(f, "primary_key", False)
+        ]
+        if pk_candidates:
+            pk_name = pk_candidates[0]
+        elif "id" in fields:
+            pk_name = "id"
+        else:
+            raise ModelDefinitionError(f"{cls.__name__} missing primary key field")
+
+        cls._meta = ModelMetaInfo(
+            table_name=table_name,
+            schema=schema,
+            connect_name=connect_name,
+            fields=fields,
+            pk_name=pk_name,
+            table=None,
+            metadata=None,
+            relations=relations,
+            m2m=m2m,
+        )
+
+        cls._registry[cls.__name__] = cls
+        cls._registry[f"{cls.__module__}.{cls.__name__}"] = cls
+        from ns_orm.manager import Manager
+
+        cls.objects = Manager(model=cls)