Initial SFERA platform baseline
This commit is contained in:
@@ -0,0 +1,237 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
import json
|
||||
|
||||
from sir import EdgeKind, SemanticEdge, SemanticNode, SirDelta, SirSnapshot
|
||||
|
||||
|
||||
@dataclass
|
||||
class InMemoryProjection:
|
||||
nodes: dict[str, SemanticNode] = field(default_factory=dict)
|
||||
edges: dict[str, SemanticEdge] = field(default_factory=dict)
|
||||
|
||||
def project_snapshot(self, snapshot: SirSnapshot) -> None:
|
||||
self.nodes = {node.lineage_id: node for node in snapshot.nodes}
|
||||
self.edges = {edge.edge_id: edge for edge in snapshot.edges}
|
||||
|
||||
def apply_delta(self, delta: SirDelta) -> None:
|
||||
for lineage_id in delta.removed_nodes:
|
||||
self.nodes.pop(lineage_id, None)
|
||||
self.edges = {
|
||||
edge_id: edge
|
||||
for edge_id, edge in self.edges.items()
|
||||
if edge.source_lineage != lineage_id and edge.target_lineage != lineage_id
|
||||
}
|
||||
for edge_id in delta.removed_edges:
|
||||
self.edges.pop(edge_id, None)
|
||||
for node in [*delta.added_nodes, *delta.updated_nodes]:
|
||||
self.nodes[node.lineage_id] = node
|
||||
for edge in delta.added_edges:
|
||||
self.edges[edge.edge_id] = edge
|
||||
|
||||
def find_procedures(self) -> list[SemanticNode]:
|
||||
return [
|
||||
node
|
||||
for node in self.nodes.values()
|
||||
if node.kind.value in {"PROCEDURE", "FUNCTION"}
|
||||
]
|
||||
|
||||
def find_callers(self, routine_name: str) -> list[SemanticNode]:
|
||||
target_ids = self._routine_lineages(routine_name)
|
||||
caller_ids = [
|
||||
edge.source_lineage
|
||||
for edge in self.edges.values()
|
||||
if edge.kind == EdgeKind.CALLS and edge.target_lineage in target_ids
|
||||
]
|
||||
return [self.nodes[lineage] for lineage in caller_ids if lineage in self.nodes]
|
||||
|
||||
def find_callees(self, routine_name: str) -> list[SemanticNode]:
|
||||
source_ids = self._routine_lineages(routine_name)
|
||||
callee_ids = [
|
||||
edge.target_lineage
|
||||
for edge in self.edges.values()
|
||||
if edge.kind == EdgeKind.CALLS and edge.source_lineage in source_ids
|
||||
]
|
||||
return [self.nodes[lineage] for lineage in callee_ids if lineage in self.nodes]
|
||||
|
||||
def find_query_tables(self, routine_name: str) -> list[SemanticNode]:
|
||||
routine_ids = self._routine_lineages(routine_name)
|
||||
query_ids = [
|
||||
edge.target_lineage
|
||||
for edge in self.edges.values()
|
||||
if edge.kind == EdgeKind.OWNS_QUERY and edge.source_lineage in routine_ids
|
||||
]
|
||||
table_ids = [
|
||||
edge.target_lineage
|
||||
for edge in self.edges.values()
|
||||
if edge.kind == EdgeKind.READS_TABLE and edge.source_lineage in query_ids
|
||||
]
|
||||
return [self.nodes[lineage] for lineage in table_ids if lineage in self.nodes]
|
||||
|
||||
def find_writes(self, routine_name: str) -> list[SemanticNode]:
|
||||
routine_ids = self._routine_lineages(routine_name)
|
||||
write_ids = [
|
||||
edge.target_lineage
|
||||
for edge in self.edges.values()
|
||||
if edge.kind == EdgeKind.WRITES and edge.source_lineage in routine_ids
|
||||
]
|
||||
return [self.nodes[lineage] for lineage in write_ids if lineage in self.nodes]
|
||||
|
||||
def _routine_lineages(self, routine_name: str) -> set[str]:
|
||||
wanted = routine_name.lower()
|
||||
return {
|
||||
node.lineage_id
|
||||
for node in self.nodes.values()
|
||||
if node.name.lower() == wanted and node.kind.value in {"PROCEDURE", "FUNCTION"}
|
||||
}
|
||||
|
||||
|
||||
class Neo4jProjection:
|
||||
def __init__(self, driver) -> None:
|
||||
self._driver = driver
|
||||
|
||||
async def ensure_schema(self) -> None:
|
||||
async with self._driver.session() as session:
|
||||
await session.run(
|
||||
"""
|
||||
CREATE CONSTRAINT sfera_node_lineage IF NOT EXISTS
|
||||
FOR (n:SferaNode) REQUIRE n.lineage_id IS UNIQUE
|
||||
"""
|
||||
)
|
||||
await session.run(
|
||||
"""
|
||||
CREATE INDEX sfera_node_project_name IF NOT EXISTS
|
||||
FOR (n:SferaNode) ON (n.project_id, n.name)
|
||||
"""
|
||||
)
|
||||
await session.run(
|
||||
"""
|
||||
CREATE INDEX sfera_edge_kind IF NOT EXISTS
|
||||
FOR ()-[r:SEMANTIC_EDGE]-() ON (r.kind)
|
||||
"""
|
||||
)
|
||||
|
||||
async def project_snapshot(self, snapshot: SirSnapshot) -> None:
|
||||
await self.ensure_schema()
|
||||
async with self._driver.session() as session:
|
||||
for node in snapshot.nodes:
|
||||
await self._merge_node(session, node, snapshot.project_id)
|
||||
for edge in snapshot.edges:
|
||||
await self._merge_edge(session, edge)
|
||||
|
||||
async def apply_delta(self, delta: SirDelta, *, project_id: str) -> None:
|
||||
await self.ensure_schema()
|
||||
async with self._driver.session() as session:
|
||||
for lineage_id in delta.removed_nodes:
|
||||
await session.run(
|
||||
"""
|
||||
MATCH (n:SferaNode {project_id: $project_id, lineage_id: $lineage_id})
|
||||
DETACH DELETE n
|
||||
""",
|
||||
project_id=project_id,
|
||||
lineage_id=lineage_id,
|
||||
)
|
||||
for edge_id in delta.removed_edges:
|
||||
await session.run(
|
||||
"""
|
||||
MATCH ()-[r:SEMANTIC_EDGE {edge_id: $edge_id}]->()
|
||||
DELETE r
|
||||
""",
|
||||
edge_id=edge_id,
|
||||
)
|
||||
for node in [*delta.added_nodes, *delta.updated_nodes]:
|
||||
await self._merge_node(session, node, project_id)
|
||||
for edge in delta.added_edges:
|
||||
await self._merge_edge(session, edge)
|
||||
|
||||
async def clear_project(self, project_id: str) -> None:
|
||||
async with self._driver.session() as session:
|
||||
await session.run(
|
||||
"""
|
||||
MATCH (n:SferaNode {project_id: $project_id})
|
||||
DETACH DELETE n
|
||||
""",
|
||||
project_id=project_id,
|
||||
)
|
||||
|
||||
async def counts(self, project_id: str | None = None) -> dict[str, int]:
|
||||
async with self._driver.session() as session:
|
||||
if project_id is None:
|
||||
node_result = await session.run("MATCH (n:SferaNode) RETURN count(n) AS count")
|
||||
edge_result = await session.run(
|
||||
"MATCH ()-[r:SEMANTIC_EDGE]->() RETURN count(r) AS count"
|
||||
)
|
||||
else:
|
||||
node_result = await session.run(
|
||||
"""
|
||||
MATCH (n:SferaNode {project_id: $project_id})
|
||||
RETURN count(n) AS count
|
||||
""",
|
||||
project_id=project_id,
|
||||
)
|
||||
edge_result = await session.run(
|
||||
"""
|
||||
MATCH (:SferaNode {project_id: $project_id})-[r:SEMANTIC_EDGE]->
|
||||
(:SferaNode {project_id: $project_id})
|
||||
RETURN count(r) AS count
|
||||
""",
|
||||
project_id=project_id,
|
||||
)
|
||||
node_record = await node_result.single()
|
||||
edge_record = await edge_result.single()
|
||||
return {
|
||||
"nodes": int(node_record["count"]) if node_record else 0,
|
||||
"edges": int(edge_record["count"]) if edge_record else 0,
|
||||
}
|
||||
|
||||
async def _merge_node(self, session, node: SemanticNode, project_id: str) -> None:
|
||||
await session.run(
|
||||
"""
|
||||
MERGE (n:SferaNode {lineage_id: $lineage_id})
|
||||
SET n.semantic_id = $semantic_id,
|
||||
n.project_id = $project_id,
|
||||
n.kind = $kind,
|
||||
n.name = $name,
|
||||
n.qualified_name = $qualified_name,
|
||||
n.attributes_json = $attributes_json,
|
||||
n.source_ref_json = $source_ref_json
|
||||
""",
|
||||
lineage_id=node.lineage_id,
|
||||
semantic_id=node.semantic_id,
|
||||
project_id=project_id,
|
||||
kind=node.kind.value,
|
||||
name=node.name,
|
||||
qualified_name=node.qualified_name,
|
||||
attributes_json=json.dumps(node.attributes, ensure_ascii=False, sort_keys=True),
|
||||
source_ref_json=node.source_ref.model_dump_json(exclude_none=True),
|
||||
)
|
||||
|
||||
async def _merge_edge(self, session, edge: SemanticEdge) -> None:
|
||||
await session.run(
|
||||
"""
|
||||
MATCH (source:SferaNode {lineage_id: $source_lineage})
|
||||
MATCH (target:SferaNode {lineage_id: $target_lineage})
|
||||
MERGE (source)-[r:SEMANTIC_EDGE {edge_id: $edge_id}]->(target)
|
||||
SET r.kind = $kind,
|
||||
r.attributes_json = $attributes_json,
|
||||
r.source_ref_json = $source_ref_json
|
||||
""",
|
||||
source_lineage=edge.source_lineage,
|
||||
target_lineage=edge.target_lineage,
|
||||
edge_id=edge.edge_id,
|
||||
kind=edge.kind.value,
|
||||
attributes_json=json.dumps(edge.attributes, ensure_ascii=False, sort_keys=True),
|
||||
source_ref_json=edge.source_ref.model_dump_json(exclude_none=True) if edge.source_ref else None,
|
||||
)
|
||||
|
||||
|
||||
def build_adjacency(snapshot: SirSnapshot) -> dict[str, list[SemanticEdge]]:
|
||||
adjacency: dict[str, list[SemanticEdge]] = defaultdict(list)
|
||||
for edge in snapshot.edges:
|
||||
adjacency[edge.source_lineage].append(edge)
|
||||
return dict(adjacency)
|
||||
|
||||
|
||||
__all__ = ["InMemoryProjection", "Neo4jProjection", "build_adjacency"]
|
||||
Reference in New Issue
Block a user