graflo 1.3.7__py3-none-any.whl

graflo/db/arango/query.py

@@ -0,0 +1,180 @@
+"""ArangoDB query utilities for graph operations.
+
+This module provides utility functions for executing and profiling AQL queries
+in ArangoDB. It includes functions for basic query execution, query profiling,
+and field fetching operations.
+
+Key Functions:
+    - basic_query: Execute a basic AQL query with configurable parameters
+    - profile_query: Profile query execution and save results
+    - fetch_fields_query: Generate and execute field-fetching queries
+
+Example:
+    >>> cursor = basic_query("FOR doc IN users RETURN doc", db_name="mydb")
+    >>> profile_query("FOR doc IN users RETURN doc", nq=1, profile_times=3, fpath=".")
+"""
+
+import gzip
+import json
+import logging
+from os.path import join
+
+from arango import ArangoClient
+
+from graflo.filter.onto import Expression
+from graflo.onto import DBFlavor
+
+logger = logging.getLogger(__name__)
+
+
+def basic_query(
+    query,
+    port=8529,
+    hostname="127.0.0.1",
+    cred_name="root",
+    cred_pass="123",
+    db_name="_system",
+    profile=False,
+    batch_size=10000,
+    bind_vars=None,
+):
+    """Execute a basic AQL query in ArangoDB.
+
+    This function provides a simple interface for executing AQL queries with
+    configurable connection parameters and query options.
+
+    Args:
+        query: AQL query string to execute
+        port: ArangoDB server port
+        hostname: ArangoDB server hostname
+        cred_name: Database username
+        cred_pass: Database password
+        db_name: Database name
+        profile: Whether to enable query profiling
+        batch_size: Size of result batches
+        bind_vars: Query bind variables
+
+    Returns:
+        Cursor: ArangoDB cursor for the query results
+    """
+    hosts = f"http://{hostname}:{port}"
+    client = ArangoClient(hosts=hosts)
+
+    sys_db = client.db(db_name, username=cred_name, password=cred_pass)
+    cursor = sys_db.aql.execute(
+        query,
+        profile=profile,
+        stream=True,
+        batch_size=batch_size,
+        bind_vars=bind_vars,
+    )
+    return cursor
+
+
+def profile_query(query, nq, profile_times, fpath, limit=None, **kwargs):
+    """Profile AQL query execution and save results.
+
+    This function executes a query multiple times with profiling enabled and
+    saves both the profiling results and query results to files.
+
+    Args:
+        query: AQL query string to profile
+        nq: Query number for file naming
+        profile_times: Number of times to profile the query
+        fpath: Path to save results
+        limit: Optional limit on query results
+        **kwargs: Additional query parameters passed to basic_query
+
+    Note:
+        Results are saved in two formats:
+        - Profiling results: query{nq}_profile{limit}.json
+        - Query results: query{nq}_result{limit}_batch_{n}.json.gz
+    """
+    limit_str = f"_limit_{limit}" if limit else ""
+    if profile_times:
+        logger.info(f"starting profiling: {limit}")
+        profiling = []
+        for n in range(profile_times):
+            cursor = basic_query(query, profile=True, **kwargs)
+            profiling += [cursor.profile()]
+            cursor.close()
+        with open(join(fpath, f"query{nq}_profile{limit_str}.json"), "w") as fp:
+            json.dump(profiling, fp, indent=4)
+
+    logger.info(f"starting actual query at {limit}")
+
+    cnt = 0
+    cursor = basic_query(query, **kwargs)
+    chunk = list(cursor.batch())
+    with gzip.open(
+        join(fpath, f"./query{nq}_result{limit_str}_batch_{cnt}.json.gz"),
+        "wt",
+        encoding="ascii",
+    ) as fp:
+        json.dump(chunk, fp, indent=4)
+
+    while cursor.has_more():
+        cnt += 1
+        with gzip.open(
+            join(fpath, f"./query{nq}_result{limit_str}_batch_{cnt}.json.gz"),
+            "wt",
+            encoding="ascii",
+        ) as fp:
+            chunk = list(cursor.fetch()["batch"])
+            json.dump(chunk, fp, indent=4)
+
+
+def fetch_fields_query(
+    collection_name,
+    docs,
+    match_keys,
+    keep_keys,
+    filters: list | dict | None = None,
+):
+    """Generate and execute a field-fetching AQL query.
+
+    This function generates an AQL query to fetch specific fields from documents
+    that match the given criteria. It supports filtering and field projection.
+
+    Args:
+        collection_name: Collection to query
+        docs: List of documents to match against
+        match_keys: Keys to use for matching documents
+        keep_keys: Keys to return in the result
+        filters: Additional query filters
+
+    Returns:
+        str: Generated AQL query string
+
+    Example:
+        >>> query = fetch_fields_query(
+        ...     "users",
+        ...     [{"email": "user@example.com"}],
+        ...     ["email"],
+        ...     ["name", "age"]
+        ... )
+    """
+    docs_ = [{k: doc[k] for k in match_keys if k in doc} for doc in docs]
+    for i, doc in enumerate(docs_):
+        doc.update({"__i": i})
+
+    docs_str = json.dumps(docs_)
+
+    match_str = " &&".join([f" _cdoc['{key}'] == _doc['{key}']" for key in match_keys])
+
+    keep_clause = f"KEEP(_x, {list(keep_keys)})" if keep_keys is not None else "_x"
+
+    if filters is not None:
+        ff = Expression.from_dict(filters)
+        extrac_filter_clause = f" && {ff(doc_name='_cdoc', kind=DBFlavor.ARANGO)}"
+    else:
+        extrac_filter_clause = ""
+
+    q0 = f"""
+        FOR _cdoc in {collection_name}
+            FOR _doc in {docs_str}
+                FILTER {match_str} {extrac_filter_clause}      
+                COLLECT i = _doc['__i'] into _group = _cdoc 
+                LET gp = (for _x in _group return {keep_clause})                                
+                    RETURN {{'__i' : i, '_group': gp}}"""
+    return q0

graflo/db/arango/util.py

@@ -0,0 +1,88 @@
+"""ArangoDB utility functions for graph operations.
+
+This module provides utility functions for working with ArangoDB graphs and
+queries. It includes functions for edge definition, filter rendering, and
+query generation.
+
+Key Functions:
+    - define_extra_edges: Generate queries for creating derived edges
+    - render_filters: Convert filter expressions to AQL filter clauses
+
+Example:
+    >>> query = define_extra_edges(edge_config)
+    >>> filter_clause = render_filters({"field": "value"}, doc_name="d")
+"""
+
+import logging
+
+from graflo.architecture.edge import Edge
+from graflo.filter.onto import Clause, Expression
+from graflo.onto import ExpressionFlavor
+
+logger = logging.getLogger(__name__)
+
+
+def define_extra_edges(g: Edge):
+    """Generate AQL query for creating derived edges.
+
+    This function creates a query to generate edges from source to target
+    vertices through an intermediate vertex, copying properties from the
+    intermediate vertex to the new edge.
+
+    Args:
+        g: Edge configuration containing source, target, and intermediate
+            vertex information
+
+    Returns:
+        str: AQL query string for creating the derived edges
+
+    Example:
+        >>> edge = Edge(source="user", target="post", by="comment")
+        >>> query = define_extra_edges(edge)
+        >>> # Generates query to create user->post edges through comments
+    """
+    ucol, vcol, wcol = g.source, g.target, g.by
+    weight = g.weight_dict
+    s = f"""FOR w IN {wcol}
+        LET uset = (FOR u IN 1..1 INBOUND w {ucol}_{wcol}_edges RETURN u)
+        LET vset = (FOR v IN 1..1 INBOUND w {vcol}_{wcol}_edges RETURN v)
+        FOR u in uset
+        FOR v in vset
+    """
+    s_ins_ = ", ".join([f"{v}: w.{k}" for k, v in weight.items()])
+    s_ins_ = f"_from: u._id, _to: v._id, {s_ins_}"
+    s_ins = f"          INSERT {{{s_ins_}}} "
+    s_last = f"IN {ucol}_{vcol}_edges"
+    query0 = s + s_ins + s_last
+    return query0
+
+
+def render_filters(filters: None | list | dict | Clause = None, doc_name="d") -> str:
+    """Convert filter expressions to AQL filter clauses.
+
+    This function converts filter expressions into AQL filter clauses that
+    can be used in queries. It supports various filter types and formats.
+
+    Args:
+        filters: Filter expression to convert
+        doc_name: Name of the document variable in the query
+
+    Returns:
+        str: AQL filter clause string
+
+    Example:
+        >>> filters = {"field": "value", "age": {"$gt": 18}}
+        >>> clause = render_filters(filters, doc_name="user")
+        >>> # Returns: "FILTER user.field == 'value' && user.age > 18"
+    """
+    if filters is not None:
+        if not isinstance(filters, Clause):
+            ff = Expression.from_dict(filters)
+        else:
+            ff = filters
+        literal_condition = ff(doc_name=doc_name, kind=ExpressionFlavor.ARANGO)
+        filter_clause = f"FILTER {literal_condition}"
+    else:
+        filter_clause = ""
+
+    return filter_clause

graflo/db/conn.py

@@ -0,0 +1,377 @@
+"""Abstract database connection interface for graph databases.
+
+This module defines the abstract interface for database connections, providing
+a unified API for different graph database implementations. It includes methods
+for database management, graph structure operations, and data manipulation.
+
+Key Components:
+    - Connection: Abstract base class for database connections
+    - ConnectionType: Type variable for connection implementations
+
+The connection interface supports:
+    - Database/Graph creation and deletion
+    - Graph structure management (vertex types, edge types)
+    - Index definition
+    - Document operations (insert, update, fetch)
+    - Edge operations
+    - Aggregation queries
+
+Database Organization Terminology:
+    Different databases organize graph data differently:
+
+    - ArangoDB:
+        * Database: Top-level container (like a schema)
+        * Collections: Container for vertices (vertex collections)
+        * Edge Collections: Container for edges
+        * Graph: Named graph that connects vertex and edge collections
+
+    - Neo4j:
+        * Database: Top-level container
+        * 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
+
+    - 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
+
+    When using the Connection interface, the terms "vertex type" and "edge type"
+    are used generically to refer to the appropriate concept in each database.
+
+Example:
+    >>> class MyConnection(Connection):
+    ...     def create_database(self, name: str):
+    ...         # Implementation
+    ...     def execute(self, query, **kwargs):
+    ...         # Implementation
+"""
+
+import abc
+import logging
+from typing import TypeVar
+
+from graflo.architecture.edge import Edge
+from graflo.architecture.schema import Schema
+from graflo.architecture.vertex import VertexConfig
+from graflo.onto import AggregationType
+
+logger = logging.getLogger(__name__)
+ConnectionType = TypeVar("ConnectionType", bound="Connection")
+
+
+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
+
+    # @abc.abstractmethod
+    # def define_vertex_collections(self, graph_config, vertex_config):
+    #     pass
+    #
+    # @abc.abstractmethod
+    # def define_edge_collections(self, graph_config):
+    #     pass
+
+    # @abc.abstractmethod
+    # def create_collection_if_absent(self, g, vcol, index, unique=True):
+    #     pass

graflo/db/connection/__init__.py

@@ -0,0 +1,6 @@
+from .onto import DBConfig, DBType
+
+__all__ = [
+    "DBConfig",
+    "DBType",
+]

graflo/db/connection/config_mapping.py

@@ -0,0 +1,18 @@
+from typing import Dict, Type
+
+from .onto import (
+    ArangoConfig,
+    DBConfig,
+    DBType,
+    Neo4jConfig,
+    PostgresConfig,
+    TigergraphConfig,
+)
+
+# Define this mapping in a separate file to avoid circular imports
+DB_TYPE_MAPPING: Dict[DBType, Type[DBConfig]] = {
+    DBType.ARANGO: ArangoConfig,
+    DBType.NEO4J: Neo4jConfig,
+    DBType.TIGERGRAPH: TigergraphConfig,
+    DBType.POSTGRES: PostgresConfig,
+}