graflo 1.1.0__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 DBFlavor
+
+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=DBFlavor.ARANGO)
+        filter_clause = f"FILTER {literal_condition}"
+    else:
+        filter_clause = ""
+
+    return filter_clause

graflo/db/connection.py

@@ -0,0 +1,304 @@
+"""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, collection operations, and data manipulation.
+
+Key Components:
+    - Connection: Abstract base class for database connections
+    - ConnectionType: Type variable for connection implementations
+
+The connection interface supports:
+    - Database creation and deletion
+    - Collection management
+    - Index definition
+    - Document operations (insert, update, fetch)
+    - Edge operations
+    - Aggregation queries
+
+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 operations, collection management,
+    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_collections(self, schema: Schema):
+        """Define collections based on the schema.
+
+        Args:
+            schema: Schema containing collection definitions
+        """
+        pass
+
+    @abc.abstractmethod
+    def delete_collections(self, cnames=(), gnames=(), delete_all=False):
+        """Delete collections from the database.
+
+        Args:
+            cnames: Collection names to delete
+            gnames: Graph names to delete
+            delete_all: Whether to delete all collections
+        """
+        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 collection
+            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 class
+            target_class: Target vertex class
+            relation_name: Name of the relation
+            collection_name: Name of the edge collection
+            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: Collections to consider for uniqueness
+            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 collection
+
+        Returns:
+            list: Inserted documents
+        """
+        pass
+
+    @abc.abstractmethod
+    def fetch_docs(self, class_name, filters, limit, return_keys, unset_keys):
+        """Fetch documents from a collection.
+
+        Args:
+            class_name: Name of the collection
+            filters: Query filters
+            limit: Maximum number of documents to return
+            return_keys: Keys to return
+            unset_keys: Keys to unset
+
+        Returns:
+            list: Fetched documents
+        """
+        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/manager.py

@@ -0,0 +1,104 @@
+"""Database connection manager for graph databases.
+
+This module provides a connection manager for handling database connections
+to different graph database implementations (ArangoDB, Neo4j). It manages
+connection lifecycle and configuration.
+
+Key Components:
+    - ConnectionManager: Main class for managing database connections
+    - ConnectionKind: Enum for supported database types
+
+The manager supports:
+    - Multiple database types (ArangoDB, Neo4j)
+    - Connection configuration
+    - Context manager interface
+    - Automatic connection cleanup
+
+Example:
+    >>> with ConnectionManager(secret_path="config.json") as conn:
+    ...     conn.execute("FOR doc IN collection RETURN doc")
+"""
+
+from typing import Optional
+
+from suthing import ConfigFactory, ConnectionKind, ProtoConnectionConfig
+
+from graflo.db.arango.conn import ArangoConnection
+from graflo.db.neo4j.conn import Neo4jConnection
+
+
+class ConnectionManager:
+    """Manager for database connections.
+
+    This class manages database connections to different graph database
+    implementations. It provides a context manager interface for safe
+    connection handling and automatic cleanup.
+
+    Attributes:
+        conn_class_mapping: Mapping of connection types to connection classes
+        config: Connection configuration
+        working_db: Current working database name
+        conn: Active database connection
+    """
+
+    conn_class_mapping = {
+        ConnectionKind.ARANGO: ArangoConnection,
+        ConnectionKind.NEO4J: Neo4jConnection,
+    }
+
+    def __init__(
+        self,
+        secret_path=None,
+        args=None,
+        connection_config: Optional[ProtoConnectionConfig] = None,
+        **kwargs,
+    ):
+        """Initialize the connection manager.
+
+        Args:
+            secret_path: Path to configuration file
+            args: Command line arguments
+            connection_config: Optional connection configuration
+            **kwargs: Additional configuration parameters
+        """
+        self.config: ProtoConnectionConfig = (
+            ConfigFactory.create_config(secret_path, args)
+            if connection_config is None
+            else 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
+        """
+        cls = self.conn_class_mapping[self.config.connection_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()

graflo/db/neo4j/__init__.py

@@ -0,0 +1,16 @@
+"""Neo4j database implementation.
+
+This package provides Neo4j-specific implementations of the database interface,
+including connection management, query execution, and utility functions.
+
+Key Components:
+    - Neo4jConnection: Neo4j connection implementation
+    - Query: Cypher query execution and profiling
+    - Util: Neo4j-specific utility functions
+
+Example:
+    >>> from graflo.db.neo4j import Neo4jConnection
+    >>> conn = Neo4jConnection(config)
+    >>> result = conn.execute("MATCH (n:User) RETURN n")
+    >>> nodes = result.data()
+"""