cypher-graphdb 0.2.2__tar.gz → 0.2.4__tar.gz

{cypher_graphdb-0.2.2/src/cypher_graphdb.egg-info → cypher_graphdb-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cypher_graphdb
-Version: 0.2.2
+Version: 0.2.4
 Summary: ORM like library and CLI for cypher query language supporting graph databases.
 Author-email: Wolfgang Miller <wolfgang.miller@petrarca-labs.com>
 License-Expression: Apache-2.0
@@ -15,7 +15,7 @@ Requires-Dist: art
 Requires-Dist: python-dotenv
 Requires-Dist: prompt-toolkit
 Requires-Dist: rich
-Requires-Dist: apache-age-python
+Requires-Dist: antlr4-python3-runtime==4.11.1
 Requires-Dist: openpyxl
 Requires-Dist: numpy
 Requires-Dist: pydantic
@@ -57,9 +57,6 @@ Requires-Dist: mkdocs-git-revision-date-localized-plugin; extra == "docs"
 Requires-Dist: pymdown-extensions; extra == "docs"
 Requires-Dist: termynal; extra == "docs"
 Requires-Dist: black; extra == "docs"
-Provides-Extra: examples
-Requires-Dist: fastapi; extra == "examples"
-Requires-Dist: uvicorn; extra == "examples"
 Dynamic: license-file
 
 # CypherGraphDB Library

{cypher_graphdb-0.2.2 → cypher_graphdb-0.2.4}/pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
     "python-dotenv",
     "prompt-toolkit",
     "rich",
-    "apache-age-python",
+    "antlr4-python3-runtime==4.11.1",
     "openpyxl",
     "numpy",
     "pydantic",
@@ -76,13 +76,10 @@ docs = [
     "black",
 ]
 
-examples = [
-    "fastapi",
-    "uvicorn",
-]
+
 
 [tool.uv.sources]
-apache-age-python = { git = "https://github.com/apache/age.git", tag = "PG18/v1.7.0-rc0", subdirectory = "drivers/python" }
+
 
 [project.scripts]
 cgdb-cli = "cypher_graphdb.args:typer_main"

cypher_graphdb-0.2.4/src/cypher_graphdb/backends/age/agebulkwriter.py

@@ -0,0 +1,322 @@
+"""Direct SQL bulk writer for Apache AGE label tables.
+
+Bypasses the Cypher parser by INSERTing directly into AGE's PostgreSQL label
+tables using _graphid() + nextval() for ID generation and ::agtype for
+properties. This is the same pattern AGE's own test suite uses
+(regress/sql/age_global_graph.sql) and is 10-30x faster than Cypher UNWIND
+for large bulk loads.
+
+The writer is stateless: it receives a connection and graph_name, performs
+the INSERT, and returns. Transaction control (commit/rollback) is the
+caller's responsibility, with one exception: when a label does not yet exist,
+``_ensure_label`` creates it via DDL inside a savepoint so that the DDL is
+committed immediately (DDL cannot run inside a user transaction in PostgreSQL)
+without affecting any surrounding data already written by the caller.
+"""
+
+from typing import Literal
+
+import psycopg
+from loguru import logger
+from psycopg.sql import SQL, Identifier
+
+from cypher_graphdb.utils import chunk_list
+
+from .ageserializer import escape_value
+from .agesqlbuilder import SQLBuilder
+
+
+class AGEBulkWriter:
+    """Direct SQL bulk writer for AGE label tables.
+
+    Args:
+        connection: An open psycopg connection to the AGE database.
+        graph_name: The AGE graph name (= PostgreSQL schema name).
+    """
+
+    def __init__(self, connection: psycopg.Connection, graph_name: str) -> None:
+        self._conn = connection
+        self._graph_name = graph_name
+        # Cache for {(label, ref_prop): {ref_value: graphid}} maps.
+        # Avoids repeated full-table scans when bulk_insert_edges is called
+        # many times with the same src/dst labels (e.g. 400+ batches of CALLS
+        # edges all resolving against Method.qualified_name).
+        # Invalidated by invalidate_graphid_cache() after node inserts that
+        # change the graphid mapping.
+        self._graphid_cache: dict[tuple[str, str], dict[str, int]] = {}
+
+    # -- Public API -----------------------------------------------------------
+
+    def bulk_insert_nodes(self, label: str, rows: list[dict], batch_size: int = 200) -> int:
+        """Insert nodes via direct SQL INSERT into the AGE label table.
+
+        Creates the vlabel if it doesn't exist yet (AGE normally does this
+        lazily on first Cypher CREATE).
+
+        Uses multi-row INSERT with inline agtype literals (not parameterized)
+        because AGE's agtype_in text-input function does not interpret JSON
+        escape sequences correctly when received as psycopg text parameters.
+        The agtype values are built using the same escape_value function as
+        the Cypher UNWIND path, ensuring identical storage.
+
+        Args:
+            label: Vertex label name (e.g. "Method").
+            rows: List of property dicts, one per node.
+            batch_size: Number of rows per multi-row INSERT.
+
+        Returns:
+            Total number of nodes inserted.
+        """
+        if not rows:
+            return 0
+
+        label_id = self._ensure_label(label, kind="v")
+
+        total = 0
+        with self._conn.cursor() as cursor:
+            for batch in chunk_list(rows, batch_size):
+                values = ", ".join(
+                    f"(ag_catalog._graphid({label_id}, nextval('{self._seq_name(label)}'::regclass)), "
+                    f"{self._to_agtype_literal(row)}::ag_catalog.agtype)"
+                    for row in batch
+                )
+                cursor.execute(
+                    SQL("INSERT INTO {schema}.{table} (id, properties) VALUES {vals}").format(
+                        schema=Identifier(self._graph_name),
+                        table=Identifier(label),
+                        vals=SQL(values),
+                    )
+                )
+                total += len(batch)
+
+        logger.debug("Direct SQL: {} {} nodes inserted", total, label)
+        # Invalidate cached graphid maps for this label so subsequent edge
+        # inserts see the newly created nodes.
+        self.invalidate_graphid_cache(label)
+        return total
+
+    def bulk_insert_edges(
+        self,
+        label: str,
+        edges: list[dict],
+        src_label: str,
+        dst_label: str,
+        src_ref_prop: str = "id",
+        dst_ref_prop: str = "id",
+        batch_size: int = 500,
+    ) -> int:
+        """Insert edges via direct SQL INSERT into the AGE edge label table.
+
+        Pre-resolves source and destination graphids via a single SQL query
+        per label, then INSERTs edges with known start_id/end_id -- no Cypher
+        MATCH at write time.
+
+        Edges whose src or dst reference cannot be resolved are silently
+        skipped (same behaviour as Cypher MATCH -- unmatched patterns produce
+        no rows).
+
+        Args:
+            label: Edge label name (e.g. "CALLS").
+            edges: List of dicts with at least "src" and "dst" keys.
+            src_label: Vertex label of source nodes.
+            dst_label: Vertex label of destination nodes.
+            src_ref_prop: Property name on source nodes to match "src" values.
+            dst_ref_prop: Property name on destination nodes to match "dst" values.
+            batch_size: Number of rows per executemany batch.
+
+        Returns:
+            Total number of edges inserted (excludes skipped).
+        """
+        if not edges:
+            return 0
+
+        label_id = self._ensure_label(label, kind="e")
+
+        # Collect the unique ref values we actually need -- avoids loading
+        # the entire label table (which can be 200K+ rows in a shared graph).
+        src_refs = {e["src"] for e in edges}
+        dst_refs = {e["dst"] for e in edges}
+
+        # Build {ref_value: graphid} maps for only the referenced nodes
+        src_map = self._build_graphid_index(src_label, src_ref_prop, ref_values=src_refs)
+        if src_label == dst_label and src_ref_prop == dst_ref_prop:
+            dst_map = self._build_graphid_index(dst_label, dst_ref_prop, ref_values=src_refs | dst_refs)
+        else:
+            dst_map = self._build_graphid_index(dst_label, dst_ref_prop, ref_values=dst_refs)
+
+        total = 0
+        skipped = 0
+        with self._conn.cursor() as cursor:
+            for batch in chunk_list(edges, batch_size):
+                value_parts = []
+                for edge in batch:
+                    src_gid = src_map.get(edge["src"])
+                    dst_gid = dst_map.get(edge["dst"])
+                    if src_gid is None or dst_gid is None:
+                        skipped += 1
+                        continue
+                    props = {k: v for k, v in edge.items() if k not in ("src", "dst")}
+                    value_parts.append(
+                        f"(ag_catalog._graphid({label_id}, nextval('{self._seq_name(label)}'::regclass)), "
+                        f"'{src_gid}'::ag_catalog.graphid, '{dst_gid}'::ag_catalog.graphid, "
+                        f"{self._to_agtype_literal(props)}::ag_catalog.agtype)"
+                    )
+                if value_parts:
+                    values = ", ".join(value_parts)
+                    cursor.execute(
+                        SQL("INSERT INTO {schema}.{table} (id, start_id, end_id, properties) VALUES {vals}").format(
+                            schema=Identifier(self._graph_name),
+                            table=Identifier(label),
+                            vals=SQL(values),
+                        )
+                    )
+                    total += len(value_parts)
+
+        if skipped:
+            logger.debug("Direct SQL: {} {} edges inserted, {} skipped (unresolved endpoints)", total, label, skipped)
+        else:
+            logger.debug("Direct SQL: {} {} edges inserted", total, label)
+        return total
+
+    # -- Internal helpers -----------------------------------------------------
+
+    def _ensure_label(self, label: str, kind: Literal["v", "e"] = "v") -> int:
+        """Look up the numeric label_id, creating the label if needed.
+
+        When the label does not exist, creates it via a savepoint so that the
+        DDL commit is isolated. PostgreSQL DDL (CREATE TABLE, which AGE's
+        create_vlabel/create_elabel issue internally) cannot run inside an open
+        transaction; committing via savepoint releases the DDL without
+        discarding any data already written by the caller in the outer
+        transaction.
+
+        Args:
+            label: Label name (e.g. "Method", "CALLS").
+            kind: ``"v"`` for vertex label, ``"e"`` for edge label.
+
+        Returns:
+            The integer label_id for the label.
+
+        Raises:
+            ValueError: If the label cannot be created or looked up.
+        """
+        lookup_sql = SQLBuilder.lookup_label_id_sql()
+        with self._conn.cursor() as cursor:
+            cursor.execute(lookup_sql, (self._graph_name, label))
+            row = cursor.fetchone()
+            if row:
+                return row[0]
+
+            # Label doesn't exist -- create it inside a savepoint so that the
+            # DDL commit does not affect the caller's surrounding transaction.
+            cursor.execute("SAVEPOINT _age_ensure_label")
+            try:
+                if kind == "v":
+                    cursor.execute(SQLBuilder.create_vlabel(self._graph_name, label))
+                else:
+                    cursor.execute(SQLBuilder.create_elabel(self._graph_name, label))
+                cursor.execute("RELEASE SAVEPOINT _age_ensure_label")
+                self._conn.commit()
+            except Exception:
+                cursor.execute("ROLLBACK TO SAVEPOINT _age_ensure_label")
+                raise
+
+            cursor.execute(lookup_sql, (self._graph_name, label))
+            row = cursor.fetchone()
+
+        if not row:
+            raise ValueError(f"Failed to create label '{label}' in graph '{self._graph_name}'")
+        return row[0]
+
+    def _seq_name(self, label: str) -> str:
+        """Return the sequence name for a label's entry IDs."""
+        return f'{self._graph_name}."{label}_id_seq"'
+
+    @staticmethod
+    def _to_agtype_literal(props: dict) -> str:
+        """Serialize a dict to a PostgreSQL string literal for ``::agtype`` cast.
+
+        The agtype text-input parser interprets ``\\"`` as an escaped double
+        quote inside strings (same as JSON). When the literal is embedded in a
+        standard PostgreSQL ``'...'`` string, backslashes are passed through
+        verbatim — so the ``\\"`` from ``escape_value`` reaches the agtype
+        parser intact. This produces identical stored values to the Cypher
+        UNWIND path.
+
+        Single quotes inside values are escaped as ``''`` (PostgreSQL standard).
+
+        Note: standard ``'...'`` strings (not ``E'...'``) are used intentionally.
+        ``E'...'`` would consume the ``\\`` escape layer before agtype sees it,
+        breaking ``\\"`` → ``"`` prematurely.
+        """
+        parts = ", ".join(f"{escape_value(k)}: {escape_value(v)}" for k, v in props.items())
+        inner = "{" + parts + "}"
+        # Escape single quotes for PostgreSQL standard string literal
+        inner = inner.replace("'", "''")
+        return "'" + inner + "'"
+
+    def invalidate_graphid_cache(self, label: str | None = None) -> None:
+        """Drop cached graphid maps so the next lookup re-reads from the database.
+
+        Call after inserting nodes into a label that is later used as an edge
+        endpoint, so the graphid map includes the newly inserted nodes.
+
+        Args:
+            label: If given, only invalidate caches for this label.
+                   If None, clear the entire cache.
+        """
+        if label is None:
+            self._graphid_cache.clear()
+        else:
+            keys_to_drop = [k for k in self._graphid_cache if k[0] == label]
+            for k in keys_to_drop:
+                del self._graphid_cache[k]
+
+    def _build_graphid_index(self, label: str, ref_prop: str, ref_values: set[str] | None = None) -> dict[str, int]:
+        """Load a {ref_prop_value: graphid} mapping for nodes of a label.
+
+        When ``ref_values`` is provided, only loads graphids for those specific
+        property values (using a WHERE IN clause). This is critical for shared
+        graphs where a label table may contain hundreds of thousands of rows
+        across all sources, but only a small subset is needed for the current
+        edge batch.
+
+        When ``ref_values`` is None, loads the entire table (used for small
+        labels like structural edges).
+
+        Results are cached per (label, ref_prop) and incrementally extended
+        when new ref_values are requested.
+
+        Values returned by agtype_access_operator are agtype-quoted strings
+        (e.g. '"mth:foo.bar"'), so we strip the surrounding quotes.
+        """
+        cache_key = (label, ref_prop)
+        cached = self._graphid_cache.get(cache_key)
+
+        if ref_values is not None and cached is not None:
+            # Check if all requested values are already in the cache
+            missing = ref_values - cached.keys()
+            if not missing:
+                return cached
+            # Only query the missing values
+            ref_values = missing
+
+        if ref_values is not None and ref_values:
+            sql_stmt = SQLBuilder.lookup_node_graphids_filtered_sql(self._graph_name, label, ref_prop, ref_values)
+        else:
+            sql_stmt = SQLBuilder.lookup_node_graphids_sql(self._graph_name, label, ref_prop)
+
+        with self._conn.cursor() as cursor:
+            cursor.execute(sql_stmt)
+            rows = cursor.fetchall()
+        result = {str(r[0]).strip('"'): r[1] for r in rows}
+
+        # Merge into cache
+        if cached is not None:
+            cached.update(result)
+            logger.debug("Graphid index extended for {}.{}: +{} entries (total {})", label, ref_prop, len(result), len(cached))
+            return cached
+
+        self._graphid_cache[cache_key] = result
+        logger.debug("Graphid index built for {}.{}: {} entries", label, ref_prop, len(result))
+        return result

{cypher_graphdb-0.2.2 → cypher_graphdb-0.2.4}/src/cypher_graphdb/backends/age/agegraphdb.py

@@ -29,6 +29,7 @@ from cypher_graphdb.models import GraphObject, GraphObjectType, TabularResult
 from cypher_graphdb.statistics import IndexInfo, IndexType, LabelStatistics
 from cypher_graphdb.utils import chunk_list, sanitize_connection_params_for_logging, sanitize_connection_string_for_logging
 
+from .agebulkwriter import AGEBulkWriter
 from .agerowfactories import age_row_factory
 from .agesearch import convert_to_fts_query
 from .ageserializer import to_cypher_list
@@ -83,6 +84,17 @@ class AGEGraphDB(CypherBackend):
         self._prepared_statements = {}  # query_hash -> stmt_name
         self._max_cached_statements = 10
 
+        # When True, bulk_create_nodes and bulk_create_edges use direct SQL
+        # INSERT into AGE label tables instead of Cypher UNWIND. This bypasses
+        # the Cypher parser overhead (~3ms per statement) and is 10-30x faster
+        # for large bulk loads. The Cypher path stays as the fallback.
+        # Set to False to revert to the original Cypher UNWIND path.
+        self.direct_bulk_insert: bool = True
+
+        # Lazy AGEBulkWriter instance -- created on first bulk operation and
+        # reused for the lifetime of the connection. Cleared on disconnect/reconnect.
+        self._bulk_writer: AGEBulkWriter | None = None
+
         if args or kwargs:
             self.connect(*args, **kwargs)
 
@@ -412,12 +424,14 @@ class AGEGraphDB(CypherBackend):
             self._connection = None
         self._cinfo = None
         self._ckwargs = None
+        self._bulk_writer = None
 
     def reconnect(self):
         """Reconnect to the database using stored connection info."""
         assert self._cinfo is not None, "Can only reconnect if already successfully connected!"
 
         logger.debug("Try to reconnect")
+        self._bulk_writer = None
         self.connect_to_db()
 
     def commit(self):
@@ -556,17 +570,22 @@ class AGEGraphDB(CypherBackend):
     # ── Index management ────────────────────────────────────────────────
 
     def create_property_index(self, label: str, *property_names: str) -> None:
-        """Create a GIN index on the properties column of a label table.
+        """Create indexes on a label table for fast property lookups.
 
-        AGE stores all node properties in a single agtype JSON column.
-        A GIN index on this column accelerates all property-based MATCH
-        lookups (the @> containment operator). The property_names parameter
-        is accepted for API compatibility but ignored -- one GIN index
-        covers all properties.
+        Creates two kinds of indexes:
+
+        1. **GIN index** on the whole ``properties`` column -- covers the
+           ``@>`` containment operator used by some internal AGE operations.
+        2. **Btree expression indexes** on each specified property -- covers
+           the ``agtype_access_operator`` expression that AGE generates for
+           Cypher ``WHERE n.prop = ...`` clauses. Without these, every property
+           lookup is a sequential scan regardless of the GIN index.
 
         Args:
             label: Node label to index (e.g. "Method").
-            *property_names: Ignored for AGE (GIN covers all properties).
+            *property_names: Property names to create expression indexes for.
+                Each gets its own btree index. If empty, only the GIN index
+                is created.
         """
         self._require_connection()
         graph_name = self._graph_name
@@ -579,17 +598,24 @@ class AGEGraphDB(CypherBackend):
             return
 
         with self._fetch_cursor(row_factory=None) as cursor:
+            # GIN on whole properties column (covers @> containment operator)
             cursor.execute(SQLBuilder.create_gin_index(graph_name, label))
+            # Btree expression index per property (covers agtype_access_operator in WHERE)
+            for prop in property_names:
+                cursor.execute(SQLBuilder.create_expression_index(graph_name, label, prop))
             self._connection.commit()
 
-        logger.debug("GIN property index created for label '{}' in graph '{}'", label, graph_name)
+        logger.debug("Property indexes created for label '{}': GIN + {} expression indexes", label, len(property_names))
 
     def drop_index(self, label: str, *property_names: str) -> None:
-        """Drop the GIN property index on a label table.
+        """Drop the GIN and expression indexes on a label table.
+
+        Drops the GIN index on the whole ``properties`` column, plus any
+        btree expression indexes for the specified properties.
 
         Args:
-            label: Node label whose index to drop.
-            *property_names: Ignored for AGE.
+            label: Node label whose indexes to drop.
+            *property_names: Property names whose expression indexes to drop.
         """
         self._require_connection()
         graph_name = self._graph_name
@@ -597,9 +623,11 @@ class AGEGraphDB(CypherBackend):
 
         with self._fetch_cursor(row_factory=None) as cursor:
             cursor.execute(SQLBuilder.drop_gin_index(graph_name, label))
+            for prop in property_names:
+                cursor.execute(SQLBuilder.drop_expression_index(graph_name, label, prop))
             self._connection.commit()
 
-        logger.debug("GIN property index dropped for label '{}' in graph '{}'", label, graph_name)
+        logger.debug("Property indexes dropped for label '{}': GIN + {} expression indexes", label, len(property_names))
 
     def list_indexes(self, include_internal: bool = False) -> list[IndexInfo]:
         """List all indexes on the current graph.
@@ -649,15 +677,16 @@ class AGEGraphDB(CypherBackend):
     # ── Bulk write operations ─────────────────────────────────────────────
 
     def bulk_create_nodes(self, label: str, rows: list[dict], batch_size: int = 200) -> int:
-        """Create nodes in batches using UNWIND with inline Cypher literals.
+        """Create nodes in batches.
 
-        AGE does not support $params in UNWIND. Rows are serialized as
-        inline Cypher map literals and executed in batches.
+        When ``direct_bulk_insert`` is True (default), uses direct SQL INSERT
+        via :class:`AGEBulkWriter` -- bypasses the Cypher parser for 10-30x
+        faster bulk loads. Falls back to Cypher UNWIND when disabled.
 
         Args:
             label: Node label for all created nodes.
             rows: List of property dicts, one per node.
-            batch_size: Number of nodes per UNWIND batch.
+            batch_size: Number of nodes per batch.
 
         Returns:
             Total number of nodes created.
@@ -666,11 +695,13 @@ class AGEGraphDB(CypherBackend):
         if not rows:
             return 0
 
+        if self.direct_bulk_insert:
+            return self._get_bulk_writer().bulk_insert_nodes(label, rows, batch_size)
+
         total = 0
         for batch in chunk_list(rows, batch_size):
             self._write_node_batch(label, batch)
             total += len(batch)
-
         return total
 
     def bulk_create_edges(
@@ -685,8 +716,9 @@ class AGEGraphDB(CypherBackend):
     ) -> int:
         """Create edges in batches by matching src/dst nodes on a reference property.
 
-        Each dict in edges must have "src" and "dst" keys whose values
-        match the src_ref_prop/dst_ref_prop properties on source/destination nodes.
+        When ``direct_bulk_insert`` is True and both ``src_label`` and ``dst_label``
+        are specified, uses direct SQL INSERT via :class:`AGEBulkWriter` with
+        pre-resolved graphids. Falls back to Cypher UNWIND MATCH CREATE otherwise.
 
         Args:
             label: Edge label for all created edges.
@@ -695,7 +727,7 @@ class AGEGraphDB(CypherBackend):
             dst_label: Label of destination nodes (empty string for any label).
             src_ref_prop: Property name on source nodes to match against "src".
             dst_ref_prop: Property name on destination nodes to match against "dst".
-            batch_size: Number of edges per UNWIND batch.
+            batch_size: Number of edges per batch.
 
         Returns:
             Total number of edges created.
@@ -704,7 +736,12 @@ class AGEGraphDB(CypherBackend):
         if not edges:
             return 0
 
-        # Build MATCH patterns from label/prop names (not user values -- safe to inline)
+        if self.direct_bulk_insert and src_label and dst_label:
+            return self._get_bulk_writer().bulk_insert_edges(
+                label, edges, src_label, dst_label, src_ref_prop, dst_ref_prop, batch_size
+            )
+
+        # Cypher UNWIND fallback (when labels not specified or direct insert disabled)
         src_pat = f"(a:{src_label} {{{src_ref_prop}: e.src}})" if src_label else f"(a {{{src_ref_prop}: e.src}})"
         dst_pat = f"(b:{dst_label} {{{dst_ref_prop}: e.dst}})" if dst_label else f"(b {{{dst_ref_prop}: e.dst}})"
 
@@ -712,17 +749,22 @@ class AGEGraphDB(CypherBackend):
         for batch in chunk_list(edges, batch_size):
             self._write_edge_batch(label, batch, src_pat, dst_pat)
             total += len(batch)
-
         return total
 
-    # ── Private helpers for bulk write ────────────────────────────────────
-
-    def _write_node_batch(self, label: str, batch: list[dict]) -> None:
-        """Write one batch of node dicts via UNWIND CREATE.
+    def _get_bulk_writer(self) -> AGEBulkWriter:
+        """Return the shared AGEBulkWriter, creating it lazily on first call.
 
-        All dicts in batch must have the same set of keys -- the SET clause
-        is derived from the first row and applied to all rows.
+        The writer is tied to the current connection and graph. It is cleared
+        automatically on disconnect() and reconnect() so it is never stale.
         """
+        if self._bulk_writer is None:
+            self._bulk_writer = AGEBulkWriter(self._connection, self._graph_name)
+        return self._bulk_writer
+
+    # ── Cypher UNWIND helpers (fallback path) ─────────────────────────────
+
+    def _write_node_batch(self, label: str, batch: list[dict]) -> None:
+        """Write one batch of node dicts via Cypher UNWIND CREATE."""
         props_keys = list(batch[0].keys())
         set_clause = ", ".join(f"n.{k} = props.{k}" for k in props_keys)
         cypher = f"UNWIND {to_cypher_list(batch)} AS props CREATE (n:{label}) SET {set_clause}"
@@ -731,11 +773,7 @@ class AGEGraphDB(CypherBackend):
         self.execute_cypher(parsed)
 
     def _write_edge_batch(self, label: str, batch: list[dict], src_pat: str, dst_pat: str) -> None:
-        """Write one batch of edges via UNWIND MATCH CREATE.
-
-        Any keys in batch dicts beyond 'src' and 'dst' are set as edge properties.
-        """
-        # Detect extra property keys (beyond src/dst) from first row
+        """Write one batch of edges via Cypher UNWIND MATCH CREATE."""
         edge_prop_keys = [k for k in batch[0] if k not in ("src", "dst")]
         if edge_prop_keys:
             set_clause = " SET " + ", ".join(f"r.{k} = e.{k}" for k in edge_prop_keys)

{cypher_graphdb-0.2.2 → cypher_graphdb-0.2.4}/src/cypher_graphdb/backends/age/agerowfactories.py

@@ -8,7 +8,6 @@ providers.
 from collections.abc import Sequence
 from typing import Any
 
-import age
 from psycopg.cursor import BaseCursor
 from psycopg.rows import BaseRowFactory, RowMaker, T
 
@@ -16,15 +15,17 @@ from cypher_graphdb.backend import ExecStatistics
 from cypher_graphdb.modelprovider import ModelProvider
 from cypher_graphdb.models import GraphEdge, GraphNode, GraphPath
 
+from . import agtype_parser
 
-def _create_graph_node(value: age.models.Vertex, stats: ExecStatistics, provider: ModelProvider) -> GraphNode:
+
+def _create_graph_node(value: agtype_parser.Vertex, stats: ExecStatistics, provider: ModelProvider) -> GraphNode:
     """Create a GraphNode from an AGE Vertex."""
     node = provider.create_node(value.label, value.properties, id_=value.id)
     stats.track_node(node)
     return node
 
 
-def _create_graph_edge(value: age.models.Edge, stats: ExecStatistics, provider: ModelProvider) -> GraphEdge:
+def _create_graph_edge(value: agtype_parser.Edge, stats: ExecStatistics, provider: ModelProvider) -> GraphEdge:
     """Create a GraphEdge from an AGE Edge."""
     edge = provider.create_edge(value.label, value.start_id, value.end_id, value.properties, id_=value.id)
     stats.track_edge(edge)
@@ -36,11 +37,11 @@ def _create_graph_path(value, stats: ExecStatistics, provider: ModelProvider) ->
     path = GraphPath()
 
     for entity in value.entities:
-        if isinstance(entity, age.models.Vertex):
+        if isinstance(entity, agtype_parser.Vertex):
             # Don't update stats here, will be updated when path stats are updated
             node = provider.create_node(entity.label, entity.properties, id_=entity.id)
             path.append(node)
-        elif isinstance(entity, age.models.Edge):
+        elif isinstance(entity, agtype_parser.Edge):
             # Don't update stats here, will be updated when path stats are updated
             edge = provider.create_edge(entity.label, entity.start_id, entity.end_id, entity.properties, id_=entity.id)
             path.append(edge)
@@ -53,11 +54,11 @@ def _create_graph_path(value, stats: ExecStatistics, provider: ModelProvider) ->
 
 def _map_age_value(value, stats: ExecStatistics, provider: ModelProvider):
     """Map an AGE value to the appropriate graph object type."""
-    if isinstance(value, age.models.Vertex):
+    if isinstance(value, agtype_parser.Vertex):
         return _create_graph_node(value, stats, provider)
-    elif isinstance(value, age.models.Edge):
+    elif isinstance(value, agtype_parser.Edge):
         return _create_graph_edge(value, stats, provider)
-    elif isinstance(value, age.models.Path):
+    elif isinstance(value, agtype_parser.Path):
         return _create_graph_path(value, stats, provider)
     else:
         # Only count non-null values (matching Memgraph behavior)