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

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
    """

    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):
        """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):
        """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, **kwargs):
        """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):
        """Close the ArangoDB connection."""
        # self.conn.close()
        pass

    def init_db(self, schema: Schema, clean_start):
        """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.

        Args:
            schema: Schema containing graph structure definitions
            clean_start: If True, delete all existing collections before initialization
        """
        # 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:
            if clean_start:
                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 clean_start 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 Exception as e:
            logger.error(
                f"Error during database schema initialization for '{db_name}': {e}",
                exc_info=True,
            )
            raise

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

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

    def define_vertex_collections(self, schema: Schema):
        """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 = self.conn.graph(gname)
            else:
                g = self.conn.create_graph(gname)  # type: ignore

            _ = 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_collections(self, edges: list[Edge]):
        """Define edge collections 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 = self.conn.graph(gname)
            else:
                g = self.conn.create_graph(gname)  # type: ignore
            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):
                _ = g.create_edge_definition(
                    edge_collection=collection_name,
                    from_vertex_collections=[item._source],
                    to_vertex_collections=[item._target],
                )

    def _add_index(self, general_collection, index: Index):
        """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
        """
        data = index.db_form(DBFlavor.ARANGO)
        if index.type == IndexType.PERSISTENT:
            ih = general_collection.add_index(data)
        if 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"}
            )
        else:
            ih = None
        return ih

    def define_vertex_indices(self, vertex_config: VertexConfig):
        """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 = [tuple(ix["fields"]) for ix in ixs]
            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]):
        """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):
        """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
        """
        if db_class_name is None:
            classes = self.conn.collections()
        elif self.conn.has_collection(db_class_name):
            classes = [self.conn.collection(db_class_name)]
        else:
            classes = []

        r = {}
        for cname in classes:
            assert isinstance(cname["name"], str)
            c = self.conn.collection(cname["name"])
            r[cname["name"]] = c.indexes()
        return r

    def create_collection(self, db_class_name, index: None | Index = None, g=None):
        """Create a new ArangoDB collection.

        Args:
            db_class_name: Name of the collection to create
            index: Optional index to create on the collection
            g: Optional graph to create the collection in

        Returns:
            IndexHandle: Handle to the created index if one was created
        """
        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=(), graph_names=(), delete_all=False):
        """Delete graph structure (collections and graphs) from ArangoDB.

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

        Args:
            vertex_types: Collection names to delete (vertex or edge collections)
            graph_names: Graph names to delete
            delete_all: If True, delete all non-system collections and graphs
        """
        cnames = vertex_types
        gnames = graph_names
        logger.info("collections (non system):")
        logger.info([c for c in self.conn.collections() if c["name"][0] != "_"])

        if delete_all:
            cnames = [c["name"] for c in self.conn.collections() if c["name"][0] != "_"]
            gnames = [g["name"] for g in self.conn.graphs()]

        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("collections (after delete operation):")
        logger.info([c for c in self.conn.collections() if c["name"][0] != "_"])

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

    def get_collections(self):
        """Get all collections in the database.

        Returns:
            list: List of collection information dictionaries
        """
        return self.conn.collections()

    def upsert_docs_batch(
        self,
        docs,
        class_name,
        match_keys: list[str] | None = None,
        **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.

        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 isinstance(docs, list):
            if filter_uniques:
                docs = pick_unique_dict(docs)
            docs = json.dumps(docs, default=_json_serializer)
        if match_keys is None:
            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}
                            {upsert_clause}
                            INSERT doc
                            {update_clause} 
                                IN {class_name} {options}"""
        if not dry:
            self.execute(q_update)

    def insert_edges_batch(
        self,
        docs_edges,
        source_class,
        target_class,
        relation_name=None,
        collection_name=None,
        match_keys_source=("_key",),
        match_keys_target=("_key",),
        filter_uniques=True,
        uniq_weight_fields=None,
        uniq_weight_collections=None,
        upsert_option=False,
        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.

        Args:
            docs_edges: List of edge documents in format [{_source_aux: source_doc, _target_aux: target_doc}]
            source_class: Source vertex collection name
            target_class: Target vertex collection name
            relation_name: Optional relation name for the edges
            collection_name: Edge collection name
            match_keys_source: Keys to match source vertices
            match_keys_target: Keys to match target vertices
            filter_uniques: If True, filter duplicate edges
            uniq_weight_fields: Fields to consider for uniqueness
            uniq_weight_collections: Collections to consider for uniqueness
            upsert_option: If True, use upsert instead of insert
            head: Optional limit on number of edges to insert
            **kwargs: Additional options:
                - dry: If True, don't execute the query
        """
        dry = kwargs.pop("dry", 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, class_name):
        """Insert documents and return their keys.

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

        Returns:
            str: AQL query string for the operation
        """
        docs = json.dumps(docs)
        query0 = f"""FOR doc in {docs}
              INSERT doc
              INTO {class_name}
              LET inserted = NEW
              RETURN {{_key: inserted._key}}
        """
        return query0

    def fetch_present_documents(
        self,
        batch,
        class_name,
        match_keys,
        keep_keys,
        flatten=False,
        filters: None | Clause | list | dict = None,
    ) -> list | dict:
        """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 = {}
            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,
        filters: None | Clause | list | dict = None,
        limit: int | None = None,
        return_keys: list | None = None,
        unset_keys: list | None = None,
        **kwargs,
    ):
        """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 | dict | Clause | None = None,
        limit: int | None = None,
        return_keys: list | None = None,
        unset_keys: list | None = None,
        **kwargs,
    ):
        """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,
        aggregation_function: AggregationType,
        discriminant: str | None = None,
        aggregated_field: str | None = None,
        filters: None | Clause | list | dict = None,
    ):
        """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,
        class_name,
        match_keys,
        keep_keys,
        filters: None | Clause | list | dict = None,
    ):
        """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, field):
        """Update a field to numeric type in all documents.

        Args:
            collection_name: Collection to update
            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		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 | Clause | list | dict	Query filters	None

Returns:

Name	Type	Description
list		Aggregation results

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

def aggregate(
    self,
    class_name,
    aggregation_function: AggregationType,
    discriminant: str | None = None,
    aggregated_field: str | None = None,
    filters: None | Clause | list | dict = None,
):
    """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

close()

Close the ArangoDB connection.

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

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

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

Create a new ArangoDB collection.

Parameters:

Name	Type	Description	Default
db_class_name		Name of the collection to create	required
index	None | Index	Optional index to create on the collection	None
g		Optional graph to create the collection in	None

Returns:

Name	Type	Description
IndexHandle		Handle to the created index if one was created

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

def create_collection(self, db_class_name, index: None | Index = None, g=None):
    """Create a new ArangoDB collection.

    Args:
        db_class_name: Name of the collection to create
        index: Optional index to create on the collection
        g: Optional graph to create the collection in

    Returns:
        IndexHandle: Handle to the created index if one was created
    """
    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):
    """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_collections(edges)

Define edge collections 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_collections(self, edges: list[Edge]):
    """Define edge collections 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 = self.conn.graph(gname)
        else:
            g = self.conn.create_graph(gname)  # type: ignore
        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):
            _ = 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]):
    """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):
    """Define ArangoDB collections based on schema.

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

define_vertex_collections(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_collections(self, schema: Schema):
    """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 = self.conn.graph(gname)
        else:
            g = self.conn.create_graph(gname)  # type: ignore

        _ = 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):
    """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 = [tuple(ix["fields"]) for ix in ixs]
        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):
    """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 (collections and graphs) from ArangoDB.

In ArangoDB: - Collections: 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		Collection names to delete (vertex or edge collections)	()
graph_names		Graph names to delete	()
delete_all		If True, delete all non-system collections and graphs	False

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

def delete_graph_structure(self, vertex_types=(), graph_names=(), delete_all=False):
    """Delete graph structure (collections and graphs) from ArangoDB.

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

    Args:
        vertex_types: Collection names to delete (vertex or edge collections)
        graph_names: Graph names to delete
        delete_all: If True, delete all non-system collections and graphs
    """
    cnames = vertex_types
    gnames = graph_names
    logger.info("collections (non system):")
    logger.info([c for c in self.conn.collections() if c["name"][0] != "_"])

    if delete_all:
        cnames = [c["name"] for c in self.conn.collections() if c["name"][0] != "_"]
        gnames = [g["name"] for g in self.conn.graphs()]

    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("collections (after delete operation):")
    logger.info([c for c in self.conn.collections() if c["name"][0] != "_"])

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

execute(query, **kwargs)

Execute an AQL query.

Parameters:

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

Returns:

Name	Type	Description
Cursor		ArangoDB cursor for the query results

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

def execute(self, query, **kwargs):
    """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		Collection to fetch from	required
filters	None | Clause | list | dict	Query filters	None
limit	int | None	Maximum number of documents to return	None
return_keys	list | None	Keys to return	None
unset_keys	list | None	Keys to unset	None

Returns:

Name	Type	Description
list		Fetched documents

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

def fetch_docs(
    self,
    class_name,
    filters: None | Clause | list | dict = None,
    limit: int | None = None,
    return_keys: list | None = None,
    unset_keys: list | None = None,
    **kwargs,
):
    """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 | dict | Clause | 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)	None
**kwargs		Additional parameters	{}

Returns:

Name	Type	Description
list		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 | dict | Clause | None = None,
    limit: int | None = None,
    return_keys: list | None = None,
    unset_keys: list | None = None,
    **kwargs,
):
    """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		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):
    """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
    """
    if db_class_name is None:
        classes = self.conn.collections()
    elif self.conn.has_collection(db_class_name):
        classes = [self.conn.collection(db_class_name)]
    else:
        classes = []

    r = {}
    for cname in classes:
        assert isinstance(cname["name"], str)
        c = self.conn.collection(cname["name"])
        r[cname["name"]] = c.indexes()
    return r

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

Fetch documents that exist in the database.

Parameters:

Name	Type	Description	Default
batch		Batch of documents to check	required
class_name		Collection to check in	required
match_keys		Keys to match documents	required
keep_keys		Keys to keep in result	required
flatten		If True, flatten the result into a list	False
filters	None | Clause | list | dict	Additional query filters	None

Returns:

Type	Description
list | dict	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,
    class_name,
    match_keys,
    keep_keys,
    flatten=False,
    filters: None | Clause | list | dict = None,
) -> list | dict:
    """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 = {}
        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 collections in the database.

Returns:

Name	Type	Description
list		List of collection information dictionaries

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

def get_collections(self):
    """Get all collections in the database.

    Returns:
        list: List of collection information dictionaries
    """
    return self.conn.collections()

init_db(schema, clean_start)

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.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing graph structure definitions	required
clean_start		If True, delete all existing collections before initialization	required

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

def init_db(self, schema: Schema, clean_start):
    """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.

    Args:
        schema: Schema containing graph structure definitions
        clean_start: If True, delete all existing collections before initialization
    """
    # 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:
        if clean_start:
            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 clean_start 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 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=None, collection_name=None, match_keys_source=('_key',), match_keys_target=('_key',), filter_uniques=True, uniq_weight_fields=None, uniq_weight_collections=None, upsert_option=False, 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 of edge documents in format [{_source_aux: source_doc, _target_aux: target_doc}]	required
source_class		Source vertex collection name	required
target_class		Target vertex collection name	required
relation_name		Optional relation name for the edges	None
collection_name		Edge collection name	None
match_keys_source		Keys to match source vertices	('_key',)
match_keys_target		Keys to match target vertices	('_key',)
filter_uniques		If True, filter duplicate edges	True
uniq_weight_fields		Fields to consider for uniqueness	None
uniq_weight_collections		Collections to consider for uniqueness	None
upsert_option		If True, use upsert instead of insert	False
head		Optional limit on number of edges to insert	None
**kwargs		Additional options: - dry: If True, don't execute the query	{}

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

def insert_edges_batch(
    self,
    docs_edges,
    source_class,
    target_class,
    relation_name=None,
    collection_name=None,
    match_keys_source=("_key",),
    match_keys_target=("_key",),
    filter_uniques=True,
    uniq_weight_fields=None,
    uniq_weight_collections=None,
    upsert_option=False,
    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.

    Args:
        docs_edges: List of edge documents in format [{_source_aux: source_doc, _target_aux: target_doc}]
        source_class: Source vertex collection name
        target_class: Target vertex collection name
        relation_name: Optional relation name for the edges
        collection_name: Edge collection name
        match_keys_source: Keys to match source vertices
        match_keys_target: Keys to match target vertices
        filter_uniques: If True, filter duplicate edges
        uniq_weight_fields: Fields to consider for uniqueness
        uniq_weight_collections: Collections to consider for uniqueness
        upsert_option: If True, use upsert instead of insert
        head: Optional limit on number of edges to insert
        **kwargs: Additional options:
            - dry: If True, don't execute the query
    """
    dry = kwargs.pop("dry", 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 their keys.

Parameters:

Name	Type	Description	Default
docs		Documents to insert	required
class_name		Collection to insert into	required

Returns:

Name	Type	Description
str		AQL query string for the operation

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

def insert_return_batch(self, docs, class_name):
    """Insert documents and return their keys.

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

    Returns:
        str: AQL query string for the operation
    """
    docs = json.dumps(docs)
    query0 = f"""FOR doc in {docs}
          INSERT doc
          INTO {class_name}
          LET inserted = NEW
          RETURN {{_key: inserted._key}}
    """
    return query0

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

Keep documents that don't exist in the database.

Parameters:

Name	Type	Description	Default
batch		Batch of documents to check	required
class_name		Collection to check in	required
match_keys		Keys to match documents	required
keep_keys		Keys to keep in result	required
filters	None | Clause | list | dict	Additional query filters	None

Returns:

Name	Type	Description
list		Documents that don't exist in the database

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

def keep_absent_documents(
    self,
    batch,
    class_name,
    match_keys,
    keep_keys,
    filters: None | Clause | list | dict = None,
):
    """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		Collection to update	required
field		Field to convert to numeric	required

Returns:

Name	Type	Description
str		AQL query string for the operation

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

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

    Args:
        collection_name: Collection to update
        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=None, **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 of documents to upsert	required
class_name		Collection name to upsert into	required
match_keys	list[str] | None	Keys to match for upsert operation	None
**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	{}

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

def upsert_docs_batch(
    self,
    docs,
    class_name,
    match_keys: list[str] | None = None,
    **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.

    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 isinstance(docs, list):
        if filter_uniques:
            docs = pick_unique_dict(docs)
        docs = json.dumps(docs, default=_json_serializer)
    if match_keys is None:
        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}
                        {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.

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.
    """

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

    @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, **kwargs):
        """Execute a database query.

        Args:
            query: Query to execute
            **kwargs: Additional query parameters
        """
        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 collections based on the schema.

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

    @abc.abstractmethod
    def delete_graph_structure(self, vertex_types=(), graph_names=(), delete_all=False):
        """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, clean_start):
        """Initialize the database with the given schema.

        Args:
            schema: Schema to initialize the database with
            clean_start: Whether to clean existing data
        """
        pass

    @abc.abstractmethod
    def upsert_docs_batch(self, docs, class_name, match_keys, **kwargs):
        """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,
        source_class,
        target_class,
        relation_name,
        collection_name,
        match_keys_source,
        match_keys_target,
        filter_uniques=True,
        uniq_weight_fields=None,
        uniq_weight_collections=None,
        upsert_option=False,
        head=None,
        **kwargs,
    ):
        """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
            collection_name: Name of the edge type (database-specific: collection/relationship type)
            match_keys_source: Keys to match source vertices
            match_keys_target: Keys to match target vertices
            filter_uniques: Whether to filter unique edges
            uniq_weight_fields: Fields to consider for uniqueness
            uniq_weight_collections: Vertex/edge types to consider for uniqueness (database-specific)
            upsert_option: Whether to upsert existing edges
            head: Optional head document
            **kwargs: Additional insertion parameters
        """
        pass

    @abc.abstractmethod
    def insert_return_batch(self, docs, class_name):
        """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: Inserted documents
        """
        pass

    @abc.abstractmethod
    def fetch_docs(
        self,
        class_name,
        filters,
        limit,
        return_keys,
        unset_keys,
        **kwargs,
    ):
        """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 | dict | None = None,
        limit: int | None = None,
        return_keys: list | None = None,
        unset_keys: list | None = None,
        **kwargs,
    ):
        """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,
        class_name,
        match_keys,
        keep_keys,
        flatten=False,
        filters: list | dict | None = None,
    ):
        """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
            filters: Additional query filters

        Returns:
            list: Documents that exist in the database
        """
        pass

    @abc.abstractmethod
    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 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:
            dict: Aggregation results
        """
        pass

    @abc.abstractmethod
    def keep_absent_documents(
        self,
        batch,
        class_name,
        match_keys,
        keep_keys,
        filters: list | dict | None = None,
    ):
        """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 collections.

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

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

        Args:
            edges: List of edge configurations containing index definitions
        """
        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		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 | dict | None	Query filters	None

Returns:

Name	Type	Description
dict		Aggregation results

Source code in graflo/db/conn.py

@abc.abstractmethod
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 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:
        dict: Aggregation results
    """
    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_indices(edges) abstractmethod

Define indices for edge collections.

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 collections.

    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 collections based on the schema.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing collection definitions	required

Source code in graflo/db/conn.py

@abc.abstractmethod
def define_schema(self, schema: Schema):
    """Define collections based on the schema.

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

define_vertex_indices(vertex_config) abstractmethod

Define indices for vertex collections.

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 collections.

    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		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	False

Source code in graflo/db/conn.py

@abc.abstractmethod
def delete_graph_structure(self, vertex_types=(), graph_names=(), delete_all=False):
    """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		Query to execute	required
**kwargs		Additional query parameters	{}

Source code in graflo/db/conn.py

@abc.abstractmethod
def execute(self, query, **kwargs):
    """Execute a database query.

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

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

Fetch documents from a vertex type.

Parameters:

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

Returns:

Name	Type	Description
list		Fetched documents

Source code in graflo/db/conn.py

@abc.abstractmethod
def fetch_docs(
    self,
    class_name,
    filters,
    limit,
    return_keys,
    unset_keys,
    **kwargs,
):
    """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 | 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)	None
**kwargs		Additional database-specific parameters	{}

Returns:

Name	Type	Description
list		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 | dict | None = None,
    limit: int | None = None,
    return_keys: list | None = None,
    unset_keys: list | None = None,
    **kwargs,
):
    """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, flatten=False, filters=None) abstractmethod

Fetch documents that exist in the database.

Parameters:

Name	Type	Description	Default
batch		Batch of documents to check	required
class_name		Name of the collection	required
match_keys		Keys to match	required
keep_keys		Keys to keep in result	required
flatten		Whether to flatten the result	False
filters	list | dict | None	Additional query filters	None

Returns:

Name	Type	Description
list		Documents that exist in the database

Source code in graflo/db/conn.py

@abc.abstractmethod
def fetch_present_documents(
    self,
    batch,
    class_name,
    match_keys,
    keep_keys,
    flatten=False,
    filters: list | dict | None = None,
):
    """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
        filters: Additional query filters

    Returns:
        list: Documents that exist in the database
    """
    pass

init_db(schema, clean_start) abstractmethod

Initialize the database with the given schema.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema to initialize the database with	required
clean_start		Whether to clean existing data	required

Source code in graflo/db/conn.py

@abc.abstractmethod
def init_db(self, schema: Schema, clean_start):
    """Initialize the database with the given schema.

    Args:
        schema: Schema to initialize the database with
        clean_start: Whether to clean existing data
    """
    pass

insert_edges_batch(docs_edges, source_class, target_class, relation_name, collection_name, match_keys_source, match_keys_target, filter_uniques=True, uniq_weight_fields=None, uniq_weight_collections=None, upsert_option=False, head=None, **kwargs) abstractmethod

Insert a batch of edges.

Parameters:

Name	Type	Description	Default
docs_edges		Edge documents to insert	required
source_class		Source vertex type/class	required
target_class		Target vertex type/class	required
relation_name		Name of the edge type/relation	required
collection_name		Name of the edge type (database-specific: collection/relationship type)	required
match_keys_source		Keys to match source vertices	required
match_keys_target		Keys to match target vertices	required
filter_uniques		Whether to filter unique edges	True
uniq_weight_fields		Fields to consider for uniqueness	None
uniq_weight_collections		Vertex/edge types to consider for uniqueness (database-specific)	None
upsert_option		Whether to upsert existing edges	False
head		Optional head document	None
**kwargs		Additional insertion parameters	{}

Source code in graflo/db/conn.py

@abc.abstractmethod
def insert_edges_batch(
    self,
    docs_edges,
    source_class,
    target_class,
    relation_name,
    collection_name,
    match_keys_source,
    match_keys_target,
    filter_uniques=True,
    uniq_weight_fields=None,
    uniq_weight_collections=None,
    upsert_option=False,
    head=None,
    **kwargs,
):
    """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
        collection_name: Name of the edge type (database-specific: collection/relationship type)
        match_keys_source: Keys to match source vertices
        match_keys_target: Keys to match target vertices
        filter_uniques: Whether to filter unique edges
        uniq_weight_fields: Fields to consider for uniqueness
        uniq_weight_collections: Vertex/edge types to consider for uniqueness (database-specific)
        upsert_option: Whether to upsert existing edges
        head: Optional head document
        **kwargs: Additional insertion parameters
    """
    pass

insert_return_batch(docs, class_name) abstractmethod

Insert documents and return the inserted documents.

Parameters:

Name	Type	Description	Default
docs		Documents to insert	required
class_name		Name of the vertex type (or collection/label in database-specific terms)	required

Returns:

Name	Type	Description
list		Inserted documents

Source code in graflo/db/conn.py

@abc.abstractmethod
def insert_return_batch(self, docs, class_name):
    """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: Inserted documents
    """
    pass

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

Keep documents that don't exist in the database.

Parameters:

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

Returns:

Name	Type	Description
list		Documents that don't exist in the database

Source code in graflo/db/conn.py

@abc.abstractmethod
def keep_absent_documents(
    self,
    batch,
    class_name,
    match_keys,
    keep_keys,
    filters: list | dict | None = None,
):
    """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		Documents to upsert	required
class_name		Name of the vertex type (or collection/label in database-specific terms)	required
match_keys		Keys to match for upsert	required
**kwargs		Additional upsert parameters	{}

Source code in graflo/db/conn.py

@abc.abstractmethod
def upsert_docs_batch(self, docs, class_name, match_keys, **kwargs):
    """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,
    }

    # 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 _add_default_port_to_uri(self):
        """Add default port to URI if missing."""
        if self.uri is None:
            return self

        parsed = urlparse(self.uri)
        if parsed.port is not None:
            return self

        # Add default port
        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

    @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
        config_class = conn_type.config_class
        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) -> T:
        """Load config from environment variables using Pydantic BaseSettings.

        Supports custom prefixes for multiple configs:
        - Default (prefix=None): Uses {BASE_PREFIX}URI, {BASE_PREFIX}USERNAME, etc.
        - With prefix (prefix="USER"): Uses USER_{BASE_PREFIX}URI, USER_{BASE_PREFIX}USERNAME, etc.

        Args:
            prefix: Optional prefix for environment variables (e.g., "USER", "LAKE", "KG").
                   If None, uses default {BASE_PREFIX}* variables.

        Returns:
            DBConfig instance loaded from environment variables using Pydantic BaseSettings

        Examples:
            # Load default config (ARANGO_URI, ARANGO_USERNAME, etc.)
            config = ArangoConfig.from_env()

            # Load config with prefix (USER_ARANGO_URI, USER_ARANGO_USERNAME, etc.)
            user_config = ArangoConfig.from_env(prefix="USER")
        """
        if prefix:
            # Get the base prefix from the class's model_config
            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"
                )
            # Create a new model class with modified env_prefix
            new_prefix = f"{prefix.upper()}_{base_prefix}"
            case_sensitive = cls.model_config.get("case_sensitive", False)
            model_config = SettingsConfigDict(
                env_prefix=new_prefix,
                case_sensitive=case_sensitive,
            )
            # Create a new class dynamically with the modified prefix
            temp_class = type(
                f"{cls.__name__}WithPrefix", (cls,), {"model_config": model_config}
            )
            return temp_class()
        else:
            # Use default prefix - Pydantic will read from environment automatically
            return cls()

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
    config_class = conn_type.config_class
    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) classmethod

Load config from environment variables using Pydantic BaseSettings.

Supports custom prefixes for multiple configs: - Default (prefix=None): Uses {BASE_PREFIX}URI, {BASE_PREFIX}USERNAME, etc. - With prefix (prefix="USER"): Uses USER_{BASE_PREFIX}URI, USER_{BASE_PREFIX}USERNAME, etc.

Parameters:

Name	Type	Description	Default
prefix	str | None	Optional prefix for environment variables (e.g., "USER", "LAKE", "KG"). If None, uses default {BASE_PREFIX}* variables.	None

Returns:

Type	Description
T	DBConfig instance loaded from environment variables using Pydantic BaseSettings

Examples:

Load default config (ARANGO_URI, ARANGO_USERNAME, etc.)

config = ArangoConfig.from_env()

Load config with prefix (USER_ARANGO_URI, USER_ARANGO_USERNAME, etc.)

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) -> T:
    """Load config from environment variables using Pydantic BaseSettings.

    Supports custom prefixes for multiple configs:
    - Default (prefix=None): Uses {BASE_PREFIX}URI, {BASE_PREFIX}USERNAME, etc.
    - With prefix (prefix="USER"): Uses USER_{BASE_PREFIX}URI, USER_{BASE_PREFIX}USERNAME, etc.

    Args:
        prefix: Optional prefix for environment variables (e.g., "USER", "LAKE", "KG").
               If None, uses default {BASE_PREFIX}* variables.

    Returns:
        DBConfig instance loaded from environment variables using Pydantic BaseSettings

    Examples:
        # Load default config (ARANGO_URI, ARANGO_USERNAME, etc.)
        config = ArangoConfig.from_env()

        # Load config with prefix (USER_ARANGO_URI, USER_ARANGO_USERNAME, etc.)
        user_config = ArangoConfig.from_env(prefix="USER")
    """
    if prefix:
        # Get the base prefix from the class's model_config
        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"
            )
        # Create a new model class with modified env_prefix
        new_prefix = f"{prefix.upper()}_{base_prefix}"
        case_sensitive = cls.model_config.get("case_sensitive", False)
        model_config = SettingsConfigDict(
            env_prefix=new_prefix,
            case_sensitive=case_sensitive,
        )
        # Create a new class dynamically with the modified prefix
        temp_class = type(
            f"{cls.__name__}WithPrefix", (cls,), {"model_config": model_config}
        )
        return temp_class()
    else:
        # Use default prefix - Pydantic will read from environment automatically
        return cls()

DBType

Bases: StrEnum

Enum representing different types of databases.

Includes both graph databases and source databases (SQL, NoSQL, etc.).

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

class DBType(StrEnum, metaclass=EnumMetaWithContains):
    """Enum representing different types of databases.

    Includes both graph databases and source databases (SQL, NoSQL, etc.).
    """

    # Graph databases
    ARANGO = "arango"
    NEO4J = "neo4j"
    TIGERGRAPH = "tigergraph"
    FALKORDB = "falkordb"

    # Source databases (SQL, NoSQL)
    POSTGRES = "postgres"
    MYSQL = "mysql"
    MONGODB = "mongodb"
    SQLITE = "sqlite"

    @property
    def config_class(self) -> Type["DBConfig"]:
        """Get the appropriate config class for this database type."""
        from .config_mapping import DB_TYPE_MAPPING

        return DB_TYPE_MAPPING[self]

config_class property

Get the appropriate config class for this database type.

FalkordbConnection

Bases: Connection

FalkorDB connector implementing the graflo Connection interface.

Provides complete graph database operations for FalkorDB 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 redis.exceptions.ResponseError
• Invalid inputs raise ValueError with descriptive messages

Attributes

flavor : DBFlavor Database type identifier (DBFlavor.FALKORDB) config : FalkordbConfig Connection configuration (URI, database, credentials) client : FalkorDB Underlying FalkorDB client instance graph : Graph Active graph object for query execution _graph_name : str Name of the currently selected graph

Examples

Direct instantiation (prefer ConnectionManager for production)::

config = FalkordbConfig(uri="redis://localhost:6379")
conn = FalkordbConnection(config)
try:
    result = conn.execute("MATCH (n) RETURN count(n)")
finally:
    conn.close()

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

class FalkordbConnection(Connection):
    """FalkorDB connector implementing the graflo Connection interface.

    Provides complete graph database operations for FalkorDB 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 redis.exceptions.ResponseError
    - Invalid inputs raise ValueError with descriptive messages

    Attributes
    ----------
    flavor : DBFlavor
        Database type identifier (DBFlavor.FALKORDB)
    config : FalkordbConfig
        Connection configuration (URI, database, credentials)
    client : FalkorDB
        Underlying FalkorDB client instance
    graph : Graph
        Active graph object for query execution
    _graph_name : str
        Name of the currently selected graph

    Examples
    --------
    Direct instantiation (prefer ConnectionManager for production)::

        config = FalkordbConfig(uri="redis://localhost:6379")
        conn = FalkordbConnection(config)
        try:
            result = conn.execute("MATCH (n) RETURN count(n)")
        finally:
            conn.close()
    """

    flavor = DBFlavor.FALKORDB

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

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

        Establishes connection to the FalkorDB instance and selects
        the specified graph for subsequent operations.

        Parameters
        ----------
        config : FalkordbConfig
            Connection configuration with the following fields:
            - uri: Redis URI (redis://host:port)
            - database: Graph name (optional, defaults to "default")
            - password: Redis password (optional)

        Raises
        ------
        ValueError
            If URI is not provided in configuration
        redis.exceptions.ConnectionError
            If unable to connect to Redis instance
        """
        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.

        Executes the provided Cypher query with optional parameters.
        Parameters are safely injected using FalkorDB's parameterized
        query mechanism to prevent injection attacks.

        Parameters
        ----------
        query : str
            OpenCypher query string. Can include parameter placeholders
            using $name syntax (e.g., "MATCH (n) WHERE n.id = $id")
        **kwargs
            Query parameters as keyword arguments. Values are safely
            escaped by the driver.

        Returns
        -------
        QueryResult
            FalkorDB result object containing:
            - result_set: List of result rows
            - statistics: Query execution statistics

        Examples
        --------
        Simple query::

            result = conn.execute("MATCH (n:Person) RETURN n.name")

        Parameterized query::

            result = conn.execute(
                "MATCH (n:Person) WHERE n.age > $min_age RETURN n",
                min_age=21
            )
        """
        assert self.graph is not None, "Connection is closed"
        result = self.graph.query(query, kwargs if kwargs else None)
        return result

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

        Note: FalkorDB client uses Redis connection pooling,
        so explicit close is not always necessary.
        """
        # 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.

        FalkorDB (like most databases) cannot store special float values.
        This method rejects values that would cause query failures.

        Parameters
        ----------
        value : Any
            Value to validate

        Returns
        -------
        bool
            True if value can be safely stored, False otherwise

        Notes
        -----
        Rejected values:
        - float('nan'): Not a Number
        - float('inf'): Positive infinity
        - float('-inf'): Negative infinity
        """
        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.

        Null bytes (\\x00) cause FalkorDB's Cypher parser to fail with
        cryptic errors. This method strips them from string values.

        Parameters
        ----------
        value : str
            String value to sanitize

        Returns
        -------
        str
            Sanitized string with problematic characters removed

        Notes
        -----
        Currently handles:
        - Null bytes (\\x00): Break Cypher parser tokenization
        """
        if "\x00" in value:
            value = value.replace("\x00", "")
        return value

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

        Performs comprehensive input validation and sanitization to ensure
        documents can be safely inserted without query errors or injection.

        Sanitization Steps
        ------------------
        1. Filter non-string property keys (log warning)
        2. Remove properties with invalid float values (NaN, Inf)
        3. Strip null bytes from string values
        4. Validate presence of required match keys

        Parameters
        ----------
        doc : dict
            Document to sanitize. Modified values are logged as warnings.
        match_keys : list[str], optional
            Keys that must be present with valid (non-None) values.
            Typically the fields used for MERGE matching.

        Returns
        -------
        dict
            Sanitized copy of the document

        Raises
        ------
        ValueError
            If a required match_key is missing or has None value

        Examples
        --------
        >>> doc = {"id": "1", "name": "test\\x00", 123: "bad_key"}
        >>> sanitized = conn._sanitize_document(doc, match_keys=["id"])
        # Logs: Skipping non-string property key: 123
        # Logs: Sanitized property 'name': removed null bytes
        >>> sanitized
        {"id": "1", "name": "test"}
        """
        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

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

            # Sanitize string values (remove null bytes that break Cypher)
            if isinstance(value, str):
                original = value
                value = self._sanitize_string_value(value)
                if value != original:
                    logger.warning(
                        f"Sanitized property '{key}': removed null bytes from value"
                    )

            sanitized[key] = 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[dict]: 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.

        In FalkorDB, creating a database means selecting a new graph.
        The graph is 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.

        FalkorDB uses CREATE INDEX syntax similar to Neo4j but with some differences.

        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 collections based on schema.

        Note: This is a no-op in FalkorDB as collections are implicit.
        Labels and relationship types are created when data is inserted.

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

    def define_vertex_collections(self, schema: Schema):
        """Define vertex collections based on schema.

        Note: This is a no-op in FalkorDB as vertex collections are implicit.

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

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

        Note: This is a no-op in FalkorDB as edge collections are implicit.

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

    def delete_graph_structure(self, vertex_types=(), graph_names=(), delete_all=False):
        """Delete graph structure (nodes and relationships) from FalkorDB.

        In FalkorDB:
        - Labels: Categories for nodes (equivalent to vertex types)
        - Relationship Types: Types of relationships (equivalent to edge types)
        - Graph: Redis key containing all nodes and relationships

        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, clean_start: bool):
        """Initialize FalkorDB with the given schema.

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

        Args:
            schema: Schema containing graph structure definitions
            clean_start: If True, delete all existing data before initialization
        """
        # 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}'")

        if clean_start:
            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"Clean start 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 upsert_docs_batch(
        self, docs: list[dict], class_name: str, match_keys: list[str], **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.

        The operation:
        1. Sanitizes all documents (removes invalid keys/values)
        2. For each document, attempts to MERGE on match_keys
        3. If node exists: updates all properties
        4. If node doesn't exist: creates with all properties

        Parameters
        ----------
        docs : list[dict]
            Documents to upsert. Each document must contain all match_keys.
        class_name : str
            Node label (e.g., "Person", "Product")
        match_keys : list[str]
            Properties used to identify existing nodes. These form the
            MERGE pattern: ``MERGE (n:Label {key1: val1, key2: val2})``
        **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

        Examples
        --------
        Insert or update users by email::

            docs = [
                {"email": "alice@example.com", "name": "Alice", "age": 30},
                {"email": "bob@example.com", "name": "Bob", "age": 25}
            ]
            conn.upsert_docs_batch(docs, "User", match_keys=["email"])

        Notes
        -----
        The generated Cypher query uses UNWIND for batch efficiency::

            UNWIND $batch AS row
            MERGE (n:Label {match_key: row.match_key})
            ON MATCH SET n += row
            ON CREATE SET n += row
        """
        dry = kwargs.pop("dry", False)

        if not docs:
            return

        # Sanitize documents: filter invalid keys/values, validate match_keys
        sanitized_docs = self._sanitize_batch(docs, match_keys)

        if not sanitized_docs:
            return

        # Build the MERGE clause with match keys
        index_str = ", ".join([f"{k}: row.{k}" for k in match_keys])
        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:
            self.execute(q, batch=sanitized_docs)

    def insert_edges_batch(
        self,
        docs_edges: list,
        source_class: str,
        target_class: str,
        relation_name: str,
        collection_name: str | None = None,
        match_keys_source: tuple[str, ...] = ("_key",),
        match_keys_target: tuple[str, ...] = ("_key",),
        filter_uniques: bool = True,
        uniq_weight_fields=None,
        uniq_weight_collections=None,
        upsert_option: bool = False,
        head: int | None = None,
        **kwargs,
    ):
        """Create relationships between existing nodes using Cypher MERGE.

        Efficiently creates relationships in batch by matching source and
        target nodes, then creating or updating the relationship between them.

        Parameters
        ----------
        docs_edges : list
            Edge specifications as list of [source, target, props] triples:
            ``[[{source_props}, {target_props}, {edge_props}], ...]``
        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
            Unused in FalkorDB (kept for interface compatibility)
        match_keys_source : tuple[str, ...]
            Properties to match source nodes (default: ("_key",))
        match_keys_target : tuple[str, ...]
            Properties to match target nodes (default: ("_key",))
        filter_uniques : bool
            Unused in FalkorDB (kept for interface compatibility)
        uniq_weight_fields
            Unused in FalkorDB (kept for interface compatibility)
        uniq_weight_collections
            Unused in FalkorDB (kept for interface compatibility)
        upsert_option : bool
            Unused in FalkorDB (kept for interface compatibility)
        head : int, optional
            Unused in FalkorDB (kept for interface compatibility)
        **kwargs
            Additional options:
            - dry (bool): If True, build query but don't execute

        Examples
        --------
        Create KNOWS relationships between people::

            edges = [
                [{"id": "1"}, {"id": "2"}, {"since": 2020}],
                [{"id": "1"}, {"id": "3"}, {"since": 2021}]
            ]
            conn.insert_edges_batch(
                edges,
                source_class="Person",
                target_class="Person",
                relation_name="KNOWS",
                match_keys_source=["id"],
                match_keys_target=["id"]
            )

        Notes
        -----
        Generated Cypher pattern::

            UNWIND $batch AS row
            MATCH (source:Label), (target:Label)
            WHERE source.key = row[0].key AND target.key = row[1].key
            MERGE (source)-[r:REL_TYPE]->(target)
            SET r += row[2]
        """
        dry = kwargs.pop("dry", False)

        if not docs_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)

        q = f"""
            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, class_name):
        """Insert nodes and return their properties.

        Note: Limited implementation in FalkorDB.

        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

        Returns:
            list: Fetched nodes as dictionaries
        """
        # Build filter clause
        if filters is not None:
            ff = Expression.from_dict(filters)
            # Use NEO4J flavor since FalkorDB uses OpenCypher
            filter_clause = f"WHERE {ff(doc_name='n', kind=DBFlavor.NEO4J)}"
        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:
            dict: 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: 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 = Expression.from_dict(filters)
            filter_clause = ff(doc_name="r", kind=ExpressionFlavor.NEO4J)
            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:
            dict: 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,
        class_name,
        match_keys,
        keep_keys,
        flatten=False,
        filters: list | dict | None = None,
    ):
        """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:
            list: 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:
            dict or int: Aggregation results
        """
        # Build filter clause
        if filters is not None:
            ff = Expression.from_dict(filters)
            filter_clause = f"WHERE {ff(doc_name='n', kind=DBFlavor.NEO4J)}"
        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,
        class_name,
        match_keys,
        keep_keys,
        filters: list | dict | None = None,
    ):
        """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:
            list: 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 and select graph.

Establishes connection to the FalkorDB instance and selects the specified graph for subsequent operations.

Parameters

config : FalkordbConfig Connection configuration with the following fields: - uri: Redis URI (redis://host:port) - database: Graph name (optional, defaults to "default") - password: Redis password (optional)

Raises

ValueError If URI is not provided in configuration redis.exceptions.ConnectionError If unable to connect to Redis instance

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

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

    Establishes connection to the FalkorDB instance and selects
    the specified graph for subsequent operations.

    Parameters
    ----------
    config : FalkordbConfig
        Connection configuration with the following fields:
        - uri: Redis URI (redis://host:port)
        - database: Graph name (optional, defaults to "default")
        - password: Redis password (optional)

    Raises
    ------
    ValueError
        If URI is not provided in configuration
    redis.exceptions.ConnectionError
        If unable to connect to Redis instance
    """
    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
	dict or int: Aggregation results

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:
        dict or int: Aggregation results
    """
    # Build filter clause
    if filters is not None:
        ff = Expression.from_dict(filters)
        filter_clause = f"WHERE {ff(doc_name='n', kind=DBFlavor.NEO4J)}"
    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}"
        )

close()

Close the FalkorDB connection.

Note: FalkorDB client uses Redis connection pooling, so explicit close is not always necessary.

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

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

    Note: FalkorDB client uses Redis connection pooling,
    so explicit close is not always necessary.
    """
    # 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.

In FalkorDB, creating a database means selecting a new graph. The graph is 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.

    In FalkorDB, creating a database means selecting a new graph.
    The graph is 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_collections(edges)

Define edge collections based on schema.

Note: This is a no-op in FalkorDB as edge collections 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_collections(self, edges: list[Edge]):
    """Define edge collections based on schema.

    Note: This is a no-op in FalkorDB as edge collections 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 collections based on schema.

Note: This is a no-op in FalkorDB as collections are implicit. Labels and relationship types are created when data is inserted.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing collection definitions	required

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

def define_schema(self, schema: Schema):
    """Define collections based on schema.

    Note: This is a no-op in FalkorDB as collections are implicit.
    Labels and relationship types are created when data is inserted.

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

define_vertex_collections(schema)

Define vertex collections based on schema.

Note: This is a no-op in FalkorDB as vertex collections 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_collections(self, schema: Schema):
    """Define vertex collections based on schema.

    Note: This is a no-op in FalkorDB as vertex collections 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.

In FalkorDB: - Labels: Categories for nodes (equivalent to vertex types) - Relationship Types: Types of relationships (equivalent to edge types) - Graph: Redis key containing all nodes and relationships

Parameters:

Name	Type	Description	Default
vertex_types		Label names to delete nodes for	()
graph_names		Graph names to delete entirely	()
delete_all		If True, delete all nodes and relationships	False

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

def delete_graph_structure(self, vertex_types=(), graph_names=(), delete_all=False):
    """Delete graph structure (nodes and relationships) from FalkorDB.

    In FalkorDB:
    - Labels: Categories for nodes (equivalent to vertex types)
    - Relationship Types: Types of relationships (equivalent to edge types)
    - Graph: Redis key containing all nodes and relationships

    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.

Executes the provided Cypher query with optional parameters. Parameters are safely injected using FalkorDB's parameterized query mechanism to prevent injection attacks.

Parameters

query : str OpenCypher query string. Can include parameter placeholders using $name syntax (e.g., "MATCH (n) WHERE n.id = $id") **kwargs Query parameters as keyword arguments. Values are safely escaped by the driver.

Returns

QueryResult FalkorDB result object containing: - result_set: List of result rows - statistics: Query execution statistics

Examples

Simple query::

result = conn.execute("MATCH (n:Person) RETURN n.name")

Parameterized query::

result = conn.execute(
    "MATCH (n:Person) WHERE n.age > $min_age RETURN n",
    min_age=21
)

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

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

    Executes the provided Cypher query with optional parameters.
    Parameters are safely injected using FalkorDB's parameterized
    query mechanism to prevent injection attacks.

    Parameters
    ----------
    query : str
        OpenCypher query string. Can include parameter placeholders
        using $name syntax (e.g., "MATCH (n) WHERE n.id = $id")
    **kwargs
        Query parameters as keyword arguments. Values are safely
        escaped by the driver.

    Returns
    -------
    QueryResult
        FalkorDB result object containing:
        - result_set: List of result rows
        - statistics: Query execution statistics

    Examples
    --------
    Simple query::

        result = conn.execute("MATCH (n:Person) RETURN n.name")

    Parameterized query::

        result = conn.execute(
            "MATCH (n:Person) WHERE n.age > $min_age RETURN n",
            min_age=21
        )
    """
    assert self.graph is not None, "Connection is closed"
    result = self.graph.query(query, kwargs if kwargs else None)
    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

Returns:

Name	Type	Description
list		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

    Returns:
        list: Fetched nodes as dictionaries
    """
    # Build filter clause
    if filters is not None:
        ff = Expression.from_dict(filters)
        # Use NEO4J flavor since FalkorDB uses OpenCypher
        filter_clause = f"WHERE {ff(doc_name='n', kind=DBFlavor.NEO4J)}"
    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:

Name	Type	Description
list		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: 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 = Expression.from_dict(filters)
        filter_clause = ff(doc_name="r", kind=ExpressionFlavor.NEO4J)
        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, flatten=False, filters=None)

Fetch nodes that exist in the database.

Parameters:

Name	Type	Description	Default
batch		Batch of documents to check	required
class_name		Label to check in	required
match_keys		Keys to match nodes	required
keep_keys		Keys to keep in result	required
flatten		Unused in FalkorDB	False
filters	list | dict | None	Additional query filters	None

Returns:

Name	Type	Description
list		Documents that exist in the database

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

def fetch_present_documents(
    self,
    batch,
    class_name,
    match_keys,
    keep_keys,
    flatten=False,
    filters: list | dict | None = None,
):
    """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:
        list: 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, clean_start)

Initialize FalkorDB with the given schema.

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

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing graph structure definitions	required
clean_start	bool	If True, delete all existing data before initialization	required

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

def init_db(self, schema: Schema, clean_start: bool):
    """Initialize FalkorDB with the given schema.

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

    Args:
        schema: Schema containing graph structure definitions
        clean_start: If True, delete all existing data before initialization
    """
    # 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}'")

    if clean_start:
        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"Clean start 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, collection_name=None, match_keys_source=('_key',), match_keys_target=('_key',), filter_uniques=True, uniq_weight_fields=None, uniq_weight_collections=None, upsert_option=False, head=None, **kwargs)

Create relationships between existing nodes using Cypher MERGE.

Efficiently creates relationships in batch by matching source and target nodes, then creating or updating the relationship between them.

Parameters

docs_edges : list Edge specifications as list of [source, target, props] triples: [[{source_props}, {target_props}, {edge_props}], ...] 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 Unused in FalkorDB (kept for interface compatibility) match_keys_source : tuple[str, ...] Properties to match source nodes (default: ("_key",)) match_keys_target : tuple[str, ...] Properties to match target nodes (default: ("_key",)) filter_uniques : bool Unused in FalkorDB (kept for interface compatibility) uniq_weight_fields Unused in FalkorDB (kept for interface compatibility) uniq_weight_collections Unused in FalkorDB (kept for interface compatibility) upsert_option : bool Unused in FalkorDB (kept for interface compatibility) head : int, optional Unused in FalkorDB (kept for interface compatibility) **kwargs Additional options: - dry (bool): If True, build query but don't execute

Examples

Create KNOWS relationships between people::

edges = [
    [{"id": "1"}, {"id": "2"}, {"since": 2020}],
    [{"id": "1"}, {"id": "3"}, {"since": 2021}]
]
conn.insert_edges_batch(
    edges,
    source_class="Person",
    target_class="Person",
    relation_name="KNOWS",
    match_keys_source=["id"],
    match_keys_target=["id"]
)

Notes

Generated Cypher pattern::

UNWIND $batch AS row
MATCH (source:Label), (target:Label)
WHERE source.key = row[0].key AND target.key = row[1].key
MERGE (source)-[r:REL_TYPE]->(target)
SET r += row[2]

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

def insert_edges_batch(
    self,
    docs_edges: list,
    source_class: str,
    target_class: str,
    relation_name: str,
    collection_name: str | None = None,
    match_keys_source: tuple[str, ...] = ("_key",),
    match_keys_target: tuple[str, ...] = ("_key",),
    filter_uniques: bool = True,
    uniq_weight_fields=None,
    uniq_weight_collections=None,
    upsert_option: bool = False,
    head: int | None = None,
    **kwargs,
):
    """Create relationships between existing nodes using Cypher MERGE.

    Efficiently creates relationships in batch by matching source and
    target nodes, then creating or updating the relationship between them.

    Parameters
    ----------
    docs_edges : list
        Edge specifications as list of [source, target, props] triples:
        ``[[{source_props}, {target_props}, {edge_props}], ...]``
    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
        Unused in FalkorDB (kept for interface compatibility)
    match_keys_source : tuple[str, ...]
        Properties to match source nodes (default: ("_key",))
    match_keys_target : tuple[str, ...]
        Properties to match target nodes (default: ("_key",))
    filter_uniques : bool
        Unused in FalkorDB (kept for interface compatibility)
    uniq_weight_fields
        Unused in FalkorDB (kept for interface compatibility)
    uniq_weight_collections
        Unused in FalkorDB (kept for interface compatibility)
    upsert_option : bool
        Unused in FalkorDB (kept for interface compatibility)
    head : int, optional
        Unused in FalkorDB (kept for interface compatibility)
    **kwargs
        Additional options:
        - dry (bool): If True, build query but don't execute

    Examples
    --------
    Create KNOWS relationships between people::

        edges = [
            [{"id": "1"}, {"id": "2"}, {"since": 2020}],
            [{"id": "1"}, {"id": "3"}, {"since": 2021}]
        ]
        conn.insert_edges_batch(
            edges,
            source_class="Person",
            target_class="Person",
            relation_name="KNOWS",
            match_keys_source=["id"],
            match_keys_target=["id"]
        )

    Notes
    -----
    Generated Cypher pattern::

        UNWIND $batch AS row
        MATCH (source:Label), (target:Label)
        WHERE source.key = row[0].key AND target.key = row[1].key
        MERGE (source)-[r:REL_TYPE]->(target)
        SET r += row[2]
    """
    dry = kwargs.pop("dry", False)

    if not docs_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)

    q = f"""
        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: Limited implementation in FalkorDB.

Parameters:

Name	Type	Description	Default
docs		Documents to insert	required
class_name		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, class_name):
    """Insert nodes and return their properties.

    Note: Limited implementation in FalkorDB.

    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, filters=None)

Keep documents that don't exist in the database.

Parameters:

Name	Type	Description	Default
batch		Batch of documents to check	required
class_name		Label to check in	required
match_keys		Keys to match nodes	required
keep_keys		Keys to keep in result	required
filters	list | dict | None	Additional query filters	None

Returns:

Name	Type	Description
list		Documents that don't exist in the database

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

def keep_absent_documents(
    self,
    batch,
    class_name,
    match_keys,
    keep_keys,
    filters: list | dict | None = None,
):
    """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:
        list: 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.

Performs atomic upsert (update-or-insert) operations on a batch of documents. Uses Cypher MERGE with ON MATCH/ON CREATE for efficiency.

The operation: 1. Sanitizes all documents (removes invalid keys/values) 2. For each document, attempts to MERGE on match_keys 3. If node exists: updates all properties 4. If node doesn't exist: creates with all properties

Parameters

docs : list[dict] Documents to upsert. Each document must contain all match_keys. class_name : str Node label (e.g., "Person", "Product") match_keys : list[str] Properties used to identify existing nodes. These form the MERGE pattern: MERGE (n:Label {key1: val1, key2: val2}) **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

Examples

Insert or update users by email::

docs = [
    {"email": "alice@example.com", "name": "Alice", "age": 30},
    {"email": "bob@example.com", "name": "Bob", "age": 25}
]
conn.upsert_docs_batch(docs, "User", match_keys=["email"])

Notes

The generated Cypher query uses UNWIND for batch efficiency::

UNWIND $batch AS row
MERGE (n:Label {match_key: row.match_key})
ON MATCH SET n += row
ON CREATE SET n += row

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

def upsert_docs_batch(
    self, docs: list[dict], class_name: str, match_keys: list[str], **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.

    The operation:
    1. Sanitizes all documents (removes invalid keys/values)
    2. For each document, attempts to MERGE on match_keys
    3. If node exists: updates all properties
    4. If node doesn't exist: creates with all properties

    Parameters
    ----------
    docs : list[dict]
        Documents to upsert. Each document must contain all match_keys.
    class_name : str
        Node label (e.g., "Person", "Product")
    match_keys : list[str]
        Properties used to identify existing nodes. These form the
        MERGE pattern: ``MERGE (n:Label {key1: val1, key2: val2})``
    **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

    Examples
    --------
    Insert or update users by email::

        docs = [
            {"email": "alice@example.com", "name": "Alice", "age": 30},
            {"email": "bob@example.com", "name": "Bob", "age": 25}
        ]
        conn.upsert_docs_batch(docs, "User", match_keys=["email"])

    Notes
    -----
    The generated Cypher query uses UNWIND for batch efficiency::

        UNWIND $batch AS row
        MERGE (n:Label {match_key: row.match_key})
        ON MATCH SET n += row
        ON CREATE SET n += row
    """
    dry = kwargs.pop("dry", False)

    if not docs:
        return

    # Sanitize documents: filter invalid keys/values, validate match_keys
    sanitized_docs = self._sanitize_batch(docs, match_keys)

    if not sanitized_docs:
        return

    # Build the MERGE clause with match keys
    index_str = ", ".join([f"{k}: row.{k}" for k in match_keys])
    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:
        self.execute(q, batch=sanitized_docs)

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 = DBFlavor.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")
        self._driver = GraphDatabase.driver(
            uri=config.url, auth=(config.username, config.password)
        )
        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
        """
        try:
            self.execute(f"CREATE DATABASE {name}")
            logger.info(f"Successfully created Neo4j database '{name}'")
        except Exception as e:
            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 collections based on schema.

        Note: This is a no-op in Neo4j as collections are implicit.

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

    def define_vertex_collections(self, schema: Schema):
        """Define vertex collections based on schema.

        Note: This is a no-op in Neo4j as vertex collections are implicit.

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

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

        Note: This is a no-op in Neo4j as edge collections are implicit.

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

    def delete_graph_structure(self, 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

        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, clean_start):
        """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.

        Args:
            schema: Schema containing graph structure definitions
            clean_start: If True, delete all existing data before initialization
        """
        # 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, we'll try to create it but it may fail 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..."
                    )
                    try:
                        self.create_database(db_name)
                        logger.info(f"Successfully created database '{db_name}'")
                    except Exception as create_error:
                        logger.info(
                            f"Neo4j Community Edition? Could not create database '{db_name}': {create_error}. "
                            f"This may be Neo4j Community Edition which only supports one database per instance.",
                            exc_info=True,
                        )
                        # Continue with default database for Community Edition
            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}"
                )
                try:
                    self.create_database(db_name)
                    logger.info(f"Successfully created database '{db_name}'")
                except Exception as create_error:
                    logger.info(
                        f"Neo4j Community Edition? Could not create database '{db_name}': {create_error}. "
                        f"This may be Neo4j Community Edition which only supports one database per instance. "
                        f"Continuing with default database.",
                        exc_info=True,
                    )
                    # Continue with default database for Community Edition
        except Exception as e:
            logger.error(
                f"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:
            if clean_start:
                try:
                    self.delete_database("")
                    logger.debug(f"Cleaned database '{db_name}' for fresh start")
                except Exception as clean_error:
                    logger.warning(
                        f"Error during clean_start 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 Exception as e:
            logger.error(
                f"Error during database schema initialization for '{db_name}': {e}",
                exc_info=True,
            )
            raise

    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,
        source_class,
        target_class,
        relation_name,
        collection_name=None,
        match_keys_source=("_key",),
        match_keys_target=("_key",),
        filter_uniques=True,
        uniq_weight_fields=None,
        uniq_weight_collections=None,
        upsert_option=False,
        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.

        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
            collection_name: Unused in Neo4j
            match_keys_source: Keys to match source nodes
            match_keys_target: Keys to match target nodes
            filter_uniques: Unused in Neo4j
            uniq_weight_fields: Unused in Neo4j
            uniq_weight_collections: Unused in Neo4j
            upsert_option: Unused in Neo4j
            head: Optional limit on number of relationships to insert
            **kwargs: Additional options:
                - dry: If True, don't execute the query
        """
        dry = kwargs.pop("dry", False)

        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, class_name):
        """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 = Expression.from_dict(filters)
            filter_clause = f"WHERE {ff(doc_name='n', kind=DBFlavor.NEO4J)}"
        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:
            from graflo.filter.onto import Expression

            ff = Expression.from_dict(filters)
            filter_clause = ff(doc_name="r", kind=ExpressionFlavor.NEO4J)
            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,
        class_name,
        match_keys,
        keep_keys,
        flatten=False,
        filters: list | dict | None = None,
    ):
        """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,
        class_name,
        match_keys,
        keep_keys,
        filters: list | dict | None = None,
    ):
        """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")
    self._driver = GraphDatabase.driver(
        uri=config.url, auth=(config.username, config.password)
    )
    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

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

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
    """
    try:
        self.execute(f"CREATE DATABASE {name}")
        logger.info(f"Successfully created Neo4j database '{name}'")
    except Exception as e:
        raise e

define_edge_collections(edges)

Define edge collections based on schema.

Note: This is a no-op in Neo4j as edge collections 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_collections(self, edges: list[Edge]):
    """Define edge collections based on schema.

    Note: This is a no-op in Neo4j as edge collections 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 collections based on schema.

Note: This is a no-op in Neo4j as collections are implicit.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing collection definitions	required

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

def define_schema(self, schema: Schema):
    """Define collections based on schema.

    Note: This is a no-op in Neo4j as collections are implicit.

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

define_vertex_collections(schema)

Define vertex collections based on schema.

Note: This is a no-op in Neo4j as vertex collections 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_collections(self, schema: Schema):
    """Define vertex collections based on schema.

    Note: This is a no-op in Neo4j as vertex collections 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		Label names to delete nodes for	()
graph_names		Unused in Neo4j (no explicit graph concept)	()
delete_all		If True, delete all nodes and relationships	False

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

def delete_graph_structure(self, 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

    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 = Expression.from_dict(filters)
        filter_clause = f"WHERE {ff(doc_name='n', kind=DBFlavor.NEO4J)}"
    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:
        from graflo.filter.onto import Expression

        ff = Expression.from_dict(filters)
        filter_clause = ff(doc_name="r", kind=ExpressionFlavor.NEO4J)
        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, flatten=False, filters=None)

Fetch nodes that exist in the database.

Note: Not implemented in Neo4j.

Parameters:

Name	Type	Description	Default
batch		Batch of documents to check	required
class_name		Label to check in	required
match_keys		Keys to match nodes	required
keep_keys		Keys to keep in result	required
flatten		Unused in Neo4j	False
filters	list | dict | 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,
    class_name,
    match_keys,
    keep_keys,
    flatten=False,
    filters: list | dict | None = None,
):
    """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, clean_start)

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.

Parameters:

Name	Type	Description	Default
schema	Schema	Schema containing graph structure definitions	required
clean_start		If True, delete all existing data before initialization	required

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

def init_db(self, schema: Schema, clean_start):
    """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.

    Args:
        schema: Schema containing graph structure definitions
        clean_start: If True, delete all existing data before initialization
    """
    # 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, we'll try to create it but it may fail 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..."
                )
                try:
                    self.create_database(db_name)
                    logger.info(f"Successfully created database '{db_name}'")
                except Exception as create_error:
                    logger.info(
                        f"Neo4j Community Edition? Could not create database '{db_name}': {create_error}. "
                        f"This may be Neo4j Community Edition which only supports one database per instance.",
                        exc_info=True,
                    )
                    # Continue with default database for Community Edition
        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}"
            )
            try:
                self.create_database(db_name)
                logger.info(f"Successfully created database '{db_name}'")
            except Exception as create_error:
                logger.info(
                    f"Neo4j Community Edition? Could not create database '{db_name}': {create_error}. "
                    f"This may be Neo4j Community Edition which only supports one database per instance. "
                    f"Continuing with default database.",
                    exc_info=True,
                )
                # Continue with default database for Community Edition
    except Exception as e:
        logger.error(
            f"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:
        if clean_start:
            try:
                self.delete_database("")
                logger.debug(f"Cleaned database '{db_name}' for fresh start")
            except Exception as clean_error:
                logger.warning(
                    f"Error during clean_start 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 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, collection_name=None, match_keys_source=('_key',), match_keys_target=('_key',), filter_uniques=True, uniq_weight_fields=None, uniq_weight_collections=None, upsert_option=False, 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 of edge documents in format [{__source: source_doc, __target: target_doc}]	required
source_class		Source node label	required
target_class		Target node label	required
relation_name		Relationship type name	required
collection_name		Unused in Neo4j	None
match_keys_source		Keys to match source nodes	('_key',)
match_keys_target		Keys to match target nodes	('_key',)
filter_uniques		Unused in Neo4j	True
uniq_weight_fields		Unused in Neo4j	None
uniq_weight_collections		Unused in Neo4j	None
upsert_option		Unused in Neo4j	False
head		Optional limit on number of relationships to insert	None
**kwargs		Additional options: - dry: If True, don't execute the query	{}

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

def insert_edges_batch(
    self,
    docs_edges,
    source_class,
    target_class,
    relation_name,
    collection_name=None,
    match_keys_source=("_key",),
    match_keys_target=("_key",),
    filter_uniques=True,
    uniq_weight_fields=None,
    uniq_weight_collections=None,
    upsert_option=False,
    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.

    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
        collection_name: Unused in Neo4j
        match_keys_source: Keys to match source nodes
        match_keys_target: Keys to match target nodes
        filter_uniques: Unused in Neo4j
        uniq_weight_fields: Unused in Neo4j
        uniq_weight_collections: Unused in Neo4j
        upsert_option: Unused in Neo4j
        head: Optional limit on number of relationships to insert
        **kwargs: Additional options:
            - dry: If True, don't execute the query
    """
    dry = kwargs.pop("dry", False)

    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		Documents to insert	required
class_name		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, class_name):
    """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, filters=None)

Keep nodes that don't exist in the database.

Note: Not implemented in Neo4j.

Parameters:

Name	Type	Description	Default
batch		Batch of documents to check	required
class_name		Label to check in	required
match_keys		Keys to match nodes	required
keep_keys		Keys to keep in result	required
filters	list | dict | 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,
    class_name,
    match_keys,
    keep_keys,
    filters: list | dict | None = None,
):
    """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
            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
                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_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 _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 columns and convert to ColumnInfo
                pk_set = set(pk_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,
                        )
                    )

                # 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 match cache once for all tables (significant performance improvement)
        match_cache = FuzzyMatchCache(vertex_table_names)

        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 columns and convert to ColumnInfo
            pk_set = set(pk_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,
                    )
                )

            # 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, match_cache
                )
            # 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,
                        match_cache,
                    )
                )

                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, match_cache
                    )
                    target_table = infer_vertex_from_column_name(
                        target_column, vertex_table_names, match_cache
                    )

            # 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 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)

        result = SchemaIntrospectionResult(
            vertex_tables=vertex_tables,
            edge_tables=edge_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 match cache once for all tables (significant performance improvement)
    match_cache = FuzzyMatchCache(vertex_table_names)

    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 columns and convert to ColumnInfo
        pk_set = set(pk_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,
                )
            )

        # 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, match_cache
            )
        # 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,
                    match_cache,
                )
            )

            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, match_cache
                )
                target_table = infer_vertex_from_column_name(
                    target_column, vertex_table_names, match_cache
                )

        # 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 columns and convert to ColumnInfo
            pk_set = set(pk_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,
                    )
                )

            # 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
            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_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)

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)

    result = SchemaIntrospectionResult(
        vertex_tables=vertex_tables,
        edge_tables=edge_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 collection creation - vertices and edges must be pre-defined 4. Different query syntax and execution model 5. Token-based authentication for some operations

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 collection creation - vertices and edges must be pre-defined
    4. Different query syntax and execution model
    5. Token-based authentication for some operations
    """

    flavor = DBFlavor.TIGERGRAPH

    def __init__(self, config: TigergraphConfig):
        super().__init__()
        self.config = config
        # Store base URLs for REST++ and GSQL endpoints
        self.restpp_url = f"{config.url_without_port}:{config.port}"
        self.gsql_url = f"{config.url_without_port}:{config.gs_port}"

        # Initialize pyTigerGraph connection for most operations
        # Use type narrowing to help type checker understand non-None values
        # PyTigerGraphConnection has defaults for all parameters, so None values are acceptable
        restpp_port: int | str = config.port if config.port is not None else "9000"
        gs_port: int | str = config.gs_port if config.gs_port is not None else "14240"
        graphname: str = (
            config.database if config.database is not None else "DefaultGraph"
        )
        username: str = config.username if config.username is not None else "tigergraph"
        password: str = config.password if config.password is not None else "tigergraph"
        cert_path: str | None = getattr(config, "certPath", None)

        # Build connection kwargs, only include certPath if it's not None
        conn_kwargs: dict[str, Any] = {
            "host": config.url_without_port,
            "restppPort": restpp_port,
            "gsPort": gs_port,
            "graphname": graphname,
            "username": username,
            "password": password,
        }
        if cert_path is not None:
            conn_kwargs["certPath"] = cert_path

        self.conn = PyTigerGraphConnection(**conn_kwargs)

        # Get authentication token if secret is provided
        if config.secret:
            try:
                self.conn.getToken(config.secret)
            except Exception as e:
                logger.warning(f"Failed to get authentication token: {e}")

    def _get_auth_headers(self) -> dict[str, str]:
        """Get HTTP Basic Auth headers if credentials are available.

        Returns:
            Dictionary with Authorization header if credentials exist
        """
        headers = {}
        if self.config.username and self.config.password:
            import base64

            credentials = f"{self.config.username}:{self.config.password}"
            encoded_credentials = base64.b64encode(credentials.encode()).decode()
            headers["Authorization"] = f"Basic {encoded_credentials}"
        return headers

    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
                )
            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,
                )
            elif method.upper() == "DELETE":
                response = requests.delete(
                    url, headers=headers, params=params, timeout=120
                )
            else:
                raise ValueError(f"Unsupported HTTP method: {method}")

            response.raise_for_status()
            return response.json()

        except requests_exceptions.HTTPError as errh:
            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.

        Updates conn.graphname for PyTigerGraph metadata operations that rely on it
        (e.g., getVertexTypes(), getEdgeTypes()).

        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.conn.graphname
        self.conn.graphname = graph_name

        try:
            yield graph_name
        finally:
            # Restore original graphname
            self.conn.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.conn.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

    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 the pyTigerGraph gsql() method 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:
            Exception: If graph creation fails
        """
        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 pyTigerGraph's gsql method which handles authentication
            logger.debug(f"Creating graph '{name}' via GSQL: {gsql_commands}")
            try:
                result = self.conn.gsql(gsql_commands)
                logger.info(
                    f"Successfully created graph '{name}' with types {all_types}: {result}"
                )
                return result
            except Exception as e:
                error_msg = str(e).lower()
                # Check if graph already exists (might be acceptable)
                if "already exists" in error_msg or "duplicate" in error_msg:
                    logger.info(f"Graph '{name}' may already exist: {e}")
                    return str(e)
                logger.error(f"Failed to create graph '{name}': {e}")
                raise

        except Exception as e:
            logger.error(f"Error creating graph '{name}' via GSQL: {e}")
            raise

    def delete_database(self, name: str):
        """
        Delete a TigerGraph database (graph).

        This method attempts to drop the graph using GSQL DROP GRAPH.
        If that fails (e.g., dependencies), it will:
          1) Remove associations and drop all edge types
          2) Drop all vertex types
          3) Clear remaining data as a last resort

        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}'")
            try:
                # Use the graph first to ensure we're working with the right graph
                drop_command = f"USE GRAPH {name}\nDROP GRAPH {name}"
                result = self.conn.gsql(drop_command)
                logger.info(f"Successfully dropped graph '{name}': {result}")
                return result
            except Exception as e:
                logger.debug(
                    f"Could not drop graph '{name}' (may not exist or have dependencies): {e}"
                )

            # Fallback 1: Attempt to disassociate edge and vertex types from graph
            # DO NOT drop global vertex/edge types as they might be used by other graphs
            try:
                with self._ensure_graph_context(name):
                    # Disassociate edge types from graph (but don't drop them globally)
                    try:
                        edge_types = self.conn.getEdgeTypes(force=True)
                    except Exception:
                        edge_types = []

                    for e_type in edge_types:
                        # Only disassociate from graph, don't drop globally
                        # ALTER GRAPH requires USE GRAPH context
                        try:
                            drop_edge_cmd = f"USE GRAPH {name}\nALTER GRAPH {name} DROP DIRECTED EDGE {e_type}"
                            self.conn.gsql(drop_edge_cmd)
                            logger.debug(
                                f"Disassociated edge type '{e_type}' from graph '{name}'"
                            )
                        except Exception as e:
                            logger.debug(
                                f"Could not disassociate edge type '{e_type}' from graph '{name}': {e}"
                            )
                            # Continue - edge might not be associated or graph might not exist

                    # Disassociate vertex types from graph (but don't drop them globally)
                    try:
                        vertex_types = self.conn.getVertexTypes(force=True)
                    except Exception:
                        vertex_types = []

                    for v_type in vertex_types:
                        # Only clear data from this graph's vertices, don't drop vertex type globally
                        # Clear data first to avoid dependency issues
                        try:
                            self.conn.delVertices(v_type)
                            logger.debug(
                                f"Cleared vertices of type '{v_type}' from graph '{name}'"
                            )
                        except Exception as e:
                            logger.debug(
                                f"Could not clear vertices of type '{v_type}' from graph '{name}': {e}"
                            )
                        # Disassociate from graph (best-effort)
                        # ALTER GRAPH requires USE GRAPH context
                        try:
                            drop_vertex_cmd = f"USE GRAPH {name}\nALTER GRAPH {name} DROP VERTEX {v_type}"
                            self.conn.gsql(drop_vertex_cmd)
                            logger.debug(
                                f"Disassociated vertex type '{v_type}' from graph '{name}'"
                            )
                        except Exception as e:
                            logger.debug(
                                f"Could not disassociate vertex type '{v_type}' from graph '{name}': {e}"
                            )
                            # Continue - vertex might not be associated or graph might not exist
            except Exception as e3:
                logger.warning(
                    f"Could not disassociate schema types from graph '{name}': {e3}. Proceeding to data clear."
                )

            # Fallback 2: Clear all data (if any remain)
            try:
                with self._ensure_graph_context(name):
                    vertex_types = self.conn.getVertexTypes()
                    for v_type in vertex_types:
                        result = self.conn.delVertices(v_type)
                        logger.debug(f"Cleared vertices of type {v_type}: {result}")
                    logger.info(f"Cleared all data from graph '{name}'")
            except Exception as e2:
                logger.warning(
                    f"Could not clear data from graph '{name}': {e2}. Graph may not exist."
                )

        except Exception as e:
            logger.error(f"Error deleting database '{name}': {e}")

    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.conn.runInstalledQuery(query_name, **kwargs)
            else:
                # Execute as raw GSQL
                result = self.conn.gsql(query)
            return result
        except Exception as e:
            logger.error(f"Error executing query '{query}': {e}")
            raise

    def close(self):
        """Close connection - pyTigerGraph handles cleanup automatically."""
        pass

    def init_db(self, schema: Schema, clean_start=False):
        """
        Initialize database with schema definition.

        Follows the same pattern as ArangoDB:
        1. Clean if needed
        2. Create vertex and edge types globally (required before CREATE GRAPH)
        3. Create graph with vertices and edges explicitly attached
        4. 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")

        try:
            if clean_start:
                try:
                    # Only delete the current graph, not all graphs or global vertex/edge types
                    # This ensures we don't affect other graphs that might share vertex/edge types
                    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 clean_start for graph '{graph_name}': {clean_error}",
                        exc_info=True,
                    )
                    # Continue - may be first run or already clean, schema will be recreated anyway

            # Step 1: Create vertex and edge types globally first
            # These must exist before they can be included in CREATE GRAPH
            logger.debug(
                f"Creating vertex and edge types globally for graph '{graph_name}'"
            )
            try:
                vertex_names = self._create_vertex_types_global(schema.vertex_config)

                # Initialize edges before creating edge types
                # This sets edge._source and edge._target to dbnames (required for GSQL)
                edges_to_create = list(schema.edge_config.edges_list(include_aux=True))
                for edge in edges_to_create:
                    edge.finish_init(schema.vertex_config)

                # Verify all vertices referenced by edges were created
                created_vertex_set = set(vertex_names)
                for edge in edges_to_create:
                    if edge._source not in created_vertex_set:
                        raise ValueError(
                            f"Edge '{edge.relation}' references source vertex '{edge._source}' "
                            f"which was not created. Created vertices: {vertex_names}"
                        )
                    if edge._target not in created_vertex_set:
                        raise ValueError(
                            f"Edge '{edge.relation}' references target vertex '{edge._target}' "
                            f"which was not created. Created vertices: {vertex_names}"
                        )

                edge_names = self._create_edge_types_global(edges_to_create)
                logger.debug(
                    f"Created {len(vertex_names)} vertex types and {len(edge_names)} edge types"
                )
            except Exception as type_error:
                logger.error(
                    f"Failed to create vertex/edge types for graph '{graph_name}': {type_error}",
                    exc_info=True,
                )
                raise

            # Step 2: Create graph with vertices and edges explicitly attached
            try:
                if not self.graph_exists(graph_name):
                    logger.debug(f"Creating graph '{graph_name}' with types in init_db")
                    try:
                        self.create_database(
                            graph_name,
                            vertex_names=vertex_names,
                            edge_names=edge_names,
                        )
                        graph_created = True
                        logger.info(f"Successfully created 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")
                    # If graph already exists, associate types via ALTER GRAPH
                    try:
                        self.define_vertex_collections(schema.vertex_config)
                        # Ensure edges are initialized before defining collections
                        edges_for_collections = list(
                            schema.edge_config.edges_list(include_aux=True)
                        )
                        for edge in edges_for_collections:
                            if edge._source is None or edge._target is None:
                                edge.finish_init(schema.vertex_config)
                        self.define_edge_collections(edges_for_collections)
                    except Exception as define_error:
                        logger.warning(
                            f"Could not define collections for existing graph '{graph_name}': {define_error}",
                            exc_info=True,
                        )
                        # Continue - graph exists, collections may already be defined
            except Exception as graph_error:
                logger.error(
                    f"Error during graph creation/verification for '{graph_name}': {graph_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

    def define_schema(self, schema: Schema):
        """
        Define TigerGraph schema with proper GSQL syntax.

        Assumes graph already exists (created in init_db). This method:
        1. Uses the graph from config.database
        2. Defines vertex types within the graph
        3. Defines edge types within the graph
        """
        try:
            # Define vertex and edge types within the graph
            # Graph context is ensured by _ensure_graph_context in the called methods
            self.define_vertex_collections(schema.vertex_config)
            # Ensure edges are initialized before defining collections
            edges_for_collections = list(
                schema.edge_config.edges_list(include_aux=True)
            )
            for edge in edges_for_collections:
                if edge._source is None or edge._target is None:
                    edge.finish_init(schema.vertex_config)
            self.define_edge_collections(edges_for_collections)

        except Exception as e:
            logger.error(f"Error defining schema: {e}")
            raise

    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(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 _create_vertex_types_global(self, vertex_config: VertexConfig) -> list[str]:
        """Create TigerGraph vertex types globally (without graph association).

        Vertices are global in TigerGraph and must be created before they can be
        included in a CREATE GRAPH statement.

        Creates vertices with PRIMARY_ID (single field) or PRIMARY KEY (composite) syntax.
        For single-field indexes, uses PRIMARY_ID syntax (required by GraphStudio).
        For composite keys, uses PRIMARY KEY syntax (works in GSQL but not GraphStudio).
        According to TigerGraph documentation, fields used in PRIMARY KEY/PRIMARY_ID must be
        defined as regular attributes first, and they remain accessible as attributes.

        Note: GraphStudio does not support composite keys. Use PRIMARY_ID for single fields
        to ensure compatibility with GraphStudio.

        Reference: https://docs.tigergraph.com/gsql-ref/4.2/ddl-and-loading/defining-a-graph-schema

        Args:
            vertex_config: Vertex configuration containing vertices to create

        Returns:
            list[str]: List of vertex type names that were created (or already existed)
        """
        vertex_names = []
        for vertex in vertex_config.vertices:
            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)
                # Format: PRIMARY_ID field_name field_type, other_field1 TYPE, other_field2 TYPE, ...
                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)
            elif len(index_fields) > 1:
                # Composite key: use PRIMARY KEY syntax (works in GSQL but not GraphStudio UI)
                # Format: field1 TYPE, field2 TYPE, ..., PRIMARY KEY (field1, field2, ...)
                logger.warning(
                    f"Vertex '{vertex_dbname}' has composite primary key {index_fields}. "
                    f"GraphStudio UI does not support composite keys. "
                    f"Consider using a single-field PRIMARY_ID instead."
                )

                # List all fields first
                field_parts = [f"{name} {ftype}" for name, ftype in all_fields]
                # Then add PRIMARY KEY at the end
                vindex = "(" + ", ".join(index_fields) + ")"
                field_parts.append(f"PRIMARY KEY {vindex}")

                field_definitions = ",\n    ".join(field_parts)
            else:
                raise ValueError(
                    f"Vertex '{vertex_dbname}' must have at least one index field"
                )

            # Create the vertex type globally (ignore if exists)
            # Vertices are global in TigerGraph, so no USE GRAPH needed
            # Note: For PRIMARY_ID, the ID field is listed first with PRIMARY_ID keyword
            # For PRIMARY KEY, all fields are listed first, then PRIMARY KEY clause at the end
            # When using PRIMARY_ID, we need primary_id_as_attribute="true" to make the ID
            # accessible as an attribute (required for REST++ API upserts)
            if len(index_fields) == 1:
                # Single field with PRIMARY_ID: enable primary_id_as_attribute so ID is accessible
                create_vertex_cmd = (
                    f"CREATE VERTEX {vertex_dbname} (\n"
                    f"    {field_definitions}\n"
                    f') WITH STATS="OUTDEGREE_BY_EDGETYPE", primary_id_as_attribute="true"'
                )
            else:
                # Composite key with PRIMARY KEY: key fields are automatically accessible as attributes
                create_vertex_cmd = (
                    f"CREATE VERTEX {vertex_dbname} (\n"
                    f"    {field_definitions}\n"
                    f') WITH STATS="OUTDEGREE_BY_EDGETYPE"'
                )
            logger.debug(f"Executing GSQL: {create_vertex_cmd}")
            try:
                result = self.conn.gsql(create_vertex_cmd)
                logger.debug(f"Result: {result}")
                vertex_names.append(vertex_dbname)
                logger.info(f"Successfully created vertex type '{vertex_dbname}'")
            except Exception as e:
                err = str(e).lower()
                if (
                    "used by another object" in err
                    or "duplicate" in err
                    or "already exists" in err
                ):
                    logger.debug(
                        f"Vertex type '{vertex_dbname}' already exists; will include in graph"
                    )
                    vertex_names.append(vertex_dbname)
                else:
                    logger.error(
                        f"Failed to create vertex type '{vertex_dbname}': {e}\n"
                        f"GSQL command was: {create_vertex_cmd}"
                    )
                    raise
        return vertex_names

    def define_vertex_collections(self, vertex_config: VertexConfig):
        """Define TigerGraph vertex types and associate them with the current graph.

        Flow per vertex type:
        1) Try to CREATE VERTEX (idempotent: ignore "already exists" errors)
        2) Associate the vertex with the graph via ALTER GRAPH <graph> ADD VERTEX <vertex>

        Args:
            vertex_config: Vertex configuration containing vertices to create
        """
        # First create all vertex types globally
        vertex_names = self._create_vertex_types_global(vertex_config)

        # Then associate them with the graph (if graph already exists)
        graph_name = self.config.database
        if graph_name:
            for vertex_name in vertex_names:
                alter_graph_cmd = f"USE GRAPH {graph_name}\nALTER GRAPH {graph_name} ADD VERTEX {vertex_name}"
                logger.debug(f"Executing GSQL: {alter_graph_cmd}")
                try:
                    result = self.conn.gsql(alter_graph_cmd)
                    logger.debug(f"Result: {result}")
                except Exception as e:
                    err = str(e).lower()
                    # If already associated, ignore
                    if "already" in err and ("added" in err or "exists" in err):
                        logger.debug(
                            f"Vertex '{vertex_name}' already associated with graph '{graph_name}'"
                        )
                    else:
                        raise

    def _create_edge_types_global(self, edges: list[Edge]) -> list[str]:
        """Create TigerGraph edge types globally (without graph association).

        Edges are global in TigerGraph and must be created before they can be
        included in a CREATE GRAPH statement.

        Args:
            edges: List of edges to create (should have _source_collection and _target_collection populated)

        Returns:
            list[str]: List of edge type names (relation names) that were created (or already existed)
        """
        edge_names = []
        for edge in edges:
            edge_attrs = self._format_edge_attributes(edge)

            # Create the edge type globally (ignore if exists/used elsewhere)
            # Edges are global in TigerGraph, so no USE GRAPH needed
            create_edge_cmd = (
                f"CREATE DIRECTED EDGE {edge.relation} (\n"
                f"    FROM {edge._source},\n"
                f"    TO {edge._target}{edge_attrs}\n"
                f")"
            )
            logger.debug(f"Executing GSQL: {create_edge_cmd}")
            try:
                result = self.conn.gsql(create_edge_cmd)
                logger.debug(f"Result: {result}")
                edge_names.append(edge.relation)
            except Exception as e:
                err = str(e).lower()
                # If the edge name is already used by another object or duplicates exist, continue
                if (
                    "used by another object" in err
                    or "duplicate" in err
                    or "already exists" in err
                ):
                    logger.debug(
                        f"Edge type '{edge.relation}' already defined; will include in graph"
                    )
                    edge_names.append(edge.relation)
                else:
                    raise
        return edge_names

    def define_edge_collections(self, edges: list[Edge]):
        """Define TigerGraph edge types and associate them with the current graph.

        Flow per edge type:
        1) Try to CREATE DIRECTED EDGE (idempotent: ignore "used by another object"/"duplicate"/"already exists")
        2) Associate the edge with the graph via ALTER GRAPH <graph> ADD DIRECTED EDGE <edge>

        Args:
            edges: List of edges to create (should have _source_collection and _target_collection populated)
        """
        # First create all edge types globally
        edge_names = self._create_edge_types_global(edges)

        # Then associate them with the graph (if graph already exists)
        graph_name = self.config.database
        if graph_name:
            for edge_name in edge_names:
                alter_graph_cmd = (
                    f"USE GRAPH {graph_name}\n"
                    f"ALTER GRAPH {graph_name} ADD DIRECTED EDGE {edge_name}"
                )
                logger.debug(f"Executing GSQL: {alter_graph_cmd}")
                try:
                    result = self.conn.gsql(alter_graph_cmd)
                    logger.debug(f"Result: {result}")
                except Exception as e:
                    err = str(e).lower()
                    # If already associated, ignore
                    if "already" in err and ("added" in err or "exists" in err):
                        logger.debug(
                            f"Edge '{edge_name}' already associated with graph '{graph_name}'"
                        )
                    else:
                        raise

    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."""
        logger.warning("TigerGraph edge indices not implemented yet [version 4.2.2]")

    def _add_index(self, obj_name, index: Index, is_vertex_index=True):
        """
        Create an index on a vertex or edge type using GSQL schema change jobs.

        TigerGraph requires indexes to be created through schema change jobs:
        1. CREATE GLOBAL SCHEMA_CHANGE job job_name {ALTER VERTEX ... ADD INDEX ... ON (...);}
        2. RUN GLOBAL SCHEMA_CHANGE job job_name

        Note: TigerGraph only supports secondary indexes on a single field.
        Indexes with multiple fields will be skipped with a warning.
        Edge indexes are not supported in TigerGraph and will be skipped with a warning.

        Args:
            obj_name: Name of the vertex type or edge type
            index: Index configuration object
            is_vertex_index: Whether this is a vertex index (True) or edge index (False)
        """
        try:
            # TigerGraph doesn't support indexes on edges
            if not is_vertex_index:
                logger.warning(
                    f"Edge indexes are not supported in TigerGraph [current version 4.2.2]"
                    f"Skipping index creation for edge '{obj_name}' on field(s) '{index.fields}'"
                )
                return

            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 (single field in parentheses)
            # Note: Only vertex indexes are supported - edge indexes are handled earlier
            alter_stmt = (
                f"ALTER VERTEX {obj_name} ADD INDEX {index_name} ON ({field_name})"
            )

            # Step 1: Create the schema change job
            # only global changes are supported by tigergraph
            create_job_cmd = (
                f"USE GLOBAL \n"
                f"CREATE GLOBAL SCHEMA_CHANGE job {job_name} {{{alter_stmt};}}"
            )

            logger.debug(f"Executing GSQL (create job): {create_job_cmd}")
            try:
                result = self.conn.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 GLOBAL SCHEMA_CHANGE job {job_name}"

            logger.debug(f"Executing GSQL (run job): {run_job_cmd}")
            try:
                result = self.conn.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]:
        """
        Generic parser for SHOW * output commands.

        Extracts names from lines matching the pattern: "- PREFIX name(...)"

        Args:
            result_str: String output from SHOW * GSQL command
            prefix: The prefix to look for (e.g., "VERTEX", "GRAPH", "JOB")

        Returns:
            List of extracted names
        """
        names = []
        lines = result_str.split("\n")

        for line in lines:
            line = line.strip()
            # Skip empty lines and headers
            if not line or line.startswith("*"):
                continue

            # Remove leading "- " if present
            if line.startswith("- "):
                line = line[2:].strip()

            # Look for prefix pattern
            prefix_upper = prefix.upper()
            if line.upper().startswith(f"{prefix_upper} "):
                # Extract name (after prefix and before opening parenthesis or whitespace)
                after_prefix = line[len(prefix_upper) + 1 :].strip()
                # Name is the first word (before space or parenthesis)
                if "(" in after_prefix:
                    name = after_prefix.split("(")[0].strip()
                else:
                    # No parenthesis, take the first word
                    name = (
                        after_prefix.split()[0].strip()
                        if after_prefix.split()
                        else None
                    )

                if name:
                    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)
        """
        edge_types = []
        lines = result_str.split("\n")

        for line in lines:
            line = line.strip()
            # Skip empty lines and headers
            if not line or line.startswith("*"):
                continue

            # Remove leading "- " if present
            if line.startswith("- "):
                line = line[2:].strip()

            # Look for "DIRECTED EDGE" or "UNDIRECTED EDGE" pattern
            is_directed = None
            prefix = None
            if "DIRECTED EDGE" in line.upper():
                prefix = "DIRECTED EDGE "
                is_directed = True
            elif "UNDIRECTED EDGE" in line.upper():
                prefix = "UNDIRECTED EDGE "
                is_directed = False

            if prefix:
                idx = line.upper().find(prefix)
                if idx >= 0:
                    after_prefix = line[idx + len(prefix) :].strip()
                    # Extract name before opening parenthesis
                    if "(" in after_prefix:
                        edge_name = after_prefix.split("(")[0].strip()
                        if edge_name:
                            edge_types.append((edge_name, is_directed))

        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.conn.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.conn.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.conn.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.conn.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.conn.delVertices(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.conn.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.conn.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.conn.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.conn.delVertices(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 _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()
                }

                # 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],
        host: str,
        graph_name: str,
        username: str | None = None,
        password: str | None = None,
    ) -> dict[str, Any]:
        """
        Sends the generated JSON payload to the TigerGraph REST++ upsert endpoint.

        Args:
            payload: The JSON payload in TigerGraph REST++ format
            host: Base host URL (e.g., "http://localhost:9000")
            graph_name: Name of the graph
            username: Optional username for authentication
            password: Optional password for authentication

        Returns:
            Dictionary containing the response from TigerGraph
        """
        url = f"{host}/graph/{graph_name}"

        headers = {
            "Content-Type": "application/json",
        }

        logger.debug(f"Attempting batch upsert to: {url}")

        try:
            # Use HTTP Basic Auth if username and password are provided
            auth = None
            if username and password:
                auth = (username, password)

            response = requests.post(
                url,
                headers=headers,
                data=json.dumps(payload, default=_json_serializer),
                auth=auth,
                # Increase timeout for large batches
                timeout=120,
            )
            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:
            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)}

    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

            # Build REST++ endpoint URL
            host = f"{self.config.url_without_port}:{self.config.port}"
            graph_name = self.config.database
            if not graph_name:
                raise ValueError("Graph name (database) must be configured")

            # Send the upsert request with username/password authentication
            result = self._upsert_data(
                payload,
                host,
                graph_name,
                username=self.config.username,
                password=self.config.password,
            )

            if result.get("error"):
                logger.error(
                    f"Error upserting vertices to {class_name}: {result.get('message')}"
                )
                # Fallback to individual operations
                self._fallback_individual_upsert(docs, class_name, match_keys)
            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}")
            # Fallback to individual operations
            self._fallback_individual_upsert(docs, class_name, match_keys)

    def _fallback_individual_upsert(self, docs, class_name, match_keys):
        """Fallback method for individual vertex upserts."""
        for doc in docs:
            try:
                vertex_id = self._extract_id(doc, match_keys)
                if vertex_id:
                    clean_doc = self._clean_document(doc)
                    # Serialize datetime objects before passing to pyTigerGraph
                    # pyTigerGraph's upsertVertex expects JSON-serializable data
                    serialized_doc = json.loads(
                        json.dumps(clean_doc, default=_json_serializer)
                    )
                    self.conn.upsertVertex(class_name, vertex_id, serialized_doc)
            except Exception as e:
                logger.error(f"Error upserting individual vertex {vertex_id}: {e}")

    def _generate_edge_upsert_payload(
        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, ...],
    ) -> dict[str, Any]:
        """
        Transforms edge data into the TigerGraph REST++ batch upsert JSON format.

        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
            match_keys_source: Tuple of index fields for source vertex
            match_keys_target: Tuple of index fields for target vertex

        Returns:
            Dictionary in TigerGraph REST++ batch upsert format for edges
        """
        # Initialize the required JSON structure for edges
        payload: dict[str, Any] = {"edges": {source_class: {}}}
        source_map = payload["edges"][source_class]

        for source_doc, target_doc, edge_props in edges_data:
            try:
                # Extract source ID (composite if needed)
                if isinstance(match_keys_source, tuple) and len(match_keys_source) > 1:
                    source_id_components = [
                        str(source_doc[key]) for key in match_keys_source
                    ]
                    source_id = "_".join(source_id_components)
                else:
                    source_id = self._extract_id(source_doc, match_keys_source)

                # Extract target ID (composite if needed)
                if isinstance(match_keys_target, tuple) and len(match_keys_target) > 1:
                    target_id_components = [
                        str(target_doc[key]) for key in match_keys_target
                    ]
                    target_id = "_".join(target_id_components)
                else:
                    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

                # Initialize source vertex entry if not exists
                if source_id not in source_map:
                    source_map[source_id] = {edge_type: {}}

                # Initialize edge type entry if not exists
                if edge_type not in source_map[source_id]:
                    source_map[source_id][edge_type] = {}

                # Initialize target vertex type entry if not exists
                if target_class not in source_map[source_id][edge_type]:
                    source_map[source_id][edge_type][target_class] = {}

                # Format edge attributes for TigerGraph REST++ API
                # Clean edge properties (remove internal keys)
                clean_edge_props = self._clean_document(edge_props)

                # Format attributes with {"value": ...} wrapper
                formatted_attributes = {
                    k: {"value": v} for k, v in clean_edge_props.items()
                }

                # Add target vertex with edge attributes under target vertex type
                source_map[source_id][edge_type][target_class][target_id] = (
                    formatted_attributes
                )

            except KeyError as e:
                logger.warning(
                    f"Edge is missing a required field: {e}. Skipping edge: {source_doc}, {target_doc}"
                )
                continue
            except Exception as e:
                logger.error(f"Error processing edge: {e}")
                continue

        return payload

    def insert_edges_batch(
        self,
        docs_edges,
        source_class,
        target_class,
        relation_name,
        collection_name=None,
        match_keys_source=("_key",),
        match_keys_target=("_key",),
        filter_uniques=True,
        uniq_weight_fields=None,
        uniq_weight_collections=None,
        upsert_option=False,
        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": {...}}, ...]

        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
            collection_name: Alternative edge collection name (used if relation_name is None)
            match_keys_source: Keys to match source vertices
            match_keys_target: Keys to match target vertices
            filter_uniques: If True, filter duplicate edges
            uniq_weight_fields: Fields to consider for uniqueness (not used in TigerGraph)
            uniq_weight_collections: Collections to consider for uniqueness (not used in TigerGraph)
            upsert_option: If True, use upsert (default behavior in TigerGraph)
            head: Optional limit on number of edges to insert
            **kwargs: Additional options:
                - dry: If True, don't execute the query
        """
        dry = kwargs.pop("dry", False)
        if dry:
            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
        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 the edge upsert payload
            payload = self._generate_edge_upsert_payload(
                normalized_edges,
                source_class,
                target_class,
                edge_type,
                match_keys_src,
                match_keys_tgt,
            )

            # Check if payload has any edges
            source_vertices = payload.get("edges", {}).get(source_class, {})
            if not source_vertices:
                logger.warning(f"No valid edges to upsert for edge type {edge_type}")
                return

            # Build REST++ endpoint URL
            host = f"{self.config.url_without_port}:{self.config.port}"
            graph_name = self.config.database
            if not graph_name:
                raise ValueError("Graph name (database) must be configured")

            # Send the upsert request with username/password authentication
            result = self._upsert_data(
                payload,
                host,
                graph_name,
                username=self.config.username,
                password=self.config.password,
            )

            if result.get("error"):
                logger.error(
                    f"Error upserting edges of type {edge_type}: {result.get('message')}"
                )
            else:
                # Count edges in payload
                edge_count = 0
                for source_edges in source_vertices.values():
                    if edge_type in source_edges:
                        if target_class in source_edges[edge_type]:
                            edge_count += len(source_edges[edge_type][target_class])
                logger.debug(
                    f"Upserted {edge_count} edges of type {edge_type}: {result}"
                )
                return result

        except Exception as e:
            logger.error(f"Error batch inserting edges: {e}")

    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, class_name):
        """
        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 | Clause | 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, Clause):
                ff = Expression.from_dict(filters)
            else:
                ff = filters

            # Use ExpressionFlavor.TIGERGRAPH with empty doc_name to trigger REST++ format
            # Pass field_types to help with proper value quoting
            filter_str = ff(
                doc_name="",
                kind=ExpressionFlavor.TIGERGRAPH,
                field_types=field_types,
            )
            return filter_str
        else:
            return ""

    def fetch_docs(
        self,
        class_name,
        filters: list | dict | Clause | None = None,
        limit: int | None = None,
        return_keys: list | None = None,
        unset_keys: list | None = None,
        **kwargs,
    ):
        """
        Fetch documents (vertices) with filtering and projection using REST++ API.

        Args:
            class_name: Vertex type name (or dbname)
            filters: Filter expression (list, dict, or Clause)
            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 | dict | Clause | None = None,
        limit: int | None = None,
        return_keys: list | None = None,
        unset_keys: list | None = None,
        **kwargs,
    ):
        """
        Fetch edges from TigerGraph using pyTigerGraph's getEdges method.

        In TigerGraph, you must know at least one vertex ID before you can fetch edges.
        Uses pyTigerGraph's getEdges method 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 pyTigerGraph)
            to_id: Optional target vertex ID to filter by (not used in pyTigerGraph)
            filters: Additional query filters (not supported by pyTigerGraph getEdges)
            limit: Maximum number of edges to return (not supported by pyTigerGraph getEdges)
            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 pyTigerGraph's getEdges method
            # Signature: getEdges(sourceVertexType, sourceVertexId, edgeType=None)
            # Returns: list of edge dictionaries
            logger.debug(
                f"Fetching edges using pyTigerGraph: from_type={from_type}, from_id={from_id}, edge_type={edge_type}"
            )

            # Handle None edge_type by passing empty string (default behavior)
            edge_type_str = edge_type if edge_type is not None else ""
            edges = self.conn.getEdges(from_type, from_id, edge_type_str, fmt="py")

            # Parse pyTigerGraph response format
            # getEdges returns list of dicts with format like:
            # [{"e_type": "...", "from": {...}, "to": {...}, "attributes": {...}}, ...]
            # Type annotation: result is list[dict[str, Any]]
            # getEdges can return dict, str, or DataFrame, but with fmt="py" it returns dict
            if isinstance(edges, list):
                # Type narrowing: after isinstance check, we know it's a list
                # Use cast to help type checker understand the elements are dicts
                result = cast(list[dict[str, Any]], edges)
            elif isinstance(edges, dict):
                # If it's a single dict, wrap it in a list
                result = [cast(dict[str, Any], edges)]
            else:
                # Fallback for unexpected types
                result: list[dict[str, Any]] = []

            # Apply limit if specified (client-side since pyTigerGraph 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 pyTigerGraph: {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,
        class_name,
        match_keys,
        keep_keys,
        flatten=False,
        filters: list | dict | None = None,
    ):
        """
        Check which documents from batch are present in the database.
        """
        try:
            present_docs = {}

            for i, doc in enumerate(batch):
                vertex_id = self._extract_id(doc, match_keys)
                if not vertex_id:
                    continue

                try:
                    vertex_data = self.conn.getVerticesById(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 = {}

                        for key in keep_keys:
                            if key == "id":
                                filtered_doc[key] = vertex_id
                            elif key in vertex_attrs:
                                filtered_doc[key] = vertex_attrs[key]

                        present_docs[i] = [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.conn.getVertexCount(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,
        class_name,
        match_keys,
        keep_keys,
        filters: list | dict | None = None,
    ):
        """
        Return documents from batch that are NOT present in database.
        """
        present_docs_indices = self.fetch_present_documents(
            batch=batch,
            class_name=class_name,
            match_keys=match_keys,
            keep_keys=keep_keys,
            flatten=False,
            filters=filters,
        )

        absent_indices = sorted(
            set(range(len(batch))) - set(present_docs_indices.keys())
        )
        return [batch[i] for i in absent_indices]

    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.conn.getVertexTypes(force=True)

                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.conn.getVertexCount(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 []

close()

Close connection - pyTigerGraph handles cleanup automatically.

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

def close(self):
    """Close connection - pyTigerGraph handles cleanup automatically."""
    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 the pyTigerGraph gsql() method 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
Exception	If graph creation fails

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

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 the pyTigerGraph gsql() method 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:
        Exception: If graph creation fails
    """
    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 pyTigerGraph's gsql method which handles authentication
        logger.debug(f"Creating graph '{name}' via GSQL: {gsql_commands}")
        try:
            result = self.conn.gsql(gsql_commands)
            logger.info(
                f"Successfully created graph '{name}' with types {all_types}: {result}"
            )
            return result
        except Exception as e:
            error_msg = str(e).lower()
            # Check if graph already exists (might be acceptable)
            if "already exists" in error_msg or "duplicate" in error_msg:
                logger.info(f"Graph '{name}' may already exist: {e}")
                return str(e)
            logger.error(f"Failed to create graph '{name}': {e}")
            raise

    except Exception as e:
        logger.error(f"Error creating graph '{name}' via GSQL: {e}")
        raise

define_edge_collections(edges)

Define TigerGraph edge types and associate them with the current graph.

Flow per edge type: 1) Try to CREATE DIRECTED EDGE (idempotent: ignore "used by another object"/"duplicate"/"already exists") 2) Associate the edge with the graph via ALTER GRAPH ADD DIRECTED EDGE

Parameters:

Name	Type	Description	Default
edges	list[Edge]	List of edges to create (should have _source_collection and _target_collection populated)	required

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

def define_edge_collections(self, edges: list[Edge]):
    """Define TigerGraph edge types and associate them with the current graph.

    Flow per edge type:
    1) Try to CREATE DIRECTED EDGE (idempotent: ignore "used by another object"/"duplicate"/"already exists")
    2) Associate the edge with the graph via ALTER GRAPH <graph> ADD DIRECTED EDGE <edge>

    Args:
        edges: List of edges to create (should have _source_collection and _target_collection populated)
    """
    # First create all edge types globally
    edge_names = self._create_edge_types_global(edges)

    # Then associate them with the graph (if graph already exists)
    graph_name = self.config.database
    if graph_name:
        for edge_name in edge_names:
            alter_graph_cmd = (
                f"USE GRAPH {graph_name}\n"
                f"ALTER GRAPH {graph_name} ADD DIRECTED EDGE {edge_name}"
            )
            logger.debug(f"Executing GSQL: {alter_graph_cmd}")
            try:
                result = self.conn.gsql(alter_graph_cmd)
                logger.debug(f"Result: {result}")
            except Exception as e:
                err = str(e).lower()
                # If already associated, ignore
                if "already" in err and ("added" in err or "exists" in err):
                    logger.debug(
                        f"Edge '{edge_name}' already associated with graph '{graph_name}'"
                    )
                else:
                    raise

define_edge_indices(edges)

Define indices for edges if specified.

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

def define_edge_indices(self, edges: list[Edge]):
    """Define indices for edges if specified."""
    logger.warning("TigerGraph edge indices not implemented yet [version 4.2.2]")

define_indexes(schema)

Define all indexes from schema.

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

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 with proper GSQL syntax.

Assumes graph already exists (created in init_db). This method: 1. Uses the graph from config.database 2. Defines vertex types within the graph 3. Defines edge types within the graph

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

def define_schema(self, schema: Schema):
    """
    Define TigerGraph schema with proper GSQL syntax.

    Assumes graph already exists (created in init_db). This method:
    1. Uses the graph from config.database
    2. Defines vertex types within the graph
    3. Defines edge types within the graph
    """
    try:
        # Define vertex and edge types within the graph
        # Graph context is ensured by _ensure_graph_context in the called methods
        self.define_vertex_collections(schema.vertex_config)
        # Ensure edges are initialized before defining collections
        edges_for_collections = list(
            schema.edge_config.edges_list(include_aux=True)
        )
        for edge in edges_for_collections:
            if edge._source is None or edge._target is None:
                edge.finish_init(schema.vertex_config)
        self.define_edge_collections(edges_for_collections)

    except Exception as e:
        logger.error(f"Error defining schema: {e}")
        raise

define_vertex_collections(vertex_config)

Define TigerGraph vertex types and associate them with the current graph.

Flow per vertex type: 1) Try to CREATE VERTEX (idempotent: ignore "already exists" errors) 2) Associate the vertex with the graph via ALTER GRAPH ADD VERTEX

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_collections(self, vertex_config: VertexConfig):
    """Define TigerGraph vertex types and associate them with the current graph.

    Flow per vertex type:
    1) Try to CREATE VERTEX (idempotent: ignore "already exists" errors)
    2) Associate the vertex with the graph via ALTER GRAPH <graph> ADD VERTEX <vertex>

    Args:
        vertex_config: Vertex configuration containing vertices to create
    """
    # First create all vertex types globally
    vertex_names = self._create_vertex_types_global(vertex_config)

    # Then associate them with the graph (if graph already exists)
    graph_name = self.config.database
    if graph_name:
        for vertex_name in vertex_names:
            alter_graph_cmd = f"USE GRAPH {graph_name}\nALTER GRAPH {graph_name} ADD VERTEX {vertex_name}"
            logger.debug(f"Executing GSQL: {alter_graph_cmd}")
            try:
                result = self.conn.gsql(alter_graph_cmd)
                logger.debug(f"Result: {result}")
            except Exception as e:
                err = str(e).lower()
                # If already associated, ignore
                if "already" in err and ("added" in err or "exists" in err):
                    logger.debug(
                        f"Vertex '{vertex_name}' already associated with graph '{graph_name}'"
                    )
                else:
                    raise

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 GSQL DROP GRAPH. If that fails (e.g., dependencies), it will: 1) Remove associations and drop all edge types 2) Drop all vertex types 3) Clear remaining data as a last resort

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

def delete_database(self, name: str):
    """
    Delete a TigerGraph database (graph).

    This method attempts to drop the graph using GSQL DROP GRAPH.
    If that fails (e.g., dependencies), it will:
      1) Remove associations and drop all edge types
      2) Drop all vertex types
      3) Clear remaining data as a last resort

    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}'")
        try:
            # Use the graph first to ensure we're working with the right graph
            drop_command = f"USE GRAPH {name}\nDROP GRAPH {name}"
            result = self.conn.gsql(drop_command)
            logger.info(f"Successfully dropped graph '{name}': {result}")
            return result
        except Exception as e:
            logger.debug(
                f"Could not drop graph '{name}' (may not exist or have dependencies): {e}"
            )

        # Fallback 1: Attempt to disassociate edge and vertex types from graph
        # DO NOT drop global vertex/edge types as they might be used by other graphs
        try:
            with self._ensure_graph_context(name):
                # Disassociate edge types from graph (but don't drop them globally)
                try:
                    edge_types = self.conn.getEdgeTypes(force=True)
                except Exception:
                    edge_types = []

                for e_type in edge_types:
                    # Only disassociate from graph, don't drop globally
                    # ALTER GRAPH requires USE GRAPH context
                    try:
                        drop_edge_cmd = f"USE GRAPH {name}\nALTER GRAPH {name} DROP DIRECTED EDGE {e_type}"
                        self.conn.gsql(drop_edge_cmd)
                        logger.debug(
                            f"Disassociated edge type '{e_type}' from graph '{name}'"
                        )
                    except Exception as e:
                        logger.debug(
                            f"Could not disassociate edge type '{e_type}' from graph '{name}': {e}"
                        )
                        # Continue - edge might not be associated or graph might not exist

                # Disassociate vertex types from graph (but don't drop them globally)
                try:
                    vertex_types = self.conn.getVertexTypes(force=True)
                except Exception:
                    vertex_types = []

                for v_type in vertex_types:
                    # Only clear data from this graph's vertices, don't drop vertex type globally
                    # Clear data first to avoid dependency issues
                    try:
                        self.conn.delVertices(v_type)
                        logger.debug(
                            f"Cleared vertices of type '{v_type}' from graph '{name}'"
                        )
                    except Exception as e:
                        logger.debug(
                            f"Could not clear vertices of type '{v_type}' from graph '{name}': {e}"
                        )
                    # Disassociate from graph (best-effort)
                    # ALTER GRAPH requires USE GRAPH context
                    try:
                        drop_vertex_cmd = f"USE GRAPH {name}\nALTER GRAPH {name} DROP VERTEX {v_type}"
                        self.conn.gsql(drop_vertex_cmd)
                        logger.debug(
                            f"Disassociated vertex type '{v_type}' from graph '{name}'"
                        )
                    except Exception as e:
                        logger.debug(
                            f"Could not disassociate vertex type '{v_type}' from graph '{name}': {e}"
                        )
                        # Continue - vertex might not be associated or graph might not exist
        except Exception as e3:
            logger.warning(
                f"Could not disassociate schema types from graph '{name}': {e3}. Proceeding to data clear."
            )

        # Fallback 2: Clear all data (if any remain)
        try:
            with self._ensure_graph_context(name):
                vertex_types = self.conn.getVertexTypes()
                for v_type in vertex_types:
                    result = self.conn.delVertices(v_type)
                    logger.debug(f"Cleared vertices of type {v_type}: {result}")
                logger.info(f"Cleared all data from graph '{name}'")
        except Exception as e2:
            logger.warning(
                f"Could not clear data from graph '{name}': {e2}. Graph may not exist."
            )

    except Exception as e:
        logger.error(f"Error deleting database '{name}': {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.conn.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.conn.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.conn.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.conn.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.conn.delVertices(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.conn.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.conn.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.conn.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.conn.delVertices(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

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.conn.runInstalledQuery(query_name, **kwargs)
        else:
            # Execute as raw GSQL
            result = self.conn.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		Vertex type name (or dbname)	required
filters	list | dict | Clause | None	Filter expression (list, dict, or Clause)	None
limit	int | None	Maximum number of documents to return	None
return_keys	list | None	Keys to return (projection)	None
unset_keys	list | None	Keys to exclude (projection)	None
**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:

Name	Type	Description
list		List of fetched documents

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

def fetch_docs(
    self,
    class_name,
    filters: list | dict | Clause | None = None,
    limit: int | None = None,
    return_keys: list | None = None,
    unset_keys: list | None = None,
    **kwargs,
):
    """
    Fetch documents (vertices) with filtering and projection using REST++ API.

    Args:
        class_name: Vertex type name (or dbname)
        filters: Filter expression (list, dict, or Clause)
        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 pyTigerGraph's getEdges method.

In TigerGraph, you must know at least one vertex ID before you can fetch edges. Uses pyTigerGraph's getEdges method 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 pyTigerGraph)	None
to_id	str | None	Optional target vertex ID to filter by (not used in pyTigerGraph)	None
filters	list | dict | Clause | None	Additional query filters (not supported by pyTigerGraph getEdges)	None
limit	int | None	Maximum number of edges to return (not supported by pyTigerGraph getEdges)	None
return_keys	list | None	Keys to return (projection)	None
unset_keys	list | None	Keys to exclude (projection)	None
**kwargs		Additional parameters	{}

Returns:

Name	Type	Description
list		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 | dict | Clause | None = None,
    limit: int | None = None,
    return_keys: list | None = None,
    unset_keys: list | None = None,
    **kwargs,
):
    """
    Fetch edges from TigerGraph using pyTigerGraph's getEdges method.

    In TigerGraph, you must know at least one vertex ID before you can fetch edges.
    Uses pyTigerGraph's getEdges method 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 pyTigerGraph)
        to_id: Optional target vertex ID to filter by (not used in pyTigerGraph)
        filters: Additional query filters (not supported by pyTigerGraph getEdges)
        limit: Maximum number of edges to return (not supported by pyTigerGraph getEdges)
        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 pyTigerGraph's getEdges method
        # Signature: getEdges(sourceVertexType, sourceVertexId, edgeType=None)
        # Returns: list of edge dictionaries
        logger.debug(
            f"Fetching edges using pyTigerGraph: from_type={from_type}, from_id={from_id}, edge_type={edge_type}"
        )

        # Handle None edge_type by passing empty string (default behavior)
        edge_type_str = edge_type if edge_type is not None else ""
        edges = self.conn.getEdges(from_type, from_id, edge_type_str, fmt="py")

        # Parse pyTigerGraph response format
        # getEdges returns list of dicts with format like:
        # [{"e_type": "...", "from": {...}, "to": {...}, "attributes": {...}}, ...]
        # Type annotation: result is list[dict[str, Any]]
        # getEdges can return dict, str, or DataFrame, but with fmt="py" it returns dict
        if isinstance(edges, list):
            # Type narrowing: after isinstance check, we know it's a list
            # Use cast to help type checker understand the elements are dicts
            result = cast(list[dict[str, Any]], edges)
        elif isinstance(edges, dict):
            # If it's a single dict, wrap it in a list
            result = [cast(dict[str, Any], edges)]
        else:
            # Fallback for unexpected types
            result: list[dict[str, Any]] = []

        # Apply limit if specified (client-side since pyTigerGraph 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 pyTigerGraph: {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.conn.getVertexTypes(force=True)

            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, 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,
    class_name,
    match_keys,
    keep_keys,
    flatten=False,
    filters: list | dict | None = None,
):
    """
    Check which documents from batch are present in the database.
    """
    try:
        present_docs = {}

        for i, doc in enumerate(batch):
            vertex_id = self._extract_id(doc, match_keys)
            if not vertex_id:
                continue

            try:
                vertex_data = self.conn.getVerticesById(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 = {}

                    for key in keep_keys:
                        if key == "id":
                            filtered_doc[key] = vertex_id
                        elif key in vertex_attrs:
                            filtered_doc[key] = vertex_attrs[key]

                    present_docs[i] = [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.conn.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, clean_start=False)

Initialize database with schema definition.

Follows the same pattern as ArangoDB: 1. Clean if needed 2. Create vertex and edge types globally (required before CREATE GRAPH) 3. Create graph with vertices and edges explicitly attached 4. Define indexes

If any step fails, the graph will be cleaned up gracefully.

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

def init_db(self, schema: Schema, clean_start=False):
    """
    Initialize database with schema definition.

    Follows the same pattern as ArangoDB:
    1. Clean if needed
    2. Create vertex and edge types globally (required before CREATE GRAPH)
    3. Create graph with vertices and edges explicitly attached
    4. 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")

    try:
        if clean_start:
            try:
                # Only delete the current graph, not all graphs or global vertex/edge types
                # This ensures we don't affect other graphs that might share vertex/edge types
                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 clean_start for graph '{graph_name}': {clean_error}",
                    exc_info=True,
                )
                # Continue - may be first run or already clean, schema will be recreated anyway

        # Step 1: Create vertex and edge types globally first
        # These must exist before they can be included in CREATE GRAPH
        logger.debug(
            f"Creating vertex and edge types globally for graph '{graph_name}'"
        )
        try:
            vertex_names = self._create_vertex_types_global(schema.vertex_config)

            # Initialize edges before creating edge types
            # This sets edge._source and edge._target to dbnames (required for GSQL)
            edges_to_create = list(schema.edge_config.edges_list(include_aux=True))
            for edge in edges_to_create:
                edge.finish_init(schema.vertex_config)

            # Verify all vertices referenced by edges were created
            created_vertex_set = set(vertex_names)
            for edge in edges_to_create:
                if edge._source not in created_vertex_set:
                    raise ValueError(
                        f"Edge '{edge.relation}' references source vertex '{edge._source}' "
                        f"which was not created. Created vertices: {vertex_names}"
                    )
                if edge._target not in created_vertex_set:
                    raise ValueError(
                        f"Edge '{edge.relation}' references target vertex '{edge._target}' "
                        f"which was not created. Created vertices: {vertex_names}"
                    )

            edge_names = self._create_edge_types_global(edges_to_create)
            logger.debug(
                f"Created {len(vertex_names)} vertex types and {len(edge_names)} edge types"
            )
        except Exception as type_error:
            logger.error(
                f"Failed to create vertex/edge types for graph '{graph_name}': {type_error}",
                exc_info=True,
            )
            raise

        # Step 2: Create graph with vertices and edges explicitly attached
        try:
            if not self.graph_exists(graph_name):
                logger.debug(f"Creating graph '{graph_name}' with types in init_db")
                try:
                    self.create_database(
                        graph_name,
                        vertex_names=vertex_names,
                        edge_names=edge_names,
                    )
                    graph_created = True
                    logger.info(f"Successfully created 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")
                # If graph already exists, associate types via ALTER GRAPH
                try:
                    self.define_vertex_collections(schema.vertex_config)
                    # Ensure edges are initialized before defining collections
                    edges_for_collections = list(
                        schema.edge_config.edges_list(include_aux=True)
                    )
                    for edge in edges_for_collections:
                        if edge._source is None or edge._target is None:
                            edge.finish_init(schema.vertex_config)
                    self.define_edge_collections(edges_for_collections)
                except Exception as define_error:
                    logger.warning(
                        f"Could not define collections for existing graph '{graph_name}': {define_error}",
                        exc_info=True,
                    )
                    # Continue - graph exists, collections may already be defined
        except Exception as graph_error:
            logger.error(
                f"Error during graph creation/verification for '{graph_name}': {graph_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, collection_name=None, match_keys_source=('_key',), match_keys_target=('_key',), filter_uniques=True, uniq_weight_fields=None, uniq_weight_collections=None, upsert_option=False, 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 of edge documents (tuples or dicts)	required
source_class		Source vertex type name	required
target_class		Target vertex type name	required
relation_name		Edge type/relation name	required
collection_name		Alternative edge collection name (used if relation_name is None)	None
match_keys_source		Keys to match source vertices	('_key',)
match_keys_target		Keys to match target vertices	('_key',)
filter_uniques		If True, filter duplicate edges	True
uniq_weight_fields		Fields to consider for uniqueness (not used in TigerGraph)	None
uniq_weight_collections		Collections to consider for uniqueness (not used in TigerGraph)	None
upsert_option		If True, use upsert (default behavior in TigerGraph)	False
head		Optional limit on number of edges to insert	None
**kwargs		Additional options: - dry: If True, don't execute the query	{}

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

def insert_edges_batch(
    self,
    docs_edges,
    source_class,
    target_class,
    relation_name,
    collection_name=None,
    match_keys_source=("_key",),
    match_keys_target=("_key",),
    filter_uniques=True,
    uniq_weight_fields=None,
    uniq_weight_collections=None,
    upsert_option=False,
    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": {...}}, ...]

    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
        collection_name: Alternative edge collection name (used if relation_name is None)
        match_keys_source: Keys to match source vertices
        match_keys_target: Keys to match target vertices
        filter_uniques: If True, filter duplicate edges
        uniq_weight_fields: Fields to consider for uniqueness (not used in TigerGraph)
        uniq_weight_collections: Collections to consider for uniqueness (not used in TigerGraph)
        upsert_option: If True, use upsert (default behavior in TigerGraph)
        head: Optional limit on number of edges to insert
        **kwargs: Additional options:
            - dry: If True, don't execute the query
    """
    dry = kwargs.pop("dry", False)
    if dry:
        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
    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 the edge upsert payload
        payload = self._generate_edge_upsert_payload(
            normalized_edges,
            source_class,
            target_class,
            edge_type,
            match_keys_src,
            match_keys_tgt,
        )

        # Check if payload has any edges
        source_vertices = payload.get("edges", {}).get(source_class, {})
        if not source_vertices:
            logger.warning(f"No valid edges to upsert for edge type {edge_type}")
            return

        # Build REST++ endpoint URL
        host = f"{self.config.url_without_port}:{self.config.port}"
        graph_name = self.config.database
        if not graph_name:
            raise ValueError("Graph name (database) must be configured")

        # Send the upsert request with username/password authentication
        result = self._upsert_data(
            payload,
            host,
            graph_name,
            username=self.config.username,
            password=self.config.password,
        )

        if result.get("error"):
            logger.error(
                f"Error upserting edges of type {edge_type}: {result.get('message')}"
            )
        else:
            # Count edges in payload
            edge_count = 0
            for source_edges in source_vertices.values():
                if edge_type in source_edges:
                    if target_class in source_edges[edge_type]:
                        edge_count += len(source_edges[edge_type][target_class])
            logger.debug(
                f"Upserted {edge_count} edges of type {edge_type}: {result}"
            )
            return result

    except Exception as e:
        logger.error(f"Error batch inserting edges: {e}")

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, class_name):
    """
    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, 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,
    class_name,
    match_keys,
    keep_keys,
    filters: list | dict | None = None,
):
    """
    Return documents from batch that are NOT present in database.
    """
    present_docs_indices = self.fetch_present_documents(
        batch=batch,
        class_name=class_name,
        match_keys=match_keys,
        keep_keys=keep_keys,
        flatten=False,
        filters=filters,
    )

    absent_indices = sorted(
        set(range(len(batch))) - set(present_docs_indices.keys())
    )
    return [batch[i] for i in absent_indices]

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

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

        # Build REST++ endpoint URL
        host = f"{self.config.url_without_port}:{self.config.port}"
        graph_name = self.config.database
        if not graph_name:
            raise ValueError("Graph name (database) must be configured")

        # Send the upsert request with username/password authentication
        result = self._upsert_data(
            payload,
            host,
            graph_name,
            username=self.config.username,
            password=self.config.password,
        )

        if result.get("error"):
            logger.error(
                f"Error upserting vertices to {class_name}: {result.get('message')}"
            )
            # Fallback to individual operations
            self._fallback_individual_upsert(docs, class_name, match_keys)
        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}")
        # Fallback to individual operations
        self._fallback_individual_upsert(docs, class_name, match_keys)