stackthread 0.1.0.dev2__py3-none-any.whl

stackthread/lang/python/orm.py

@@ -0,0 +1,300 @@
+"""Data-layer detectors: Django ORM, SQLAlchemy and SQLModel models, MongoDB collections (design §8.4).
+
+Model classes become `model` nodes in the data layer. Functions that query or change them get `reads` and
+`writes` edges. MongoDB collections named by literals become `store` nodes (heuristic confidence).
+"""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING, Any
+
+from stackthread.lang.python.resolve import bare, split_dotted
+from stackthread.model import Confidence, Edge, EdgeKind, Node
+
+if TYPE_CHECKING:
+    from stackthread.lang.python.build import PyBuild
+    from stackthread.lang.python.resolve import Ref
+
+DJANGO_MODEL = ("django.db.models.Model", "django.db.models.base.Model")
+DJANGO_MANAGERS = ("objects", "_default_manager", "_base_manager")
+DJANGO_WRITES = frozenset(
+    {
+        "create",
+        "update",
+        "delete",
+        "bulk_create",
+        "bulk_update",
+        "get_or_create",
+        "update_or_create",
+        "acreate",
+        "aupdate",
+        "adelete",
+        "abulk_create",
+        "abulk_update",
+        "aget_or_create",
+        "aupdate_or_create",
+    }
+)
+SQLA_BASES = (
+    "sqlalchemy.orm.DeclarativeBase",
+    "sqlalchemy.orm.decl_api.DeclarativeBase",
+    "sqlalchemy.ext.declarative.declarative_base",
+    "sqlalchemy.orm.declarative_base",
+    "sqlalchemy.orm.DeclarativeBaseNoMeta",
+    "sqlmodel.SQLModel",
+    "sqlmodel.main.SQLModel",
+)
+SQLA_STATEMENTS: dict[str, EdgeKind] = {"select": "reads", "insert": "writes", "update": "writes", "delete": "writes"}
+MONGO_OPS: dict[str, EdgeKind] = {
+    "find": "reads",
+    "find_one": "reads",
+    "aggregate": "reads",
+    "count_documents": "reads",
+    "distinct": "reads",
+    "estimated_document_count": "reads",
+    "insert_one": "writes",
+    "insert_many": "writes",
+    "update_one": "writes",
+    "update_many": "writes",
+    "replace_one": "writes",
+    "delete_one": "writes",
+    "delete_many": "writes",
+    "find_one_and_update": "writes",
+    "find_one_and_replace": "writes",
+    "find_one_and_delete": "writes",
+    "bulk_write": "writes",
+}
+
+
+class OrmDetector:
+    def __init__(self, b: PyBuild) -> None:
+        self.b = b
+        self.idx = b.idx
+        self.models: dict[str, str] = {}  # class key -> orm name
+        self.sqla_bases: set[str] = set()
+
+    def run(self) -> None:
+        self._find_models()
+        for key, d in self.idx.defs.items():
+            for call in d.get("calls", []):
+                self._call(key, d, call)
+            for assignment in d.get("sets", []):
+                self._set(key, d, assignment)
+
+    def _set(self, key: str, d: dict[str, Any], assignment: dict[str, Any]) -> None:
+        """`task.done = True` where `task` is a model instance: a write to that model."""
+        receiver = assignment["f"].split(".")[0]
+        hint = self.idx.local_hint(d, receiver)
+        ref = self.idx.instance_type(d["module"], hint) if hint else None
+        if ref is not None and ref.kind == "def" and ref.key in self.models:
+            self._edge(key, d, ref.key, "writes", assignment["line"])
+
+    # ---- models ----
+
+    def _find_models(self) -> None:
+        idx = self.idx
+        for key, var in idx.vars.items():
+            value = var.get("value")
+            if value and value.get("t") == "call":
+                ref = idx.resolve_dotted(var["module"], value["f"])
+                if ref is not None and ref.kind == "external" and ref.key.endswith("declarative_base"):
+                    self.sqla_bases.add(key)
+        changed = True
+        while changed:
+            changed = False
+            for key, d in idx.defs.items():
+                if d["kind"] != "class" or key in self.models:
+                    continue
+                orm = self._model_kind(key, d)
+                if orm:
+                    self.models[key] = orm
+                    changed = True
+        for key, orm in self.models.items():
+            node = self.b.graph.nodes.get(self.b.node_of(key))
+            if node is None:
+                continue
+            d = idx.defs[key]
+            node.kind = "model"
+            node.layer = "data"
+            node.tags = [*node.tags, orm]
+            node.attrs["orm"] = orm
+            table = self._table(key, d, orm)
+            if table:
+                node.attrs["table"] = table
+
+    def _model_kind(self, key: str, d: dict[str, Any]) -> str | None:
+        for ref in self.idx.bases(key):
+            if ref.kind == "external":
+                if ref.key in DJANGO_MODEL or (ref.key.endswith("models.Model") and ref.key.startswith("django")):
+                    return "django"
+                if ref.key in SQLA_BASES:
+                    if ref.key.startswith("sqlmodel"):
+                        return "sqlalchemy" if self._sqlmodel_table(d) else None
+                    # `class Base(DeclarativeBase)` is the declarative base, not a table
+                    if self._is_sqla_model(d):
+                        return "sqlalchemy"
+                    self.sqla_bases.add(key)
+                    return None
+            elif ref.kind == "var" and ref.key in self.sqla_bases:
+                return "sqlalchemy"
+            elif ref.kind == "def":
+                if ref.key in self.models:
+                    return self.models[ref.key]
+                if ref.key in self.sqla_bases or self._declarative(ref.key):
+                    return "sqlalchemy"
+        return None
+
+    def _declarative(self, key: str) -> bool:
+        """A project class that is itself a declarative base (subclasses DeclarativeBase without a table)."""
+        d = self.idx.defs.get(key, {})
+        return not self._is_sqla_model(d) and any(
+            r.kind == "external" and r.key in SQLA_BASES and not r.key.startswith("sqlmodel")
+            for r in self.idx.bases(key)
+        )
+
+    def _sqlmodel_table(self, d: dict[str, Any]) -> bool:
+        table = d.get("kw", {}).get("table")
+        return bool(table and table.get("t") == "lit" and table.get("v") is True)
+
+    def _is_sqla_model(self, d: dict[str, Any]) -> bool:
+        return "__tablename__" in d.get("attrs", {})
+
+    def _table(self, key: str, d: dict[str, Any], orm: str) -> str | None:
+        if orm == "sqlalchemy":
+            value = self.idx.const(d["module"], d.get("attrs", {}).get("__tablename__"))
+            return value if isinstance(value, str) else None
+        meta = self.idx.children.get(key, {}).get("Meta")
+        if meta:
+            value = self.idx.const(d["module"], self.idx.defs[meta].get("attrs", {}).get("db_table"))
+            if isinstance(value, str):
+                return value
+        parts = d["module"].split(".")
+        app = parts[-2] if len(parts) >= 2 and parts[-1] == "models" else parts[0]
+        return f"{app}_{d['name'].lower()}"
+
+    def _model_ref(self, module: str, name: str, scope: dict[str, Any] | None = None) -> str | None:
+        ref: Ref | None = self.idx.resolve_dotted(module, name, scope)
+        if ref is not None and ref.kind == "def" and ref.key in self.models:
+            return ref.key
+        return None
+
+    # ---- calls ----
+
+    def _edge(
+        self,
+        scope_key: str,
+        d: dict[str, Any],
+        model_key: str,
+        kind: EdgeKind,
+        line: int,
+        confidence: Confidence = "exact",
+    ) -> None:
+        self.b.graph.add_edge(
+            Edge(
+                self.b.node_of(scope_key),
+                self.b.node_of(model_key),
+                kind,
+                confidence,
+                site=self.b.site(d["module"], line),
+            )
+        )
+
+    def _call(self, key: str, d: dict[str, Any], call: dict[str, Any]) -> None:
+        f, line, module = call["f"], call["line"], d["module"]
+        segments = split_dotted(f)
+        names = [bare(s) for s in segments]
+        # Django: Model.objects.<op>(...).<op>(...)
+        for i, name in enumerate(names):
+            if name in DJANGO_MANAGERS and i > 0:
+                model = self._model_ref(module, ".".join(names[:i]), d)
+                if model and i + 1 < len(names):
+                    ops = names[i + 1 :]
+                    manager_kind: EdgeKind = "writes" if any(op in DJANGO_WRITES for op in ops) else "reads"
+                    self._edge(key, d, model, manager_kind, line)
+                return
+        last = names[-1] if names else ""
+        if last in ("get_object_or_404", "get_list_or_404", "aget_object_or_404") and call["args"]:
+            arg = call["args"][0]
+            if arg and arg.get("t") == "name" and (model := self._model_ref(module, arg["v"], d)):
+                self._edge(key, d, model, "reads", line)
+            return
+        # instance.save() / .delete() / .refresh_from_db() on a typed local or self attribute
+        if len(names) >= 2 and last in ("save", "delete", "asave", "adelete", "refresh_from_db", "arefresh_from_db"):
+            receiver = ".".join(segments[:-1])
+            ref = self.idx.resolve_dotted(module, receiver, d) if "(" not in receiver else None
+            if ref is None and len(names) == 2:
+                hint = self.idx.local_hint(d, names[0])
+                ref = self.idx.instance_type(module, hint) if hint else None
+            if ref is not None and ref.kind == "def" and ref.key in self.models:
+                self._edge(key, d, ref.key, "reads" if "refresh" in last else "writes", line)
+                return
+            # not a model instance: maybe a SQLAlchemy session (`db.delete(obj)`); fall through
+        # SQLAlchemy statements: select(Model), insert(Model), update(Model), delete(Model)
+        if len(names) == 1 or (len(names) == 2 and names[0] in ("sa", "sqlalchemy", "sqlmodel")):
+            statement = SQLA_STATEMENTS.get(last)
+            if statement and call["args"]:
+                ref = self.idx.resolve_dotted(module, f)
+                if ref is not None and ref.kind == "external" and ref.key.split(".")[0] in ("sqlalchemy", "sqlmodel"):
+                    for arg in call["args"]:
+                        if arg and arg.get("t") == "name" and (model := self._model_ref(module, arg["v"], d)):
+                            self._edge(key, d, model, statement, line)
+                return
+        # SQLAlchemy session API: session.query(M), session.get(M, id), session.add(M(...)), session.delete(obj)
+        if len(names) >= 2 and last in ("query", "get", "add", "merge", "delete", "add_all") and call["args"]:
+            arg = call["args"][0]
+            session_kind: EdgeKind = "reads" if last in ("query", "get") else "writes"
+            if arg and arg.get("t") == "name":
+                model = self._model_ref(module, arg["v"], d)
+                if model is None and last in ("add", "merge", "delete"):
+                    hint = self.idx.local_hint(d, arg["v"])
+                    ref = self.idx.instance_type(module, hint) if hint else None
+                    model = ref.key if ref is not None and ref.kind == "def" and ref.key in self.models else None
+                if model:
+                    self._edge(key, d, model, session_kind, line)
+            elif arg and arg.get("t") == "call" and (model := self._model_ref(module, arg["f"], d)):
+                self._edge(key, d, model, session_kind, line)
+            return
+        # Flask-SQLAlchemy style: Model.query.<op>
+        if "query" in names[1:] and (model := self._model_ref(module, ".".join(names[: names.index("query", 1)]), d)):
+            self._edge(key, d, model, "reads", line)
+            return
+        # MongoDB: <db>.<collection>.<op>(...) or <db>[<collection>].<op>(...)
+        if last in MONGO_OPS and len(segments) >= 3 and "(" not in segments[-2]:
+            coll_segment = segments[-2]
+            sub = re.fullmatch(r"\w+\[(\w+)\]", coll_segment)
+            collection = sub.group(1) if sub else bare(coll_segment)
+            if collection and collection not in ("self", "cls") and not collection.startswith("_"):
+                self._store(key, d, collection, MONGO_OPS[last], line)
+
+    def _store(self, key: str, d: dict[str, Any], collection: str, kind: EdgeKind, line: int) -> None:
+        g, src = self.b.graph, self.b.source.name
+        node_id = f"store:{src}:mongo:{collection}"
+        g.add_node(
+            Node(
+                id=node_id,
+                kind="store",
+                name=collection,
+                lang="python",
+                source=src,
+                layer="data",
+                tags=["mongo"],
+                attrs={"engine": "mongodb", "collection": collection},
+            )
+        )
+        g.add_edge(
+            Edge(
+                self.b.node_of(key),
+                node_id,
+                kind,
+                "heuristic",
+                site=self.b.site(d["module"], line),
+            )
+        )
+
+
+def detect(b: PyBuild) -> None:
+    enabled = set(b.source.frameworks)
+    if not ({"auto", "django-orm", "django", "drf", "ninja", "sqlalchemy", "mongo"} & enabled):
+        return
+    OrmDetector(b).run()