devmem-agents 0.1.0__py3-none-any.whl

devmem/storage/neptune_store.py

@@ -0,0 +1,194 @@
+"""Neptune adapter for devmem repository ingestion."""
+
+from __future__ import annotations
+
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any
+
+import boto3
+import urllib3
+from botocore import UNSIGNED
+from botocore.config import Config as BotoConfig
+
+from devmem.live_backend import LiveBackendConfig, neptune_tunnel
+
+logger = logging.getLogger(__name__)
+
+
+class NeptuneStore:
+    """Write project/file lineage facts into Neptune (openCypher)."""
+
+    def __init__(self, cfg: LiveBackendConfig, *, namespace: str) -> None:
+        self.cfg = cfg
+        self.namespace = namespace
+        self._client: Any | None = None
+        self._tunnel_ctx = None
+        self._endpoint_host: str | None = None
+        self._endpoint_port: int | None = None
+
+    def connect(self) -> None:
+        """Connect to Neptune, optionally through SSH tunnel."""
+        self._tunnel_ctx = neptune_tunnel(self.cfg)
+        host, port = self._tunnel_ctx.__enter__()
+        self._endpoint_host = host
+        self._endpoint_port = port
+
+        endpoint_url = f"{'https' if self.cfg.neptune_use_https else 'http'}://{host}:{port}"
+        session = boto3.Session()
+        kwargs: dict[str, Any] = {
+            "endpoint_url": endpoint_url,
+            "region_name": self.cfg.neptune_region,
+        }
+
+        if self.cfg.neptune_use_ssh_tunnel:
+            kwargs["verify"] = False
+            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+        if not self.cfg.neptune_iam_auth:
+            kwargs["config"] = BotoConfig(signature_version=UNSIGNED)
+
+        client = session.client("neptunedata", **kwargs)
+
+        if self.cfg.neptune_use_ssh_tunnel and self.cfg.neptune_iam_auth and self.cfg.neptune_endpoint:
+            sign_host = self.cfg.neptune_endpoint
+
+            def _inject_host_header(request, **_kwargs):
+                request.headers["Host"] = f"{sign_host}:{self.cfg.neptune_port}"
+
+            client.meta.events.register("before-sign.neptunedata.*", _inject_host_header)
+
+        self._client = client
+
+    def close(self) -> None:
+        """Close tunnel context when opened."""
+        if self._tunnel_ctx is not None:
+            try:
+                self._tunnel_ctx.__exit__(None, None, None)
+            finally:
+                self._tunnel_ctx = None
+        self._client = None
+
+    def _exec(self, query: str, params: dict[str, Any]) -> dict[str, Any]:
+        if self._client is None:
+            raise RuntimeError("Neptune client is not initialized")
+        return self._client.execute_open_cypher_query(
+            openCypherQuery=query,
+            parameters=json.dumps(params, separators=(",", ":")),
+        )
+
+    @staticmethod
+    def _coerce_scalar(value: Any) -> Any:
+        if isinstance(value, dict):
+            for key in ("stringValue", "intValue", "longValue", "doubleValue", "booleanValue"):
+                if key in value:
+                    return value[key]
+        return value
+
+    def health_check(self) -> None:
+        """Validate openCypher query execution."""
+        self._exec("RETURN 1 AS ok", {})
+
+    @staticmethod
+    def _file_id(project_id: str, path: str) -> str:
+        import hashlib
+
+        h = hashlib.sha256()
+        h.update(project_id.encode("utf-8"))
+        h.update(path.encode("utf-8"))
+        return h.hexdigest()
+
+    def upsert_project(self, *, project_id: str, name: str, repo_path: str, repo_id: str) -> None:
+        ts = datetime.now(timezone.utc).isoformat()
+        query = """
+        MERGE (p:DevMemProject {project_id: $project_id, ns: $ns})
+        SET p.name = $name,
+            p.repo_path = $repo_path,
+            p.repo_id = $repo_id,
+            p.updated_at = $updated_at
+        RETURN p.project_id AS project_id
+        """
+        self._exec(
+            query,
+            {
+                "project_id": project_id,
+                "ns": self.namespace,
+                "name": name,
+                "repo_path": repo_path,
+                "repo_id": repo_id,
+                "updated_at": ts,
+            },
+        )
+
+    def upsert_file(
+        self,
+        *,
+        project_id: str,
+        repo_id: str,
+        path: str,
+        ext: str,
+        sha: str,
+        size_bytes: int,
+    ) -> None:
+        ts = datetime.now(timezone.utc).isoformat()
+        file_id = self._file_id(project_id, path)
+
+        query = """
+        MERGE (p:DevMemProject {project_id: $project_id, ns: $ns})
+        MERGE (f:DevMemFile {file_id: $file_id, ns: $ns})
+        SET f.path = $path,
+            f.ext = $ext,
+            f.sha = $sha,
+            f.size_bytes = $size_bytes,
+            f.repo_id = $repo_id,
+            f.updated_at = $updated_at
+        MERGE (p)-[r:HAS_FILE {ns: $ns}]->(f)
+        SET r.updated_at = $updated_at
+        RETURN f.file_id AS file_id
+        """
+
+        self._exec(
+            query,
+            {
+                "project_id": project_id,
+                "repo_id": repo_id,
+                "file_id": file_id,
+                "ns": self.namespace,
+                "path": path,
+                "ext": ext,
+                "sha": sha,
+                "size_bytes": int(size_bytes),
+                "updated_at": ts,
+            },
+        )
+
+    def list_project_files(self, *, project_id: str, repo_id: str | None = None, limit: int = 20000) -> list[str]:
+        """Return file paths linked to a project in this namespace."""
+        query = f"""
+        MATCH (p:DevMemProject {{project_id: $project_id, ns: $ns}})-[:HAS_FILE {{ns: $ns}}]->(f:DevMemFile {{ns: $ns}})
+        WHERE $repo_id IS NULL OR f.repo_id = $repo_id
+        RETURN f.path AS path
+        LIMIT {int(limit)}
+        """
+        response = self._exec(
+            query,
+            {
+                "project_id": project_id,
+                "repo_id": repo_id,
+                "ns": self.namespace,
+            },
+        )
+        results = response.get("results") or []
+        paths: list[str] = []
+        for row in results:
+            if isinstance(row, dict):
+                path = self._coerce_scalar(row.get("path"))
+                if path:
+                    paths.append(str(path))
+        return paths
+
+    def endpoint_summary(self) -> str:
+        if self._endpoint_host and self._endpoint_port:
+            return f"{self._endpoint_host}:{self._endpoint_port}"
+        return f"{self.cfg.neptune_endpoint}:{self.cfg.neptune_port}"