graflo.db

Database connection and management components.

This package provides database connection implementations and management utilities for different graph databases (ArangoDB, Neo4j, TigerGraph). It includes connection interfaces, query execution, and database operations.

Key Components

• Connection: Abstract database connection interface
• ConnectionManager: Database connection management
• ArangoDB: ArangoDB-specific implementation
• Neo4j: Neo4j-specific implementation
• TigerGraph: TigerGraph-specific implementation
• Query: Query generation and execution utilities

Example

from graflo.db import ConnectionManager from graflo.db.arango import ArangoConnection manager = ConnectionManager( ... connection_config={"url": "http://localhost:8529"}, ... conn_class=ArangoConnection ... ) with manager as conn: ... conn.init_db(schema)

ArangoConnection

Bases: Connection

ArangoDB-specific implementation of the Connection interface.

This class provides ArangoDB-specific implementations for all database operations, including graph management, document operations, and query execution. It uses the ArangoDB Python driver for all operations.

Attributes:

Name	Type	Description
conn		ArangoDB database connection instance
flavor		Database type (ARANGO) for expression flavor mapping

Source code in graflo/db/arango/conn.py

class ArangoConnection(Connection):
    """ArangoDB-specific implementation of the Connection interface.

    This class provides ArangoDB-specific implementations for all database
    operations, including graph management, document operations, and query
    execution. It uses the ArangoDB Python driver for all operations.

    Attributes:
        conn: ArangoDB database connection instance
        flavor: Database type (ARANGO) for expression flavor mapping
    """

    flavor = DBType.ARANGO

    def __init__(self, config: ArangoConfig):
        """Initialize ArangoDB connection.

        Args:
            config: ArangoDB connection configuration containing URL, credentials,
                and database name
        """
        super().__init__()
        # Store config for later use
        self.config = config
        # Validate required config values
        if config.url is None:
            raise ValueError("ArangoDB connection requires a URL to be configured")
        if config.database is None:
            raise ValueError(
                "ArangoDB connection requires a database name to be configured"
            )

        # ArangoDB accepts empty string for password if None
        password = config.password if config.password is not None else ""
        # ArangoDB has default username "root" if None
        username = config.username if config.username is not None else "root"

        # Store client for system operations
        self.client = ArangoClient(
            hosts=config.url, request_timeout=config.request_timeout
        )
        # Connect to the configured database for regular operations
        self.conn = self.client.db(
            config.database,
            username=username,
            password=password,
        )
        # Store credentials for system operations
        self._username = username
        self._password = password

    def create_database(self, name: str) -> None:
        """Create a new ArangoDB database.

        Database creation/deletion operations must be performed from the _system database.

        Args:
            name: Name of the database to create
        """
        try:
            # Connect to _system database for system operations
            system_db = self.client.db(
                "_system", username=self._username, password=self._password
            )
            if not system_db.has_database(name):
                try:
                    system_db.create_database(name)
                    logger.info(f"Successfully created ArangoDB database '{name}'")
                except Exception as create_error:
                    logger.error(
                        f"Failed to create ArangoDB database '{name}': {create_error}",
                        exc_info=True,
                    )
                    raise
            else:
                logger.debug(f"ArangoDB database '{name}' already exists")
        except Exception as e:
            logger.error(
                f"Error creating ArangoDB database '{name}': {e}",
                exc_info=True,
            )
            raise

    def delete_database(self, name: str) -> None:
        """Delete an ArangoDB database.

        Database creation/deletion operations must be performed from the _system database.

        Args:
            name: Name of the database to delete
        """
        try:
            # Connect to _system database for system operations
            system_db = self.client.db(
                "_system", username=self._username, password=self._password
            )
            if system_db.has_database(name):
                try:
                    system_db.delete_database(name)
                    logger.info(f"Successfully deleted ArangoDB database '{name}'")
                except Exception as delete_error:
                    logger.error(
                        f"Failed to delete ArangoDB database '{name}': {delete_error}",
                        exc_info=True,
                    )
                    raise
            else:
                logger.debug(
                    f"ArangoDB database '{name}' does not exist, skipping deletion"
                )
        except Exception as e:
            logger.error(
                f"Error deleting ArangoDB database '{name}': {e}",
                exc_info=True,
            )
            raise

    def execute(self, query: str, **kwargs: Any) -> Any:
        """Execute an AQL query.

        Args:
            query: AQL query string to execute
            **kwargs: Additional query parameters

        Returns:
            Cursor: ArangoDB cursor for the query results
        """
        cursor = self.conn.aql.execute(query)
        return cursor

    def close(self) -> None:
        """Close the ArangoDB connection."""
        # self.conn.close()
        pass

    def init_db(self, schema: Schema, recreate_schema: bool) -> None:
        """Initialize ArangoDB with the given schema.

        Checks if the database exists and creates it if it doesn't.
        Uses schema.general.name if database is not set in config.

        If the schema/graph already exists and recreate_schema is False, raises
        SchemaExistsError and the script halts.

        Args:
            schema: Schema containing graph structure definitions
            recreate_schema: If True, drop existing vertex/edge classes and define new ones.
                If False and any collections or graphs exist, raises SchemaExistsError.
        """
        # Determine database name: use config.database if set, otherwise use schema.general.name
        db_name = self.config.database
        if not db_name:
            db_name = schema.general.name
            # Update config for subsequent operations
            self.config.database = db_name

        # Check if database exists and create it if it doesn't
        # Use context manager pattern for system database operations
        try:
            system_db = self.client.db(
                "_system", username=self._username, password=self._password
            )
            if not system_db.has_database(db_name):
                logger.info(f"Database '{db_name}' does not exist, creating it...")
                try:
                    system_db.create_database(db_name)
                    logger.info(f"Successfully created database '{db_name}'")
                except Exception as create_error:
                    logger.error(
                        f"Failed to create database '{db_name}': {create_error}",
                        exc_info=True,
                    )
                    raise

            # Reconnect to the target database (newly created or existing)
            if (
                self.config.database != db_name
                or not hasattr(self, "_db_connected")
                or self._db_connected != db_name
            ):
                try:
                    self.conn = self.client.db(
                        db_name, username=self._username, password=self._password
                    )
                    self._db_connected = db_name
                    logger.debug(f"Connected to database '{db_name}'")
                except Exception as conn_error:
                    logger.error(
                        f"Failed to connect to database '{db_name}': {conn_error}",
                        exc_info=True,
                    )
                    raise
        except Exception as e:
            logger.error(
                f"Error during database initialization for '{db_name}': {e}",
                exc_info=True,
            )
            raise

        try:
            # Check if schema/graph already exists (any non-system collection or graph)
            graphs_result = self.conn.graphs()
            collections_result = self.conn.collections()
            has_graphs = isinstance(graphs_result, list) and len(graphs_result) > 0
            non_system = []
            if isinstance(collections_result, list):
                for c in collections_result:
                    if isinstance(c, dict):
                        name_value = cast(dict[str, Any], c).get("name")
                        if isinstance(name_value, str) and name_value[0] != "_":
                            non_system.append(name_value)
            has_collections = len(non_system) > 0
            if (has_graphs or has_collections) and not recreate_schema:
                raise SchemaExistsError(
                    f"Schema/graph already exists in database '{db_name}'. "
                    "Set recreate_schema=True to replace, or use clear_data=True before ingestion."
                )

            if recreate_schema:
                try:
                    self.delete_graph_structure((), (), delete_all=True)
                    logger.debug(f"Cleaned database '{db_name}' for fresh start")
                except Exception as clean_error:
                    logger.warning(
                        f"Error during recreate_schema for database '{db_name}': {clean_error}",
                        exc_info=True,
                    )
                    # Continue - may be first run or already clean

            try:
                self.define_schema(schema)
                logger.debug(f"Defined schema for database '{db_name}'")
            except Exception as schema_error:
                logger.error(
                    f"Failed to define schema for database '{db_name}': {schema_error}",
                    exc_info=True,
                )
                raise

            try:
                self.define_indexes(schema)
                logger.debug(f"Defined indexes for database '{db_name}'")
            except Exception as index_error:
                logger.error(
                    f"Failed to define indexes for database '{db_name}': {index_error}",
                    exc_info=True,
                )
                raise
        except SchemaExistsError:
            raise
        except Exception as e:
            logger.error(
                f"Error during database schema initialization for '{db_name}': {e}",
                exc_info=True,
            )
            raise

    def clear_data(self, schema: Schema) -> None:
        """Remove all data from collections without dropping the schema.

        Truncates vertex and edge collections that belong to the schema.
        """
        vc = schema.vertex_config
        for v in vc.vertex_set:
            cname = vc.vertex_dbname(v)
            if self.conn.has_collection(cname):
                self.conn.collection(cname).truncate()
                logger.debug(f"Truncated vertex collection '{cname}'")
        for edge in schema.edge_config.edges_list(include_aux=True):
            cname = edge.database_name
            if cname and self.conn.has_collection(cname):
                self.conn.collection(cname).truncate()
                logger.debug(f"Truncated edge collection '{cname}'")

    def define_schema(self, schema: Schema) -> None:
        """Define ArangoDB collections based on schema.

        Args:
            schema: Schema containing collection definitions
        """
        self.define_vertex_classes(schema)
        self.define_edge_classes(schema.edge_config.edges_list(include_aux=True))

    def define_vertex_classes(self, schema: Schema) -> None:
        """Define vertex collections in ArangoDB.

        Creates vertex collections for both connected and disconnected vertices,
        organizing them into appropriate graphs.

        Args:
            schema: Schema containing vertex definitions
        """
        vertex_config = schema.vertex_config
        disconnected_vertex_collections = (
            set(vertex_config.vertex_set) - schema.edge_config.vertices
        )
        for item in schema.edge_config.edges_list():
            u, v = item.source, item.target
            gname = item.graph_name
            if not gname:
                logger.warning(
                    f"Edge {item.source} -> {item.target} has no graph_name, skipping"
                )
                continue
            logger.info(f"{item.source}, {item.target}, {gname}")
            if self.conn.has_graph(gname):
                g_result = self.conn.graph(gname)
            else:
                g_result = self.conn.create_graph(gname)  # type: ignore

            # Type narrowing: ensure g is a Graph instance
            g: Graph | None = None
            if isinstance(g_result, Graph):
                g = g_result
            elif g_result is not None:
                # If it's not a Graph, log warning and skip
                logger.warning(f"Graph {gname} is not a Graph instance, skipping")
                continue

            _ = self.create_collection(
                vertex_config.vertex_dbname(u), vertex_config.index(u), g
            )

            _ = self.create_collection(
                vertex_config.vertex_dbname(v), vertex_config.index(v), g
            )
        for v in disconnected_vertex_collections:
            _ = self.create_collection(
                vertex_config.vertex_dbname(v), vertex_config.index(v), None
            )

    def define_edge_classes(self, edges: list[Edge]) -> None:
        """Define edge classes in ArangoDB.

        Creates edge collections and their definitions in the appropriate graphs.

        Args:
            edges: List of edge configurations to create
        """
        for item in edges:
            gname = item.graph_name
            if not gname:
                logger.warning("Edge has no graph_name, skipping")
                continue
            if self.conn.has_graph(gname):
                g_result = self.conn.graph(gname)
            else:
                g_result = self.conn.create_graph(gname)
            # Type guard: ensure g is a Graph instance
            if not isinstance(g_result, Graph):
                logger.warning(f"Graph {gname} is not a Graph instance, skipping")
                continue
            g = g_result
            collection_name = item.database_name
            if not collection_name:
                logger.warning("Edge has no database_name, skipping")
                continue
            if not g.has_edge_definition(collection_name):
                if item._source is None or item._target is None:
                    logger.warning(
                        "Edge has no _source or _target, skipping edge definition"
                    )
                    continue
                _ = g.create_edge_definition(
                    edge_collection=collection_name,
                    from_vertex_collections=[item._source],
                    to_vertex_collections=[item._target],
                )

    def _add_index(self, general_collection: Any, index: Index) -> Any | None:
        """Add an index to an ArangoDB collection.

        Supports persistent, hash, skiplist, and fulltext indices.

        Args:
            general_collection: ArangoDB collection to add index to
            index: Index configuration to create

        Returns:
            IndexHandle: Handle to the created index, or None if index type is not supported
        """
        data = index.db_form(DBType.ARANGO)
        ih: Any | None = None
        if index.type == IndexType.PERSISTENT:
            ih = general_collection.add_index(data)
        elif index.type == IndexType.HASH:
            ih = general_collection.add_index(data)
        elif index.type == IndexType.SKIPLIST:
            ih = general_collection.add_skiplist_index(
                fields=index.fields, unique=index.unique
            )
        elif index.type == IndexType.FULLTEXT:
            ih = general_collection.add_index(
                data={"fields": index.fields, "type": "fulltext"}
            )
        return ih

    def define_vertex_indices(self, vertex_config: VertexConfig) -> None:
        """Define indices for vertex collections.

        Creates indices for each vertex collection based on the configuration.

        Args:
            vertex_config: Vertex configuration containing index definitions
        """
        for c in vertex_config.vertex_set:
            general_collection = self.conn.collection(vertex_config.vertex_dbname(c))
            ixs = general_collection.indexes()
            field_combinations: list[tuple[Any, ...]] = []
            if isinstance(ixs, list):
                for ix in ixs:
                    if isinstance(ix, dict):
                        ix_dict = cast(dict[str, Any], ix)
                        fields_value = ix_dict.get("fields")
                        if isinstance(fields_value, (list, tuple)):
                            field_combinations.append(tuple(fields_value))
            for index_obj in vertex_config.indexes(c):
                if tuple(index_obj.fields) not in field_combinations:
                    self._add_index(general_collection, index_obj)

    def define_edge_indices(self, edges: list[Edge]) -> None:
        """Define indices for edge collections.

        Creates indices for each edge collection based on the configuration.

        Args:
            edges: List of edge configurations containing index definitions
        """
        for edge in edges:
            collection_name = edge.database_name
            if not collection_name:
                logger.warning("Edge has no database_name, skipping index creation")
                continue
            general_collection = self.conn.collection(collection_name)
            for index_obj in edge.indexes:
                self._add_index(general_collection, index_obj)

    def fetch_indexes(self, db_class_name: str | None = None) -> dict[str, Any]:
        """Fetch all indices from the database.

        Args:
            db_class_name: Optional collection name to fetch indices for

        Returns:
            dict: Mapping of collection names to their indices
        """
        classes: list[Any] = []
        if db_class_name is None:
            classes_result = self.conn.collections()
            if isinstance(classes_result, list):
                classes = classes_result
        elif self.conn.has_collection(db_class_name):
            classes = [self.conn.collection(db_class_name)]

        r: dict[str, Any] = {}
        for cname in classes:
            if isinstance(cname, dict):
                cname_dict = cast(dict[str, Any], cname)
                name_value = cname_dict.get("name")
                if isinstance(name_value, str):
                    c = self.conn.collection(name_value)
                    r[name_value] = c.indexes()
        return r

    def create_collection(
        self,
        db_class_name: str,
        index: None | Index = None,
        g: Graph | None = None,
    ) -> Any | None:
        """Create a new vertex or edge class (ArangoDB uses collections internally).

        Args:
            db_class_name: Name of the vertex/edge class to create (ArangoDB collection name)
            index: Optional index to create on the class
            g: Optional graph to create the class in

        Returns:
            IndexHandle: Handle to the created index if one was created, None otherwise
        """
        if not self.conn.has_collection(db_class_name):
            if g is not None:
                _ = g.create_vertex_collection(db_class_name)
            else:
                self.conn.create_collection(db_class_name)
            general_collection = self.conn.collection(db_class_name)
            if index is not None and index.fields != ["_key"]:
                ih = self._add_index(general_collection, index)
                return ih
            else:
                return None

    def delete_graph_structure(
        self,
        vertex_types: tuple[str, ...] | list[str] = (),
        graph_names: tuple[str, ...] | list[str] = (),
        delete_all: bool = False,
    ) -> None:
        """Delete graph structure (vertex/edge classes and graphs) from ArangoDB.

        In ArangoDB:
        - Collections (internal): Container for vertices (vertex collections) and edges (edge collections)
        - Graphs: Named graphs that connect vertex and edge collections

        Args:
            vertex_types: Vertex/edge class names to delete (ArangoDB collection names)
            graph_names: Graph names to delete
            delete_all: If True, delete all non-system vertex/edge classes and graphs
        """
        cnames: list[str] = list(vertex_types)
        gnames: list[str] = list(graph_names)
        logger.info("vertex/edge classes (non system, ArangoDB collections):")
        collections_result = self.conn.collections()
        if isinstance(collections_result, list):
            filtered_collections: list[dict[str, Any]] = []
            for c in collections_result:
                if isinstance(c, dict):
                    c_dict = cast(dict[str, Any], c)
                    name_value = c_dict.get("name")
                    if isinstance(name_value, str) and name_value[0] != "_":
                        filtered_collections.append(c_dict)
            logger.info(filtered_collections)
        else:
            logger.info([])

        if delete_all:
            collections_result = self.conn.collections()
            graphs_result = self.conn.graphs()
            cnames = []
            if isinstance(collections_result, list):
                for c in collections_result:
                    if isinstance(c, dict):
                        c_dict = cast(dict[str, Any], c)
                        name_value = c_dict.get("name")
                        if isinstance(name_value, str) and name_value[0] != "_":
                            cnames.append(name_value)
            gnames = []
            if isinstance(graphs_result, list):
                for g in graphs_result:
                    if isinstance(g, dict):
                        g_dict = cast(dict[str, Any], g)
                        name_value = g_dict.get("name")
                        if isinstance(name_value, str):
                            gnames.append(name_value)

        for gn in gnames:
            if self.conn.has_graph(gn):
                self.conn.delete_graph(gn)

        logger.info("graphs (after delete operation):")
        logger.info(self.conn.graphs())

        for cn in cnames:
            if self.conn.has_collection(cn):
                self.conn.delete_collection(cn)

        logger.info(
            "vertex/edge classes (after delete operation, ArangoDB collections):"
        )
        collections_result = self.conn.collections()
        if isinstance(collections_result, list):
            collection_names: list[str] = []
            for c in collections_result:
                if isinstance(c, dict):
                    c_dict = cast(dict[str, Any], c)
                    name_value = c_dict.get("name")
                    if isinstance(name_value, str) and name_value[0] != "_":
                        collection_names.append(name_value)
            logger.info(collection_names)
        else:
            logger.info([])

        logger.info("graphs:")
        logger.info(self.conn.graphs())

    def get_collections(self) -> list[dict[str, Any]]:
        """Get all vertex and edge classes in the database (ArangoDB collections).

        Returns:
            list: List of class information dictionaries (ArangoDB collection info)
        """
        result = self.conn.collections()
        if isinstance(result, list):
            return [
                cast(dict[str, Any], item) if isinstance(item, dict) else {}
                for item in result
            ]
        return []

    def upsert_docs_batch(
        self,
        docs: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...] = (),
        **kwargs: Any,
    ) -> None:
        """Upsert a batch of documents using AQL.

        Performs an upsert operation on a batch of documents, using the specified
        match keys to determine whether to update existing documents or insert new ones.

        Args:
            docs: List of documents to upsert
            class_name: Collection name to upsert into
            match_keys: Keys to match for upsert operation
            **kwargs: Additional options:
                - dry: If True, don't execute the query
                - update_keys: Keys to update on match
                - filter_uniques: If True, filter duplicate documents
        """
        dry = kwargs.pop("dry", False)
        update_keys = kwargs.pop("update_keys", None)
        filter_uniques = kwargs.pop("filter_uniques", True)

        if not docs:
            return
        if filter_uniques:
            docs = pick_unique_dict(docs)
        docs_json = json.dumps(docs, default=json_serializer)
        if not match_keys:
            upsert_clause = ""
            update_clause = ""
        else:
            upsert_clause = ", ".join([f'"{k}": doc.{k}' for k in match_keys])
            upsert_clause = f"UPSERT {{{upsert_clause}}}"

            if isinstance(update_keys, list):
                update_clause = ", ".join([f'"{k}": doc.{k}' for k in update_keys])
                update_clause = f"{{{update_clause}}}"
            elif update_keys == "doc":
                update_clause = "doc"
            else:
                update_clause = "{}"
            update_clause = f"UPDATE {update_clause}"

        options = "OPTIONS {exclusive: true, ignoreErrors: true}"

        q_update = f"""FOR doc in {docs_json}
                            {upsert_clause}
                            INSERT doc
                            {update_clause} 
                                IN {class_name} {options}"""
        if not dry:
            self.execute(q_update)

    def insert_edges_batch(
        self,
        docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
        source_class: str,
        target_class: str,
        relation_name: str,
        match_keys_source: tuple[str, ...] = ("_key",),
        match_keys_target: tuple[str, ...] = ("_key",),
        filter_uniques: bool = True,
        head: int | None = None,
        **kwargs: Any,
    ) -> None:
        """Insert a batch of edges using AQL.

        Creates edges between source and target vertices, with support for
        weight fields and unique constraints.

        Args:
            docs_edges: List of edge documents in format [{_source_aux: source_doc, _target_aux: target_doc}]
            source_class: Source vertex class name
            target_class: Target vertex class name
            relation_name: Optional relation name for the edges
            match_keys_source: Keys to match source vertices
            match_keys_target: Keys to match target vertices
            filter_uniques: If True, filter duplicate edges
            head: Optional limit on number of edges to insert
            **kwargs: Additional options:
                - dry: If True, don't execute the query
                - collection_name: Edge collection name (defaults to {source_class}_{target_class}_edges if not provided)
                - uniq_weight_fields: Fields to consider for uniqueness
                - uniq_weight_collections: Classes to consider for uniqueness
                - upsert_option: If True, use upsert instead of insert
        """
        dry = kwargs.pop("dry", False)

        # Extract collection_name from kwargs, with default generation
        collection_name = kwargs.pop("collection_name", None)
        if collection_name is None:
            collection_name = f"{source_class}_{target_class}_edges"

        # Extract ArangoDB-specific parameters from kwargs
        uniq_weight_fields = kwargs.pop("uniq_weight_fields", None)
        uniq_weight_collections = kwargs.pop("uniq_weight_collections", None)
        upsert_option = kwargs.pop("upsert_option", False)

        if isinstance(docs_edges, list):
            if docs_edges:
                logger.debug(f" docs_edges[0] = {docs_edges[0]}")
            if head is not None:
                docs_edges = docs_edges[:head]
            if filter_uniques:
                docs_edges = pick_unique_dict(docs_edges)
            docs_edges_str = json.dumps(docs_edges)
        else:
            return

        if match_keys_source[0] == "_key":
            result_from = f'CONCAT("{source_class}/", edge[0]._key)'
            source_filter = ""
        else:
            result_from = "sources[0]._id"
            filter_source = " && ".join(
                [f"v.{k} == edge[0].{k}" for k in match_keys_source]
            )
            source_filter = (
                f"LET sources = (FOR v IN {source_class} FILTER"
                f" {filter_source} LIMIT 1 RETURN v)"
            )

        if match_keys_target[0] == "_key":
            result_to = f'CONCAT("{target_class}/", edge[1]._key)'
            target_filter = ""
        else:
            result_to = "targets[0]._id"
            filter_target = " && ".join(
                [f"v.{k} == edge[1].{k}" for k in match_keys_target]
            )
            target_filter = (
                f"LET targets = (FOR v IN {target_class} FILTER"
                f" {filter_target} LIMIT 1 RETURN v)"
            )

        doc_definition = f"MERGE({{_from : {result_from}, _to : {result_to}}}, edge[2])"

        logger.debug(f" source_filter = {source_filter}")
        logger.debug(f" target_filter = {target_filter}")
        logger.debug(f" doc = {doc_definition}")

        if upsert_option:
            ups_from = result_from if source_filter else "doc._from"
            ups_to = result_to if target_filter else "doc._to"

            weight_fs = []
            if uniq_weight_fields is not None:
                weight_fs += uniq_weight_fields
            if uniq_weight_collections is not None:
                weight_fs += uniq_weight_collections
            if relation_name is not None:
                weight_fs += ["relation"]

            if weight_fs:
                weights_clause = ", " + ", ".join(
                    [f"'{x}' : edge.{x}" for x in weight_fs]
                )
            else:
                weights_clause = ""

            upsert = f"{{'_from': {ups_from}, '_to': {ups_to}" + weights_clause + "}"
            logger.debug(f" upsert clause: {upsert}")
            clauses = f"UPSERT {upsert} INSERT doc UPDATE {{}}"
            options = "OPTIONS {exclusive: true}"
        else:
            if relation_name is None:
                doc_clause = "doc"
            else:
                doc_clause = f"MERGE(doc, {{'relation': '{relation_name}' }})"
            clauses = f"INSERT {doc_clause}"
            options = "OPTIONS {exclusive: true, ignoreErrors: true}"

        q_update = f"""
            FOR edge in {docs_edges_str} {source_filter} {target_filter}
                LET doc = {doc_definition}
                {clauses}
                in {collection_name} {options}"""
        if not dry:
            self.execute(q_update)

    def insert_return_batch(self, docs: list[dict[str, Any]], class_name: str) -> str:
        """Insert documents and return the AQL query string.

        Note: ArangoDB-specific behavior - returns query string instead of executing.
        This allows for deferred execution and is used by caster.py and tests.

        Args:
            docs: Documents to insert
            class_name: Collection to insert into

        Returns:
            str: AQL query string for the operation (can be executed with execute())
        """
        docs_str = json.dumps(docs, default=json_serializer)
        query0 = f"""FOR doc in {docs_str}
              INSERT doc
              INTO {class_name}
              LET inserted = NEW
              RETURN {{_key: inserted._key}}
        """
        return query0

    def fetch_present_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        flatten: bool = False,
        filters: None | FilterExpression | list[Any] | dict[str, Any] = None,
    ) -> list[dict[str, Any]] | dict[int, list[dict[str, Any]]]:
        """Fetch documents that exist in the database.

        Args:
            batch: Batch of documents to check
            class_name: Collection to check in
            match_keys: Keys to match documents
            keep_keys: Keys to keep in result
            flatten: If True, flatten the result into a list
            filters: Additional query filters

        Returns:
            list | dict: Documents that exist in the database, either as a
                flat list or a dictionary mapping batch indices to documents
        """
        q0 = fetch_fields_query(
            collection_name=class_name,
            docs=batch,
            match_keys=match_keys,
            keep_keys=keep_keys,
            filters=filters,
        )
        # {"__i": i, "_group": [doc]}
        cursor = self.execute(q0)

        if flatten:
            rdata = []
            for item in get_data_from_cursor(cursor):
                group = item.pop("_group", [])
                rdata += [sub_item for sub_item in group]
            return rdata
        else:
            rdata_dict: dict[int, list[dict[str, Any]]] = {}
            for item in get_data_from_cursor(cursor):
                __i = item.pop("__i")
                group = item.pop("_group", [])
                rdata_dict[__i] = group
            return rdata_dict

    def fetch_docs(
        self,
        class_name: str,
        filters: None | FilterExpression | list[Any] | dict[str, Any] = None,
        limit: int | None = None,
        return_keys: list[str] | None = None,
        unset_keys: list[str] | None = None,
        **kwargs: Any,
    ) -> list[dict[str, Any]]:
        """Fetch documents from a collection.

        Args:
            class_name: Collection to fetch from
            filters: Query filters
            limit: Maximum number of documents to return
            return_keys: Keys to return
            unset_keys: Keys to unset

        Returns:
            list: Fetched documents
        """
        filter_clause = render_filters(filters, doc_name="d")

        if return_keys is None:
            if unset_keys is None:
                return_clause = "d"
            else:
                tmp_clause = ", ".join([f'"{item}"' for item in unset_keys])
                return_clause = f"UNSET(d, {tmp_clause})"
        else:
            if unset_keys is None:
                tmp_clause = ", ".join([f'"{item}"' for item in return_keys])
                return_clause = f"KEEP(d, {tmp_clause})"
            else:
                raise ValueError("both return_keys and unset_keys are set")

        if limit is not None and isinstance(limit, int):
            limit_clause = f"LIMIT {limit}"
        else:
            limit_clause = ""

        q = (
            f"FOR d in {class_name}"
            f"  {filter_clause}"
            f"  {limit_clause}"
            f"  RETURN {return_clause}"
        )
        cursor = self.execute(q)
        return get_data_from_cursor(cursor)

    # TODO test
    def fetch_edges(
        self,
        from_type: str,
        from_id: str,
        edge_type: str | None = None,
        to_type: str | None = None,
        to_id: str | None = None,
        filters: list[Any] | dict[str, Any] | FilterExpression | None = None,
        limit: int | None = None,
        return_keys: list[str] | None = None,
        unset_keys: list[str] | None = None,
        **kwargs: Any,
    ) -> list[dict[str, Any]]:
        """Fetch edges from ArangoDB using AQL.

        Args:
            from_type: Source vertex collection name
            from_id: Source vertex ID (can be _key or _id)
            edge_type: Optional edge collection name to filter by
            to_type: Optional target vertex collection name to filter by
            to_id: Optional target vertex ID to filter by
            filters: Additional query filters
            limit: Maximum number of edges to return
            return_keys: Keys to return (projection)
            unset_keys: Keys to exclude (projection)
            **kwargs: Additional parameters

        Returns:
            list: List of fetched edges
        """
        # Convert from_id to _id format if needed
        if not from_id.startswith(from_type):
            # Assume it's a _key, convert to _id
            from_vertex_id = f"{from_type}/{from_id}"
        else:
            from_vertex_id = from_id

        # Build AQL query to fetch edges
        # Start with basic edge traversal
        if edge_type:
            edge_collection = edge_type
        else:
            # If no edge_type specified, we need to search all edge collections
            # This is a simplified version - in practice you might want to list all edge collections
            raise ValueError("edge_type is required for ArangoDB edge fetching")

        filter_clause = render_filters(filters, doc_name="e")
        filter_parts = []

        if to_type:
            filter_parts.append(f"e._to LIKE '{to_type}/%'")
        if to_id and to_type:
            if not to_id.startswith(to_type):
                to_vertex_id = f"{to_type}/{to_id}"
            else:
                to_vertex_id = to_id
            filter_parts.append(f"e._to == '{to_vertex_id}'")

        additional_filters = " && ".join(filter_parts)
        if filter_clause and additional_filters:
            filter_clause = f"{filter_clause} && {additional_filters}"
        elif additional_filters:
            filter_clause = additional_filters

        query = f"""
            FOR e IN {edge_collection}
                FILTER e._from == '{from_vertex_id}'
                {f"FILTER {filter_clause}" if filter_clause else ""}
                {f"LIMIT {limit}" if limit else ""}
                RETURN e
        """

        cursor = self.execute(query)
        result = list(get_data_from_cursor(cursor))

        # Apply projection
        if return_keys is not None:
            result = [
                {k: doc.get(k) for k in return_keys if k in doc} for doc in result
            ]
        elif unset_keys is not None:
            result = [
                {k: v for k, v in doc.items() if k not in unset_keys} for doc in result
            ]

        return result

    def aggregate(
        self,
        class_name: str,
        aggregation_function: AggregationType,
        discriminant: str | None = None,
        aggregated_field: str | None = None,
        filters: None | FilterExpression | list[Any] | dict[str, Any] = None,
    ) -> list[dict[str, Any]]:
        """Perform aggregation on a collection.

        Args:
            class_name: Collection to aggregate
            aggregation_function: Type of aggregation to perform
            discriminant: Field to group by
            aggregated_field: Field to aggregate
            filters: Query filters

        Returns:
            list: Aggregation results
        """
        filter_clause = render_filters(filters, doc_name="doc")

        if (
            aggregated_field is not None
            and aggregation_function != AggregationType.COUNT
        ):
            group_unit = f"g[*].doc.{aggregated_field}"
        else:
            group_unit = "g"

        if discriminant is not None:
            collect_clause = f"COLLECT value = doc['{discriminant}'] INTO g"
            return_clause = f"""{{ '{discriminant}' : value, '_value': {aggregation_function}({group_unit})}}"""
        else:
            if (
                aggregated_field is None
                and aggregation_function == AggregationType.COUNT
            ):
                collect_clause = (
                    f"COLLECT AGGREGATE value =  {aggregation_function} (doc)"
                )
            else:
                collect_clause = (
                    "COLLECT AGGREGATE value ="
                    f" {aggregation_function}(doc['{aggregated_field}'])"
                )
            return_clause = """{ '_value' : value }"""

        q = f"""FOR doc IN {class_name} 
                    {filter_clause}
                    {collect_clause}
                    RETURN {return_clause}"""

        cursor = self.execute(q)
        data = get_data_from_cursor(cursor)
        return data

    def keep_absent_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        filters: None | FilterExpression | list[Any] | dict[str, Any] = None,
    ) -> list[dict[str, Any]]:
        """Keep documents that don't exist in the database.

        Args:
            batch: Batch of documents to check
            class_name: Collection to check in
            match_keys: Keys to match documents
            keep_keys: Keys to keep in result
            filters: Additional query filters

        Returns:
            list: Documents that don't exist in the database
        """
        present_docs_keys = self.fetch_present_documents(
            batch=batch,
            class_name=class_name,
            match_keys=match_keys,
            keep_keys=keep_keys,
            flatten=False,
            filters=filters,
        )

        assert isinstance(present_docs_keys, dict)

        if any([len(v) > 1 for v in present_docs_keys.values()]):
            logger.warning(
                "fetch_present_documents returned multiple docs per filtering condition"
            )

        absent_indices = sorted(set(range(len(batch))) - set(present_docs_keys.keys()))
        batch_absent = [batch[j] for j in absent_indices]
        return batch_absent

    def update_to_numeric(self, collection_name: str, field: str) -> str:
        """Update a field to numeric type in all documents.

        Args:
            collection_name: Vertex/edge class name to update (ArangoDB collection name)
            field: Field to convert to numeric

        Returns:
            str: AQL query string for the operation
        """
        s1 = f"FOR p IN {collection_name} FILTER p.{field} update p with {{"
        s2 = f"{field}: TO_NUMBER(p.{field}) "
        s3 = f"}} in {collection_name}"
        q0 = s1 + s2 + s3
        return q0

__init__(config)

Initialize ArangoDB connection.

Parameters:

Name	Type	Description	Default
config	ArangoConfig	ArangoDB connection configuration containing URL, credentials, and database name	required

Source code in graflo/db/arango/conn.py

def __init__(self, config: ArangoConfig):
    """Initialize ArangoDB connection.

    Args:
        config: ArangoDB connection configuration containing URL, credentials,
            and database name
    """
    super().__init__()
    # Store config for later use
    self.config = config
    # Validate required config values
    if config.url is None:
        raise ValueError("ArangoDB connection requires a URL to be configured")
    if config.database is None:
        raise ValueError(
            "ArangoDB connection requires a database name to be configured"
        )

    # ArangoDB accepts empty string for password if None
    password = config.password if config.password is not None else ""
    # ArangoDB has default username "root" if None
    username = config.username if config.username is not None else "root"

    # Store client for system operations
    self.client = ArangoClient(
        hosts=config.url, request_timeout=config.request_timeout
    )
    # Connect to the configured database for regular operations
    self.conn = self.client.db(
        config.database,
        username=username,
        password=password,
    )
    # Store credentials for system operations
    self._username = username
    self._password = password

aggregate(class_name, aggregation_function, discriminant=None, aggregated_field=None, filters=None)

Perform aggregation on a collection.

Parameters:

Name	Type	Description	Default
class_name	str	Collection to aggregate	required
aggregation_function	AggregationType	Type of aggregation to perform	required
discriminant	str | None	Field to group by	None
aggregated_field	str | None	Field to aggregate	None
filters	None | FilterExpression | list[Any] | dict[str, Any]	Query filters	None

Returns:

Name	Type	Description
list	list[dict[str, Any]]	Aggregation results

Source code in graflo/db/arango/conn.py

def aggregate(
    self,
    class_name: str,
    aggregation_function: AggregationType,
    discriminant: str | None = None,
    aggregated_field: str | None = None,
    filters: None | FilterExpression | list[Any] | dict[str, Any] = None,
) -> list[dict[str, Any]]:
    """Perform aggregation on a collection.

    Args:
        class_name: Collection to aggregate
        aggregation_function: Type of aggregation to perform
        discriminant: Field to group by
        aggregated_field: Field to aggregate
        filters: Query filters

    Returns:
        list: Aggregation results
    """
    filter_clause = render_filters(filters, doc_name="doc")

    if (
        aggregated_field is not None
        and aggregation_function != AggregationType.COUNT
    ):
        group_unit = f"g[*].doc.{aggregated_field}"
    else:
        group_unit = "g"

    if discriminant is not None:
        collect_clause = f"COLLECT value = doc['{discriminant}'] INTO g"
        return_clause = f"""{{ '{discriminant}' : value, '_value': {aggregation_function}({group_unit})}}"""
    else:
        if (
            aggregated_field is None
            and aggregation_function == AggregationType.COUNT
        ):
            collect_clause = (
                f"COLLECT AGGREGATE value =  {aggregation_function} (doc)"
            )
        else:
            collect_clause = (
                "COLLECT AGGREGATE value ="
                f" {aggregation_function}(doc['{aggregated_field}'])"
            )
        return_clause = """{ '_value' : value }"""

    q = f"""FOR doc IN {class_name} 
                {filter_clause}
                {collect_clause}
                RETURN {return_clause}"""

    cursor = self.execute(q)
    data = get_data_from_cursor(cursor)
    return data

clear_data(schema)

Remove all data from collections without dropping the schema.

Truncates vertex and edge collections that belong to the schema.

Source code in graflo/db/arango/conn.py

def clear_data(self, schema: Schema) -> None:
    """Remove all data from collections without dropping the schema.

    Truncates vertex and edge collections that belong to the schema.
    """
    vc = schema.vertex_config
    for v in vc.vertex_set:
        cname = vc.vertex_dbname(v)
        if self.conn.has_collection(cname):
            self.conn.collection(cname).truncate()
            logger.debug(f"Truncated vertex collection '{cname}'")
    for edge in schema.edge_config.edges_list(include_aux=True):
        cname = edge.database_name
        if cname and self.conn.has_collection(cname):
            self.conn.collection(cname).truncate()
            logger.debug(f"Truncated edge collection '{cname}'")

close()

Close the ArangoDB connection.

Source code in graflo/db/arango/conn.py

def close(self) -> None:
    """Close the ArangoDB connection."""
    # self.conn.close()
    pass

create_collection(db_class_name, index=None, g=None)

Create a new vertex or edge class (ArangoDB uses collections internally).

Parameters:

Name	Type	Description	Default
db_class_name	str	Name of the vertex/edge class to create (ArangoDB collection name)	required
index	None | Index	Optional index to create on the class	None
g	Graph | None	Optional graph to create the class in	None

Returns:

Name	Type	Description
IndexHandle	Any | None	Handle to the created index if one was created, None otherwise

Source code in graflo/db/arango/conn.py

def create_collection(
    self,
    db_class_name: str,
    index: None | Index = None,
    g: Graph | None = None,
) -> Any | None:
    """Create a new vertex or edge class (ArangoDB uses collections internally).

    Args:
        db_class_name: Name of the vertex/edge class to create (ArangoDB collection name)
        index: Optional index to create on the class
        g: Optional graph to create the class in

    Returns:
        IndexHandle: Handle to the created index if one was created, None otherwise
    """
    if not self.conn.has_collection(db_class_name):
        if g is not None:
            _ = g.create_vertex_collection(db_class_name)
        else:
            self.conn.create_collection(db_class_name)
        general_collection = self.conn.collection(db_class_name)
        if index is not None and index.fields != ["_key"]:
            ih = self._add_index(general_collection, index)
            return ih
        else:
            return None

create_database(name)

Create a new ArangoDB database.

Database creation/deletion operations must be performed from the _system database.

Parameters:

Name	Type	Description	Default
name	str	Name of the database to create	required

Source code in graflo/db/arango/conn.py

def create_database(self, name: str) -> None:
    """Create a new ArangoDB database.

    Database creation/deletion operations must be performed from the _system database.

    Args:
        name: Name of the database to create
    """
    try:
        # Connect to _system database for system operations
        system_db = self.client.db(
            "_system", username=self._username, password=self._password
        )
        if not system_db.has_database(name):
            try:
                system_db.create_database(name)
                logger.info(f"Successfully created ArangoDB database '{name}'")
            except Exception as create_error:
                logger.error(
                    f"Failed to create ArangoDB database '{name}': {create_error}",
                    exc_info=True,
                )
                raise
        else:
            logger.debug(f"ArangoDB database '{name}' already exists")
    except Exception as e:
        logger.error(
            f"Error creating ArangoDB database '{name}': {e}",
            exc_info=True,
        )
        raise

define_edge_classes(edges)

Define edge classes in ArangoDB.

Creates edge collections and their definitions in the appropriate graphs.

Parameters:

Name	Type	Description	Default
edges	list[Edge]	List of edge configurations to create	required

Source code in graflo/db/arango/conn.py

def define_edge_classes(self, edges: list[Edge]) -> None:
    """Define edge classes in ArangoDB.

    Creates edge collections and their definitions in the appropriate graphs.

    Args:
        edges: List of edge configurations to create
    """
    for item in edges:
        gname = item.graph_name
        if not gname:
            logger.warning("Edge has no graph_name, skipping")
            continue
        if self.conn.has_graph(gname):
            g_result = self.conn.graph(gname)
        else:
            g_result = self.conn.create_graph(gname)
        # Type guard: ensure g is a Graph instance
        if not isinstance(g_result, Graph):
            logger.warning(f"Graph {gname} is not a Graph instance, skipping")
            continue
        g = g_result
        collection_name = item.database_name
        if not collection_name:
            logger.warning("Edge has no database_name, skipping")
            continue
        if not g.has_edge_definition(collection_name):
            if item._source is None or item._target is None:
                logger.warning(
                    "Edge has no _source or _target, skipping edge definition"
                )
                continue
            _ = g.create_edge_definition(
                edge_collection=collection_name,
                from_vertex_collections=[item._source],
                to_vertex_collections=[item._target],
            )

define_edge_indices(edges)

Define indices for edge collections.

Creates indices for each edge collection based on the configuration.

Parameters:

Name	Type	Description	Default
edges	list[Edge]	List of edge configurations containing index definitions	required

Source code in graflo/db/arango/conn.py

def define_edge_indices(self, edges: list[Edge]) -> None:
    """Define indices for edge collections.

    Creates indices for each edge collection based on the configuration.

    Args:
        edges: List of edge configurations containing index definitions
    """
    for edge in edges:
        collection_name = edge.database_name
        if not collection_name:
            logger.warning("Edge has no database_name, skipping index creation")
            continue
        general_collection = self.conn.collection(collection_name)
        for index_obj in edge.indexes:
            self._add_index(general_collection, index_obj)

define_schema(schema)

Define ArangoDB collections based on schema.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing collection definitions	required

Source code in graflo/db/arango/conn.py

def define_schema(self, schema: Schema) -> None:
    """Define ArangoDB collections based on schema.

    Args:
        schema: Schema containing collection definitions
    """
    self.define_vertex_classes(schema)
    self.define_edge_classes(schema.edge_config.edges_list(include_aux=True))

define_vertex_classes(schema)

Define vertex collections in ArangoDB.

Creates vertex collections for both connected and disconnected vertices, organizing them into appropriate graphs.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing vertex definitions	required

Source code in graflo/db/arango/conn.py

def define_vertex_classes(self, schema: Schema) -> None:
    """Define vertex collections in ArangoDB.

    Creates vertex collections for both connected and disconnected vertices,
    organizing them into appropriate graphs.

    Args:
        schema: Schema containing vertex definitions
    """
    vertex_config = schema.vertex_config
    disconnected_vertex_collections = (
        set(vertex_config.vertex_set) - schema.edge_config.vertices
    )
    for item in schema.edge_config.edges_list():
        u, v = item.source, item.target
        gname = item.graph_name
        if not gname:
            logger.warning(
                f"Edge {item.source} -> {item.target} has no graph_name, skipping"
            )
            continue
        logger.info(f"{item.source}, {item.target}, {gname}")
        if self.conn.has_graph(gname):
            g_result = self.conn.graph(gname)
        else:
            g_result = self.conn.create_graph(gname)  # type: ignore

        # Type narrowing: ensure g is a Graph instance
        g: Graph | None = None
        if isinstance(g_result, Graph):
            g = g_result
        elif g_result is not None:
            # If it's not a Graph, log warning and skip
            logger.warning(f"Graph {gname} is not a Graph instance, skipping")
            continue

        _ = self.create_collection(
            vertex_config.vertex_dbname(u), vertex_config.index(u), g
        )

        _ = self.create_collection(
            vertex_config.vertex_dbname(v), vertex_config.index(v), g
        )
    for v in disconnected_vertex_collections:
        _ = self.create_collection(
            vertex_config.vertex_dbname(v), vertex_config.index(v), None
        )

define_vertex_indices(vertex_config)

Define indices for vertex collections.

Creates indices for each vertex collection based on the configuration.

Parameters:

Name	Type	Description	Default
vertex_config	VertexConfig	Vertex configuration containing index definitions	required

Source code in graflo/db/arango/conn.py

def define_vertex_indices(self, vertex_config: VertexConfig) -> None:
    """Define indices for vertex collections.

    Creates indices for each vertex collection based on the configuration.

    Args:
        vertex_config: Vertex configuration containing index definitions
    """
    for c in vertex_config.vertex_set:
        general_collection = self.conn.collection(vertex_config.vertex_dbname(c))
        ixs = general_collection.indexes()
        field_combinations: list[tuple[Any, ...]] = []
        if isinstance(ixs, list):
            for ix in ixs:
                if isinstance(ix, dict):
                    ix_dict = cast(dict[str, Any], ix)
                    fields_value = ix_dict.get("fields")
                    if isinstance(fields_value, (list, tuple)):
                        field_combinations.append(tuple(fields_value))
        for index_obj in vertex_config.indexes(c):
            if tuple(index_obj.fields) not in field_combinations:
                self._add_index(general_collection, index_obj)

delete_database(name)

Delete an ArangoDB database.

Database creation/deletion operations must be performed from the _system database.

Parameters:

Name	Type	Description	Default
name	str	Name of the database to delete	required

Source code in graflo/db/arango/conn.py

def delete_database(self, name: str) -> None:
    """Delete an ArangoDB database.

    Database creation/deletion operations must be performed from the _system database.

    Args:
        name: Name of the database to delete
    """
    try:
        # Connect to _system database for system operations
        system_db = self.client.db(
            "_system", username=self._username, password=self._password
        )
        if system_db.has_database(name):
            try:
                system_db.delete_database(name)
                logger.info(f"Successfully deleted ArangoDB database '{name}'")
            except Exception as delete_error:
                logger.error(
                    f"Failed to delete ArangoDB database '{name}': {delete_error}",
                    exc_info=True,
                )
                raise
        else:
            logger.debug(
                f"ArangoDB database '{name}' does not exist, skipping deletion"
            )
    except Exception as e:
        logger.error(
            f"Error deleting ArangoDB database '{name}': {e}",
            exc_info=True,
        )
        raise

delete_graph_structure(vertex_types=(), graph_names=(), delete_all=False)

Delete graph structure (vertex/edge classes and graphs) from ArangoDB.

In ArangoDB: - Collections (internal): Container for vertices (vertex collections) and edges (edge collections) - Graphs: Named graphs that connect vertex and edge collections

Parameters:

Name	Type	Description	Default
vertex_types	tuple[str, ...] | list[str]	Vertex/edge class names to delete (ArangoDB collection names)	()
graph_names	tuple[str, ...] | list[str]	Graph names to delete	()
delete_all	bool	If True, delete all non-system vertex/edge classes and graphs	False

Source code in graflo/db/arango/conn.py

def delete_graph_structure(
    self,
    vertex_types: tuple[str, ...] | list[str] = (),
    graph_names: tuple[str, ...] | list[str] = (),
    delete_all: bool = False,
) -> None:
    """Delete graph structure (vertex/edge classes and graphs) from ArangoDB.

    In ArangoDB:
    - Collections (internal): Container for vertices (vertex collections) and edges (edge collections)
    - Graphs: Named graphs that connect vertex and edge collections

    Args:
        vertex_types: Vertex/edge class names to delete (ArangoDB collection names)
        graph_names: Graph names to delete
        delete_all: If True, delete all non-system vertex/edge classes and graphs
    """
    cnames: list[str] = list(vertex_types)
    gnames: list[str] = list(graph_names)
    logger.info("vertex/edge classes (non system, ArangoDB collections):")
    collections_result = self.conn.collections()
    if isinstance(collections_result, list):
        filtered_collections: list[dict[str, Any]] = []
        for c in collections_result:
            if isinstance(c, dict):
                c_dict = cast(dict[str, Any], c)
                name_value = c_dict.get("name")
                if isinstance(name_value, str) and name_value[0] != "_":
                    filtered_collections.append(c_dict)
        logger.info(filtered_collections)
    else:
        logger.info([])

    if delete_all:
        collections_result = self.conn.collections()
        graphs_result = self.conn.graphs()
        cnames = []
        if isinstance(collections_result, list):
            for c in collections_result:
                if isinstance(c, dict):
                    c_dict = cast(dict[str, Any], c)
                    name_value = c_dict.get("name")
                    if isinstance(name_value, str) and name_value[0] != "_":
                        cnames.append(name_value)
        gnames = []
        if isinstance(graphs_result, list):
            for g in graphs_result:
                if isinstance(g, dict):
                    g_dict = cast(dict[str, Any], g)
                    name_value = g_dict.get("name")
                    if isinstance(name_value, str):
                        gnames.append(name_value)

    for gn in gnames:
        if self.conn.has_graph(gn):
            self.conn.delete_graph(gn)

    logger.info("graphs (after delete operation):")
    logger.info(self.conn.graphs())

    for cn in cnames:
        if self.conn.has_collection(cn):
            self.conn.delete_collection(cn)

    logger.info(
        "vertex/edge classes (after delete operation, ArangoDB collections):"
    )
    collections_result = self.conn.collections()
    if isinstance(collections_result, list):
        collection_names: list[str] = []
        for c in collections_result:
            if isinstance(c, dict):
                c_dict = cast(dict[str, Any], c)
                name_value = c_dict.get("name")
                if isinstance(name_value, str) and name_value[0] != "_":
                    collection_names.append(name_value)
        logger.info(collection_names)
    else:
        logger.info([])

    logger.info("graphs:")
    logger.info(self.conn.graphs())

execute(query, **kwargs)

Execute an AQL query.

Parameters:

Name	Type	Description	Default
query	str	AQL query string to execute	required
**kwargs	Any	Additional query parameters	{}

Returns:

Name	Type	Description
Cursor	Any	ArangoDB cursor for the query results

Source code in graflo/db/arango/conn.py

def execute(self, query: str, **kwargs: Any) -> Any:
    """Execute an AQL query.

    Args:
        query: AQL query string to execute
        **kwargs: Additional query parameters

    Returns:
        Cursor: ArangoDB cursor for the query results
    """
    cursor = self.conn.aql.execute(query)
    return cursor

fetch_docs(class_name, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs)

Fetch documents from a collection.

Parameters:

Name	Type	Description	Default
class_name	str	Collection to fetch from	required
filters	None | FilterExpression | list[Any] | dict[str, Any]	Query filters	None
limit	int | None	Maximum number of documents to return	None
return_keys	list[str] | None	Keys to return	None
unset_keys	list[str] | None	Keys to unset	None

Returns:

Name	Type	Description
list	list[dict[str, Any]]	Fetched documents

Source code in graflo/db/arango/conn.py

def fetch_docs(
    self,
    class_name: str,
    filters: None | FilterExpression | list[Any] | dict[str, Any] = None,
    limit: int | None = None,
    return_keys: list[str] | None = None,
    unset_keys: list[str] | None = None,
    **kwargs: Any,
) -> list[dict[str, Any]]:
    """Fetch documents from a collection.

    Args:
        class_name: Collection to fetch from
        filters: Query filters
        limit: Maximum number of documents to return
        return_keys: Keys to return
        unset_keys: Keys to unset

    Returns:
        list: Fetched documents
    """
    filter_clause = render_filters(filters, doc_name="d")

    if return_keys is None:
        if unset_keys is None:
            return_clause = "d"
        else:
            tmp_clause = ", ".join([f'"{item}"' for item in unset_keys])
            return_clause = f"UNSET(d, {tmp_clause})"
    else:
        if unset_keys is None:
            tmp_clause = ", ".join([f'"{item}"' for item in return_keys])
            return_clause = f"KEEP(d, {tmp_clause})"
        else:
            raise ValueError("both return_keys and unset_keys are set")

    if limit is not None and isinstance(limit, int):
        limit_clause = f"LIMIT {limit}"
    else:
        limit_clause = ""

    q = (
        f"FOR d in {class_name}"
        f"  {filter_clause}"
        f"  {limit_clause}"
        f"  RETURN {return_clause}"
    )
    cursor = self.execute(q)
    return get_data_from_cursor(cursor)

fetch_edges(from_type, from_id, edge_type=None, to_type=None, to_id=None, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs)

Fetch edges from ArangoDB using AQL.

Parameters:

Name	Type	Description	Default
from_type	str	Source vertex collection name	required
from_id	str	Source vertex ID (can be _key or _id)	required
edge_type	str | None	Optional edge collection name to filter by	None
to_type	str | None	Optional target vertex collection name to filter by	None
to_id	str | None	Optional target vertex ID to filter by	None
filters	list[Any] | dict[str, Any] | FilterExpression | None	Additional query filters	None
limit	int | None	Maximum number of edges to return	None
return_keys	list[str] | None	Keys to return (projection)	None
unset_keys	list[str] | None	Keys to exclude (projection)	None
**kwargs	Any	Additional parameters	{}

Returns:

Name	Type	Description
list	list[dict[str, Any]]	List of fetched edges

Source code in graflo/db/arango/conn.py

def fetch_edges(
    self,
    from_type: str,
    from_id: str,
    edge_type: str | None = None,
    to_type: str | None = None,
    to_id: str | None = None,
    filters: list[Any] | dict[str, Any] | FilterExpression | None = None,
    limit: int | None = None,
    return_keys: list[str] | None = None,
    unset_keys: list[str] | None = None,
    **kwargs: Any,
) -> list[dict[str, Any]]:
    """Fetch edges from ArangoDB using AQL.

    Args:
        from_type: Source vertex collection name
        from_id: Source vertex ID (can be _key or _id)
        edge_type: Optional edge collection name to filter by
        to_type: Optional target vertex collection name to filter by
        to_id: Optional target vertex ID to filter by
        filters: Additional query filters
        limit: Maximum number of edges to return
        return_keys: Keys to return (projection)
        unset_keys: Keys to exclude (projection)
        **kwargs: Additional parameters

    Returns:
        list: List of fetched edges
    """
    # Convert from_id to _id format if needed
    if not from_id.startswith(from_type):
        # Assume it's a _key, convert to _id
        from_vertex_id = f"{from_type}/{from_id}"
    else:
        from_vertex_id = from_id

    # Build AQL query to fetch edges
    # Start with basic edge traversal
    if edge_type:
        edge_collection = edge_type
    else:
        # If no edge_type specified, we need to search all edge collections
        # This is a simplified version - in practice you might want to list all edge collections
        raise ValueError("edge_type is required for ArangoDB edge fetching")

    filter_clause = render_filters(filters, doc_name="e")
    filter_parts = []

    if to_type:
        filter_parts.append(f"e._to LIKE '{to_type}/%'")
    if to_id and to_type:
        if not to_id.startswith(to_type):
            to_vertex_id = f"{to_type}/{to_id}"
        else:
            to_vertex_id = to_id
        filter_parts.append(f"e._to == '{to_vertex_id}'")

    additional_filters = " && ".join(filter_parts)
    if filter_clause and additional_filters:
        filter_clause = f"{filter_clause} && {additional_filters}"
    elif additional_filters:
        filter_clause = additional_filters

    query = f"""
        FOR e IN {edge_collection}
            FILTER e._from == '{from_vertex_id}'
            {f"FILTER {filter_clause}" if filter_clause else ""}
            {f"LIMIT {limit}" if limit else ""}
            RETURN e
    """

    cursor = self.execute(query)
    result = list(get_data_from_cursor(cursor))

    # Apply projection
    if return_keys is not None:
        result = [
            {k: doc.get(k) for k in return_keys if k in doc} for doc in result
        ]
    elif unset_keys is not None:
        result = [
            {k: v for k, v in doc.items() if k not in unset_keys} for doc in result
        ]

    return result

fetch_indexes(db_class_name=None)

Fetch all indices from the database.

Parameters:

Name	Type	Description	Default
db_class_name	str | None	Optional collection name to fetch indices for	None

Returns:

Name	Type	Description
dict	dict[str, Any]	Mapping of collection names to their indices

Source code in graflo/db/arango/conn.py

def fetch_indexes(self, db_class_name: str | None = None) -> dict[str, Any]:
    """Fetch all indices from the database.

    Args:
        db_class_name: Optional collection name to fetch indices for

    Returns:
        dict: Mapping of collection names to their indices
    """
    classes: list[Any] = []
    if db_class_name is None:
        classes_result = self.conn.collections()
        if isinstance(classes_result, list):
            classes = classes_result
    elif self.conn.has_collection(db_class_name):
        classes = [self.conn.collection(db_class_name)]

    r: dict[str, Any] = {}
    for cname in classes:
        if isinstance(cname, dict):
            cname_dict = cast(dict[str, Any], cname)
            name_value = cname_dict.get("name")
            if isinstance(name_value, str):
                c = self.conn.collection(name_value)
                r[name_value] = c.indexes()
    return r

fetch_present_documents(batch, class_name, match_keys, keep_keys=None, flatten=False, filters=None)

Fetch documents that exist in the database.

Parameters:

Name	Type	Description	Default
batch	list[dict[str, Any]]	Batch of documents to check	required
class_name	str	Collection to check in	required
match_keys	list[str] | tuple[str, ...]	Keys to match documents	required
keep_keys	list[str] | tuple[str, ...] | None	Keys to keep in result	None
flatten	bool	If True, flatten the result into a list	False
filters	None | FilterExpression | list[Any] | dict[str, Any]	Additional query filters	None

Returns:

Type	Description
list[dict[str, Any]] | dict[int, list[dict[str, Any]]]	list | dict: Documents that exist in the database, either as a flat list or a dictionary mapping batch indices to documents

Source code in graflo/db/arango/conn.py

def fetch_present_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    flatten: bool = False,
    filters: None | FilterExpression | list[Any] | dict[str, Any] = None,
) -> list[dict[str, Any]] | dict[int, list[dict[str, Any]]]:
    """Fetch documents that exist in the database.

    Args:
        batch: Batch of documents to check
        class_name: Collection to check in
        match_keys: Keys to match documents
        keep_keys: Keys to keep in result
        flatten: If True, flatten the result into a list
        filters: Additional query filters

    Returns:
        list | dict: Documents that exist in the database, either as a
            flat list or a dictionary mapping batch indices to documents
    """
    q0 = fetch_fields_query(
        collection_name=class_name,
        docs=batch,
        match_keys=match_keys,
        keep_keys=keep_keys,
        filters=filters,
    )
    # {"__i": i, "_group": [doc]}
    cursor = self.execute(q0)

    if flatten:
        rdata = []
        for item in get_data_from_cursor(cursor):
            group = item.pop("_group", [])
            rdata += [sub_item for sub_item in group]
        return rdata
    else:
        rdata_dict: dict[int, list[dict[str, Any]]] = {}
        for item in get_data_from_cursor(cursor):
            __i = item.pop("__i")
            group = item.pop("_group", [])
            rdata_dict[__i] = group
        return rdata_dict

get_collections()

Get all vertex and edge classes in the database (ArangoDB collections).

Returns:

Name	Type	Description
list	list[dict[str, Any]]	List of class information dictionaries (ArangoDB collection info)

Source code in graflo/db/arango/conn.py

def get_collections(self) -> list[dict[str, Any]]:
    """Get all vertex and edge classes in the database (ArangoDB collections).

    Returns:
        list: List of class information dictionaries (ArangoDB collection info)
    """
    result = self.conn.collections()
    if isinstance(result, list):
        return [
            cast(dict[str, Any], item) if isinstance(item, dict) else {}
            for item in result
        ]
    return []

init_db(schema, recreate_schema)

Initialize ArangoDB with the given schema.

Checks if the database exists and creates it if it doesn't. Uses schema.general.name if database is not set in config.

If the schema/graph already exists and recreate_schema is False, raises SchemaExistsError and the script halts.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing graph structure definitions	required
recreate_schema	bool	If True, drop existing vertex/edge classes and define new ones. If False and any collections or graphs exist, raises SchemaExistsError.	required

Source code in graflo/db/arango/conn.py

def init_db(self, schema: Schema, recreate_schema: bool) -> None:
    """Initialize ArangoDB with the given schema.

    Checks if the database exists and creates it if it doesn't.
    Uses schema.general.name if database is not set in config.

    If the schema/graph already exists and recreate_schema is False, raises
    SchemaExistsError and the script halts.

    Args:
        schema: Schema containing graph structure definitions
        recreate_schema: If True, drop existing vertex/edge classes and define new ones.
            If False and any collections or graphs exist, raises SchemaExistsError.
    """
    # Determine database name: use config.database if set, otherwise use schema.general.name
    db_name = self.config.database
    if not db_name:
        db_name = schema.general.name
        # Update config for subsequent operations
        self.config.database = db_name

    # Check if database exists and create it if it doesn't
    # Use context manager pattern for system database operations
    try:
        system_db = self.client.db(
            "_system", username=self._username, password=self._password
        )
        if not system_db.has_database(db_name):
            logger.info(f"Database '{db_name}' does not exist, creating it...")
            try:
                system_db.create_database(db_name)
                logger.info(f"Successfully created database '{db_name}'")
            except Exception as create_error:
                logger.error(
                    f"Failed to create database '{db_name}': {create_error}",
                    exc_info=True,
                )
                raise

        # Reconnect to the target database (newly created or existing)
        if (
            self.config.database != db_name
            or not hasattr(self, "_db_connected")
            or self._db_connected != db_name
        ):
            try:
                self.conn = self.client.db(
                    db_name, username=self._username, password=self._password
                )
                self._db_connected = db_name
                logger.debug(f"Connected to database '{db_name}'")
            except Exception as conn_error:
                logger.error(
                    f"Failed to connect to database '{db_name}': {conn_error}",
                    exc_info=True,
                )
                raise
    except Exception as e:
        logger.error(
            f"Error during database initialization for '{db_name}': {e}",
            exc_info=True,
        )
        raise

    try:
        # Check if schema/graph already exists (any non-system collection or graph)
        graphs_result = self.conn.graphs()
        collections_result = self.conn.collections()
        has_graphs = isinstance(graphs_result, list) and len(graphs_result) > 0
        non_system = []
        if isinstance(collections_result, list):
            for c in collections_result:
                if isinstance(c, dict):
                    name_value = cast(dict[str, Any], c).get("name")
                    if isinstance(name_value, str) and name_value[0] != "_":
                        non_system.append(name_value)
        has_collections = len(non_system) > 0
        if (has_graphs or has_collections) and not recreate_schema:
            raise SchemaExistsError(
                f"Schema/graph already exists in database '{db_name}'. "
                "Set recreate_schema=True to replace, or use clear_data=True before ingestion."
            )

        if recreate_schema:
            try:
                self.delete_graph_structure((), (), delete_all=True)
                logger.debug(f"Cleaned database '{db_name}' for fresh start")
            except Exception as clean_error:
                logger.warning(
                    f"Error during recreate_schema for database '{db_name}': {clean_error}",
                    exc_info=True,
                )
                # Continue - may be first run or already clean

        try:
            self.define_schema(schema)
            logger.debug(f"Defined schema for database '{db_name}'")
        except Exception as schema_error:
            logger.error(
                f"Failed to define schema for database '{db_name}': {schema_error}",
                exc_info=True,
            )
            raise

        try:
            self.define_indexes(schema)
            logger.debug(f"Defined indexes for database '{db_name}'")
        except Exception as index_error:
            logger.error(
                f"Failed to define indexes for database '{db_name}': {index_error}",
                exc_info=True,
            )
            raise
    except SchemaExistsError:
        raise
    except Exception as e:
        logger.error(
            f"Error during database schema initialization for '{db_name}': {e}",
            exc_info=True,
        )
        raise

insert_edges_batch(docs_edges, source_class, target_class, relation_name, match_keys_source=('_key',), match_keys_target=('_key',), filter_uniques=True, head=None, **kwargs)

Insert a batch of edges using AQL.

Creates edges between source and target vertices, with support for weight fields and unique constraints.

Parameters:

Name	Type	Description	Default
docs_edges	list[list[dict[str, Any]]] | list[Any] | None	List of edge documents in format [{_source_aux: source_doc, _target_aux: target_doc}]	required
source_class	str	Source vertex class name	required
target_class	str	Target vertex class name	required
relation_name	str	Optional relation name for the edges	required
match_keys_source	tuple[str, ...]	Keys to match source vertices	('_key',)
match_keys_target	tuple[str, ...]	Keys to match target vertices	('_key',)
filter_uniques	bool	If True, filter duplicate edges	True
head	int | None	Optional limit on number of edges to insert	None
**kwargs	Any	Additional options: - dry: If True, don't execute the query - collection_name: Edge collection name (defaults to {source_class}_{target_class}_edges if not provided) - uniq_weight_fields: Fields to consider for uniqueness - uniq_weight_collections: Classes to consider for uniqueness - upsert_option: If True, use upsert instead of insert	{}

Source code in graflo/db/arango/conn.py

def insert_edges_batch(
    self,
    docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
    source_class: str,
    target_class: str,
    relation_name: str,
    match_keys_source: tuple[str, ...] = ("_key",),
    match_keys_target: tuple[str, ...] = ("_key",),
    filter_uniques: bool = True,
    head: int | None = None,
    **kwargs: Any,
) -> None:
    """Insert a batch of edges using AQL.

    Creates edges between source and target vertices, with support for
    weight fields and unique constraints.

    Args:
        docs_edges: List of edge documents in format [{_source_aux: source_doc, _target_aux: target_doc}]
        source_class: Source vertex class name
        target_class: Target vertex class name
        relation_name: Optional relation name for the edges
        match_keys_source: Keys to match source vertices
        match_keys_target: Keys to match target vertices
        filter_uniques: If True, filter duplicate edges
        head: Optional limit on number of edges to insert
        **kwargs: Additional options:
            - dry: If True, don't execute the query
            - collection_name: Edge collection name (defaults to {source_class}_{target_class}_edges if not provided)
            - uniq_weight_fields: Fields to consider for uniqueness
            - uniq_weight_collections: Classes to consider for uniqueness
            - upsert_option: If True, use upsert instead of insert
    """
    dry = kwargs.pop("dry", False)

    # Extract collection_name from kwargs, with default generation
    collection_name = kwargs.pop("collection_name", None)
    if collection_name is None:
        collection_name = f"{source_class}_{target_class}_edges"

    # Extract ArangoDB-specific parameters from kwargs
    uniq_weight_fields = kwargs.pop("uniq_weight_fields", None)
    uniq_weight_collections = kwargs.pop("uniq_weight_collections", None)
    upsert_option = kwargs.pop("upsert_option", False)

    if isinstance(docs_edges, list):
        if docs_edges:
            logger.debug(f" docs_edges[0] = {docs_edges[0]}")
        if head is not None:
            docs_edges = docs_edges[:head]
        if filter_uniques:
            docs_edges = pick_unique_dict(docs_edges)
        docs_edges_str = json.dumps(docs_edges)
    else:
        return

    if match_keys_source[0] == "_key":
        result_from = f'CONCAT("{source_class}/", edge[0]._key)'
        source_filter = ""
    else:
        result_from = "sources[0]._id"
        filter_source = " && ".join(
            [f"v.{k} == edge[0].{k}" for k in match_keys_source]
        )
        source_filter = (
            f"LET sources = (FOR v IN {source_class} FILTER"
            f" {filter_source} LIMIT 1 RETURN v)"
        )

    if match_keys_target[0] == "_key":
        result_to = f'CONCAT("{target_class}/", edge[1]._key)'
        target_filter = ""
    else:
        result_to = "targets[0]._id"
        filter_target = " && ".join(
            [f"v.{k} == edge[1].{k}" for k in match_keys_target]
        )
        target_filter = (
            f"LET targets = (FOR v IN {target_class} FILTER"
            f" {filter_target} LIMIT 1 RETURN v)"
        )

    doc_definition = f"MERGE({{_from : {result_from}, _to : {result_to}}}, edge[2])"

    logger.debug(f" source_filter = {source_filter}")
    logger.debug(f" target_filter = {target_filter}")
    logger.debug(f" doc = {doc_definition}")

    if upsert_option:
        ups_from = result_from if source_filter else "doc._from"
        ups_to = result_to if target_filter else "doc._to"

        weight_fs = []
        if uniq_weight_fields is not None:
            weight_fs += uniq_weight_fields
        if uniq_weight_collections is not None:
            weight_fs += uniq_weight_collections
        if relation_name is not None:
            weight_fs += ["relation"]

        if weight_fs:
            weights_clause = ", " + ", ".join(
                [f"'{x}' : edge.{x}" for x in weight_fs]
            )
        else:
            weights_clause = ""

        upsert = f"{{'_from': {ups_from}, '_to': {ups_to}" + weights_clause + "}"
        logger.debug(f" upsert clause: {upsert}")
        clauses = f"UPSERT {upsert} INSERT doc UPDATE {{}}"
        options = "OPTIONS {exclusive: true}"
    else:
        if relation_name is None:
            doc_clause = "doc"
        else:
            doc_clause = f"MERGE(doc, {{'relation': '{relation_name}' }})"
        clauses = f"INSERT {doc_clause}"
        options = "OPTIONS {exclusive: true, ignoreErrors: true}"

    q_update = f"""
        FOR edge in {docs_edges_str} {source_filter} {target_filter}
            LET doc = {doc_definition}
            {clauses}
            in {collection_name} {options}"""
    if not dry:
        self.execute(q_update)

insert_return_batch(docs, class_name)

Insert documents and return the AQL query string.

Note: ArangoDB-specific behavior - returns query string instead of executing. This allows for deferred execution and is used by caster.py and tests.

Parameters:

Name	Type	Description	Default
docs	list[dict[str, Any]]	Documents to insert	required
class_name	str	Collection to insert into	required

Returns:

Name	Type	Description
str	str	AQL query string for the operation (can be executed with execute())

Source code in graflo/db/arango/conn.py

def insert_return_batch(self, docs: list[dict[str, Any]], class_name: str) -> str:
    """Insert documents and return the AQL query string.

    Note: ArangoDB-specific behavior - returns query string instead of executing.
    This allows for deferred execution and is used by caster.py and tests.

    Args:
        docs: Documents to insert
        class_name: Collection to insert into

    Returns:
        str: AQL query string for the operation (can be executed with execute())
    """
    docs_str = json.dumps(docs, default=json_serializer)
    query0 = f"""FOR doc in {docs_str}
          INSERT doc
          INTO {class_name}
          LET inserted = NEW
          RETURN {{_key: inserted._key}}
    """
    return query0

keep_absent_documents(batch, class_name, match_keys, keep_keys=None, filters=None)

Keep documents that don't exist in the database.

Parameters:

Name	Type	Description	Default
batch	list[dict[str, Any]]	Batch of documents to check	required
class_name	str	Collection to check in	required
match_keys	list[str] | tuple[str, ...]	Keys to match documents	required
keep_keys	list[str] | tuple[str, ...] | None	Keys to keep in result	None
filters	None | FilterExpression | list[Any] | dict[str, Any]	Additional query filters	None

Returns:

Name	Type	Description
list	list[dict[str, Any]]	Documents that don't exist in the database

Source code in graflo/db/arango/conn.py

def keep_absent_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    filters: None | FilterExpression | list[Any] | dict[str, Any] = None,
) -> list[dict[str, Any]]:
    """Keep documents that don't exist in the database.

    Args:
        batch: Batch of documents to check
        class_name: Collection to check in
        match_keys: Keys to match documents
        keep_keys: Keys to keep in result
        filters: Additional query filters

    Returns:
        list: Documents that don't exist in the database
    """
    present_docs_keys = self.fetch_present_documents(
        batch=batch,
        class_name=class_name,
        match_keys=match_keys,
        keep_keys=keep_keys,
        flatten=False,
        filters=filters,
    )

    assert isinstance(present_docs_keys, dict)

    if any([len(v) > 1 for v in present_docs_keys.values()]):
        logger.warning(
            "fetch_present_documents returned multiple docs per filtering condition"
        )

    absent_indices = sorted(set(range(len(batch))) - set(present_docs_keys.keys()))
    batch_absent = [batch[j] for j in absent_indices]
    return batch_absent

update_to_numeric(collection_name, field)

Update a field to numeric type in all documents.

Parameters:

Name	Type	Description	Default
collection_name	str	Vertex/edge class name to update (ArangoDB collection name)	required
field	str	Field to convert to numeric	required

Returns:

Name	Type	Description
str	str	AQL query string for the operation

Source code in graflo/db/arango/conn.py

def update_to_numeric(self, collection_name: str, field: str) -> str:
    """Update a field to numeric type in all documents.

    Args:
        collection_name: Vertex/edge class name to update (ArangoDB collection name)
        field: Field to convert to numeric

    Returns:
        str: AQL query string for the operation
    """
    s1 = f"FOR p IN {collection_name} FILTER p.{field} update p with {{"
    s2 = f"{field}: TO_NUMBER(p.{field}) "
    s3 = f"}} in {collection_name}"
    q0 = s1 + s2 + s3
    return q0

upsert_docs_batch(docs, class_name, match_keys=(), **kwargs)

Upsert a batch of documents using AQL.

Performs an upsert operation on a batch of documents, using the specified match keys to determine whether to update existing documents or insert new ones.

Parameters:

Name	Type	Description	Default
docs	list[dict[str, Any]]	List of documents to upsert	required
class_name	str	Collection name to upsert into	required
match_keys	list[str] | tuple[str, ...]	Keys to match for upsert operation	()
**kwargs	Any	Additional options: - dry: If True, don't execute the query - update_keys: Keys to update on match - filter_uniques: If True, filter duplicate documents	{}

Source code in graflo/db/arango/conn.py

def upsert_docs_batch(
    self,
    docs: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...] = (),
    **kwargs: Any,
) -> None:
    """Upsert a batch of documents using AQL.

    Performs an upsert operation on a batch of documents, using the specified
    match keys to determine whether to update existing documents or insert new ones.

    Args:
        docs: List of documents to upsert
        class_name: Collection name to upsert into
        match_keys: Keys to match for upsert operation
        **kwargs: Additional options:
            - dry: If True, don't execute the query
            - update_keys: Keys to update on match
            - filter_uniques: If True, filter duplicate documents
    """
    dry = kwargs.pop("dry", False)
    update_keys = kwargs.pop("update_keys", None)
    filter_uniques = kwargs.pop("filter_uniques", True)

    if not docs:
        return
    if filter_uniques:
        docs = pick_unique_dict(docs)
    docs_json = json.dumps(docs, default=json_serializer)
    if not match_keys:
        upsert_clause = ""
        update_clause = ""
    else:
        upsert_clause = ", ".join([f'"{k}": doc.{k}' for k in match_keys])
        upsert_clause = f"UPSERT {{{upsert_clause}}}"

        if isinstance(update_keys, list):
            update_clause = ", ".join([f'"{k}": doc.{k}' for k in update_keys])
            update_clause = f"{{{update_clause}}}"
        elif update_keys == "doc":
            update_clause = "doc"
        else:
            update_clause = "{}"
        update_clause = f"UPDATE {update_clause}"

    options = "OPTIONS {exclusive: true, ignoreErrors: true}"

    q_update = f"""FOR doc in {docs_json}
                        {upsert_clause}
                        INSERT doc
                        {update_clause} 
                            IN {class_name} {options}"""
    if not dry:
        self.execute(q_update)

Connection

Bases: ABC

Abstract base class for database connections.

This class defines the interface that all database connection implementations must follow. It provides methods for database/graph operations, graph structure management (vertex types, edge types), and data manipulation.

Note

All methods marked with @abc.abstractmethod must be implemented by concrete connection classes. Subclasses must set the class attribute flavor to their DBType.

Source code in graflo/db/conn.py

class Connection(abc.ABC):
    """Abstract base class for database connections.

    This class defines the interface that all database connection implementations
    must follow. It provides methods for database/graph operations, graph structure
    management (vertex types, edge types), and data manipulation.

    Note:
        All methods marked with @abc.abstractmethod must be implemented by
        concrete connection classes. Subclasses must set the class attribute
        `flavor` to their DBType.
    """

    flavor: ClassVar[DBType] = DBType.ARANGO  # Overridden by subclasses

    def __init__(self):
        """Initialize the connection."""
        pass

    @classmethod
    def expression_flavor(cls) -> ExpressionFlavor:
        """Expression flavor for filter rendering (AQL, CYPHER, GSQL).

        Graph connection subclasses must set class attribute `flavor` to a
        DBType present in DB_TYPE_TO_EXPRESSION_FLAVOR.
        """
        return DB_TYPE_TO_EXPRESSION_FLAVOR[cls.flavor]

    @abc.abstractmethod
    def create_database(self, name: str):
        """Create a new database.

        Args:
            name: Name of the database to create
        """
        pass

    @abc.abstractmethod
    def delete_database(self, name: str):
        """Delete a database.

        Args:
            name: Name of the database to delete
        """
        pass

    @abc.abstractmethod
    def execute(self, query: str | Any, **kwargs: Any) -> Any:
        """Execute a database query.

        Args:
            query: Query to execute
            **kwargs: Additional query parameters

        Returns:
            Query result (database-specific)
        """
        pass

    @abc.abstractmethod
    def close(self):
        """Close the database connection."""
        pass

    def define_indexes(self, schema: Schema):
        """Define indexes for vertices and edges in the schema.

        Args:
            schema: Schema containing vertex and edge configurations
        """
        self.define_vertex_indices(schema.vertex_config)
        self.define_edge_indices(schema.edge_config.edges_list(include_aux=True))

    @abc.abstractmethod
    def define_schema(self, schema: Schema):
        """Define vertex and edge classes based on the schema.

        Args:
            schema: Schema containing vertex and edge class definitions
        """
        pass

    @abc.abstractmethod
    def delete_graph_structure(
        self,
        vertex_types: tuple[str, ...] | list[str] = (),
        graph_names: tuple[str, ...] | list[str] = (),
        delete_all: bool = False,
    ) -> None:
        """Delete graph structure (graphs, vertex types, edge types) from the database.

        This method deletes graphs and their associated vertex/edge types.
        The exact behavior depends on the database implementation:

        - ArangoDB: Deletes graphs and collections (vertex/edge collections)
        - Neo4j: Deletes nodes from labels (vertex types) and relationships
        - TigerGraph: Deletes graphs, vertex types, edge types, and jobs

        Args:
            vertex_types: Vertex type names to delete (database-specific interpretation)
            graph_names: Graph/database names to delete
            delete_all: If True, delete all graphs and their associated structures
        """
        pass

    @abc.abstractmethod
    def init_db(self, schema: Schema, recreate_schema: bool) -> None:
        """Initialize the database with the given schema.

        If the schema/graph already exists and recreate_schema is False, raises
        SchemaExistsError and the script halts.

        Args:
            schema: Schema to initialize the database with
            recreate_schema: If True, drop existing schema and define new one.
                If False and schema/graph already exists, raises SchemaExistsError.
        """
        pass

    @abc.abstractmethod
    def clear_data(self, schema: Schema) -> None:
        """Remove all data from the graph without dropping or changing the schema.

        Args:
            schema: Schema describing the graph (used to identify collections/labels).
        """
        pass

    @abc.abstractmethod
    def upsert_docs_batch(
        self,
        docs: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        **kwargs: Any,
    ) -> None:
        """Upsert a batch of documents.

        Args:
            docs: Documents to upsert
            class_name: Name of the vertex type (or collection/label in database-specific terms)
            match_keys: Keys to match for upsert
            **kwargs: Additional upsert parameters
        """
        pass

    @abc.abstractmethod
    def insert_edges_batch(
        self,
        docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
        source_class: str,
        target_class: str,
        relation_name: str,
        match_keys_source: tuple[str, ...],
        match_keys_target: tuple[str, ...],
        filter_uniques: bool = True,
        head: int | None = None,
        **kwargs: Any,
    ) -> None:
        """Insert a batch of edges.

        Args:
            docs_edges: Edge documents to insert
            source_class: Source vertex type/class
            target_class: Target vertex type/class
            relation_name: Name of the edge type/relation
            match_keys_source: Keys to match source vertices
            match_keys_target: Keys to match target vertices
            filter_uniques: Whether to filter unique edges
            head: Optional limit on number of edges to insert
            **kwargs: Additional insertion parameters, including:
                - collection_name: Name of the edge type (database-specific: collection/relationship type).
                  Required for ArangoDB (defaults to {source_class}_{target_class}_edges if not provided),
                  optional for other databases.
                - uniq_weight_fields: Fields to consider for uniqueness (ArangoDB-specific)
                - uniq_weight_collections: Vertex/edge types to consider for uniqueness (ArangoDB-specific)
                - upsert_option: Whether to upsert existing edges (ArangoDB-specific)
        """
        pass

    @abc.abstractmethod
    def insert_return_batch(
        self, docs: list[dict[str, Any]], class_name: str
    ) -> list[dict[str, Any]] | str:
        """Insert documents and return the inserted documents.

        Args:
            docs: Documents to insert
            class_name: Name of the vertex type (or collection/label in database-specific terms)

        Returns:
            list | str: Inserted documents, or a query string (database-specific behavior).
                Most implementations return a list of inserted documents. ArangoDB returns
                an AQL query string for deferred execution.
        """
        pass

    @abc.abstractmethod
    def fetch_docs(
        self,
        class_name: str,
        filters: list[Any] | dict[str, Any] | None = None,
        limit: int | None = None,
        return_keys: list[str] | None = None,
        unset_keys: list[str] | None = None,
        **kwargs: Any,
    ) -> list[dict[str, Any]]:
        """Fetch documents from a vertex type.

        Args:
            class_name: Name of the vertex type (or collection/label in database-specific terms)
            filters: Query filters
            limit: Maximum number of documents to return
            return_keys: Keys to return
            unset_keys: Keys to unset
            **kwargs: Additional database-specific parameters (e.g., field_types for TigerGraph)

        Returns:
            list: Fetched documents
        """
        pass

    @abc.abstractmethod
    def fetch_edges(
        self,
        from_type: str,
        from_id: str,
        edge_type: str | None = None,
        to_type: str | None = None,
        to_id: str | None = None,
        filters: list[Any] | dict[str, Any] | None = None,
        limit: int | None = None,
        return_keys: list[str] | None = None,
        unset_keys: list[str] | None = None,
        **kwargs: Any,
    ) -> list[dict[str, Any]]:
        """Fetch edges from the database.

        Args:
            from_type: Source vertex type
            from_id: Source vertex ID (required)
            edge_type: Optional edge type to filter by
            to_type: Optional target vertex type to filter by
            to_id: Optional target vertex ID to filter by
            filters: Additional query filters
            limit: Maximum number of edges to return
            return_keys: Keys to return (projection)
            unset_keys: Keys to exclude (projection)
            **kwargs: Additional database-specific parameters

        Returns:
            list: List of fetched edges
        """
        pass

    @abc.abstractmethod
    def fetch_present_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        flatten: bool = False,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> list[dict[str, Any]] | dict[int, list[dict[str, Any]]]:
        """Fetch documents that exist in the database.

        Args:
            batch: Batch of documents to check
            class_name: Name of the collection
            match_keys: Keys to match
            keep_keys: Keys to keep in result
            flatten: Whether to flatten the result. If True, returns a flat list.
                If False, returns a dict mapping batch indices to matching documents.
            filters: Additional query filters

        Returns:
            list | dict: Documents that exist in the database. Returns a list if
                flatten=True, otherwise returns a dict mapping batch indices to documents.
        """
        pass

    @abc.abstractmethod
    def aggregate(
        self,
        class_name: str,
        aggregation_function: AggregationType,
        discriminant: str | None = None,
        aggregated_field: str | None = None,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> int | float | list[dict[str, Any]] | dict[str, int | float] | None:
        """Perform aggregation on a collection.

        Args:
            class_name: Name of the collection
            aggregation_function: Type of aggregation to perform
            discriminant: Field to group by
            aggregated_field: Field to aggregate
            filters: Query filters

        Returns:
            Aggregation results (type depends on aggregation function)
        """
        pass

    @abc.abstractmethod
    def keep_absent_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Keep documents that don't exist in the database.

        Args:
            batch: Batch of documents to check
            class_name: Name of the collection
            match_keys: Keys to match
            keep_keys: Keys to keep in result
            filters: Additional query filters

        Returns:
            list: Documents that don't exist in the database
        """
        pass

    @abc.abstractmethod
    def define_vertex_indices(self, vertex_config: VertexConfig):
        """Define indices for vertex classes.

        Args:
            vertex_config: Vertex configuration containing index definitions
        """
        pass

    @abc.abstractmethod
    def define_edge_indices(self, edges: list[Edge]):
        """Define indices for edge classes.

        Args:
            edges: List of edge configurations containing index definitions
        """
        pass

    def define_vertex_classes(self, schema: Schema) -> None:
        """Define vertex classes based on schema.

        This method is called from define_schema() to create vertex types/collections.
        Most implementations take a Schema. Some implementations (like TigerGraph)
        may override with a more specific signature (VertexConfig).

        Default implementation is a no-op. Override in subclasses as needed.

        Args:
            schema: Schema containing vertex definitions
        """
        pass

    def define_edge_classes(self, edges: list[Edge]) -> None:
        """Define edge classes based on edge configurations.

        This method is called from define_schema() to create edge types/collections.

        Default implementation is a no-op. Override in subclasses as needed.

        Args:
            edges: List of edge configurations to create
        """
        pass

__init__()

Initialize the connection.

Source code in graflo/db/conn.py

def __init__(self):
    """Initialize the connection."""
    pass

aggregate(class_name, aggregation_function, discriminant=None, aggregated_field=None, filters=None) abstractmethod

Perform aggregation on a collection.

Parameters:

Name	Type	Description	Default
class_name	str	Name of the collection	required
aggregation_function	AggregationType	Type of aggregation to perform	required
discriminant	str | None	Field to group by	None
aggregated_field	str | None	Field to aggregate	None
filters	list[Any] | dict[str, Any] | None	Query filters	None

Returns:

Type	Description
int | float | list[dict[str, Any]] | dict[str, int | float] | None	Aggregation results (type depends on aggregation function)

Source code in graflo/db/conn.py

@abc.abstractmethod
def aggregate(
    self,
    class_name: str,
    aggregation_function: AggregationType,
    discriminant: str | None = None,
    aggregated_field: str | None = None,
    filters: list[Any] | dict[str, Any] | None = None,
) -> int | float | list[dict[str, Any]] | dict[str, int | float] | None:
    """Perform aggregation on a collection.

    Args:
        class_name: Name of the collection
        aggregation_function: Type of aggregation to perform
        discriminant: Field to group by
        aggregated_field: Field to aggregate
        filters: Query filters

    Returns:
        Aggregation results (type depends on aggregation function)
    """
    pass

clear_data(schema) abstractmethod

Remove all data from the graph without dropping or changing the schema.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema describing the graph (used to identify collections/labels).	required

Source code in graflo/db/conn.py

@abc.abstractmethod
def clear_data(self, schema: Schema) -> None:
    """Remove all data from the graph without dropping or changing the schema.

    Args:
        schema: Schema describing the graph (used to identify collections/labels).
    """
    pass

close() abstractmethod

Close the database connection.

Source code in graflo/db/conn.py

@abc.abstractmethod
def close(self):
    """Close the database connection."""
    pass

create_database(name) abstractmethod

Create a new database.

Parameters:

Name	Type	Description	Default
name	str	Name of the database to create	required

Source code in graflo/db/conn.py

@abc.abstractmethod
def create_database(self, name: str):
    """Create a new database.

    Args:
        name: Name of the database to create
    """
    pass

define_edge_classes(edges)

Define edge classes based on edge configurations.

This method is called from define_schema() to create edge types/collections.

Default implementation is a no-op. Override in subclasses as needed.

Parameters:

Name	Type	Description	Default
edges	list[Edge]	List of edge configurations to create	required

Source code in graflo/db/conn.py

def define_edge_classes(self, edges: list[Edge]) -> None:
    """Define edge classes based on edge configurations.

    This method is called from define_schema() to create edge types/collections.

    Default implementation is a no-op. Override in subclasses as needed.

    Args:
        edges: List of edge configurations to create
    """
    pass

define_edge_indices(edges) abstractmethod

Define indices for edge classes.

Parameters:

Name	Type	Description	Default
edges	list[Edge]	List of edge configurations containing index definitions	required

Source code in graflo/db/conn.py

@abc.abstractmethod
def define_edge_indices(self, edges: list[Edge]):
    """Define indices for edge classes.

    Args:
        edges: List of edge configurations containing index definitions
    """
    pass

define_indexes(schema)

Define indexes for vertices and edges in the schema.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing vertex and edge configurations	required

Source code in graflo/db/conn.py

def define_indexes(self, schema: Schema):
    """Define indexes for vertices and edges in the schema.

    Args:
        schema: Schema containing vertex and edge configurations
    """
    self.define_vertex_indices(schema.vertex_config)
    self.define_edge_indices(schema.edge_config.edges_list(include_aux=True))

define_schema(schema) abstractmethod

Define vertex and edge classes based on the schema.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing vertex and edge class definitions	required

Source code in graflo/db/conn.py

@abc.abstractmethod
def define_schema(self, schema: Schema):
    """Define vertex and edge classes based on the schema.

    Args:
        schema: Schema containing vertex and edge class definitions
    """
    pass

define_vertex_classes(schema)

Define vertex classes based on schema.

This method is called from define_schema() to create vertex types/collections. Most implementations take a Schema. Some implementations (like TigerGraph) may override with a more specific signature (VertexConfig).

Default implementation is a no-op. Override in subclasses as needed.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing vertex definitions	required

Source code in graflo/db/conn.py

def define_vertex_classes(self, schema: Schema) -> None:
    """Define vertex classes based on schema.

    This method is called from define_schema() to create vertex types/collections.
    Most implementations take a Schema. Some implementations (like TigerGraph)
    may override with a more specific signature (VertexConfig).

    Default implementation is a no-op. Override in subclasses as needed.

    Args:
        schema: Schema containing vertex definitions
    """
    pass

define_vertex_indices(vertex_config) abstractmethod

Define indices for vertex classes.

Parameters:

Name	Type	Description	Default
vertex_config	VertexConfig	Vertex configuration containing index definitions	required

Source code in graflo/db/conn.py

@abc.abstractmethod
def define_vertex_indices(self, vertex_config: VertexConfig):
    """Define indices for vertex classes.

    Args:
        vertex_config: Vertex configuration containing index definitions
    """
    pass

delete_database(name) abstractmethod

Delete a database.

Parameters:

Name	Type	Description	Default
name	str	Name of the database to delete	required

Source code in graflo/db/conn.py

@abc.abstractmethod
def delete_database(self, name: str):
    """Delete a database.

    Args:
        name: Name of the database to delete
    """
    pass

delete_graph_structure(vertex_types=(), graph_names=(), delete_all=False) abstractmethod

Delete graph structure (graphs, vertex types, edge types) from the database.

This method deletes graphs and their associated vertex/edge types. The exact behavior depends on the database implementation:

• ArangoDB: Deletes graphs and collections (vertex/edge collections)
• Neo4j: Deletes nodes from labels (vertex types) and relationships
• TigerGraph: Deletes graphs, vertex types, edge types, and jobs

Parameters:

Name	Type	Description	Default
vertex_types	tuple[str, ...] | list[str]	Vertex type names to delete (database-specific interpretation)	()
graph_names	tuple[str, ...] | list[str]	Graph/database names to delete	()
delete_all	bool	If True, delete all graphs and their associated structures	False

Source code in graflo/db/conn.py

@abc.abstractmethod
def delete_graph_structure(
    self,
    vertex_types: tuple[str, ...] | list[str] = (),
    graph_names: tuple[str, ...] | list[str] = (),
    delete_all: bool = False,
) -> None:
    """Delete graph structure (graphs, vertex types, edge types) from the database.

    This method deletes graphs and their associated vertex/edge types.
    The exact behavior depends on the database implementation:

    - ArangoDB: Deletes graphs and collections (vertex/edge collections)
    - Neo4j: Deletes nodes from labels (vertex types) and relationships
    - TigerGraph: Deletes graphs, vertex types, edge types, and jobs

    Args:
        vertex_types: Vertex type names to delete (database-specific interpretation)
        graph_names: Graph/database names to delete
        delete_all: If True, delete all graphs and their associated structures
    """
    pass

execute(query, **kwargs) abstractmethod

Execute a database query.

Parameters:

Name	Type	Description	Default
query	str | Any	Query to execute	required
**kwargs	Any	Additional query parameters	{}

Returns:

Type	Description
Any	Query result (database-specific)

Source code in graflo/db/conn.py

@abc.abstractmethod
def execute(self, query: str | Any, **kwargs: Any) -> Any:
    """Execute a database query.

    Args:
        query: Query to execute
        **kwargs: Additional query parameters

    Returns:
        Query result (database-specific)
    """
    pass

expression_flavor() classmethod

Expression flavor for filter rendering (AQL, CYPHER, GSQL).

Graph connection subclasses must set class attribute flavor to a DBType present in DB_TYPE_TO_EXPRESSION_FLAVOR.

Source code in graflo/db/conn.py

@classmethod
def expression_flavor(cls) -> ExpressionFlavor:
    """Expression flavor for filter rendering (AQL, CYPHER, GSQL).

    Graph connection subclasses must set class attribute `flavor` to a
    DBType present in DB_TYPE_TO_EXPRESSION_FLAVOR.
    """
    return DB_TYPE_TO_EXPRESSION_FLAVOR[cls.flavor]

fetch_docs(class_name, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs) abstractmethod

Fetch documents from a vertex type.

Parameters:

Name	Type	Description	Default
class_name	str	Name of the vertex type (or collection/label in database-specific terms)	required
filters	list[Any] | dict[str, Any] | None	Query filters	None
limit	int | None	Maximum number of documents to return	None
return_keys	list[str] | None	Keys to return	None
unset_keys	list[str] | None	Keys to unset	None
**kwargs	Any	Additional database-specific parameters (e.g., field_types for TigerGraph)	{}

Returns:

Name	Type	Description
list	list[dict[str, Any]]	Fetched documents

Source code in graflo/db/conn.py

@abc.abstractmethod
def fetch_docs(
    self,
    class_name: str,
    filters: list[Any] | dict[str, Any] | None = None,
    limit: int | None = None,
    return_keys: list[str] | None = None,
    unset_keys: list[str] | None = None,
    **kwargs: Any,
) -> list[dict[str, Any]]:
    """Fetch documents from a vertex type.

    Args:
        class_name: Name of the vertex type (or collection/label in database-specific terms)
        filters: Query filters
        limit: Maximum number of documents to return
        return_keys: Keys to return
        unset_keys: Keys to unset
        **kwargs: Additional database-specific parameters (e.g., field_types for TigerGraph)

    Returns:
        list: Fetched documents
    """
    pass

fetch_edges(from_type, from_id, edge_type=None, to_type=None, to_id=None, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs) abstractmethod

Fetch edges from the database.

Parameters:

Name	Type	Description	Default
from_type	str	Source vertex type	required
from_id	str	Source vertex ID (required)	required
edge_type	str | None	Optional edge type to filter by	None
to_type	str | None	Optional target vertex type to filter by	None
to_id	str | None	Optional target vertex ID to filter by	None
filters	list[Any] | dict[str, Any] | None	Additional query filters	None
limit	int | None	Maximum number of edges to return	None
return_keys	list[str] | None	Keys to return (projection)	None
unset_keys	list[str] | None	Keys to exclude (projection)	None
**kwargs	Any	Additional database-specific parameters	{}

Returns:

Name	Type	Description
list	list[dict[str, Any]]	List of fetched edges

Source code in graflo/db/conn.py

@abc.abstractmethod
def fetch_edges(
    self,
    from_type: str,
    from_id: str,
    edge_type: str | None = None,
    to_type: str | None = None,
    to_id: str | None = None,
    filters: list[Any] | dict[str, Any] | None = None,
    limit: int | None = None,
    return_keys: list[str] | None = None,
    unset_keys: list[str] | None = None,
    **kwargs: Any,
) -> list[dict[str, Any]]:
    """Fetch edges from the database.

    Args:
        from_type: Source vertex type
        from_id: Source vertex ID (required)
        edge_type: Optional edge type to filter by
        to_type: Optional target vertex type to filter by
        to_id: Optional target vertex ID to filter by
        filters: Additional query filters
        limit: Maximum number of edges to return
        return_keys: Keys to return (projection)
        unset_keys: Keys to exclude (projection)
        **kwargs: Additional database-specific parameters

    Returns:
        list: List of fetched edges
    """
    pass

fetch_present_documents(batch, class_name, match_keys, keep_keys=None, flatten=False, filters=None) abstractmethod

Fetch documents that exist in the database.

Parameters:

Name	Type	Description	Default
batch	list[dict[str, Any]]	Batch of documents to check	required
class_name	str	Name of the collection	required
match_keys	list[str] | tuple[str, ...]	Keys to match	required
keep_keys	list[str] | tuple[str, ...] | None	Keys to keep in result	None
flatten	bool	Whether to flatten the result. If True, returns a flat list. If False, returns a dict mapping batch indices to matching documents.	False
filters	list[Any] | dict[str, Any] | None	Additional query filters	None

Returns:

Type	Description
list[dict[str, Any]] | dict[int, list[dict[str, Any]]]	list | dict: Documents that exist in the database. Returns a list if flatten=True, otherwise returns a dict mapping batch indices to documents.

Source code in graflo/db/conn.py

@abc.abstractmethod
def fetch_present_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    flatten: bool = False,
    filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]] | dict[int, list[dict[str, Any]]]:
    """Fetch documents that exist in the database.

    Args:
        batch: Batch of documents to check
        class_name: Name of the collection
        match_keys: Keys to match
        keep_keys: Keys to keep in result
        flatten: Whether to flatten the result. If True, returns a flat list.
            If False, returns a dict mapping batch indices to matching documents.
        filters: Additional query filters

    Returns:
        list | dict: Documents that exist in the database. Returns a list if
            flatten=True, otherwise returns a dict mapping batch indices to documents.
    """
    pass

init_db(schema, recreate_schema) abstractmethod

Initialize the database with the given schema.

If the schema/graph already exists and recreate_schema is False, raises SchemaExistsError and the script halts.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema to initialize the database with	required
recreate_schema	bool	If True, drop existing schema and define new one. If False and schema/graph already exists, raises SchemaExistsError.	required

Source code in graflo/db/conn.py

@abc.abstractmethod
def init_db(self, schema: Schema, recreate_schema: bool) -> None:
    """Initialize the database with the given schema.

    If the schema/graph already exists and recreate_schema is False, raises
    SchemaExistsError and the script halts.

    Args:
        schema: Schema to initialize the database with
        recreate_schema: If True, drop existing schema and define new one.
            If False and schema/graph already exists, raises SchemaExistsError.
    """
    pass

insert_edges_batch(docs_edges, source_class, target_class, relation_name, match_keys_source, match_keys_target, filter_uniques=True, head=None, **kwargs) abstractmethod

Insert a batch of edges.

Parameters:

Name	Type	Description	Default
docs_edges	list[list[dict[str, Any]]] | list[Any] | None	Edge documents to insert	required
source_class	str	Source vertex type/class	required
target_class	str	Target vertex type/class	required
relation_name	str	Name of the edge type/relation	required
match_keys_source	tuple[str, ...]	Keys to match source vertices	required
match_keys_target	tuple[str, ...]	Keys to match target vertices	required
filter_uniques	bool	Whether to filter unique edges	True
head	int | None	Optional limit on number of edges to insert	None
**kwargs	Any	Additional insertion parameters, including: - collection_name: Name of the edge type (database-specific: collection/relationship type). Required for ArangoDB (defaults to {source_class}_{target_class}_edges if not provided), optional for other databases. - uniq_weight_fields: Fields to consider for uniqueness (ArangoDB-specific) - uniq_weight_collections: Vertex/edge types to consider for uniqueness (ArangoDB-specific) - upsert_option: Whether to upsert existing edges (ArangoDB-specific)	{}

Source code in graflo/db/conn.py

@abc.abstractmethod
def insert_edges_batch(
    self,
    docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
    source_class: str,
    target_class: str,
    relation_name: str,
    match_keys_source: tuple[str, ...],
    match_keys_target: tuple[str, ...],
    filter_uniques: bool = True,
    head: int | None = None,
    **kwargs: Any,
) -> None:
    """Insert a batch of edges.

    Args:
        docs_edges: Edge documents to insert
        source_class: Source vertex type/class
        target_class: Target vertex type/class
        relation_name: Name of the edge type/relation
        match_keys_source: Keys to match source vertices
        match_keys_target: Keys to match target vertices
        filter_uniques: Whether to filter unique edges
        head: Optional limit on number of edges to insert
        **kwargs: Additional insertion parameters, including:
            - collection_name: Name of the edge type (database-specific: collection/relationship type).
              Required for ArangoDB (defaults to {source_class}_{target_class}_edges if not provided),
              optional for other databases.
            - uniq_weight_fields: Fields to consider for uniqueness (ArangoDB-specific)
            - uniq_weight_collections: Vertex/edge types to consider for uniqueness (ArangoDB-specific)
            - upsert_option: Whether to upsert existing edges (ArangoDB-specific)
    """
    pass

insert_return_batch(docs, class_name) abstractmethod

Insert documents and return the inserted documents.

Parameters:

Name	Type	Description	Default
docs	list[dict[str, Any]]	Documents to insert	required
class_name	str	Name of the vertex type (or collection/label in database-specific terms)	required

Returns:

Type	Description
list[dict[str, Any]] | str	list | str: Inserted documents, or a query string (database-specific behavior). Most implementations return a list of inserted documents. ArangoDB returns an AQL query string for deferred execution.

Source code in graflo/db/conn.py

@abc.abstractmethod
def insert_return_batch(
    self, docs: list[dict[str, Any]], class_name: str
) -> list[dict[str, Any]] | str:
    """Insert documents and return the inserted documents.

    Args:
        docs: Documents to insert
        class_name: Name of the vertex type (or collection/label in database-specific terms)

    Returns:
        list | str: Inserted documents, or a query string (database-specific behavior).
            Most implementations return a list of inserted documents. ArangoDB returns
            an AQL query string for deferred execution.
    """
    pass

keep_absent_documents(batch, class_name, match_keys, keep_keys=None, filters=None) abstractmethod

Keep documents that don't exist in the database.

Parameters:

Name	Type	Description	Default
batch	list[dict[str, Any]]	Batch of documents to check	required
class_name	str	Name of the collection	required
match_keys	list[str] | tuple[str, ...]	Keys to match	required
keep_keys	list[str] | tuple[str, ...] | None	Keys to keep in result	None
filters	list[Any] | dict[str, Any] | None	Additional query filters	None

Returns:

Name	Type	Description
list	list[dict[str, Any]]	Documents that don't exist in the database

Source code in graflo/db/conn.py

@abc.abstractmethod
def keep_absent_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """Keep documents that don't exist in the database.

    Args:
        batch: Batch of documents to check
        class_name: Name of the collection
        match_keys: Keys to match
        keep_keys: Keys to keep in result
        filters: Additional query filters

    Returns:
        list: Documents that don't exist in the database
    """
    pass

upsert_docs_batch(docs, class_name, match_keys, **kwargs) abstractmethod

Upsert a batch of documents.

Parameters:

Name	Type	Description	Default
docs	list[dict[str, Any]]	Documents to upsert	required
class_name	str	Name of the vertex type (or collection/label in database-specific terms)	required
match_keys	list[str] | tuple[str, ...]	Keys to match for upsert	required
**kwargs	Any	Additional upsert parameters	{}

Source code in graflo/db/conn.py

@abc.abstractmethod
def upsert_docs_batch(
    self,
    docs: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    **kwargs: Any,
) -> None:
    """Upsert a batch of documents.

    Args:
        docs: Documents to upsert
        class_name: Name of the vertex type (or collection/label in database-specific terms)
        match_keys: Keys to match for upsert
        **kwargs: Additional upsert parameters
    """
    pass

ConnectionManager

Manager for database connections (both graph and source databases).

This class manages database connections to different database implementations. It provides a context manager interface for safe connection handling and automatic cleanup.

Supports: - Target databases (OUTPUT): ArangoDB, Neo4j, TigerGraph - Source databases (INPUT): PostgreSQL, MySQL, MongoDB, etc.

Attributes:

Name	Type	Description
target_conn_mapping		Mapping of target database types to connection classes
config	DBConfig	Connection configuration
working_db		Current working database name
conn		Active database connection

Source code in graflo/db/manager.py

class ConnectionManager:
    """Manager for database connections (both graph and source databases).

    This class manages database connections to different database
    implementations. It provides a context manager interface for safe
    connection handling and automatic cleanup.

    Supports:
    - Target databases (OUTPUT): ArangoDB, Neo4j, TigerGraph
    - Source databases (INPUT): PostgreSQL, MySQL, MongoDB, etc.

    Attributes:
        target_conn_mapping: Mapping of target database types to connection classes
        config: Connection configuration
        working_db: Current working database name
        conn: Active database connection
    """

    # Target database connections (OUTPUT)
    target_conn_mapping = {
        DBType.ARANGO: ArangoConnection,
        DBType.NEO4J: Neo4jConnection,
        DBType.TIGERGRAPH: TigerGraphConnection,
        DBType.FALKORDB: FalkordbConnection,
        DBType.MEMGRAPH: MemgraphConnection,
    }

    # Source database connections (INPUT) - to be implemented
    # source_conn_mapping = {
    #     DBType.POSTGRES: PostgresConnection,
    #     DBType.MYSQL: MySQLConnection,
    #     DBType.MONGODB: MongoDBConnection,
    # }

    def __init__(
        self,
        connection_config: DBConfig,
        **kwargs,
    ):
        """Initialize the connection manager.

        Args:
            connection_config: Database connection configuration
            **kwargs: Additional configuration parameters
        """
        self.config: DBConfig = connection_config
        self.working_db = kwargs.pop("working_db", None)
        self.conn = None

    def __enter__(self):
        """Enter the context manager.

        Creates and returns a new database connection.

        Returns:
            Connection: Database connection instance
        """
        # Check if database can be used as target
        if not self.config.can_be_target():
            raise ValueError(
                f"Database type '{self.config.connection_type}' cannot be used as a target. "
                f"Only these types can be targets: {[t.value for t in TARGET_DATABASES]}"
            )

        db_type = self.config.connection_type
        cls = self.target_conn_mapping[db_type]

        if self.working_db is not None:
            self.config.database = self.working_db
        self.conn = cls(config=self.config)
        return self.conn

    def close(self):
        """Close the database connection.

        Closes the active connection and performs any necessary cleanup.
        """
        if self.conn is not None:
            self.conn.close()

    def __exit__(self, exc_type, exc_value, exc_traceback):
        """Exit the context manager.

        Ensures the connection is properly closed when exiting the context.

        Args:
            exc_type: Exception type if an exception occurred
            exc_value: Exception value if an exception occurred
            exc_traceback: Exception traceback if an exception occurred
        """
        self.close()

__enter__()

Enter the context manager.

Creates and returns a new database connection.

Returns:

Name	Type	Description
Connection		Database connection instance

Source code in graflo/db/manager.py

def __enter__(self):
    """Enter the context manager.

    Creates and returns a new database connection.

    Returns:
        Connection: Database connection instance
    """
    # Check if database can be used as target
    if not self.config.can_be_target():
        raise ValueError(
            f"Database type '{self.config.connection_type}' cannot be used as a target. "
            f"Only these types can be targets: {[t.value for t in TARGET_DATABASES]}"
        )

    db_type = self.config.connection_type
    cls = self.target_conn_mapping[db_type]

    if self.working_db is not None:
        self.config.database = self.working_db
    self.conn = cls(config=self.config)
    return self.conn

__exit__(exc_type, exc_value, exc_traceback)

Exit the context manager.

Ensures the connection is properly closed when exiting the context.

Parameters:

Name	Type	Description	Default
exc_type		Exception type if an exception occurred	required
exc_value		Exception value if an exception occurred	required
exc_traceback		Exception traceback if an exception occurred	required

Source code in graflo/db/manager.py

def __exit__(self, exc_type, exc_value, exc_traceback):
    """Exit the context manager.

    Ensures the connection is properly closed when exiting the context.

    Args:
        exc_type: Exception type if an exception occurred
        exc_value: Exception value if an exception occurred
        exc_traceback: Exception traceback if an exception occurred
    """
    self.close()

__init__(connection_config, **kwargs)

Initialize the connection manager.

Parameters:

Name	Type	Description	Default
connection_config	DBConfig	Database connection configuration	required
**kwargs		Additional configuration parameters	{}

Source code in graflo/db/manager.py

def __init__(
    self,
    connection_config: DBConfig,
    **kwargs,
):
    """Initialize the connection manager.

    Args:
        connection_config: Database connection configuration
        **kwargs: Additional configuration parameters
    """
    self.config: DBConfig = connection_config
    self.working_db = kwargs.pop("working_db", None)
    self.conn = None

close()

Close the database connection.

Closes the active connection and performs any necessary cleanup.

Source code in graflo/db/manager.py

def close(self):
    """Close the database connection.

    Closes the active connection and performs any necessary cleanup.
    """
    if self.conn is not None:
        self.conn.close()

DBConfig

Bases: BaseSettings, ABC

Abstract base class for all database connection configurations using Pydantic BaseSettings.

Source code in graflo/db/connection/onto.py

class DBConfig(BaseSettings, abc.ABC):
    """Abstract base class for all database connection configurations using Pydantic BaseSettings."""

    uri: str | None = Field(default=None, description="Backend URI")
    username: str | None = Field(default=None, description="Authentication username")
    password: str | None = Field(default=None, description="Authentication Password")
    database: str | None = Field(
        default=None,
        description="Database name (backward compatibility, DB-specific mapping)",
    )
    schema_name: str | None = Field(
        default=None,
        validation_alias=AliasChoices("schema", "schema_name"),
        description="Schema/graph name (unified internal structure)",
    )
    request_timeout: float = Field(
        default=60.0, description="Request timeout in seconds"
    )

    @abc.abstractmethod
    def _get_default_port(self) -> int:
        """Get the default port for this db type."""
        pass

    @abc.abstractmethod
    def _get_effective_database(self) -> str | None:
        """Get the effective database name based on DB type.

        For SQL databases: returns the database name
        For graph databases: returns None (they don't have a database level)

        Returns:
            Database name or None
        """
        pass

    @abc.abstractmethod
    def _get_effective_schema(self) -> str | None:
        """Get the effective schema/graph name based on DB type.

        For SQL databases: returns the schema name
        For graph databases: returns the graph/database name (mapped from user-facing field)

        Returns:
            Schema/graph name or None
        """
        pass

    @property
    def effective_database(self) -> str | None:
        """Get the effective database name (delegates to concrete class)."""
        return self._get_effective_database()

    @property
    def effective_schema(self) -> str | None:
        """Get the effective schema/graph name (delegates to concrete class)."""
        return self._get_effective_schema()

    @model_validator(mode="after")
    def _normalize_uri(self):
        """Normalize URI: handle URIs without scheme and add default port if missing."""
        if self.uri is None:
            return self

        # Valid URL schemes (common database protocols)
        valid_schemes = {
            "http",
            "https",
            "bolt",
            "bolt+s",
            "bolt+ssc",
            "neo4j",
            "neo4j+s",
            "neo4j+ssc",
            "mongodb",
            "postgresql",
            "postgres",
            "mysql",
            "nebula",
            "redis",  # FalkorDB uses redis:// protocol
            "rediss",  # Redis with SSL
        }

        # Try to parse as-is first
        parsed = urlparse(self.uri)

        # Check if parsed scheme is actually a valid scheme or if it's a hostname
        # urlparse treats "localhost:14240" as scheme="localhost", path="14240"
        # We need to detect this case
        has_valid_scheme = parsed.scheme.lower() in valid_schemes
        has_netloc = bool(parsed.netloc)

        # If scheme doesn't look like a valid scheme and we have a colon, treat as host:port
        if not has_valid_scheme and ":" in self.uri and not self.uri.startswith("//"):
            # Check if it looks like host:port format
            parts = self.uri.split(":", 1)
            if len(parts) == 2:
                potential_host = parts[0]
                port_and_rest = parts[1]
                # Extract port (may have path/query after it)
                port_part = port_and_rest.split("/")[0].split("?")[0].split("#")[0]
                try:
                    # Validate port is numeric
                    int(port_part)
                    # If hostname doesn't look like a scheme (contains dots, is localhost, etc.)
                    # or if the parsed scheme is not in valid schemes, treat as host:port
                    if (
                        "." in potential_host
                        or potential_host.lower() in {"localhost", "127.0.0.1"}
                        or not has_valid_scheme
                    ):
                        # Reconstruct as proper URI with default scheme
                        default_scheme = "http"  # Default to http for most DBs
                        rest = port_and_rest[len(port_part) :]  # Everything after port
                        self.uri = (
                            f"{default_scheme}://{potential_host}:{port_part}{rest}"
                        )
                        parsed = urlparse(self.uri)
                except ValueError:
                    # Not a valid port, treat as regular URI - add scheme if needed
                    if not has_valid_scheme:
                        default_scheme = "http"
                        self.uri = f"{default_scheme}://{self.uri}"
                        parsed = urlparse(self.uri)
        elif not has_valid_scheme and not has_netloc:
            # No valid scheme and no netloc - add default scheme
            default_scheme = "http"
            self.uri = f"{default_scheme}://{self.uri}"
            parsed = urlparse(self.uri)

        # Add default port if missing
        if parsed.port is None:
            default_port = self._get_default_port()
            if parsed.scheme and parsed.hostname:
                # Reconstruct URI with port
                port_part = f":{default_port}" if default_port else ""
                path_part = parsed.path or ""
                query_part = f"?{parsed.query}" if parsed.query else ""
                fragment_part = f"#{parsed.fragment}" if parsed.fragment else ""
                self.uri = f"{parsed.scheme}://{parsed.hostname}{port_part}{path_part}{query_part}{fragment_part}"

        return self

    @model_validator(mode="after")
    def _extract_port_from_uri(self):
        """Extract port from URI and set it as gs_port for TigerGraph (if applicable).

        For TigerGraph 4+, gs_port is the primary port. If URI has a port but gs_port
        is not set, automatically extract and set gs_port from URI port.
        This simplifies configuration - users can just provide URI with port.
        """
        # Only apply to configs that have gs_port field (TigerGraph)
        if not hasattr(self, "gs_port"):
            return self

        if self.uri and self.gs_port is None:
            uri_port = self.port  # Get port from URI (property from base class)
            if uri_port:
                try:
                    self.gs_port = int(uri_port)
                    logger.debug(
                        f"Automatically set gs_port={self.gs_port} from URI port"
                    )
                except (ValueError, TypeError):
                    # Port couldn't be converted to int, skip auto-setting
                    pass

        return self

    @model_validator(mode="after")
    def _check_port_conflicts(self):
        """Check for port conflicts between URI and separate port fields.

        If port is provided both in URI and as a separate field, warn and prefer URI port.
        This ensures consistency and avoids confusion.
        """
        if self.uri is None:
            return self

        uri_port = self.port  # Get port from URI
        if uri_port is None:
            return self

        # Check for port fields in subclasses
        # Get model fields to check for port-related fields
        port_fields = []

        # Check for specific port fields that might exist in subclasses
        # Use getattr with None default to avoid AttributeError
        if hasattr(self, "gs_port"):
            gs_port_val = getattr(self, "gs_port", None)
            if gs_port_val is not None:
                port_fields.append(("gs_port", gs_port_val))

        if hasattr(self, "bolt_port"):
            bolt_port_val = getattr(self, "bolt_port", None)
            if bolt_port_val is not None:
                port_fields.append(("bolt_port", bolt_port_val))

        # Check each port field for conflicts
        port_conflicts = []
        for field_name, field_port in port_fields:
            # Compare as strings to handle int vs str differences
            if str(field_port) != str(uri_port):
                port_conflicts.append((field_name, field_port, uri_port))

        # Warn about conflicts and prefer URI port
        if port_conflicts:
            conflict_msgs = [
                f"{field_name}={field_port} (URI has port={uri_port})"
                for field_name, field_port, _ in port_conflicts
            ]
            warning_msg = (
                f"Port conflict detected: Port specified both in URI ({uri_port}) "
                f"and as separate field(s): {', '.join(conflict_msgs)}. "
                f"Using port from URI ({uri_port}). Consider removing the separate port field(s)."
            )
            warnings.warn(warning_msg, UserWarning, stacklevel=2)
            logger.warning(warning_msg)

            # Update port fields to match URI port (prefer URI)
            for field_name, _, _ in port_conflicts:
                try:
                    setattr(self, field_name, int(uri_port))
                except (ValueError, AttributeError):
                    # Field might be read-only or not settable, that's okay
                    pass

        return self

    @property
    def url(self) -> str | None:
        """Backward compatibility property: alias for uri."""
        return self.uri

    @property
    def url_without_port(self) -> str:
        """Get URL without port."""
        if self.uri is None:
            raise ValueError("URI is not set")
        parsed = urlparse(self.uri)
        return f"{parsed.scheme}://{parsed.hostname}"

    @property
    def port(self) -> str | None:
        """Get port from URI."""
        if self.uri is None:
            return None
        parsed = urlparse(self.uri)
        return str(parsed.port) if parsed.port else None

    @property
    def protocol(self) -> str:
        """Get protocol/scheme from URI."""
        if self.uri is None:
            return "http"
        parsed = urlparse(self.uri)
        return parsed.scheme or "http"

    @property
    def hostname(self) -> str | None:
        """Get hostname from URI."""
        if self.uri is None:
            return None
        parsed = urlparse(self.uri)
        return parsed.hostname

    @property
    def connection_type(self) -> "DBType":
        """Get database type from class."""
        # Map class to DBType - need to import here to avoid circular import
        from .config_mapping import DB_TYPE_MAPPING

        # Reverse lookup: find DBType for this class
        for db_type, config_class in DB_TYPE_MAPPING.items():
            if type(self) is config_class:
                return db_type

        # Fallback (shouldn't happen)
        return DBType.ARANGO

    def can_be_source(self) -> bool:
        """Check if this database type can be used as a source."""
        return self.connection_type in SOURCE_DATABASES

    def can_be_target(self) -> bool:
        """Check if this database type can be used as a target."""
        return self.connection_type in TARGET_DATABASES

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "DBConfig":
        """Create a connection config from a dictionary."""
        if not isinstance(data, dict):
            raise TypeError(f"Expected dict, got {type(data)}")

        # Copy the data to avoid modifying the original
        config_data = data.copy()

        db_type = config_data.pop("db_type", None) or config_data.pop(
            "connection_type", None
        )
        if not db_type:
            raise ValueError("Missing 'db_type' or 'connection_type' in configuration")

        try:
            conn_type = DBType(db_type)
        except ValueError:
            raise ValueError(
                f"Database type '{db_type}' not supported. "
                f"Should be one of: {list(DBType)}"
            )

        # Map old 'url' field to 'uri' for backward compatibility
        if "url" in config_data and "uri" not in config_data:
            config_data["uri"] = config_data.pop("url")

        # Map old credential fields
        if "cred_name" in config_data and "username" not in config_data:
            config_data["username"] = config_data.pop("cred_name")
        if "cred_pass" in config_data and "password" not in config_data:
            config_data["password"] = config_data.pop("cred_pass")

        # Construct URI from protocol/hostname/port if uri is not provided
        if "uri" not in config_data:
            protocol = config_data.pop("protocol", "http")
            hostname = config_data.pop("hostname", None)
            port = config_data.pop("port", None)
            hosts = config_data.pop("hosts", None)

            if hosts:
                # Use hosts as URI
                config_data["uri"] = hosts
            elif hostname:
                # Construct URI from components
                if port:
                    config_data["uri"] = f"{protocol}://{hostname}:{port}"
                else:
                    config_data["uri"] = f"{protocol}://{hostname}"

        # Get the appropriate config class and initialize it
        from .config_mapping import get_config_class

        config_class = get_config_class(conn_type)
        return config_class(**config_data)

    @classmethod
    def from_docker_env(cls, docker_dir: str | Path | None = None) -> "DBConfig":
        """Load config from docker .env file.

        Args:
            docker_dir: Path to docker directory. If None, uses default based on db type.

        Returns:
            DBConfig instance loaded from .env file
        """
        raise NotImplementedError("Subclasses must implement from_docker_env")

    @classmethod
    def from_env(
        cls: Type[T],
        *,
        prefix: str | None = None,
        profile: str | None = None,
        suffix: str | None = None,
    ) -> T:
        """Load config from environment variables using Pydantic BaseSettings.

        Supports qualifiers for multiple configs from the same env:

        - **prefix**: outer prefix → ``{prefix}_{BASE_PREFIX}URI`` (e.g. ``USER_ARANGO_URI``).
        - **profile**: segment after base → ``{BASE_PREFIX}{profile}_URI`` (e.g. ``ARANGO_DEV_URI``).
        - **suffix**: after field name → ``{BASE_PREFIX}URI_{suffix}`` (e.g. ``ARANGO_URI_DEV``).

        At most one of ``prefix``, ``profile``, ``suffix`` should be set.

        Args:
            prefix: Outer env prefix (e.g. ``"USER"`` → ``USER_ARANGO_URI``).
            profile: Env segment after base (e.g. ``"DEV"`` → ``ARANGO_DEV_URI``).
            suffix: Env segment after field name (e.g. ``"DEV"`` → ``ARANGO_URI_DEV``).

        Returns:
            DBConfig instance loaded from environment variables.

        Examples:
            # Default (ARANGO_URI, ARANGO_USERNAME, ...)
            config = ArangoConfig.from_env()

            # By profile: ARANGO_DEV_URI, ARANGO_DEV_USERNAME, ...
            dev = ArangoConfig.from_env(profile="DEV")

            # By suffix: ARANGO_URI_DEV, ARANGO_USERNAME_DEV, ...
            dev2 = ArangoConfig.from_env(suffix="DEV")

            # Outer prefix: USER_ARANGO_URI, ...
            user_config = ArangoConfig.from_env(prefix="USER")
        """
        base_prefix = cls.model_config.get("env_prefix")
        if not base_prefix:
            raise ValueError(
                f"Class {cls.__name__} does not have env_prefix configured in model_config"
            )
        case_sensitive = cls.model_config.get("case_sensitive", False)
        qualifiers = sum(1 for q in (prefix, profile, suffix) if q is not None)
        if qualifiers > 1:
            raise ValueError("At most one of prefix, profile, suffix may be set")

        if suffix:
            # Pydantic doesn't support env_suffix; read suffixed vars manually.
            data: Dict[str, Any] = {}
            suf = suffix if case_sensitive else suffix.upper()
            for name in cls.model_fields:
                env_name = f"{base_prefix}{name.upper()}_{suf}"
                if not case_sensitive:
                    # Match pydantic-settings: first try exact, then uppercase
                    val = os.environ.get(env_name) or os.environ.get(env_name.lower())
                else:
                    val = os.environ.get(env_name)
                if val is not None:
                    data[name] = val
            return cls(**data)

        if prefix:
            new_prefix = f"{prefix.upper()}_{base_prefix}"
        elif profile:
            new_prefix = f"{base_prefix}{profile.upper()}_"
        else:
            return cls()

        model_config = SettingsConfigDict(
            env_prefix=new_prefix,
            case_sensitive=case_sensitive,
        )
        temp_class = type(
            f"{cls.__name__}WithPrefix", (cls,), {"model_config": model_config}
        )
        return temp_class()

connection_type property

Get database type from class.

effective_database property

Get the effective database name (delegates to concrete class).

effective_schema property

Get the effective schema/graph name (delegates to concrete class).

hostname property

Get hostname from URI.

port property

Get port from URI.

protocol property

Get protocol/scheme from URI.

url property

Backward compatibility property: alias for uri.

url_without_port property

Get URL without port.

can_be_source()

Check if this database type can be used as a source.

Source code in graflo/db/connection/onto.py

def can_be_source(self) -> bool:
    """Check if this database type can be used as a source."""
    return self.connection_type in SOURCE_DATABASES

can_be_target()

Check if this database type can be used as a target.

Source code in graflo/db/connection/onto.py

def can_be_target(self) -> bool:
    """Check if this database type can be used as a target."""
    return self.connection_type in TARGET_DATABASES

from_dict(data) classmethod

Create a connection config from a dictionary.

Source code in graflo/db/connection/onto.py

@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "DBConfig":
    """Create a connection config from a dictionary."""
    if not isinstance(data, dict):
        raise TypeError(f"Expected dict, got {type(data)}")

    # Copy the data to avoid modifying the original
    config_data = data.copy()

    db_type = config_data.pop("db_type", None) or config_data.pop(
        "connection_type", None
    )
    if not db_type:
        raise ValueError("Missing 'db_type' or 'connection_type' in configuration")

    try:
        conn_type = DBType(db_type)
    except ValueError:
        raise ValueError(
            f"Database type '{db_type}' not supported. "
            f"Should be one of: {list(DBType)}"
        )

    # Map old 'url' field to 'uri' for backward compatibility
    if "url" in config_data and "uri" not in config_data:
        config_data["uri"] = config_data.pop("url")

    # Map old credential fields
    if "cred_name" in config_data and "username" not in config_data:
        config_data["username"] = config_data.pop("cred_name")
    if "cred_pass" in config_data and "password" not in config_data:
        config_data["password"] = config_data.pop("cred_pass")

    # Construct URI from protocol/hostname/port if uri is not provided
    if "uri" not in config_data:
        protocol = config_data.pop("protocol", "http")
        hostname = config_data.pop("hostname", None)
        port = config_data.pop("port", None)
        hosts = config_data.pop("hosts", None)

        if hosts:
            # Use hosts as URI
            config_data["uri"] = hosts
        elif hostname:
            # Construct URI from components
            if port:
                config_data["uri"] = f"{protocol}://{hostname}:{port}"
            else:
                config_data["uri"] = f"{protocol}://{hostname}"

    # Get the appropriate config class and initialize it
    from .config_mapping import get_config_class

    config_class = get_config_class(conn_type)
    return config_class(**config_data)

from_docker_env(docker_dir=None) classmethod

Load config from docker .env file.

Parameters:

Name	Type	Description	Default
docker_dir	str | Path | None	Path to docker directory. If None, uses default based on db type.	None

Returns:

Type	Description
DBConfig	DBConfig instance loaded from .env file

Source code in graflo/db/connection/onto.py

@classmethod
def from_docker_env(cls, docker_dir: str | Path | None = None) -> "DBConfig":
    """Load config from docker .env file.

    Args:
        docker_dir: Path to docker directory. If None, uses default based on db type.

    Returns:
        DBConfig instance loaded from .env file
    """
    raise NotImplementedError("Subclasses must implement from_docker_env")

from_env(*, prefix=None, profile=None, suffix=None) classmethod

Load config from environment variables using Pydantic BaseSettings.

Supports qualifiers for multiple configs from the same env:

• prefix: outer prefix → {prefix}_{BASE_PREFIX}URI (e.g. USER_ARANGO_URI).
• profile: segment after base → {BASE_PREFIX}{profile}_URI (e.g. ARANGO_DEV_URI).
• suffix: after field name → {BASE_PREFIX}URI_{suffix} (e.g. ARANGO_URI_DEV).

At most one of prefix, profile, suffix should be set.

Parameters:

Name	Type	Description	Default
prefix	str | None	Outer env prefix (e.g. "USER" → USER_ARANGO_URI).	None
profile	str | None	Env segment after base (e.g. "DEV" → ARANGO_DEV_URI).	None
suffix	str | None	Env segment after field name (e.g. "DEV" → ARANGO_URI_DEV).	None

Returns:

Type	Description
T	DBConfig instance loaded from environment variables.

Examples:

Default (ARANGO_URI, ARANGO_USERNAME, ...)

config = ArangoConfig.from_env()

By profile: ARANGO_DEV_URI, ARANGO_DEV_USERNAME, ...

dev = ArangoConfig.from_env(profile="DEV")

By suffix: ARANGO_URI_DEV, ARANGO_USERNAME_DEV, ...

dev2 = ArangoConfig.from_env(suffix="DEV")

Outer prefix: USER_ARANGO_URI, ...

user_config = ArangoConfig.from_env(prefix="USER")

Source code in graflo/db/connection/onto.py

@classmethod
def from_env(
    cls: Type[T],
    *,
    prefix: str | None = None,
    profile: str | None = None,
    suffix: str | None = None,
) -> T:
    """Load config from environment variables using Pydantic BaseSettings.

    Supports qualifiers for multiple configs from the same env:

    - **prefix**: outer prefix → ``{prefix}_{BASE_PREFIX}URI`` (e.g. ``USER_ARANGO_URI``).
    - **profile**: segment after base → ``{BASE_PREFIX}{profile}_URI`` (e.g. ``ARANGO_DEV_URI``).
    - **suffix**: after field name → ``{BASE_PREFIX}URI_{suffix}`` (e.g. ``ARANGO_URI_DEV``).

    At most one of ``prefix``, ``profile``, ``suffix`` should be set.

    Args:
        prefix: Outer env prefix (e.g. ``"USER"`` → ``USER_ARANGO_URI``).
        profile: Env segment after base (e.g. ``"DEV"`` → ``ARANGO_DEV_URI``).
        suffix: Env segment after field name (e.g. ``"DEV"`` → ``ARANGO_URI_DEV``).

    Returns:
        DBConfig instance loaded from environment variables.

    Examples:
        # Default (ARANGO_URI, ARANGO_USERNAME, ...)
        config = ArangoConfig.from_env()

        # By profile: ARANGO_DEV_URI, ARANGO_DEV_USERNAME, ...
        dev = ArangoConfig.from_env(profile="DEV")

        # By suffix: ARANGO_URI_DEV, ARANGO_USERNAME_DEV, ...
        dev2 = ArangoConfig.from_env(suffix="DEV")

        # Outer prefix: USER_ARANGO_URI, ...
        user_config = ArangoConfig.from_env(prefix="USER")
    """
    base_prefix = cls.model_config.get("env_prefix")
    if not base_prefix:
        raise ValueError(
            f"Class {cls.__name__} does not have env_prefix configured in model_config"
        )
    case_sensitive = cls.model_config.get("case_sensitive", False)
    qualifiers = sum(1 for q in (prefix, profile, suffix) if q is not None)
    if qualifiers > 1:
        raise ValueError("At most one of prefix, profile, suffix may be set")

    if suffix:
        # Pydantic doesn't support env_suffix; read suffixed vars manually.
        data: Dict[str, Any] = {}
        suf = suffix if case_sensitive else suffix.upper()
        for name in cls.model_fields:
            env_name = f"{base_prefix}{name.upper()}_{suf}"
            if not case_sensitive:
                # Match pydantic-settings: first try exact, then uppercase
                val = os.environ.get(env_name) or os.environ.get(env_name.lower())
            else:
                val = os.environ.get(env_name)
            if val is not None:
                data[name] = val
        return cls(**data)

    if prefix:
        new_prefix = f"{prefix.upper()}_{base_prefix}"
    elif profile:
        new_prefix = f"{base_prefix}{profile.upper()}_"
    else:
        return cls()

    model_config = SettingsConfigDict(
        env_prefix=new_prefix,
        case_sensitive=case_sensitive,
    )
    temp_class = type(
        f"{cls.__name__}WithPrefix", (cls,), {"model_config": model_config}
    )
    return temp_class()

FalkordbConnection

Bases: Connection

FalkorDB-specific implementation of the Connection interface.

This class provides FalkorDB-specific implementations for all database operations, including node management, relationship operations, and Cypher query execution. It uses the FalkorDB Python client for all operations.

Attributes:

Name	Type	Description
flavor		Database flavor identifier (FALKORDB)
config		FalkorDB connection configuration (URI, database, credentials)
client	FalkorDB | None	Underlying FalkorDB client instance
graph	Graph | None	Active graph object for query execution
_graph_name	str	Name of the currently selected graph

Source code in graflo/db/falkordb/conn.py

class FalkordbConnection(Connection):
    """FalkorDB-specific implementation of the Connection interface.

    This class provides FalkorDB-specific implementations for all database
    operations, including node management, relationship operations, and
    Cypher query execution. It uses the FalkorDB Python client for all operations.

    Attributes:
        flavor: Database flavor identifier (FALKORDB)
        config: FalkorDB connection configuration (URI, database, credentials)
        client: Underlying FalkorDB client instance
        graph: Active graph object for query execution
        _graph_name: Name of the currently selected graph
    """

    flavor = DBType.FALKORDB

    # Type annotations for instance attributes
    client: FalkorDB | None
    graph: Graph | None
    _graph_name: str

    def __init__(self, config: FalkordbConfig):
        """Initialize FalkorDB connection.

        Args:
            config: FalkorDB connection configuration containing URI, database,
                and optional password
        """
        super().__init__()
        self.config = config

        if config.uri is None:
            raise ValueError("FalkorDB connection requires a URI to be configured")

        # Parse URI to extract host and port
        parsed = urlparse(config.uri)
        host = parsed.hostname or "localhost"
        port = parsed.port or 6379

        # Initialize FalkorDB client
        if config.password:
            self.client = FalkorDB(host=host, port=port, password=config.password)
        else:
            self.client = FalkorDB(host=host, port=port)

        # Select the graph (database in config maps to graph name)
        graph_name = config.database or "default"
        self.graph = self.client.select_graph(graph_name)
        self._graph_name = graph_name

    def execute(self, query: str, **kwargs):
        """Execute a raw OpenCypher query against the graph.

        Args:
            query: OpenCypher query string. Can include parameter placeholders
                using $name syntax (e.g., "MATCH (n) WHERE n.id = $id")
            **kwargs: Query parameters as keyword arguments

        Returns:
            QueryResult: FalkorDB result object containing result_set and statistics
        """
        assert self.graph is not None, "Connection is closed"
        # Pass params as keyword argument if the client supports it, otherwise as positional
        # Try params keyword first, fall back to positional
        if kwargs:
            try:
                result = self.graph.query(query, params=kwargs)
            except TypeError:
                # Fall back to positional argument if params keyword not supported
                result = self.graph.query(query, kwargs)
        else:
            result = self.graph.query(query)
        return result

    def close(self):
        """Close the FalkorDB connection."""
        # FalkorDB client handles connection pooling internally
        # No explicit close needed, but we can delete the reference
        self.graph = None
        self.client = None

    @staticmethod
    def _is_valid_property_value(value) -> bool:
        """Validate that a value can be stored as a FalkorDB property.

        Rejects NaN and infinity values that cannot be stored.

        Args:
            value: Value to validate

        Returns:
            True if value can be safely stored, False otherwise
        """
        import math

        if isinstance(value, float):
            if math.isnan(value) or math.isinf(value):
                return False
        return True

    @staticmethod
    def _sanitize_string_value(value: str) -> str:
        """Remove characters that break the Cypher parser.

        Args:
            value: String value to sanitize

        Returns:
            Sanitized string with problematic characters removed
        """
        if "\x00" in value:
            value = value.replace("\x00", "")
        return value

    def _sanitize_value(self, value):
        """Recursively sanitize a value, handling nested structures.

        Args:
            value: Value to sanitize (can be dict, list, or primitive)

        Returns:
            Sanitized value with datetime objects serialized to epoch microseconds
        """
        # Handle nested dictionaries recursively
        if isinstance(value, dict):
            sanitized_dict = {}
            for k, v in value.items():
                # Filter non-string keys in nested dicts too
                if not isinstance(k, str):
                    logger.warning(
                        f"Skipping non-string nested key: {k!r} (type: {type(k).__name__})"
                    )
                    continue
                sanitized_v = self._sanitize_value(v)
                # Check for invalid float values
                if self._is_valid_property_value(sanitized_v):
                    sanitized_dict[k] = sanitized_v
            return sanitized_dict

        # Handle lists recursively
        elif isinstance(value, list):
            sanitized_list = []
            for item in value:
                sanitized_item = self._sanitize_value(item)
                # Check for invalid float values
                if self._is_valid_property_value(sanitized_item):
                    sanitized_list.append(sanitized_item)
            return sanitized_list

        # Handle primitive values
        else:
            # Convert datetime objects to ISO-8601 strings for FalkorDB
            # These will be wrapped with datetime() function in Cypher queries
            from datetime import date, datetime, time

            if isinstance(value, (datetime, date, time)):
                # Use ISO format - will be wrapped with datetime() in Cypher
                serialized = serialize_value(value)  # This returns ISO-8601 string
            else:
                # Use shared serialize_value for other types (Decimal, etc.)
                serialized = serialize_value(value)

            # Sanitize string values (remove null bytes that break Cypher)
            if isinstance(serialized, str):
                serialized = self._sanitize_string_value(serialized)

            return serialized

    def _sanitize_document(
        self, doc: dict, match_keys: list[str] | None = None
    ) -> dict:
        """Sanitize a document for safe FalkorDB insertion.

        Filters invalid keys/values, serializes datetime objects (including nested ones),
        and validates required match keys.

        Args:
            doc: Document to sanitize
            match_keys: Optional list of keys that must be present with valid values

        Returns:
            Sanitized copy of the document

        Raises:
            ValueError: If a required match_key is missing or has None value
        """
        sanitized = {}

        for key, value in doc.items():
            # Filter non-string keys
            if not isinstance(key, str):
                logger.warning(
                    f"Skipping non-string property key: {key!r} (type: {type(key).__name__})"
                )
                continue

            # Recursively sanitize the value (handles nested dicts/lists and datetime objects)
            sanitized_value = self._sanitize_value(value)

            # Check for invalid float values
            if not self._is_valid_property_value(sanitized_value):
                logger.warning(
                    f"Skipping property '{key}' with invalid value: {sanitized_value}"
                )
                continue

            sanitized[key] = sanitized_value

        # Validate match_keys presence
        if match_keys:
            for key in match_keys:
                if key not in sanitized:
                    raise ValueError(
                        f"Required match key '{key}' is missing or has invalid value in document: {doc}"
                    )
                if sanitized[key] is None:
                    raise ValueError(
                        f"Match key '{key}' cannot be None in document: {doc}"
                    )

        return sanitized

    def _sanitize_batch(
        self, docs: list[dict], match_keys: list[str] | None = None
    ) -> list[dict]:
        """Sanitize a batch of documents.

        Args:
            docs: List of documents to sanitize
            match_keys: Optional list of required keys to validate

        Returns:
            List of sanitized documents
        """
        return [self._sanitize_document(doc, match_keys) for doc in docs]

    def create_database(self, name: str):
        """Create a new graph in FalkorDB.

        Note: In FalkorDB, graphs are created implicitly when data is first inserted.

        Args:
            name: Name of the graph to create
        """
        # In FalkorDB, graphs are created implicitly when you first insert data
        # We just need to select the graph
        assert self.client is not None, "Connection is closed"
        self.graph = self.client.select_graph(name)
        self._graph_name = name
        logger.info(f"Selected FalkorDB graph '{name}'")

    def delete_database(self, name: str):
        """Delete a graph from FalkorDB.

        Args:
            name: Name of the graph to delete (if empty, uses current graph)
        """
        graph_to_delete = name if name else self._graph_name
        assert self.client is not None, "Connection is closed"
        try:
            # Delete the graph using the FalkorDB API
            graph = self.client.select_graph(graph_to_delete)
            graph.delete()
            logger.info(f"Successfully deleted FalkorDB graph '{graph_to_delete}'")
        except Exception as e:
            logger.error(
                f"Failed to delete FalkorDB graph '{graph_to_delete}': {e}",
                exc_info=True,
            )
            raise

    def define_vertex_indices(self, vertex_config: VertexConfig):
        """Define indices for vertex labels.

        Creates indices for each vertex label based on the configuration.
        FalkorDB supports range indices on node properties.

        Args:
            vertex_config: Vertex configuration containing index definitions
        """
        for c in vertex_config.vertex_set:
            for index_obj in vertex_config.indexes(c):
                self._add_index(c, index_obj)

    def define_edge_indices(self, edges: list[Edge]):
        """Define indices for relationship types.

        Creates indices for each relationship type based on the configuration.
        FalkorDB supports range indices on relationship properties.

        Args:
            edges: List of edge configurations containing index definitions
        """
        for edge in edges:
            for index_obj in edge.indexes:
                if edge.relation is not None:
                    self._add_index(edge.relation, index_obj, is_vertex_index=False)

    def _add_index(self, obj_name: str, index: Index, is_vertex_index: bool = True):
        """Add an index to a label or relationship type.

        Args:
            obj_name: Label or relationship type name
            index: Index configuration to create
            is_vertex_index: If True, create index on nodes, otherwise on relationships
        """
        for field in index.fields:
            try:
                if is_vertex_index:
                    # FalkorDB node index syntax
                    q = f"CREATE INDEX FOR (n:{obj_name}) ON (n.{field})"
                else:
                    # FalkorDB relationship index syntax
                    q = f"CREATE INDEX FOR ()-[r:{obj_name}]-() ON (r.{field})"

                self.execute(q)
                logger.debug(f"Created index on {obj_name}.{field}")
            except Exception as e:
                # Index may already exist, log and continue
                logger.debug(f"Index creation note for {obj_name}.{field}: {e}")

    def define_schema(self, schema: Schema):
        """Define vertex and edge classes based on schema.

        Note: This is a no-op in FalkorDB as vertex/edge classes (labels/relationship types) are implicit.

        Args:
            schema: Schema containing vertex and edge class definitions
        """
        pass

    def define_vertex_classes(self, schema: Schema):
        """Define vertex classes based on schema.

        Note: This is a no-op in FalkorDB as vertex classes (labels) are implicit.

        Args:
            schema: Schema containing vertex definitions
        """
        pass

    def define_edge_classes(self, edges: list[Edge]):
        """Define edge classes based on schema.

        Note: This is a no-op in FalkorDB as edge classes (relationship types) are implicit.

        Args:
            edges: List of edge configurations
        """
        pass

    def delete_graph_structure(
        self,
        vertex_types: tuple[str, ...] | list[str] = (),
        graph_names: tuple[str, ...] | list[str] = (),
        delete_all: bool = False,
    ) -> None:
        """Delete graph structure (nodes and relationships) from FalkorDB.

        Args:
            vertex_types: Label names to delete nodes for
            graph_names: Graph names to delete entirely
            delete_all: If True, delete all nodes and relationships
        """
        if delete_all or (not vertex_types and not graph_names):
            # Delete all nodes and relationships in current graph
            try:
                self.execute("MATCH (n) DETACH DELETE n")
                logger.debug("Deleted all nodes and relationships from graph")
            except Exception as e:
                logger.debug(f"Graph may be empty or not exist: {e}")
        elif vertex_types:
            # Delete nodes with specific labels
            for label in vertex_types:
                try:
                    self.execute(f"MATCH (n:{label}) DETACH DELETE n")
                    logger.debug(f"Deleted all nodes with label '{label}'")
                except Exception as e:
                    logger.warning(f"Failed to delete nodes with label '{label}': {e}")

        # Delete specific graphs
        assert self.client is not None, "Connection is closed"
        for graph_name in graph_names:
            try:
                graph = self.client.select_graph(graph_name)
                graph.delete()
                logger.debug(f"Deleted graph '{graph_name}'")
            except Exception as e:
                logger.warning(f"Failed to delete graph '{graph_name}': {e}")

    def init_db(self, schema: Schema, recreate_schema: bool) -> None:
        """Initialize FalkorDB with the given schema.

        Uses schema.general.name if database is not set in config.

        If the graph already has nodes and recreate_schema is False, raises
        SchemaExistsError and the script halts.

        Args:
            schema: Schema containing graph structure definitions
            recreate_schema: If True, delete all existing data before initialization.
                If False and graph has nodes, raises SchemaExistsError.
        """
        # Determine graph name: use config.database if set, otherwise use schema.general.name
        graph_name = self.config.database
        if not graph_name:
            graph_name = schema.general.name
            self.config.database = graph_name

        # Select/create the graph
        assert self.client is not None, "Connection is closed"
        self.graph = self.client.select_graph(graph_name)
        self._graph_name = graph_name
        logger.info(f"Initialized FalkorDB graph '{graph_name}'")

        # Check if graph already has nodes (schema/graph exists)
        try:
            result = self.execute("MATCH (n) RETURN count(n) AS c")
            count = 0
            if hasattr(result, "data") and result.data():
                count = result.data()[0].get("c", 0) or 0
            elif result is not None and hasattr(result, "__iter__"):
                for record in result:
                    count = (
                        record.get("c", 0)
                        if hasattr(record, "get")
                        else getattr(record, "c", 0)
                    ) or 0
                    break
            if count > 0 and not recreate_schema:
                raise SchemaExistsError(
                    f"Schema/graph already exists in graph '{graph_name}' ({count} nodes). "
                    "Set recreate_schema=True to replace, or use clear_data=True before ingestion."
                )
        except SchemaExistsError:
            raise
        except Exception as e:
            logger.debug(f"Could not check graph node count: {e}")

        if recreate_schema:
            try:
                self.delete_graph_structure(delete_all=True)
                logger.debug(f"Cleaned graph '{graph_name}' for fresh start")
            except Exception as e:
                logger.debug(f"Recreate schema note for graph '{graph_name}': {e}")

        try:
            self.define_indexes(schema)
            logger.debug(f"Defined indexes for graph '{graph_name}'")
        except Exception as e:
            logger.error(
                f"Failed to define indexes for graph '{graph_name}': {e}",
                exc_info=True,
            )
            raise

    def clear_data(self, schema: Schema) -> None:
        """Remove all data from the graph without dropping the schema.

        Deletes all nodes and relationships; labels (schema) remain.
        """
        self.delete_graph_structure(delete_all=True)

    def upsert_docs_batch(
        self,
        docs: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        **kwargs: Any,
    ) -> None:
        """Upsert a batch of nodes using Cypher MERGE.

        Args:
            docs: List of node documents to upsert
            class_name: Label to upsert into
            match_keys: Keys to match for upsert operation
            **kwargs: Additional options:
                - dry (bool): If True, build query but don't execute

        Raises:
            ValueError: If any document is missing a required match_key or has None value
        """
        dry = kwargs.pop("dry", False)

        if not docs:
            return

        # Sanitize documents: filter invalid keys/values, validate match_keys
        match_keys_list = (
            list(match_keys) if isinstance(match_keys, tuple) else match_keys
        )
        sanitized_docs = self._sanitize_batch(docs, match_keys_list)

        if not sanitized_docs:
            return

        # Build the MERGE clause with match keys
        index_str = ", ".join([f"{k}: row.{k}" for k in match_keys])
        # Use 'data' instead of 'batch' as parameter name (batch might be reserved)
        # Try direct UNWIND - FalkorDB may not support WITH $param AS alias
        # Datetime objects are converted to ISO-8601 strings during sanitization
        q = f"""
            UNWIND $data AS row
            MERGE (n:{class_name} {{ {index_str} }})
            ON MATCH SET n += row
            ON CREATE SET n += row
        """
        if not dry:
            self.execute(q, data=sanitized_docs)

    def insert_edges_batch(
        self,
        docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
        source_class: str,
        target_class: str,
        relation_name: str,
        match_keys_source: tuple[str, ...],
        match_keys_target: tuple[str, ...],
        filter_uniques: bool = True,
        head: int | None = None,
        **kwargs: Any,
    ) -> None:
        """Create relationships between existing nodes using Cypher MERGE.

        Args:
            docs_edges: Edge specifications as list of [source, target, props] triples
            source_class: Label of source nodes
            target_class: Label of target nodes
            relation_name: Relationship type name
            match_keys_source: Properties to match source nodes
            match_keys_target: Properties to match target nodes
            filter_uniques: Unused in FalkorDB (MERGE handles uniqueness automatically)
            head: Optional limit on number of relationships to insert
            **kwargs: Additional options:
                - dry (bool): If True, build query but don't execute
        """
        dry = kwargs.pop("dry", False)
        # Extract and ignore unused parameters (kept for interface compatibility)
        kwargs.pop("collection_name", None)
        kwargs.pop("uniq_weight_fields", None)
        kwargs.pop("uniq_weight_collections", None)
        kwargs.pop("upsert_option", None)

        # Apply head limit if specified
        if head is not None and isinstance(docs_edges, list):
            docs_edges = docs_edges[:head]

        if not docs_edges:
            return

        # Note: filter_uniques is unused because FalkorDB's MERGE handles uniqueness automatically

        # Sanitize edge data: each edge is [source_dict, target_dict, props_dict]
        # We need to sanitize source, target, and props dictionaries
        sanitized_edges = []
        for edge in docs_edges:
            if len(edge) != 3:
                logger.warning(
                    f"Skipping invalid edge format: expected [source, target, props], got {edge}"
                )
                continue

            source_dict, target_dict, props_dict = edge

            # Sanitize source and target dictionaries (for match keys)
            sanitized_source = self._sanitize_document(
                source_dict if isinstance(source_dict, dict) else {},
                match_keys=list(match_keys_source),
            )
            sanitized_target = self._sanitize_document(
                target_dict if isinstance(target_dict, dict) else {},
                match_keys=list(match_keys_target),
            )

            # Sanitize props dictionary (may contain datetime objects)
            sanitized_props = self._sanitize_document(
                props_dict if isinstance(props_dict, dict) else {}
            )

            sanitized_edges.append(
                (sanitized_source, sanitized_target, sanitized_props)
            )

        if not sanitized_edges:
            return

        # Build match conditions for source and target nodes
        source_match_str = [f"source.{key} = row[0].{key}" for key in match_keys_source]
        target_match_str = [f"target.{key} = row[1].{key}" for key in match_keys_target]

        match_clause = "WHERE " + " AND ".join(source_match_str + target_match_str)

        # Datetime objects are converted to ISO-8601 strings during sanitization
        q = f"""
            UNWIND $data AS row
            MATCH (source:{source_class}),
                  (target:{target_class}) {match_clause}
            MERGE (source)-[r:{relation_name}]->(target)
            SET r += row[2]
        """
        if not dry:
            self.execute(q, data=sanitized_edges)

    def insert_return_batch(
        self, docs: list[dict[str, Any]], class_name: str
    ) -> list[dict[str, Any]] | str:
        """Insert nodes and return their properties.

        Args:
            docs: Documents to insert
            class_name: Label to insert into

        Raises:
            NotImplementedError: This method is not fully implemented for FalkorDB
        """
        raise NotImplementedError("insert_return_batch is not implemented for FalkorDB")

    def fetch_docs(
        self,
        class_name,
        filters: list | dict | None = None,
        limit: int | None = None,
        return_keys: list | None = None,
        unset_keys: list | None = None,
        **kwargs,
    ):
        """Fetch nodes from a label.

        Args:
            class_name: Label to fetch from
            filters: Query filters
            limit: Maximum number of nodes to return
            return_keys: Keys to return
            unset_keys: Unused in FalkorDB
            **kwargs: Additional parameters

        Returns:
            List of fetched nodes as dictionaries
        """
        # Build filter clause
        if filters is not None:
            ff = FilterExpression.from_dict(filters)
            filter_clause = f"WHERE {ff(doc_name='n', kind=self.expression_flavor())}"
        else:
            filter_clause = ""

        # Build return clause
        if return_keys is not None:
            # Project specific keys
            keep_clause_ = ", ".join([f"n.{item} AS {item}" for item in return_keys])
            return_clause = f"RETURN {keep_clause_}"
        else:
            return_clause = "RETURN n"

        # Build limit clause (must be positive integer)
        if limit is not None and isinstance(limit, int) and limit > 0:
            limit_clause = f"LIMIT {limit}"
        else:
            limit_clause = ""

        q = f"""
            MATCH (n:{class_name})
            {filter_clause}
            {return_clause}
            {limit_clause}
        """

        result = self.execute(q)

        # Convert FalkorDB results to list of dictionaries
        if return_keys is not None:
            # Results are already projected
            return [dict(zip(return_keys, row)) for row in result.result_set]
        else:
            # Results contain node objects
            return [self._node_to_dict(row[0]) for row in result.result_set]

    def _node_to_dict(self, node) -> dict:
        """Convert a FalkorDB node to a dictionary.

        Args:
            node: FalkorDB node object

        Returns:
            Node properties as dictionary
        """
        if hasattr(node, "properties"):
            return dict(node.properties)
        elif isinstance(node, dict):
            return node
        else:
            # Try to convert to dict
            return dict(node) if node else {}

    def fetch_edges(
        self,
        from_type: str,
        from_id: str,
        edge_type: str | None = None,
        to_type: str | None = None,
        to_id: str | None = None,
        filters: list | dict | None = None,
        limit: int | None = None,
        return_keys: list | None = None,
        unset_keys: list | None = None,
        **kwargs,
    ):
        """Fetch edges from FalkorDB using Cypher.

        Args:
            from_type: Source node label
            from_id: Source node ID (property name depends on match_keys used)
            edge_type: Optional relationship type to filter by
            to_type: Optional target node label to filter by
            to_id: Optional target node ID to filter by
            filters: Additional query filters
            limit: Maximum number of edges to return
            return_keys: Keys to return (projection)
            unset_keys: Keys to exclude (projection) - not supported in FalkorDB
            **kwargs: Additional parameters

        Returns:
            List of fetched edges as dictionaries
        """
        # Build source node match
        source_match = f"(source:{from_type} {{id: '{from_id}'}})"

        # Build relationship pattern
        if edge_type:
            rel_pattern = f"-[r:{edge_type}]->"
        else:
            rel_pattern = "-[r]->"

        # Build target node match
        if to_type:
            target_match = f"(target:{to_type})"
        else:
            target_match = "(target)"

        # Build WHERE clauses
        where_clauses = []
        if to_id:
            where_clauses.append(f"target.id = '{to_id}'")

        # Add additional filters if provided
        if filters is not None:
            ff = FilterExpression.from_dict(filters)
            filter_clause = ff(doc_name="r", kind=self.expression_flavor())
            where_clauses.append(filter_clause)

        where_clause = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""

        # Build return clause
        if return_keys is not None:
            return_parts = ", ".join([f"r.{key} AS {key}" for key in return_keys])
            return_clause = f"RETURN {return_parts}"
        else:
            return_clause = "RETURN r"

        limit_clause = f"LIMIT {limit}" if limit and limit > 0 else ""

        query = f"""
            MATCH {source_match}{rel_pattern}{target_match}
            {where_clause}
            {return_clause}
            {limit_clause}
        """

        result = self.execute(query)

        # Convert results
        if return_keys is not None:
            return [dict(zip(return_keys, row)) for row in result.result_set]
        else:
            return [self._edge_to_dict(row[0]) for row in result.result_set]

    def _edge_to_dict(self, edge) -> dict:
        """Convert a FalkorDB edge to a dictionary.

        Args:
            edge: FalkorDB edge object

        Returns:
            Edge properties as dictionary
        """
        if hasattr(edge, "properties"):
            return dict(edge.properties)
        elif isinstance(edge, dict):
            return edge
        else:
            return dict(edge) if edge else {}

    def fetch_present_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        flatten: bool = False,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Fetch nodes that exist in the database.

        Args:
            batch: Batch of documents to check
            class_name: Label to check in
            match_keys: Keys to match nodes
            keep_keys: Keys to keep in result
            flatten: Unused in FalkorDB
            filters: Additional query filters

        Returns:
            Documents that exist in the database
        """
        if not batch:
            return []

        # Build match conditions for each document in batch
        results = []
        for doc in batch:
            match_conditions = " AND ".join([f"n.{key} = ${key}" for key in match_keys])
            params = {key: doc.get(key) for key in match_keys}

            q = f"""
                MATCH (n:{class_name})
                WHERE {match_conditions}
                RETURN n
                LIMIT 1
            """

            try:
                result = self.execute(q, **params)
                if result.result_set:
                    node_dict = self._node_to_dict(result.result_set[0][0])
                    if keep_keys:
                        node_dict = {k: node_dict.get(k) for k in keep_keys}
                    results.append(node_dict)
            except Exception as e:
                logger.debug(f"Error checking document presence: {e}")

        return results

    def aggregate(
        self,
        class_name,
        aggregation_function: AggregationType,
        discriminant: str | None = None,
        aggregated_field: str | None = None,
        filters: list | dict | None = None,
    ):
        """Perform aggregation on nodes.

        Args:
            class_name: Label to aggregate
            aggregation_function: Type of aggregation to perform
            discriminant: Field to group by
            aggregated_field: Field to aggregate
            filters: Query filters

        Returns:
            Aggregation results (dict for grouped aggregations, int/float for single value)
        """
        # Build filter clause
        if filters is not None:
            ff = FilterExpression.from_dict(filters)
            filter_clause = f"WHERE {ff(doc_name='n', kind=self.expression_flavor())}"
        else:
            filter_clause = ""

        # Build aggregation query based on function type
        if aggregation_function == AggregationType.COUNT:
            if discriminant:
                q = f"""
                    MATCH (n:{class_name})
                    {filter_clause}
                    RETURN n.{discriminant} AS key, count(*) AS count
                """
                result = self.execute(q)
                return {row[0]: row[1] for row in result.result_set}
            else:
                q = f"""
                    MATCH (n:{class_name})
                    {filter_clause}
                    RETURN count(*) AS count
                """
                result = self.execute(q)
                return result.result_set[0][0] if result.result_set else 0

        elif aggregation_function == AggregationType.MAX:
            if not aggregated_field:
                raise ValueError("aggregated_field is required for MAX aggregation")
            q = f"""
                MATCH (n:{class_name})
                {filter_clause}
                RETURN max(n.{aggregated_field}) AS max_value
            """
            result = self.execute(q)
            return result.result_set[0][0] if result.result_set else None

        elif aggregation_function == AggregationType.MIN:
            if not aggregated_field:
                raise ValueError("aggregated_field is required for MIN aggregation")
            q = f"""
                MATCH (n:{class_name})
                {filter_clause}
                RETURN min(n.{aggregated_field}) AS min_value
            """
            result = self.execute(q)
            return result.result_set[0][0] if result.result_set else None

        elif aggregation_function == AggregationType.AVERAGE:
            if not aggregated_field:
                raise ValueError("aggregated_field is required for AVERAGE aggregation")
            q = f"""
                MATCH (n:{class_name})
                {filter_clause}
                RETURN avg(n.{aggregated_field}) AS avg_value
            """
            result = self.execute(q)
            return result.result_set[0][0] if result.result_set else None

        elif aggregation_function == AggregationType.SORTED_UNIQUE:
            if not aggregated_field:
                raise ValueError(
                    "aggregated_field is required for SORTED_UNIQUE aggregation"
                )
            q = f"""
                MATCH (n:{class_name})
                {filter_clause}
                RETURN DISTINCT n.{aggregated_field} AS value
                ORDER BY value
            """
            result = self.execute(q)
            return [row[0] for row in result.result_set]

        else:
            raise ValueError(
                f"Unsupported aggregation function: {aggregation_function}"
            )

    def keep_absent_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Keep documents that don't exist in the database.

        Args:
            batch: Batch of documents to check
            class_name: Label to check in
            match_keys: Keys to match nodes
            keep_keys: Keys to keep in result
            filters: Additional query filters

        Returns:
            Documents that don't exist in the database
        """
        if not batch:
            return []

        # Find documents that exist
        present_docs = self.fetch_present_documents(
            batch, class_name, match_keys, match_keys, filters=filters
        )

        # Create a set of present document keys for efficient lookup
        present_keys = set()
        for doc in present_docs:
            key_tuple = tuple(doc.get(k) for k in match_keys)
            present_keys.add(key_tuple)

        # Filter out documents that exist
        absent_docs = []
        for doc in batch:
            key_tuple = tuple(doc.get(k) for k in match_keys)
            if key_tuple not in present_keys:
                if keep_keys:
                    absent_docs.append({k: doc.get(k) for k in keep_keys})
                else:
                    absent_docs.append(doc)

        return absent_docs

__init__(config)

Initialize FalkorDB connection.

Parameters:

Name	Type	Description	Default
config	FalkordbConfig	FalkorDB connection configuration containing URI, database, and optional password	required

Source code in graflo/db/falkordb/conn.py

def __init__(self, config: FalkordbConfig):
    """Initialize FalkorDB connection.

    Args:
        config: FalkorDB connection configuration containing URI, database,
            and optional password
    """
    super().__init__()
    self.config = config

    if config.uri is None:
        raise ValueError("FalkorDB connection requires a URI to be configured")

    # Parse URI to extract host and port
    parsed = urlparse(config.uri)
    host = parsed.hostname or "localhost"
    port = parsed.port or 6379

    # Initialize FalkorDB client
    if config.password:
        self.client = FalkorDB(host=host, port=port, password=config.password)
    else:
        self.client = FalkorDB(host=host, port=port)

    # Select the graph (database in config maps to graph name)
    graph_name = config.database or "default"
    self.graph = self.client.select_graph(graph_name)
    self._graph_name = graph_name

aggregate(class_name, aggregation_function, discriminant=None, aggregated_field=None, filters=None)

Perform aggregation on nodes.

Parameters:

Name	Type	Description	Default
class_name		Label to aggregate	required
aggregation_function	AggregationType	Type of aggregation to perform	required
discriminant	str | None	Field to group by	None
aggregated_field	str | None	Field to aggregate	None
filters	list | dict | None	Query filters	None

Returns:

Type	Description
	Aggregation results (dict for grouped aggregations, int/float for single value)

Source code in graflo/db/falkordb/conn.py

def aggregate(
    self,
    class_name,
    aggregation_function: AggregationType,
    discriminant: str | None = None,
    aggregated_field: str | None = None,
    filters: list | dict | None = None,
):
    """Perform aggregation on nodes.

    Args:
        class_name: Label to aggregate
        aggregation_function: Type of aggregation to perform
        discriminant: Field to group by
        aggregated_field: Field to aggregate
        filters: Query filters

    Returns:
        Aggregation results (dict for grouped aggregations, int/float for single value)
    """
    # Build filter clause
    if filters is not None:
        ff = FilterExpression.from_dict(filters)
        filter_clause = f"WHERE {ff(doc_name='n', kind=self.expression_flavor())}"
    else:
        filter_clause = ""

    # Build aggregation query based on function type
    if aggregation_function == AggregationType.COUNT:
        if discriminant:
            q = f"""
                MATCH (n:{class_name})
                {filter_clause}
                RETURN n.{discriminant} AS key, count(*) AS count
            """
            result = self.execute(q)
            return {row[0]: row[1] for row in result.result_set}
        else:
            q = f"""
                MATCH (n:{class_name})
                {filter_clause}
                RETURN count(*) AS count
            """
            result = self.execute(q)
            return result.result_set[0][0] if result.result_set else 0

    elif aggregation_function == AggregationType.MAX:
        if not aggregated_field:
            raise ValueError("aggregated_field is required for MAX aggregation")
        q = f"""
            MATCH (n:{class_name})
            {filter_clause}
            RETURN max(n.{aggregated_field}) AS max_value
        """
        result = self.execute(q)
        return result.result_set[0][0] if result.result_set else None

    elif aggregation_function == AggregationType.MIN:
        if not aggregated_field:
            raise ValueError("aggregated_field is required for MIN aggregation")
        q = f"""
            MATCH (n:{class_name})
            {filter_clause}
            RETURN min(n.{aggregated_field}) AS min_value
        """
        result = self.execute(q)
        return result.result_set[0][0] if result.result_set else None

    elif aggregation_function == AggregationType.AVERAGE:
        if not aggregated_field:
            raise ValueError("aggregated_field is required for AVERAGE aggregation")
        q = f"""
            MATCH (n:{class_name})
            {filter_clause}
            RETURN avg(n.{aggregated_field}) AS avg_value
        """
        result = self.execute(q)
        return result.result_set[0][0] if result.result_set else None

    elif aggregation_function == AggregationType.SORTED_UNIQUE:
        if not aggregated_field:
            raise ValueError(
                "aggregated_field is required for SORTED_UNIQUE aggregation"
            )
        q = f"""
            MATCH (n:{class_name})
            {filter_clause}
            RETURN DISTINCT n.{aggregated_field} AS value
            ORDER BY value
        """
        result = self.execute(q)
        return [row[0] for row in result.result_set]

    else:
        raise ValueError(
            f"Unsupported aggregation function: {aggregation_function}"
        )

clear_data(schema)

Remove all data from the graph without dropping the schema.

Deletes all nodes and relationships; labels (schema) remain.

Source code in graflo/db/falkordb/conn.py

def clear_data(self, schema: Schema) -> None:
    """Remove all data from the graph without dropping the schema.

    Deletes all nodes and relationships; labels (schema) remain.
    """
    self.delete_graph_structure(delete_all=True)

close()

Close the FalkorDB connection.

Source code in graflo/db/falkordb/conn.py

def close(self):
    """Close the FalkorDB connection."""
    # FalkorDB client handles connection pooling internally
    # No explicit close needed, but we can delete the reference
    self.graph = None
    self.client = None

create_database(name)

Create a new graph in FalkorDB.

Note: In FalkorDB, graphs are created implicitly when data is first inserted.

Parameters:

Name	Type	Description	Default
name	str	Name of the graph to create	required

Source code in graflo/db/falkordb/conn.py

def create_database(self, name: str):
    """Create a new graph in FalkorDB.

    Note: In FalkorDB, graphs are created implicitly when data is first inserted.

    Args:
        name: Name of the graph to create
    """
    # In FalkorDB, graphs are created implicitly when you first insert data
    # We just need to select the graph
    assert self.client is not None, "Connection is closed"
    self.graph = self.client.select_graph(name)
    self._graph_name = name
    logger.info(f"Selected FalkorDB graph '{name}'")

define_edge_classes(edges)

Define edge classes based on schema.

Note: This is a no-op in FalkorDB as edge classes (relationship types) are implicit.

Parameters:

Name	Type	Description	Default
edges	list[Edge]	List of edge configurations	required

Source code in graflo/db/falkordb/conn.py

def define_edge_classes(self, edges: list[Edge]):
    """Define edge classes based on schema.

    Note: This is a no-op in FalkorDB as edge classes (relationship types) are implicit.

    Args:
        edges: List of edge configurations
    """
    pass

define_edge_indices(edges)

Define indices for relationship types.

Creates indices for each relationship type based on the configuration. FalkorDB supports range indices on relationship properties.

Parameters:

Name	Type	Description	Default
edges	list[Edge]	List of edge configurations containing index definitions	required

Source code in graflo/db/falkordb/conn.py

def define_edge_indices(self, edges: list[Edge]):
    """Define indices for relationship types.

    Creates indices for each relationship type based on the configuration.
    FalkorDB supports range indices on relationship properties.

    Args:
        edges: List of edge configurations containing index definitions
    """
    for edge in edges:
        for index_obj in edge.indexes:
            if edge.relation is not None:
                self._add_index(edge.relation, index_obj, is_vertex_index=False)

define_schema(schema)

Define vertex and edge classes based on schema.

Note: This is a no-op in FalkorDB as vertex/edge classes (labels/relationship types) are implicit.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing vertex and edge class definitions	required

Source code in graflo/db/falkordb/conn.py

def define_schema(self, schema: Schema):
    """Define vertex and edge classes based on schema.

    Note: This is a no-op in FalkorDB as vertex/edge classes (labels/relationship types) are implicit.

    Args:
        schema: Schema containing vertex and edge class definitions
    """
    pass

define_vertex_classes(schema)

Define vertex classes based on schema.

Note: This is a no-op in FalkorDB as vertex classes (labels) are implicit.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing vertex definitions	required

Source code in graflo/db/falkordb/conn.py

def define_vertex_classes(self, schema: Schema):
    """Define vertex classes based on schema.

    Note: This is a no-op in FalkorDB as vertex classes (labels) are implicit.

    Args:
        schema: Schema containing vertex definitions
    """
    pass

define_vertex_indices(vertex_config)

Define indices for vertex labels.

Creates indices for each vertex label based on the configuration. FalkorDB supports range indices on node properties.

Parameters:

Name	Type	Description	Default
vertex_config	VertexConfig	Vertex configuration containing index definitions	required

Source code in graflo/db/falkordb/conn.py

def define_vertex_indices(self, vertex_config: VertexConfig):
    """Define indices for vertex labels.

    Creates indices for each vertex label based on the configuration.
    FalkorDB supports range indices on node properties.

    Args:
        vertex_config: Vertex configuration containing index definitions
    """
    for c in vertex_config.vertex_set:
        for index_obj in vertex_config.indexes(c):
            self._add_index(c, index_obj)

delete_database(name)

Delete a graph from FalkorDB.

Parameters:

Name	Type	Description	Default
name	str	Name of the graph to delete (if empty, uses current graph)	required

Source code in graflo/db/falkordb/conn.py

def delete_database(self, name: str):
    """Delete a graph from FalkorDB.

    Args:
        name: Name of the graph to delete (if empty, uses current graph)
    """
    graph_to_delete = name if name else self._graph_name
    assert self.client is not None, "Connection is closed"
    try:
        # Delete the graph using the FalkorDB API
        graph = self.client.select_graph(graph_to_delete)
        graph.delete()
        logger.info(f"Successfully deleted FalkorDB graph '{graph_to_delete}'")
    except Exception as e:
        logger.error(
            f"Failed to delete FalkorDB graph '{graph_to_delete}': {e}",
            exc_info=True,
        )
        raise

delete_graph_structure(vertex_types=(), graph_names=(), delete_all=False)

Delete graph structure (nodes and relationships) from FalkorDB.

Parameters:

Name	Type	Description	Default
vertex_types	tuple[str, ...] | list[str]	Label names to delete nodes for	()
graph_names	tuple[str, ...] | list[str]	Graph names to delete entirely	()
delete_all	bool	If True, delete all nodes and relationships	False

Source code in graflo/db/falkordb/conn.py

def delete_graph_structure(
    self,
    vertex_types: tuple[str, ...] | list[str] = (),
    graph_names: tuple[str, ...] | list[str] = (),
    delete_all: bool = False,
) -> None:
    """Delete graph structure (nodes and relationships) from FalkorDB.

    Args:
        vertex_types: Label names to delete nodes for
        graph_names: Graph names to delete entirely
        delete_all: If True, delete all nodes and relationships
    """
    if delete_all or (not vertex_types and not graph_names):
        # Delete all nodes and relationships in current graph
        try:
            self.execute("MATCH (n) DETACH DELETE n")
            logger.debug("Deleted all nodes and relationships from graph")
        except Exception as e:
            logger.debug(f"Graph may be empty or not exist: {e}")
    elif vertex_types:
        # Delete nodes with specific labels
        for label in vertex_types:
            try:
                self.execute(f"MATCH (n:{label}) DETACH DELETE n")
                logger.debug(f"Deleted all nodes with label '{label}'")
            except Exception as e:
                logger.warning(f"Failed to delete nodes with label '{label}': {e}")

    # Delete specific graphs
    assert self.client is not None, "Connection is closed"
    for graph_name in graph_names:
        try:
            graph = self.client.select_graph(graph_name)
            graph.delete()
            logger.debug(f"Deleted graph '{graph_name}'")
        except Exception as e:
            logger.warning(f"Failed to delete graph '{graph_name}': {e}")

execute(query, **kwargs)

Execute a raw OpenCypher query against the graph.

Parameters:

Name	Type	Description	Default
query	str	OpenCypher query string. Can include parameter placeholders using $name syntax (e.g., "MATCH (n) WHERE n.id = $id")	required
**kwargs		Query parameters as keyword arguments	{}

Returns:

Name	Type	Description
QueryResult		FalkorDB result object containing result_set and statistics

Source code in graflo/db/falkordb/conn.py

def execute(self, query: str, **kwargs):
    """Execute a raw OpenCypher query against the graph.

    Args:
        query: OpenCypher query string. Can include parameter placeholders
            using $name syntax (e.g., "MATCH (n) WHERE n.id = $id")
        **kwargs: Query parameters as keyword arguments

    Returns:
        QueryResult: FalkorDB result object containing result_set and statistics
    """
    assert self.graph is not None, "Connection is closed"
    # Pass params as keyword argument if the client supports it, otherwise as positional
    # Try params keyword first, fall back to positional
    if kwargs:
        try:
            result = self.graph.query(query, params=kwargs)
        except TypeError:
            # Fall back to positional argument if params keyword not supported
            result = self.graph.query(query, kwargs)
    else:
        result = self.graph.query(query)
    return result

fetch_docs(class_name, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs)

Fetch nodes from a label.

Parameters:

Name	Type	Description	Default
class_name		Label to fetch from	required
filters	list | dict | None	Query filters	None
limit	int | None	Maximum number of nodes to return	None
return_keys	list | None	Keys to return	None
unset_keys	list | None	Unused in FalkorDB	None
**kwargs		Additional parameters	{}

Returns:

Type	Description
	List of fetched nodes as dictionaries

Source code in graflo/db/falkordb/conn.py

def fetch_docs(
    self,
    class_name,
    filters: list | dict | None = None,
    limit: int | None = None,
    return_keys: list | None = None,
    unset_keys: list | None = None,
    **kwargs,
):
    """Fetch nodes from a label.

    Args:
        class_name: Label to fetch from
        filters: Query filters
        limit: Maximum number of nodes to return
        return_keys: Keys to return
        unset_keys: Unused in FalkorDB
        **kwargs: Additional parameters

    Returns:
        List of fetched nodes as dictionaries
    """
    # Build filter clause
    if filters is not None:
        ff = FilterExpression.from_dict(filters)
        filter_clause = f"WHERE {ff(doc_name='n', kind=self.expression_flavor())}"
    else:
        filter_clause = ""

    # Build return clause
    if return_keys is not None:
        # Project specific keys
        keep_clause_ = ", ".join([f"n.{item} AS {item}" for item in return_keys])
        return_clause = f"RETURN {keep_clause_}"
    else:
        return_clause = "RETURN n"

    # Build limit clause (must be positive integer)
    if limit is not None and isinstance(limit, int) and limit > 0:
        limit_clause = f"LIMIT {limit}"
    else:
        limit_clause = ""

    q = f"""
        MATCH (n:{class_name})
        {filter_clause}
        {return_clause}
        {limit_clause}
    """

    result = self.execute(q)

    # Convert FalkorDB results to list of dictionaries
    if return_keys is not None:
        # Results are already projected
        return [dict(zip(return_keys, row)) for row in result.result_set]
    else:
        # Results contain node objects
        return [self._node_to_dict(row[0]) for row in result.result_set]

fetch_edges(from_type, from_id, edge_type=None, to_type=None, to_id=None, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs)

Fetch edges from FalkorDB using Cypher.

Parameters:

Name	Type	Description	Default
from_type	str	Source node label	required
from_id	str	Source node ID (property name depends on match_keys used)	required
edge_type	str | None	Optional relationship type to filter by	None
to_type	str | None	Optional target node label to filter by	None
to_id	str | None	Optional target node ID to filter by	None
filters	list | dict | None	Additional query filters	None
limit	int | None	Maximum number of edges to return	None
return_keys	list | None	Keys to return (projection)	None
unset_keys	list | None	Keys to exclude (projection) - not supported in FalkorDB	None
**kwargs		Additional parameters	{}

Returns:

Type	Description
	List of fetched edges as dictionaries

Source code in graflo/db/falkordb/conn.py

def fetch_edges(
    self,
    from_type: str,
    from_id: str,
    edge_type: str | None = None,
    to_type: str | None = None,
    to_id: str | None = None,
    filters: list | dict | None = None,
    limit: int | None = None,
    return_keys: list | None = None,
    unset_keys: list | None = None,
    **kwargs,
):
    """Fetch edges from FalkorDB using Cypher.

    Args:
        from_type: Source node label
        from_id: Source node ID (property name depends on match_keys used)
        edge_type: Optional relationship type to filter by
        to_type: Optional target node label to filter by
        to_id: Optional target node ID to filter by
        filters: Additional query filters
        limit: Maximum number of edges to return
        return_keys: Keys to return (projection)
        unset_keys: Keys to exclude (projection) - not supported in FalkorDB
        **kwargs: Additional parameters

    Returns:
        List of fetched edges as dictionaries
    """
    # Build source node match
    source_match = f"(source:{from_type} {{id: '{from_id}'}})"

    # Build relationship pattern
    if edge_type:
        rel_pattern = f"-[r:{edge_type}]->"
    else:
        rel_pattern = "-[r]->"

    # Build target node match
    if to_type:
        target_match = f"(target:{to_type})"
    else:
        target_match = "(target)"

    # Build WHERE clauses
    where_clauses = []
    if to_id:
        where_clauses.append(f"target.id = '{to_id}'")

    # Add additional filters if provided
    if filters is not None:
        ff = FilterExpression.from_dict(filters)
        filter_clause = ff(doc_name="r", kind=self.expression_flavor())
        where_clauses.append(filter_clause)

    where_clause = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""

    # Build return clause
    if return_keys is not None:
        return_parts = ", ".join([f"r.{key} AS {key}" for key in return_keys])
        return_clause = f"RETURN {return_parts}"
    else:
        return_clause = "RETURN r"

    limit_clause = f"LIMIT {limit}" if limit and limit > 0 else ""

    query = f"""
        MATCH {source_match}{rel_pattern}{target_match}
        {where_clause}
        {return_clause}
        {limit_clause}
    """

    result = self.execute(query)

    # Convert results
    if return_keys is not None:
        return [dict(zip(return_keys, row)) for row in result.result_set]
    else:
        return [self._edge_to_dict(row[0]) for row in result.result_set]

fetch_present_documents(batch, class_name, match_keys, keep_keys=None, flatten=False, filters=None)

Fetch nodes that exist in the database.

Parameters:

Name	Type	Description	Default
batch	list[dict[str, Any]]	Batch of documents to check	required
class_name	str	Label to check in	required
match_keys	list[str] | tuple[str, ...]	Keys to match nodes	required
keep_keys	list[str] | tuple[str, ...] | None	Keys to keep in result	None
flatten	bool	Unused in FalkorDB	False
filters	list[Any] | dict[str, Any] | None	Additional query filters	None

Returns:

Type	Description
list[dict[str, Any]]	Documents that exist in the database

Source code in graflo/db/falkordb/conn.py

def fetch_present_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    flatten: bool = False,
    filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """Fetch nodes that exist in the database.

    Args:
        batch: Batch of documents to check
        class_name: Label to check in
        match_keys: Keys to match nodes
        keep_keys: Keys to keep in result
        flatten: Unused in FalkorDB
        filters: Additional query filters

    Returns:
        Documents that exist in the database
    """
    if not batch:
        return []

    # Build match conditions for each document in batch
    results = []
    for doc in batch:
        match_conditions = " AND ".join([f"n.{key} = ${key}" for key in match_keys])
        params = {key: doc.get(key) for key in match_keys}

        q = f"""
            MATCH (n:{class_name})
            WHERE {match_conditions}
            RETURN n
            LIMIT 1
        """

        try:
            result = self.execute(q, **params)
            if result.result_set:
                node_dict = self._node_to_dict(result.result_set[0][0])
                if keep_keys:
                    node_dict = {k: node_dict.get(k) for k in keep_keys}
                results.append(node_dict)
        except Exception as e:
            logger.debug(f"Error checking document presence: {e}")

    return results

init_db(schema, recreate_schema)

Initialize FalkorDB with the given schema.

Uses schema.general.name if database is not set in config.

If the graph already has nodes and recreate_schema is False, raises SchemaExistsError and the script halts.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing graph structure definitions	required
recreate_schema	bool	If True, delete all existing data before initialization. If False and graph has nodes, raises SchemaExistsError.	required

Source code in graflo/db/falkordb/conn.py

def init_db(self, schema: Schema, recreate_schema: bool) -> None:
    """Initialize FalkorDB with the given schema.

    Uses schema.general.name if database is not set in config.

    If the graph already has nodes and recreate_schema is False, raises
    SchemaExistsError and the script halts.

    Args:
        schema: Schema containing graph structure definitions
        recreate_schema: If True, delete all existing data before initialization.
            If False and graph has nodes, raises SchemaExistsError.
    """
    # Determine graph name: use config.database if set, otherwise use schema.general.name
    graph_name = self.config.database
    if not graph_name:
        graph_name = schema.general.name
        self.config.database = graph_name

    # Select/create the graph
    assert self.client is not None, "Connection is closed"
    self.graph = self.client.select_graph(graph_name)
    self._graph_name = graph_name
    logger.info(f"Initialized FalkorDB graph '{graph_name}'")

    # Check if graph already has nodes (schema/graph exists)
    try:
        result = self.execute("MATCH (n) RETURN count(n) AS c")
        count = 0
        if hasattr(result, "data") and result.data():
            count = result.data()[0].get("c", 0) or 0
        elif result is not None and hasattr(result, "__iter__"):
            for record in result:
                count = (
                    record.get("c", 0)
                    if hasattr(record, "get")
                    else getattr(record, "c", 0)
                ) or 0
                break
        if count > 0 and not recreate_schema:
            raise SchemaExistsError(
                f"Schema/graph already exists in graph '{graph_name}' ({count} nodes). "
                "Set recreate_schema=True to replace, or use clear_data=True before ingestion."
            )
    except SchemaExistsError:
        raise
    except Exception as e:
        logger.debug(f"Could not check graph node count: {e}")

    if recreate_schema:
        try:
            self.delete_graph_structure(delete_all=True)
            logger.debug(f"Cleaned graph '{graph_name}' for fresh start")
        except Exception as e:
            logger.debug(f"Recreate schema note for graph '{graph_name}': {e}")

    try:
        self.define_indexes(schema)
        logger.debug(f"Defined indexes for graph '{graph_name}'")
    except Exception as e:
        logger.error(
            f"Failed to define indexes for graph '{graph_name}': {e}",
            exc_info=True,
        )
        raise

insert_edges_batch(docs_edges, source_class, target_class, relation_name, match_keys_source, match_keys_target, filter_uniques=True, head=None, **kwargs)

Create relationships between existing nodes using Cypher MERGE.

Parameters:

Name	Type	Description	Default
docs_edges	list[list[dict[str, Any]]] | list[Any] | None	Edge specifications as list of [source, target, props] triples	required
source_class	str	Label of source nodes	required
target_class	str	Label of target nodes	required
relation_name	str	Relationship type name	required
match_keys_source	tuple[str, ...]	Properties to match source nodes	required
match_keys_target	tuple[str, ...]	Properties to match target nodes	required
filter_uniques	bool	Unused in FalkorDB (MERGE handles uniqueness automatically)	True
head	int | None	Optional limit on number of relationships to insert	None
**kwargs	Any	Additional options: - dry (bool): If True, build query but don't execute	{}

Source code in graflo/db/falkordb/conn.py

def insert_edges_batch(
    self,
    docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
    source_class: str,
    target_class: str,
    relation_name: str,
    match_keys_source: tuple[str, ...],
    match_keys_target: tuple[str, ...],
    filter_uniques: bool = True,
    head: int | None = None,
    **kwargs: Any,
) -> None:
    """Create relationships between existing nodes using Cypher MERGE.

    Args:
        docs_edges: Edge specifications as list of [source, target, props] triples
        source_class: Label of source nodes
        target_class: Label of target nodes
        relation_name: Relationship type name
        match_keys_source: Properties to match source nodes
        match_keys_target: Properties to match target nodes
        filter_uniques: Unused in FalkorDB (MERGE handles uniqueness automatically)
        head: Optional limit on number of relationships to insert
        **kwargs: Additional options:
            - dry (bool): If True, build query but don't execute
    """
    dry = kwargs.pop("dry", False)
    # Extract and ignore unused parameters (kept for interface compatibility)
    kwargs.pop("collection_name", None)
    kwargs.pop("uniq_weight_fields", None)
    kwargs.pop("uniq_weight_collections", None)
    kwargs.pop("upsert_option", None)

    # Apply head limit if specified
    if head is not None and isinstance(docs_edges, list):
        docs_edges = docs_edges[:head]

    if not docs_edges:
        return

    # Note: filter_uniques is unused because FalkorDB's MERGE handles uniqueness automatically

    # Sanitize edge data: each edge is [source_dict, target_dict, props_dict]
    # We need to sanitize source, target, and props dictionaries
    sanitized_edges = []
    for edge in docs_edges:
        if len(edge) != 3:
            logger.warning(
                f"Skipping invalid edge format: expected [source, target, props], got {edge}"
            )
            continue

        source_dict, target_dict, props_dict = edge

        # Sanitize source and target dictionaries (for match keys)
        sanitized_source = self._sanitize_document(
            source_dict if isinstance(source_dict, dict) else {},
            match_keys=list(match_keys_source),
        )
        sanitized_target = self._sanitize_document(
            target_dict if isinstance(target_dict, dict) else {},
            match_keys=list(match_keys_target),
        )

        # Sanitize props dictionary (may contain datetime objects)
        sanitized_props = self._sanitize_document(
            props_dict if isinstance(props_dict, dict) else {}
        )

        sanitized_edges.append(
            (sanitized_source, sanitized_target, sanitized_props)
        )

    if not sanitized_edges:
        return

    # Build match conditions for source and target nodes
    source_match_str = [f"source.{key} = row[0].{key}" for key in match_keys_source]
    target_match_str = [f"target.{key} = row[1].{key}" for key in match_keys_target]

    match_clause = "WHERE " + " AND ".join(source_match_str + target_match_str)

    # Datetime objects are converted to ISO-8601 strings during sanitization
    q = f"""
        UNWIND $data AS row
        MATCH (source:{source_class}),
              (target:{target_class}) {match_clause}
        MERGE (source)-[r:{relation_name}]->(target)
        SET r += row[2]
    """
    if not dry:
        self.execute(q, data=sanitized_edges)

insert_return_batch(docs, class_name)

Insert nodes and return their properties.

Parameters:

Name	Type	Description	Default
docs	list[dict[str, Any]]	Documents to insert	required
class_name	str	Label to insert into	required

Raises:

Type	Description
NotImplementedError	This method is not fully implemented for FalkorDB

Source code in graflo/db/falkordb/conn.py

def insert_return_batch(
    self, docs: list[dict[str, Any]], class_name: str
) -> list[dict[str, Any]] | str:
    """Insert nodes and return their properties.

    Args:
        docs: Documents to insert
        class_name: Label to insert into

    Raises:
        NotImplementedError: This method is not fully implemented for FalkorDB
    """
    raise NotImplementedError("insert_return_batch is not implemented for FalkorDB")

keep_absent_documents(batch, class_name, match_keys, keep_keys=None, filters=None)

Keep documents that don't exist in the database.

Parameters:

Name	Type	Description	Default
batch	list[dict[str, Any]]	Batch of documents to check	required
class_name	str	Label to check in	required
match_keys	list[str] | tuple[str, ...]	Keys to match nodes	required
keep_keys	list[str] | tuple[str, ...] | None	Keys to keep in result	None
filters	list[Any] | dict[str, Any] | None	Additional query filters	None

Returns:

Type	Description
list[dict[str, Any]]	Documents that don't exist in the database

Source code in graflo/db/falkordb/conn.py

def keep_absent_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """Keep documents that don't exist in the database.

    Args:
        batch: Batch of documents to check
        class_name: Label to check in
        match_keys: Keys to match nodes
        keep_keys: Keys to keep in result
        filters: Additional query filters

    Returns:
        Documents that don't exist in the database
    """
    if not batch:
        return []

    # Find documents that exist
    present_docs = self.fetch_present_documents(
        batch, class_name, match_keys, match_keys, filters=filters
    )

    # Create a set of present document keys for efficient lookup
    present_keys = set()
    for doc in present_docs:
        key_tuple = tuple(doc.get(k) for k in match_keys)
        present_keys.add(key_tuple)

    # Filter out documents that exist
    absent_docs = []
    for doc in batch:
        key_tuple = tuple(doc.get(k) for k in match_keys)
        if key_tuple not in present_keys:
            if keep_keys:
                absent_docs.append({k: doc.get(k) for k in keep_keys})
            else:
                absent_docs.append(doc)

    return absent_docs

upsert_docs_batch(docs, class_name, match_keys, **kwargs)

Upsert a batch of nodes using Cypher MERGE.

Parameters:

Name	Type	Description	Default
docs	list[dict[str, Any]]	List of node documents to upsert	required
class_name	str	Label to upsert into	required
match_keys	list[str] | tuple[str, ...]	Keys to match for upsert operation	required
**kwargs	Any	Additional options: - dry (bool): If True, build query but don't execute	{}

Raises:

Type	Description
ValueError	If any document is missing a required match_key or has None value

Source code in graflo/db/falkordb/conn.py

def upsert_docs_batch(
    self,
    docs: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    **kwargs: Any,
) -> None:
    """Upsert a batch of nodes using Cypher MERGE.

    Args:
        docs: List of node documents to upsert
        class_name: Label to upsert into
        match_keys: Keys to match for upsert operation
        **kwargs: Additional options:
            - dry (bool): If True, build query but don't execute

    Raises:
        ValueError: If any document is missing a required match_key or has None value
    """
    dry = kwargs.pop("dry", False)

    if not docs:
        return

    # Sanitize documents: filter invalid keys/values, validate match_keys
    match_keys_list = (
        list(match_keys) if isinstance(match_keys, tuple) else match_keys
    )
    sanitized_docs = self._sanitize_batch(docs, match_keys_list)

    if not sanitized_docs:
        return

    # Build the MERGE clause with match keys
    index_str = ", ".join([f"{k}: row.{k}" for k in match_keys])
    # Use 'data' instead of 'batch' as parameter name (batch might be reserved)
    # Try direct UNWIND - FalkorDB may not support WITH $param AS alias
    # Datetime objects are converted to ISO-8601 strings during sanitization
    q = f"""
        UNWIND $data AS row
        MERGE (n:{class_name} {{ {index_str} }})
        ON MATCH SET n += row
        ON CREATE SET n += row
    """
    if not dry:
        self.execute(q, data=sanitized_docs)

MemgraphConnection

Bases: Connection

Memgraph connector implementing the graflo Connection interface.

Provides complete graph database operations for Memgraph including node/relationship CRUD, batch operations, aggregations, and raw Cypher query execution.

Thread Safety

This class is NOT thread-safe. Each thread should use its own connection instance. For concurrent access, use ConnectionManager with separate instances per thread.

Error Handling

• Connection errors raise on instantiation
• Query errors propagate as mgclient.DatabaseError
• Invalid inputs raise ValueError with descriptive messages

Attributes

flavor : DBType Database type identifier (DBType.MEMGRAPH) config : MemgraphConfig Connection configuration (URI, credentials) conn : mgclient.Connection Underlying Memgraph connection instance

Examples

Direct instantiation (prefer ConnectionManager for production)::

config = MemgraphConfig(uri="bolt://localhost:7687")
conn = MemgraphConnection(config)
try:
    result = conn.execute("MATCH (n) RETURN count(n)")
finally:
    conn.close()

Source code in graflo/db/memgraph/conn.py

class MemgraphConnection(Connection):
    """Memgraph connector implementing the graflo Connection interface.

    Provides complete graph database operations for Memgraph including
    node/relationship CRUD, batch operations, aggregations, and raw
    Cypher query execution.

    Thread Safety
    -------------
    This class is NOT thread-safe. Each thread should use its own
    connection instance. For concurrent access, use ConnectionManager
    with separate instances per thread.

    Error Handling
    --------------
    - Connection errors raise on instantiation
    - Query errors propagate as mgclient.DatabaseError
    - Invalid inputs raise ValueError with descriptive messages

    Attributes
    ----------
    flavor : DBType
        Database type identifier (DBType.MEMGRAPH)
    config : MemgraphConfig
        Connection configuration (URI, credentials)
    conn : mgclient.Connection
        Underlying Memgraph connection instance

    Examples
    --------
    Direct instantiation (prefer ConnectionManager for production)::

        config = MemgraphConfig(uri="bolt://localhost:7687")
        conn = MemgraphConnection(config)
        try:
            result = conn.execute("MATCH (n) RETURN count(n)")
        finally:
            conn.close()
    """

    flavor = DBType.MEMGRAPH

    # Type annotations for instance attributes
    conn: mgclient.Connection | None
    _database_name: str

    def __init__(self, config: MemgraphConfig):
        """Initialize Memgraph connection.

        Establishes connection to the Memgraph instance.

        Parameters
        ----------
        config : MemgraphConfig
            Connection configuration with the following fields:
            - uri: Bolt URI (bolt://host:port)
            - username: Username (optional)
            - password: Password (optional)

        Raises
        ------
        ValueError
            If URI is not provided in configuration
        mgclient.DatabaseError
            If unable to connect to Memgraph instance
        """
        super().__init__()
        self.config = config

        if config.uri is None:
            raise ValueError("Memgraph connection requires a URI to be configured")

        # Parse URI to extract host and port
        parsed = urlparse(config.uri)
        host = parsed.hostname or "localhost"
        port = parsed.port or 7687

        # Initialize Memgraph connection
        connect_kwargs: dict[str, Any] = {
            "host": host,
            "port": port,
        }

        if config.username:
            connect_kwargs["username"] = config.username
        if config.password:
            connect_kwargs["password"] = config.password

        self.conn = mgclient.connect(**connect_kwargs)
        self.conn.autocommit = True
        self._database_name = config.database or "memgraph"

    def execute(self, query: str, **kwargs) -> QueryResult:
        """Execute a raw OpenCypher query against the database.

        Executes the provided Cypher query with optional parameters.
        Parameters are safely injected using Memgraph's parameterized
        query mechanism to prevent injection attacks.

        Parameters
        ----------
        query : str
            Cypher query string to execute
        **kwargs
            Query parameters to be safely injected

        Returns
        -------
        QueryResult
            Result object with result_set (list of tuples) and columns

        Examples
        --------
        Simple query::

            result = conn.execute("MATCH (n:Person) RETURN n.name")
            for row in result.result_set:
                print(row[0])  # Access by index

        Parameterized query::

            result = conn.execute(
                "MATCH (n:Person) WHERE n.age > $min_age RETURN n",
                min_age=21
            )
        """
        assert self.conn is not None, "Connection is closed"
        cursor = self.conn.cursor()
        try:
            if kwargs:
                cursor.execute(query, kwargs)
            else:
                cursor.execute(query)
            # mgclient uses Column objects with .name attribute, not tuples
            columns = (
                [col.name for col in cursor.description] if cursor.description else []
            )
            rows = []
            for row in cursor.fetchall():
                processed_row = []
                for value in row:
                    # Convert Memgraph Node/Relationship objects to dicts
                    if hasattr(value, "properties"):
                        processed_row.append(dict(value.properties))
                    else:
                        processed_row.append(value)
                rows.append(tuple(processed_row))
            return QueryResult(columns, rows)
        finally:
            cursor.close()

    def close(self):
        """Close the Memgraph connection."""
        if self.conn is not None:
            self.conn.close()
        self.conn = None

    def create_database(self, name: str):
        """Create a new database (no-op for Memgraph).

        Memgraph uses a single database per instance.
        This method is provided for interface compatibility.

        Args:
            name: Database name (stored for reference)
        """
        self._database_name = name
        logger.info(f"Database name set to '{name}' (Memgraph uses single database)")

    def delete_database(self, name: str):
        """Delete all data from the database.

        Since Memgraph uses a single database, this clears all data.

        Args:
            name: Database name (ignored, clears current database)
        """
        assert self.conn is not None, "Connection is closed"
        try:
            cursor = self.conn.cursor()
            cursor.execute("MATCH (n) DETACH DELETE n")
            cursor.close()
            logger.info("Successfully cleared all data from Memgraph")
        except Exception as e:
            logger.error(f"Failed to clear Memgraph data: {e}", exc_info=True)
            raise

    @staticmethod
    def _is_valid_property_value(value: Any) -> bool:
        """Validate that a value can be stored as a Memgraph property.

        Memgraph cannot store special float values (NaN, Inf) or null bytes.

        Parameters
        ----------
        value
            Value to validate

        Returns
        -------
        bool
            True if value is valid for storage
        """
        if isinstance(value, float):
            if math.isnan(value) or math.isinf(value):
                return False
        if isinstance(value, str):
            if "\x00" in value:
                return False
        return True

    @staticmethod
    def _is_valid_property_key(key: Any) -> bool:
        """Validate that a key can be used as a property name.

        Property keys must be non-empty strings that don't start with
        reserved prefixes.

        Parameters
        ----------
        key
            Key to validate

        Returns
        -------
        bool
            True if key is valid
        """
        if not isinstance(key, str):
            return False
        if not key:
            return False
        if key.startswith("_"):
            return True
        return True

    def _sanitize_doc(self, doc: dict) -> dict:
        """Sanitize a document for safe storage.

        Removes invalid keys and values, logs warnings for skipped items.

        Parameters
        ----------
        doc : dict
            Document to sanitize

        Returns
        -------
        dict
            Sanitized document
        """
        sanitized = {}
        for key, value in doc.items():
            if not self._is_valid_property_key(key):
                logger.warning(f"Skipping invalid property key: {key}")
                continue
            if not self._is_valid_property_value(value):
                logger.warning(f"Skipping property '{key}' with invalid value: {value}")
                continue
            sanitized[key] = value
        return sanitized

    def _sanitize_batch(self, docs: list[dict], match_keys: list[str]) -> list[dict]:
        """Sanitize a batch of documents.

        Parameters
        ----------
        docs : list[dict]
            Documents to sanitize
        match_keys : list[str]
            Required keys that must be present

        Returns
        -------
        list[dict]
            List of sanitized documents (invalid documents are skipped)
        """
        sanitized = []
        for doc in docs:
            clean_doc = self._sanitize_doc(doc)
            # Verify all match keys are present and valid
            valid = True
            for key in match_keys:
                if key not in clean_doc or clean_doc[key] is None:
                    logger.warning(
                        f"Document missing required match_key '{key}': {doc}"
                    )
                    valid = False
                    break
            if valid:
                sanitized.append(clean_doc)
        return sanitized

    def define_vertex_indices(self, vertex_config: VertexConfig):
        """Create indices for vertex labels based on configuration.

        Iterates through all vertices defined in the configuration and creates
        indices for each specified field. Memgraph supports indices on node
        properties for faster lookups.

        Parameters
        ----------
        vertex_config : VertexConfig
            Vertex configuration containing vertices and their index definitions.
            Each vertex may have multiple indices, each covering one or more fields.

        Notes
        -----
        - Index creation is idempotent (existing indices are skipped)
        - Uses Memgraph syntax: ``CREATE INDEX ON :Label(property)``
        - Errors are logged but don't stop processing of other indices
        """
        assert self.conn is not None, "Connection is closed"

        for label in vertex_config.vertex_set:
            for index_obj in vertex_config.indexes(label):
                for field in index_obj.fields:
                    try:
                        query = f"CREATE INDEX ON :{label}({field})"
                        cursor = self.conn.cursor()
                        cursor.execute(query)
                        cursor.close()
                        logger.debug(f"Created index on {label}.{field}")
                    except Exception as e:
                        if "already exists" in str(e).lower():
                            logger.debug(f"Index on {label}.{field} already exists")
                        else:
                            logger.warning(
                                f"Failed to create index on {label}.{field}: {e}"
                            )

    def define_edge_indices(self, edges: list[Edge]):
        """Create indices for edge types.

        Memgraph doesn't support relationship property indices in the same way,
        so this creates indices on the relationship properties if defined.

        Parameters
        ----------
        edges : list[Edge]
            List of edge configurations
        """
        assert self.conn is not None, "Connection is closed"
        for edge in edges:
            if edge.relation is None:
                continue
            for idx in edge.indexes:
                for field in idx.fields:
                    try:
                        # Create index on relationship type
                        query = f"CREATE INDEX ON :{edge.relation}({field})"
                        cursor = self.conn.cursor()
                        cursor.execute(query)
                        cursor.close()
                        logger.debug(
                            f"Created index on relationship {edge.relation}.{field}"
                        )
                    except Exception as e:
                        if "already exists" in str(e).lower():
                            logger.debug(
                                f"Index on {edge.relation}.{field} already exists"
                            )
                        else:
                            logger.debug(
                                f"Could not create index on {edge.relation}.{field}: {e}"
                            )

    def delete_graph_structure(
        self,
        vertex_types: tuple[str, ...] | list[str] = (),
        graph_names: tuple[str, ...] | list[str] = (),
        delete_all: bool = False,
    ) -> None:
        """Delete graph structure (nodes and relationships).

        Parameters
        ----------
        vertex_types : list[str], optional
            Specific node labels to delete
        edge_types : list[str], optional
            Specific relationship types to delete (not used, deletes via nodes)
        graph_names : list[str], optional
            Not applicable for Memgraph (single database)
        delete_all : bool
            If True, delete all nodes and relationships
        """
        assert self.conn is not None, "Connection is closed"

        if delete_all:
            cursor = self.conn.cursor()
            cursor.execute("MATCH (n) DETACH DELETE n")
            cursor.close()
            logger.info("Deleted all nodes and relationships")
            return

        # Convert tuple to list if needed
        vertex_types_list = (
            list(vertex_types) if isinstance(vertex_types, tuple) else vertex_types
        )
        if vertex_types_list:
            for label in vertex_types_list:
                try:
                    cursor = self.conn.cursor()
                    cursor.execute(f"MATCH (n:{label}) DETACH DELETE n")
                    cursor.close()
                    logger.debug(f"Deleted all nodes with label '{label}'")
                except Exception as e:
                    logger.warning(f"Failed to delete nodes with label '{label}': {e}")

    def init_db(self, schema: Schema, recreate_schema: bool) -> None:
        """Initialize Memgraph with the given schema.

        If the database already has nodes and recreate_schema is False, raises
        SchemaExistsError and the script halts.

        Parameters
        ----------
        schema : Schema
            Schema containing graph structure definitions
        recreate_schema : bool
            If True, delete all existing data before initialization.
            If False and database has nodes, raises SchemaExistsError.
        """
        assert self.conn is not None, "Connection is closed"

        self._database_name = schema.general.name
        logger.info(f"Initialized Memgraph with schema '{self._database_name}'")

        # Check if database already has nodes (schema/graph exists)
        cursor = self.conn.cursor()
        cursor.execute("MATCH (n) RETURN count(n) AS c")
        row = cursor.fetchone()
        cursor.close()
        count = 0
        if row is not None:
            count = (
                row[0]
                if isinstance(row, (list, tuple))
                else getattr(row, "c", row.get("c", 0) if hasattr(row, "get") else 0)
            )
        if count > 0 and not recreate_schema:
            raise SchemaExistsError(
                f"Schema/graph already exists ({count} nodes). "
                "Set recreate_schema=True to replace, or use clear_data=True before ingestion."
            )

        if recreate_schema:
            try:
                self.delete_graph_structure(delete_all=True)
            except Exception as e:
                logger.warning(f"Error clearing data on recreate_schema: {e}")

    def clear_data(self, schema: Schema) -> None:
        """Remove all data from the graph without dropping the schema.

        Deletes all nodes and relationships; labels (schema) remain.
        """
        self.delete_graph_structure(delete_all=True)

    def upsert_docs_batch(
        self,
        docs: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        **kwargs: Any,
    ) -> None:
        """Upsert a batch of nodes using Cypher MERGE.

        Performs atomic upsert (update-or-insert) operations on a batch of
        documents. Uses Cypher MERGE with ON MATCH/ON CREATE for efficiency.
        Existing properties not in the document are preserved on update.

        Parameters
        ----------
        docs : list[dict]
            Documents to upsert. Each document must contain all match_keys
            with non-None values. Invalid documents are skipped with a warning.
        class_name : str
            Node label to create/update (e.g., "Person", "Product")
        match_keys : list[str]
            Property keys used to identify existing nodes for update.
            Supports composite keys (multiple fields).
        **kwargs
            Additional options:
            - dry (bool): If True, build query but don't execute (for debugging)

        Raises
        ------
        ValueError
            If any document is missing a required match_key or has None value.

        Notes
        -----
        - Documents are sanitized before insertion (invalid keys/values removed)
        - NaN, Inf, and null bytes are automatically filtered with warnings
        - The operation is atomic per batch (all succeed or fail together)

        Examples
        --------
        Insert or update Person nodes::

            db.upsert_docs_batch(
                [
                    {"id": "1", "name": "Alice", "age": 30},
                    {"id": "2", "name": "Bob", "age": 25},
                ],
                class_name="Person",
                match_keys=["id"]
            )

        With composite key::

            db.upsert_docs_batch(
                [{"tenant": "acme", "user_id": "u1", "email": "a@b.com"}],
                class_name="User",
                match_keys=["tenant", "user_id"]
            )
        """
        assert self.conn is not None, "Connection is closed"
        dry = kwargs.pop("dry", False)

        if not docs:
            return

        # Convert tuple to list if needed
        match_keys_list = (
            list(match_keys) if isinstance(match_keys, tuple) else match_keys
        )
        # Sanitize documents
        sanitized_docs = self._sanitize_batch(docs, match_keys_list)

        if not sanitized_docs:
            return

        # Auto-create index on match_keys for MERGE performance (idempotent)
        cursor = self.conn.cursor()
        for key in match_keys:
            try:
                cursor.execute(f"CREATE INDEX ON :{class_name}({key})")
                logger.debug(f"Created index on {class_name}.{key}")
            except Exception as e:
                if "already exists" not in str(e).lower():
                    logger.debug(f"Index on {class_name}.{key}: {e}")
        cursor.close()

        # Build the MERGE clause with match keys
        index_str = ", ".join([f"{k}: row.{k}" for k in match_keys_list])
        q = f"""
            UNWIND $batch AS row
            MERGE (n:{class_name} {{ {index_str} }})
            ON MATCH SET n += row
            ON CREATE SET n += row
        """
        if not dry:
            cursor = self.conn.cursor()
            cursor.execute(q, {"batch": sanitized_docs})
            cursor.close()

    def insert_edges_batch(
        self,
        docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
        source_class: str,
        target_class: str,
        relation_name: str,
        match_keys_source: tuple[str, ...],
        match_keys_target: tuple[str, ...],
        filter_uniques: bool = True,
        head: int | None = None,
        **kwargs: Any,
    ) -> None:
        """Insert a batch of edges using Cypher MERGE.

        Creates relationships between existing nodes by matching source and
        target nodes using the specified match keys, then creating or updating
        the relationship between them.

        Parameters
        ----------
        docs_edges : list[list[dict]]
            Edge specifications as list of [source_doc, target_doc, edge_props]:
            ``[[{source_props}, {target_props}, {edge_props}], ...]``
            - source_props: Properties to match the source node
            - target_props: Properties to match the target node
            - edge_props: Properties to set on the relationship (optional)
        source_class : str
            Label of source nodes (e.g., "Person")
        target_class : str
            Label of target nodes (e.g., "Company")
        relation_name : str
            Relationship type name (e.g., "WORKS_AT")
        collection_name : str, optional
            Not used for Memgraph (kept for interface compatibility)
        match_keys_source : tuple[str, ...]
            Property keys used to identify source nodes (default: ("_key",))
        match_keys_target : tuple[str, ...]
            Property keys used to identify target nodes (default: ("_key",))
        filter_uniques : bool
            Not used for Memgraph (kept for interface compatibility)
        uniq_weight_fields : list[str], optional
            Not used for Memgraph (kept for interface compatibility)
        uniq_weight_collections : list[str], optional
            Not used for Memgraph (kept for interface compatibility)
        **kwargs
            Additional options (currently unused)

        Notes
        -----
        - Edges are created with MERGE, preventing duplicates
        - If source or target node doesn't exist, the edge is silently skipped
        - Edge properties are merged on update (existing props preserved)

        Examples
        --------
        Create relationships between Person and Company nodes::

            db.insert_edges_batch(
                [
                    [{"id": "alice"}, {"id": "acme"}, {"role": "engineer"}],
                    [{"id": "bob"}, {"id": "acme"}, {"role": "manager"}],
                ],
                source_class="Person",
                target_class="Company",
                relation_name="WORKS_AT",
                match_keys_source=("id",),
                match_keys_target=("id",),
            )
        """
        assert self.conn is not None, "Connection is closed"

        if not docs_edges:
            return

        # Handle head limit if specified
        if head is not None and head > 0:
            docs_edges = docs_edges[:head]

        # Build batch data
        batch = []
        for edge_data in docs_edges:
            if len(edge_data) < 2:
                continue

            source_doc = edge_data[0]
            target_doc = edge_data[1]
            edge_props = edge_data[2] if len(edge_data) > 2 else {}

            # Sanitize
            source_doc = self._sanitize_doc(source_doc)
            target_doc = self._sanitize_doc(target_doc)
            edge_props = self._sanitize_doc(edge_props) if edge_props else {}

            batch.append(
                {
                    "source": source_doc,
                    "target": target_doc,
                    "props": edge_props,
                }
            )

        if not batch:
            return

        # Build match patterns
        source_match = ", ".join([f"{k}: row.source.{k}" for k in match_keys_source])
        target_match = ", ".join([f"{k}: row.target.{k}" for k in match_keys_target])

        q = f"""
            UNWIND $batch AS row
            MATCH (s:{source_class} {{ {source_match} }})
            MATCH (t:{target_class} {{ {target_match} }})
            MERGE (s)-[r:{relation_name}]->(t)
            ON CREATE SET r = row.props
            ON MATCH SET r += row.props
        """

        cursor = self.conn.cursor()
        cursor.execute(q, {"batch": batch})
        cursor.close()

    def fetch_docs(
        self,
        class_name: str,
        filters: list | dict | None = None,
        limit: int | None = None,
        return_keys: list[str] | None = None,
        unset_keys: list[str] | None = None,
        **kwargs,
    ) -> list[dict]:
        """Fetch nodes from the database with optional filtering and projection.

        Retrieves nodes matching the specified label and optional filter
        conditions. Supports field projection to return only specific properties.

        Parameters
        ----------
        class_name : str
            Node label to fetch (e.g., "Person", "Product")
        filters : list | dict, optional
            Query filters in graflo expression format.
            Examples: ``["==", "Alice", "name"]`` or ``["AND", [...], [...]]``
        limit : int, optional
            Maximum number of results to return. If None or <= 0, returns all.
        return_keys : list[str], optional
            Properties to include in results (projection). If None, returns
            all properties. Example: ``["id", "name"]``
        unset_keys : list[str], optional
            Not used for Memgraph (kept for interface compatibility)
        **kwargs
            Additional options (currently unused)

        Returns
        -------
        list[dict]
            List of node property dictionaries. Each dict contains the
            requested properties (or all properties if no projection).

        Examples
        --------
        Fetch all Person nodes::

            results = db.fetch_docs("Person")

        Fetch with filter and projection::

            results = db.fetch_docs(
                "Person",
                filters=["==", "Alice", "name"],
                return_keys=["id", "name", "email"],
                limit=10
            )

        Fetch with complex filter::

            results = db.fetch_docs(
                "Product",
                filters=["AND", [">", 100, "price"], ["==", "active", "status"]]
            )
        """
        assert self.conn is not None, "Connection is closed"

        q = f"MATCH (n:{class_name})"

        if filters is not None:
            ff = FilterExpression.from_dict(filters)
            filter_str = ff(doc_name="n", kind=self.expression_flavor())
            q += f" WHERE {filter_str}"

        # Handle projection
        if return_keys:
            return_clause = ", ".join([f"n.{k} AS {k}" for k in return_keys])
            q += f" RETURN {return_clause}"
        else:
            q += " RETURN n"

        if limit is not None and limit > 0:
            q += f" LIMIT {limit}"

        cursor = self.conn.cursor()
        cursor.execute(q)
        results = []

        if return_keys:
            # With projection, build dict from column values
            for row in cursor.fetchall():
                result = {return_keys[i]: row[i] for i in range(len(return_keys))}
                results.append(result)
        else:
            # Without projection, extract node properties
            for row in cursor.fetchall():
                node = row[0]
                if hasattr(node, "properties"):
                    results.append(dict(node.properties))
                else:
                    results.append(node)

        cursor.close()
        return results

    def fetch_edges(
        self,
        from_type: str,
        from_id: str,
        edge_type: str | None = None,
        to_type: str | None = None,
        to_id: str | None = None,
        filters: list[Any] | dict[str, Any] | None = None,
        limit: int | None = None,
        return_keys: list[str] | None = None,
        unset_keys: list[str] | None = None,
        **kwargs: Any,
    ) -> list[dict[str, Any]]:
        """Fetch edges from Memgraph using Cypher.

        Retrieves relationships starting from a specific node, optionally filtered
        by edge type, target node type, and target node ID.

        Args:
            from_type: Source node label (e.g., "Person")
            from_id: Source node ID (property value, typically "id" or "_key")
            edge_type: Optional relationship type to filter by (e.g., "WORKS_AT")
            to_type: Optional target node label to filter by
            to_id: Optional target node ID to filter by
            filters: Additional query filters applied to relationship properties
            limit: Maximum number of edges to return
            return_keys: Keys to return (projection) - not fully supported
            unset_keys: Keys to exclude (projection) - not fully supported
            **kwargs: Additional options

        Returns:
            list: List of fetched edges as dictionaries
        """
        assert self.conn is not None, "Connection is closed"

        # Build Cypher query starting from the source node
        # Use id property (common in Memgraph) or _key if needed
        q = f"MATCH (s:{from_type} {{id: $from_id}})"

        # Build relationship pattern
        if edge_type:
            rel_pattern = f"-[r:{edge_type}]->"
        else:
            rel_pattern = "-[r]->"

        # Build target node match
        if to_type:
            target_match = f"(t:{to_type})"
        else:
            target_match = "(t)"

        q += f" {rel_pattern} {target_match}"

        # Build WHERE clauses
        where_clauses = []
        if to_id:
            where_clauses.append("t.id = $to_id")

        # Add relationship property filters
        if filters is not None:
            ff = FilterExpression.from_dict(filters)
            filter_str = ff(doc_name="r", kind=self.expression_flavor())
            where_clauses.append(filter_str)

        if where_clauses:
            q += f" WHERE {' AND '.join(where_clauses)}"

        # Build RETURN clause
        # Default: return relationship properties and basic node info
        if return_keys:
            # If return_keys specified, try to return those fields
            return_fields = []
            for key in return_keys:
                if key.startswith("from_") or key.startswith("source_"):
                    return_fields.append(
                        f"s.{key.replace('from_', '').replace('source_', '')} AS {key}"
                    )
                elif key.startswith("to_") or key.startswith("target_"):
                    return_fields.append(
                        f"t.{key.replace('to_', '').replace('target_', '')} AS {key}"
                    )
                else:
                    return_fields.append(f"r.{key} AS {key}")
            q += f" RETURN {', '.join(return_fields)}"
        else:
            # Default: return relationship properties and node IDs
            q += " RETURN properties(r) AS props, s.id AS from_id, t.id AS to_id"

        if limit is not None and limit > 0:
            q += f" LIMIT {limit}"

        # Execute query with parameters
        params: dict[str, Any] = {"from_id": from_id}
        if to_id:
            params["to_id"] = to_id

        cursor = self.conn.cursor()
        cursor.execute(q, params)
        columns = [desc[0] for desc in cursor.description] if cursor.description else []
        results = []
        for row in cursor.fetchall():
            result = {}
            for i, col in enumerate(columns):
                result[col] = row[i]
            # Apply unset_keys if specified
            if unset_keys:
                for key in unset_keys:
                    result.pop(key, None)
            results.append(result)
        cursor.close()
        return results

    def aggregate(
        self,
        class_name: str,
        aggregation_function: AggregationType,
        discriminant: str | None = None,
        aggregated_field: str | None = None,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> int | float | list[dict[str, Any]] | dict[str, int | float] | None:
        """Perform aggregation operations on nodes.

        Computes aggregate statistics over nodes matching the specified label
        and optional filters. Supports counting, min/max, average, and distinct
        value extraction.

        Parameters
        ----------
        class_name : str
            Node label to aggregate (e.g., "Person", "Product")
        aggregation_function : AggregationType
            Type of aggregation to perform:
            - COUNT: Count matching nodes (with optional GROUP BY)
            - MAX: Maximum value of a field
            - MIN: Minimum value of a field
            - AVERAGE: Average value of a field
            - SORTED_UNIQUE: Distinct values sorted ascending
        discriminant : str, optional
            Field to group by when using COUNT. Returns a dictionary mapping
            each distinct value to its count. Example: ``"status"``
        aggregated_field : str, optional
            Field to aggregate. Required for MAX, MIN, AVERAGE, SORTED_UNIQUE.
            Example: ``"price"``
        filters : list | dict, optional
            Query filters to apply before aggregation.
            Example: ``["==", "active", "status"]``
        **kwargs
            Additional options (currently unused)

        Returns
        -------
        int | float | list | dict[str, int] | None
            - COUNT without discriminant: int (total count)
            - COUNT with discriminant: dict mapping values to counts
            - MAX/MIN/AVERAGE: numeric value or None if no data
            - SORTED_UNIQUE: list of distinct values

        Raises
        ------
        ValueError
            If aggregated_field is missing for non-COUNT aggregations,
            or if an unsupported aggregation type is specified.

        Examples
        --------
        Count all Person nodes::

            count = db.aggregate("Person", AggregationType.COUNT)
            # Returns: 42

        Count by status (GROUP BY)::

            by_status = db.aggregate(
                "Product",
                AggregationType.COUNT,
                discriminant="status"
            )
            # Returns: {"active": 100, "inactive": 25, "draft": 5}

        Get maximum price::

            max_price = db.aggregate(
                "Product",
                AggregationType.MAX,
                aggregated_field="price"
            )
            # Returns: 999.99

        Get distinct categories::

            categories = db.aggregate(
                "Product",
                AggregationType.SORTED_UNIQUE,
                aggregated_field="category"
            )
            # Returns: ["electronics", "furniture", "toys"]
        """
        assert self.conn is not None, "Connection is closed"

        # Build filter clause
        filter_clause = ""
        if filters is not None:
            ff = FilterExpression.from_dict(filters)
            filter_str = ff(doc_name="n", kind=self.expression_flavor())
            filter_clause = f" WHERE {filter_str}"

        q = f"MATCH (n:{class_name}){filter_clause}"

        if aggregation_function == AggregationType.COUNT:
            if discriminant:
                q += f" RETURN n.{discriminant} AS key, count(*) AS count"
                cursor = self.conn.cursor()
                cursor.execute(q)
                rows = cursor.fetchall()
                cursor.close()
                return {row[0]: row[1] for row in rows}
            else:
                q += " RETURN count(n)"
        elif aggregation_function == AggregationType.MAX:
            q += f" RETURN max(n.{aggregated_field})"
        elif aggregation_function == AggregationType.MIN:
            q += f" RETURN min(n.{aggregated_field})"
        elif aggregation_function == AggregationType.AVERAGE:
            q += f" RETURN avg(n.{aggregated_field})"
        elif aggregation_function == AggregationType.SORTED_UNIQUE:
            q += f" RETURN DISTINCT n.{aggregated_field} ORDER BY n.{aggregated_field}"
        else:
            raise ValueError(f"Unsupported aggregation type: {aggregation_function}")

        cursor = self.conn.cursor()
        cursor.execute(q)
        rows = cursor.fetchall()
        cursor.close()

        if aggregation_function == AggregationType.SORTED_UNIQUE:
            return [row[0] for row in rows]
        return rows[0][0] if rows else None

    def define_schema(self, schema: Schema):
        """Define collections based on schema.

        Note: This is a no-op in Memgraph as collections are implicit.
        Labels and relationship types are created when data is inserted.

        Parameters
        ----------
        schema : Schema
            Schema containing collection definitions
        """
        pass

    def insert_return_batch(
        self, docs: list[dict[str, Any]], class_name: str
    ) -> list[dict[str, Any]] | str:
        """Insert nodes and return their properties.

        Parameters
        ----------
        docs : list[dict]
            Documents to insert
        class_name : str
            Label to insert into

        Returns
        -------
        list[dict] | str
            Inserted documents with their properties, or query string

        Raises
        ------
        NotImplementedError
            This method is not fully implemented for Memgraph
        """
        raise NotImplementedError("insert_return_batch is not implemented for Memgraph")

    def fetch_present_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        flatten: bool = False,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Fetch nodes that exist in the database.

        Parameters
        ----------
        batch : list[dict]
            Batch of documents to check
        class_name : str
            Label to check in
        match_keys : list[str]
            Keys to match nodes
        keep_keys : list[str]
            Keys to keep in result
        flatten : bool
            Unused in Memgraph
        filters : list | dict, optional
            Additional query filters

        Returns
        -------
        list[dict]
            Documents that exist in the database
        """
        if not batch:
            return []

        assert self.conn is not None, "Connection is closed"
        results = []

        for doc in batch:
            # Build match conditions
            match_conditions = " AND ".join([f"n.{key} = ${key}" for key in match_keys])
            params = {key: doc.get(key) for key in match_keys}

            # Build return clause with keep_keys
            if keep_keys:
                return_clause = ", ".join([f"n.{k} AS {k}" for k in keep_keys])
            else:
                return_clause = "n"

            q = f"MATCH (n:{class_name}) WHERE {match_conditions} RETURN {return_clause} LIMIT 1"

            cursor = self.conn.cursor()
            cursor.execute(q, params)
            rows = cursor.fetchall()
            cursor.close()

            if rows:
                if keep_keys:
                    result = {keep_keys[i]: rows[0][i] for i in range(len(keep_keys))}
                else:
                    node = rows[0][0]
                    if hasattr(node, "properties"):
                        result = dict(node.properties)
                    else:
                        result = node
                results.append(result)

        return results

    def keep_absent_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Keep documents that don't exist in the database.

        Parameters
        ----------
        batch : list[dict]
            Batch of documents to check
        class_name : str
            Label to check in
        match_keys : list[str]
            Keys to match nodes
        keep_keys : list[str]
            Keys to keep in result
        filters : list | dict, optional
            Additional query filters

        Returns
        -------
        list[dict]
            Documents that don't exist in the database
        """
        if not batch:
            return []

        # Find documents that exist
        present_docs = self.fetch_present_documents(
            batch, class_name, match_keys, match_keys, filters=filters
        )

        # Create a set of present document keys for efficient lookup
        present_keys = set()
        for doc in present_docs:
            key_tuple = tuple(doc.get(k) for k in match_keys)
            present_keys.add(key_tuple)

        # Keep documents that don't exist
        absent = []
        for doc in batch:
            key_tuple = tuple(doc.get(k) for k in match_keys)
            if key_tuple not in present_keys:
                if keep_keys:
                    absent.append({k: doc.get(k) for k in keep_keys})
                else:
                    absent.append(doc)

        return absent

__init__(config)

Initialize Memgraph connection.

Establishes connection to the Memgraph instance.

Parameters

config : MemgraphConfig Connection configuration with the following fields: - uri: Bolt URI (bolt://host:port) - username: Username (optional) - password: Password (optional)

Raises

ValueError If URI is not provided in configuration mgclient.DatabaseError If unable to connect to Memgraph instance

Source code in graflo/db/memgraph/conn.py

def __init__(self, config: MemgraphConfig):
    """Initialize Memgraph connection.

    Establishes connection to the Memgraph instance.

    Parameters
    ----------
    config : MemgraphConfig
        Connection configuration with the following fields:
        - uri: Bolt URI (bolt://host:port)
        - username: Username (optional)
        - password: Password (optional)

    Raises
    ------
    ValueError
        If URI is not provided in configuration
    mgclient.DatabaseError
        If unable to connect to Memgraph instance
    """
    super().__init__()
    self.config = config

    if config.uri is None:
        raise ValueError("Memgraph connection requires a URI to be configured")

    # Parse URI to extract host and port
    parsed = urlparse(config.uri)
    host = parsed.hostname or "localhost"
    port = parsed.port or 7687

    # Initialize Memgraph connection
    connect_kwargs: dict[str, Any] = {
        "host": host,
        "port": port,
    }

    if config.username:
        connect_kwargs["username"] = config.username
    if config.password:
        connect_kwargs["password"] = config.password

    self.conn = mgclient.connect(**connect_kwargs)
    self.conn.autocommit = True
    self._database_name = config.database or "memgraph"

aggregate(class_name, aggregation_function, discriminant=None, aggregated_field=None, filters=None)

Perform aggregation operations on nodes.

Computes aggregate statistics over nodes matching the specified label and optional filters. Supports counting, min/max, average, and distinct value extraction.

Parameters

class_name : str Node label to aggregate (e.g., "Person", "Product") aggregation_function : AggregationType Type of aggregation to perform: - COUNT: Count matching nodes (with optional GROUP BY) - MAX: Maximum value of a field - MIN: Minimum value of a field - AVERAGE: Average value of a field - SORTED_UNIQUE: Distinct values sorted ascending discriminant : str, optional Field to group by when using COUNT. Returns a dictionary mapping each distinct value to its count. Example: "status" aggregated_field : str, optional Field to aggregate. Required for MAX, MIN, AVERAGE, SORTED_UNIQUE. Example: "price" filters : list | dict, optional Query filters to apply before aggregation. Example: ["==", "active", "status"] **kwargs Additional options (currently unused)

Returns

int | float | list | dict[str, int] | None - COUNT without discriminant: int (total count) - COUNT with discriminant: dict mapping values to counts - MAX/MIN/AVERAGE: numeric value or None if no data - SORTED_UNIQUE: list of distinct values

Raises

ValueError If aggregated_field is missing for non-COUNT aggregations, or if an unsupported aggregation type is specified.

Examples

Count all Person nodes::

count = db.aggregate("Person", AggregationType.COUNT)
# Returns: 42

Count by status (GROUP BY)::

by_status = db.aggregate(
    "Product",
    AggregationType.COUNT,
    discriminant="status"
)
# Returns: {"active": 100, "inactive": 25, "draft": 5}

Get maximum price::

max_price = db.aggregate(
    "Product",
    AggregationType.MAX,
    aggregated_field="price"
)
# Returns: 999.99

Get distinct categories::

categories = db.aggregate(
    "Product",
    AggregationType.SORTED_UNIQUE,
    aggregated_field="category"
)
# Returns: ["electronics", "furniture", "toys"]

Source code in graflo/db/memgraph/conn.py

def aggregate(
    self,
    class_name: str,
    aggregation_function: AggregationType,
    discriminant: str | None = None,
    aggregated_field: str | None = None,
    filters: list[Any] | dict[str, Any] | None = None,
) -> int | float | list[dict[str, Any]] | dict[str, int | float] | None:
    """Perform aggregation operations on nodes.

    Computes aggregate statistics over nodes matching the specified label
    and optional filters. Supports counting, min/max, average, and distinct
    value extraction.

    Parameters
    ----------
    class_name : str
        Node label to aggregate (e.g., "Person", "Product")
    aggregation_function : AggregationType
        Type of aggregation to perform:
        - COUNT: Count matching nodes (with optional GROUP BY)
        - MAX: Maximum value of a field
        - MIN: Minimum value of a field
        - AVERAGE: Average value of a field
        - SORTED_UNIQUE: Distinct values sorted ascending
    discriminant : str, optional
        Field to group by when using COUNT. Returns a dictionary mapping
        each distinct value to its count. Example: ``"status"``
    aggregated_field : str, optional
        Field to aggregate. Required for MAX, MIN, AVERAGE, SORTED_UNIQUE.
        Example: ``"price"``
    filters : list | dict, optional
        Query filters to apply before aggregation.
        Example: ``["==", "active", "status"]``
    **kwargs
        Additional options (currently unused)

    Returns
    -------
    int | float | list | dict[str, int] | None
        - COUNT without discriminant: int (total count)
        - COUNT with discriminant: dict mapping values to counts
        - MAX/MIN/AVERAGE: numeric value or None if no data
        - SORTED_UNIQUE: list of distinct values

    Raises
    ------
    ValueError
        If aggregated_field is missing for non-COUNT aggregations,
        or if an unsupported aggregation type is specified.

    Examples
    --------
    Count all Person nodes::

        count = db.aggregate("Person", AggregationType.COUNT)
        # Returns: 42

    Count by status (GROUP BY)::

        by_status = db.aggregate(
            "Product",
            AggregationType.COUNT,
            discriminant="status"
        )
        # Returns: {"active": 100, "inactive": 25, "draft": 5}

    Get maximum price::

        max_price = db.aggregate(
            "Product",
            AggregationType.MAX,
            aggregated_field="price"
        )
        # Returns: 999.99

    Get distinct categories::

        categories = db.aggregate(
            "Product",
            AggregationType.SORTED_UNIQUE,
            aggregated_field="category"
        )
        # Returns: ["electronics", "furniture", "toys"]
    """
    assert self.conn is not None, "Connection is closed"

    # Build filter clause
    filter_clause = ""
    if filters is not None:
        ff = FilterExpression.from_dict(filters)
        filter_str = ff(doc_name="n", kind=self.expression_flavor())
        filter_clause = f" WHERE {filter_str}"

    q = f"MATCH (n:{class_name}){filter_clause}"

    if aggregation_function == AggregationType.COUNT:
        if discriminant:
            q += f" RETURN n.{discriminant} AS key, count(*) AS count"
            cursor = self.conn.cursor()
            cursor.execute(q)
            rows = cursor.fetchall()
            cursor.close()
            return {row[0]: row[1] for row in rows}
        else:
            q += " RETURN count(n)"
    elif aggregation_function == AggregationType.MAX:
        q += f" RETURN max(n.{aggregated_field})"
    elif aggregation_function == AggregationType.MIN:
        q += f" RETURN min(n.{aggregated_field})"
    elif aggregation_function == AggregationType.AVERAGE:
        q += f" RETURN avg(n.{aggregated_field})"
    elif aggregation_function == AggregationType.SORTED_UNIQUE:
        q += f" RETURN DISTINCT n.{aggregated_field} ORDER BY n.{aggregated_field}"
    else:
        raise ValueError(f"Unsupported aggregation type: {aggregation_function}")

    cursor = self.conn.cursor()
    cursor.execute(q)
    rows = cursor.fetchall()
    cursor.close()

    if aggregation_function == AggregationType.SORTED_UNIQUE:
        return [row[0] for row in rows]
    return rows[0][0] if rows else None

clear_data(schema)

Remove all data from the graph without dropping the schema.

Deletes all nodes and relationships; labels (schema) remain.

Source code in graflo/db/memgraph/conn.py

def clear_data(self, schema: Schema) -> None:
    """Remove all data from the graph without dropping the schema.

    Deletes all nodes and relationships; labels (schema) remain.
    """
    self.delete_graph_structure(delete_all=True)

close()

Close the Memgraph connection.

Source code in graflo/db/memgraph/conn.py

def close(self):
    """Close the Memgraph connection."""
    if self.conn is not None:
        self.conn.close()
    self.conn = None

create_database(name)

Create a new database (no-op for Memgraph).

Memgraph uses a single database per instance. This method is provided for interface compatibility.

Parameters:

Name	Type	Description	Default
name	str	Database name (stored for reference)	required

Source code in graflo/db/memgraph/conn.py

def create_database(self, name: str):
    """Create a new database (no-op for Memgraph).

    Memgraph uses a single database per instance.
    This method is provided for interface compatibility.

    Args:
        name: Database name (stored for reference)
    """
    self._database_name = name
    logger.info(f"Database name set to '{name}' (Memgraph uses single database)")

define_edge_indices(edges)

Create indices for edge types.

Memgraph doesn't support relationship property indices in the same way, so this creates indices on the relationship properties if defined.

Parameters

edges : list[Edge] List of edge configurations

Source code in graflo/db/memgraph/conn.py

def define_edge_indices(self, edges: list[Edge]):
    """Create indices for edge types.

    Memgraph doesn't support relationship property indices in the same way,
    so this creates indices on the relationship properties if defined.

    Parameters
    ----------
    edges : list[Edge]
        List of edge configurations
    """
    assert self.conn is not None, "Connection is closed"
    for edge in edges:
        if edge.relation is None:
            continue
        for idx in edge.indexes:
            for field in idx.fields:
                try:
                    # Create index on relationship type
                    query = f"CREATE INDEX ON :{edge.relation}({field})"
                    cursor = self.conn.cursor()
                    cursor.execute(query)
                    cursor.close()
                    logger.debug(
                        f"Created index on relationship {edge.relation}.{field}"
                    )
                except Exception as e:
                    if "already exists" in str(e).lower():
                        logger.debug(
                            f"Index on {edge.relation}.{field} already exists"
                        )
                    else:
                        logger.debug(
                            f"Could not create index on {edge.relation}.{field}: {e}"
                        )

define_schema(schema)

Define collections based on schema.

Note: This is a no-op in Memgraph as collections are implicit. Labels and relationship types are created when data is inserted.

Parameters

schema : Schema Schema containing collection definitions

Source code in graflo/db/memgraph/conn.py

def define_schema(self, schema: Schema):
    """Define collections based on schema.

    Note: This is a no-op in Memgraph as collections are implicit.
    Labels and relationship types are created when data is inserted.

    Parameters
    ----------
    schema : Schema
        Schema containing collection definitions
    """
    pass

define_vertex_indices(vertex_config)

Create indices for vertex labels based on configuration.

Iterates through all vertices defined in the configuration and creates indices for each specified field. Memgraph supports indices on node properties for faster lookups.

Parameters

vertex_config : VertexConfig Vertex configuration containing vertices and their index definitions. Each vertex may have multiple indices, each covering one or more fields.

Notes

• Index creation is idempotent (existing indices are skipped)
• Uses Memgraph syntax: CREATE INDEX ON :Label(property)
• Errors are logged but don't stop processing of other indices

Source code in graflo/db/memgraph/conn.py

def define_vertex_indices(self, vertex_config: VertexConfig):
    """Create indices for vertex labels based on configuration.

    Iterates through all vertices defined in the configuration and creates
    indices for each specified field. Memgraph supports indices on node
    properties for faster lookups.

    Parameters
    ----------
    vertex_config : VertexConfig
        Vertex configuration containing vertices and their index definitions.
        Each vertex may have multiple indices, each covering one or more fields.

    Notes
    -----
    - Index creation is idempotent (existing indices are skipped)
    - Uses Memgraph syntax: ``CREATE INDEX ON :Label(property)``
    - Errors are logged but don't stop processing of other indices
    """
    assert self.conn is not None, "Connection is closed"

    for label in vertex_config.vertex_set:
        for index_obj in vertex_config.indexes(label):
            for field in index_obj.fields:
                try:
                    query = f"CREATE INDEX ON :{label}({field})"
                    cursor = self.conn.cursor()
                    cursor.execute(query)
                    cursor.close()
                    logger.debug(f"Created index on {label}.{field}")
                except Exception as e:
                    if "already exists" in str(e).lower():
                        logger.debug(f"Index on {label}.{field} already exists")
                    else:
                        logger.warning(
                            f"Failed to create index on {label}.{field}: {e}"
                        )

delete_database(name)

Delete all data from the database.

Since Memgraph uses a single database, this clears all data.

Parameters:

Name	Type	Description	Default
name	str	Database name (ignored, clears current database)	required

Source code in graflo/db/memgraph/conn.py

def delete_database(self, name: str):
    """Delete all data from the database.

    Since Memgraph uses a single database, this clears all data.

    Args:
        name: Database name (ignored, clears current database)
    """
    assert self.conn is not None, "Connection is closed"
    try:
        cursor = self.conn.cursor()
        cursor.execute("MATCH (n) DETACH DELETE n")
        cursor.close()
        logger.info("Successfully cleared all data from Memgraph")
    except Exception as e:
        logger.error(f"Failed to clear Memgraph data: {e}", exc_info=True)
        raise

delete_graph_structure(vertex_types=(), graph_names=(), delete_all=False)

Delete graph structure (nodes and relationships).

Parameters

vertex_types : list[str], optional Specific node labels to delete edge_types : list[str], optional Specific relationship types to delete (not used, deletes via nodes) graph_names : list[str], optional Not applicable for Memgraph (single database) delete_all : bool If True, delete all nodes and relationships

Source code in graflo/db/memgraph/conn.py

def delete_graph_structure(
    self,
    vertex_types: tuple[str, ...] | list[str] = (),
    graph_names: tuple[str, ...] | list[str] = (),
    delete_all: bool = False,
) -> None:
    """Delete graph structure (nodes and relationships).

    Parameters
    ----------
    vertex_types : list[str], optional
        Specific node labels to delete
    edge_types : list[str], optional
        Specific relationship types to delete (not used, deletes via nodes)
    graph_names : list[str], optional
        Not applicable for Memgraph (single database)
    delete_all : bool
        If True, delete all nodes and relationships
    """
    assert self.conn is not None, "Connection is closed"

    if delete_all:
        cursor = self.conn.cursor()
        cursor.execute("MATCH (n) DETACH DELETE n")
        cursor.close()
        logger.info("Deleted all nodes and relationships")
        return

    # Convert tuple to list if needed
    vertex_types_list = (
        list(vertex_types) if isinstance(vertex_types, tuple) else vertex_types
    )
    if vertex_types_list:
        for label in vertex_types_list:
            try:
                cursor = self.conn.cursor()
                cursor.execute(f"MATCH (n:{label}) DETACH DELETE n")
                cursor.close()
                logger.debug(f"Deleted all nodes with label '{label}'")
            except Exception as e:
                logger.warning(f"Failed to delete nodes with label '{label}': {e}")

execute(query, **kwargs)

Execute a raw OpenCypher query against the database.

Executes the provided Cypher query with optional parameters. Parameters are safely injected using Memgraph's parameterized query mechanism to prevent injection attacks.

Parameters

query : str Cypher query string to execute **kwargs Query parameters to be safely injected

Returns

QueryResult Result object with result_set (list of tuples) and columns

Examples

Simple query::

result = conn.execute("MATCH (n:Person) RETURN n.name")
for row in result.result_set:
    print(row[0])  # Access by index

Parameterized query::

result = conn.execute(
    "MATCH (n:Person) WHERE n.age > $min_age RETURN n",
    min_age=21
)

Source code in graflo/db/memgraph/conn.py

def execute(self, query: str, **kwargs) -> QueryResult:
    """Execute a raw OpenCypher query against the database.

    Executes the provided Cypher query with optional parameters.
    Parameters are safely injected using Memgraph's parameterized
    query mechanism to prevent injection attacks.

    Parameters
    ----------
    query : str
        Cypher query string to execute
    **kwargs
        Query parameters to be safely injected

    Returns
    -------
    QueryResult
        Result object with result_set (list of tuples) and columns

    Examples
    --------
    Simple query::

        result = conn.execute("MATCH (n:Person) RETURN n.name")
        for row in result.result_set:
            print(row[0])  # Access by index

    Parameterized query::

        result = conn.execute(
            "MATCH (n:Person) WHERE n.age > $min_age RETURN n",
            min_age=21
        )
    """
    assert self.conn is not None, "Connection is closed"
    cursor = self.conn.cursor()
    try:
        if kwargs:
            cursor.execute(query, kwargs)
        else:
            cursor.execute(query)
        # mgclient uses Column objects with .name attribute, not tuples
        columns = (
            [col.name for col in cursor.description] if cursor.description else []
        )
        rows = []
        for row in cursor.fetchall():
            processed_row = []
            for value in row:
                # Convert Memgraph Node/Relationship objects to dicts
                if hasattr(value, "properties"):
                    processed_row.append(dict(value.properties))
                else:
                    processed_row.append(value)
            rows.append(tuple(processed_row))
        return QueryResult(columns, rows)
    finally:
        cursor.close()

fetch_docs(class_name, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs)

Fetch nodes from the database with optional filtering and projection.

Retrieves nodes matching the specified label and optional filter conditions. Supports field projection to return only specific properties.

Parameters

class_name : str Node label to fetch (e.g., "Person", "Product") filters : list | dict, optional Query filters in graflo expression format. Examples: ["==", "Alice", "name"] or ["AND", [...], [...]] limit : int, optional Maximum number of results to return. If None or <= 0, returns all. return_keys : list[str], optional Properties to include in results (projection). If None, returns all properties. Example: ["id", "name"] unset_keys : list[str], optional Not used for Memgraph (kept for interface compatibility) **kwargs Additional options (currently unused)

Returns

list[dict] List of node property dictionaries. Each dict contains the requested properties (or all properties if no projection).

Examples

Fetch all Person nodes::

results = db.fetch_docs("Person")

Fetch with filter and projection::

results = db.fetch_docs(
    "Person",
    filters=["==", "Alice", "name"],
    return_keys=["id", "name", "email"],
    limit=10
)

Fetch with complex filter::

results = db.fetch_docs(
    "Product",
    filters=["AND", [">", 100, "price"], ["==", "active", "status"]]
)

Source code in graflo/db/memgraph/conn.py

def fetch_docs(
    self,
    class_name: str,
    filters: list | dict | None = None,
    limit: int | None = None,
    return_keys: list[str] | None = None,
    unset_keys: list[str] | None = None,
    **kwargs,
) -> list[dict]:
    """Fetch nodes from the database with optional filtering and projection.

    Retrieves nodes matching the specified label and optional filter
    conditions. Supports field projection to return only specific properties.

    Parameters
    ----------
    class_name : str
        Node label to fetch (e.g., "Person", "Product")
    filters : list | dict, optional
        Query filters in graflo expression format.
        Examples: ``["==", "Alice", "name"]`` or ``["AND", [...], [...]]``
    limit : int, optional
        Maximum number of results to return. If None or <= 0, returns all.
    return_keys : list[str], optional
        Properties to include in results (projection). If None, returns
        all properties. Example: ``["id", "name"]``
    unset_keys : list[str], optional
        Not used for Memgraph (kept for interface compatibility)
    **kwargs
        Additional options (currently unused)

    Returns
    -------
    list[dict]
        List of node property dictionaries. Each dict contains the
        requested properties (or all properties if no projection).

    Examples
    --------
    Fetch all Person nodes::

        results = db.fetch_docs("Person")

    Fetch with filter and projection::

        results = db.fetch_docs(
            "Person",
            filters=["==", "Alice", "name"],
            return_keys=["id", "name", "email"],
            limit=10
        )

    Fetch with complex filter::

        results = db.fetch_docs(
            "Product",
            filters=["AND", [">", 100, "price"], ["==", "active", "status"]]
        )
    """
    assert self.conn is not None, "Connection is closed"

    q = f"MATCH (n:{class_name})"

    if filters is not None:
        ff = FilterExpression.from_dict(filters)
        filter_str = ff(doc_name="n", kind=self.expression_flavor())
        q += f" WHERE {filter_str}"

    # Handle projection
    if return_keys:
        return_clause = ", ".join([f"n.{k} AS {k}" for k in return_keys])
        q += f" RETURN {return_clause}"
    else:
        q += " RETURN n"

    if limit is not None and limit > 0:
        q += f" LIMIT {limit}"

    cursor = self.conn.cursor()
    cursor.execute(q)
    results = []

    if return_keys:
        # With projection, build dict from column values
        for row in cursor.fetchall():
            result = {return_keys[i]: row[i] for i in range(len(return_keys))}
            results.append(result)
    else:
        # Without projection, extract node properties
        for row in cursor.fetchall():
            node = row[0]
            if hasattr(node, "properties"):
                results.append(dict(node.properties))
            else:
                results.append(node)

    cursor.close()
    return results

fetch_edges(from_type, from_id, edge_type=None, to_type=None, to_id=None, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs)

Fetch edges from Memgraph using Cypher.

Retrieves relationships starting from a specific node, optionally filtered by edge type, target node type, and target node ID.

Parameters:

Name	Type	Description	Default
from_type	str	Source node label (e.g., "Person")	required
from_id	str	Source node ID (property value, typically "id" or "_key")	required
edge_type	str | None	Optional relationship type to filter by (e.g., "WORKS_AT")	None
to_type	str | None	Optional target node label to filter by	None
to_id	str | None	Optional target node ID to filter by	None
filters	list[Any] | dict[str, Any] | None	Additional query filters applied to relationship properties	None
limit	int | None	Maximum number of edges to return	None
return_keys	list[str] | None	Keys to return (projection) - not fully supported	None
unset_keys	list[str] | None	Keys to exclude (projection) - not fully supported	None
**kwargs	Any	Additional options	{}

Returns:

Name	Type	Description
list	list[dict[str, Any]]	List of fetched edges as dictionaries

Source code in graflo/db/memgraph/conn.py

def fetch_edges(
    self,
    from_type: str,
    from_id: str,
    edge_type: str | None = None,
    to_type: str | None = None,
    to_id: str | None = None,
    filters: list[Any] | dict[str, Any] | None = None,
    limit: int | None = None,
    return_keys: list[str] | None = None,
    unset_keys: list[str] | None = None,
    **kwargs: Any,
) -> list[dict[str, Any]]:
    """Fetch edges from Memgraph using Cypher.

    Retrieves relationships starting from a specific node, optionally filtered
    by edge type, target node type, and target node ID.

    Args:
        from_type: Source node label (e.g., "Person")
        from_id: Source node ID (property value, typically "id" or "_key")
        edge_type: Optional relationship type to filter by (e.g., "WORKS_AT")
        to_type: Optional target node label to filter by
        to_id: Optional target node ID to filter by
        filters: Additional query filters applied to relationship properties
        limit: Maximum number of edges to return
        return_keys: Keys to return (projection) - not fully supported
        unset_keys: Keys to exclude (projection) - not fully supported
        **kwargs: Additional options

    Returns:
        list: List of fetched edges as dictionaries
    """
    assert self.conn is not None, "Connection is closed"

    # Build Cypher query starting from the source node
    # Use id property (common in Memgraph) or _key if needed
    q = f"MATCH (s:{from_type} {{id: $from_id}})"

    # Build relationship pattern
    if edge_type:
        rel_pattern = f"-[r:{edge_type}]->"
    else:
        rel_pattern = "-[r]->"

    # Build target node match
    if to_type:
        target_match = f"(t:{to_type})"
    else:
        target_match = "(t)"

    q += f" {rel_pattern} {target_match}"

    # Build WHERE clauses
    where_clauses = []
    if to_id:
        where_clauses.append("t.id = $to_id")

    # Add relationship property filters
    if filters is not None:
        ff = FilterExpression.from_dict(filters)
        filter_str = ff(doc_name="r", kind=self.expression_flavor())
        where_clauses.append(filter_str)

    if where_clauses:
        q += f" WHERE {' AND '.join(where_clauses)}"

    # Build RETURN clause
    # Default: return relationship properties and basic node info
    if return_keys:
        # If return_keys specified, try to return those fields
        return_fields = []
        for key in return_keys:
            if key.startswith("from_") or key.startswith("source_"):
                return_fields.append(
                    f"s.{key.replace('from_', '').replace('source_', '')} AS {key}"
                )
            elif key.startswith("to_") or key.startswith("target_"):
                return_fields.append(
                    f"t.{key.replace('to_', '').replace('target_', '')} AS {key}"
                )
            else:
                return_fields.append(f"r.{key} AS {key}")
        q += f" RETURN {', '.join(return_fields)}"
    else:
        # Default: return relationship properties and node IDs
        q += " RETURN properties(r) AS props, s.id AS from_id, t.id AS to_id"

    if limit is not None and limit > 0:
        q += f" LIMIT {limit}"

    # Execute query with parameters
    params: dict[str, Any] = {"from_id": from_id}
    if to_id:
        params["to_id"] = to_id

    cursor = self.conn.cursor()
    cursor.execute(q, params)
    columns = [desc[0] for desc in cursor.description] if cursor.description else []
    results = []
    for row in cursor.fetchall():
        result = {}
        for i, col in enumerate(columns):
            result[col] = row[i]
        # Apply unset_keys if specified
        if unset_keys:
            for key in unset_keys:
                result.pop(key, None)
        results.append(result)
    cursor.close()
    return results

fetch_present_documents(batch, class_name, match_keys, keep_keys=None, flatten=False, filters=None)

Fetch nodes that exist in the database.

Parameters

batch : list[dict] Batch of documents to check class_name : str Label to check in match_keys : list[str] Keys to match nodes keep_keys : list[str] Keys to keep in result flatten : bool Unused in Memgraph filters : list | dict, optional Additional query filters

Returns

list[dict] Documents that exist in the database

Source code in graflo/db/memgraph/conn.py

def fetch_present_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    flatten: bool = False,
    filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """Fetch nodes that exist in the database.

    Parameters
    ----------
    batch : list[dict]
        Batch of documents to check
    class_name : str
        Label to check in
    match_keys : list[str]
        Keys to match nodes
    keep_keys : list[str]
        Keys to keep in result
    flatten : bool
        Unused in Memgraph
    filters : list | dict, optional
        Additional query filters

    Returns
    -------
    list[dict]
        Documents that exist in the database
    """
    if not batch:
        return []

    assert self.conn is not None, "Connection is closed"
    results = []

    for doc in batch:
        # Build match conditions
        match_conditions = " AND ".join([f"n.{key} = ${key}" for key in match_keys])
        params = {key: doc.get(key) for key in match_keys}

        # Build return clause with keep_keys
        if keep_keys:
            return_clause = ", ".join([f"n.{k} AS {k}" for k in keep_keys])
        else:
            return_clause = "n"

        q = f"MATCH (n:{class_name}) WHERE {match_conditions} RETURN {return_clause} LIMIT 1"

        cursor = self.conn.cursor()
        cursor.execute(q, params)
        rows = cursor.fetchall()
        cursor.close()

        if rows:
            if keep_keys:
                result = {keep_keys[i]: rows[0][i] for i in range(len(keep_keys))}
            else:
                node = rows[0][0]
                if hasattr(node, "properties"):
                    result = dict(node.properties)
                else:
                    result = node
            results.append(result)

    return results

init_db(schema, recreate_schema)

Initialize Memgraph with the given schema.

If the database already has nodes and recreate_schema is False, raises SchemaExistsError and the script halts.

Parameters

schema : Schema Schema containing graph structure definitions recreate_schema : bool If True, delete all existing data before initialization. If False and database has nodes, raises SchemaExistsError.

Source code in graflo/db/memgraph/conn.py

def init_db(self, schema: Schema, recreate_schema: bool) -> None:
    """Initialize Memgraph with the given schema.

    If the database already has nodes and recreate_schema is False, raises
    SchemaExistsError and the script halts.

    Parameters
    ----------
    schema : Schema
        Schema containing graph structure definitions
    recreate_schema : bool
        If True, delete all existing data before initialization.
        If False and database has nodes, raises SchemaExistsError.
    """
    assert self.conn is not None, "Connection is closed"

    self._database_name = schema.general.name
    logger.info(f"Initialized Memgraph with schema '{self._database_name}'")

    # Check if database already has nodes (schema/graph exists)
    cursor = self.conn.cursor()
    cursor.execute("MATCH (n) RETURN count(n) AS c")
    row = cursor.fetchone()
    cursor.close()
    count = 0
    if row is not None:
        count = (
            row[0]
            if isinstance(row, (list, tuple))
            else getattr(row, "c", row.get("c", 0) if hasattr(row, "get") else 0)
        )
    if count > 0 and not recreate_schema:
        raise SchemaExistsError(
            f"Schema/graph already exists ({count} nodes). "
            "Set recreate_schema=True to replace, or use clear_data=True before ingestion."
        )

    if recreate_schema:
        try:
            self.delete_graph_structure(delete_all=True)
        except Exception as e:
            logger.warning(f"Error clearing data on recreate_schema: {e}")

insert_edges_batch(docs_edges, source_class, target_class, relation_name, match_keys_source, match_keys_target, filter_uniques=True, head=None, **kwargs)

Insert a batch of edges using Cypher MERGE.

Creates relationships between existing nodes by matching source and target nodes using the specified match keys, then creating or updating the relationship between them.

Parameters

docs_edges : list[list[dict]] Edge specifications as list of [source_doc, target_doc, edge_props]: [[{source_props}, {target_props}, {edge_props}], ...] - source_props: Properties to match the source node - target_props: Properties to match the target node - edge_props: Properties to set on the relationship (optional) source_class : str Label of source nodes (e.g., "Person") target_class : str Label of target nodes (e.g., "Company") relation_name : str Relationship type name (e.g., "WORKS_AT") collection_name : str, optional Not used for Memgraph (kept for interface compatibility) match_keys_source : tuple[str, ...] Property keys used to identify source nodes (default: ("_key",)) match_keys_target : tuple[str, ...] Property keys used to identify target nodes (default: ("_key",)) filter_uniques : bool Not used for Memgraph (kept for interface compatibility) uniq_weight_fields : list[str], optional Not used for Memgraph (kept for interface compatibility) uniq_weight_collections : list[str], optional Not used for Memgraph (kept for interface compatibility) **kwargs Additional options (currently unused)

Notes

• Edges are created with MERGE, preventing duplicates
• If source or target node doesn't exist, the edge is silently skipped
• Edge properties are merged on update (existing props preserved)

Examples

Create relationships between Person and Company nodes::

db.insert_edges_batch(
    [
        [{"id": "alice"}, {"id": "acme"}, {"role": "engineer"}],
        [{"id": "bob"}, {"id": "acme"}, {"role": "manager"}],
    ],
    source_class="Person",
    target_class="Company",
    relation_name="WORKS_AT",
    match_keys_source=("id",),
    match_keys_target=("id",),
)

Source code in graflo/db/memgraph/conn.py

def insert_edges_batch(
    self,
    docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
    source_class: str,
    target_class: str,
    relation_name: str,
    match_keys_source: tuple[str, ...],
    match_keys_target: tuple[str, ...],
    filter_uniques: bool = True,
    head: int | None = None,
    **kwargs: Any,
) -> None:
    """Insert a batch of edges using Cypher MERGE.

    Creates relationships between existing nodes by matching source and
    target nodes using the specified match keys, then creating or updating
    the relationship between them.

    Parameters
    ----------
    docs_edges : list[list[dict]]
        Edge specifications as list of [source_doc, target_doc, edge_props]:
        ``[[{source_props}, {target_props}, {edge_props}], ...]``
        - source_props: Properties to match the source node
        - target_props: Properties to match the target node
        - edge_props: Properties to set on the relationship (optional)
    source_class : str
        Label of source nodes (e.g., "Person")
    target_class : str
        Label of target nodes (e.g., "Company")
    relation_name : str
        Relationship type name (e.g., "WORKS_AT")
    collection_name : str, optional
        Not used for Memgraph (kept for interface compatibility)
    match_keys_source : tuple[str, ...]
        Property keys used to identify source nodes (default: ("_key",))
    match_keys_target : tuple[str, ...]
        Property keys used to identify target nodes (default: ("_key",))
    filter_uniques : bool
        Not used for Memgraph (kept for interface compatibility)
    uniq_weight_fields : list[str], optional
        Not used for Memgraph (kept for interface compatibility)
    uniq_weight_collections : list[str], optional
        Not used for Memgraph (kept for interface compatibility)
    **kwargs
        Additional options (currently unused)

    Notes
    -----
    - Edges are created with MERGE, preventing duplicates
    - If source or target node doesn't exist, the edge is silently skipped
    - Edge properties are merged on update (existing props preserved)

    Examples
    --------
    Create relationships between Person and Company nodes::

        db.insert_edges_batch(
            [
                [{"id": "alice"}, {"id": "acme"}, {"role": "engineer"}],
                [{"id": "bob"}, {"id": "acme"}, {"role": "manager"}],
            ],
            source_class="Person",
            target_class="Company",
            relation_name="WORKS_AT",
            match_keys_source=("id",),
            match_keys_target=("id",),
        )
    """
    assert self.conn is not None, "Connection is closed"

    if not docs_edges:
        return

    # Handle head limit if specified
    if head is not None and head > 0:
        docs_edges = docs_edges[:head]

    # Build batch data
    batch = []
    for edge_data in docs_edges:
        if len(edge_data) < 2:
            continue

        source_doc = edge_data[0]
        target_doc = edge_data[1]
        edge_props = edge_data[2] if len(edge_data) > 2 else {}

        # Sanitize
        source_doc = self._sanitize_doc(source_doc)
        target_doc = self._sanitize_doc(target_doc)
        edge_props = self._sanitize_doc(edge_props) if edge_props else {}

        batch.append(
            {
                "source": source_doc,
                "target": target_doc,
                "props": edge_props,
            }
        )

    if not batch:
        return

    # Build match patterns
    source_match = ", ".join([f"{k}: row.source.{k}" for k in match_keys_source])
    target_match = ", ".join([f"{k}: row.target.{k}" for k in match_keys_target])

    q = f"""
        UNWIND $batch AS row
        MATCH (s:{source_class} {{ {source_match} }})
        MATCH (t:{target_class} {{ {target_match} }})
        MERGE (s)-[r:{relation_name}]->(t)
        ON CREATE SET r = row.props
        ON MATCH SET r += row.props
    """

    cursor = self.conn.cursor()
    cursor.execute(q, {"batch": batch})
    cursor.close()

insert_return_batch(docs, class_name)

Insert nodes and return their properties.

Parameters

docs : list[dict] Documents to insert class_name : str Label to insert into

Returns

list[dict] | str Inserted documents with their properties, or query string

Raises

NotImplementedError This method is not fully implemented for Memgraph

Source code in graflo/db/memgraph/conn.py

def insert_return_batch(
    self, docs: list[dict[str, Any]], class_name: str
) -> list[dict[str, Any]] | str:
    """Insert nodes and return their properties.

    Parameters
    ----------
    docs : list[dict]
        Documents to insert
    class_name : str
        Label to insert into

    Returns
    -------
    list[dict] | str
        Inserted documents with their properties, or query string

    Raises
    ------
    NotImplementedError
        This method is not fully implemented for Memgraph
    """
    raise NotImplementedError("insert_return_batch is not implemented for Memgraph")

keep_absent_documents(batch, class_name, match_keys, keep_keys=None, filters=None)

Keep documents that don't exist in the database.

Parameters

batch : list[dict] Batch of documents to check class_name : str Label to check in match_keys : list[str] Keys to match nodes keep_keys : list[str] Keys to keep in result filters : list | dict, optional Additional query filters

Returns

list[dict] Documents that don't exist in the database

Source code in graflo/db/memgraph/conn.py

def keep_absent_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """Keep documents that don't exist in the database.

    Parameters
    ----------
    batch : list[dict]
        Batch of documents to check
    class_name : str
        Label to check in
    match_keys : list[str]
        Keys to match nodes
    keep_keys : list[str]
        Keys to keep in result
    filters : list | dict, optional
        Additional query filters

    Returns
    -------
    list[dict]
        Documents that don't exist in the database
    """
    if not batch:
        return []

    # Find documents that exist
    present_docs = self.fetch_present_documents(
        batch, class_name, match_keys, match_keys, filters=filters
    )

    # Create a set of present document keys for efficient lookup
    present_keys = set()
    for doc in present_docs:
        key_tuple = tuple(doc.get(k) for k in match_keys)
        present_keys.add(key_tuple)

    # Keep documents that don't exist
    absent = []
    for doc in batch:
        key_tuple = tuple(doc.get(k) for k in match_keys)
        if key_tuple not in present_keys:
            if keep_keys:
                absent.append({k: doc.get(k) for k in keep_keys})
            else:
                absent.append(doc)

    return absent

upsert_docs_batch(docs, class_name, match_keys, **kwargs)

Upsert a batch of nodes using Cypher MERGE.

Performs atomic upsert (update-or-insert) operations on a batch of documents. Uses Cypher MERGE with ON MATCH/ON CREATE for efficiency. Existing properties not in the document are preserved on update.

Parameters

docs : list[dict] Documents to upsert. Each document must contain all match_keys with non-None values. Invalid documents are skipped with a warning. class_name : str Node label to create/update (e.g., "Person", "Product") match_keys : list[str] Property keys used to identify existing nodes for update. Supports composite keys (multiple fields). **kwargs Additional options: - dry (bool): If True, build query but don't execute (for debugging)

Raises

ValueError If any document is missing a required match_key or has None value.

Notes

• Documents are sanitized before insertion (invalid keys/values removed)
• NaN, Inf, and null bytes are automatically filtered with warnings
• The operation is atomic per batch (all succeed or fail together)

Examples

Insert or update Person nodes::

db.upsert_docs_batch(
    [
        {"id": "1", "name": "Alice", "age": 30},
        {"id": "2", "name": "Bob", "age": 25},
    ],
    class_name="Person",
    match_keys=["id"]
)

With composite key::

db.upsert_docs_batch(
    [{"tenant": "acme", "user_id": "u1", "email": "a@b.com"}],
    class_name="User",
    match_keys=["tenant", "user_id"]
)

Source code in graflo/db/memgraph/conn.py

def upsert_docs_batch(
    self,
    docs: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    **kwargs: Any,
) -> None:
    """Upsert a batch of nodes using Cypher MERGE.

    Performs atomic upsert (update-or-insert) operations on a batch of
    documents. Uses Cypher MERGE with ON MATCH/ON CREATE for efficiency.
    Existing properties not in the document are preserved on update.

    Parameters
    ----------
    docs : list[dict]
        Documents to upsert. Each document must contain all match_keys
        with non-None values. Invalid documents are skipped with a warning.
    class_name : str
        Node label to create/update (e.g., "Person", "Product")
    match_keys : list[str]
        Property keys used to identify existing nodes for update.
        Supports composite keys (multiple fields).
    **kwargs
        Additional options:
        - dry (bool): If True, build query but don't execute (for debugging)

    Raises
    ------
    ValueError
        If any document is missing a required match_key or has None value.

    Notes
    -----
    - Documents are sanitized before insertion (invalid keys/values removed)
    - NaN, Inf, and null bytes are automatically filtered with warnings
    - The operation is atomic per batch (all succeed or fail together)

    Examples
    --------
    Insert or update Person nodes::

        db.upsert_docs_batch(
            [
                {"id": "1", "name": "Alice", "age": 30},
                {"id": "2", "name": "Bob", "age": 25},
            ],
            class_name="Person",
            match_keys=["id"]
        )

    With composite key::

        db.upsert_docs_batch(
            [{"tenant": "acme", "user_id": "u1", "email": "a@b.com"}],
            class_name="User",
            match_keys=["tenant", "user_id"]
        )
    """
    assert self.conn is not None, "Connection is closed"
    dry = kwargs.pop("dry", False)

    if not docs:
        return

    # Convert tuple to list if needed
    match_keys_list = (
        list(match_keys) if isinstance(match_keys, tuple) else match_keys
    )
    # Sanitize documents
    sanitized_docs = self._sanitize_batch(docs, match_keys_list)

    if not sanitized_docs:
        return

    # Auto-create index on match_keys for MERGE performance (idempotent)
    cursor = self.conn.cursor()
    for key in match_keys:
        try:
            cursor.execute(f"CREATE INDEX ON :{class_name}({key})")
            logger.debug(f"Created index on {class_name}.{key}")
        except Exception as e:
            if "already exists" not in str(e).lower():
                logger.debug(f"Index on {class_name}.{key}: {e}")
    cursor.close()

    # Build the MERGE clause with match keys
    index_str = ", ".join([f"{k}: row.{k}" for k in match_keys_list])
    q = f"""
        UNWIND $batch AS row
        MERGE (n:{class_name} {{ {index_str} }})
        ON MATCH SET n += row
        ON CREATE SET n += row
    """
    if not dry:
        cursor = self.conn.cursor()
        cursor.execute(q, {"batch": sanitized_docs})
        cursor.close()

Neo4jConnection

Bases: Connection

Neo4j-specific implementation of the Connection interface.

This class provides Neo4j-specific implementations for all database operations, including node management, relationship operations, and Cypher query execution. It uses the Neo4j Python driver for all operations.

Attributes:

Name	Type	Description
flavor		Database flavor identifier (NEO4J)
conn		Neo4j session instance

Source code in graflo/db/neo4j/conn.py

class Neo4jConnection(Connection):
    """Neo4j-specific implementation of the Connection interface.

    This class provides Neo4j-specific implementations for all database
    operations, including node management, relationship operations, and
    Cypher query execution. It uses the Neo4j Python driver for all operations.

    Attributes:
        flavor: Database flavor identifier (NEO4J)
        conn: Neo4j session instance
    """

    flavor = DBType.NEO4J

    def __init__(self, config: Neo4jConfig):
        """Initialize Neo4j connection.

        Args:
            config: Neo4j connection configuration containing URL and credentials
        """
        super().__init__()
        # Store config for later use
        self.config = config
        # Ensure url is not None - GraphDatabase.driver requires a non-None URI
        if config.url is None:
            raise ValueError("Neo4j connection requires a URL to be configured")
        # Handle None values in auth tuple
        auth = None
        if config.username is not None and config.password is not None:
            auth = (config.username, config.password)
        self._driver = GraphDatabase.driver(uri=config.url, auth=auth)
        self.conn = self._driver.session()

    def execute(self, query, **kwargs):
        """Execute a Cypher query.

        Args:
            query: Cypher query string to execute
            **kwargs: Additional query parameters

        Returns:
            Result: Neo4j query result
        """
        cursor = self.conn.run(query, **kwargs)
        return cursor

    def close(self):
        """Close the Neo4j connection and session."""
        # Close session first, then the underlying driver
        try:
            self.conn.close()
        finally:
            # Ensure the driver is also closed to release resources
            self._driver.close()

    def create_database(self, name: str):
        """Create a new Neo4j database.

        Note: This operation is only supported in Neo4j Enterprise Edition.
        Community Edition only supports one database per instance.

        Args:
            name: Name of the database to create

        Raises:
            Exception: If database creation fails and it's not a Community Edition limitation
        """
        try:
            self.execute(f"CREATE DATABASE {name}")
            logger.info(f"Successfully created Neo4j database '{name}'")
        except Exception as e:
            # Check if this is a Neo4j Community Edition limitation
            error_str = str(e).lower()
            if (
                "unsupported administration command" in error_str
                or "create database" in error_str
            ):
                # This is likely Neo4j Community Edition - don't raise, just log
                logger.info(
                    f"Neo4j Community Edition detected: Cannot create database '{name}'. "
                    f"Community Edition only supports one database per instance. "
                    f"Continuing with default database."
                )
                # Don't raise - allow operation to continue with default database
                return
            # For other errors, raise them
            raise e

    def delete_database(self, name: str):
        """Delete a Neo4j database.

        Note: This operation is only supported in Neo4j Enterprise Edition.
        As a fallback, it deletes all nodes and relationships.

        Args:
            name: Name of the database to delete (unused, deletes all data)
        """
        try:
            self.execute("MATCH (n) DETACH DELETE n")
            logger.info("Successfully cleaned Neo4j database")
        except Exception as e:
            logger.error(
                f"Failed to clean Neo4j database: {e}",
                exc_info=True,
            )
            raise

    def define_vertex_indices(self, vertex_config: VertexConfig):
        """Define indices for vertex labels.

        Creates indices for each vertex label based on the configuration.

        Args:
            vertex_config: Vertex configuration containing index definitions
        """
        for c in vertex_config.vertex_set:
            for index_obj in vertex_config.indexes(c):
                self._add_index(c, index_obj)

    def define_edge_indices(self, edges: list[Edge]):
        """Define indices for relationship types.

        Creates indices for each relationship type based on the configuration.

        Args:
            edges: List of edge configurations containing index definitions
        """
        for edge in edges:
            for index_obj in edge.indexes:
                if edge.relation is not None:
                    self._add_index(edge.relation, index_obj, is_vertex_index=False)

    def _add_index(self, obj_name, index: Index, is_vertex_index=True):
        """Add an index to a label or relationship type.

        Args:
            obj_name: Label or relationship type name
            index: Index configuration to create
            is_vertex_index: If True, create index on nodes, otherwise on relationships
        """
        fields_str = ", ".join([f"x.{f}" for f in index.fields])
        fields_str2 = "_".join(index.fields)
        index_name = f"{obj_name}_{fields_str2}"
        if is_vertex_index:
            formula = f"(x:{obj_name})"
        else:
            formula = f"()-[x:{obj_name}]-()"

        q = f"CREATE INDEX {index_name} IF NOT EXISTS FOR {formula} ON ({fields_str});"

        self.execute(q)

    def define_schema(self, schema: Schema):
        """Define vertex and edge classes based on schema.

        Note: This is a no-op in Neo4j as vertex/edge classes (labels/relationship types) are implicit.

        Args:
            schema: Schema containing vertex and edge class definitions
        """
        pass

    def define_vertex_classes(self, schema: Schema):
        """Define vertex classes based on schema.

        Note: This is a no-op in Neo4j as vertex classes (labels) are implicit.

        Args:
            schema: Schema containing vertex definitions
        """
        pass

    def define_edge_classes(self, edges: list[Edge]):
        """Define edge classes based on schema.

        Note: This is a no-op in Neo4j as edge classes (relationship types) are implicit.

        Args:
            edges: List of edge configurations
        """
        pass

    def delete_graph_structure(
        self,
        vertex_types: tuple[str, ...] | list[str] = (),
        graph_names: tuple[str, ...] | list[str] = (),
        delete_all: bool = False,
    ) -> None:
        """Delete graph structure (nodes and relationships) from Neo4j.

        In Neo4j:
        - Labels: Categories for nodes (equivalent to vertex types)
        - Relationship Types: Types of relationships (equivalent to edge types)
        - No explicit "graph" concept - all nodes/relationships are in the database

        Args:
            vertex_types: Label names to delete nodes for
            graph_names: Unused in Neo4j (no explicit graph concept)
            delete_all: If True, delete all nodes and relationships
        """
        cnames = vertex_types
        if cnames:
            for c in cnames:
                q = f"MATCH (n:{c}) DELETE n"
                self.execute(q)
        else:
            q = "MATCH (n) DELETE n"
            self.execute(q)

    def init_db(self, schema: Schema, recreate_schema: bool) -> None:
        """Initialize Neo4j with the given schema.

        Checks if the database exists and creates it if it doesn't.
        Uses schema.general.name if database is not set in config.
        Note: Database creation is only supported in Neo4j Enterprise Edition.

        If the database already has nodes and recreate_schema is False, raises
        SchemaExistsError and the script halts.

        Args:
            schema: Schema containing graph structure definitions
            recreate_schema: If True, delete all existing data before initialization.
                If False and database has nodes, raises SchemaExistsError.
        """
        # Determine database name: use config.database if set, otherwise use schema.general.name
        db_name = self.config.database
        if not db_name:
            db_name = schema.general.name
            # Update config for subsequent operations
            self.config.database = db_name

        # Check if database exists and create it if it doesn't
        # Note: This only works in Neo4j Enterprise Edition
        # For Community Edition, create_database will handle it gracefully
        # Community Edition only allows one database per instance
        try:
            # Try to check if database exists (Enterprise feature)
            try:
                result = self.execute("SHOW DATABASES")
                # Neo4j result is a cursor-like object, iterate to get records
                databases = []
                for record in result:
                    # Record structure may vary, try common field names
                    if hasattr(record, "get"):
                        db_name_field = (
                            record.get("name")
                            or record.get("database")
                            or record.get("db")
                        )
                    else:
                        # If record is a dict-like object, try direct access
                        db_name_field = getattr(record, "name", None) or getattr(
                            record, "database", None
                        )
                    if db_name_field:
                        databases.append(db_name_field)

                if db_name not in databases:
                    logger.info(
                        f"Database '{db_name}' does not exist, attempting to create it..."
                    )
                    # create_database handles Community Edition errors gracefully
                    self.create_database(db_name)
            except Exception as show_error:
                # If SHOW DATABASES fails (Community Edition or older versions), try to create anyway
                logger.debug(
                    f"Could not check database existence (may be Community Edition): {show_error}"
                )
                # create_database handles Community Edition errors gracefully
                self.create_database(db_name)
        except Exception as e:
            # Only log unexpected errors (create_database handles Community Edition gracefully)
            logger.error(
                f"Unexpected error during database initialization for '{db_name}': {e}",
                exc_info=True,
            )
            # Don't raise - allow operation to continue with default database
            logger.warning(
                "Continuing with default database due to initialization error"
            )

        try:
            # Check if database already has nodes (schema/graph exists)
            result = self.execute("MATCH (n) RETURN count(n) AS c")
            count = 0
            if hasattr(result, "data"):
                data = result.data()
                if data:
                    first = data[0]
                    count = first.get("c", 0) or 0
            if count == 0 and hasattr(result, "__iter__"):
                for record in result:
                    if hasattr(record, "get"):
                        count = record.get("c", 0) or 0
                    else:
                        count = getattr(record, "c", 0) or 0
                    break
            if count > 0 and not recreate_schema:
                raise SchemaExistsError(
                    f"Schema/graph already exists in database '{db_name}' ({count} nodes). "
                    "Set recreate_schema=True to replace, or use clear_data=True before ingestion."
                )

            if recreate_schema:
                try:
                    self.delete_database("")
                    logger.debug(f"Cleaned database '{db_name}' for fresh start")
                except Exception as clean_error:
                    logger.warning(
                        f"Error during recreate_schema for database '{db_name}': {clean_error}",
                        exc_info=True,
                    )
                    # Continue - may be first run or already clean

            try:
                self.define_indexes(schema)
                logger.debug(f"Defined indexes for database '{db_name}'")
            except Exception as index_error:
                logger.error(
                    f"Failed to define indexes for database '{db_name}': {index_error}",
                    exc_info=True,
                )
                raise
        except SchemaExistsError:
            raise
        except Exception as e:
            logger.error(
                f"Error during database schema initialization for '{db_name}': {e}",
                exc_info=True,
            )
            raise

    def clear_data(self, schema: Schema) -> None:
        """Remove all data from the graph without dropping the schema.

        Deletes all nodes and relationships; labels (schema) remain.
        """
        self.delete_graph_structure((), (), delete_all=True)

    def upsert_docs_batch(self, docs, class_name, match_keys, **kwargs):
        """Upsert a batch of nodes using Cypher.

        Performs an upsert operation on a batch of nodes, using the specified
        match keys to determine whether to update existing nodes or create new ones.

        Args:
            docs: List of node documents to upsert
            class_name: Label to upsert into
            match_keys: Keys to match for upsert operation
            **kwargs: Additional options:
                - dry: If True, don't execute the query
        """
        dry = kwargs.pop("dry", False)

        index_str = ", ".join([f"{k}: row.{k}" for k in match_keys])
        q = f"""
            WITH $batch AS batch 
            UNWIND batch as row 
            MERGE (n:{class_name} {{ {index_str} }}) 
            ON MATCH set n += row 
            ON CREATE set n += row
        """
        if not dry:
            self.execute(q, batch=docs)

    def insert_edges_batch(
        self,
        docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
        source_class: str,
        target_class: str,
        relation_name: str,
        match_keys_source: tuple[str, ...],
        match_keys_target: tuple[str, ...],
        filter_uniques: bool = True,
        head: int | None = None,
        **kwargs: Any,
    ) -> None:
        """Insert a batch of relationships using Cypher.

        Creates relationships between source and target nodes, with support for
        property matching and unique constraints.

        Args:
            docs_edges: List of edge documents in format [{__source: source_doc, __target: target_doc}]
            source_class: Source node label
            target_class: Target node label
            relation_name: Relationship type name
            match_keys_source: Keys to match source nodes
            match_keys_target: Keys to match target nodes
            filter_uniques: Unused in Neo4j (MERGE handles uniqueness automatically)
            head: Optional limit on number of relationships to insert
            **kwargs: Additional options:
                - dry: If True, don't execute the query
                - collection_name: Unused in Neo4j (kept for interface compatibility)
                - uniq_weight_fields: Unused in Neo4j (ArangoDB-specific)
                - uniq_weight_collections: Unused in Neo4j (ArangoDB-specific)
                - upsert_option: Unused in Neo4j (ArangoDB-specific, MERGE is always upsert)
        """
        dry = kwargs.pop("dry", False)
        # Extract and ignore unused parameters (kept for interface compatibility)
        kwargs.pop("collection_name", None)
        kwargs.pop("uniq_weight_fields", None)
        kwargs.pop("uniq_weight_collections", None)
        kwargs.pop("upsert_option", None)

        # Apply head limit if specified
        if head is not None and isinstance(docs_edges, list):
            docs_edges = docs_edges[:head]

        # Note: filter_uniques is unused because Neo4j's MERGE handles uniqueness automatically

        source_match_str = [f"source.{key} = row[0].{key}" for key in match_keys_source]
        target_match_str = [f"target.{key} = row[1].{key}" for key in match_keys_target]

        match_clause = "WHERE " + " AND ".join(source_match_str + target_match_str)

        q = f"""
            WITH $batch AS batch 
            UNWIND batch as row 
            MATCH (source:{source_class}), 
                  (target:{target_class}) {match_clause} 
                        MERGE (source)-[r:{relation_name}]->(target)
                SET r += row[2]

        """
        if not dry:
            self.execute(q, batch=docs_edges)

    def insert_return_batch(
        self, docs: list[dict[str, Any]], class_name: str
    ) -> list[dict[str, Any]] | str:
        """Insert nodes and return their properties.

        Note: Not implemented in Neo4j.

        Args:
            docs: Documents to insert
            class_name: Label to insert into

        Raises:
            NotImplementedError: This method is not implemented for Neo4j
        """
        raise NotImplementedError()

    def fetch_docs(
        self,
        class_name,
        filters: list | dict | None = None,
        limit: int | None = None,
        return_keys: list | None = None,
        unset_keys: list | None = None,
        **kwargs,
    ):
        """Fetch nodes from a label.

        Args:
            class_name: Label to fetch from
            filters: Query filters
            limit: Maximum number of nodes to return
            return_keys: Keys to return
            unset_keys: Unused in Neo4j

        Returns:
            list: Fetched nodes
        """
        if filters is not None:
            ff = FilterExpression.from_dict(filters)
            filter_clause = f"WHERE {ff(doc_name='n', kind=self.expression_flavor())}"
        else:
            filter_clause = ""

        if return_keys is not None:
            keep_clause_ = ", ".join([f".{item}" for item in return_keys])
            keep_clause = f"{{ {keep_clause_} }}"
        else:
            keep_clause = ""

        if limit is not None and isinstance(limit, int):
            limit_clause = f"LIMIT {limit}"
        else:
            limit_clause = ""

        q = (
            f"MATCH (n:{class_name})"
            f"  {filter_clause}"
            f"  RETURN n {keep_clause}"
            f"  {limit_clause}"
        )
        cursor = self.execute(q)
        r = [item["n"] for item in cursor.data()]
        return r

    # TODO test
    def fetch_edges(
        self,
        from_type: str,
        from_id: str,
        edge_type: str | None = None,
        to_type: str | None = None,
        to_id: str | None = None,
        filters: list | dict | None = None,
        limit: int | None = None,
        return_keys: list | None = None,
        unset_keys: list | None = None,
        **kwargs,
    ):
        """Fetch edges from Neo4j using Cypher.

        Args:
            from_type: Source node label
            from_id: Source node ID (property name depends on match_keys used)
            edge_type: Optional relationship type to filter by
            to_type: Optional target node label to filter by
            to_id: Optional target node ID to filter by
            filters: Additional query filters
            limit: Maximum number of edges to return
            return_keys: Keys to return (projection)
            unset_keys: Keys to exclude (projection) - not supported in Neo4j
            **kwargs: Additional parameters

        Returns:
            list: List of fetched edges
        """
        # Build Cypher query to fetch edges
        # Match source node first
        source_match = f"(source:{from_type} {{id: '{from_id}'}})"

        # Build relationship pattern
        if edge_type:
            rel_pattern = f"-[r:{edge_type}]->"
        else:
            rel_pattern = "-[r]->"

        # Build target node match
        if to_type:
            target_match = f"(target:{to_type})"
        else:
            target_match = "(target)"

        # Add target ID filter if provided
        where_clauses = []
        if to_id:
            where_clauses.append(f"target.id = '{to_id}'")

        # Add additional filters if provided
        if filters is not None:
            ff = FilterExpression.from_dict(filters)
            filter_clause = ff(doc_name="r", kind=self.expression_flavor())
            where_clauses.append(filter_clause)

        where_clause = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""

        # Build return clause
        if return_keys is not None:
            return_clause = ", ".join([f"r.{key} as {key}" for key in return_keys])
            return_clause = f"RETURN {return_clause}"
        else:
            return_clause = "RETURN r"

        limit_clause = f"LIMIT {limit}" if limit else ""

        query = f"""
            MATCH {source_match}{rel_pattern}{target_match}
            {where_clause}
            {return_clause}
            {limit_clause}
        """

        cursor = self.execute(query)
        result = [item["r"] for item in cursor.data()]

        # Note: unset_keys is not supported in Neo4j as we can't modify the result structure
        # after the query

        return result

    def fetch_present_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        flatten: bool = False,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Fetch nodes that exist in the database.

        Note: Not implemented in Neo4j.

        Args:
            batch: Batch of documents to check
            class_name: Label to check in
            match_keys: Keys to match nodes
            keep_keys: Keys to keep in result
            flatten: Unused in Neo4j
            filters: Additional query filters

        Raises:
            NotImplementedError: This method is not implemented for Neo4j
        """
        raise NotImplementedError

    def aggregate(
        self,
        class_name,
        aggregation_function: AggregationType,
        discriminant: str | None = None,
        aggregated_field: str | None = None,
        filters: list | dict | None = None,
    ):
        """Perform aggregation on nodes.

        Note: Not implemented in Neo4j.

        Args:
            class_name: Label to aggregate
            aggregation_function: Type of aggregation to perform
            discriminant: Field to group by
            aggregated_field: Field to aggregate
            filters: Query filters

        Raises:
            NotImplementedError: This method is not implemented for Neo4j
        """
        raise NotImplementedError

    def keep_absent_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Keep nodes that don't exist in the database.

        Note: Not implemented in Neo4j.

        Args:
            batch: Batch of documents to check
            class_name: Label to check in
            match_keys: Keys to match nodes
            keep_keys: Keys to keep in result
            filters: Additional query filters

        Raises:
            NotImplementedError: This method is not implemented for Neo4j
        """
        raise NotImplementedError

__init__(config)

Initialize Neo4j connection.

Parameters:

Name	Type	Description	Default
config	Neo4jConfig	Neo4j connection configuration containing URL and credentials	required

Source code in graflo/db/neo4j/conn.py

def __init__(self, config: Neo4jConfig):
    """Initialize Neo4j connection.

    Args:
        config: Neo4j connection configuration containing URL and credentials
    """
    super().__init__()
    # Store config for later use
    self.config = config
    # Ensure url is not None - GraphDatabase.driver requires a non-None URI
    if config.url is None:
        raise ValueError("Neo4j connection requires a URL to be configured")
    # Handle None values in auth tuple
    auth = None
    if config.username is not None and config.password is not None:
        auth = (config.username, config.password)
    self._driver = GraphDatabase.driver(uri=config.url, auth=auth)
    self.conn = self._driver.session()

aggregate(class_name, aggregation_function, discriminant=None, aggregated_field=None, filters=None)

Perform aggregation on nodes.

Note: Not implemented in Neo4j.

Parameters:

Name	Type	Description	Default
class_name		Label to aggregate	required
aggregation_function	AggregationType	Type of aggregation to perform	required
discriminant	str | None	Field to group by	None
aggregated_field	str | None	Field to aggregate	None
filters	list | dict | None	Query filters	None

Raises:

Type	Description
NotImplementedError	This method is not implemented for Neo4j

Source code in graflo/db/neo4j/conn.py

def aggregate(
    self,
    class_name,
    aggregation_function: AggregationType,
    discriminant: str | None = None,
    aggregated_field: str | None = None,
    filters: list | dict | None = None,
):
    """Perform aggregation on nodes.

    Note: Not implemented in Neo4j.

    Args:
        class_name: Label to aggregate
        aggregation_function: Type of aggregation to perform
        discriminant: Field to group by
        aggregated_field: Field to aggregate
        filters: Query filters

    Raises:
        NotImplementedError: This method is not implemented for Neo4j
    """
    raise NotImplementedError

clear_data(schema)

Remove all data from the graph without dropping the schema.

Deletes all nodes and relationships; labels (schema) remain.

Source code in graflo/db/neo4j/conn.py

def clear_data(self, schema: Schema) -> None:
    """Remove all data from the graph without dropping the schema.

    Deletes all nodes and relationships; labels (schema) remain.
    """
    self.delete_graph_structure((), (), delete_all=True)

close()

Close the Neo4j connection and session.

Source code in graflo/db/neo4j/conn.py

def close(self):
    """Close the Neo4j connection and session."""
    # Close session first, then the underlying driver
    try:
        self.conn.close()
    finally:
        # Ensure the driver is also closed to release resources
        self._driver.close()

create_database(name)

Create a new Neo4j database.

Note: This operation is only supported in Neo4j Enterprise Edition. Community Edition only supports one database per instance.

Parameters:

Name	Type	Description	Default
name	str	Name of the database to create	required

Raises:

Type	Description
Exception	If database creation fails and it's not a Community Edition limitation

Source code in graflo/db/neo4j/conn.py

def create_database(self, name: str):
    """Create a new Neo4j database.

    Note: This operation is only supported in Neo4j Enterprise Edition.
    Community Edition only supports one database per instance.

    Args:
        name: Name of the database to create

    Raises:
        Exception: If database creation fails and it's not a Community Edition limitation
    """
    try:
        self.execute(f"CREATE DATABASE {name}")
        logger.info(f"Successfully created Neo4j database '{name}'")
    except Exception as e:
        # Check if this is a Neo4j Community Edition limitation
        error_str = str(e).lower()
        if (
            "unsupported administration command" in error_str
            or "create database" in error_str
        ):
            # This is likely Neo4j Community Edition - don't raise, just log
            logger.info(
                f"Neo4j Community Edition detected: Cannot create database '{name}'. "
                f"Community Edition only supports one database per instance. "
                f"Continuing with default database."
            )
            # Don't raise - allow operation to continue with default database
            return
        # For other errors, raise them
        raise e

define_edge_classes(edges)

Define edge classes based on schema.

Note: This is a no-op in Neo4j as edge classes (relationship types) are implicit.

Parameters:

Name	Type	Description	Default
edges	list[Edge]	List of edge configurations	required

Source code in graflo/db/neo4j/conn.py

def define_edge_classes(self, edges: list[Edge]):
    """Define edge classes based on schema.

    Note: This is a no-op in Neo4j as edge classes (relationship types) are implicit.

    Args:
        edges: List of edge configurations
    """
    pass

define_edge_indices(edges)

Define indices for relationship types.

Creates indices for each relationship type based on the configuration.

Parameters:

Name	Type	Description	Default
edges	list[Edge]	List of edge configurations containing index definitions	required

Source code in graflo/db/neo4j/conn.py

def define_edge_indices(self, edges: list[Edge]):
    """Define indices for relationship types.

    Creates indices for each relationship type based on the configuration.

    Args:
        edges: List of edge configurations containing index definitions
    """
    for edge in edges:
        for index_obj in edge.indexes:
            if edge.relation is not None:
                self._add_index(edge.relation, index_obj, is_vertex_index=False)

define_schema(schema)

Define vertex and edge classes based on schema.

Note: This is a no-op in Neo4j as vertex/edge classes (labels/relationship types) are implicit.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing vertex and edge class definitions	required

Source code in graflo/db/neo4j/conn.py

def define_schema(self, schema: Schema):
    """Define vertex and edge classes based on schema.

    Note: This is a no-op in Neo4j as vertex/edge classes (labels/relationship types) are implicit.

    Args:
        schema: Schema containing vertex and edge class definitions
    """
    pass

define_vertex_classes(schema)

Define vertex classes based on schema.

Note: This is a no-op in Neo4j as vertex classes (labels) are implicit.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing vertex definitions	required

Source code in graflo/db/neo4j/conn.py

def define_vertex_classes(self, schema: Schema):
    """Define vertex classes based on schema.

    Note: This is a no-op in Neo4j as vertex classes (labels) are implicit.

    Args:
        schema: Schema containing vertex definitions
    """
    pass

define_vertex_indices(vertex_config)

Define indices for vertex labels.

Creates indices for each vertex label based on the configuration.

Parameters:

Name	Type	Description	Default
vertex_config	VertexConfig	Vertex configuration containing index definitions	required

Source code in graflo/db/neo4j/conn.py

def define_vertex_indices(self, vertex_config: VertexConfig):
    """Define indices for vertex labels.

    Creates indices for each vertex label based on the configuration.

    Args:
        vertex_config: Vertex configuration containing index definitions
    """
    for c in vertex_config.vertex_set:
        for index_obj in vertex_config.indexes(c):
            self._add_index(c, index_obj)

delete_database(name)

Delete a Neo4j database.

Note: This operation is only supported in Neo4j Enterprise Edition. As a fallback, it deletes all nodes and relationships.

Parameters:

Name	Type	Description	Default
name	str	Name of the database to delete (unused, deletes all data)	required

Source code in graflo/db/neo4j/conn.py

def delete_database(self, name: str):
    """Delete a Neo4j database.

    Note: This operation is only supported in Neo4j Enterprise Edition.
    As a fallback, it deletes all nodes and relationships.

    Args:
        name: Name of the database to delete (unused, deletes all data)
    """
    try:
        self.execute("MATCH (n) DETACH DELETE n")
        logger.info("Successfully cleaned Neo4j database")
    except Exception as e:
        logger.error(
            f"Failed to clean Neo4j database: {e}",
            exc_info=True,
        )
        raise

delete_graph_structure(vertex_types=(), graph_names=(), delete_all=False)

Delete graph structure (nodes and relationships) from Neo4j.

In Neo4j: - Labels: Categories for nodes (equivalent to vertex types) - Relationship Types: Types of relationships (equivalent to edge types) - No explicit "graph" concept - all nodes/relationships are in the database

Parameters:

Name	Type	Description	Default
vertex_types	tuple[str, ...] | list[str]	Label names to delete nodes for	()
graph_names	tuple[str, ...] | list[str]	Unused in Neo4j (no explicit graph concept)	()
delete_all	bool	If True, delete all nodes and relationships	False

Source code in graflo/db/neo4j/conn.py

def delete_graph_structure(
    self,
    vertex_types: tuple[str, ...] | list[str] = (),
    graph_names: tuple[str, ...] | list[str] = (),
    delete_all: bool = False,
) -> None:
    """Delete graph structure (nodes and relationships) from Neo4j.

    In Neo4j:
    - Labels: Categories for nodes (equivalent to vertex types)
    - Relationship Types: Types of relationships (equivalent to edge types)
    - No explicit "graph" concept - all nodes/relationships are in the database

    Args:
        vertex_types: Label names to delete nodes for
        graph_names: Unused in Neo4j (no explicit graph concept)
        delete_all: If True, delete all nodes and relationships
    """
    cnames = vertex_types
    if cnames:
        for c in cnames:
            q = f"MATCH (n:{c}) DELETE n"
            self.execute(q)
    else:
        q = "MATCH (n) DELETE n"
        self.execute(q)

execute(query, **kwargs)

Execute a Cypher query.

Parameters:

Name	Type	Description	Default
query		Cypher query string to execute	required
**kwargs		Additional query parameters	{}

Returns:

Name	Type	Description
Result		Neo4j query result

Source code in graflo/db/neo4j/conn.py

def execute(self, query, **kwargs):
    """Execute a Cypher query.

    Args:
        query: Cypher query string to execute
        **kwargs: Additional query parameters

    Returns:
        Result: Neo4j query result
    """
    cursor = self.conn.run(query, **kwargs)
    return cursor

fetch_docs(class_name, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs)

Fetch nodes from a label.

Parameters:

Name	Type	Description	Default
class_name		Label to fetch from	required
filters	list | dict | None	Query filters	None
limit	int | None	Maximum number of nodes to return	None
return_keys	list | None	Keys to return	None
unset_keys	list | None	Unused in Neo4j	None

Returns:

Name	Type	Description
list		Fetched nodes

Source code in graflo/db/neo4j/conn.py

def fetch_docs(
    self,
    class_name,
    filters: list | dict | None = None,
    limit: int | None = None,
    return_keys: list | None = None,
    unset_keys: list | None = None,
    **kwargs,
):
    """Fetch nodes from a label.

    Args:
        class_name: Label to fetch from
        filters: Query filters
        limit: Maximum number of nodes to return
        return_keys: Keys to return
        unset_keys: Unused in Neo4j

    Returns:
        list: Fetched nodes
    """
    if filters is not None:
        ff = FilterExpression.from_dict(filters)
        filter_clause = f"WHERE {ff(doc_name='n', kind=self.expression_flavor())}"
    else:
        filter_clause = ""

    if return_keys is not None:
        keep_clause_ = ", ".join([f".{item}" for item in return_keys])
        keep_clause = f"{{ {keep_clause_} }}"
    else:
        keep_clause = ""

    if limit is not None and isinstance(limit, int):
        limit_clause = f"LIMIT {limit}"
    else:
        limit_clause = ""

    q = (
        f"MATCH (n:{class_name})"
        f"  {filter_clause}"
        f"  RETURN n {keep_clause}"
        f"  {limit_clause}"
    )
    cursor = self.execute(q)
    r = [item["n"] for item in cursor.data()]
    return r

fetch_edges(from_type, from_id, edge_type=None, to_type=None, to_id=None, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs)

Fetch edges from Neo4j using Cypher.

Parameters:

Name	Type	Description	Default
from_type	str	Source node label	required
from_id	str	Source node ID (property name depends on match_keys used)	required
edge_type	str | None	Optional relationship type to filter by	None
to_type	str | None	Optional target node label to filter by	None
to_id	str | None	Optional target node ID to filter by	None
filters	list | dict | None	Additional query filters	None
limit	int | None	Maximum number of edges to return	None
return_keys	list | None	Keys to return (projection)	None
unset_keys	list | None	Keys to exclude (projection) - not supported in Neo4j	None
**kwargs		Additional parameters	{}

Returns:

Name	Type	Description
list		List of fetched edges

Source code in graflo/db/neo4j/conn.py

def fetch_edges(
    self,
    from_type: str,
    from_id: str,
    edge_type: str | None = None,
    to_type: str | None = None,
    to_id: str | None = None,
    filters: list | dict | None = None,
    limit: int | None = None,
    return_keys: list | None = None,
    unset_keys: list | None = None,
    **kwargs,
):
    """Fetch edges from Neo4j using Cypher.

    Args:
        from_type: Source node label
        from_id: Source node ID (property name depends on match_keys used)
        edge_type: Optional relationship type to filter by
        to_type: Optional target node label to filter by
        to_id: Optional target node ID to filter by
        filters: Additional query filters
        limit: Maximum number of edges to return
        return_keys: Keys to return (projection)
        unset_keys: Keys to exclude (projection) - not supported in Neo4j
        **kwargs: Additional parameters

    Returns:
        list: List of fetched edges
    """
    # Build Cypher query to fetch edges
    # Match source node first
    source_match = f"(source:{from_type} {{id: '{from_id}'}})"

    # Build relationship pattern
    if edge_type:
        rel_pattern = f"-[r:{edge_type}]->"
    else:
        rel_pattern = "-[r]->"

    # Build target node match
    if to_type:
        target_match = f"(target:{to_type})"
    else:
        target_match = "(target)"

    # Add target ID filter if provided
    where_clauses = []
    if to_id:
        where_clauses.append(f"target.id = '{to_id}'")

    # Add additional filters if provided
    if filters is not None:
        ff = FilterExpression.from_dict(filters)
        filter_clause = ff(doc_name="r", kind=self.expression_flavor())
        where_clauses.append(filter_clause)

    where_clause = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""

    # Build return clause
    if return_keys is not None:
        return_clause = ", ".join([f"r.{key} as {key}" for key in return_keys])
        return_clause = f"RETURN {return_clause}"
    else:
        return_clause = "RETURN r"

    limit_clause = f"LIMIT {limit}" if limit else ""

    query = f"""
        MATCH {source_match}{rel_pattern}{target_match}
        {where_clause}
        {return_clause}
        {limit_clause}
    """

    cursor = self.execute(query)
    result = [item["r"] for item in cursor.data()]

    # Note: unset_keys is not supported in Neo4j as we can't modify the result structure
    # after the query

    return result

fetch_present_documents(batch, class_name, match_keys, keep_keys=None, flatten=False, filters=None)

Fetch nodes that exist in the database.

Note: Not implemented in Neo4j.

Parameters:

Name	Type	Description	Default
batch	list[dict[str, Any]]	Batch of documents to check	required
class_name	str	Label to check in	required
match_keys	list[str] | tuple[str, ...]	Keys to match nodes	required
keep_keys	list[str] | tuple[str, ...] | None	Keys to keep in result	None
flatten	bool	Unused in Neo4j	False
filters	list[Any] | dict[str, Any] | None	Additional query filters	None

Raises:

Type	Description
NotImplementedError	This method is not implemented for Neo4j

Source code in graflo/db/neo4j/conn.py

def fetch_present_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    flatten: bool = False,
    filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """Fetch nodes that exist in the database.

    Note: Not implemented in Neo4j.

    Args:
        batch: Batch of documents to check
        class_name: Label to check in
        match_keys: Keys to match nodes
        keep_keys: Keys to keep in result
        flatten: Unused in Neo4j
        filters: Additional query filters

    Raises:
        NotImplementedError: This method is not implemented for Neo4j
    """
    raise NotImplementedError

init_db(schema, recreate_schema)

Initialize Neo4j with the given schema.

Checks if the database exists and creates it if it doesn't. Uses schema.general.name if database is not set in config. Note: Database creation is only supported in Neo4j Enterprise Edition.

If the database already has nodes and recreate_schema is False, raises SchemaExistsError and the script halts.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing graph structure definitions	required
recreate_schema	bool	If True, delete all existing data before initialization. If False and database has nodes, raises SchemaExistsError.	required

Source code in graflo/db/neo4j/conn.py

def init_db(self, schema: Schema, recreate_schema: bool) -> None:
    """Initialize Neo4j with the given schema.

    Checks if the database exists and creates it if it doesn't.
    Uses schema.general.name if database is not set in config.
    Note: Database creation is only supported in Neo4j Enterprise Edition.

    If the database already has nodes and recreate_schema is False, raises
    SchemaExistsError and the script halts.

    Args:
        schema: Schema containing graph structure definitions
        recreate_schema: If True, delete all existing data before initialization.
            If False and database has nodes, raises SchemaExistsError.
    """
    # Determine database name: use config.database if set, otherwise use schema.general.name
    db_name = self.config.database
    if not db_name:
        db_name = schema.general.name
        # Update config for subsequent operations
        self.config.database = db_name

    # Check if database exists and create it if it doesn't
    # Note: This only works in Neo4j Enterprise Edition
    # For Community Edition, create_database will handle it gracefully
    # Community Edition only allows one database per instance
    try:
        # Try to check if database exists (Enterprise feature)
        try:
            result = self.execute("SHOW DATABASES")
            # Neo4j result is a cursor-like object, iterate to get records
            databases = []
            for record in result:
                # Record structure may vary, try common field names
                if hasattr(record, "get"):
                    db_name_field = (
                        record.get("name")
                        or record.get("database")
                        or record.get("db")
                    )
                else:
                    # If record is a dict-like object, try direct access
                    db_name_field = getattr(record, "name", None) or getattr(
                        record, "database", None
                    )
                if db_name_field:
                    databases.append(db_name_field)

            if db_name not in databases:
                logger.info(
                    f"Database '{db_name}' does not exist, attempting to create it..."
                )
                # create_database handles Community Edition errors gracefully
                self.create_database(db_name)
        except Exception as show_error:
            # If SHOW DATABASES fails (Community Edition or older versions), try to create anyway
            logger.debug(
                f"Could not check database existence (may be Community Edition): {show_error}"
            )
            # create_database handles Community Edition errors gracefully
            self.create_database(db_name)
    except Exception as e:
        # Only log unexpected errors (create_database handles Community Edition gracefully)
        logger.error(
            f"Unexpected error during database initialization for '{db_name}': {e}",
            exc_info=True,
        )
        # Don't raise - allow operation to continue with default database
        logger.warning(
            "Continuing with default database due to initialization error"
        )

    try:
        # Check if database already has nodes (schema/graph exists)
        result = self.execute("MATCH (n) RETURN count(n) AS c")
        count = 0
        if hasattr(result, "data"):
            data = result.data()
            if data:
                first = data[0]
                count = first.get("c", 0) or 0
        if count == 0 and hasattr(result, "__iter__"):
            for record in result:
                if hasattr(record, "get"):
                    count = record.get("c", 0) or 0
                else:
                    count = getattr(record, "c", 0) or 0
                break
        if count > 0 and not recreate_schema:
            raise SchemaExistsError(
                f"Schema/graph already exists in database '{db_name}' ({count} nodes). "
                "Set recreate_schema=True to replace, or use clear_data=True before ingestion."
            )

        if recreate_schema:
            try:
                self.delete_database("")
                logger.debug(f"Cleaned database '{db_name}' for fresh start")
            except Exception as clean_error:
                logger.warning(
                    f"Error during recreate_schema for database '{db_name}': {clean_error}",
                    exc_info=True,
                )
                # Continue - may be first run or already clean

        try:
            self.define_indexes(schema)
            logger.debug(f"Defined indexes for database '{db_name}'")
        except Exception as index_error:
            logger.error(
                f"Failed to define indexes for database '{db_name}': {index_error}",
                exc_info=True,
            )
            raise
    except SchemaExistsError:
        raise
    except Exception as e:
        logger.error(
            f"Error during database schema initialization for '{db_name}': {e}",
            exc_info=True,
        )
        raise

insert_edges_batch(docs_edges, source_class, target_class, relation_name, match_keys_source, match_keys_target, filter_uniques=True, head=None, **kwargs)

Insert a batch of relationships using Cypher.

Creates relationships between source and target nodes, with support for property matching and unique constraints.

Parameters:

Name	Type	Description	Default
docs_edges	list[list[dict[str, Any]]] | list[Any] | None	List of edge documents in format [{__source: source_doc, __target: target_doc}]	required
source_class	str	Source node label	required
target_class	str	Target node label	required
relation_name	str	Relationship type name	required
match_keys_source	tuple[str, ...]	Keys to match source nodes	required
match_keys_target	tuple[str, ...]	Keys to match target nodes	required
filter_uniques	bool	Unused in Neo4j (MERGE handles uniqueness automatically)	True
head	int | None	Optional limit on number of relationships to insert	None
**kwargs	Any	Additional options: - dry: If True, don't execute the query - collection_name: Unused in Neo4j (kept for interface compatibility) - uniq_weight_fields: Unused in Neo4j (ArangoDB-specific) - uniq_weight_collections: Unused in Neo4j (ArangoDB-specific) - upsert_option: Unused in Neo4j (ArangoDB-specific, MERGE is always upsert)	{}

Source code in graflo/db/neo4j/conn.py

def insert_edges_batch(
    self,
    docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
    source_class: str,
    target_class: str,
    relation_name: str,
    match_keys_source: tuple[str, ...],
    match_keys_target: tuple[str, ...],
    filter_uniques: bool = True,
    head: int | None = None,
    **kwargs: Any,
) -> None:
    """Insert a batch of relationships using Cypher.

    Creates relationships between source and target nodes, with support for
    property matching and unique constraints.

    Args:
        docs_edges: List of edge documents in format [{__source: source_doc, __target: target_doc}]
        source_class: Source node label
        target_class: Target node label
        relation_name: Relationship type name
        match_keys_source: Keys to match source nodes
        match_keys_target: Keys to match target nodes
        filter_uniques: Unused in Neo4j (MERGE handles uniqueness automatically)
        head: Optional limit on number of relationships to insert
        **kwargs: Additional options:
            - dry: If True, don't execute the query
            - collection_name: Unused in Neo4j (kept for interface compatibility)
            - uniq_weight_fields: Unused in Neo4j (ArangoDB-specific)
            - uniq_weight_collections: Unused in Neo4j (ArangoDB-specific)
            - upsert_option: Unused in Neo4j (ArangoDB-specific, MERGE is always upsert)
    """
    dry = kwargs.pop("dry", False)
    # Extract and ignore unused parameters (kept for interface compatibility)
    kwargs.pop("collection_name", None)
    kwargs.pop("uniq_weight_fields", None)
    kwargs.pop("uniq_weight_collections", None)
    kwargs.pop("upsert_option", None)

    # Apply head limit if specified
    if head is not None and isinstance(docs_edges, list):
        docs_edges = docs_edges[:head]

    # Note: filter_uniques is unused because Neo4j's MERGE handles uniqueness automatically

    source_match_str = [f"source.{key} = row[0].{key}" for key in match_keys_source]
    target_match_str = [f"target.{key} = row[1].{key}" for key in match_keys_target]

    match_clause = "WHERE " + " AND ".join(source_match_str + target_match_str)

    q = f"""
        WITH $batch AS batch 
        UNWIND batch as row 
        MATCH (source:{source_class}), 
              (target:{target_class}) {match_clause} 
                    MERGE (source)-[r:{relation_name}]->(target)
            SET r += row[2]

    """
    if not dry:
        self.execute(q, batch=docs_edges)

insert_return_batch(docs, class_name)

Insert nodes and return their properties.

Note: Not implemented in Neo4j.

Parameters:

Name	Type	Description	Default
docs	list[dict[str, Any]]	Documents to insert	required
class_name	str	Label to insert into	required

Raises:

Type	Description
NotImplementedError	This method is not implemented for Neo4j

Source code in graflo/db/neo4j/conn.py

def insert_return_batch(
    self, docs: list[dict[str, Any]], class_name: str
) -> list[dict[str, Any]] | str:
    """Insert nodes and return their properties.

    Note: Not implemented in Neo4j.

    Args:
        docs: Documents to insert
        class_name: Label to insert into

    Raises:
        NotImplementedError: This method is not implemented for Neo4j
    """
    raise NotImplementedError()

keep_absent_documents(batch, class_name, match_keys, keep_keys=None, filters=None)

Keep nodes that don't exist in the database.

Note: Not implemented in Neo4j.

Parameters:

Name	Type	Description	Default
batch	list[dict[str, Any]]	Batch of documents to check	required
class_name	str	Label to check in	required
match_keys	list[str] | tuple[str, ...]	Keys to match nodes	required
keep_keys	list[str] | tuple[str, ...] | None	Keys to keep in result	None
filters	list[Any] | dict[str, Any] | None	Additional query filters	None

Raises:

Type	Description
NotImplementedError	This method is not implemented for Neo4j

Source code in graflo/db/neo4j/conn.py

def keep_absent_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """Keep nodes that don't exist in the database.

    Note: Not implemented in Neo4j.

    Args:
        batch: Batch of documents to check
        class_name: Label to check in
        match_keys: Keys to match nodes
        keep_keys: Keys to keep in result
        filters: Additional query filters

    Raises:
        NotImplementedError: This method is not implemented for Neo4j
    """
    raise NotImplementedError

upsert_docs_batch(docs, class_name, match_keys, **kwargs)

Upsert a batch of nodes using Cypher.

Performs an upsert operation on a batch of nodes, using the specified match keys to determine whether to update existing nodes or create new ones.

Parameters:

Name	Type	Description	Default
docs		List of node documents to upsert	required
class_name		Label to upsert into	required
match_keys		Keys to match for upsert operation	required
**kwargs		Additional options: - dry: If True, don't execute the query	{}

Source code in graflo/db/neo4j/conn.py

def upsert_docs_batch(self, docs, class_name, match_keys, **kwargs):
    """Upsert a batch of nodes using Cypher.

    Performs an upsert operation on a batch of nodes, using the specified
    match keys to determine whether to update existing nodes or create new ones.

    Args:
        docs: List of node documents to upsert
        class_name: Label to upsert into
        match_keys: Keys to match for upsert operation
        **kwargs: Additional options:
            - dry: If True, don't execute the query
    """
    dry = kwargs.pop("dry", False)

    index_str = ", ".join([f"{k}: row.{k}" for k in match_keys])
    q = f"""
        WITH $batch AS batch 
        UNWIND batch as row 
        MERGE (n:{class_name} {{ {index_str} }}) 
        ON MATCH set n += row 
        ON CREATE set n += row
    """
    if not dry:
        self.execute(q, batch=docs)

PostgresConnection

PostgreSQL connection for schema introspection.

This class provides PostgreSQL-specific functionality for connecting to databases and introspecting 3NF schemas to identify vertex-like and edge-like tables.

Attributes:

Name	Type	Description
config		PostgreSQL connection configuration
conn		psycopg2 connection instance

Source code in graflo/db/postgres/conn.py

class PostgresConnection:
    """PostgreSQL connection for schema introspection.

    This class provides PostgreSQL-specific functionality for connecting to databases
    and introspecting 3NF schemas to identify vertex-like and edge-like tables.

    Attributes:
        config: PostgreSQL connection configuration
        conn: psycopg2 connection instance
    """

    def __init__(self, config: PostgresConfig):
        """Initialize PostgreSQL connection.

        Args:
            config: PostgreSQL connection configuration containing URI and credentials
        """
        self.config = config

        # Validate required config values
        if config.uri is None:
            raise ValueError("PostgreSQL connection requires a URI to be configured")
        if config.database is None:
            raise ValueError(
                "PostgreSQL connection requires a database name to be configured"
            )

        # Use config properties directly - all fallbacks are handled in PostgresConfig
        host = config.hostname or "localhost"
        port = int(config.port) if config.port else 5432
        database = config.database
        user = config.username or "postgres"
        password = config.password

        # Build connection parameters dict
        conn_params = {
            "host": host,
            "port": port,
            "database": database,
            "user": user,
        }

        if password:
            conn_params["password"] = password

        try:
            self.conn = psycopg2.connect(**conn_params)
            logger.info(f"Successfully connected to PostgreSQL database '{database}'")
        except Exception as e:
            logger.error(f"Failed to connect to PostgreSQL: {e}", exc_info=True)
            raise

    def read(self, query: str, params: tuple | None = None) -> list[dict[str, Any]]:
        """Execute a SELECT query and return results as a list of dictionaries.

        Args:
            query: SQL SELECT query to execute
            params: Optional tuple of parameters for parameterized queries

        Returns:
            List of dictionaries, where each dictionary represents a row with column names as keys.
            Decimal values are converted to float for compatibility with graph databases.
        """
        from decimal import Decimal

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            if params:
                cursor.execute(query, params)
            else:
                cursor.execute(query)

            # Convert rows to dictionaries and convert Decimal to float
            results = []
            for row in cursor.fetchall():
                row_dict = dict(row)
                # Convert Decimal to float for JSON/graph database compatibility
                for key, value in row_dict.items():
                    if isinstance(value, Decimal):
                        row_dict[key] = float(value)
                results.append(row_dict)

            return results

    def __enter__(self):
        """Enter the context manager.

        Returns:
            PostgresConnection: Self for use in 'with' statements
        """
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        """Exit the context manager.

        Ensures the connection is properly closed when exiting the context.

        Args:
            exc_type: Exception type if an exception occurred
            exc_value: Exception value if an exception occurred
            exc_traceback: Exception traceback if an exception occurred
        """
        self.close()
        return False  # Don't suppress exceptions

    def close(self):
        """Close the PostgreSQL connection."""
        if hasattr(self, "conn") and self.conn:
            try:
                self.conn.close()
                logger.debug("PostgreSQL connection closed")
            except Exception as e:
                logger.warning(
                    f"Error closing PostgreSQL connection: {e}", exc_info=True
                )

    def _check_information_schema_reliable(self, schema_name: str) -> bool:
        """Check if information_schema is reliable for the given schema.

        Args:
            schema_name: Schema name to check

        Returns:
            True if information_schema appears reliable, False otherwise
        """
        try:
            # Try to query information_schema.tables
            query = """
                SELECT COUNT(*) as count
                FROM information_schema.tables
                WHERE table_schema = %s
                  AND table_type = 'BASE TABLE'
            """
            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (schema_name,))
                result = cursor.fetchone()
                # If query succeeds, check if we can also query constraints
                pk_query = """
                    SELECT COUNT(*) as count
                    FROM information_schema.table_constraints tc
                    JOIN information_schema.key_column_usage kcu
                        ON tc.constraint_name = kcu.constraint_name
                        AND tc.table_schema = kcu.table_schema
                    WHERE tc.constraint_type = 'PRIMARY KEY'
                      AND tc.table_schema = %s
                """
                cursor.execute(pk_query, (schema_name,))
                pk_result = cursor.fetchone()
                # If both queries work, information_schema seems reliable
                return result is not None and pk_result is not None
        except Exception as e:
            logger.debug(f"information_schema check failed: {e}")
            return False

    def _get_tables_pg_catalog(self, schema_name: str) -> list[dict[str, Any]]:
        """Get all tables using pg_catalog (fallback method).

        Args:
            schema_name: Schema name to query

        Returns:
            List of table information dictionaries with keys: table_name, table_schema
        """
        query = """
            SELECT
                c.relname as table_name,
                n.nspname as table_schema
            FROM pg_catalog.pg_class c
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            WHERE n.nspname = %s
              AND c.relkind = 'r'
              AND NOT c.relispartition
            ORDER BY c.relname;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name,))
            return [dict(row) for row in cursor.fetchall()]

    def get_tables(self, schema_name: str | None = None) -> list[dict[str, Any]]:
        """Get all tables in the specified schema.

        Tries information_schema first, falls back to pg_catalog if needed.

        Args:
            schema_name: Schema name to query. If None, uses 'public' or config schema_name.

        Returns:
            List of table information dictionaries with keys: table_name, table_schema
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        # Try information_schema first
        try:
            query = """
                SELECT table_name, table_schema
                FROM information_schema.tables
                WHERE table_schema = %s
                  AND table_type = 'BASE TABLE'
                ORDER BY table_name;
            """

            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (schema_name,))
                results = [dict(row) for row in cursor.fetchall()]
                # If we got results, check if information_schema is reliable
                if results and self._check_information_schema_reliable(schema_name):
                    return results
                # If no results or unreliable, fall back to pg_catalog
                logger.debug(
                    f"information_schema returned no results or is unreliable, "
                    f"falling back to pg_catalog for schema '{schema_name}'"
                )
        except Exception as e:
            logger.debug(
                f"information_schema query failed: {e}, falling back to pg_catalog"
            )

        # Fallback to pg_catalog
        return self._get_tables_pg_catalog(schema_name)

    def _get_table_columns_pg_catalog(
        self, table_name: str, schema_name: str
    ) -> list[dict[str, Any]]:
        """Get columns using pg_catalog (fallback method).

        Args:
            table_name: Name of the table
            schema_name: Schema name

        Returns:
            List of column information dictionaries with keys:
            name, type, description, is_nullable, column_default
        """
        query = """
            SELECT
                a.attname as name,
                pg_catalog.format_type(a.atttypid, a.atttypmod) as type,
                CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END as is_nullable,
                pg_catalog.pg_get_expr(d.adbin, d.adrelid) as column_default,
                COALESCE(dsc.description, '') as description,
                a.attnum as ordinal_position
            FROM pg_catalog.pg_attribute a
            JOIN pg_catalog.pg_class c ON c.oid = a.attrelid
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            LEFT JOIN pg_catalog.pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum
            LEFT JOIN pg_catalog.pg_description dsc ON dsc.objoid = a.attrelid AND dsc.objsubid = a.attnum
            WHERE n.nspname = %s
              AND c.relname = %s
              AND a.attnum > 0
              AND NOT a.attisdropped
            ORDER BY a.attnum;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name, table_name))
            columns = []
            for row in cursor.fetchall():
                col_dict = dict(row)
                # Normalize type format
                if col_dict["type"]:
                    # Remove length info from type if present (e.g., "character varying(255)" -> "varchar")
                    type_str = col_dict["type"]
                    if "(" in type_str:
                        base_type = type_str.split("(")[0]
                        # Map common types
                        type_mapping = {
                            "character varying": "varchar",
                            "character": "char",
                            "double precision": "float8",
                            "real": "float4",
                            "integer": "int4",
                            "bigint": "int8",
                            "smallint": "int2",
                        }
                        col_dict["type"] = type_mapping.get(
                            base_type.lower(), base_type.lower()
                        )
                    else:
                        type_mapping = {
                            "character varying": "varchar",
                            "character": "char",
                            "double precision": "float8",
                            "real": "float4",
                            "integer": "int4",
                            "bigint": "int8",
                            "smallint": "int2",
                        }
                        col_dict["type"] = type_mapping.get(
                            type_str.lower(), type_str.lower()
                        )
                columns.append(col_dict)
            return columns

    def get_table_columns(
        self, table_name: str, schema_name: str | None = None
    ) -> list[dict[str, Any]]:
        """Get columns for a specific table with types and descriptions.

        Tries information_schema first, falls back to pg_catalog if needed.

        Args:
            table_name: Name of the table
            schema_name: Schema name. If None, uses 'public' or config schema_name.

        Returns:
            List of column information dictionaries with keys:
            name, type, description, is_nullable, column_default
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        # Try information_schema first
        try:
            query = """
                SELECT
                    c.column_name as name,
                    c.data_type as type,
                    c.udt_name as udt_name,
                    c.character_maximum_length,
                    c.is_nullable,
                    c.column_default,
                    COALESCE(d.description, '') as description,
                    c.ordinal_position as ordinal_position
                FROM information_schema.columns c
                LEFT JOIN pg_catalog.pg_statio_all_tables st
                    ON st.schemaname = c.table_schema
                    AND st.relname = c.table_name
                LEFT JOIN pg_catalog.pg_description d
                    ON d.objoid = st.relid
                    AND d.objsubid = c.ordinal_position
                WHERE c.table_schema = %s
                  AND c.table_name = %s
                ORDER BY c.ordinal_position;
            """

            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (schema_name, table_name))
                columns = []
                for row in cursor.fetchall():
                    col_dict = dict(row)
                    # Format type with length if applicable
                    if col_dict["character_maximum_length"]:
                        col_dict["type"] = (
                            f"{col_dict['type']}({col_dict['character_maximum_length']})"
                        )
                    # Use udt_name if it's more specific (e.g., varchar, int4)
                    if (
                        col_dict["udt_name"]
                        and col_dict["udt_name"] != col_dict["type"]
                    ):
                        col_dict["type"] = col_dict["udt_name"]
                    # Remove helper fields
                    col_dict.pop("character_maximum_length", None)
                    col_dict.pop("udt_name", None)
                    columns.append(col_dict)

                # If we got results and information_schema is reliable, return them
                if columns and self._check_information_schema_reliable(schema_name):
                    return columns
                # Otherwise fall back to pg_catalog
                logger.debug(
                    f"information_schema returned no results or is unreliable, "
                    f"falling back to pg_catalog for table '{schema_name}.{table_name}'"
                )
        except Exception as e:
            logger.debug(
                f"information_schema query failed: {e}, falling back to pg_catalog"
            )

        # Fallback to pg_catalog
        return self._get_table_columns_pg_catalog(table_name, schema_name)

    def _get_primary_keys_pg_catalog(
        self, table_name: str, schema_name: str
    ) -> list[str]:
        """Get primary key columns using pg_catalog (fallback method).

        Args:
            table_name: Name of the table
            schema_name: Schema name

        Returns:
            List of primary key column names
        """
        query = """
            SELECT a.attname
            FROM pg_catalog.pg_constraint con
            JOIN pg_catalog.pg_class c ON c.oid = con.conrelid
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            JOIN pg_catalog.pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = ANY(con.conkey)
            WHERE n.nspname = %s
              AND c.relname = %s
              AND con.contype = 'p'
            ORDER BY array_position(con.conkey, a.attnum);
        """

        with self.conn.cursor() as cursor:
            cursor.execute(query, (schema_name, table_name))
            return [row[0] for row in cursor.fetchall()]

    def get_primary_keys(
        self, table_name: str, schema_name: str | None = None
    ) -> list[str]:
        """Get primary key columns for a table.

        Tries information_schema first, falls back to pg_catalog if needed.

        Args:
            table_name: Name of the table
            schema_name: Schema name. If None, uses 'public' or config schema_name.

        Returns:
            List of primary key column names
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        # Try information_schema first
        try:
            query = """
                SELECT kcu.column_name
                FROM information_schema.table_constraints tc
                JOIN information_schema.key_column_usage kcu
                    ON tc.constraint_name = kcu.constraint_name
                    AND tc.table_schema = kcu.table_schema
                WHERE tc.constraint_type = 'PRIMARY KEY'
                  AND tc.table_schema = %s
                  AND tc.table_name = %s
                ORDER BY kcu.ordinal_position;
            """

            with self.conn.cursor() as cursor:
                cursor.execute(query, (schema_name, table_name))
                results = [row[0] for row in cursor.fetchall()]
                # If we got results and information_schema is reliable, return them
                if results and self._check_information_schema_reliable(schema_name):
                    return results
                # Otherwise fall back to pg_catalog
                logger.debug(
                    f"information_schema returned no results or is unreliable, "
                    f"falling back to pg_catalog for primary keys of '{schema_name}.{table_name}'"
                )
        except Exception as e:
            logger.debug(
                f"information_schema query failed: {e}, falling back to pg_catalog"
            )

        # Fallback to pg_catalog
        return self._get_primary_keys_pg_catalog(table_name, schema_name)

    def _get_unique_columns_pg_catalog(
        self, table_name: str, schema_name: str
    ) -> list[str]:
        """Get columns in UNIQUE constraints using pg_catalog (fallback method)."""
        query = """
            SELECT a.attname
            FROM pg_catalog.pg_constraint con
            JOIN pg_catalog.pg_class c ON c.oid = con.conrelid
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            JOIN pg_catalog.pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = ANY(con.conkey)
            WHERE n.nspname = %s
              AND c.relname = %s
              AND con.contype = 'u'
              AND a.attnum > 0
              AND NOT a.attisdropped
            ORDER BY array_position(con.conkey, a.attnum);
        """
        with self.conn.cursor() as cursor:
            cursor.execute(query, (schema_name, table_name))
            return list(dict.fromkeys(row[0] for row in cursor.fetchall()))

    def get_unique_columns(
        self, table_name: str, schema_name: str | None = None
    ) -> list[str]:
        """Get column names that participate in any UNIQUE constraint.

        Tries information_schema first, falls back to pg_catalog if needed.
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        try:
            query = """
                SELECT kcu.column_name
                FROM information_schema.table_constraints tc
                JOIN information_schema.key_column_usage kcu
                    ON tc.constraint_name = kcu.constraint_name
                    AND tc.table_schema = kcu.table_schema
                WHERE tc.constraint_type = 'UNIQUE'
                  AND tc.table_schema = %s
                  AND tc.table_name = %s
                ORDER BY kcu.ordinal_position;
            """
            with self.conn.cursor() as cursor:
                cursor.execute(query, (schema_name, table_name))
                results = list(dict.fromkeys(row[0] for row in cursor.fetchall()))
                if results and self._check_information_schema_reliable(schema_name):
                    return results
        except Exception as e:
            logger.debug(
                f"information_schema query failed for unique columns: {e}, "
                "falling back to pg_catalog"
            )
        return self._get_unique_columns_pg_catalog(table_name, schema_name)

    def _get_foreign_keys_pg_catalog(
        self, table_name: str, schema_name: str
    ) -> list[dict[str, Any]]:
        """Get foreign key relationships using pg_catalog (fallback method).

        Handles both single-column and multi-column foreign keys.
        For multi-column foreign keys, returns one row per column.

        Args:
            table_name: Name of the table
            schema_name: Schema name

        Returns:
            List of foreign key dictionaries with keys:
            column, references_table, references_column, constraint_name
        """
        # Use generate_subscripts for better compatibility with older PostgreSQL versions
        query = """
            SELECT
                a.attname as column,
                ref_c.relname as references_table,
                ref_a.attname as references_column,
                con.conname as constraint_name
            FROM pg_catalog.pg_constraint con
            JOIN pg_catalog.pg_class c ON c.oid = con.conrelid
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            JOIN pg_catalog.pg_class ref_c ON ref_c.oid = con.confrelid
            JOIN generate_subscripts(con.conkey, 1) AS i ON true
            JOIN pg_catalog.pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = con.conkey[i]
            JOIN pg_catalog.pg_attribute ref_a ON ref_a.attrelid = con.confrelid AND ref_a.attnum = con.confkey[i]
            WHERE n.nspname = %s
              AND c.relname = %s
              AND con.contype = 'f'
            ORDER BY con.conname, i;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name, table_name))
            return [dict(row) for row in cursor.fetchall()]

    def get_foreign_keys(
        self, table_name: str, schema_name: str | None = None
    ) -> list[dict[str, Any]]:
        """Get foreign key relationships for a table.

        Tries information_schema first, falls back to pg_catalog if needed.

        Args:
            table_name: Name of the table
            schema_name: Schema name. If None, uses 'public' or config schema_name.

        Returns:
            List of foreign key dictionaries with keys:
            column, references_table, references_column, constraint_name
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        # Try information_schema first
        try:
            query = """
                SELECT
                    kcu.column_name as column,
                    ccu.table_name as references_table,
                    ccu.column_name as references_column,
                    tc.constraint_name
                FROM information_schema.table_constraints tc
                JOIN information_schema.key_column_usage kcu
                    ON tc.constraint_name = kcu.constraint_name
                    AND tc.table_schema = kcu.table_schema
                JOIN information_schema.constraint_column_usage ccu
                    ON ccu.constraint_name = tc.constraint_name
                    AND ccu.table_schema = tc.table_schema
                WHERE tc.constraint_type = 'FOREIGN KEY'
                  AND tc.table_schema = %s
                  AND tc.table_name = %s
                ORDER BY kcu.ordinal_position;
            """

            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (schema_name, table_name))
                results = [dict(row) for row in cursor.fetchall()]
                # If we got results and information_schema is reliable, return them
                if self._check_information_schema_reliable(schema_name):
                    return results
                # Otherwise fall back to pg_catalog
                logger.debug(
                    f"information_schema returned no results or is unreliable, "
                    f"falling back to pg_catalog for foreign keys of '{schema_name}.{table_name}'"
                )
        except Exception as e:
            logger.debug(
                f"information_schema query failed: {e}, falling back to pg_catalog"
            )

        # Fallback to pg_catalog
        return self._get_foreign_keys_pg_catalog(table_name, schema_name)

    def get_table_row_count_estimate(
        self, table_name: str, schema_name: str | None = None
    ) -> int | None:
        """Return approximate row count from pg_class.reltuples (updated by ANALYZE).

        Avoids full table scan; may be stale until next ANALYZE.
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"
        query = """
            SELECT c.reltuples::bigint AS estimate
            FROM pg_catalog.pg_class c
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            WHERE n.nspname = %s AND c.relname = %s AND c.relkind = 'r';
        """
        with self.conn.cursor() as cursor:
            cursor.execute(query, (schema_name, table_name))
            row = cursor.fetchone()
            if row is None:
                return None
            val = row[0]
            return int(val) if val is not None else None

    def get_table_sample_rows(
        self,
        table_name: str,
        schema_name: str | None = None,
        limit: int = 5,
    ) -> list[dict[str, Any]]:
        """Return first `limit` rows from the table (no ORDER BY for speed)."""
        if schema_name is None:
            schema_name = self.config.schema_name or "public"
        query = sql.SQL("SELECT * FROM {}.{} LIMIT %s").format(
            sql.Identifier(schema_name),
            sql.Identifier(table_name),
        )
        try:
            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (limit,))
                return [dict(row) for row in cursor.fetchall()]
        except Exception as e:
            logger.debug(
                f"Could not fetch sample rows for '{schema_name}.{table_name}': {e}"
            )
            return []

    def _is_edge_like_table(
        self, table_name: str, pk_columns: list[str], fk_columns: list[dict[str, Any]]
    ) -> bool:
        """Determine if a table is edge-like based on heuristics.

        Heuristics:
        1. Tables with 2 or more primary keys are likely edge tables
        2. Tables with exactly 2 foreign keys are likely edge tables
        3. Tables with names starting with 'rel_' are likely edge tables
        4. Tables where primary key columns match foreign key columns are likely edge tables

        Args:
            table_name: Name of the table
            pk_columns: List of primary key column names
            fk_columns: List of foreign key dictionaries

        Returns:
            True if table appears to be edge-like, False otherwise
        """
        # Heuristic 1: Tables with 2 or more primary keys are likely edge tables
        if len(pk_columns) >= 2:
            return True

        # Heuristic 2: Tables with exactly 2 foreign keys are likely edge tables
        if len(fk_columns) == 2:
            return True

        # Heuristic 3: Tables with names starting with 'rel_' are likely edge tables
        if table_name.startswith("rel_"):
            return True

        # Heuristic 4: If primary key columns match foreign key columns, it's likely an edge table
        fk_column_names = {fk["column"] for fk in fk_columns}
        pk_set = set(pk_columns)
        # If all PK columns are FK columns and we have at least 2 FKs, it's likely an edge table
        if pk_set.issubset(fk_column_names) and len(fk_columns) >= 2:
            return True

        return False

    def detect_vertex_tables(
        self, schema_name: str | None = None
    ) -> list[VertexTableInfo]:
        """Detect vertex-like tables in the schema.

        Heuristic: Tables with a primary key and descriptive columns
        (not just foreign keys). These represent entities.

        Note: Tables identified as edge-like are excluded from vertex tables.

        Args:
            schema_name: Schema name. If None, uses 'public' or config schema_name.

        Returns:
            List of vertex table information dictionaries
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        tables = self.get_tables(schema_name)
        vertex_tables = []

        for table_info in tables:
            table_name = table_info["table_name"]
            pk_columns = self.get_primary_keys(table_name, schema_name)
            fk_columns = self.get_foreign_keys(table_name, schema_name)
            all_columns = self.get_table_columns(table_name, schema_name)

            # Vertex-like tables have:
            # 1. A primary key
            # 2. Not identified as edge-like tables
            # 3. Descriptive columns beyond just foreign keys

            if not pk_columns:
                continue  # Skip tables without primary keys

            # Skip edge-like tables
            if self._is_edge_like_table(table_name, pk_columns, fk_columns):
                continue

            # Count non-FK, non-PK columns (descriptive columns)
            fk_column_names = {fk["column"] for fk in fk_columns}
            pk_column_names = set(pk_columns)
            descriptive_columns = [
                col
                for col in all_columns
                if col["name"] not in fk_column_names
                and col["name"] not in pk_column_names
            ]

            # If table has descriptive columns, consider it vertex-like
            if descriptive_columns:
                # Mark primary key and unique columns and convert to ColumnInfo
                pk_set = set(pk_columns)
                unique_columns = self.get_unique_columns(table_name, schema_name)
                unique_set = set(unique_columns)
                column_infos = []
                for col in all_columns:
                    column_infos.append(
                        ColumnInfo(
                            name=col["name"],
                            type=col["type"],
                            description=col.get("description", ""),
                            is_nullable=col.get("is_nullable", "YES"),
                            column_default=col.get("column_default"),
                            is_pk=col["name"] in pk_set,
                            is_unique=col["name"] in unique_set,
                            ordinal_position=col.get("ordinal_position"),
                        )
                    )

                # Convert foreign keys to ForeignKeyInfo
                fk_infos = []
                for fk in fk_columns:
                    fk_infos.append(
                        ForeignKeyInfo(
                            column=fk["column"],
                            references_table=fk["references_table"],
                            references_column=fk.get("references_column"),
                            constraint_name=fk.get("constraint_name"),
                        )
                    )

                vertex_tables.append(
                    VertexTableInfo(
                        name=table_name,
                        schema_name=schema_name,
                        columns=column_infos,
                        primary_key=pk_columns,
                        foreign_keys=fk_infos,
                    )
                )

        return vertex_tables

    def detect_edge_tables(
        self,
        schema_name: str | None = None,
        vertex_table_names: list[str] | None = None,
    ) -> list[EdgeTableInfo]:
        """Detect edge-like tables in the schema.

        Heuristic: Tables with 2 or more primary keys, or exactly 2 foreign keys,
        or names starting with 'rel_'. These represent relationships between entities.

        Args:
            schema_name: Schema name. If None, uses 'public' or config schema_name.
            vertex_table_names: Optional list of vertex table names for fuzzy matching.
                              If None, will be inferred from detect_vertex_tables().

        Returns:
            List of edge table information dictionaries with source_table and target_table
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        # Get vertex table names if not provided
        if vertex_table_names is None:
            vertex_tables = self.detect_vertex_tables(schema_name)
            vertex_table_names = [vt.name for vt in vertex_tables]

        # Create fuzzy matcher once for all tables (significant performance improvement)
        # Caching is enabled by default for better performance
        matcher = FuzzyMatcher(vertex_table_names, threshold=0.6, enable_cache=True)

        tables = self.get_tables(schema_name)
        edge_tables = []

        for table_info in tables:
            table_name = table_info["table_name"]
            pk_columns = self.get_primary_keys(table_name, schema_name)
            fk_columns = self.get_foreign_keys(table_name, schema_name)

            # Skip tables without primary keys
            if not pk_columns:
                continue

            # Check if table is edge-like
            if not self._is_edge_like_table(table_name, pk_columns, fk_columns):
                continue

            all_columns = self.get_table_columns(table_name, schema_name)

            # Mark primary key and unique columns and convert to ColumnInfo
            pk_set = set(pk_columns)
            unique_columns = self.get_unique_columns(table_name, schema_name)
            unique_set = set(unique_columns)
            column_infos = []
            for col in all_columns:
                column_infos.append(
                    ColumnInfo(
                        name=col["name"],
                        type=col["type"],
                        description=col.get("description", ""),
                        is_nullable=col.get("is_nullable", "YES"),
                        column_default=col.get("column_default"),
                        is_pk=col["name"] in pk_set,
                        is_unique=col["name"] in unique_set,
                        ordinal_position=col.get("ordinal_position"),
                    )
                )

            # Convert foreign keys to ForeignKeyInfo
            fk_infos = []
            for fk in fk_columns:
                fk_infos.append(
                    ForeignKeyInfo(
                        column=fk["column"],
                        references_table=fk["references_table"],
                        references_column=fk.get("references_column"),
                        constraint_name=fk.get("constraint_name"),
                    )
                )

            # Determine source and target tables
            source_table = None
            target_table = None
            source_column = None
            target_column = None
            relation_name = None

            # If we have exactly 2 foreign keys, use them directly
            if len(fk_infos) == 2:
                source_fk = fk_infos[0]
                target_fk = fk_infos[1]
                source_table = source_fk.references_table
                target_table = target_fk.references_table
                source_column = source_fk.column
                target_column = target_fk.column
                # Still try to infer relation from table name
                fk_dicts = [
                    {
                        "column": fk.column,
                        "references_table": fk.references_table,
                    }
                    for fk in fk_infos
                ]
                _, _, relation_name = infer_edge_vertices_from_table_name(
                    table_name, pk_columns, fk_dicts, vertex_table_names, matcher
                )
            # If we have 2 or more primary keys, try to infer from table name and structure
            elif len(pk_columns) >= 2:
                # Convert fk_infos to dicts for _infer_edge_vertices_from_table_name
                fk_dicts = [
                    {
                        "column": fk.column,
                        "references_table": fk.references_table,
                    }
                    for fk in fk_infos
                ]

                # Try to infer from table name pattern
                inferred_source, inferred_target, relation_name = (
                    infer_edge_vertices_from_table_name(
                        table_name,
                        pk_columns,
                        fk_dicts,
                        vertex_table_names,
                        matcher,
                    )
                )

                if inferred_source and inferred_target:
                    source_table = inferred_source
                    target_table = inferred_target
                    # Try to match PK columns to FK columns for source/target columns
                    if fk_infos:
                        # Use first FK for source, second for target if available
                        if len(fk_infos) >= 2:
                            source_column = fk_infos[0].column
                            target_column = fk_infos[1].column
                        elif len(fk_infos) == 1:
                            # Self-reference case
                            source_column = fk_infos[0].column
                            target_column = fk_infos[0].column
                    else:
                        # Use PK columns as source/target columns
                        source_column = pk_columns[0]
                        target_column = (
                            pk_columns[1] if len(pk_columns) > 1 else pk_columns[0]
                        )
                elif fk_infos:
                    # Fallback: use FK references if available
                    if len(fk_infos) >= 2:
                        source_table = fk_infos[0].references_table
                        target_table = fk_infos[1].references_table
                        source_column = fk_infos[0].column
                        target_column = fk_infos[1].column
                    elif len(fk_infos) == 1:
                        source_table = fk_infos[0].references_table
                        target_table = fk_infos[0].references_table
                        source_column = fk_infos[0].column
                        target_column = fk_infos[0].column
                else:
                    # Last resort: use PK columns and infer table names from column names
                    source_column = pk_columns[0]
                    target_column = (
                        pk_columns[1] if len(pk_columns) > 1 else pk_columns[0]
                    )
                    # Use robust inference logic to extract vertex names from column names
                    source_table = infer_vertex_from_column_name(
                        source_column, vertex_table_names, matcher
                    )
                    target_table = infer_vertex_from_column_name(
                        target_column, vertex_table_names, matcher
                    )

            # Only add if we have source and target information
            if source_table and target_table:
                edge_tables.append(
                    EdgeTableInfo(
                        name=table_name,
                        schema_name=schema_name,
                        columns=column_infos,
                        primary_key=pk_columns,
                        foreign_keys=fk_infos,
                        source_table=source_table,
                        target_table=target_table,
                        source_column=source_column or pk_columns[0],
                        target_column=target_column
                        or (pk_columns[1] if len(pk_columns) > 1 else pk_columns[0]),
                        relation=relation_name,
                    )
                )
            else:
                logger.warning(
                    f"Could not determine source/target tables for edge-like table '{table_name}'. "
                    f"Skipping."
                )

        return edge_tables

    def _build_raw_tables(self, schema_name: str) -> list[RawTableInfo]:
        """Build raw table metadata for all tables in the schema."""
        tables = self.get_tables(schema_name)
        raw_tables = []
        for table_info in tables:
            table_name = table_info["table_name"]
            pk_columns = self.get_primary_keys(table_name, schema_name)
            fk_columns = self.get_foreign_keys(table_name, schema_name)
            unique_columns = self.get_unique_columns(table_name, schema_name)
            all_columns = self.get_table_columns(table_name, schema_name)
            row_count_estimate = self.get_table_row_count_estimate(
                table_name, schema_name
            )
            sample_rows = self.get_table_sample_rows(table_name, schema_name, limit=5)

            pk_set = set(pk_columns)
            unique_set = set(unique_columns)
            # Per-column sample values: list of values from first 5 rows (stringified)
            column_names = [c["name"] for c in all_columns]
            sample_by_col: dict[str, list[str]] = {c: [] for c in column_names}
            for row in sample_rows:
                for col_name in column_names:
                    if col_name in row and len(sample_by_col[col_name]) < 5:
                        v = row[col_name]
                        sample_by_col[col_name].append("NULL" if v is None else str(v))

            column_infos = []
            for col in all_columns:
                column_infos.append(
                    ColumnInfo(
                        name=col["name"],
                        type=col["type"],
                        description=col.get("description", ""),
                        is_nullable=col.get("is_nullable", "YES"),
                        column_default=col.get("column_default"),
                        is_pk=col["name"] in pk_set,
                        is_unique=col["name"] in unique_set,
                        ordinal_position=col.get("ordinal_position"),
                        sample_values=sample_by_col.get(col["name"], [])[:5],
                    )
                )

            fk_infos = [
                ForeignKeyInfo(
                    column=fk["column"],
                    references_table=fk["references_table"],
                    references_column=fk.get("references_column"),
                    constraint_name=fk.get("constraint_name"),
                )
                for fk in fk_columns
            ]

            raw_tables.append(
                RawTableInfo(
                    name=table_name,
                    schema_name=schema_name,
                    columns=column_infos,
                    primary_key=pk_columns,
                    foreign_keys=fk_infos,
                    row_count_estimate=row_count_estimate,
                )
            )
        return raw_tables

    def introspect_schema(
        self, schema_name: str | None = None
    ) -> SchemaIntrospectionResult:
        """Introspect the database schema and return structured information.

        This is the main method that analyzes the schema and returns information
        about vertex-like and edge-like tables.

        Args:
            schema_name: Schema name. If None, uses 'public' or config schema_name.

        Returns:
            SchemaIntrospectionResult with vertex_tables, edge_tables, and schema_name
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        logger.info(f"Introspecting PostgreSQL schema '{schema_name}'")

        vertex_tables = self.detect_vertex_tables(schema_name)
        edge_tables = self.detect_edge_tables(schema_name)
        raw_tables = self._build_raw_tables(schema_name)

        result = SchemaIntrospectionResult(
            vertex_tables=vertex_tables,
            edge_tables=edge_tables,
            raw_tables=raw_tables,
            schema_name=schema_name,
        )

        logger.info(
            f"Found {len(vertex_tables)} vertex-like tables and {len(edge_tables)} edge-like tables"
        )

        return result

__enter__()

Enter the context manager.

Returns:

Name	Type	Description
PostgresConnection		Self for use in 'with' statements

Source code in graflo/db/postgres/conn.py

def __enter__(self):
    """Enter the context manager.

    Returns:
        PostgresConnection: Self for use in 'with' statements
    """
    return self

__exit__(exc_type, exc_value, exc_traceback)

Exit the context manager.

Ensures the connection is properly closed when exiting the context.

Parameters:

Name	Type	Description	Default
exc_type		Exception type if an exception occurred	required
exc_value		Exception value if an exception occurred	required
exc_traceback		Exception traceback if an exception occurred	required

Source code in graflo/db/postgres/conn.py

def __exit__(self, exc_type, exc_value, exc_traceback):
    """Exit the context manager.

    Ensures the connection is properly closed when exiting the context.

    Args:
        exc_type: Exception type if an exception occurred
        exc_value: Exception value if an exception occurred
        exc_traceback: Exception traceback if an exception occurred
    """
    self.close()
    return False  # Don't suppress exceptions

__init__(config)

Initialize PostgreSQL connection.

Parameters:

Name	Type	Description	Default
config	PostgresConfig	PostgreSQL connection configuration containing URI and credentials	required

Source code in graflo/db/postgres/conn.py

def __init__(self, config: PostgresConfig):
    """Initialize PostgreSQL connection.

    Args:
        config: PostgreSQL connection configuration containing URI and credentials
    """
    self.config = config

    # Validate required config values
    if config.uri is None:
        raise ValueError("PostgreSQL connection requires a URI to be configured")
    if config.database is None:
        raise ValueError(
            "PostgreSQL connection requires a database name to be configured"
        )

    # Use config properties directly - all fallbacks are handled in PostgresConfig
    host = config.hostname or "localhost"
    port = int(config.port) if config.port else 5432
    database = config.database
    user = config.username or "postgres"
    password = config.password

    # Build connection parameters dict
    conn_params = {
        "host": host,
        "port": port,
        "database": database,
        "user": user,
    }

    if password:
        conn_params["password"] = password

    try:
        self.conn = psycopg2.connect(**conn_params)
        logger.info(f"Successfully connected to PostgreSQL database '{database}'")
    except Exception as e:
        logger.error(f"Failed to connect to PostgreSQL: {e}", exc_info=True)
        raise

close()

Close the PostgreSQL connection.

Source code in graflo/db/postgres/conn.py

def close(self):
    """Close the PostgreSQL connection."""
    if hasattr(self, "conn") and self.conn:
        try:
            self.conn.close()
            logger.debug("PostgreSQL connection closed")
        except Exception as e:
            logger.warning(
                f"Error closing PostgreSQL connection: {e}", exc_info=True
            )

detect_edge_tables(schema_name=None, vertex_table_names=None)

Detect edge-like tables in the schema.

Heuristic: Tables with 2 or more primary keys, or exactly 2 foreign keys, or names starting with 'rel_'. These represent relationships between entities.

Parameters:

Name	Type	Description	Default
schema_name	str | None	Schema name. If None, uses 'public' or config schema_name.	None
vertex_table_names	list[str] | None	Optional list of vertex table names for fuzzy matching. If None, will be inferred from detect_vertex_tables().	None

Returns:

Type	Description
list[EdgeTableInfo]	List of edge table information dictionaries with source_table and target_table

Source code in graflo/db/postgres/conn.py

def detect_edge_tables(
    self,
    schema_name: str | None = None,
    vertex_table_names: list[str] | None = None,
) -> list[EdgeTableInfo]:
    """Detect edge-like tables in the schema.

    Heuristic: Tables with 2 or more primary keys, or exactly 2 foreign keys,
    or names starting with 'rel_'. These represent relationships between entities.

    Args:
        schema_name: Schema name. If None, uses 'public' or config schema_name.
        vertex_table_names: Optional list of vertex table names for fuzzy matching.
                          If None, will be inferred from detect_vertex_tables().

    Returns:
        List of edge table information dictionaries with source_table and target_table
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    # Get vertex table names if not provided
    if vertex_table_names is None:
        vertex_tables = self.detect_vertex_tables(schema_name)
        vertex_table_names = [vt.name for vt in vertex_tables]

    # Create fuzzy matcher once for all tables (significant performance improvement)
    # Caching is enabled by default for better performance
    matcher = FuzzyMatcher(vertex_table_names, threshold=0.6, enable_cache=True)

    tables = self.get_tables(schema_name)
    edge_tables = []

    for table_info in tables:
        table_name = table_info["table_name"]
        pk_columns = self.get_primary_keys(table_name, schema_name)
        fk_columns = self.get_foreign_keys(table_name, schema_name)

        # Skip tables without primary keys
        if not pk_columns:
            continue

        # Check if table is edge-like
        if not self._is_edge_like_table(table_name, pk_columns, fk_columns):
            continue

        all_columns = self.get_table_columns(table_name, schema_name)

        # Mark primary key and unique columns and convert to ColumnInfo
        pk_set = set(pk_columns)
        unique_columns = self.get_unique_columns(table_name, schema_name)
        unique_set = set(unique_columns)
        column_infos = []
        for col in all_columns:
            column_infos.append(
                ColumnInfo(
                    name=col["name"],
                    type=col["type"],
                    description=col.get("description", ""),
                    is_nullable=col.get("is_nullable", "YES"),
                    column_default=col.get("column_default"),
                    is_pk=col["name"] in pk_set,
                    is_unique=col["name"] in unique_set,
                    ordinal_position=col.get("ordinal_position"),
                )
            )

        # Convert foreign keys to ForeignKeyInfo
        fk_infos = []
        for fk in fk_columns:
            fk_infos.append(
                ForeignKeyInfo(
                    column=fk["column"],
                    references_table=fk["references_table"],
                    references_column=fk.get("references_column"),
                    constraint_name=fk.get("constraint_name"),
                )
            )

        # Determine source and target tables
        source_table = None
        target_table = None
        source_column = None
        target_column = None
        relation_name = None

        # If we have exactly 2 foreign keys, use them directly
        if len(fk_infos) == 2:
            source_fk = fk_infos[0]
            target_fk = fk_infos[1]
            source_table = source_fk.references_table
            target_table = target_fk.references_table
            source_column = source_fk.column
            target_column = target_fk.column
            # Still try to infer relation from table name
            fk_dicts = [
                {
                    "column": fk.column,
                    "references_table": fk.references_table,
                }
                for fk in fk_infos
            ]
            _, _, relation_name = infer_edge_vertices_from_table_name(
                table_name, pk_columns, fk_dicts, vertex_table_names, matcher
            )
        # If we have 2 or more primary keys, try to infer from table name and structure
        elif len(pk_columns) >= 2:
            # Convert fk_infos to dicts for _infer_edge_vertices_from_table_name
            fk_dicts = [
                {
                    "column": fk.column,
                    "references_table": fk.references_table,
                }
                for fk in fk_infos
            ]

            # Try to infer from table name pattern
            inferred_source, inferred_target, relation_name = (
                infer_edge_vertices_from_table_name(
                    table_name,
                    pk_columns,
                    fk_dicts,
                    vertex_table_names,
                    matcher,
                )
            )

            if inferred_source and inferred_target:
                source_table = inferred_source
                target_table = inferred_target
                # Try to match PK columns to FK columns for source/target columns
                if fk_infos:
                    # Use first FK for source, second for target if available
                    if len(fk_infos) >= 2:
                        source_column = fk_infos[0].column
                        target_column = fk_infos[1].column
                    elif len(fk_infos) == 1:
                        # Self-reference case
                        source_column = fk_infos[0].column
                        target_column = fk_infos[0].column
                else:
                    # Use PK columns as source/target columns
                    source_column = pk_columns[0]
                    target_column = (
                        pk_columns[1] if len(pk_columns) > 1 else pk_columns[0]
                    )
            elif fk_infos:
                # Fallback: use FK references if available
                if len(fk_infos) >= 2:
                    source_table = fk_infos[0].references_table
                    target_table = fk_infos[1].references_table
                    source_column = fk_infos[0].column
                    target_column = fk_infos[1].column
                elif len(fk_infos) == 1:
                    source_table = fk_infos[0].references_table
                    target_table = fk_infos[0].references_table
                    source_column = fk_infos[0].column
                    target_column = fk_infos[0].column
            else:
                # Last resort: use PK columns and infer table names from column names
                source_column = pk_columns[0]
                target_column = (
                    pk_columns[1] if len(pk_columns) > 1 else pk_columns[0]
                )
                # Use robust inference logic to extract vertex names from column names
                source_table = infer_vertex_from_column_name(
                    source_column, vertex_table_names, matcher
                )
                target_table = infer_vertex_from_column_name(
                    target_column, vertex_table_names, matcher
                )

        # Only add if we have source and target information
        if source_table and target_table:
            edge_tables.append(
                EdgeTableInfo(
                    name=table_name,
                    schema_name=schema_name,
                    columns=column_infos,
                    primary_key=pk_columns,
                    foreign_keys=fk_infos,
                    source_table=source_table,
                    target_table=target_table,
                    source_column=source_column or pk_columns[0],
                    target_column=target_column
                    or (pk_columns[1] if len(pk_columns) > 1 else pk_columns[0]),
                    relation=relation_name,
                )
            )
        else:
            logger.warning(
                f"Could not determine source/target tables for edge-like table '{table_name}'. "
                f"Skipping."
            )

    return edge_tables

detect_vertex_tables(schema_name=None)

Detect vertex-like tables in the schema.

Heuristic: Tables with a primary key and descriptive columns (not just foreign keys). These represent entities.

Note: Tables identified as edge-like are excluded from vertex tables.

Parameters:

Name	Type	Description	Default
schema_name	str | None	Schema name. If None, uses 'public' or config schema_name.	None

Returns:

Type	Description
list[VertexTableInfo]	List of vertex table information dictionaries

Source code in graflo/db/postgres/conn.py

def detect_vertex_tables(
    self, schema_name: str | None = None
) -> list[VertexTableInfo]:
    """Detect vertex-like tables in the schema.

    Heuristic: Tables with a primary key and descriptive columns
    (not just foreign keys). These represent entities.

    Note: Tables identified as edge-like are excluded from vertex tables.

    Args:
        schema_name: Schema name. If None, uses 'public' or config schema_name.

    Returns:
        List of vertex table information dictionaries
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    tables = self.get_tables(schema_name)
    vertex_tables = []

    for table_info in tables:
        table_name = table_info["table_name"]
        pk_columns = self.get_primary_keys(table_name, schema_name)
        fk_columns = self.get_foreign_keys(table_name, schema_name)
        all_columns = self.get_table_columns(table_name, schema_name)

        # Vertex-like tables have:
        # 1. A primary key
        # 2. Not identified as edge-like tables
        # 3. Descriptive columns beyond just foreign keys

        if not pk_columns:
            continue  # Skip tables without primary keys

        # Skip edge-like tables
        if self._is_edge_like_table(table_name, pk_columns, fk_columns):
            continue

        # Count non-FK, non-PK columns (descriptive columns)
        fk_column_names = {fk["column"] for fk in fk_columns}
        pk_column_names = set(pk_columns)
        descriptive_columns = [
            col
            for col in all_columns
            if col["name"] not in fk_column_names
            and col["name"] not in pk_column_names
        ]

        # If table has descriptive columns, consider it vertex-like
        if descriptive_columns:
            # Mark primary key and unique columns and convert to ColumnInfo
            pk_set = set(pk_columns)
            unique_columns = self.get_unique_columns(table_name, schema_name)
            unique_set = set(unique_columns)
            column_infos = []
            for col in all_columns:
                column_infos.append(
                    ColumnInfo(
                        name=col["name"],
                        type=col["type"],
                        description=col.get("description", ""),
                        is_nullable=col.get("is_nullable", "YES"),
                        column_default=col.get("column_default"),
                        is_pk=col["name"] in pk_set,
                        is_unique=col["name"] in unique_set,
                        ordinal_position=col.get("ordinal_position"),
                    )
                )

            # Convert foreign keys to ForeignKeyInfo
            fk_infos = []
            for fk in fk_columns:
                fk_infos.append(
                    ForeignKeyInfo(
                        column=fk["column"],
                        references_table=fk["references_table"],
                        references_column=fk.get("references_column"),
                        constraint_name=fk.get("constraint_name"),
                    )
                )

            vertex_tables.append(
                VertexTableInfo(
                    name=table_name,
                    schema_name=schema_name,
                    columns=column_infos,
                    primary_key=pk_columns,
                    foreign_keys=fk_infos,
                )
            )

    return vertex_tables

get_foreign_keys(table_name, schema_name=None)

Get foreign key relationships for a table.

Tries information_schema first, falls back to pg_catalog if needed.

Parameters:

Name	Type	Description	Default
table_name	str	Name of the table	required
schema_name	str | None	Schema name. If None, uses 'public' or config schema_name.	None

Returns:

Type	Description
list[dict[str, Any]]	List of foreign key dictionaries with keys:
list[dict[str, Any]]	column, references_table, references_column, constraint_name

Source code in graflo/db/postgres/conn.py

def get_foreign_keys(
    self, table_name: str, schema_name: str | None = None
) -> list[dict[str, Any]]:
    """Get foreign key relationships for a table.

    Tries information_schema first, falls back to pg_catalog if needed.

    Args:
        table_name: Name of the table
        schema_name: Schema name. If None, uses 'public' or config schema_name.

    Returns:
        List of foreign key dictionaries with keys:
        column, references_table, references_column, constraint_name
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    # Try information_schema first
    try:
        query = """
            SELECT
                kcu.column_name as column,
                ccu.table_name as references_table,
                ccu.column_name as references_column,
                tc.constraint_name
            FROM information_schema.table_constraints tc
            JOIN information_schema.key_column_usage kcu
                ON tc.constraint_name = kcu.constraint_name
                AND tc.table_schema = kcu.table_schema
            JOIN information_schema.constraint_column_usage ccu
                ON ccu.constraint_name = tc.constraint_name
                AND ccu.table_schema = tc.table_schema
            WHERE tc.constraint_type = 'FOREIGN KEY'
              AND tc.table_schema = %s
              AND tc.table_name = %s
            ORDER BY kcu.ordinal_position;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name, table_name))
            results = [dict(row) for row in cursor.fetchall()]
            # If we got results and information_schema is reliable, return them
            if self._check_information_schema_reliable(schema_name):
                return results
            # Otherwise fall back to pg_catalog
            logger.debug(
                f"information_schema returned no results or is unreliable, "
                f"falling back to pg_catalog for foreign keys of '{schema_name}.{table_name}'"
            )
    except Exception as e:
        logger.debug(
            f"information_schema query failed: {e}, falling back to pg_catalog"
        )

    # Fallback to pg_catalog
    return self._get_foreign_keys_pg_catalog(table_name, schema_name)

get_primary_keys(table_name, schema_name=None)

Get primary key columns for a table.

Tries information_schema first, falls back to pg_catalog if needed.

Parameters:

Name	Type	Description	Default
table_name	str	Name of the table	required
schema_name	str | None	Schema name. If None, uses 'public' or config schema_name.	None

Returns:

Type	Description
list[str]	List of primary key column names

Source code in graflo/db/postgres/conn.py

def get_primary_keys(
    self, table_name: str, schema_name: str | None = None
) -> list[str]:
    """Get primary key columns for a table.

    Tries information_schema first, falls back to pg_catalog if needed.

    Args:
        table_name: Name of the table
        schema_name: Schema name. If None, uses 'public' or config schema_name.

    Returns:
        List of primary key column names
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    # Try information_schema first
    try:
        query = """
            SELECT kcu.column_name
            FROM information_schema.table_constraints tc
            JOIN information_schema.key_column_usage kcu
                ON tc.constraint_name = kcu.constraint_name
                AND tc.table_schema = kcu.table_schema
            WHERE tc.constraint_type = 'PRIMARY KEY'
              AND tc.table_schema = %s
              AND tc.table_name = %s
            ORDER BY kcu.ordinal_position;
        """

        with self.conn.cursor() as cursor:
            cursor.execute(query, (schema_name, table_name))
            results = [row[0] for row in cursor.fetchall()]
            # If we got results and information_schema is reliable, return them
            if results and self._check_information_schema_reliable(schema_name):
                return results
            # Otherwise fall back to pg_catalog
            logger.debug(
                f"information_schema returned no results or is unreliable, "
                f"falling back to pg_catalog for primary keys of '{schema_name}.{table_name}'"
            )
    except Exception as e:
        logger.debug(
            f"information_schema query failed: {e}, falling back to pg_catalog"
        )

    # Fallback to pg_catalog
    return self._get_primary_keys_pg_catalog(table_name, schema_name)

get_table_columns(table_name, schema_name=None)

Get columns for a specific table with types and descriptions.

Tries information_schema first, falls back to pg_catalog if needed.

Parameters:

Name	Type	Description	Default
table_name	str	Name of the table	required
schema_name	str | None	Schema name. If None, uses 'public' or config schema_name.	None

Returns:

Type	Description
list[dict[str, Any]]	List of column information dictionaries with keys:
list[dict[str, Any]]	name, type, description, is_nullable, column_default

Source code in graflo/db/postgres/conn.py

def get_table_columns(
    self, table_name: str, schema_name: str | None = None
) -> list[dict[str, Any]]:
    """Get columns for a specific table with types and descriptions.

    Tries information_schema first, falls back to pg_catalog if needed.

    Args:
        table_name: Name of the table
        schema_name: Schema name. If None, uses 'public' or config schema_name.

    Returns:
        List of column information dictionaries with keys:
        name, type, description, is_nullable, column_default
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    # Try information_schema first
    try:
        query = """
            SELECT
                c.column_name as name,
                c.data_type as type,
                c.udt_name as udt_name,
                c.character_maximum_length,
                c.is_nullable,
                c.column_default,
                COALESCE(d.description, '') as description,
                c.ordinal_position as ordinal_position
            FROM information_schema.columns c
            LEFT JOIN pg_catalog.pg_statio_all_tables st
                ON st.schemaname = c.table_schema
                AND st.relname = c.table_name
            LEFT JOIN pg_catalog.pg_description d
                ON d.objoid = st.relid
                AND d.objsubid = c.ordinal_position
            WHERE c.table_schema = %s
              AND c.table_name = %s
            ORDER BY c.ordinal_position;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name, table_name))
            columns = []
            for row in cursor.fetchall():
                col_dict = dict(row)
                # Format type with length if applicable
                if col_dict["character_maximum_length"]:
                    col_dict["type"] = (
                        f"{col_dict['type']}({col_dict['character_maximum_length']})"
                    )
                # Use udt_name if it's more specific (e.g., varchar, int4)
                if (
                    col_dict["udt_name"]
                    and col_dict["udt_name"] != col_dict["type"]
                ):
                    col_dict["type"] = col_dict["udt_name"]
                # Remove helper fields
                col_dict.pop("character_maximum_length", None)
                col_dict.pop("udt_name", None)
                columns.append(col_dict)

            # If we got results and information_schema is reliable, return them
            if columns and self._check_information_schema_reliable(schema_name):
                return columns
            # Otherwise fall back to pg_catalog
            logger.debug(
                f"information_schema returned no results or is unreliable, "
                f"falling back to pg_catalog for table '{schema_name}.{table_name}'"
            )
    except Exception as e:
        logger.debug(
            f"information_schema query failed: {e}, falling back to pg_catalog"
        )

    # Fallback to pg_catalog
    return self._get_table_columns_pg_catalog(table_name, schema_name)

get_table_row_count_estimate(table_name, schema_name=None)

Return approximate row count from pg_class.reltuples (updated by ANALYZE).

Avoids full table scan; may be stale until next ANALYZE.

Source code in graflo/db/postgres/conn.py

def get_table_row_count_estimate(
    self, table_name: str, schema_name: str | None = None
) -> int | None:
    """Return approximate row count from pg_class.reltuples (updated by ANALYZE).

    Avoids full table scan; may be stale until next ANALYZE.
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"
    query = """
        SELECT c.reltuples::bigint AS estimate
        FROM pg_catalog.pg_class c
        JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
        WHERE n.nspname = %s AND c.relname = %s AND c.relkind = 'r';
    """
    with self.conn.cursor() as cursor:
        cursor.execute(query, (schema_name, table_name))
        row = cursor.fetchone()
        if row is None:
            return None
        val = row[0]
        return int(val) if val is not None else None

get_table_sample_rows(table_name, schema_name=None, limit=5)

Return first limit rows from the table (no ORDER BY for speed).

Source code in graflo/db/postgres/conn.py

def get_table_sample_rows(
    self,
    table_name: str,
    schema_name: str | None = None,
    limit: int = 5,
) -> list[dict[str, Any]]:
    """Return first `limit` rows from the table (no ORDER BY for speed)."""
    if schema_name is None:
        schema_name = self.config.schema_name or "public"
    query = sql.SQL("SELECT * FROM {}.{} LIMIT %s").format(
        sql.Identifier(schema_name),
        sql.Identifier(table_name),
    )
    try:
        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (limit,))
            return [dict(row) for row in cursor.fetchall()]
    except Exception as e:
        logger.debug(
            f"Could not fetch sample rows for '{schema_name}.{table_name}': {e}"
        )
        return []

get_tables(schema_name=None)

Get all tables in the specified schema.

Tries information_schema first, falls back to pg_catalog if needed.

Parameters:

Name	Type	Description	Default
schema_name	str | None	Schema name to query. If None, uses 'public' or config schema_name.	None

Returns:

Type	Description
list[dict[str, Any]]	List of table information dictionaries with keys: table_name, table_schema

Source code in graflo/db/postgres/conn.py

def get_tables(self, schema_name: str | None = None) -> list[dict[str, Any]]:
    """Get all tables in the specified schema.

    Tries information_schema first, falls back to pg_catalog if needed.

    Args:
        schema_name: Schema name to query. If None, uses 'public' or config schema_name.

    Returns:
        List of table information dictionaries with keys: table_name, table_schema
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    # Try information_schema first
    try:
        query = """
            SELECT table_name, table_schema
            FROM information_schema.tables
            WHERE table_schema = %s
              AND table_type = 'BASE TABLE'
            ORDER BY table_name;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name,))
            results = [dict(row) for row in cursor.fetchall()]
            # If we got results, check if information_schema is reliable
            if results and self._check_information_schema_reliable(schema_name):
                return results
            # If no results or unreliable, fall back to pg_catalog
            logger.debug(
                f"information_schema returned no results or is unreliable, "
                f"falling back to pg_catalog for schema '{schema_name}'"
            )
    except Exception as e:
        logger.debug(
            f"information_schema query failed: {e}, falling back to pg_catalog"
        )

    # Fallback to pg_catalog
    return self._get_tables_pg_catalog(schema_name)

get_unique_columns(table_name, schema_name=None)

Get column names that participate in any UNIQUE constraint.

Tries information_schema first, falls back to pg_catalog if needed.

Source code in graflo/db/postgres/conn.py

def get_unique_columns(
    self, table_name: str, schema_name: str | None = None
) -> list[str]:
    """Get column names that participate in any UNIQUE constraint.

    Tries information_schema first, falls back to pg_catalog if needed.
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    try:
        query = """
            SELECT kcu.column_name
            FROM information_schema.table_constraints tc
            JOIN information_schema.key_column_usage kcu
                ON tc.constraint_name = kcu.constraint_name
                AND tc.table_schema = kcu.table_schema
            WHERE tc.constraint_type = 'UNIQUE'
              AND tc.table_schema = %s
              AND tc.table_name = %s
            ORDER BY kcu.ordinal_position;
        """
        with self.conn.cursor() as cursor:
            cursor.execute(query, (schema_name, table_name))
            results = list(dict.fromkeys(row[0] for row in cursor.fetchall()))
            if results and self._check_information_schema_reliable(schema_name):
                return results
    except Exception as e:
        logger.debug(
            f"information_schema query failed for unique columns: {e}, "
            "falling back to pg_catalog"
        )
    return self._get_unique_columns_pg_catalog(table_name, schema_name)

introspect_schema(schema_name=None)

Introspect the database schema and return structured information.

This is the main method that analyzes the schema and returns information about vertex-like and edge-like tables.

Parameters:

Name	Type	Description	Default
schema_name	str | None	Schema name. If None, uses 'public' or config schema_name.	None

Returns:

Type	Description
SchemaIntrospectionResult	SchemaIntrospectionResult with vertex_tables, edge_tables, and schema_name

Source code in graflo/db/postgres/conn.py

def introspect_schema(
    self, schema_name: str | None = None
) -> SchemaIntrospectionResult:
    """Introspect the database schema and return structured information.

    This is the main method that analyzes the schema and returns information
    about vertex-like and edge-like tables.

    Args:
        schema_name: Schema name. If None, uses 'public' or config schema_name.

    Returns:
        SchemaIntrospectionResult with vertex_tables, edge_tables, and schema_name
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    logger.info(f"Introspecting PostgreSQL schema '{schema_name}'")

    vertex_tables = self.detect_vertex_tables(schema_name)
    edge_tables = self.detect_edge_tables(schema_name)
    raw_tables = self._build_raw_tables(schema_name)

    result = SchemaIntrospectionResult(
        vertex_tables=vertex_tables,
        edge_tables=edge_tables,
        raw_tables=raw_tables,
        schema_name=schema_name,
    )

    logger.info(
        f"Found {len(vertex_tables)} vertex-like tables and {len(edge_tables)} edge-like tables"
    )

    return result

read(query, params=None)

Execute a SELECT query and return results as a list of dictionaries.

Parameters:

Name	Type	Description	Default
query	str	SQL SELECT query to execute	required
params	tuple | None	Optional tuple of parameters for parameterized queries	None

Returns:

Type	Description
list[dict[str, Any]]	List of dictionaries, where each dictionary represents a row with column names as keys.
list[dict[str, Any]]	Decimal values are converted to float for compatibility with graph databases.

Source code in graflo/db/postgres/conn.py

def read(self, query: str, params: tuple | None = None) -> list[dict[str, Any]]:
    """Execute a SELECT query and return results as a list of dictionaries.

    Args:
        query: SQL SELECT query to execute
        params: Optional tuple of parameters for parameterized queries

    Returns:
        List of dictionaries, where each dictionary represents a row with column names as keys.
        Decimal values are converted to float for compatibility with graph databases.
    """
    from decimal import Decimal

    with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
        if params:
            cursor.execute(query, params)
        else:
            cursor.execute(query)

        # Convert rows to dictionaries and convert Decimal to float
        results = []
        for row in cursor.fetchall():
            row_dict = dict(row)
            # Convert Decimal to float for JSON/graph database compatibility
            for key, value in row_dict.items():
                if isinstance(value, Decimal):
                    row_dict[key] = float(value)
            results.append(row_dict)

        return results

TigerGraphConnection

Bases: Connection

TigerGraph database connection implementation.

Key conceptual differences from ArangoDB: 1. TigerGraph uses GSQL (Graph Query Language) instead of AQL 2. Schema must be defined explicitly before data insertion 3. No automatic vertex/edge class creation - vertices and edges must be pre-defined 4. Different query syntax and execution model 5. Token-based authentication recommended for TigerGraph 4+

Authentication (recommended for TG 4+): For best results, provide BOTH username/password AND secret: - username/password: Required for initial connection and GSQL operations - secret: Generates token that works for both GSQL and REST API operations

Token-based authentication using secrets is the most robust and recommended
approach for TigerGraph 4+. The connection will:
1. Use username/password for initial connection
2. Generate a token from the secret
3. Use the token for both GSQL operations (via REST API) and REST API calls

Example:
    >>> config = TigergraphConfig(
    ...     uri="http://localhost:14240",
    ...     username="tigergraph",      # Required for initial connection
    ...     password="tigergraph",      # Required for initial connection
    ...     secret="your_secret_here",  # Generates token for GSQL + REST API
    ...     database="my_graph"
    ... )
    >>> conn = TigerGraphConnection(config)

Port Configuration for TigerGraph 4+: TigerGraph 4.1+ uses port 14240 (GSQL server) as the primary interface. Port 9000 (REST++) is for internal use only in TG 4.1+.

Standard ports:
- Port 14240: GSQL server (primary interface for all API requests)
- Port 9000: REST++ (internal-only in TG 4.1+)

For custom Docker deployments with port mapping, ports are configured via
environment variables (e.g., TG_WEB, TG_REST) and loaded automatically
when using TigergraphConfig.from_docker_env().

Version Compatibility

• All TigerGraph versions use /restpp prefix for REST++ endpoints
• Version is auto-detected, or can be manually specified in config

Source code in graflo/db/tigergraph/conn.py

class TigerGraphConnection(Connection):
    """
    TigerGraph database connection implementation.

    Key conceptual differences from ArangoDB:
    1. TigerGraph uses GSQL (Graph Query Language) instead of AQL
    2. Schema must be defined explicitly before data insertion
    3. No automatic vertex/edge class creation - vertices and edges must be pre-defined
    4. Different query syntax and execution model
    5. Token-based authentication recommended for TigerGraph 4+

    Authentication (recommended for TG 4+):
        For best results, provide BOTH username/password AND secret:
        - username/password: Required for initial connection and GSQL operations
        - secret: Generates token that works for both GSQL and REST API operations

        Token-based authentication using secrets is the most robust and recommended
        approach for TigerGraph 4+. The connection will:
        1. Use username/password for initial connection
        2. Generate a token from the secret
        3. Use the token for both GSQL operations (via REST API) and REST API calls

        Example:
            >>> config = TigergraphConfig(
            ...     uri="http://localhost:14240",
            ...     username="tigergraph",      # Required for initial connection
            ...     password="tigergraph",      # Required for initial connection
            ...     secret="your_secret_here",  # Generates token for GSQL + REST API
            ...     database="my_graph"
            ... )
            >>> conn = TigerGraphConnection(config)

    Port Configuration for TigerGraph 4+:
        TigerGraph 4.1+ uses port 14240 (GSQL server) as the primary interface.
        Port 9000 (REST++) is for internal use only in TG 4.1+.

        Standard ports:
        - Port 14240: GSQL server (primary interface for all API requests)
        - Port 9000: REST++ (internal-only in TG 4.1+)

        For custom Docker deployments with port mapping, ports are configured via
        environment variables (e.g., TG_WEB, TG_REST) and loaded automatically
        when using TigergraphConfig.from_docker_env().

    Version Compatibility:
        - All TigerGraph versions use /restpp prefix for REST++ endpoints
        - Version is auto-detected, or can be manually specified in config
    """

    flavor = DBType.TIGERGRAPH

    def __init__(self, config: TigergraphConfig):
        super().__init__()
        self.config = config
        self.ssl_verify = getattr(config, "ssl_verify", True)

        # Store connection configuration (no longer using pyTigerGraph)
        # For TigerGraph 4+, both ports typically route through the GSQL server
        # Port 9000 (REST++) is internal-only in TG 4.1+
        self.graphname: str = (
            config.database if config.database is not None else "DefaultGraph"
        )

        # Initialize URLs (ports come from config, no hardcoded defaults)
        # Set GSQL URL first as it's needed for token generation
        # For TigerGraph 4+, gs_port is the primary port (extracted from URI if not explicitly set)
        # Fall back to port from URI if gs_port is not set
        gs_port: int | str | None = config.gs_port
        if gs_port is None:
            # Try to get port from URI
            uri_port = config.port
            if uri_port:
                try:
                    gs_port = int(uri_port)
                    logger.debug(f"Using port {gs_port} from URI for GSQL endpoint")
                except (ValueError, TypeError):
                    pass

        if gs_port is None:
            raise ValueError(
                "gs_port or URI with port must be set in TigergraphConfig. "
                "Standard ports: 14240 (GSQL), 9000 (REST++)."
            )
        self.gsql_url = f"{config.url_without_port}:{gs_port}"

        # Detect TigerGraph version for compatibility (needed before token generation)
        self.tg_version: str | None = None
        self._use_restpp_prefix = False  # Default for 4.2.2+

        # Check if version is manually configured first
        if hasattr(config, "version") and config.version:
            version_str = config.version
            logger.info(f"Using manually configured TigerGraph version: {version_str}")
        else:
            # Auto-detect version using REST API
            try:
                version_str = self._get_version()
            except Exception as e:
                logger.warning(
                    f"Failed to detect TigerGraph version: {e}. "
                    f"Defaulting to 4.2.2+ behavior (no /restpp prefix)"
                )
                version_str = None

        # Parse version string if we have one
        if version_str:
            # Extract version from strings like "release_4.2.2_09-29-2025" or "4.2.1" or "v4.2.1"
            import re

            version_match = re.search(r"(\d+)\.(\d+)\.(\d+)", version_str)
            if version_match:
                major = int(version_match.group(1))
                minor = int(version_match.group(2))
                patch = int(version_match.group(3))
                self.tg_version = f"{major}.{minor}.{patch}"

                # All TigerGraph versions use /restpp prefix for REST++ endpoints
                # Even 4.2.2+ requires /restpp prefix (despite some documentation suggesting otherwise)
                self._use_restpp_prefix = True
                logger.info(
                    f"TigerGraph version {self.tg_version} detected, "
                    f"using /restpp prefix for REST API"
                )
            else:
                logger.warning(
                    f"Could not extract version number from '{version_str}'. "
                    f"Defaulting to using /restpp prefix for REST API"
                )
                self._use_restpp_prefix = True

        # Store base URLs for REST++ and GSQL endpoints
        # For TigerGraph 4.1+, REST++ endpoints use the GSQL port with /restpp prefix
        # Port 9000 is internal-only in TG 4.1+, so we use the same port as GSQL
        # Use the GSQL port we already determined to ensure consistency
        base_url = f"{config.url_without_port}:{gs_port}"
        # Always use /restpp prefix for REST++ endpoints (required for all TG versions)
        self.restpp_url = f"{base_url}/restpp"

        # Get authentication token if secret is provided
        # Token-based auth is the recommended approach for TigerGraph 4+
        # IMPORTANT: You should provide BOTH username/password AND secret:
        # - username/password: Used for initial connection and GSQL operations
        # - secret: Generates token that works for both GSQL and REST API operations
        # Use graph-specific token (is_global=False) for better security
        self.api_token: str | None = None
        if config.secret:
            try:
                token, expiration = self._get_token_from_secret(
                    config.secret,
                    self.graphname,  # Pass graph name for graph-specific token
                )
                self.api_token = token
                if expiration:
                    logger.info(
                        f"Successfully obtained API token for graph '{self.graphname}' "
                        f"(expires: {expiration})"
                    )
                else:
                    logger.info(
                        f"Successfully obtained API token for graph '{self.graphname}'"
                    )
            except Exception as e:
                # Log and fall back to username/password authentication
                logger.warning(f"Failed to get authentication token: {e}")
                logger.warning("Falling back to username/password authentication")
                logger.warning(
                    "Note: For best results, provide both username/password AND secret. "
                    "Username/password is used for GSQL operations, secret generates token for REST API."
                )

    def _get_auth_headers(self, use_basic_auth: bool = False) -> dict[str, str]:
        """Get authentication headers for REST API calls.

        Args:
            use_basic_auth: If True, always use Basic Auth (required for GSQL endpoints).
                           If False, prioritize token-based auth for REST++ endpoints.

        Prioritizes token-based authentication over Basic Auth for REST++ endpoints:
        1. If API token is available (from secret), use Bearer token (recommended for TG 4+)
        2. Otherwise, fall back to HTTP Basic Auth with username/password

        For GSQL endpoints, always use Basic Auth as they don't support Bearer tokens.

        Returns:
            Dictionary with Authorization header
        """
        headers = {}

        # GSQL endpoints require Basic Auth, not Bearer tokens
        if use_basic_auth or not self.api_token:
            # Use default username "tigergraph" if username is None but password is set
            username = self.config.username if self.config.username else "tigergraph"
            password = self.config.password

            if password:
                import base64

                credentials = f"{username}:{password}"
                encoded_credentials = base64.b64encode(credentials.encode()).decode()
                headers["Authorization"] = f"Basic {encoded_credentials}"
            else:
                logger.warning(
                    f"No password configured for Basic Auth. "
                    f"Username: {username}, Password: {password}"
                )
        else:
            # Use Bearer token for REST++ endpoints
            headers["Authorization"] = f"Bearer {self.api_token}"

        return headers

    def _get_token_from_secret(
        self, secret: str, graph_name: str | None = None, lifetime: int = 3600 * 24 * 30
    ) -> tuple[str, str | None]:
        """
        Generate authentication token from secret using TigerGraph REST API.

        Implements robust token generation with fallback logic for different TG 4.x versions:
        - TigerGraph 4.2.2+: POST /gsql/v1/tokens (lifetime in milliseconds)
        - TigerGraph 4.0-4.2.1: POST /gsql/v1/auth/token (lifetime in seconds)

        Based on pyTigerGraph's token generation mechanism with version-specific endpoint handling.

        Args:
            secret: Secret string created via CREATE SECRET in GSQL
            graph_name: Name of the graph (None for global token)
            lifetime: Token lifetime in seconds (default: 30 days)

        Returns:
            Tuple of (token, expiration_timestamp) or (token, None) if expiration not provided

        Raises:
            RuntimeError: If token generation fails after all retry attempts
        """
        auth_headers = self._get_auth_headers(use_basic_auth=True)
        headers = {
            "Content-Type": "application/json",
            **auth_headers,
        }

        # Determine which endpoint to try based on version
        # For TG 4.2.2+, use /gsql/v1/tokens (lifetime in milliseconds)
        # For TG 4.0-4.2.1, use /gsql/v1/auth/token (lifetime in seconds)
        use_new_endpoint = False
        if self.tg_version:
            import re

            version_match = re.search(r"(\d+)\.(\d+)\.(\d+)", self.tg_version)
            if version_match:
                major = int(version_match.group(1))
                minor = int(version_match.group(2))
                patch = int(version_match.group(3))
                # Use new endpoint for 4.2.2+
                use_new_endpoint = (major, minor, patch) >= (4, 2, 2)

        # Try endpoints in order: new endpoint first (if version >= 4.2.2), then fallback
        endpoints_to_try = []
        if use_new_endpoint:
            # Try new endpoint first for 4.2.2+
            endpoints_to_try.append(
                (
                    f"{self.gsql_url}/gsql/v1/tokens",
                    {
                        "secret": secret,
                        "graph": graph_name,
                        "lifetime": lifetime * 1000,  # Convert to milliseconds
                    },
                    True,  # lifetime in milliseconds
                )
            )
            # Fallback to old endpoint if new one fails
            endpoints_to_try.append(
                (
                    f"{self.gsql_url}/gsql/v1/auth/token",
                    {
                        "secret": secret,
                        "graph": graph_name,
                        "lifetime": lifetime,  # In seconds
                    },
                    False,  # lifetime in seconds
                )
            )
        else:
            # For older versions or unknown version, try old endpoint first
            endpoints_to_try.append(
                (
                    f"{self.gsql_url}/gsql/v1/auth/token",
                    {
                        "secret": secret,
                        "graph": graph_name,
                        "lifetime": lifetime,  # In seconds
                    },
                    False,  # lifetime in seconds
                )
            )
            # Fallback to new endpoint (in case version detection was wrong)
            endpoints_to_try.append(
                (
                    f"{self.gsql_url}/gsql/v1/tokens",
                    {
                        "secret": secret,
                        "graph": graph_name,
                        "lifetime": lifetime * 1000,  # Convert to milliseconds
                    },
                    True,  # lifetime in milliseconds
                )
            )

        last_error: Exception | None = None
        all_404_errors = True  # Track if all failures were 404 errors

        for url, payload, _is_milliseconds in endpoints_to_try:
            try:
                # Remove None values from payload
                clean_payload = {k: v for k, v in payload.items() if v is not None}

                response = requests.post(
                    url,
                    headers=headers,
                    json=clean_payload,  # Use json parameter instead of data
                    timeout=30,
                    verify=self.ssl_verify,
                )

                # Check for 404 - might indicate wrong endpoint or port issue
                if response.status_code == 404:
                    # Try port fallback (similar to pyTigerGraph's _req method)
                    # If using wrong port, try GSQL port
                    if (
                        "/gsql" in url
                        and self.config.port is not None
                        and self.config.gs_port is not None
                        and self.config.port != self.config.gs_port
                    ):
                        logger.debug(f"404 on {url}, trying GSQL port fallback...")
                        # Replace port in URL with GSQL port
                        fallback_url = url.replace(
                            f":{self.config.port}", f":{self.config.gs_port}"
                        )
                        try:
                            response = requests.post(
                                fallback_url,
                                headers=headers,
                                json=clean_payload,
                                timeout=30,
                                verify=self.ssl_verify,
                            )
                            if response.status_code == 200:
                                url = fallback_url  # Update URL for logging
                        except Exception:
                            pass  # Continue to next endpoint

                response.raise_for_status()
                result = response.json()

                # Parse response (both endpoints return similar format)
                # Format: {"token": "...", "expiration": "...", "error": false, "message": "..."}
                # or {"token": "..."} for older versions
                if result.get("error") is True:
                    error_msg = result.get("message", "Unknown error")
                    raise RuntimeError(f"Token generation failed: {error_msg}")

                token = result.get("token")
                expiration = result.get("expiration")

                if token:
                    logger.debug(
                        f"Successfully obtained token from {url} "
                        f"(expiration: {expiration or 'not provided'})"
                    )
                    return (token, expiration)
                else:
                    raise ValueError(f"No token in response: {result}")

            except requests.exceptions.HTTPError as e:
                # Track if this was a 404 error
                if e.response.status_code != 404:
                    all_404_errors = False

                # If 404 and we have more endpoints to try, continue
                if e.response.status_code == 404 and len(endpoints_to_try) > 1:
                    logger.debug(
                        f"Endpoint {url} returned 404, trying next endpoint..."
                    )
                    last_error = e
                    continue
                # For other HTTP errors, log and try next endpoint if available
                logger.debug(
                    f"HTTP error {e.response.status_code} on {url}: {e.response.text}"
                )
                last_error = e
                continue
            except Exception as e:
                all_404_errors = False  # Non-HTTP errors are not 404s
                logger.debug(f"Error trying {url}: {e}")
                last_error = e
                continue

        # All graph-specific endpoints failed
        # If all failures were 404 errors and we have a graph_name, try generating a global token
        # This handles cases where the graph doesn't exist yet (e.g., "DefaultGraph" at init time)
        # For TigerGraph 4.2.1, /gsql/v1/tokens requires the graph to exist, but /gsql/v1/auth/token
        # can generate a global token without a graph parameter
        if all_404_errors and graph_name is not None and last_error:
            logger.debug(
                f"All graph-specific token attempts failed with 404. "
                f"Graph '{graph_name}' may not exist yet. "
                f"Trying to generate a global token (without graph parameter)..."
            )

            # Try generating a global token using /gsql/v1/auth/token (works for TG 4.0-4.2.1)
            global_token_endpoints = [
                (
                    f"{self.gsql_url}/gsql/v1/auth/token",
                    {
                        "secret": secret,
                        "lifetime": lifetime,  # In seconds
                        # No graph parameter = global token
                    },
                    False,  # lifetime in seconds
                )
            ]

            # Also try /gsql/v1/tokens without graph parameter (for TG 4.2.2+)
            global_token_endpoints.append(
                (
                    f"{self.gsql_url}/gsql/v1/tokens",
                    {
                        "secret": secret,
                        "lifetime": lifetime * 1000,  # In milliseconds
                        # No graph parameter = global token
                    },
                    True,  # lifetime in milliseconds
                )
            )

            for url, payload, _is_milliseconds in global_token_endpoints:
                try:
                    clean_payload = {k: v for k, v in payload.items() if v is not None}

                    response = requests.post(
                        url,
                        headers=headers,
                        json=clean_payload,
                        timeout=30,
                        verify=self.ssl_verify,
                    )

                    response.raise_for_status()
                    result = response.json()

                    if result.get("error") is True:
                        error_msg = result.get("message", "Unknown error")
                        logger.debug(f"Global token generation failed: {error_msg}")
                        continue

                    token = result.get("token")
                    expiration = result.get("expiration")

                    if token:
                        logger.info(
                            f"Successfully obtained global token from {url} "
                            f"(graph '{graph_name}' may not exist yet, using global token). "
                            f"Expiration: {expiration or 'not provided'}"
                        )
                        return (token, expiration)

                except Exception as e:
                    logger.debug(f"Error trying global token endpoint {url}: {e}")
                    continue

        # All endpoints failed (including global token fallback)
        error_msg = f"Failed to get token from secret after trying {len(endpoints_to_try)} endpoint(s)"
        if all_404_errors and graph_name:
            error_msg += f" (all returned 404, graph '{graph_name}' may not exist yet)"
        if last_error:
            error_msg += f": {last_error}"
        logger.error(error_msg)
        raise RuntimeError(error_msg)

    def _get_version(self) -> str | None:
        """
        Get TigerGraph version using REST API.

        Tries multiple endpoints in order:
        1. GET /gsql/v1/version (GSQL server, port 14240) - primary for TG 4+
        2. GET /version (REST++ server, port 9000) - fallback for older versions

        Note: The /version endpoint does NOT exist on GSQL port (14240).
        It only exists on REST++ port (9000) for older versions.

        Returns:
            Version string (e.g., "4.2.1") or None if detection fails
        """
        import re

        if self.config.gs_port is None:
            raise ValueError("gs_port must be set in config for version detection")

        # Try GSQL endpoint first (primary for TigerGraph 4+)
        # Note: /gsql/v1/version exists on GSQL port, but /version does NOT
        # Response format: plain text like "GSQL version: 4.2.2\n"
        gsql_url = f"{self.gsql_url}/gsql/v1/version"
        headers = self._get_auth_headers(use_basic_auth=True)

        try:
            response = requests.get(
                gsql_url, headers=headers, timeout=10, verify=self.ssl_verify
            )
            response.raise_for_status()

            if not response.text.strip():
                # Empty response
                logger.debug("GSQL version endpoint returned empty response")
                raise ValueError("Empty response from GSQL version endpoint")

            # GSQL /gsql/v1/version returns plain text, not JSON
            # Format: "GSQL version: 4.2.2\n" or similar
            response_text = response.text.strip()

            # Try to parse version from text response
            # Format: "GSQL version: 4.2.2" or "version: 4.2.2" or "4.2.2"
            version_match = re.search(
                r"version:\s*(\d+)\.(\d+)\.(\d+)", response_text, re.IGNORECASE
            )
            if version_match:
                version_str = f"{version_match.group(1)}.{version_match.group(2)}.{version_match.group(3)}"
                logger.debug(
                    f"Detected TigerGraph version: {version_str} from GSQL endpoint (text format)"
                )
                return version_str

            # Try alternative: just look for version number pattern
            version_match = re.search(r"(\d+)\.(\d+)\.(\d+)", response_text)
            if version_match:
                version_str = f"{version_match.group(1)}.{version_match.group(2)}.{version_match.group(3)}"
                logger.debug(
                    f"Detected TigerGraph version: {version_str} from GSQL endpoint (text format)"
                )
                return version_str

            # If text parsing failed, try JSON as fallback (some versions might return JSON)
            try:
                result = response.json()
                message = result.get("message", "")
                if message:
                    version_match = re.search(r"release_(\d+)\.(\d+)\.(\d+)", message)
                    if version_match:
                        version_str = f"{version_match.group(1)}.{version_match.group(2)}.{version_match.group(3)}"
                        logger.debug(
                            f"Detected TigerGraph version: {version_str} from GSQL endpoint (JSON format)"
                        )
                        return version_str
            except ValueError:
                # Not JSON, that's fine - we already tried text parsing
                pass

        except Exception as e:
            logger.debug(f"Failed to get version from GSQL endpoint: {e}")

        # Fallback: Try REST++ /version endpoint (for older versions or if GSQL endpoint fails)
        # Note: /version only exists on REST++ port (9000), not GSQL port (14240)
        try:
            # Use REST++ port if different from GSQL port
            restpp_port = self.config.port if self.config.port else self.config.gs_port
            if restpp_port is None:
                return None

            restpp_url = f"{self.config.url_without_port}:{restpp_port}/version"
            headers = self._get_auth_headers(use_basic_auth=True)

            response = requests.get(
                restpp_url, headers=headers, timeout=10, verify=self.ssl_verify
            )
            response.raise_for_status()

            # Check content type and response
            if not response.text.strip():
                logger.debug("REST++ version endpoint returned empty response")
                return None

            try:
                result = response.json()
            except ValueError:
                logger.debug(
                    f"REST++ version endpoint returned non-JSON response: "
                    f"status={response.status_code}, text={response.text[:200]}"
                )
                return None

            # Parse version from REST++ response
            message = result.get("message", "")
            if message:
                version_match = re.search(r"release_(\d+)\.(\d+)\.(\d+)", message)
                if version_match:
                    version_str = f"{version_match.group(1)}.{version_match.group(2)}.{version_match.group(3)}"
                    logger.debug(
                        f"Detected TigerGraph version: {version_str} from REST++ endpoint"
                    )
                    return version_str

        except Exception as e:
            logger.debug(f"Failed to get version from REST++ endpoint: {e}")

        return None

    def _execute_gsql(self, gsql_command: str) -> str:
        """
        Execute GSQL command using REST API.

        For TigerGraph 4.0-4.2.1, uses POST /gsql/v1/statements endpoint.

        Note: GSQL endpoints require Basic Auth (username/password), not Bearer tokens.

        Args:
            gsql_command: GSQL command string to execute

        Returns:
            Response string from GSQL execution
        """
        url = f"{self.gsql_url}/gsql/v1/statements"
        auth_headers = self._get_auth_headers(use_basic_auth=True)
        headers = {
            "Content-Type": "text/plain",
            **auth_headers,
        }

        # Debug: Log if Authorization header is missing
        if "Authorization" not in headers:
            logger.error(
                f"No Authorization header generated. "
                f"Username: {self.config.username}, Password: {'***' if self.config.password else None}"
            )

        try:
            response = requests.post(
                url,
                headers=headers,
                data=gsql_command,
                timeout=120,
                verify=self.ssl_verify,
            )
            response.raise_for_status()

            # Try to parse JSON response, fallback to text
            try:
                result = response.json()
                # Extract message or result from JSON response
                if isinstance(result, dict):
                    return result.get("message", str(result))
                return str(result)
            except ValueError:
                # Not JSON, return text
                return response.text
        except requests_exceptions.HTTPError as e:
            error_msg = str(e)
            # Try to extract error message from response
            try:
                error_details = e.response.json() if e.response else {}
                error_msg = error_details.get("message", error_msg)
            except Exception:
                pass
            raise RuntimeError(f"GSQL execution failed: {error_msg}") from e

    def _get_vertex_types(self, graph_name: str | None = None) -> list[str]:
        """
        Get list of vertex types using GSQL.

        Args:
            graph_name: Name of the graph (defaults to self.graphname)

        Returns:
            List of vertex type names
        """
        graph_name = graph_name or self.graphname
        try:
            result = self._execute_gsql(f"USE GRAPH {graph_name}\nSHOW VERTEX *")
            # Parse GSQL output using the proper parser
            if isinstance(result, str):
                return self._parse_show_output(result, "VERTEX")
            return []
        except Exception as e:
            logger.debug(f"Failed to get vertex types via GSQL: {e}")
            return []

    def _parse_show_edge_output_with_vertices(
        self, output: str
    ) -> dict[str, list[tuple[str, str]]]:
        """
        Parse SHOW EDGE * output (compact TigerGraph format).

        Returns:
            dict mapping edge_name -> list of (source_vertex, target_vertex)
        """
        edge_map: dict[str, list[tuple[str, str]]] = defaultdict(list)

        # Match lines like:
        # - DIRECTED EDGE contains(FROM Author, TO ResearchField|FROM ResearchField, TO ResearchField)
        edge_line_pattern = re.compile(
            r"-\s+(?:DIRECTED|UNDIRECTED)\s+EDGE\s+(\w+)\(([^)]+)\)"
        )

        # Match FROM X, TO Y
        from_to_pattern = re.compile(r"FROM\s+(\w+)\s*,\s*TO\s+(\w+)")

        for line in output.splitlines():
            line = line.strip()
            if not line.startswith("-"):
                continue

            edge_match = edge_line_pattern.search(line)
            if not edge_match:
                continue

            edge_name = edge_match.group(1)
            endpoints_blob = edge_match.group(2)

            # Split multiple vertex pairs
            for endpoint in endpoints_blob.split("|"):
                ft_match = from_to_pattern.search(endpoint)
                if ft_match:
                    source, target = ft_match.groups()
                    edge_map[edge_name].append((source, target))

        return dict(edge_map)

    def _get_edge_types(
        self, graph_name: str | None = None
    ) -> dict[str, list[tuple[str, str]]]:
        """
        Get edge types and their (source, target) vertex pairs using GSQL.

        Args:
            graph_name: Name of the graph (defaults to self.graphname)

        Returns:
            Dict mapping edge_type -> list of (source_vertex, target_vertex)
        """
        graph_name = graph_name or self.graphname
        try:
            result = self._execute_gsql(f"USE GRAPH {graph_name}\nSHOW EDGE *")

            if isinstance(result, str):
                return self._parse_show_edge_output_with_vertices(result)

            return {}

        except Exception as e:
            logger.error(f"Failed to get edge types via GSQL: {e}")
            return {}

    def _get_installed_queries(self, graph_name: str | None = None) -> list[str]:
        """
        Get list of installed queries using REST API.

        Uses the /endpoints endpoint with dynamic=true to get all installed query endpoints,
        then extracts query names from the endpoint paths.

        Args:
            graph_name: Name of the graph (defaults to self.graphname)

        Returns:
            List of query names
        """
        graph_name = graph_name or self.graphname
        try:
            # Use REST API endpoint to get dynamic endpoints (installed queries)
            # Format: GET /endpoints?dynamic=true
            endpoint = "/endpoints"
            params = {"dynamic": "true"}
            result = self._call_restpp_api(endpoint, method="GET", params=params)

            # Parse the response to extract query names
            # The response is a dict where keys are endpoint paths like:
            # "POST /query/{graph_name}/{query_name}" or "GET /query/{graph_name}/{query_name}"
            queries = []
            if isinstance(result, dict):
                query_prefix = f"/query/{graph_name}/"
                for endpoint_path in result.keys():
                    # Extract query name from endpoint path
                    # Format: "POST /query/{graph_name}/{query_name}" or "GET /query/{graph_name}/{query_name}"
                    if query_prefix in endpoint_path:
                        # Extract the query name after the graph name
                        # Handle both "POST /query/..." and "/query/..." formats
                        idx = endpoint_path.find(query_prefix)
                        if idx >= 0:
                            query_part = endpoint_path[idx + len(query_prefix) :]
                            # Extract query name (everything up to first space, newline, or end)
                            query_name = query_part.split()[0] if query_part else ""
                            # Remove any trailing slashes or special characters
                            query_name = query_name.rstrip("/").strip()
                            if query_name and query_name not in queries:
                                queries.append(query_name)

            return queries
        except Exception as e:
            logger.debug(f"Failed to get installed queries via REST API: {e}")
            return []

    def _run_installed_query(
        self, query_name: str, graph_name: str | None = None, **kwargs: Any
    ) -> dict[str, Any] | list[dict]:
        """
        Run an installed query using REST API.

        Args:
            query_name: Name of the installed query
            graph_name: Name of the graph (defaults to self.graphname)
            **kwargs: Query parameters

        Returns:
            Query result (dict or list)
        """
        graph_name = graph_name or self.graphname
        endpoint = f"/query/{graph_name}/{query_name}"
        return self._call_restpp_api(endpoint, method="POST", data=kwargs)

    def _upsert_vertex(
        self,
        vertex_type: str,
        vertex_id: str,
        attributes: dict[str, Any],
        graph_name: str | None = None,
    ) -> dict[str, Any] | list[dict]:
        """
        Upsert a single vertex using REST API.

        Args:
            vertex_type: Vertex type name
            vertex_id: Vertex ID
            attributes: Vertex attributes
            graph_name: Name of the graph (defaults to self.graphname)

        Returns:
            Response from API
        """
        graph_name = graph_name or self.graphname
        endpoint = f"/graph/{graph_name}/vertices/{vertex_type}/{quote(str(vertex_id))}"
        return self._call_restpp_api(endpoint, method="POST", data=attributes)

    def _upsert_edge(
        self,
        source_type: str,
        source_id: str,
        edge_type: str,
        target_type: str,
        target_id: str,
        attributes: dict[str, Any] | None = None,
        graph_name: str | None = None,
    ) -> dict[str, Any] | list[dict]:
        """
        Upsert a single edge using REST API.

        Args:
            source_type: Source vertex type
            source_id: Source vertex ID
            edge_type: Edge type name
            target_type: Target vertex type
            target_id: Target vertex ID
            attributes: Edge attributes (optional)
            graph_name: Name of the graph (defaults to self.graphname)

        Returns:
            Response from API
        """
        graph_name = graph_name or self.graphname
        endpoint = (
            f"/graph/{graph_name}/edges/{edge_type}/"
            f"{source_type}/{quote(str(source_id))}/"
            f"{target_type}/{quote(str(target_id))}"
        )
        data = attributes if attributes else {}
        return self._call_restpp_api(endpoint, method="POST", data=data)

    def _get_edges(
        self,
        source_type: str,
        source_id: str,
        edge_type: str | None = None,
        graph_name: str | None = None,
    ) -> list[dict[str, Any]]:
        """
        Get edges from a vertex using REST API.

        Based on pyTigerGraph's getEdges() implementation.
        Uses GET /graph/{graph}/edges/{source_vertex_type}/{source_vertex_id} endpoint.

        Args:
            source_type: Source vertex type
            source_id: Source vertex ID
            edge_type: Edge type to filter by (optional, filtered client-side)
            graph_name: Name of the graph (defaults to self.graphname)

        Returns:
            List of edge dictionaries
        """
        graph_name = graph_name or self.graphname

        # Use the correct endpoint format matching pyTigerGraph's _prep_get_edges:
        # GET /graph/{graph}/edges/{source_type}/{source_id}
        # If edge_type is specified, append it: /graph/{graph}/edges/{source_type}/{source_id}/{edge_type}
        if edge_type:
            endpoint = f"/graph/{graph_name}/edges/{source_type}/{quote(str(source_id))}/{edge_type}"
        else:
            endpoint = (
                f"/graph/{graph_name}/edges/{source_type}/{quote(str(source_id))}"
            )

        result = self._call_restpp_api(endpoint, method="GET")

        # Parse REST++ API response format
        # Response format: {"version": {...}, "error": false, "message": "", "results": [...]}
        if isinstance(result, dict):
            # Check for error first
            if result.get("error") is True:
                error_msg = result.get("message", "Unknown error")
                logger.error(f"Error fetching edges: {error_msg}")
                return []

            # Extract results array
            if "results" in result:
                edges = result["results"]
            else:
                logger.debug(
                    f"Unexpected response format from edges endpoint: {result.keys()}"
                )
                return []
        elif isinstance(result, list):
            edges = result
        else:
            logger.debug(
                f"Unexpected response type from edges endpoint: {type(result)}"
            )
            return []

        # Filter by edge_type if specified (client-side filtering)
        # REST API endpoint doesn't support edge_type filtering directly
        if edge_type and isinstance(edges, list):
            edges = [
                e for e in edges if isinstance(e, dict) and e.get("e_type") == edge_type
            ]

        return edges

    def _get_vertices_by_id(
        self, vertex_type: str, vertex_id: str, graph_name: str | None = None
    ) -> dict[str, dict[str, Any]]:
        """
        Get vertex by ID using REST API.

        Args:
            vertex_type: Vertex type name
            vertex_id: Vertex ID
            graph_name: Name of the graph (defaults to self.graphname)

        Returns:
            Dictionary mapping vertex_id to vertex data
        """
        graph_name = graph_name or self.graphname
        endpoint = f"/graph/{graph_name}/vertices/{vertex_type}/{quote(str(vertex_id))}"
        result = self._call_restpp_api(endpoint, method="GET")
        # Parse response format to match expected format
        # Returns {vertex_id: {"attributes": {...}}}
        if isinstance(result, dict):
            if "results" in result:
                # REST API format
                results = result["results"]
                if results and isinstance(results, list) and len(results) > 0:
                    vertex_data = results[0]
                    return {
                        vertex_id: {"attributes": vertex_data.get("attributes", {})}
                    }
            elif vertex_id in result:
                return {vertex_id: result[vertex_id]}
            else:
                # Try to extract vertex data
                return {vertex_id: {"attributes": result.get("attributes", {})}}
        return {}

    def _get_vertex_count(self, vertex_type: str, graph_name: str | None = None) -> int:
        """
        Get vertex count using REST API.

        Args:
            vertex_type: Vertex type name
            graph_name: Name of the graph (defaults to self.graphname)

        Returns:
            Number of vertices
        """
        graph_name = graph_name or self.graphname
        endpoint = f"/graph/{graph_name}/vertices/{vertex_type}"
        params = {"limit": "1", "count": "true"}
        result = self._call_restpp_api(endpoint, method="GET", params=params)
        # Parse count from response
        if isinstance(result, dict):
            return result.get("count", 0)
        return 0

    def _delete_vertices(
        self, vertex_type: str, where: str | None = None, graph_name: str | None = None
    ) -> dict[str, Any] | list[dict]:
        """
        Delete vertices using REST API.

        Args:
            vertex_type: Vertex type name
            where: WHERE clause for filtering (optional)
            graph_name: Name of the graph (defaults to self.graphname)

        Returns:
            Response from API
        """
        graph_name = graph_name or self.graphname
        endpoint = f"/graph/{graph_name}/vertices/{vertex_type}"
        params = {}
        if where:
            params["filter"] = where
        return self._call_restpp_api(endpoint, method="DELETE", params=params)

    def _call_restpp_api(
        self,
        endpoint: str,
        method: str = "GET",
        data: dict[str, Any] | None = None,
        params: dict[str, str] | None = None,
    ) -> dict[str, Any] | list[dict]:
        """Call TigerGraph REST++ API endpoint.

        Args:
            endpoint: REST++ API endpoint (e.g., "/graph/{graph_name}/vertices/{vertex_type}")
            method: HTTP method (GET, POST, etc.)
            data: Optional data to send in request body (for POST)
            params: Optional query parameters

        Returns:
            Response data (dict or list)
        """
        url = f"{self.restpp_url}{endpoint}"

        headers = {
            "Content-Type": "application/json",
            **self._get_auth_headers(),
        }

        logger.debug(f"REST++ API call: {method} {url}")

        try:
            if method.upper() == "GET":
                response = requests.get(
                    url,
                    headers=headers,
                    params=params,
                    timeout=120,
                    verify=self.ssl_verify,
                )
            elif method.upper() == "POST":
                response = requests.post(
                    url,
                    headers=headers,
                    data=json.dumps(data, default=_json_serializer) if data else None,
                    params=params,
                    timeout=120,
                    verify=self.ssl_verify,
                )
            elif method.upper() == "DELETE":
                response = requests.delete(
                    url,
                    headers=headers,
                    params=params,
                    timeout=120,
                    verify=self.ssl_verify,
                )
            else:
                raise ValueError(f"Unsupported HTTP method: {method}")

            response.raise_for_status()
            return response.json()

        except requests_exceptions.HTTPError as errh:
            # For TigerGraph 4.2.1, if token auth fails with 401/REST-10018, try Basic Auth fallback
            if (
                errh.response.status_code == 401
                and self.api_token
                and self.config.username
                and self.config.password
                and "REST-10018" in str(errh)
            ):
                logger.warning(
                    "Token authentication failed with REST-10018, "
                    "falling back to Basic Auth for TigerGraph 4.2.1 compatibility"
                )
                # Retry with Basic Auth
                import base64

                credentials = f"{self.config.username}:{self.config.password}"
                encoded_credentials = base64.b64encode(credentials.encode()).decode()
                headers["Authorization"] = f"Basic {encoded_credentials}"
                try:
                    if method.upper() == "GET":
                        response = requests.get(
                            url,
                            headers=headers,
                            params=params,
                            timeout=120,
                            verify=self.ssl_verify,
                        )
                    elif method.upper() == "POST":
                        response = requests.post(
                            url,
                            headers=headers,
                            data=json.dumps(data, default=_json_serializer)
                            if data
                            else None,
                            params=params,
                            timeout=120,
                            verify=self.ssl_verify,
                        )
                    elif method.upper() == "DELETE":
                        response = requests.delete(
                            url,
                            headers=headers,
                            params=params,
                            timeout=120,
                            verify=self.ssl_verify,
                        )
                    else:
                        raise ValueError(f"Unsupported HTTP method: {method}")
                    response.raise_for_status()
                    logger.info("Successfully authenticated using Basic Auth fallback")
                    return response.json()
                except requests_exceptions.HTTPError as errh2:
                    logger.error(f"HTTP Error (after Basic Auth fallback): {errh2}")
                    error_response = {"error": True, "message": str(errh2)}
                    try:
                        error_json = response.json()
                        if isinstance(error_json, dict):
                            error_response.update(error_json)
                        else:
                            error_response["details"] = response.text
                    except Exception:
                        error_response["details"] = response.text
                    return error_response

            logger.error(f"HTTP Error: {errh}")
            error_response = {"error": True, "message": str(errh)}
            try:
                # Try to parse error response for more details
                error_json = response.json()
                if isinstance(error_json, dict):
                    error_response.update(error_json)
                else:
                    error_response["details"] = response.text
            except Exception:
                error_response["details"] = response.text
            return error_response
        except requests_exceptions.ConnectionError as errc:
            logger.error(f"Error Connecting: {errc}")
            return {"error": True, "message": str(errc)}
        except requests_exceptions.Timeout as errt:
            logger.error(f"Timeout Error: {errt}")
            return {"error": True, "message": str(errt)}
        except requests_exceptions.RequestException as err:
            logger.error(f"An unexpected error occurred: {err}")
            return {"error": True, "message": str(err)}

    @contextlib.contextmanager
    def _ensure_graph_context(self, graph_name: str | None = None):
        """
        Context manager that ensures graph context for metadata operations.

        Stores graph name for operations that need it.

        Args:
            graph_name: Name of the graph to use. If None, uses self.config.database.

        Yields:
            The graph name that was set.
        """
        graph_name = graph_name or self.config.database
        if not graph_name:
            raise ValueError(
                "Graph name must be provided via graph_name parameter or config.database"
            )

        old_graphname = self.graphname
        self.graphname = graph_name

        try:
            yield graph_name
        finally:
            # Restore original graphname
            self.graphname = old_graphname

    def graph_exists(self, name: str) -> bool:
        """
        Check if a graph with the given name exists.

        Uses the USE GRAPH command and checks the returned message.
        If the graph doesn't exist, USE GRAPH returns an error message like
        "Graph 'name' does not exist."

        Args:
            name: Name of the graph to check

        Returns:
            bool: True if the graph exists, False otherwise
        """
        try:
            result = self._execute_gsql(f"USE GRAPH {name}")
            result_str = str(result).lower()

            # If the graph doesn't exist, USE GRAPH returns an error message
            # Check for common error messages indicating the graph doesn't exist
            error_patterns = [
                "does not exist",
                "doesn't exist",
                "doesn't exist!",
                f"graph '{name.lower()}' does not exist",
            ]

            # If any error pattern is found, the graph doesn't exist
            for pattern in error_patterns:
                if pattern in result_str:
                    return False

            # If no error pattern is found, the graph likely exists
            # (USE GRAPH succeeded or returned success message)
            return True
        except Exception as e:
            logger.debug(f"Error checking if graph '{name}' exists: {e}")
            # If there's an exception, try to parse it
            error_str = str(e).lower()
            if "does not exist" in error_str or "doesn't exist" in error_str:
                return False
            # If exception doesn't indicate "doesn't exist", assume it exists
            # (other errors might indicate connection issues, not missing graph)
            return False

    @_wrap_tg_exception
    def create_database(
        self,
        name: str,
        vertex_names: list[str] | None = None,
        edge_names: list[str] | None = None,
    ):
        """
        Create a TigerGraph database (graph) using GSQL commands.

        This method creates a graph with explicitly attached vertices and edges.
        Example: CREATE GRAPH researchGraph (author, paper, wrote)

        This method uses direct REST API calls to execute GSQL commands
        that create and use the graph. Supported in TigerGraph version 4.2.2+.

        Args:
            name: Name of the graph to create
            vertex_names: Optional list of vertex type names to attach to the graph
            edge_names: Optional list of edge type names to attach to the graph

        Raises:
            RuntimeError: If graph already exists or creation fails
        """
        # Check if graph already exists first
        if self.graph_exists(name):
            raise RuntimeError(f"Graph '{name}' already exists")

        try:
            # Build the list of types to include in CREATE GRAPH
            all_types = []
            if vertex_names:
                all_types.extend(vertex_names)
            if edge_names:
                all_types.extend(edge_names)

            # Format the CREATE GRAPH command with types
            if all_types:
                types_str = ", ".join(all_types)
                gsql_commands = f"CREATE GRAPH {name} ({types_str})\nUSE GRAPH {name}"
            else:
                # Fallback to empty graph if no types provided
                gsql_commands = f"CREATE GRAPH {name}()\nUSE GRAPH {name}"

            # Execute using direct GSQL REST API which handles authentication
            logger.debug(f"Creating graph '{name}' via GSQL: {gsql_commands}")
            try:
                result = self._execute_gsql(gsql_commands)
                logger.info(
                    f"Successfully created graph '{name}' with types {all_types}: {result}"
                )
                # Verify the result doesn't indicate the graph already existed
                result_str = str(result).lower()
                if (
                    "already exists" in result_str
                    or "duplicate" in result_str
                    or "graph already exists" in result_str
                ):
                    raise RuntimeError(f"Graph '{name}' already exists")
                return result
            except RuntimeError:
                # Re-raise RuntimeError as-is (already handled)
                raise
            except Exception as e:
                error_msg = str(e).lower()
                # Check if graph already exists - raise exception in this case
                # TigerGraph may return various error messages for existing graphs
                if (
                    "already exists" in error_msg
                    or "duplicate" in error_msg
                    or "graph already exists" in error_msg
                    or "already exist" in error_msg
                ):
                    logger.warning(f"Graph '{name}' already exists: {e}")
                    raise RuntimeError(f"Graph '{name}' already exists") from e
                logger.error(f"Failed to create graph '{name}': {e}")
                raise

        except RuntimeError:
            # Re-raise RuntimeError as-is
            raise
        except Exception as e:
            logger.error(f"Error creating graph '{name}' via GSQL: {e}")
            raise

    @_wrap_tg_exception
    def delete_database(self, name: str):
        """
        Delete a TigerGraph database (graph).

        This method attempts to drop the graph using a clean teardown sequence:
          1) Drop all queries associated with the graph
          2) Drop the graph itself

        Args:
            name: Name of the graph to delete

        Note:
            In TigerGraph, deleting a graph structure requires the graph to be empty
            or may fail if it has dependencies. This method handles both cases.
        """
        try:
            logger.debug(f"Attempting to drop graph '{name}'")

            # The order matters for a clean teardown
            cleanup_script = f"""
                USE GRAPH {name}
                DROP QUERY *
                USE GLOBAL
                DROP GRAPH {name}
            """
            result = self._execute_gsql(cleanup_script)
            logger.info(f"Successfully dropped graph '{name}': {result}")
            return result
        except Exception as e:
            error_str = str(e).lower()
            # If the clean teardown fails, try fallback approaches
            if (
                "depends on" in error_str
                or "query" in error_str
                or "not exist" in error_str
            ):
                logger.warning(
                    f"Clean teardown failed for graph '{name}': {e}. "
                    f"Attempting fallback cleanup."
                )
                # Fallback: Try to drop queries individually, then drop graph
                try:
                    with self._ensure_graph_context(name):
                        try:
                            queries = self._get_installed_queries()
                            if queries:
                                logger.info(
                                    f"Dropping {len(queries)} queries from graph '{name}'"
                                )
                                for query_name in queries:
                                    try:
                                        drop_query_cmd = f"USE GRAPH {name}\nDROP QUERY {query_name} IF EXISTS"
                                        self._execute_gsql(drop_query_cmd)
                                        logger.debug(
                                            f"Dropped query '{query_name}' from graph '{name}'"
                                        )
                                    except Exception:
                                        # Try without IF EXISTS for older TigerGraph versions
                                        try:
                                            drop_query_cmd = f"USE GRAPH {name}\nDROP QUERY {query_name}"
                                            self._execute_gsql(drop_query_cmd)
                                        except Exception as qe2:
                                            logger.debug(
                                                f"Could not drop query '{query_name}': {qe2}"
                                            )
                        except Exception as e2:
                            logger.debug(
                                f"Could not list queries for graph '{name}': {e2}"
                            )

                    # Now try to drop the graph
                    drop_command = f"USE GLOBAL\nDROP GRAPH {name}"
                    result = self._execute_gsql(drop_command)
                    logger.info(
                        f"Successfully dropped graph '{name}' via fallback: {result}"
                    )
                    return result
                except Exception as fallback_error:
                    logger.warning(
                        f"Fallback cleanup also failed for graph '{name}': {fallback_error}. "
                        f"Graph may be partially cleaned or may not exist."
                    )
                    # Don't raise - allow the process to continue
                    # The schema creation will handle existing types
                    return None
            else:
                error_msg = f"Could not drop graph '{name}'. Error: {e}"
                logger.error(error_msg)
                raise RuntimeError(error_msg) from e

    @_wrap_tg_exception
    def execute(self, query, **kwargs):
        """
        Execute GSQL query or installed query based on content.
        """
        try:
            # Check if this is an installed query call
            if query.strip().upper().startswith("RUN "):
                # Extract query name and parameters
                query_name = query.strip()[4:].split("(")[0].strip()
                result = self._run_installed_query(query_name, **kwargs)
            else:
                # Execute as raw GSQL
                result = self._execute_gsql(query)
            return result
        except Exception as e:
            logger.error(f"Error executing query '{query}': {e}")
            raise

    def close(self):
        """Close connection - no cleanup needed (using direct REST API calls)."""
        pass

    def _get_vertex_add_statement(
        self, vertex: Vertex, vertex_config: VertexConfig
    ) -> str:
        """Generate ADD VERTEX statement for a schema change job.

        Args:
            vertex: Vertex object to generate statement for
            vertex_config: Vertex configuration

        Returns:
            str: GSQL ADD VERTEX statement
        """
        vertex_dbname = vertex_config.vertex_dbname(vertex.name)
        index_fields = vertex_config.index(vertex.name).fields

        if len(index_fields) == 0:
            raise ValueError(
                f"Vertex '{vertex_dbname}' must have at least one index field"
            )

        # Get field type for primary key field(s) - convert FieldType enum to string
        field_type_map = {}
        for f in vertex.fields:
            if f.type:
                field_type_map[f.name] = (
                    f.type.value if hasattr(f.type, "value") else str(f.type)
                )
            else:
                field_type_map[f.name] = FieldType.STRING.value

        # Format all fields
        all_fields = []
        for field in vertex.fields:
            if field.type:
                field_type = (
                    field.type.value
                    if hasattr(field.type, "value")
                    else str(field.type)
                )
            else:
                field_type = FieldType.STRING.value
            all_fields.append((field.name, field_type))

        if len(index_fields) == 1:
            # Single field: use PRIMARY_ID syntax (required by GSQL)
            primary_field_name = index_fields[0]
            primary_field_type = field_type_map.get(
                primary_field_name, FieldType.STRING.value
            )

            other_fields = [
                (name, ftype)
                for name, ftype in all_fields
                if name != primary_field_name
            ]

            # Build field list: PRIMARY_ID comes first, then other fields
            field_parts = [f"PRIMARY_ID {primary_field_name} {primary_field_type}"]
            field_parts.extend([f"{name} {ftype}" for name, ftype in other_fields])

            field_definitions = ",\n        ".join(field_parts)

            return (
                f"ADD VERTEX {vertex_dbname} (\n"
                f"        {field_definitions}\n"
                f'    ) WITH STATS="OUTDEGREE_BY_EDGETYPE", PRIMARY_ID_AS_ATTRIBUTE="true"'
            )
        else:
            # Composite key: use PRIMARY KEY syntax
            field_parts = [f"{name} {ftype}" for name, ftype in all_fields]
            vindex = "(" + ", ".join(index_fields) + ")"
            field_parts.append(f"PRIMARY KEY {vindex}")

            field_definitions = ",\n        ".join(field_parts)

            return (
                f"ADD VERTEX {vertex_dbname} (\n"
                f"        {field_definitions}\n"
                f'    ) WITH STATS="OUTDEGREE_BY_EDGETYPE"'
            )

    def _format_edge_attributes(
        self, edge: Edge, exclude_fields: set[str] | None = None
    ) -> str:
        """Format edge attributes for GSQL ADD DIRECTED EDGE statement.

        Args:
            edge: Edge object to format attributes for
            exclude_fields: Optional set of field names to exclude from attributes

        Returns:
            str: Formatted attribute string (e.g., "    date STRING,\n    relation STRING")
        """
        if not edge.weights or not edge.weights.direct:
            return ""

        if exclude_fields is None:
            exclude_fields = set()

        attr_parts = []
        for field in edge.weights.direct:
            field_name = field.name
            if field_name not in exclude_fields:
                field_type = self._get_tigergraph_type(field.type)
                attr_parts.append(f"    {field_name} {field_type}")

        return ",\n".join(attr_parts)

    def _get_edge_add_statement(self, edge: Edge) -> str:
        """Generate ADD DIRECTED EDGE statement for a schema change job.

        Args:
            edge: Edge object to generate statement for

        Returns:
            str: GSQL ADD DIRECTED EDGE statement
        """
        # TigerGraph requires discriminators to support multiple edges of the same type
        # between the same pair of vertices. We add discriminators for all indexed fields.
        # Collect all indexed fields from edge.indexes
        indexed_field_names = set()
        for index in edge.indexes:
            for field_name in index.fields:
                # Skip special fields like "_from", "_to" which are ArangoDB-specific
                if field_name not in ["_from", "_to"]:
                    indexed_field_names.add(field_name)

        # Also include relation_field if it's set (for backward compatibility)
        if edge.relation_field and edge.relation_field not in indexed_field_names:
            indexed_field_names.add(edge.relation_field)

        # IMPORTANT: In TigerGraph, discriminator fields MUST also be edge attributes.
        # If an indexed field is not in weights.direct, we need to add it.
        # Initialize weights if not present
        if edge.weights is None:
            from graflo.architecture.edge import WeightConfig, Field

            edge.weights = WeightConfig()

        # Type assertion: weights is guaranteed to be WeightConfig after assignment
        assert edge.weights is not None, "weights should be initialized"
        # Get existing weight field names
        existing_weight_names = set()
        if edge.weights.direct:
            existing_weight_names = {field.name for field in edge.weights.direct}

        # Add any indexed fields that are missing from weights
        for field_name in indexed_field_names:
            if field_name not in existing_weight_names:
                # Add the field to weights with STRING type (default)
                from graflo.architecture.edge import Field

                edge.weights.direct.append(
                    Field(name=field_name, type=FieldType.STRING)
                )
                logger.info(
                    f"Added indexed field '{field_name}' to edge weights for discriminator compatibility"
                )

        # Format edge attributes, excluding discriminator fields (they're in DISCRIMINATOR clause)
        edge_attrs = self._format_edge_attributes(
            edge, exclude_fields=indexed_field_names
        )

        # Build discriminator clause with all indexed fields
        # DISCRIMINATOR goes INSIDE parentheses, on same line as FROM/TO, with types
        # Format: FROM company, TO company, DISCRIMINATOR(relation STRING), date STRING, ...

        # Get field types for discriminator fields
        field_types = {}
        if edge.weights and edge.weights.direct:
            for field in edge.weights.direct:
                field_types[field.name] = self._get_tigergraph_type(field.type)

        # Use sanitized dbname for schema names when available
        relation_db = edge.relation_dbname

        # Build FROM/TO line with discriminator
        from_to_parts = [
            f"        FROM {edge._source}",
            f"        TO {edge._target}",
        ]

        if indexed_field_names:
            # Format discriminator with types: DISCRIMINATOR(field1 TYPE1, field2 TYPE2)
            discriminator_parts = []
            for field_name in sorted(indexed_field_names):
                field_type = field_types.get(field_name, "STRING")  # Default to STRING
                discriminator_parts.append(f"{field_name} {field_type}")

            discriminator_str = f"DISCRIMINATOR({', '.join(discriminator_parts)})"
            from_to_parts.append(f"        {discriminator_str}")
            logger.info(
                f"Added discriminator for edge {relation_db}: {', '.join(discriminator_parts)}"
            )
        else:
            logger.debug(
                f"No indexed fields found for edge {relation_db}. "
                f"Indexes: {[idx.fields for idx in edge.indexes]}, "
                f"relation_field: {edge.relation_field}"
            )

        # Combine FROM/TO and discriminator with commas
        from_to_line = ",\n".join(from_to_parts)

        # Build the complete statement
        if edge_attrs:
            # Has attributes - add comma after FROM/TO line (which may include discriminator)
            # edge_attrs already has proper indentation, so we just need to add it after a comma
            return (
                f"ADD DIRECTED EDGE {relation_db} (\n"
                f"{from_to_line},\n"
                f"{edge_attrs}\n"
                f"    )"
            )
        else:
            # No attributes - FROM/TO line (which may include discriminator) is the last thing
            # No trailing comma needed
            return f"ADD DIRECTED EDGE {relation_db} (\n{from_to_line}\n    )"

    def _get_edge_group_create_statement(self, edges: list[Edge]) -> str:
        """Generate ADD DIRECTED EDGE statement for a group of edges with the same relation.

        TigerGraph requires edges of the same type to be created in a single statement
        with multiple FROM/TO pairs separated by |.

        Args:
            edges: List of Edge objects with the same relation (edge type)

        Returns:
            str: GSQL ADD DIRECTED EDGE statement with multiple FROM/TO pairs
        """
        if not edges:
            raise ValueError("Cannot create edge statement from empty edge list")

        # Use the first edge to determine attributes and discriminator
        # (all edges of the same relation should have the same schema)
        first_edge = edges[0]
        relation = first_edge.relation_dbname

        # Collect indexed fields for discriminator (same logic as _get_edge_add_statement)
        indexed_field_names = set()
        for index in first_edge.indexes:
            for field_name in index.fields:
                if field_name not in ["_from", "_to"]:
                    indexed_field_names.add(field_name)

        if (
            first_edge.relation_field
            and first_edge.relation_field not in indexed_field_names
        ):
            indexed_field_names.add(first_edge.relation_field)

        # Ensure indexed fields are in weights (same logic as _get_edge_add_statement)
        if first_edge.weights is None:
            from graflo.architecture.edge import WeightConfig

            first_edge.weights = WeightConfig()

        assert first_edge.weights is not None, "weights should be initialized"
        existing_weight_names = set()
        if first_edge.weights.direct:
            existing_weight_names = {field.name for field in first_edge.weights.direct}

        for field_name in indexed_field_names:
            if field_name not in existing_weight_names:
                from graflo.architecture.edge import Field

                first_edge.weights.direct.append(
                    Field(name=field_name, type=FieldType.STRING)
                )

        # Format edge attributes, excluding discriminator fields
        edge_attrs = self._format_edge_attributes(
            first_edge, exclude_fields=indexed_field_names
        )

        # Get field types for discriminator fields
        field_types = {}
        if first_edge.weights and first_edge.weights.direct:
            for field in first_edge.weights.direct:
                field_types[field.name] = self._get_tigergraph_type(field.type)

        # Build FROM/TO pairs for all edges, separated by |
        from_to_lines = []
        for edge in edges:
            # Build FROM/TO line: "FROM A, TO B" or "FROM A, TO B, DISCRIMINATOR(...)"
            from_to_parts = [f"FROM {edge._source}", f"TO {edge._target}"]

            # Add discriminator if needed (same for all edges of the same relation)
            if indexed_field_names:
                discriminator_parts = []
                for field_name in sorted(indexed_field_names):
                    field_type = field_types.get(field_name, "STRING")
                    discriminator_parts.append(f"{field_name} {field_type}")

                discriminator_str = f"DISCRIMINATOR({', '.join(discriminator_parts)})"
                from_to_parts.append(discriminator_str)

            # Combine FROM/TO and discriminator with commas on one line
            from_to_line = ", ".join(from_to_parts)
            from_to_lines.append(f"    {from_to_line}")

        # Join all FROM/TO pairs with |
        all_from_to = " |\n".join(from_to_lines)

        # Build the complete statement
        if edge_attrs:
            # Has attributes - add comma after FROM/TO section
            return (
                f"ADD DIRECTED EDGE {relation} (\n{all_from_to},\n{edge_attrs}\n    )"
            )
        else:
            # No attributes - FROM/TO section is the last thing
            return f"ADD DIRECTED EDGE {relation} (\n{all_from_to}\n    )"

    def _batch_schema_statements(
        self, schema_change_stmts: list[str], graph_name: str, max_job_size: int
    ) -> list[list[str]]:
        """Batch schema change statements into groups that fit within max_job_size.

        Intelligently merges small statements together while ensuring no batch
        exceeds the maximum job size limit.

        Args:
            schema_change_stmts: List of schema change statements to batch
            graph_name: Name of the graph (used for size estimation)
            max_job_size: Maximum size in characters for a single job

        Returns:
            List of batches, where each batch is a list of statements
        """
        if not schema_change_stmts:
            return []

        # Calculate base overhead for a job
        # Use worst-case job name length (multi-batch format) for conservative estimation
        worst_case_job_name = (
            f"schema_change_{graph_name}_batch_999"  # Use large number for worst case
        )
        base_template = (
            f"USE GRAPH {graph_name}\n"
            f"CREATE SCHEMA_CHANGE JOB {worst_case_job_name} FOR GRAPH {graph_name} {{\n"
            f"}}\n"
            f"RUN SCHEMA_CHANGE JOB {worst_case_job_name}"
        )
        base_overhead = len(base_template)

        # Each statement adds 5 characters: first gets "    " (4) + ";" (1),
        # subsequent get ";\n    " (5) between statements, final ";" (1) is included
        # For N statements: 4 (first indent) + (N-1)*5 (separators) + 1 (final semicolon) = 5*N

        def estimate_batch_size(stmts: list[str]) -> int:
            """Estimate the total size of a batch of statements."""
            if not stmts:
                return base_overhead
            total_stmt_size = sum(len(stmt) for stmt in stmts)
            return base_overhead + total_stmt_size + 5 * len(stmts)

        # Calculate total estimated size for all statements
        num_statements = len(schema_change_stmts)
        total_stmt_size = sum(len(stmt) for stmt in schema_change_stmts)
        estimated_size = base_overhead + total_stmt_size + 5 * num_statements

        # If everything fits in one batch, return single batch
        if estimated_size <= max_job_size:
            logger.info(
                f"Applying schema change as single job (estimated size: {estimated_size} chars)"
            )
            return [schema_change_stmts]

        # Need to split into multiple batches
        # Strategy: Use a greedy bin-packing approach that merges small statements
        # Start by creating batches, trying to pack as many statements as possible
        # into each batch without exceeding max_job_size

        batches: list[list[str]] = []

        # Sort statements by size (smallest first) to help pack efficiently
        # We'll process them in order and try to add to existing batches
        stmt_with_size = [(stmt, len(stmt)) for stmt in schema_change_stmts]
        stmt_with_size.sort(key=lambda x: x[1])  # Sort by statement size

        for stmt, stmt_size in stmt_with_size:
            # Calculate overhead for adding this statement: 5 chars (indent + semicolon)
            stmt_overhead = 5

            # Try to add to an existing batch
            added = False
            for batch in batches:
                current_batch_size = estimate_batch_size(batch)
                # Check if adding this statement would exceed the limit
                if current_batch_size + stmt_size + stmt_overhead <= max_job_size:
                    batch.append(stmt)
                    added = True
                    break

            # If couldn't add to existing batch, create a new one
            if not added:
                # Check if statement itself is too large
                single_stmt_size = estimate_batch_size([stmt])
                if single_stmt_size > max_job_size:
                    logger.warning(
                        f"Statement exceeds max_job_size ({single_stmt_size} > {max_job_size}). "
                        f"Will attempt to execute anyway, but may fail."
                    )
                batches.append([stmt])

        logger.info(
            f"Large schema detected (estimated size: {estimated_size} chars). "
            f"Splitting into {len(batches)} batches."
        )

        return batches

    @_wrap_tg_exception
    def _define_schema_local(self, schema: Schema) -> None:
        """Define TigerGraph schema locally for the current graph using a SCHEMA_CHANGE job.

        Args:
            schema: Schema definition
        """
        graph_name = self.config.database
        if not graph_name:
            raise ValueError("Graph name (database) must be configured")

        # Validate graph name
        _validate_tigergraph_schema_name(graph_name, "graph")

        vertex_config = schema.vertex_config
        edge_config = schema.edge_config

        vertex_stmts = []
        edge_stmts = []

        # Vertices
        for vertex in vertex_config.vertices:
            # Validate vertex name
            if vertex.dbname is None:
                raise ValueError(f"Vertex {vertex.name!r} has no dbname")
            _validate_tigergraph_schema_name(vertex.dbname, "vertex")
            stmt = self._get_vertex_add_statement(vertex, vertex_config)
            vertex_stmts.append(stmt)

        # Edges - group by relation since TigerGraph requires edges of the same type
        # to be created in a single statement with multiple FROM/TO pairs
        edges_to_create = list(edge_config.edges_list(include_aux=True))
        for edge in edges_to_create:
            edge.finish_init(vertex_config)
            # Validate edge name using sanitized dbname when available
            edge_dbname = edge.relation_dbname
            _validate_tigergraph_schema_name(edge_dbname, "edge")

        # Group edges by relation
        edges_by_relation: dict[str, list[Edge]] = defaultdict(list)
        for edge in edges_to_create:
            key = edge.relation_dbname
            edges_by_relation[key].append(edge)

        # Create one statement per relation with all FROM/TO pairs
        for relation, edge_group in edges_by_relation.items():
            stmt = self._get_edge_group_create_statement(edge_group)
            edge_stmts.append(stmt)

        if not vertex_stmts and not edge_stmts:
            logger.debug(f"No schema changes to apply for graph '{graph_name}'")
            return

        # Estimate the size of the GSQL command to determine if we need to split it
        # Large SCHEMA_CHANGE JOBs (>30k chars) can cause parser failures with misleading errors
        # like "Missing return statement" (which is actually a parser size limit issue)
        # We'll split into batches based on configurable max_job_size
        # Batch vertices and edges separately, then concatenate
        vertex_batches = (
            self._batch_schema_statements(
                vertex_stmts, graph_name, self.config.max_job_size
            )
            if vertex_stmts
            else []
        )
        edge_batches = (
            self._batch_schema_statements(
                edge_stmts, graph_name, self.config.max_job_size
            )
            if edge_stmts
            else []
        )
        batches = vertex_batches + edge_batches

        # Execute batches sequentially
        for batch_idx, batch_stmts in enumerate(batches):
            job_name = (
                f"schema_change_{graph_name}_batch_{batch_idx}"
                if len(batches) > 1
                else f"schema_change_{graph_name}"
            )

            # First, try to drop the job if it exists (ignore errors if it doesn't)
            try:
                drop_job_cmd = f"USE GRAPH {graph_name}\nDROP JOB {job_name}"
                self._execute_gsql(drop_job_cmd)
                logger.debug(f"Dropped existing schema change job '{job_name}'")
            except Exception as e:
                err_str = str(e).lower()
                # Ignore errors if job doesn't exist
                if "not found" in err_str or "could not be found" in err_str:
                    logger.debug(
                        f"Schema change job '{job_name}' does not exist, skipping drop"
                    )
                else:
                    logger.debug(f"Could not drop schema change job '{job_name}': {e}")

            # Create and run SCHEMA_CHANGE job for this batch
            gsql_commands = [
                f"USE GRAPH {graph_name}",
                f"CREATE SCHEMA_CHANGE JOB {job_name} FOR GRAPH {graph_name} {{",
                "    " + ";\n    ".join(batch_stmts) + ";",
                "}",
                f"RUN SCHEMA_CHANGE JOB {job_name}",
            ]

            full_gsql = "\n".join(gsql_commands)
            actual_size = len(full_gsql)

            # Safety check: warn if actual size exceeds limit (indicates estimation error)
            if actual_size > self.config.max_job_size:
                logger.warning(
                    f"Batch {batch_idx + 1} actual size ({actual_size} chars) exceeds limit ({self.config.max_job_size} chars). "
                    f"This may cause parser errors. Consider reducing max_job_size or improving estimation."
                )

            logger.info(
                f"Applying schema change batch {batch_idx + 1}/{len(batches)} for graph '{graph_name}' "
                f"({len(batch_stmts)} statements, {actual_size} chars)"
            )
            if actual_size < 5000:  # Only log full command if it's reasonably small
                logger.debug(f"GSQL command:\n{full_gsql}")
            else:
                logger.debug(f"GSQL command size: {actual_size} characters")

            try:
                result = self._execute_gsql(full_gsql)
                logger.debug(f"Schema change batch {batch_idx + 1} result: {result}")

                # Check if result indicates success - should contain "Local schema change succeeded." near the end
                result_str = str(result) if result else ""
                if result_str:
                    # Check for success message near the end (last 500 characters to handle long outputs)
                    result_tail = (
                        result_str[-500:] if len(result_str) > 500 else result_str
                    )
                    if "Local schema change succeeded." not in result_tail:
                        error_msg = (
                            f"Schema change job batch {batch_idx + 1} did not report success. "
                            f"Expected 'Local schema change succeeded.' near the end of the result. "
                            f"Result (last 500 chars): {result_tail}"
                        )
                        logger.error(error_msg)
                        logger.error(f"Full result: {result_str}")
                        raise RuntimeError(error_msg)

                # Check if result indicates an error - be more lenient with error detection
                # Only treat as error if result explicitly contains error indicators
                if (
                    result
                    and result_str
                    and (
                        "Encountered" in result_str
                        or "syntax error" in result_str.lower()
                        or "parse error" in result_str.lower()
                        or "missing return statement" in result_str.lower()
                    )
                ):
                    # "Missing return statement" is a misleading error - it's actually a parser size limit
                    # SCHEMA_CHANGE JOB doesn't require RETURN statements, so this indicates parser failure
                    if "missing return statement" in result_str.lower():
                        error_msg = (
                            f"Schema change job batch {batch_idx + 1} failed with parser error. "
                            f"This is likely due to the GSQL command size ({actual_size} chars) exceeding "
                            f"TigerGraph's parser limit (~30-40K chars). The 'Missing return statement' error "
                            f"is misleading - SCHEMA_CHANGE JOB doesn't require RETURN statements. "
                            f"Original error: {result}"
                        )
                    else:
                        error_msg = f"Schema change job batch {batch_idx + 1} reported an error: {result}"

                    logger.error(error_msg)
                    logger.error(
                        f"GSQL command that failed (first 1000 chars):\n{full_gsql[:1000]}..."
                    )
                    raise RuntimeError(error_msg)
            except Exception as e:
                logger.error(
                    f"Failed to execute schema change batch {batch_idx + 1}: {e}"
                )
                raise

        # Verify that the schema was actually created by checking vertex and edge types
        # Wait a moment for schema changes to propagate (after all batches)
        import time

        time.sleep(1.0)  # Increased wait time

        with self._ensure_graph_context(graph_name):
            vertex_types = self._get_vertex_types()
            edge_types = self._get_edge_types()

            # Use vertex_dbname instead of v.name to match what TigerGraph actually creates
            # vertex_dbname returns dbname if set, otherwise None - fallback to v.name if None
            expected_vertex_types = set()
            for v in vertex_config.vertices:
                try:
                    dbname = vertex_config.vertex_dbname(v.name)
                    # If dbname is None, use vertex name
                    expected_name = dbname if dbname is not None else v.name
                except (KeyError, AttributeError):
                    # Fallback to vertex name if vertex_dbname fails
                    expected_name = v.name
                expected_vertex_types.add(expected_name)

            expected_edge_types = {
                e.relation_dbname for e in edges_to_create if e.relation
            }

            # Convert to sets for case-insensitive comparison
            # TigerGraph may capitalize vertex names, so compare case-insensitively
            vertex_types_lower = {vt.lower() for vt in vertex_types}
            expected_vertex_types_lower = {evt.lower() for evt in expected_vertex_types}

            missing_vertices_lower = expected_vertex_types_lower - vertex_types_lower
            # Convert back to original case for error message
            missing_vertices = {
                evt
                for evt in expected_vertex_types
                if evt.lower() in missing_vertices_lower
            }

            missing_edges = expected_edge_types - set(edge_types)

            if missing_vertices or missing_edges:
                error_msg = (
                    f"Schema change job completed but types were not created correctly. "
                    f"Missing vertex types: {missing_vertices}, "
                    f"Missing edge types: {missing_edges}. "
                    f"Created vertex types: {vertex_types}, "
                    f"Created edge types: {edge_types}."
                )
                logger.error(error_msg)
                raise RuntimeError(error_msg)

            logger.info(
                f"Schema verified: {len(vertex_types)} vertex types, {len(edge_types)} edge types created"
            )

    @_wrap_tg_exception
    def init_db(self, schema: Schema, recreate_schema: bool = False) -> None:
        """
        Initialize database with schema definition.

        If the graph already exists and recreate_schema is False, raises
        SchemaExistsError and the script halts.

        Follows the same pattern as ArangoDB:
        1. Halt if graph exists and recreate_schema is False
        2. Clean (drop graph) if recreate_schema
        3. Create graph if not exists
        4. Define schema locally within the graph
        5. Define indexes

        If any step fails, the graph will be cleaned up gracefully.
        """
        # Use schema.general.name for graph creation
        graph_created = False

        # Determine graph name: use config.database if set, otherwise use schema.general.name
        graph_name = self.config.database
        if not graph_name:
            graph_name = schema.general.name
            # Update config for subsequent operations
            self.config.database = graph_name
            logger.info(f"Using schema name '{graph_name}' from schema.general.name")

        # Validate graph name
        _validate_tigergraph_schema_name(graph_name, "graph")

        try:
            if self.graph_exists(graph_name) and not recreate_schema:
                raise SchemaExistsError(
                    f"Schema/graph already exists: graph '{graph_name}'. "
                    "Set recreate_schema=True to replace, or use clear_data=True before ingestion."
                )

            if recreate_schema:
                try:
                    # Only delete the current graph
                    self.delete_database(graph_name)
                    logger.debug(f"Cleaned graph '{graph_name}' for fresh start")
                except Exception as clean_error:
                    logger.warning(
                        f"Error during recreate_schema for graph '{graph_name}': {clean_error}",
                        exc_info=True,
                    )

            # Step 1: Create graph first if it doesn't exist
            if not self.graph_exists(graph_name):
                logger.debug(f"Creating empty graph '{graph_name}'")
                try:
                    # Create empty graph
                    self.create_database(graph_name)
                    graph_created = True
                    logger.info(f"Successfully created empty graph '{graph_name}'")
                except Exception as create_error:
                    logger.error(
                        f"Failed to create graph '{graph_name}': {create_error}",
                        exc_info=True,
                    )
                    raise
            else:
                logger.debug(f"Graph '{graph_name}' already exists in init_db")

            # Step 2: Define schema locally for the graph
            # This uses a SCHEMA_CHANGE job which is the standard way to define local types
            logger.info(f"Defining local schema for graph '{graph_name}'")
            try:
                self._define_schema_local(schema)
            except Exception as schema_error:
                logger.error(
                    f"Failed to define local schema for graph '{graph_name}': {schema_error}",
                    exc_info=True,
                )
                raise

            # Step 3: Define indexes
            try:
                self.define_indexes(schema)
                logger.info(f"Index definition completed for graph '{graph_name}'")
            except Exception as index_error:
                logger.error(
                    f"Failed to define indexes for graph '{graph_name}': {index_error}",
                    exc_info=True,
                )
                raise
        except Exception as e:
            logger.error(f"Error initializing database: {e}")
            # Graceful teardown: if graph was created in this session, clean it up
            if graph_created:
                try:
                    logger.info(
                        f"Cleaning up graph '{graph_name}' after initialization failure"
                    )
                    self.delete_database(graph_name)
                except Exception as cleanup_error:
                    logger.warning(
                        f"Failed to clean up graph '{graph_name}': {cleanup_error}"
                    )
            raise

    @_wrap_tg_exception
    def define_schema(self, schema: Schema):
        """
        Define TigerGraph schema locally for the current graph.

        Assumes graph already exists (created in init_db).
        """
        try:
            self._define_schema_local(schema)
        except Exception as e:
            logger.error(f"Error defining schema: {e}")
            raise

    def define_vertex_classes(  # type: ignore[override]
        self, vertex_config: VertexConfig
    ) -> None:
        """Define TigerGraph vertex types locally for the current graph.

        Args:
            vertex_config: Vertex configuration containing vertices to create
        """
        graph_name = self.config.database
        if not graph_name:
            raise ValueError("Graph name (database) must be configured")

        schema_change_stmts = []
        for vertex in vertex_config.vertices:
            stmt = self._get_vertex_add_statement(vertex, vertex_config)
            schema_change_stmts.append(stmt)

        if not schema_change_stmts:
            return

        job_name = f"add_vertices_{graph_name}"
        gsql_commands = [
            f"USE GRAPH {graph_name}",
            f"DROP JOB {job_name}",
            f"CREATE SCHEMA_CHANGE JOB {job_name} FOR GRAPH {graph_name} {{",
            "    " + ";\n    ".join(schema_change_stmts) + ";",
            "}",
            f"RUN SCHEMA_CHANGE JOB {job_name}",
        ]

        logger.info(f"Adding vertices locally to graph '{graph_name}'")
        self._execute_gsql("\n".join(gsql_commands))

    def define_edge_classes(self, edges: list[Edge]):
        """Define TigerGraph edge types locally for the current graph.

        Args:
            edges: List of edges to create
        """
        graph_name = self.config.database
        if not graph_name:
            raise ValueError("Graph name (database) must be configured")

        # Need vertex_config for dbname lookup if finish_init hasn't been called
        # But edges should ideally already be initialized.
        # If not, this might fail or needs a vertex_config.

        schema_change_stmts = []
        for edge in edges:
            stmt = self._get_edge_add_statement(edge)
            schema_change_stmts.append(stmt)

        if not schema_change_stmts:
            return

        job_name = f"add_edges_{graph_name}"
        gsql_commands = [
            f"USE GRAPH {graph_name}",
            f"DROP JOB {job_name}",
            f"CREATE SCHEMA_CHANGE JOB {job_name} FOR GRAPH {graph_name} {{",
            "    " + ";\n    ".join(schema_change_stmts) + ";",
            "}",
            f"RUN SCHEMA_CHANGE JOB {job_name}",
        ]

        logger.info(f"Adding edges locally to graph '{graph_name}'")
        self._execute_gsql("\n".join(gsql_commands))

    def _format_vertex_fields(self, vertex: Vertex) -> str:
        """
        Format vertex fields for GSQL CREATE VERTEX statement.

        Uses Field objects with types, applying TigerGraph defaults (STRING for None types).
        Formats fields as: field_name TYPE

        Args:
            vertex: Vertex object with Field definitions

        Returns:
            str: Formatted field definitions for GSQL CREATE VERTEX statement
        """
        fields = vertex.fields

        if not fields:
            # Default fields if none specified
            return 'name STRING DEFAULT "",\n    properties MAP<STRING, STRING> DEFAULT (map())'

        field_list = []
        for field in fields:
            # Field type should already be set (STRING if was None)
            field_type = field.type or FieldType.STRING.value
            # Format as: field_name TYPE
            # TODO: Add DEFAULT clause support if needed in the future
            field_list.append(f"{field.name} {field_type}")

        return ",\n    ".join(field_list)

    def _format_edge_attributes_for_create(self, edge: Edge) -> str:
        """
        Format edge attributes for GSQL CREATE EDGE statement.

        Edge weights/attributes come from edge.weights.direct (list of Field objects).
        Each weight field needs to be included in the CREATE EDGE statement with its type.
        """
        attrs = []

        # Get weight fields from edge.weights.direct
        if edge.weights and edge.weights.direct:
            for field in edge.weights.direct:
                # Field objects have name and type attributes
                field_name = field.name
                # Get TigerGraph type - FieldType enum values are already in TigerGraph format
                tg_type = self._get_tigergraph_type(field.type)
                attrs.append(f"{field_name} {tg_type}")

        return ",\n    " + ",\n    ".join(attrs) if attrs else ""

    def _get_tigergraph_type(self, field_type: FieldType | str | None) -> str:
        """
        Convert field type to TigerGraph type string.

        FieldType enum values are already in TigerGraph format (e.g., "INT", "STRING", "DATETIME").
        This method normalizes various input formats to the correct TigerGraph type.

        Args:
            field_type: FieldType enum, string, or None

        Returns:
            str: TigerGraph type string (e.g., "INT", "STRING", "DATETIME")
        """
        if field_type is None:
            return FieldType.STRING.value

        # If it's a FieldType enum, use its value directly (already in TigerGraph format)
        if isinstance(field_type, FieldType):
            return field_type.value

        # If it's an enum-like object with a value attribute
        if hasattr(field_type, "value"):
            enum_value = field_type.value
            # Convert to string and normalize
            enum_value_str = str(enum_value).upper()
            # Check if the value matches a FieldType enum value
            if enum_value_str in VALID_TIGERGRAPH_TYPES:
                return enum_value_str
            # Return as string (normalized to uppercase)
            return enum_value_str

        # If it's a string, normalize and check against FieldType values
        field_type_str = str(field_type).upper()

        # Check if it matches a FieldType enum value directly
        if field_type_str in VALID_TIGERGRAPH_TYPES:
            return field_type_str

        # Handle TigerGraph-specific type aliases
        return TIGERGRAPH_TYPE_ALIASES.get(field_type_str, FieldType.STRING.value)

    def define_vertex_indices(self, vertex_config: VertexConfig):
        """
        TigerGraph automatically indexes primary keys.
        Secondary indices are less common but can be created.
        """
        for vertex_class in vertex_config.vertex_set:
            vertex_dbname = vertex_config.vertex_dbname(vertex_class)
            for index_obj in vertex_config.indexes(vertex_class)[1:]:
                self._add_index(vertex_dbname, index_obj)

    def define_edge_indices(self, edges: list[Edge]):
        """Define indices for edges if specified.

        Note: TigerGraph does not support creating indexes on edge attributes.
        Edge indexes are skipped with a warning. Only vertex indexes are supported.
        """
        for edge in edges:
            if edge.indexes:
                edge_db = edge.relation_dbname
                logger.info(
                    f"Skipping {len(edge.indexes)} index(es) on edge '{edge_db}': "
                    f"TigerGraph does not support indexes on edge attributes. "
                    f"Only vertex indexes are supported."
                )
                # Skip edge index creation - TigerGraph doesn't support it
                # for index_obj in edge.indexes:
                #     self._add_index(edge.relation, index_obj, is_vertex_index=False)

    def _add_index(self, obj_name, index: Index, is_vertex_index=True):
        """
        Create an index on a vertex type using GSQL schema change jobs.

        TigerGraph requires indexes to be created through schema change jobs.
        This implementation creates a local schema change job for the current graph.

        Note: TigerGraph only supports secondary indexes on vertex attributes, not on edge attributes.
        Indexes on edges are not supported and should be skipped.
        TigerGraph only supports indexes on a single field.
        Indexes with multiple fields will be skipped with a warning.

        Args:
            obj_name: Name of the vertex type
            index: Index configuration object
            is_vertex_index: Whether this is a vertex index (True) or edge index (False)
        """
        # TigerGraph does not support indexes on edge attributes
        if not is_vertex_index:
            logger.warning(
                f"Skipping index creation on edge '{obj_name}': "
                f"TigerGraph does not support indexes on edge attributes. "
                f"Only vertex indexes are supported."
            )
            return

        try:
            if not index.fields:
                logger.warning(f"No fields specified for index on {obj_name}, skipping")
                return

            # TigerGraph only supports secondary indexes on a single field
            if len(index.fields) > 1:
                logger.warning(
                    f"TigerGraph only supports indexes on a single field. "
                    f"Skipping multi-field index on {obj_name} with fields {index.fields}"
                )
                return

            # We have exactly one field - proceed with index creation
            field_name = index.fields[0]

            # Generate index name if not provided
            if index.name:
                index_name = index.name
            else:
                # Generate name from obj_name and field name
                index_name = f"{obj_name}_{field_name}_index"

            # Generate job name from obj_name and field name
            job_name = f"add_{obj_name}_{field_name}_index"

            # Build the ALTER command (single field only)
            graph_name = self.config.database

            if not graph_name:
                logger.warning(
                    f"No graph name configured, cannot create index on {obj_name}"
                )
                return

            # Build the ALTER statement inside the job
            # Note: For edges, use "EDGE" not "DIRECTED EDGE" in ALTER statements
            obj_type = "VERTEX" if is_vertex_index else "EDGE"
            alter_stmt = (
                f"ALTER {obj_type} {obj_name} ADD INDEX {index_name} ON ({field_name})"
            )

            # Step 1: Drop existing job if it exists (ignore errors)
            try:
                drop_job_cmd = f"USE GRAPH {graph_name}\nDROP JOB {job_name}"
                self._execute_gsql(drop_job_cmd)
                logger.debug(f"Dropped existing job '{job_name}'")
            except Exception as e:
                err_str = str(e).lower()
                # Ignore errors if job doesn't exist
                if "not found" in err_str or "could not be found" in err_str:
                    logger.debug(f"Job '{job_name}' does not exist, skipping drop")
                else:
                    logger.debug(f"Could not drop job '{job_name}': {e}")

            # Step 2: Create the schema change job
            # Use local schema change for the graph
            create_job_cmd = (
                f"USE GRAPH {graph_name}\n"
                f"CREATE SCHEMA_CHANGE job {job_name} FOR GRAPH {graph_name} {{{alter_stmt};}}"
            )

            logger.debug(f"Executing GSQL (create job): {create_job_cmd}")
            try:
                result = self._execute_gsql(create_job_cmd)
                logger.debug(f"Created schema change job '{job_name}': {result}")
            except Exception as e:
                err = str(e).lower()
                # Check if job already exists
                if (
                    "already exists" in err
                    or "duplicate" in err
                    or "used by another object" in err
                ):
                    logger.debug(f"Schema change job '{job_name}' already exists")
                else:
                    logger.error(
                        f"Failed to create schema change job '{job_name}': {e}"
                    )
                    raise

            # Step 2: Run the schema change job
            run_job_cmd = f"RUN SCHEMA_CHANGE job {job_name}"

            logger.debug(f"Executing GSQL (run job): {run_job_cmd}")
            try:
                result = self._execute_gsql(run_job_cmd)
                logger.debug(
                    f"Ran schema change job '{job_name}', created index '{index_name}' on {obj_name}: {result}"
                )
            except Exception as e:
                err = str(e).lower()
                # Check if index already exists or job was already run
                if (
                    "already exists" in err
                    or "duplicate" in err
                    or "used by another object" in err
                    or "already applied" in err
                ):
                    logger.debug(
                        f"Index '{index_name}' on {obj_name} already exists or job already run, skipping"
                    )
                else:
                    logger.error(f"Failed to run schema change job '{job_name}': {e}")
                    raise
        except Exception as e:
            logger.warning(f"Could not create index for {obj_name}: {e}")

    def _parse_show_output(self, result_str: str, prefix: str) -> list[str]:
        """
        Parse SHOW * output to extract type names.

        Looks for lines matching: "- PREFIX name(" or "PREFIX name("

        Args:
            result_str: String output from SHOW * GSQL command
            prefix: The prefix to look for (e.g., "VERTEX", "EDGE")

        Returns:
            List of extracted names
        """
        import re

        names = []
        # Pattern: "- VERTEX name(" or "VERTEX name("
        # Match lines that contain the prefix followed by a word (the name) and then "("
        pattern = rf"(?:^|\s)-?\s*{re.escape(prefix)}\s+(\w+)\s*\("

        for line in result_str.split("\n"):
            line = line.strip()
            if not line:
                continue

            # Use regex to find matches
            match = re.search(pattern, line, re.IGNORECASE)
            if match:
                name = match.group(1)
                if name and name not in names:
                    names.append(name)

        return names

    def _parse_show_edge_output(self, result_str: str) -> list[tuple[str, bool]]:
        """
        Parse SHOW EDGE * output to extract edge type names and direction.

        Format: "- DIRECTED EDGE belongsTo(FROM Author, TO ResearchField, ...)"
                or "- UNDIRECTED EDGE edgeName(...)"

        Args:
            result_str: String output from SHOW EDGE * GSQL command

        Returns:
            List of tuples (edge_name, is_directed)
        """
        import re

        edge_types = []
        # Pattern for DIRECTED EDGE: "- DIRECTED EDGE name("
        directed_pattern = r"(?:^|\s)-?\s*DIRECTED\s+EDGE\s+(\w+)\s*\("
        # Pattern for UNDIRECTED EDGE: "- UNDIRECTED EDGE name("
        undirected_pattern = r"(?:^|\s)-?\s*UNDIRECTED\s+EDGE\s+(\w+)\s*\("

        for line in result_str.split("\n"):
            line = line.strip()
            if not line:
                continue

            # Check for DIRECTED EDGE
            match = re.search(directed_pattern, line, re.IGNORECASE)
            if match:
                edge_name = match.group(1)
                if edge_name:
                    edge_types.append((edge_name, True))
                continue

            # Check for UNDIRECTED EDGE
            match = re.search(undirected_pattern, line, re.IGNORECASE)
            if match:
                edge_name = match.group(1)
                if edge_name:
                    edge_types.append((edge_name, False))

        return edge_types

    def _is_not_found_error(self, error: Exception | str) -> bool:
        """
        Check if an error indicates that an object doesn't exist.

        Args:
            error: Exception object or error string

        Returns:
            True if the error indicates "not found" or "does not exist"
        """
        err_str = str(error).lower()
        return "does not exist" in err_str or "not found" in err_str

    def _clean_document(self, doc: dict[str, Any]) -> dict[str, Any]:
        """
        Remove internal keys that shouldn't be stored in the database.

        Removes keys starting with "_" except "_key".

        Args:
            doc: Document dictionary to clean

        Returns:
            Cleaned document dictionary
        """
        return {k: v for k, v in doc.items() if not k.startswith("_") or k == "_key"}

    def _parse_show_vertex_output(self, result_str: str) -> list[str]:
        """Parse SHOW VERTEX * output to extract vertex type names."""
        return self._parse_show_output(result_str, "VERTEX")

    def _parse_show_graph_output(self, result_str: str) -> list[str]:
        """Parse SHOW GRAPH * output to extract graph names."""
        return self._parse_show_output(result_str, "GRAPH")

    def _parse_show_job_output(self, result_str: str) -> list[str]:
        """Parse SHOW JOB * output to extract job names."""
        return self._parse_show_output(result_str, "JOB")

    def delete_graph_structure(self, vertex_types=(), graph_names=(), delete_all=False):
        """
        Delete graph structure (graphs, vertex types, edge types) from TigerGraph.

        In TigerGraph:
        - Graph: Top-level container (functions like a database in ArangoDB)
        - Vertex Types: Global vertex type definitions (can be shared across graphs)
        - Edge Types: Global edge type definitions (can be shared across graphs)
        - Vertex and edge types are associated with graphs

        Teardown order:
        1. Drop all graphs
        2. Drop all edge types globally
        3. Drop all vertex types globally
        4. Drop all jobs globally

        Args:
            vertex_types: Vertex type names to delete (not used in TigerGraph teardown)
            graph_names: Graph names to delete (if empty and delete_all=True, deletes all)
            delete_all: If True, perform full teardown of all graphs, edges, vertices, and jobs
        """
        cnames = vertex_types
        gnames = graph_names
        try:
            if delete_all:
                # Step 1: Drop all graphs
                graphs_to_drop = list(gnames) if gnames else []

                # If no specific graphs provided, try to discover and drop all graphs
                if not graphs_to_drop:
                    try:
                        # Use GSQL to list all graphs
                        show_graphs_cmd = "SHOW GRAPH *"
                        result = self._execute_gsql(show_graphs_cmd)
                        result_str = str(result)

                        # Parse graph names using helper method
                        graphs_to_drop = self._parse_show_graph_output(result_str)
                    except Exception as e:
                        logger.debug(f"Could not list graphs: {e}")
                        graphs_to_drop = []

                # Drop each graph
                logger.info(
                    f"Found {len(graphs_to_drop)} graphs to drop: {graphs_to_drop}"
                )
                for graph_name in graphs_to_drop:
                    try:
                        self.delete_database(graph_name)
                        logger.info(f"Successfully dropped graph '{graph_name}'")
                    except Exception as e:
                        if self._is_not_found_error(e):
                            logger.debug(
                                f"Graph '{graph_name}' already dropped or doesn't exist"
                            )
                        else:
                            logger.warning(f"Failed to drop graph '{graph_name}': {e}")
                            logger.warning(
                                f"Error details: {type(e).__name__}: {str(e)}"
                            )

                # Step 2: Drop all edge types globally
                # Note: Edges must be dropped before vertices due to dependencies
                # Edges are global, so we need to query them at global level using GSQL
                try:
                    # Use GSQL to list all global edge types (not graph-scoped)
                    show_edges_cmd = "SHOW EDGE *"
                    result = self._execute_gsql(show_edges_cmd)
                    result_str = str(result)

                    # Parse edge types using helper method
                    edge_types = self._parse_show_edge_output(result_str)

                    logger.info(
                        f"Found {len(edge_types)} edge types to drop: {[name for name, _ in edge_types]}"
                    )
                    for e_type, is_directed in edge_types:
                        try:
                            # DROP EDGE works for both directed and undirected edges
                            drop_edge_cmd = f"DROP EDGE {e_type}"
                            logger.debug(f"Executing: {drop_edge_cmd}")
                            result = self._execute_gsql(drop_edge_cmd)
                            logger.info(
                                f"Successfully dropped edge type '{e_type}': {result}"
                            )
                        except Exception as e:
                            if self._is_not_found_error(e):
                                logger.debug(
                                    f"Edge type '{e_type}' already dropped or doesn't exist"
                                )
                            else:
                                logger.warning(
                                    f"Failed to drop edge type '{e_type}': {e}"
                                )
                                logger.warning(
                                    f"Error details: {type(e).__name__}: {str(e)}"
                                )
                except Exception as e:
                    logger.warning(f"Could not list or drop edge types: {e}")
                    logger.warning(f"Error details: {type(e).__name__}: {str(e)}")

                # Step 3: Drop all vertex types globally
                # Vertices are dropped after edges to avoid dependency issues
                # Vertices are global, so we need to query them at global level using GSQL
                try:
                    # Use GSQL to list all global vertex types (not graph-scoped)
                    show_vertices_cmd = "SHOW VERTEX *"
                    result = self._execute_gsql(show_vertices_cmd)
                    result_str = str(result)

                    # Parse vertex types using helper method
                    vertex_types = self._parse_show_vertex_output(result_str)

                    logger.info(
                        f"Found {len(vertex_types)} vertex types to drop: {vertex_types}"
                    )
                    for v_type in vertex_types:
                        try:
                            # Clear data first to avoid dependency issues
                            try:
                                result = self._delete_vertices(v_type)
                                logger.debug(
                                    f"Cleared data from vertex type '{v_type}': {result}"
                                )
                            except Exception as clear_err:
                                logger.debug(
                                    f"Could not clear data from vertex type '{v_type}': {clear_err}"
                                )

                            # Drop vertex type
                            drop_vertex_cmd = f"DROP VERTEX {v_type}"
                            logger.debug(f"Executing: {drop_vertex_cmd}")
                            result = self._execute_gsql(drop_vertex_cmd)
                            logger.info(
                                f"Successfully dropped vertex type '{v_type}': {result}"
                            )
                        except Exception as e:
                            if self._is_not_found_error(e):
                                logger.debug(
                                    f"Vertex type '{v_type}' already dropped or doesn't exist"
                                )
                            else:
                                logger.warning(
                                    f"Failed to drop vertex type '{v_type}': {e}"
                                )
                                logger.warning(
                                    f"Error details: {type(e).__name__}: {str(e)}"
                                )
                except Exception as e:
                    logger.warning(f"Could not list or drop vertex types: {e}")
                    logger.warning(f"Error details: {type(e).__name__}: {str(e)}")

                # Step 4: Drop all jobs globally
                # Jobs are dropped last since they may reference schema objects
                try:
                    # Use GSQL to list all global jobs
                    show_jobs_cmd = "SHOW JOB *"
                    result = self._execute_gsql(show_jobs_cmd)
                    result_str = str(result)

                    # Parse job names using helper method
                    job_names = self._parse_show_job_output(result_str)

                    logger.info(f"Found {len(job_names)} jobs to drop: {job_names}")
                    for job_name in job_names:
                        try:
                            # Drop job
                            # Jobs can be of different types (SCHEMA_CHANGE, LOADING, etc.)
                            # DROP JOB works for all job types
                            drop_job_cmd = f"DROP JOB {job_name}"
                            logger.debug(f"Executing: {drop_job_cmd}")
                            result = self._execute_gsql(drop_job_cmd)
                            logger.info(
                                f"Successfully dropped job '{job_name}': {result}"
                            )
                        except Exception as e:
                            if self._is_not_found_error(e):
                                logger.debug(
                                    f"Job '{job_name}' already dropped or doesn't exist"
                                )
                            else:
                                logger.warning(f"Failed to drop job '{job_name}': {e}")
                                logger.warning(
                                    f"Error details: {type(e).__name__}: {str(e)}"
                                )
                except Exception as e:
                    logger.warning(f"Could not list or drop jobs: {e}")
                    logger.warning(f"Error details: {type(e).__name__}: {str(e)}")

            elif gnames:
                # Drop specific graphs
                for graph_name in gnames:
                    try:
                        self.delete_database(graph_name)
                    except Exception as e:
                        logger.error(f"Error deleting graph '{graph_name}': {e}")
            elif cnames:
                # Delete vertices from specific vertex types (data only, not schema)
                with self._ensure_graph_context():
                    for class_name in cnames:
                        try:
                            result = self._delete_vertices(class_name)
                            logger.debug(
                                f"Deleted vertices from {class_name}: {result}"
                            )
                        except Exception as e:
                            logger.error(
                                f"Error deleting vertices from {class_name}: {e}"
                            )

        except Exception as e:
            logger.error(f"Error in delete_graph_structure: {e}")

    def clear_data(self, schema: Schema) -> None:
        """Remove all data from the graph without dropping the schema.

        Deletes vertices (and their edges) for all vertex types in the schema.
        """
        vc = schema.vertex_config
        vertex_types = tuple(vc.vertex_dbname(v) for v in vc.vertex_set)
        if vertex_types:
            self.delete_graph_structure(vertex_types=vertex_types)

    def _generate_upsert_payload(
        self, data: list[dict[str, Any]], vname: str, vindex: tuple[str, ...]
    ) -> dict[str, Any]:
        """
        Transforms a list of dictionaries into the TigerGraph REST++ batch upsert JSON format.

        The composite Primary ID is created by concatenating the values of the fields
        specified in vindex with an underscore '_'. Index fields are included in the
        vertex attributes since PRIMARY KEY fields are automatically accessible as
        attributes in TigerGraph queries.

        Attribute values are wrapped in {"value": ...} format as required by TigerGraph REST++ API.

        Args:
            data: List of document dictionaries to upsert
            vname: Target vertex name
            vindex: Tuple of index fields used to create the composite Primary ID

        Returns:
            Dictionary in TigerGraph REST++ batch upsert format:
            {"vertices": {vname: {vertex_id: {attr_name: {"value": attr_value}, ...}}}}
        """
        # Initialize the required JSON structure for vertices
        payload: dict[str, Any] = {"vertices": {vname: {}}}
        vertex_map = payload["vertices"][vname]

        for record in data:
            try:
                # 1. Calculate the Composite Primary ID
                # Assumes all index keys exist in the record
                primary_id_components = [str(record[key]) for key in vindex]
                vertex_id = "_".join(primary_id_components)

                # 2. Clean the record (remove internal keys that shouldn't be stored)
                clean_record = self._clean_document(record)

                # 3. Keep index fields in attributes
                # When using PRIMARY KEY (composite keys), the key fields are automatically
                # accessible as attributes in queries, so we include them in the payload

                # 4. Format attributes for TigerGraph REST++ API
                # TigerGraph requires attribute values to be wrapped in {"value": ...}
                formatted_attributes = {
                    k: {"value": v} for k, v in clean_record.items() if v
                }

                # 5. Add the record attributes to the map using the composite ID as the key
                vertex_map[vertex_id] = formatted_attributes

            except KeyError as e:
                logger.warning(
                    f"Record is missing a required index field: {e}. Skipping record: {record}"
                )
                continue

        return payload

    def _upsert_data(
        self,
        payload: dict[str, Any],
    ) -> dict[str, Any]:
        """
        Sends the generated JSON payload to the TigerGraph REST++ upsert endpoint.

        Args:
            payload: The JSON payload in TigerGraph REST++ format

        Returns:
            Dictionary containing the response from TigerGraph
        """
        graph_name = self.config.database
        if not graph_name:
            raise ValueError("Graph name (database) must be configured")

        # Use restpp_url which handles version-specific prefixes (e.g., /restpp for 4.2.1)
        url = f"{self.restpp_url}/graph/{graph_name}"

        # Use centralized auth headers (supports Bearer token for 4.2.1+)
        headers = self._get_auth_headers()
        headers["Content-Type"] = "application/json"

        logger.debug(f"Attempting batch upsert to: {url}")

        try:
            response = requests.post(
                url,
                headers=headers,
                data=json.dumps(payload, default=_json_serializer),
                # Increase timeout for large batches
                timeout=120,
                verify=self.ssl_verify,
            )
            response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)

            # TigerGraph response is a JSON object
            return response.json()

        except requests_exceptions.HTTPError as errh:
            # For TigerGraph 4.2.1, if token auth fails with 401/REST-10018, try Basic Auth fallback
            if (
                errh.response.status_code == 401
                and self.api_token
                and self.config.username
                and self.config.password
                and "REST-10018" in str(errh)
            ):
                logger.warning(
                    "Token authentication failed with REST-10018, "
                    "falling back to Basic Auth for TigerGraph 4.2.1 compatibility"
                )
                # Retry with Basic Auth
                import base64

                credentials = f"{self.config.username}:{self.config.password}"
                encoded_credentials = base64.b64encode(credentials.encode()).decode()
                headers["Authorization"] = f"Basic {encoded_credentials}"
                try:
                    response = requests.post(
                        url,
                        headers=headers,
                        data=json.dumps(payload, default=_json_serializer),
                        timeout=120,
                        verify=self.ssl_verify,
                    )
                    response.raise_for_status()
                    logger.info("Successfully authenticated using Basic Auth fallback")
                    return response.json()
                except requests_exceptions.HTTPError as errh2:
                    logger.error(f"HTTP Error (after Basic Auth fallback): {errh2}")
                    error_details = ""
                    try:
                        error_details = response.text
                    except Exception:
                        pass
                    return {
                        "error": True,
                        "message": str(errh2),
                        "details": error_details,
                    }

            logger.error(f"HTTP Error: {errh}")
            error_details = ""
            try:
                error_details = response.text
            except Exception:
                pass
            return {"error": True, "message": str(errh), "details": error_details}
        except requests_exceptions.ConnectionError as errc:
            logger.error(f"Error Connecting: {errc}")
            return {"error": True, "message": str(errc)}
        except requests_exceptions.Timeout as errt:
            logger.error(f"Timeout Error: {errt}")
            return {"error": True, "message": str(errt)}
        except requests_exceptions.RequestException as err:
            logger.error(f"An unexpected error occurred: {err}")
            return {"error": True, "message": str(err)}

    @_wrap_tg_exception
    def upsert_docs_batch(self, docs, class_name, match_keys, **kwargs):
        """
        Batch upsert documents as vertices using TigerGraph REST++ API.

        Creates a GSQL job and formats the payload for batch upsert operations.
        Uses composite Primary IDs constructed from match_keys.
        """
        dry = kwargs.pop("dry", False)
        if dry:
            logger.debug(f"Dry run: would upsert {len(docs)} documents to {class_name}")
            return

        try:
            # Convert match_keys to tuple if it's a list
            vindex = tuple(match_keys) if isinstance(match_keys, list) else match_keys

            # Generate the upsert payload
            payload = self._generate_upsert_payload(docs, class_name, vindex)

            # Check if payload has any vertices
            if not payload.get("vertices", {}).get(class_name):
                logger.warning(f"No valid vertices to upsert for {class_name}")
                return

            # Send the upsert request
            result = self._upsert_data(payload)

            if result.get("error"):
                logger.error(
                    f"Error upserting vertices to {class_name}: {result.get('message')}"
                )
            else:
                num_vertices = len(payload["vertices"][class_name])
                logger.debug(
                    f"Upserted {num_vertices} vertices to {class_name}: {result}"
                )
                return result

        except Exception as e:
            logger.error(f"Error upserting vertices to {class_name}: {e}")

    def _generate_edge_upsert_payloads(
        self,
        edges_data: list[tuple[dict, dict, dict]],
        source_class: str,
        target_class: str,
        edge_type: str,
        match_keys_source: tuple[str, ...],
        match_keys_target: tuple[str, ...],
    ) -> list[dict[str, Any]]:
        """
        Transforms edge data into multiple TigerGraph REST++ batch upsert JSON payloads.

        Groups edges by (source_id, target_id, edge_type) and collects all weight combinations
        for each triple. Then creates separate payloads by "zipping" the weight lists across
        all (source_id, target_id, edge_type) groups.

        Args:
            edges_data: List of tuples (source_doc, target_doc, edge_props)
            source_class: Source vertex type name
            target_class: Target vertex type name
            edge_type: Edge type/relation name (e.g., "relates")
            match_keys_source: Tuple of index fields for source vertex
            match_keys_target: Tuple of index fields for target vertex

        Returns:
            List of payload dictionaries in TigerGraph REST++ format:
            [{"edges": {source_v_type: {source_id: {edge_type: {target_v_type: {target_id: attributes}}}}}}, ...]
        """
        from collections import defaultdict

        # Step 1: Group edges by (source_id, target_id, edge_type) and collect weight combinations
        # Structure: {(source_id, target_id, edge_type): [weight_dict1, weight_dict2, ...]}
        uvr_weights_map: defaultdict[tuple[str, str, str], list[dict]] = defaultdict(
            list
        )

        # Also track original edge data for fallback
        uvr_edges_map: defaultdict[
            tuple[str, str, str], list[tuple[dict, dict, dict]]
        ] = defaultdict(list)

        for source_doc, target_doc, edge_props in edges_data:
            try:
                # Extract IDs
                source_id = self._extract_id(source_doc, match_keys_source)
                target_id = self._extract_id(target_doc, match_keys_target)

                if not source_id or not target_id:
                    logger.warning(
                        f"Missing source_id ({source_id}) or target_id ({target_id}) for edge"
                    )
                    continue

                # Clean and format edge attributes
                clean_edge_props = self._clean_document(edge_props)
                formatted_attributes = {
                    k: {"value": v} for k, v in clean_edge_props.items() if v
                }

                # Group by (source_id, target_id, edge_type)
                # edge_type is the actual edge type name (e.g., "relates"), not a weight value
                uvr_key = (source_id, target_id, edge_type)
                uvr_weights_map[uvr_key].append(formatted_attributes)
                uvr_edges_map[uvr_key].append((source_doc, target_doc, edge_props))

            except Exception as e:
                logger.error(f"Error processing edge: {e}")
                continue

        # Step 2: Find the maximum number of weights across all (u, v, r) groups
        # This determines how many payloads we need to create (k payloads for k max elements)
        max_weights = (
            max(len(weights_list) for weights_list in uvr_weights_map.values())
            if uvr_weights_map
            else 0
        )

        if max_weights == 0:
            return []

        # Step 3: Create k payloads by "zipping" weight lists across all (u, v, r) groups
        # Unlike Python's zip() which stops at the shortest iterable, we create k payloads
        # where k is the maximum group size. Payload i contains element i from each group
        # (if that group has an element at index i).
        payloads = []
        for weight_idx in range(max_weights):
            payload: dict[str, Any] = {"edges": {source_class: {}}}
            source_map = payload["edges"][source_class]
            payload_original_edges = []

            # Iterate through all (u, v, r) groups and take element at weight_idx
            for uvr_key, weights_list in uvr_weights_map.items():
                # Skip if this group doesn't have a weight at this index
                if weight_idx >= len(weights_list):
                    continue

                source_id, target_id, edge_type_key = uvr_key
                weight_attrs = weights_list[weight_idx]
                original_edge = uvr_edges_map[uvr_key][weight_idx]

                # Build nested structure
                if source_id not in source_map:
                    source_map[source_id] = {edge_type: {}}

                if edge_type not in source_map[source_id]:
                    source_map[source_id][edge_type] = {target_class: {}}

                if target_class not in source_map[source_id][edge_type]:
                    source_map[source_id][edge_type][target_class] = {}

                target_map = source_map[source_id][edge_type][target_class]

                # Add edge at this index from this (u, v, r) group
                target_map[target_id] = weight_attrs
                payload_original_edges.append(original_edge)

            # Only add payload if it has edges (skip empty payloads)
            if payload_original_edges:
                payload["_original_edges"] = payload_original_edges
                payloads.append(payload)

        return payloads

    def _extract_id(
        self, doc: dict[str, Any], match_keys: list[str] | tuple[str, ...]
    ) -> str | None:
        """
        Extract vertex ID from document based on match keys.

        For composite keys, concatenates values with an underscore '_'.
        Prefers '_key' if present.

        Args:
            doc: Document dictionary
            match_keys: Keys used to identify the vertex

        Returns:
            str | None: The extracted ID or None if missing required fields
        """
        if not doc:
            return None

        # Try _key first (common in ArangoDB style docs)
        if "_key" in doc and doc["_key"]:
            return str(doc["_key"])

        # If multiple match keys, create a composite ID
        if len(match_keys) > 1:
            try:
                id_parts = [str(doc[key]) for key in match_keys]
                return "_".join(id_parts)
            except KeyError:
                return None

        # Single match key
        if len(match_keys) == 1:
            key = match_keys[0]
            if key in doc and doc[key] is not None:
                return str(doc[key])

        return None

    def _fallback_individual_edge_upsert(
        self,
        edges_data: list[tuple[dict, dict, dict]],
        source_class: str,
        target_class: str,
        edge_type: str,
        match_keys_source: tuple[str, ...],
        match_keys_target: tuple[str, ...],
    ) -> None:
        """Fallback method for individual edge upserts.

        Args:
            edges_data: List of tuples (source_doc, target_doc, edge_props)
            source_class: Source vertex type name
            target_class: Target vertex type name
            edge_type: Edge type name
            match_keys_source: Keys for source vertex ID
            match_keys_target: Keys for target vertex ID
        """
        for source_doc, target_doc, edge_props in edges_data:
            try:
                source_id = self._extract_id(source_doc, match_keys_source)
                target_id = self._extract_id(target_doc, match_keys_target)

                if source_id and target_id:
                    clean_edge_props = self._clean_document(edge_props)
                    # Serialize data for REST API
                    serialized_props = json.loads(
                        json.dumps(clean_edge_props, default=_json_serializer)
                    )
                    self._upsert_edge(
                        source_class,
                        source_id,
                        edge_type,
                        target_class,
                        target_id,
                        serialized_props,
                    )
            except Exception as e:
                logger.error(f"Error upserting individual edge: {e}")

    def insert_edges_batch(
        self,
        docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
        source_class: str,
        target_class: str,
        relation_name: str,
        match_keys_source: tuple[str, ...],
        match_keys_target: tuple[str, ...],
        filter_uniques: bool = True,
        head: int | None = None,
        **kwargs: Any,
    ) -> None:
        """
        Batch insert/upsert edges using TigerGraph REST++ API.

        Handles edge data in tuple format: [(source_doc, target_doc, edge_props), ...]
        or dict format: [{"_source_aux": {...}, "_target_aux": {...}, "_edge_props": {...}}, ...]

        Args:
            docs_edges: List of edge documents (tuples or dicts)
            source_class: Source vertex type name
            target_class: Target vertex type name
            relation_name: Edge type/relation name
            match_keys_source: Keys to match source vertices
            match_keys_target: Keys to match target vertices
            filter_uniques: If True, filter duplicate edges (used)
            head: Optional limit on number of edges to insert (used)
            **kwargs: Additional options:
                - dry: If True, don't execute the query
                - collection_name: Alternative edge type name (used if relation_name is None)
                - uniq_weight_fields: Unused in TigerGraph (ArangoDB-specific)
                - uniq_weight_collections: Unused in TigerGraph (ArangoDB-specific)
                - upsert_option: Unused in TigerGraph (ArangoDB-specific, always upserts by default)
        """
        dry = kwargs.pop("dry", False)
        collection_name = kwargs.pop("collection_name", None)
        # Extract and ignore ArangoDB-specific parameters
        kwargs.pop("uniq_weight_fields", None)
        kwargs.pop("uniq_weight_collections", None)
        kwargs.pop("upsert_option", None)
        if dry:
            if docs_edges is not None:
                logger.debug(f"Dry run: would insert {len(docs_edges)} edges")
            return

        # Process edges list
        if isinstance(docs_edges, list):
            if head is not None:
                docs_edges = docs_edges[:head]
            if filter_uniques:
                docs_edges = pick_unique_dict(docs_edges)

        # Normalize edge data format - handle both tuple and dict formats
        if docs_edges is None:
            return
        normalized_edges = []
        for edge_item in docs_edges:
            try:
                if isinstance(edge_item, tuple) and len(edge_item) == 3:
                    # Tuple format: (source_doc, target_doc, edge_props)
                    source_doc, target_doc, edge_props = edge_item
                    normalized_edges.append((source_doc, target_doc, edge_props))
                elif isinstance(edge_item, dict):
                    # Dict format: {"_source_aux": {...}, "_target_aux": {...}, "_edge_props": {...}}
                    source_doc = edge_item.get("_source_aux", {})
                    target_doc = edge_item.get("_target_aux", {})
                    edge_props = edge_item.get("_edge_props", {})
                    normalized_edges.append((source_doc, target_doc, edge_props))
                else:
                    logger.warning(f"Unexpected edge format: {edge_item}")
            except Exception as e:
                logger.error(f"Error normalizing edge item: {e}")
                continue

        if not normalized_edges:
            logger.warning("No valid edges to insert")
            return

        try:
            # Convert match_keys to tuples if they're lists
            match_keys_src = (
                tuple(match_keys_source)
                if isinstance(match_keys_source, list)
                else match_keys_source
            )
            match_keys_tgt = (
                tuple(match_keys_target)
                if isinstance(match_keys_target, list)
                else match_keys_target
            )

            edge_type = relation_name or collection_name
            if not edge_type:
                logger.error(
                    "Edge type must be specified via relation_name or collection_name"
                )
                return

            # Generate multiple edge upsert payloads (one per unique attribute combination)
            payloads = self._generate_edge_upsert_payloads(
                normalized_edges,
                source_class,
                target_class,
                edge_type,
                match_keys_src,
                match_keys_tgt,
            )

            if not payloads:
                logger.warning(f"No valid edges to upsert for edge type {edge_type}")
                return

            # Send each payload in batch
            total_edges = 0
            failed_payloads = []
            for i, payload in enumerate(payloads):
                edges_payload = payload.get("edges", {})
                if not edges_payload or source_class not in edges_payload:
                    continue

                # Store original edges for fallback before removing metadata
                original_edges = payload.pop("_original_edges", [])

                # Send the batch upsert request
                result = self._upsert_data(payload)

                # Restore original edges for potential fallback
                payload["_original_edges"] = original_edges

                if result.get("error"):
                    logger.error(
                        f"Error upserting edges of type {edge_type} (payload {i + 1}/{len(payloads)}): "
                        f"{result.get('message')}"
                    )
                    # Collect failed payload for fallback
                    failed_payloads.append((payload, i))
                else:
                    # Count edges in this payload
                    edge_count = 0
                    for source_id_map in edges_payload[source_class].values():
                        if edge_type in source_id_map:
                            for target_type_map in source_id_map[edge_type].values():
                                for attrs_or_list in target_type_map.values():
                                    if isinstance(attrs_or_list, list):
                                        edge_count += len(attrs_or_list)
                                    else:
                                        edge_count += 1
                    total_edges += edge_count
                    logger.debug(
                        f"Upserted {edge_count} edges of type {edge_type} via batch "
                        f"(payload {i + 1}/{len(payloads)}): {result}"
                    )

            # Handle failed payloads with individual upserts
            if failed_payloads:
                logger.warning(
                    f"{len(failed_payloads)} payload(s) failed, falling back to individual upserts"
                )
                # Extract original edges from failed payloads for individual upsert
                failed_edges = []
                for payload, _ in failed_payloads:
                    # Use the stored original edges for this payload
                    original_edges = payload.get("_original_edges", [])
                    failed_edges.extend(original_edges)

                if failed_edges:
                    logger.debug(
                        f"Sending {len(failed_edges)} edges from failed payloads via individual upserts"
                    )
                    self._fallback_individual_edge_upsert(
                        failed_edges,
                        source_class,
                        target_class,
                        edge_type,
                        match_keys_src,
                        match_keys_tgt,
                    )

            logger.debug(
                f"Total upserted {total_edges} edges of type {edge_type} across {len(payloads)} payloads"
            )
            return

        except Exception as e:
            logger.error(f"Error batch inserting edges: {e}")
            # Fallback to individual operations
            self._fallback_individual_edge_upsert(
                normalized_edges,
                source_class,
                target_class,
                edge_type,
                match_keys_src,
                match_keys_tgt,
            )

    def _extract_id(self, doc, match_keys):
        """
        Extract vertex ID from document based on match keys.
        """
        if not doc:
            return None

        # Try _key first (common in ArangoDB style docs)
        if "_key" in doc and doc["_key"]:
            return str(doc["_key"])

        # Try other match keys
        for key in match_keys:
            if key in doc and doc[key] is not None:
                return str(doc[key])

        # Fallback: create composite ID
        id_parts = []
        for key in match_keys:
            if key in doc and doc[key] is not None:
                id_parts.append(str(doc[key]))

        return "_".join(id_parts) if id_parts else None

    def insert_return_batch(
        self, docs: list[dict[str, Any]], class_name: str
    ) -> list[dict[str, Any]] | str:
        """
        TigerGraph doesn't have INSERT...RETURN semantics like ArangoDB.
        """
        raise NotImplementedError(
            "insert_return_batch not supported in TigerGraph - use upsert_docs_batch instead"
        )

    def _render_rest_filter(
        self,
        filters: list | dict | FilterExpression | None,
        field_types: dict[str, FieldType] | None = None,
    ) -> str:
        """Convert filter expressions to REST++ filter format.

        REST++ filter format: "field=value" or "field>value" etc.
        Format: fieldoperatorvalue (no spaces, quotes for string values)
        Example: "hindex=10" or "hindex>20" or 'name="John"'

        Args:
            filters: Filter expression to convert
            field_types: Optional mapping of field names to FieldType enum values

        Returns:
            str: REST++ filter string (empty if no filters)
        """
        if filters is not None:
            if not isinstance(filters, FilterExpression):
                ff = FilterExpression.from_dict(filters)
            else:
                ff = filters

            # Use GSQL flavor with empty doc_name to trigger REST++ format
            # Pass field_types to help with proper value quoting
            result = ff(
                doc_name="",
                kind=self.expression_flavor(),
                field_types=field_types,
            )
            return result if isinstance(result, str) else ""
        else:
            return ""

    def fetch_docs(
        self,
        class_name: str,
        filters: list[Any] | dict[str, Any] | FilterExpression | None = None,
        limit: int | None = None,
        return_keys: list[str] | None = None,
        unset_keys: list[str] | None = None,
        **kwargs: Any,
    ) -> list[dict[str, Any]]:
        """
        Fetch documents (vertices) with filtering and projection using REST++ API.

        Args:
            class_name: Vertex type name (or dbname)
            filters: Filter expression (list, dict, or FilterExpression)
            limit: Maximum number of documents to return
            return_keys: Keys to return (projection)
            unset_keys: Keys to exclude (projection)
            **kwargs: Additional parameters
                field_types: Optional mapping of field names to FieldType enum values
                           Used to properly quote string values in filters
                           If not provided and vertex_config is provided, will be auto-detected
                vertex_config: Optional VertexConfig object to use for field type lookup

        Returns:
            list: List of fetched documents
        """
        try:
            graph_name = self.config.database
            if not graph_name:
                raise ValueError("Graph name (database) must be configured")

            # Get field_types from kwargs or auto-detect from vertex_config
            field_types = kwargs.get("field_types")
            vertex_config = kwargs.get("vertex_config")

            if field_types is None and vertex_config is not None:
                field_types = {f.name: f.type for f in vertex_config.fields(class_name)}

            # Build REST++ filter string with field type information
            filter_str = self._render_rest_filter(filters, field_types=field_types)

            # Build REST++ API endpoint with query parameters manually
            # Format: /graph/{graph_name}/vertices/{vertex_type}?filter=...&limit=...
            # Example: /graph/g22c97325/vertices/Author?filter=hindex>20&limit=10

            endpoint = f"/graph/{graph_name}/vertices/{class_name}"
            query_parts = []

            if filter_str:
                # URL-encode the filter string to handle special characters
                encoded_filter = quote(filter_str, safe="=<>!&|")
                query_parts.append(f"filter={encoded_filter}")
            if limit is not None:
                query_parts.append(f"limit={limit}")

            if query_parts:
                endpoint = f"{endpoint}?{'&'.join(query_parts)}"

            logger.debug(f"Calling REST++ API: {endpoint}")

            # Call REST++ API directly (no params dict, we built the URL ourselves)
            response = self._call_restpp_api(endpoint)

            # Parse REST++ response (vertices only)
            result: list[dict[str, Any]] = self._parse_restpp_response(
                response, is_edge=False
            )

            # Check for errors
            if isinstance(response, dict) and response.get("error"):
                raise Exception(
                    f"REST++ API error: {response.get('message', response)}"
                )

            # Apply projection (client-side projection is acceptable for result formatting)
            if return_keys is not None:
                result = [
                    {k: doc.get(k) for k in return_keys if k in doc}
                    for doc in result
                    if isinstance(doc, dict)
                ]
            elif unset_keys is not None:
                result = [
                    {k: v for k, v in doc.items() if k not in unset_keys}
                    for doc in result
                    if isinstance(doc, dict)
                ]

            return result

        except Exception as e:
            logger.error(f"Error fetching documents from {class_name} via REST++: {e}")
            raise

    def fetch_edges(
        self,
        from_type: str,
        from_id: str,
        edge_type: str | None = None,
        to_type: str | None = None,
        to_id: str | None = None,
        filters: list[Any] | dict[str, Any] | FilterExpression | None = None,
        limit: int | None = None,
        return_keys: list[str] | None = None,
        unset_keys: list[str] | None = None,
        **kwargs: Any,
    ) -> list[dict[str, Any]]:
        """
        Fetch edges from TigerGraph using REST API.

        In TigerGraph, you must know at least one vertex ID before you can fetch edges.
        Uses REST API which handles special characters in vertex IDs.

        Args:
            from_type: Source vertex type (required)
            from_id: Source vertex ID (required)
            edge_type: Optional edge type to filter by
            to_type: Optional target vertex type to filter by (not used in REST API)
            to_id: Optional target vertex ID to filter by (not used in REST API)
            filters: Additional query filters (not supported by REST API)
            limit: Maximum number of edges to return (not supported by REST API)
            return_keys: Keys to return (projection)
            unset_keys: Keys to exclude (projection)
            **kwargs: Additional parameters

        Returns:
            list: List of fetched edges
        """
        try:
            if not from_type or not from_id:
                raise ValueError(
                    "from_type and from_id are required for fetching edges in TigerGraph"
                )

            # Use REST API to get edges
            # Returns: list of edge dictionaries
            logger.debug(
                f"Fetching edges using REST API: from_type={from_type}, from_id={from_id}, edge_type={edge_type}"
            )

            # Handle None edge_type
            edge_type_str = edge_type if edge_type is not None else None
            edges = self._get_edges(from_type, from_id, edge_type_str)

            # Parse REST API response format
            # _get_edges() returns list of edge dicts from REST++ API
            # Format: [{"e_type": "...", "from_id": "...", "to_id": "...", "attributes": {...}}, ...]
            # The REST API returns edges in a flat format with e_type, from_id, to_id, attributes
            if isinstance(edges, list):
                # Process each edge to normalize format
                result = []
                for edge in edges:
                    if isinstance(edge, dict):
                        # Normalize edge format - REST API returns flat structure
                        normalized_edge = {}

                        # Extract edge type (rename e_type to edge_type for consistency)
                        normalized_edge["edge_type"] = edge.get(
                            "e_type", edge.get("edge_type", "")
                        )

                        # Extract from/to IDs and types
                        normalized_edge["from_id"] = edge.get("from_id", "")
                        normalized_edge["from_type"] = edge.get("from_type", "")
                        normalized_edge["to_id"] = edge.get("to_id", "")
                        normalized_edge["to_type"] = edge.get("to_type", "")

                        # Handle nested "from"/"to" objects if present (some API versions)
                        if "from" in edge and isinstance(edge["from"], dict):
                            normalized_edge["from_id"] = edge["from"].get(
                                "id",
                                edge["from"].get("v_id", normalized_edge["from_id"]),
                            )
                            normalized_edge["from_type"] = edge["from"].get(
                                "type",
                                edge["from"].get(
                                    "v_type", normalized_edge["from_type"]
                                ),
                            )

                        if "to" in edge and isinstance(edge["to"], dict):
                            normalized_edge["to_id"] = edge["to"].get(
                                "id", edge["to"].get("v_id", normalized_edge["to_id"])
                            )
                            normalized_edge["to_type"] = edge["to"].get(
                                "type",
                                edge["to"].get("v_type", normalized_edge["to_type"]),
                            )

                        # Extract attributes and merge into normalized edge
                        attributes = edge.get("attributes", {})
                        if attributes:
                            normalized_edge.update(attributes)
                        else:
                            # If no attributes key, include all other fields as attributes
                            for k, v in edge.items():
                                if k not in (
                                    "e_type",
                                    "edge_type",
                                    "from",
                                    "to",
                                    "from_id",
                                    "to_id",
                                    "from_type",
                                    "to_type",
                                    "directed",
                                ):
                                    normalized_edge[k] = v

                        result.append(normalized_edge)
            elif isinstance(edges, dict):
                # Single edge dict - normalize and wrap in list
                normalized_edge = {}
                normalized_edge["edge_type"] = edges.get(
                    "e_type", edges.get("edge_type", "")
                )
                normalized_edge["from_id"] = edges.get("from_id", "")
                normalized_edge["to_id"] = edges.get("to_id", "")

                if "from" in edges and isinstance(edges["from"], dict):
                    normalized_edge["from_id"] = edges["from"].get(
                        "id", edges["from"].get("v_id", normalized_edge["from_id"])
                    )
                if "to" in edges and isinstance(edges["to"], dict):
                    normalized_edge["to_id"] = edges["to"].get(
                        "id", edges["to"].get("v_id", normalized_edge["to_id"])
                    )

                attributes = edges.get("attributes", {})
                if attributes:
                    normalized_edge.update(attributes)
                else:
                    for k, v in edges.items():
                        if k not in (
                            "e_type",
                            "edge_type",
                            "from",
                            "to",
                            "from_id",
                            "to_id",
                        ):
                            normalized_edge[k] = v

                result = [normalized_edge]
            else:
                # Fallback for unexpected types
                result: list[dict[str, Any]] = []
                logger.debug(f"Unexpected edges type: {type(edges)}")

            # Apply limit if specified (client-side since REST API doesn't support it)
            if limit is not None and limit > 0:
                result = result[:limit]

            # Apply projection (client-side projection is acceptable for result formatting)
            if return_keys is not None:
                result = [
                    {k: doc.get(k) for k in return_keys if k in doc}
                    for doc in result
                    if isinstance(doc, dict)
                ]
            elif unset_keys is not None:
                result = [
                    {k: v for k, v in doc.items() if k not in unset_keys}
                    for doc in result
                    if isinstance(doc, dict)
                ]

            return result

        except Exception as e:
            logger.error(f"Error fetching edges via REST API: {e}")
            raise

    def _parse_restpp_response(
        self, response: dict | list, is_edge: bool = False
    ) -> list[dict]:
        """Parse REST++ API response into list of documents.

        Args:
            response: REST++ API response (dict or list)
            is_edge: Whether this is an edge response (default: False for vertices)

        Returns:
            list: List of parsed documents
        """
        result = []
        if isinstance(response, dict):
            if "results" in response:
                for data in response["results"]:
                    if is_edge:
                        # Edge response format: {"e_type": "...", "from_id": "...", "to_id": "...", "attributes": {...}}
                        edge_type = data.get("e_type", "")
                        from_id = data.get("from_id", data.get("from", ""))
                        to_id = data.get("to_id", data.get("to", ""))
                        attributes = data.get("attributes", {})
                        doc = {
                            **attributes,
                            "edge_type": edge_type,
                            "from_id": from_id,
                            "to_id": to_id,
                        }
                    else:
                        # Vertex response format: {"v_id": "...", "attributes": {...}}
                        vertex_id = data.get("v_id", data.get("id"))
                        attributes = data.get("attributes", {})
                        doc = {**attributes, "id": vertex_id}
                    result.append(doc)
        elif isinstance(response, list):
            # Direct list response
            for data in response:
                if isinstance(data, dict):
                    if is_edge:
                        edge_type = data.get("e_type", "")
                        from_id = data.get("from_id", data.get("from", ""))
                        to_id = data.get("to_id", data.get("to", ""))
                        attributes = data.get("attributes", data)
                        doc = {
                            **attributes,
                            "edge_type": edge_type,
                            "from_id": from_id,
                            "to_id": to_id,
                        }
                    else:
                        vertex_id = data.get("v_id", data.get("id"))
                        attributes = data.get("attributes", data)
                        doc = {**attributes, "id": vertex_id}
                    result.append(doc)
        return result

    def fetch_present_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        flatten: bool = False,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """
        Check which documents from batch are present in the database.
        """
        try:
            present_docs: list[dict[str, Any]] = []
            keep_keys_list: list[str] | tuple[str, ...] = (
                list(keep_keys) if keep_keys is not None else []
            )
            if isinstance(keep_keys_list, tuple):
                keep_keys_list = list(keep_keys_list)

            for doc in batch:
                vertex_id = self._extract_id(doc, match_keys)
                if not vertex_id:
                    continue

                try:
                    vertex_data = self._get_vertices_by_id(class_name, vertex_id)
                    if vertex_data and vertex_id in vertex_data:
                        # Extract requested keys
                        vertex_attrs = vertex_data[vertex_id].get("attributes", {})
                        filtered_doc: dict[str, Any] = {}

                        if keep_keys_list:
                            for key in keep_keys_list:
                                if key == "id":
                                    filtered_doc[key] = vertex_id
                                elif key in vertex_attrs:
                                    filtered_doc[key] = vertex_attrs[key]
                        else:
                            # If no keep_keys specified, return all attributes
                            filtered_doc = vertex_attrs.copy()
                            filtered_doc["id"] = vertex_id

                        present_docs.append(filtered_doc)

                except Exception:
                    # Vertex doesn't exist or error occurred
                    continue

            return present_docs

        except Exception as e:
            logger.error(f"Error fetching present documents: {e}")
            return []

    def aggregate(
        self,
        class_name,
        aggregation_function: AggregationType,
        discriminant: str | None = None,
        aggregated_field: str | None = None,
        filters: list | dict | None = None,
    ):
        """
        Perform aggregation operations.
        """
        try:
            if aggregation_function == AggregationType.COUNT and discriminant is None:
                # Simple vertex count
                count = self._get_vertex_count(class_name)
                return [{"_value": count}]
            else:
                # Complex aggregations require custom GSQL queries
                logger.warning(
                    f"Complex aggregation {aggregation_function} requires custom GSQL implementation"
                )
                return []
        except Exception as e:
            logger.error(f"Error in aggregation: {e}")
            return []

    def keep_absent_documents(
        self,
        batch: list[dict[str, Any]],
        class_name: str,
        match_keys: list[str] | tuple[str, ...],
        keep_keys: list[str] | tuple[str, ...] | None = None,
        filters: list[Any] | dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """
        Return documents from batch that are NOT present in database.
        """
        present_docs = self.fetch_present_documents(
            batch=batch,
            class_name=class_name,
            match_keys=match_keys,
            keep_keys=keep_keys,
            flatten=False,
            filters=filters,
        )

        # Create a set of IDs from present documents for efficient lookup
        present_ids = set()
        for present_doc in present_docs:
            # Extract ID from present document (it should have 'id' key)
            if "id" in present_doc:
                present_ids.add(present_doc["id"])

        # Find documents that are not present
        absent_docs: list[dict[str, Any]] = []
        keep_keys_list: list[str] | tuple[str, ...] = (
            list(keep_keys) if keep_keys is not None else []
        )
        if isinstance(keep_keys_list, tuple):
            keep_keys_list = list(keep_keys_list)

        for doc in batch:
            vertex_id = self._extract_id(doc, match_keys)
            if not vertex_id or vertex_id not in present_ids:
                if keep_keys_list:
                    # Filter to keep only requested keys
                    filtered_doc = {k: doc.get(k) for k in keep_keys_list if k in doc}
                    absent_docs.append(filtered_doc)
                else:
                    absent_docs.append(doc)

        return absent_docs

    @_wrap_tg_exception
    def define_indexes(self, schema: Schema):
        """Define all indexes from schema."""
        try:
            self.define_vertex_indices(schema.vertex_config)
            # Ensure edges are initialized before defining indices
            edges_for_indices = list(schema.edge_config.edges_list(include_aux=True))
            for edge in edges_for_indices:
                if edge._source is None or edge._target is None:
                    edge.finish_init(schema.vertex_config)
            self.define_edge_indices(edges_for_indices)
        except Exception as e:
            logger.error(f"Error defining indexes: {e}")

    def fetch_indexes(self, vertex_type: str | None = None):
        """
        Fetch indexes for vertex types using GSQL.

        In TigerGraph, indexes are associated with vertex types.
        Use DESCRIBE VERTEX to get index information.

        Args:
            vertex_type: Optional vertex type name to fetch indexes for.
                        If None, fetches indexes for all vertex types.

        Returns:
            dict: Mapping of vertex type names to their indexes.
                  Format: {vertex_type: [{"name": "index_name", "fields": ["field1", ...]}, ...]}
        """
        try:
            with self._ensure_graph_context():
                result = {}

                if vertex_type:
                    vertex_types = [vertex_type]
                else:
                    vertex_types = self._get_vertex_types()

                for v_type in vertex_types:
                    try:
                        # Parse indexes from the describe output
                        indexes = []
                        try:
                            indexes.append(
                                {"name": "stat_index", "source": "show_stat"}
                            )
                        except Exception:
                            # If SHOW STAT INDEX doesn't work, try alternative methods
                            pass

                        result[v_type] = indexes
                    except Exception as e:
                        logger.debug(
                            f"Could not fetch indexes for vertex type {v_type}: {e}"
                        )
                        result[v_type] = []

                return result
        except Exception as e:
            logger.error(f"Error fetching indexes: {e}")
            return {}

aggregate(class_name, aggregation_function, discriminant=None, aggregated_field=None, filters=None)

Perform aggregation operations.

Source code in graflo/db/tigergraph/conn.py

def aggregate(
    self,
    class_name,
    aggregation_function: AggregationType,
    discriminant: str | None = None,
    aggregated_field: str | None = None,
    filters: list | dict | None = None,
):
    """
    Perform aggregation operations.
    """
    try:
        if aggregation_function == AggregationType.COUNT and discriminant is None:
            # Simple vertex count
            count = self._get_vertex_count(class_name)
            return [{"_value": count}]
        else:
            # Complex aggregations require custom GSQL queries
            logger.warning(
                f"Complex aggregation {aggregation_function} requires custom GSQL implementation"
            )
            return []
    except Exception as e:
        logger.error(f"Error in aggregation: {e}")
        return []

clear_data(schema)

Remove all data from the graph without dropping the schema.

Deletes vertices (and their edges) for all vertex types in the schema.

Source code in graflo/db/tigergraph/conn.py

def clear_data(self, schema: Schema) -> None:
    """Remove all data from the graph without dropping the schema.

    Deletes vertices (and their edges) for all vertex types in the schema.
    """
    vc = schema.vertex_config
    vertex_types = tuple(vc.vertex_dbname(v) for v in vc.vertex_set)
    if vertex_types:
        self.delete_graph_structure(vertex_types=vertex_types)

close()

Close connection - no cleanup needed (using direct REST API calls).

Source code in graflo/db/tigergraph/conn.py

def close(self):
    """Close connection - no cleanup needed (using direct REST API calls)."""
    pass

create_database(name, vertex_names=None, edge_names=None)

Create a TigerGraph database (graph) using GSQL commands.

This method creates a graph with explicitly attached vertices and edges. Example: CREATE GRAPH researchGraph (author, paper, wrote)

This method uses direct REST API calls to execute GSQL commands that create and use the graph. Supported in TigerGraph version 4.2.2+.

Parameters:

Name	Type	Description	Default
name	str	Name of the graph to create	required
vertex_names	list[str] | None	Optional list of vertex type names to attach to the graph	None
edge_names	list[str] | None	Optional list of edge type names to attach to the graph	None

Raises:

Type	Description
RuntimeError	If graph already exists or creation fails

Source code in graflo/db/tigergraph/conn.py

@_wrap_tg_exception
def create_database(
    self,
    name: str,
    vertex_names: list[str] | None = None,
    edge_names: list[str] | None = None,
):
    """
    Create a TigerGraph database (graph) using GSQL commands.

    This method creates a graph with explicitly attached vertices and edges.
    Example: CREATE GRAPH researchGraph (author, paper, wrote)

    This method uses direct REST API calls to execute GSQL commands
    that create and use the graph. Supported in TigerGraph version 4.2.2+.

    Args:
        name: Name of the graph to create
        vertex_names: Optional list of vertex type names to attach to the graph
        edge_names: Optional list of edge type names to attach to the graph

    Raises:
        RuntimeError: If graph already exists or creation fails
    """
    # Check if graph already exists first
    if self.graph_exists(name):
        raise RuntimeError(f"Graph '{name}' already exists")

    try:
        # Build the list of types to include in CREATE GRAPH
        all_types = []
        if vertex_names:
            all_types.extend(vertex_names)
        if edge_names:
            all_types.extend(edge_names)

        # Format the CREATE GRAPH command with types
        if all_types:
            types_str = ", ".join(all_types)
            gsql_commands = f"CREATE GRAPH {name} ({types_str})\nUSE GRAPH {name}"
        else:
            # Fallback to empty graph if no types provided
            gsql_commands = f"CREATE GRAPH {name}()\nUSE GRAPH {name}"

        # Execute using direct GSQL REST API which handles authentication
        logger.debug(f"Creating graph '{name}' via GSQL: {gsql_commands}")
        try:
            result = self._execute_gsql(gsql_commands)
            logger.info(
                f"Successfully created graph '{name}' with types {all_types}: {result}"
            )
            # Verify the result doesn't indicate the graph already existed
            result_str = str(result).lower()
            if (
                "already exists" in result_str
                or "duplicate" in result_str
                or "graph already exists" in result_str
            ):
                raise RuntimeError(f"Graph '{name}' already exists")
            return result
        except RuntimeError:
            # Re-raise RuntimeError as-is (already handled)
            raise
        except Exception as e:
            error_msg = str(e).lower()
            # Check if graph already exists - raise exception in this case
            # TigerGraph may return various error messages for existing graphs
            if (
                "already exists" in error_msg
                or "duplicate" in error_msg
                or "graph already exists" in error_msg
                or "already exist" in error_msg
            ):
                logger.warning(f"Graph '{name}' already exists: {e}")
                raise RuntimeError(f"Graph '{name}' already exists") from e
            logger.error(f"Failed to create graph '{name}': {e}")
            raise

    except RuntimeError:
        # Re-raise RuntimeError as-is
        raise
    except Exception as e:
        logger.error(f"Error creating graph '{name}' via GSQL: {e}")
        raise

define_edge_classes(edges)

Define TigerGraph edge types locally for the current graph.

Parameters:

Name	Type	Description	Default
edges	list[Edge]	List of edges to create	required

Source code in graflo/db/tigergraph/conn.py

def define_edge_classes(self, edges: list[Edge]):
    """Define TigerGraph edge types locally for the current graph.

    Args:
        edges: List of edges to create
    """
    graph_name = self.config.database
    if not graph_name:
        raise ValueError("Graph name (database) must be configured")

    # Need vertex_config for dbname lookup if finish_init hasn't been called
    # But edges should ideally already be initialized.
    # If not, this might fail or needs a vertex_config.

    schema_change_stmts = []
    for edge in edges:
        stmt = self._get_edge_add_statement(edge)
        schema_change_stmts.append(stmt)

    if not schema_change_stmts:
        return

    job_name = f"add_edges_{graph_name}"
    gsql_commands = [
        f"USE GRAPH {graph_name}",
        f"DROP JOB {job_name}",
        f"CREATE SCHEMA_CHANGE JOB {job_name} FOR GRAPH {graph_name} {{",
        "    " + ";\n    ".join(schema_change_stmts) + ";",
        "}",
        f"RUN SCHEMA_CHANGE JOB {job_name}",
    ]

    logger.info(f"Adding edges locally to graph '{graph_name}'")
    self._execute_gsql("\n".join(gsql_commands))

define_edge_indices(edges)

Define indices for edges if specified.

Note: TigerGraph does not support creating indexes on edge attributes. Edge indexes are skipped with a warning. Only vertex indexes are supported.

Source code in graflo/db/tigergraph/conn.py

def define_edge_indices(self, edges: list[Edge]):
    """Define indices for edges if specified.

    Note: TigerGraph does not support creating indexes on edge attributes.
    Edge indexes are skipped with a warning. Only vertex indexes are supported.
    """
    for edge in edges:
        if edge.indexes:
            edge_db = edge.relation_dbname
            logger.info(
                f"Skipping {len(edge.indexes)} index(es) on edge '{edge_db}': "
                f"TigerGraph does not support indexes on edge attributes. "
                f"Only vertex indexes are supported."
            )

define_indexes(schema)

Define all indexes from schema.

Source code in graflo/db/tigergraph/conn.py

@_wrap_tg_exception
def define_indexes(self, schema: Schema):
    """Define all indexes from schema."""
    try:
        self.define_vertex_indices(schema.vertex_config)
        # Ensure edges are initialized before defining indices
        edges_for_indices = list(schema.edge_config.edges_list(include_aux=True))
        for edge in edges_for_indices:
            if edge._source is None or edge._target is None:
                edge.finish_init(schema.vertex_config)
        self.define_edge_indices(edges_for_indices)
    except Exception as e:
        logger.error(f"Error defining indexes: {e}")

define_schema(schema)

Define TigerGraph schema locally for the current graph.

Assumes graph already exists (created in init_db).

Source code in graflo/db/tigergraph/conn.py

@_wrap_tg_exception
def define_schema(self, schema: Schema):
    """
    Define TigerGraph schema locally for the current graph.

    Assumes graph already exists (created in init_db).
    """
    try:
        self._define_schema_local(schema)
    except Exception as e:
        logger.error(f"Error defining schema: {e}")
        raise

define_vertex_classes(vertex_config)

Define TigerGraph vertex types locally for the current graph.

Parameters:

Name	Type	Description	Default
vertex_config	VertexConfig	Vertex configuration containing vertices to create	required

Source code in graflo/db/tigergraph/conn.py

def define_vertex_classes(  # type: ignore[override]
    self, vertex_config: VertexConfig
) -> None:
    """Define TigerGraph vertex types locally for the current graph.

    Args:
        vertex_config: Vertex configuration containing vertices to create
    """
    graph_name = self.config.database
    if not graph_name:
        raise ValueError("Graph name (database) must be configured")

    schema_change_stmts = []
    for vertex in vertex_config.vertices:
        stmt = self._get_vertex_add_statement(vertex, vertex_config)
        schema_change_stmts.append(stmt)

    if not schema_change_stmts:
        return

    job_name = f"add_vertices_{graph_name}"
    gsql_commands = [
        f"USE GRAPH {graph_name}",
        f"DROP JOB {job_name}",
        f"CREATE SCHEMA_CHANGE JOB {job_name} FOR GRAPH {graph_name} {{",
        "    " + ";\n    ".join(schema_change_stmts) + ";",
        "}",
        f"RUN SCHEMA_CHANGE JOB {job_name}",
    ]

    logger.info(f"Adding vertices locally to graph '{graph_name}'")
    self._execute_gsql("\n".join(gsql_commands))

define_vertex_indices(vertex_config)

TigerGraph automatically indexes primary keys. Secondary indices are less common but can be created.

Source code in graflo/db/tigergraph/conn.py

def define_vertex_indices(self, vertex_config: VertexConfig):
    """
    TigerGraph automatically indexes primary keys.
    Secondary indices are less common but can be created.
    """
    for vertex_class in vertex_config.vertex_set:
        vertex_dbname = vertex_config.vertex_dbname(vertex_class)
        for index_obj in vertex_config.indexes(vertex_class)[1:]:
            self._add_index(vertex_dbname, index_obj)

delete_database(name)

Delete a TigerGraph database (graph).

This method attempts to drop the graph using a clean teardown sequence

1) Drop all queries associated with the graph 2) Drop the graph itself

Parameters:

Name	Type	Description	Default
name	str	Name of the graph to delete	required

Note

In TigerGraph, deleting a graph structure requires the graph to be empty or may fail if it has dependencies. This method handles both cases.

Source code in graflo/db/tigergraph/conn.py

@_wrap_tg_exception
def delete_database(self, name: str):
    """
    Delete a TigerGraph database (graph).

    This method attempts to drop the graph using a clean teardown sequence:
      1) Drop all queries associated with the graph
      2) Drop the graph itself

    Args:
        name: Name of the graph to delete

    Note:
        In TigerGraph, deleting a graph structure requires the graph to be empty
        or may fail if it has dependencies. This method handles both cases.
    """
    try:
        logger.debug(f"Attempting to drop graph '{name}'")

        # The order matters for a clean teardown
        cleanup_script = f"""
            USE GRAPH {name}
            DROP QUERY *
            USE GLOBAL
            DROP GRAPH {name}
        """
        result = self._execute_gsql(cleanup_script)
        logger.info(f"Successfully dropped graph '{name}': {result}")
        return result
    except Exception as e:
        error_str = str(e).lower()
        # If the clean teardown fails, try fallback approaches
        if (
            "depends on" in error_str
            or "query" in error_str
            or "not exist" in error_str
        ):
            logger.warning(
                f"Clean teardown failed for graph '{name}': {e}. "
                f"Attempting fallback cleanup."
            )
            # Fallback: Try to drop queries individually, then drop graph
            try:
                with self._ensure_graph_context(name):
                    try:
                        queries = self._get_installed_queries()
                        if queries:
                            logger.info(
                                f"Dropping {len(queries)} queries from graph '{name}'"
                            )
                            for query_name in queries:
                                try:
                                    drop_query_cmd = f"USE GRAPH {name}\nDROP QUERY {query_name} IF EXISTS"
                                    self._execute_gsql(drop_query_cmd)
                                    logger.debug(
                                        f"Dropped query '{query_name}' from graph '{name}'"
                                    )
                                except Exception:
                                    # Try without IF EXISTS for older TigerGraph versions
                                    try:
                                        drop_query_cmd = f"USE GRAPH {name}\nDROP QUERY {query_name}"
                                        self._execute_gsql(drop_query_cmd)
                                    except Exception as qe2:
                                        logger.debug(
                                            f"Could not drop query '{query_name}': {qe2}"
                                        )
                    except Exception as e2:
                        logger.debug(
                            f"Could not list queries for graph '{name}': {e2}"
                        )

                # Now try to drop the graph
                drop_command = f"USE GLOBAL\nDROP GRAPH {name}"
                result = self._execute_gsql(drop_command)
                logger.info(
                    f"Successfully dropped graph '{name}' via fallback: {result}"
                )
                return result
            except Exception as fallback_error:
                logger.warning(
                    f"Fallback cleanup also failed for graph '{name}': {fallback_error}. "
                    f"Graph may be partially cleaned or may not exist."
                )
                # Don't raise - allow the process to continue
                # The schema creation will handle existing types
                return None
        else:
            error_msg = f"Could not drop graph '{name}'. Error: {e}"
            logger.error(error_msg)
            raise RuntimeError(error_msg) from e

delete_graph_structure(vertex_types=(), graph_names=(), delete_all=False)

Delete graph structure (graphs, vertex types, edge types) from TigerGraph.

In TigerGraph: - Graph: Top-level container (functions like a database in ArangoDB) - Vertex Types: Global vertex type definitions (can be shared across graphs) - Edge Types: Global edge type definitions (can be shared across graphs) - Vertex and edge types are associated with graphs

Teardown order: 1. Drop all graphs 2. Drop all edge types globally 3. Drop all vertex types globally 4. Drop all jobs globally

Parameters:

Name	Type	Description	Default
vertex_types		Vertex type names to delete (not used in TigerGraph teardown)	()
graph_names		Graph names to delete (if empty and delete_all=True, deletes all)	()
delete_all		If True, perform full teardown of all graphs, edges, vertices, and jobs	False

Source code in graflo/db/tigergraph/conn.py

def delete_graph_structure(self, vertex_types=(), graph_names=(), delete_all=False):
    """
    Delete graph structure (graphs, vertex types, edge types) from TigerGraph.

    In TigerGraph:
    - Graph: Top-level container (functions like a database in ArangoDB)
    - Vertex Types: Global vertex type definitions (can be shared across graphs)
    - Edge Types: Global edge type definitions (can be shared across graphs)
    - Vertex and edge types are associated with graphs

    Teardown order:
    1. Drop all graphs
    2. Drop all edge types globally
    3. Drop all vertex types globally
    4. Drop all jobs globally

    Args:
        vertex_types: Vertex type names to delete (not used in TigerGraph teardown)
        graph_names: Graph names to delete (if empty and delete_all=True, deletes all)
        delete_all: If True, perform full teardown of all graphs, edges, vertices, and jobs
    """
    cnames = vertex_types
    gnames = graph_names
    try:
        if delete_all:
            # Step 1: Drop all graphs
            graphs_to_drop = list(gnames) if gnames else []

            # If no specific graphs provided, try to discover and drop all graphs
            if not graphs_to_drop:
                try:
                    # Use GSQL to list all graphs
                    show_graphs_cmd = "SHOW GRAPH *"
                    result = self._execute_gsql(show_graphs_cmd)
                    result_str = str(result)

                    # Parse graph names using helper method
                    graphs_to_drop = self._parse_show_graph_output(result_str)
                except Exception as e:
                    logger.debug(f"Could not list graphs: {e}")
                    graphs_to_drop = []

            # Drop each graph
            logger.info(
                f"Found {len(graphs_to_drop)} graphs to drop: {graphs_to_drop}"
            )
            for graph_name in graphs_to_drop:
                try:
                    self.delete_database(graph_name)
                    logger.info(f"Successfully dropped graph '{graph_name}'")
                except Exception as e:
                    if self._is_not_found_error(e):
                        logger.debug(
                            f"Graph '{graph_name}' already dropped or doesn't exist"
                        )
                    else:
                        logger.warning(f"Failed to drop graph '{graph_name}': {e}")
                        logger.warning(
                            f"Error details: {type(e).__name__}: {str(e)}"
                        )

            # Step 2: Drop all edge types globally
            # Note: Edges must be dropped before vertices due to dependencies
            # Edges are global, so we need to query them at global level using GSQL
            try:
                # Use GSQL to list all global edge types (not graph-scoped)
                show_edges_cmd = "SHOW EDGE *"
                result = self._execute_gsql(show_edges_cmd)
                result_str = str(result)

                # Parse edge types using helper method
                edge_types = self._parse_show_edge_output(result_str)

                logger.info(
                    f"Found {len(edge_types)} edge types to drop: {[name for name, _ in edge_types]}"
                )
                for e_type, is_directed in edge_types:
                    try:
                        # DROP EDGE works for both directed and undirected edges
                        drop_edge_cmd = f"DROP EDGE {e_type}"
                        logger.debug(f"Executing: {drop_edge_cmd}")
                        result = self._execute_gsql(drop_edge_cmd)
                        logger.info(
                            f"Successfully dropped edge type '{e_type}': {result}"
                        )
                    except Exception as e:
                        if self._is_not_found_error(e):
                            logger.debug(
                                f"Edge type '{e_type}' already dropped or doesn't exist"
                            )
                        else:
                            logger.warning(
                                f"Failed to drop edge type '{e_type}': {e}"
                            )
                            logger.warning(
                                f"Error details: {type(e).__name__}: {str(e)}"
                            )
            except Exception as e:
                logger.warning(f"Could not list or drop edge types: {e}")
                logger.warning(f"Error details: {type(e).__name__}: {str(e)}")

            # Step 3: Drop all vertex types globally
            # Vertices are dropped after edges to avoid dependency issues
            # Vertices are global, so we need to query them at global level using GSQL
            try:
                # Use GSQL to list all global vertex types (not graph-scoped)
                show_vertices_cmd = "SHOW VERTEX *"
                result = self._execute_gsql(show_vertices_cmd)
                result_str = str(result)

                # Parse vertex types using helper method
                vertex_types = self._parse_show_vertex_output(result_str)

                logger.info(
                    f"Found {len(vertex_types)} vertex types to drop: {vertex_types}"
                )
                for v_type in vertex_types:
                    try:
                        # Clear data first to avoid dependency issues
                        try:
                            result = self._delete_vertices(v_type)
                            logger.debug(
                                f"Cleared data from vertex type '{v_type}': {result}"
                            )
                        except Exception as clear_err:
                            logger.debug(
                                f"Could not clear data from vertex type '{v_type}': {clear_err}"
                            )

                        # Drop vertex type
                        drop_vertex_cmd = f"DROP VERTEX {v_type}"
                        logger.debug(f"Executing: {drop_vertex_cmd}")
                        result = self._execute_gsql(drop_vertex_cmd)
                        logger.info(
                            f"Successfully dropped vertex type '{v_type}': {result}"
                        )
                    except Exception as e:
                        if self._is_not_found_error(e):
                            logger.debug(
                                f"Vertex type '{v_type}' already dropped or doesn't exist"
                            )
                        else:
                            logger.warning(
                                f"Failed to drop vertex type '{v_type}': {e}"
                            )
                            logger.warning(
                                f"Error details: {type(e).__name__}: {str(e)}"
                            )
            except Exception as e:
                logger.warning(f"Could not list or drop vertex types: {e}")
                logger.warning(f"Error details: {type(e).__name__}: {str(e)}")

            # Step 4: Drop all jobs globally
            # Jobs are dropped last since they may reference schema objects
            try:
                # Use GSQL to list all global jobs
                show_jobs_cmd = "SHOW JOB *"
                result = self._execute_gsql(show_jobs_cmd)
                result_str = str(result)

                # Parse job names using helper method
                job_names = self._parse_show_job_output(result_str)

                logger.info(f"Found {len(job_names)} jobs to drop: {job_names}")
                for job_name in job_names:
                    try:
                        # Drop job
                        # Jobs can be of different types (SCHEMA_CHANGE, LOADING, etc.)
                        # DROP JOB works for all job types
                        drop_job_cmd = f"DROP JOB {job_name}"
                        logger.debug(f"Executing: {drop_job_cmd}")
                        result = self._execute_gsql(drop_job_cmd)
                        logger.info(
                            f"Successfully dropped job '{job_name}': {result}"
                        )
                    except Exception as e:
                        if self._is_not_found_error(e):
                            logger.debug(
                                f"Job '{job_name}' already dropped or doesn't exist"
                            )
                        else:
                            logger.warning(f"Failed to drop job '{job_name}': {e}")
                            logger.warning(
                                f"Error details: {type(e).__name__}: {str(e)}"
                            )
            except Exception as e:
                logger.warning(f"Could not list or drop jobs: {e}")
                logger.warning(f"Error details: {type(e).__name__}: {str(e)}")

        elif gnames:
            # Drop specific graphs
            for graph_name in gnames:
                try:
                    self.delete_database(graph_name)
                except Exception as e:
                    logger.error(f"Error deleting graph '{graph_name}': {e}")
        elif cnames:
            # Delete vertices from specific vertex types (data only, not schema)
            with self._ensure_graph_context():
                for class_name in cnames:
                    try:
                        result = self._delete_vertices(class_name)
                        logger.debug(
                            f"Deleted vertices from {class_name}: {result}"
                        )
                    except Exception as e:
                        logger.error(
                            f"Error deleting vertices from {class_name}: {e}"
                        )

    except Exception as e:
        logger.error(f"Error in delete_graph_structure: {e}")

execute(query, **kwargs)

Execute GSQL query or installed query based on content.

Source code in graflo/db/tigergraph/conn.py

@_wrap_tg_exception
def execute(self, query, **kwargs):
    """
    Execute GSQL query or installed query based on content.
    """
    try:
        # Check if this is an installed query call
        if query.strip().upper().startswith("RUN "):
            # Extract query name and parameters
            query_name = query.strip()[4:].split("(")[0].strip()
            result = self._run_installed_query(query_name, **kwargs)
        else:
            # Execute as raw GSQL
            result = self._execute_gsql(query)
        return result
    except Exception as e:
        logger.error(f"Error executing query '{query}': {e}")
        raise

fetch_docs(class_name, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs)

Fetch documents (vertices) with filtering and projection using REST++ API.

Parameters:

Name	Type	Description	Default
class_name	str	Vertex type name (or dbname)	required
filters	list[Any] | dict[str, Any] | FilterExpression | None	Filter expression (list, dict, or FilterExpression)	None
limit	int | None	Maximum number of documents to return	None
return_keys	list[str] | None	Keys to return (projection)	None
unset_keys	list[str] | None	Keys to exclude (projection)	None
**kwargs	Any	Additional parameters field_types: Optional mapping of field names to FieldType enum values Used to properly quote string values in filters If not provided and vertex_config is provided, will be auto-detected vertex_config: Optional VertexConfig object to use for field type lookup	{}

Returns:

Name	Type	Description
list	list[dict[str, Any]]	List of fetched documents

Source code in graflo/db/tigergraph/conn.py

def fetch_docs(
    self,
    class_name: str,
    filters: list[Any] | dict[str, Any] | FilterExpression | None = None,
    limit: int | None = None,
    return_keys: list[str] | None = None,
    unset_keys: list[str] | None = None,
    **kwargs: Any,
) -> list[dict[str, Any]]:
    """
    Fetch documents (vertices) with filtering and projection using REST++ API.

    Args:
        class_name: Vertex type name (or dbname)
        filters: Filter expression (list, dict, or FilterExpression)
        limit: Maximum number of documents to return
        return_keys: Keys to return (projection)
        unset_keys: Keys to exclude (projection)
        **kwargs: Additional parameters
            field_types: Optional mapping of field names to FieldType enum values
                       Used to properly quote string values in filters
                       If not provided and vertex_config is provided, will be auto-detected
            vertex_config: Optional VertexConfig object to use for field type lookup

    Returns:
        list: List of fetched documents
    """
    try:
        graph_name = self.config.database
        if not graph_name:
            raise ValueError("Graph name (database) must be configured")

        # Get field_types from kwargs or auto-detect from vertex_config
        field_types = kwargs.get("field_types")
        vertex_config = kwargs.get("vertex_config")

        if field_types is None and vertex_config is not None:
            field_types = {f.name: f.type for f in vertex_config.fields(class_name)}

        # Build REST++ filter string with field type information
        filter_str = self._render_rest_filter(filters, field_types=field_types)

        # Build REST++ API endpoint with query parameters manually
        # Format: /graph/{graph_name}/vertices/{vertex_type}?filter=...&limit=...
        # Example: /graph/g22c97325/vertices/Author?filter=hindex>20&limit=10

        endpoint = f"/graph/{graph_name}/vertices/{class_name}"
        query_parts = []

        if filter_str:
            # URL-encode the filter string to handle special characters
            encoded_filter = quote(filter_str, safe="=<>!&|")
            query_parts.append(f"filter={encoded_filter}")
        if limit is not None:
            query_parts.append(f"limit={limit}")

        if query_parts:
            endpoint = f"{endpoint}?{'&'.join(query_parts)}"

        logger.debug(f"Calling REST++ API: {endpoint}")

        # Call REST++ API directly (no params dict, we built the URL ourselves)
        response = self._call_restpp_api(endpoint)

        # Parse REST++ response (vertices only)
        result: list[dict[str, Any]] = self._parse_restpp_response(
            response, is_edge=False
        )

        # Check for errors
        if isinstance(response, dict) and response.get("error"):
            raise Exception(
                f"REST++ API error: {response.get('message', response)}"
            )

        # Apply projection (client-side projection is acceptable for result formatting)
        if return_keys is not None:
            result = [
                {k: doc.get(k) for k in return_keys if k in doc}
                for doc in result
                if isinstance(doc, dict)
            ]
        elif unset_keys is not None:
            result = [
                {k: v for k, v in doc.items() if k not in unset_keys}
                for doc in result
                if isinstance(doc, dict)
            ]

        return result

    except Exception as e:
        logger.error(f"Error fetching documents from {class_name} via REST++: {e}")
        raise

fetch_edges(from_type, from_id, edge_type=None, to_type=None, to_id=None, filters=None, limit=None, return_keys=None, unset_keys=None, **kwargs)

Fetch edges from TigerGraph using REST API.

In TigerGraph, you must know at least one vertex ID before you can fetch edges. Uses REST API which handles special characters in vertex IDs.

Parameters:

Name	Type	Description	Default
from_type	str	Source vertex type (required)	required
from_id	str	Source vertex ID (required)	required
edge_type	str | None	Optional edge type to filter by	None
to_type	str | None	Optional target vertex type to filter by (not used in REST API)	None
to_id	str | None	Optional target vertex ID to filter by (not used in REST API)	None
filters	list[Any] | dict[str, Any] | FilterExpression | None	Additional query filters (not supported by REST API)	None
limit	int | None	Maximum number of edges to return (not supported by REST API)	None
return_keys	list[str] | None	Keys to return (projection)	None
unset_keys	list[str] | None	Keys to exclude (projection)	None
**kwargs	Any	Additional parameters	{}

Returns:

Name	Type	Description
list	list[dict[str, Any]]	List of fetched edges

Source code in graflo/db/tigergraph/conn.py

def fetch_edges(
    self,
    from_type: str,
    from_id: str,
    edge_type: str | None = None,
    to_type: str | None = None,
    to_id: str | None = None,
    filters: list[Any] | dict[str, Any] | FilterExpression | None = None,
    limit: int | None = None,
    return_keys: list[str] | None = None,
    unset_keys: list[str] | None = None,
    **kwargs: Any,
) -> list[dict[str, Any]]:
    """
    Fetch edges from TigerGraph using REST API.

    In TigerGraph, you must know at least one vertex ID before you can fetch edges.
    Uses REST API which handles special characters in vertex IDs.

    Args:
        from_type: Source vertex type (required)
        from_id: Source vertex ID (required)
        edge_type: Optional edge type to filter by
        to_type: Optional target vertex type to filter by (not used in REST API)
        to_id: Optional target vertex ID to filter by (not used in REST API)
        filters: Additional query filters (not supported by REST API)
        limit: Maximum number of edges to return (not supported by REST API)
        return_keys: Keys to return (projection)
        unset_keys: Keys to exclude (projection)
        **kwargs: Additional parameters

    Returns:
        list: List of fetched edges
    """
    try:
        if not from_type or not from_id:
            raise ValueError(
                "from_type and from_id are required for fetching edges in TigerGraph"
            )

        # Use REST API to get edges
        # Returns: list of edge dictionaries
        logger.debug(
            f"Fetching edges using REST API: from_type={from_type}, from_id={from_id}, edge_type={edge_type}"
        )

        # Handle None edge_type
        edge_type_str = edge_type if edge_type is not None else None
        edges = self._get_edges(from_type, from_id, edge_type_str)

        # Parse REST API response format
        # _get_edges() returns list of edge dicts from REST++ API
        # Format: [{"e_type": "...", "from_id": "...", "to_id": "...", "attributes": {...}}, ...]
        # The REST API returns edges in a flat format with e_type, from_id, to_id, attributes
        if isinstance(edges, list):
            # Process each edge to normalize format
            result = []
            for edge in edges:
                if isinstance(edge, dict):
                    # Normalize edge format - REST API returns flat structure
                    normalized_edge = {}

                    # Extract edge type (rename e_type to edge_type for consistency)
                    normalized_edge["edge_type"] = edge.get(
                        "e_type", edge.get("edge_type", "")
                    )

                    # Extract from/to IDs and types
                    normalized_edge["from_id"] = edge.get("from_id", "")
                    normalized_edge["from_type"] = edge.get("from_type", "")
                    normalized_edge["to_id"] = edge.get("to_id", "")
                    normalized_edge["to_type"] = edge.get("to_type", "")

                    # Handle nested "from"/"to" objects if present (some API versions)
                    if "from" in edge and isinstance(edge["from"], dict):
                        normalized_edge["from_id"] = edge["from"].get(
                            "id",
                            edge["from"].get("v_id", normalized_edge["from_id"]),
                        )
                        normalized_edge["from_type"] = edge["from"].get(
                            "type",
                            edge["from"].get(
                                "v_type", normalized_edge["from_type"]
                            ),
                        )

                    if "to" in edge and isinstance(edge["to"], dict):
                        normalized_edge["to_id"] = edge["to"].get(
                            "id", edge["to"].get("v_id", normalized_edge["to_id"])
                        )
                        normalized_edge["to_type"] = edge["to"].get(
                            "type",
                            edge["to"].get("v_type", normalized_edge["to_type"]),
                        )

                    # Extract attributes and merge into normalized edge
                    attributes = edge.get("attributes", {})
                    if attributes:
                        normalized_edge.update(attributes)
                    else:
                        # If no attributes key, include all other fields as attributes
                        for k, v in edge.items():
                            if k not in (
                                "e_type",
                                "edge_type",
                                "from",
                                "to",
                                "from_id",
                                "to_id",
                                "from_type",
                                "to_type",
                                "directed",
                            ):
                                normalized_edge[k] = v

                    result.append(normalized_edge)
        elif isinstance(edges, dict):
            # Single edge dict - normalize and wrap in list
            normalized_edge = {}
            normalized_edge["edge_type"] = edges.get(
                "e_type", edges.get("edge_type", "")
            )
            normalized_edge["from_id"] = edges.get("from_id", "")
            normalized_edge["to_id"] = edges.get("to_id", "")

            if "from" in edges and isinstance(edges["from"], dict):
                normalized_edge["from_id"] = edges["from"].get(
                    "id", edges["from"].get("v_id", normalized_edge["from_id"])
                )
            if "to" in edges and isinstance(edges["to"], dict):
                normalized_edge["to_id"] = edges["to"].get(
                    "id", edges["to"].get("v_id", normalized_edge["to_id"])
                )

            attributes = edges.get("attributes", {})
            if attributes:
                normalized_edge.update(attributes)
            else:
                for k, v in edges.items():
                    if k not in (
                        "e_type",
                        "edge_type",
                        "from",
                        "to",
                        "from_id",
                        "to_id",
                    ):
                        normalized_edge[k] = v

            result = [normalized_edge]
        else:
            # Fallback for unexpected types
            result: list[dict[str, Any]] = []
            logger.debug(f"Unexpected edges type: {type(edges)}")

        # Apply limit if specified (client-side since REST API doesn't support it)
        if limit is not None and limit > 0:
            result = result[:limit]

        # Apply projection (client-side projection is acceptable for result formatting)
        if return_keys is not None:
            result = [
                {k: doc.get(k) for k in return_keys if k in doc}
                for doc in result
                if isinstance(doc, dict)
            ]
        elif unset_keys is not None:
            result = [
                {k: v for k, v in doc.items() if k not in unset_keys}
                for doc in result
                if isinstance(doc, dict)
            ]

        return result

    except Exception as e:
        logger.error(f"Error fetching edges via REST API: {e}")
        raise

fetch_indexes(vertex_type=None)

Fetch indexes for vertex types using GSQL.

In TigerGraph, indexes are associated with vertex types. Use DESCRIBE VERTEX to get index information.

Parameters:

Name	Type	Description	Default
vertex_type	str | None	Optional vertex type name to fetch indexes for. If None, fetches indexes for all vertex types.	None

Returns:

Name	Type	Description
dict		Mapping of vertex type names to their indexes. Format: {vertex_type: [{"name": "index_name", "fields": ["field1", ...]}, ...]}

Source code in graflo/db/tigergraph/conn.py

def fetch_indexes(self, vertex_type: str | None = None):
    """
    Fetch indexes for vertex types using GSQL.

    In TigerGraph, indexes are associated with vertex types.
    Use DESCRIBE VERTEX to get index information.

    Args:
        vertex_type: Optional vertex type name to fetch indexes for.
                    If None, fetches indexes for all vertex types.

    Returns:
        dict: Mapping of vertex type names to their indexes.
              Format: {vertex_type: [{"name": "index_name", "fields": ["field1", ...]}, ...]}
    """
    try:
        with self._ensure_graph_context():
            result = {}

            if vertex_type:
                vertex_types = [vertex_type]
            else:
                vertex_types = self._get_vertex_types()

            for v_type in vertex_types:
                try:
                    # Parse indexes from the describe output
                    indexes = []
                    try:
                        indexes.append(
                            {"name": "stat_index", "source": "show_stat"}
                        )
                    except Exception:
                        # If SHOW STAT INDEX doesn't work, try alternative methods
                        pass

                    result[v_type] = indexes
                except Exception as e:
                    logger.debug(
                        f"Could not fetch indexes for vertex type {v_type}: {e}"
                    )
                    result[v_type] = []

            return result
    except Exception as e:
        logger.error(f"Error fetching indexes: {e}")
        return {}

fetch_present_documents(batch, class_name, match_keys, keep_keys=None, flatten=False, filters=None)

Check which documents from batch are present in the database.

Source code in graflo/db/tigergraph/conn.py

def fetch_present_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    flatten: bool = False,
    filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """
    Check which documents from batch are present in the database.
    """
    try:
        present_docs: list[dict[str, Any]] = []
        keep_keys_list: list[str] | tuple[str, ...] = (
            list(keep_keys) if keep_keys is not None else []
        )
        if isinstance(keep_keys_list, tuple):
            keep_keys_list = list(keep_keys_list)

        for doc in batch:
            vertex_id = self._extract_id(doc, match_keys)
            if not vertex_id:
                continue

            try:
                vertex_data = self._get_vertices_by_id(class_name, vertex_id)
                if vertex_data and vertex_id in vertex_data:
                    # Extract requested keys
                    vertex_attrs = vertex_data[vertex_id].get("attributes", {})
                    filtered_doc: dict[str, Any] = {}

                    if keep_keys_list:
                        for key in keep_keys_list:
                            if key == "id":
                                filtered_doc[key] = vertex_id
                            elif key in vertex_attrs:
                                filtered_doc[key] = vertex_attrs[key]
                    else:
                        # If no keep_keys specified, return all attributes
                        filtered_doc = vertex_attrs.copy()
                        filtered_doc["id"] = vertex_id

                    present_docs.append(filtered_doc)

            except Exception:
                # Vertex doesn't exist or error occurred
                continue

        return present_docs

    except Exception as e:
        logger.error(f"Error fetching present documents: {e}")
        return []

graph_exists(name)

Check if a graph with the given name exists.

Uses the USE GRAPH command and checks the returned message. If the graph doesn't exist, USE GRAPH returns an error message like "Graph 'name' does not exist."

Parameters:

Name	Type	Description	Default
name	str	Name of the graph to check	required

Returns:

Name	Type	Description
bool	bool	True if the graph exists, False otherwise

Source code in graflo/db/tigergraph/conn.py

def graph_exists(self, name: str) -> bool:
    """
    Check if a graph with the given name exists.

    Uses the USE GRAPH command and checks the returned message.
    If the graph doesn't exist, USE GRAPH returns an error message like
    "Graph 'name' does not exist."

    Args:
        name: Name of the graph to check

    Returns:
        bool: True if the graph exists, False otherwise
    """
    try:
        result = self._execute_gsql(f"USE GRAPH {name}")
        result_str = str(result).lower()

        # If the graph doesn't exist, USE GRAPH returns an error message
        # Check for common error messages indicating the graph doesn't exist
        error_patterns = [
            "does not exist",
            "doesn't exist",
            "doesn't exist!",
            f"graph '{name.lower()}' does not exist",
        ]

        # If any error pattern is found, the graph doesn't exist
        for pattern in error_patterns:
            if pattern in result_str:
                return False

        # If no error pattern is found, the graph likely exists
        # (USE GRAPH succeeded or returned success message)
        return True
    except Exception as e:
        logger.debug(f"Error checking if graph '{name}' exists: {e}")
        # If there's an exception, try to parse it
        error_str = str(e).lower()
        if "does not exist" in error_str or "doesn't exist" in error_str:
            return False
        # If exception doesn't indicate "doesn't exist", assume it exists
        # (other errors might indicate connection issues, not missing graph)
        return False

init_db(schema, recreate_schema=False)

Initialize database with schema definition.

If the graph already exists and recreate_schema is False, raises SchemaExistsError and the script halts.

Follows the same pattern as ArangoDB: 1. Halt if graph exists and recreate_schema is False 2. Clean (drop graph) if recreate_schema 3. Create graph if not exists 4. Define schema locally within the graph 5. Define indexes

If any step fails, the graph will be cleaned up gracefully.

Source code in graflo/db/tigergraph/conn.py

@_wrap_tg_exception
def init_db(self, schema: Schema, recreate_schema: bool = False) -> None:
    """
    Initialize database with schema definition.

    If the graph already exists and recreate_schema is False, raises
    SchemaExistsError and the script halts.

    Follows the same pattern as ArangoDB:
    1. Halt if graph exists and recreate_schema is False
    2. Clean (drop graph) if recreate_schema
    3. Create graph if not exists
    4. Define schema locally within the graph
    5. Define indexes

    If any step fails, the graph will be cleaned up gracefully.
    """
    # Use schema.general.name for graph creation
    graph_created = False

    # Determine graph name: use config.database if set, otherwise use schema.general.name
    graph_name = self.config.database
    if not graph_name:
        graph_name = schema.general.name
        # Update config for subsequent operations
        self.config.database = graph_name
        logger.info(f"Using schema name '{graph_name}' from schema.general.name")

    # Validate graph name
    _validate_tigergraph_schema_name(graph_name, "graph")

    try:
        if self.graph_exists(graph_name) and not recreate_schema:
            raise SchemaExistsError(
                f"Schema/graph already exists: graph '{graph_name}'. "
                "Set recreate_schema=True to replace, or use clear_data=True before ingestion."
            )

        if recreate_schema:
            try:
                # Only delete the current graph
                self.delete_database(graph_name)
                logger.debug(f"Cleaned graph '{graph_name}' for fresh start")
            except Exception as clean_error:
                logger.warning(
                    f"Error during recreate_schema for graph '{graph_name}': {clean_error}",
                    exc_info=True,
                )

        # Step 1: Create graph first if it doesn't exist
        if not self.graph_exists(graph_name):
            logger.debug(f"Creating empty graph '{graph_name}'")
            try:
                # Create empty graph
                self.create_database(graph_name)
                graph_created = True
                logger.info(f"Successfully created empty graph '{graph_name}'")
            except Exception as create_error:
                logger.error(
                    f"Failed to create graph '{graph_name}': {create_error}",
                    exc_info=True,
                )
                raise
        else:
            logger.debug(f"Graph '{graph_name}' already exists in init_db")

        # Step 2: Define schema locally for the graph
        # This uses a SCHEMA_CHANGE job which is the standard way to define local types
        logger.info(f"Defining local schema for graph '{graph_name}'")
        try:
            self._define_schema_local(schema)
        except Exception as schema_error:
            logger.error(
                f"Failed to define local schema for graph '{graph_name}': {schema_error}",
                exc_info=True,
            )
            raise

        # Step 3: Define indexes
        try:
            self.define_indexes(schema)
            logger.info(f"Index definition completed for graph '{graph_name}'")
        except Exception as index_error:
            logger.error(
                f"Failed to define indexes for graph '{graph_name}': {index_error}",
                exc_info=True,
            )
            raise
    except Exception as e:
        logger.error(f"Error initializing database: {e}")
        # Graceful teardown: if graph was created in this session, clean it up
        if graph_created:
            try:
                logger.info(
                    f"Cleaning up graph '{graph_name}' after initialization failure"
                )
                self.delete_database(graph_name)
            except Exception as cleanup_error:
                logger.warning(
                    f"Failed to clean up graph '{graph_name}': {cleanup_error}"
                )
        raise

insert_edges_batch(docs_edges, source_class, target_class, relation_name, match_keys_source, match_keys_target, filter_uniques=True, head=None, **kwargs)

Batch insert/upsert edges using TigerGraph REST++ API.

Handles edge data in tuple format: [(source_doc, target_doc, edge_props), ...] or dict format: [{"_source_aux": {...}, "_target_aux": {...}, "_edge_props": {...}}, ...]

Parameters:

Name	Type	Description	Default
docs_edges	list[list[dict[str, Any]]] | list[Any] | None	List of edge documents (tuples or dicts)	required
source_class	str	Source vertex type name	required
target_class	str	Target vertex type name	required
relation_name	str	Edge type/relation name	required
match_keys_source	tuple[str, ...]	Keys to match source vertices	required
match_keys_target	tuple[str, ...]	Keys to match target vertices	required
filter_uniques	bool	If True, filter duplicate edges (used)	True
head	int | None	Optional limit on number of edges to insert (used)	None
**kwargs	Any	Additional options: - dry: If True, don't execute the query - collection_name: Alternative edge type name (used if relation_name is None) - uniq_weight_fields: Unused in TigerGraph (ArangoDB-specific) - uniq_weight_collections: Unused in TigerGraph (ArangoDB-specific) - upsert_option: Unused in TigerGraph (ArangoDB-specific, always upserts by default)	{}

Source code in graflo/db/tigergraph/conn.py

def insert_edges_batch(
    self,
    docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
    source_class: str,
    target_class: str,
    relation_name: str,
    match_keys_source: tuple[str, ...],
    match_keys_target: tuple[str, ...],
    filter_uniques: bool = True,
    head: int | None = None,
    **kwargs: Any,
) -> None:
    """
    Batch insert/upsert edges using TigerGraph REST++ API.

    Handles edge data in tuple format: [(source_doc, target_doc, edge_props), ...]
    or dict format: [{"_source_aux": {...}, "_target_aux": {...}, "_edge_props": {...}}, ...]

    Args:
        docs_edges: List of edge documents (tuples or dicts)
        source_class: Source vertex type name
        target_class: Target vertex type name
        relation_name: Edge type/relation name
        match_keys_source: Keys to match source vertices
        match_keys_target: Keys to match target vertices
        filter_uniques: If True, filter duplicate edges (used)
        head: Optional limit on number of edges to insert (used)
        **kwargs: Additional options:
            - dry: If True, don't execute the query
            - collection_name: Alternative edge type name (used if relation_name is None)
            - uniq_weight_fields: Unused in TigerGraph (ArangoDB-specific)
            - uniq_weight_collections: Unused in TigerGraph (ArangoDB-specific)
            - upsert_option: Unused in TigerGraph (ArangoDB-specific, always upserts by default)
    """
    dry = kwargs.pop("dry", False)
    collection_name = kwargs.pop("collection_name", None)
    # Extract and ignore ArangoDB-specific parameters
    kwargs.pop("uniq_weight_fields", None)
    kwargs.pop("uniq_weight_collections", None)
    kwargs.pop("upsert_option", None)
    if dry:
        if docs_edges is not None:
            logger.debug(f"Dry run: would insert {len(docs_edges)} edges")
        return

    # Process edges list
    if isinstance(docs_edges, list):
        if head is not None:
            docs_edges = docs_edges[:head]
        if filter_uniques:
            docs_edges = pick_unique_dict(docs_edges)

    # Normalize edge data format - handle both tuple and dict formats
    if docs_edges is None:
        return
    normalized_edges = []
    for edge_item in docs_edges:
        try:
            if isinstance(edge_item, tuple) and len(edge_item) == 3:
                # Tuple format: (source_doc, target_doc, edge_props)
                source_doc, target_doc, edge_props = edge_item
                normalized_edges.append((source_doc, target_doc, edge_props))
            elif isinstance(edge_item, dict):
                # Dict format: {"_source_aux": {...}, "_target_aux": {...}, "_edge_props": {...}}
                source_doc = edge_item.get("_source_aux", {})
                target_doc = edge_item.get("_target_aux", {})
                edge_props = edge_item.get("_edge_props", {})
                normalized_edges.append((source_doc, target_doc, edge_props))
            else:
                logger.warning(f"Unexpected edge format: {edge_item}")
        except Exception as e:
            logger.error(f"Error normalizing edge item: {e}")
            continue

    if not normalized_edges:
        logger.warning("No valid edges to insert")
        return

    try:
        # Convert match_keys to tuples if they're lists
        match_keys_src = (
            tuple(match_keys_source)
            if isinstance(match_keys_source, list)
            else match_keys_source
        )
        match_keys_tgt = (
            tuple(match_keys_target)
            if isinstance(match_keys_target, list)
            else match_keys_target
        )

        edge_type = relation_name or collection_name
        if not edge_type:
            logger.error(
                "Edge type must be specified via relation_name or collection_name"
            )
            return

        # Generate multiple edge upsert payloads (one per unique attribute combination)
        payloads = self._generate_edge_upsert_payloads(
            normalized_edges,
            source_class,
            target_class,
            edge_type,
            match_keys_src,
            match_keys_tgt,
        )

        if not payloads:
            logger.warning(f"No valid edges to upsert for edge type {edge_type}")
            return

        # Send each payload in batch
        total_edges = 0
        failed_payloads = []
        for i, payload in enumerate(payloads):
            edges_payload = payload.get("edges", {})
            if not edges_payload or source_class not in edges_payload:
                continue

            # Store original edges for fallback before removing metadata
            original_edges = payload.pop("_original_edges", [])

            # Send the batch upsert request
            result = self._upsert_data(payload)

            # Restore original edges for potential fallback
            payload["_original_edges"] = original_edges

            if result.get("error"):
                logger.error(
                    f"Error upserting edges of type {edge_type} (payload {i + 1}/{len(payloads)}): "
                    f"{result.get('message')}"
                )
                # Collect failed payload for fallback
                failed_payloads.append((payload, i))
            else:
                # Count edges in this payload
                edge_count = 0
                for source_id_map in edges_payload[source_class].values():
                    if edge_type in source_id_map:
                        for target_type_map in source_id_map[edge_type].values():
                            for attrs_or_list in target_type_map.values():
                                if isinstance(attrs_or_list, list):
                                    edge_count += len(attrs_or_list)
                                else:
                                    edge_count += 1
                total_edges += edge_count
                logger.debug(
                    f"Upserted {edge_count} edges of type {edge_type} via batch "
                    f"(payload {i + 1}/{len(payloads)}): {result}"
                )

        # Handle failed payloads with individual upserts
        if failed_payloads:
            logger.warning(
                f"{len(failed_payloads)} payload(s) failed, falling back to individual upserts"
            )
            # Extract original edges from failed payloads for individual upsert
            failed_edges = []
            for payload, _ in failed_payloads:
                # Use the stored original edges for this payload
                original_edges = payload.get("_original_edges", [])
                failed_edges.extend(original_edges)

            if failed_edges:
                logger.debug(
                    f"Sending {len(failed_edges)} edges from failed payloads via individual upserts"
                )
                self._fallback_individual_edge_upsert(
                    failed_edges,
                    source_class,
                    target_class,
                    edge_type,
                    match_keys_src,
                    match_keys_tgt,
                )

        logger.debug(
            f"Total upserted {total_edges} edges of type {edge_type} across {len(payloads)} payloads"
        )
        return

    except Exception as e:
        logger.error(f"Error batch inserting edges: {e}")
        # Fallback to individual operations
        self._fallback_individual_edge_upsert(
            normalized_edges,
            source_class,
            target_class,
            edge_type,
            match_keys_src,
            match_keys_tgt,
        )

insert_return_batch(docs, class_name)

TigerGraph doesn't have INSERT...RETURN semantics like ArangoDB.

Source code in graflo/db/tigergraph/conn.py

def insert_return_batch(
    self, docs: list[dict[str, Any]], class_name: str
) -> list[dict[str, Any]] | str:
    """
    TigerGraph doesn't have INSERT...RETURN semantics like ArangoDB.
    """
    raise NotImplementedError(
        "insert_return_batch not supported in TigerGraph - use upsert_docs_batch instead"
    )

keep_absent_documents(batch, class_name, match_keys, keep_keys=None, filters=None)

Return documents from batch that are NOT present in database.

Source code in graflo/db/tigergraph/conn.py

def keep_absent_documents(
    self,
    batch: list[dict[str, Any]],
    class_name: str,
    match_keys: list[str] | tuple[str, ...],
    keep_keys: list[str] | tuple[str, ...] | None = None,
    filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """
    Return documents from batch that are NOT present in database.
    """
    present_docs = self.fetch_present_documents(
        batch=batch,
        class_name=class_name,
        match_keys=match_keys,
        keep_keys=keep_keys,
        flatten=False,
        filters=filters,
    )

    # Create a set of IDs from present documents for efficient lookup
    present_ids = set()
    for present_doc in present_docs:
        # Extract ID from present document (it should have 'id' key)
        if "id" in present_doc:
            present_ids.add(present_doc["id"])

    # Find documents that are not present
    absent_docs: list[dict[str, Any]] = []
    keep_keys_list: list[str] | tuple[str, ...] = (
        list(keep_keys) if keep_keys is not None else []
    )
    if isinstance(keep_keys_list, tuple):
        keep_keys_list = list(keep_keys_list)

    for doc in batch:
        vertex_id = self._extract_id(doc, match_keys)
        if not vertex_id or vertex_id not in present_ids:
            if keep_keys_list:
                # Filter to keep only requested keys
                filtered_doc = {k: doc.get(k) for k in keep_keys_list if k in doc}
                absent_docs.append(filtered_doc)
            else:
                absent_docs.append(doc)

    return absent_docs

upsert_docs_batch(docs, class_name, match_keys, **kwargs)

Batch upsert documents as vertices using TigerGraph REST++ API.

Creates a GSQL job and formats the payload for batch upsert operations. Uses composite Primary IDs constructed from match_keys.

Source code in graflo/db/tigergraph/conn.py

@_wrap_tg_exception
def upsert_docs_batch(self, docs, class_name, match_keys, **kwargs):
    """
    Batch upsert documents as vertices using TigerGraph REST++ API.

    Creates a GSQL job and formats the payload for batch upsert operations.
    Uses composite Primary IDs constructed from match_keys.
    """
    dry = kwargs.pop("dry", False)
    if dry:
        logger.debug(f"Dry run: would upsert {len(docs)} documents to {class_name}")
        return

    try:
        # Convert match_keys to tuple if it's a list
        vindex = tuple(match_keys) if isinstance(match_keys, list) else match_keys

        # Generate the upsert payload
        payload = self._generate_upsert_payload(docs, class_name, vindex)

        # Check if payload has any vertices
        if not payload.get("vertices", {}).get(class_name):
            logger.warning(f"No valid vertices to upsert for {class_name}")
            return

        # Send the upsert request
        result = self._upsert_data(payload)

        if result.get("error"):
            logger.error(
                f"Error upserting vertices to {class_name}: {result.get('message')}"
            )
        else:
            num_vertices = len(payload["vertices"][class_name])
            logger.debug(
                f"Upserted {num_vertices} vertices to {class_name}: {result}"
            )
            return result

    except Exception as e:
        logger.error(f"Error upserting vertices to {class_name}: {e}")