cognite-neat 0.87.4__py3-none-any.whl → 0.88.0__py3-none-any.whl

cognite/neat/legacy/graph/stores/_base.py

@@ -1,400 +0,0 @@
-import logging
-import sys
-import time
-from abc import ABC, abstractmethod
-from collections.abc import Iterable, Iterator
-from pathlib import Path
-from typing import Literal, TypeAlias, cast
-
-import pandas as pd
-from prometheus_client import Gauge, Summary
-from rdflib import Graph, Namespace, URIRef
-from rdflib.query import Result, ResultRow
-
-from cognite.neat.constants import DEFAULT_NAMESPACE, get_default_prefixes
-from cognite.neat.legacy.graph.models import Triple
-from cognite.neat.legacy.graph.stores._rdf_to_graph import rdf_file_to_graph
-from cognite.neat.legacy.rules.models.rules import Rules
-
-if sys.version_info >= (3, 11):
-    from typing import Self
-else:
-    from typing_extensions import Self
-
-prom_qsm = Summary("store_query_time_summary_legacy", "Time spent processing queries", ["query"])
-prom_sq = Gauge("store_single_query_time_legacy", "Time spent processing a single query", ["query"])
-
-MIMETypes: TypeAlias = Literal[
-    "application/rdf+xml",
-    "text/turtle",
-    "application/n-triple",
-    "application/n-quads",
-    "application/trig",
-]
-
-
-class NeatGraphStoreBase(ABC):
-    """NeatGraphStore is a class that stores the graph and provides methods to read/write data it contains
-
-
-    Args:
-        graph : Instance of rdflib.Graph class for graph storage
-        base_prefix : Used as a base prefix for graph namespace,
-                      allowing querying graph data using a short form of a URI
-        namespace : Namespace (aka URI) used to resolve any relative URI in the graph
-        prefixes : Dictionary of additional prefixes used and bounded to the graph
-    """
-
-    rdf_store_type: str
-
-    def __init__(
-        self,
-        graph: Graph | None = None,
-        base_prefix: str = "",  # usually empty
-        namespace: Namespace = DEFAULT_NAMESPACE,
-        prefixes: dict[str, Namespace] | None = None,
-    ):
-        self.graph = graph or Graph()
-        self.base_prefix: str = base_prefix
-        self.namespace: Namespace = namespace
-        self.prefixes: dict[str, Namespace] = prefixes or get_default_prefixes()
-
-        self.rdf_store_query_url: str | None = None
-        self.rdf_store_update_url: str | None = None
-        self.returnFormat: str | None = None
-        self.df_cache: pd.DataFrame | None = None
-        self.internal_storage_dir: Path | None = None
-        self.graph_name: str | None = None
-        self.internal_storage_dir_orig: Path | None = None
-        self.storage_dirs_to_delete: list[Path] = []
-        self.queries = _Queries(self)
-
-    @classmethod
-    def from_rules(cls, rules: Rules) -> Self:
-        """
-        Creates a new instance of NeatGraphStore from TransformationRules and runs the .init_graph() method on it.
-
-        Args:
-            rules: TransformationRules object containing information about the graph store.
-
-        Returns:
-            An instantiated instance of NeatGraphStore
-
-        """
-        if rules.metadata.namespace is None:
-            namespace = DEFAULT_NAMESPACE
-        else:
-            namespace = rules.metadata.namespace
-        store = cls(prefixes=rules.prefixes, namespace=namespace)
-        store.init_graph(base_prefix=rules.metadata.prefix)
-        return store
-
-    @abstractmethod
-    def _set_graph(self) -> None:
-        raise NotImplementedError()
-
-    def init_graph(
-        self,
-        rdf_store_query_url: str | None = None,
-        rdf_store_update_url: str | None = None,
-        graph_name: str | None = None,
-        base_prefix: str | None = None,
-        returnFormat: str = "csv",
-        internal_storage_dir: Path | None = None,
-    ):
-        """Initializes the graph.
-
-        Args:
-            rdf_store_query_url : URL towards which SPARQL query is executed, by default None
-            rdf_store_update_url : URL towards which SPARQL update is executed, by default None
-            graph_name : Name of graph, by default None
-            base_prefix : Base prefix for graph namespace to change if needed, by default None
-            returnFormat : Transport format of graph data between, by default "csv"
-            internal_storage_dir : Path to directory where internal storage is located,
-                                   by default None (in-memory storage).
-
-        !!! note "internal_storage_dir"
-            Used only for Oxigraph
-        """
-        logging.info("Initializing NeatGraphStore")
-        self.rdf_store_query_url = rdf_store_query_url
-        self.rdf_store_update_url = rdf_store_update_url
-        self.graph_name = graph_name
-        self.returnFormat = returnFormat
-        self.internal_storage_dir = Path(internal_storage_dir) if internal_storage_dir else None
-        self.internal_storage_dir_orig = (
-            self.internal_storage_dir if self.internal_storage_dir_orig is None else self.internal_storage_dir_orig
-        )
-
-        self._set_graph()
-
-        if self.prefixes:
-            for prefix, namespace in self.prefixes.items():
-                logging.info("Adding prefix %s with namespace %s", prefix, namespace)
-                self.graph.bind(prefix, namespace)
-
-        if base_prefix:
-            self.base_prefix = base_prefix
-
-        self.graph.bind(self.base_prefix, self.namespace)
-        logging.info("Adding prefix %s with namespace %s", self.base_prefix, self.namespace)
-        logging.info("Graph initialized")
-
-    def reinitialize_graph(self):
-        """Reinitialize the graph."""
-        self.init_graph(
-            self.rdf_store_query_url,
-            self.rdf_store_update_url,
-            self.graph_name,
-            self.base_prefix,
-            self.returnFormat,
-            self.internal_storage_dir,
-        )
-
-    def upsert_prefixes(self, prefixes: dict[str, Namespace]) -> None:
-        """Adds prefixes to the graph store."""
-        self.prefixes.update(prefixes)
-        for prefix, namespace in prefixes.items():
-            logging.info("Adding prefix %s with namespace %s", prefix, namespace)
-            self.graph.bind(prefix, namespace)
-
-    def close(self) -> None:
-        """Closes the graph."""
-        # Can be overridden in subclasses
-        return None
-
-    def restart(self) -> None:
-        """Restarts the graph"""
-        # Can be overridden in subclasses
-        return None
-
-    def import_from_file(
-        self,
-        graph_file: Path,
-        mime_type: MIMETypes = "application/rdf+xml",
-        add_base_iri: bool = True,
-    ) -> None:
-        """Imports graph data from file.
-
-        Args:
-            graph_file : File path to file containing graph data, by default None
-            mime_type : MIME type of graph data, by default "application/rdf+xml"
-            add_base_iri : Add base IRI to graph, by default True
-        """
-        if add_base_iri:
-            self.graph = rdf_file_to_graph(
-                self.graph,
-                graph_file,
-                base_namespace=self.namespace,
-                prefixes=self.prefixes,
-            )
-        else:
-            self.graph = rdf_file_to_graph(self.graph, graph_file, prefixes=self.prefixes)
-        return None
-
-    def get_graph(self) -> Graph:
-        """Returns the graph."""
-        return self.graph
-
-    def set_graph(self, graph: Graph):
-        """Sets the graph."""
-        self.graph = graph
-
-    def query(self, query: str) -> Result:
-        """Returns the result of the query."""
-        start_time = time.perf_counter()
-        result = self.graph.query(query)
-        stop_time = time.perf_counter()
-        elapsed_time = stop_time - start_time
-        prom_qsm.labels("query").observe(elapsed_time)
-        prom_sq.labels("query").set(elapsed_time)
-        return result
-
-    def serialize(self, *args, **kwargs):
-        """Serializes the graph."""
-        return self.graph.serialize(*args, **kwargs)
-
-    def query_delayed(self, query) -> Iterable[Triple]:
-        """Returns the result of the query, but does not execute it immediately.
-
-        The query is not executed until the result is iterated over.
-
-        Args:
-            query: SPARQL query to execute
-
-        Returns:
-            An iterable of triples
-
-        """
-        return _DelayedQuery(self.graph, query)
-
-    @abstractmethod
-    def drop(self) -> None:
-        """Drops the graph."""
-        raise NotImplementedError()
-
-    def garbage_collector(self) -> None:
-        """Garbage collection of the graph store."""
-        # Can be overridden in subclasses
-        return None
-
-    def query_to_dataframe(
-        self,
-        query: str,
-        column_mapping: dict | None = None,
-        save_to_cache: bool = False,
-        index_column: str = "instance",
-    ) -> pd.DataFrame:
-        """Returns the result of the query as a dataframe.
-
-        Args:
-            query: SPARQL query to execute
-            column_mapping: Columns name mapping, by default None
-            save_to_cache: Save result of query to cache, by default False
-            index_column: Indexing column , by default "instance"
-
-        Returns:
-            Dataframe with result of query
-        """
-
-        if column_mapping is None:
-            column_mapping = {0: "instance", 1: "property", 2: "value"}
-
-        result = self.graph.query(query, DEBUG=False)
-        df_cache = pd.DataFrame(list(result))
-        df_cache.rename(columns=column_mapping, inplace=True)
-        df_cache[index_column] = df_cache[index_column].apply(lambda x: str(x))
-        if save_to_cache:
-            self.df_cache = df_cache
-        return df_cache
-
-    def commit(self):
-        """Commits the graph."""
-        self.graph.commit()
-
-    def get_df(self) -> pd.DataFrame:
-        """Returns the cached dataframe."""
-        if self.df_cache is None:
-            raise ValueError("Cache is empty. Run query_to_dataframe() first with save_to_cache.")
-        return self.df_cache
-
-    def get_instance_properties_from_cache(self, instance_id: str) -> pd.DataFrame:
-        """Returns the properties of an instance."""
-        if self.df_cache is None:
-            raise ValueError("Cache is empty. Run query_to_dataframe() first with save_to_cache.")
-        return self.df_cache.loc[self.df_cache["instance"] == instance_id]
-
-    def print_triples(self):
-        """Prints the triples of the graph."""
-        for subj, pred, obj in self.graph:
-            logging.info(f"Triple: {subj} {pred} {obj}")
-
-    def diagnostic_report(self):
-        """Returns the dictionary representation graph diagnostic data ."""
-        return {
-            "rdf_store_type": self.rdf_store_type,
-            "base_prefix": self.base_prefix,
-            "namespace": self.namespace,
-            "prefixes": self.prefixes,
-            "internal_storage_dir": self.internal_storage_dir,
-            "rdf_store_query_url": self.rdf_store_query_url,
-            "rdf_store_update_url": self.rdf_store_update_url,
-        }
-
-    def add_triples(
-        self,
-        triples: list[Triple] | set[Triple],
-        batch_size: int = 10_000,
-        verbose: bool = False,
-    ):
-        """Adds triples to the graph store in batches.
-
-        Args:
-            triples: list of triples to be added to the graph store
-            batch_size: Batch size of triples per commit, by default 10_000
-            verbose: Verbose mode, by default False
-        """
-
-        commit_counter = 0
-        if verbose:
-            logging.info(f"Committing total of {len(triples)} triples to knowledge graph!")
-        total_number_of_triples = len(triples)
-        number_of_uploaded_triples = 0
-
-        def check_commit(force_commit: bool = False):
-            """Commit nodes to the graph if batch counter is reached or if force_commit is True"""
-            nonlocal commit_counter
-            nonlocal number_of_uploaded_triples
-            if force_commit:
-                number_of_uploaded_triples += commit_counter
-                self.graph.commit()
-                if verbose:
-                    logging.info(f"Committed {number_of_uploaded_triples} of {total_number_of_triples} triples")
-                return
-            commit_counter += 1
-            if commit_counter >= batch_size:
-                number_of_uploaded_triples += commit_counter
-                self.graph.commit()
-                if verbose:
-                    logging.info(f"Committed {number_of_uploaded_triples} of {total_number_of_triples} triples")
-                commit_counter = 0
-
-        for triple in triples:
-            self.graph.add(triple)
-            check_commit()
-
-        check_commit(force_commit=True)
-
-
-class _DelayedQuery(Iterable):
-    def __init__(self, graph_ref: Graph, query: str):
-        self.graph_ref = graph_ref
-        self.query = query
-
-    def __iter__(self) -> Iterator[Triple]:
-        start_time = time.perf_counter()
-        result = self.graph_ref.query(self.query)
-        stop_time = time.perf_counter()
-        elapsed_time = stop_time - start_time
-        prom_qsm.labels("query").observe(elapsed_time)
-        prom_sq.labels("query").set(elapsed_time)
-        return cast(Iterator[Triple], iter(result))
-
-
-class _Queries:
-    """Helper class for storing standard queries for the graph store."""
-
-    def __init__(self, store: NeatGraphStoreBase):
-        self.store = store
-
-    def list_instances_ids_of_class(self, class_uri: URIRef, limit: int = -1) -> list[URIRef]:
-        """Get instances ids for a given class
-
-        Args:
-            class_uri: Class for which instances are to be found
-            limit: Max number of instances to return, by default -1 meaning all instances
-
-        Returns:
-            List of class instance URIs
-        """
-        query_statement = "SELECT DISTINCT ?subject WHERE { ?subject a <class> .} LIMIT X".replace(
-            "class", class_uri
-        ).replace("LIMIT X", "" if limit == -1 else f"LIMIT {limit}")
-        return [cast(tuple, res)[0] for res in list(self.store.query(query_statement))]
-
-    def list_instances_of_type(self, class_uri: URIRef) -> list[ResultRow]:
-        """Get all triples for instances of a given class
-
-        Args:
-            class_uri: Class for which instances are to be found
-
-        Returns:
-            List of triples for instances of the given class
-        """
-        query = (
-            f"SELECT ?instance ?prop ?value "
-            f"WHERE {{ ?instance rdf:type <{class_uri}> . ?instance ?prop ?value . }} order by ?instance "
-        )
-        logging.info(query)
-        # Select queries gives an iterable of result rows
-        return cast(list[ResultRow], list(self.store.query(query)))

cognite/neat/legacy/graph/stores/_graphdb_store.py

@@ -1,52 +0,0 @@
-import logging
-
-import requests
-from rdflib import Graph, Namespace
-from rdflib.plugins.stores.sparqlstore import SPARQLUpdateStore
-
-from cognite.neat.constants import DEFAULT_NAMESPACE, get_default_prefixes
-
-from ._base import NeatGraphStoreBase
-
-
-class GraphDBStore(NeatGraphStoreBase):
-    """GraphDB is a class that stores the graph in a GraphDB instances and provides methods to
-    read/write data it contains
-
-
-    Args:
-        graph : Instance of rdflib.Graph class for graph storage
-        base_prefix : Used as a base prefix for graph namespace, allowing querying graph data using a shortform of a URI
-        namespace : Namespace (aka URI) used to resolve any relative URI in the graph
-        prefixes : Dictionary of additional prefixes used and bounded to the graph
-    """
-
-    rdf_store_type = "graphdb"
-
-    def __init__(
-        self,
-        graph: Graph | None = None,
-        base_prefix: str = "",  # usually empty
-        namespace: Namespace = DEFAULT_NAMESPACE,
-        prefixes: dict[str, Namespace] | None = None,
-    ):
-        prefixes = prefixes if prefixes else get_default_prefixes()
-        super().__init__(graph, base_prefix, namespace, prefixes)
-        self.graph_db_rest_url: str = "http://localhost:7200"
-
-    def _set_graph(self) -> None:
-        logging.info("Initializing graph store with GraphDB")
-        store = SPARQLUpdateStore(
-            query_endpoint=self.rdf_store_query_url,
-            update_endpoint=self.rdf_store_update_url,
-            returnFormat=self.returnFormat,
-            context_aware=False,
-            postAsEncoded=False,
-            autocommit=False,
-        )
-        self.graph = Graph(store=store)
-
-    def drop(self):
-        """Drops the graph."""
-        r = requests.delete(f"{self.rdf_store_query_url}/rdf-graphs/service?default")
-        logging.info(f"Dropped graph with state: {r.text}")

cognite/neat/legacy/graph/stores/_memory_store.py

@@ -1,43 +0,0 @@
-import logging
-
-from rdflib import Graph, Namespace
-
-from cognite.neat.constants import DEFAULT_NAMESPACE, get_default_prefixes
-
-from ._base import NeatGraphStoreBase
-
-
-class MemoryStore(NeatGraphStoreBase):
-    """MemoryStore is a class that stores the graph in memory using rdflib and provides
-    methods to read/write data it contains.
-
-
-    Args:
-        graph : Instance of rdflib.Graph class for graph storage
-        base_prefix : Used as a base prefix for graph namespace, allowing querying graph data using a shortform of a URI
-        namespace : Namespace (aka URI) used to resolve any relative URI in the graph
-        prefixes : Dictionary of additional prefixes used and bounded to the graph
-    """
-
-    rdf_store_type: str = "memory"
-
-    def __init__(
-        self,
-        graph: Graph | None = None,
-        base_prefix: str = "",  # usually empty
-        namespace: Namespace = DEFAULT_NAMESPACE,
-        prefixes: dict[str, Namespace] | None = None,
-    ):
-        prefixes = prefixes if prefixes else get_default_prefixes()  # Init repeated to get nice docstring
-        super().__init__(graph, base_prefix, namespace, prefixes)
-
-    def _set_graph(self):
-        logging.info("Initializing graph in memory")
-        self.graph = Graph()
-
-    def drop(self):
-        """Drops the graph."""
-        # In the case of in-memory graph, we just reinitialize the graph
-        # otherwise we would lose the prefixes and bindings, which fails
-        # workflow
-        self.reinitialize_graph()

cognite/neat/legacy/graph/stores/_oxigraph_store.py

@@ -1,151 +0,0 @@
-import logging
-import os
-import shutil
-from pathlib import Path
-
-from rdflib import Graph, Namespace
-
-from cognite.neat.constants import DEFAULT_NAMESPACE, get_default_prefixes
-from cognite.neat.utils.auxiliary import local_import
-
-from ._base import MIMETypes, NeatGraphStoreBase
-
-
-class OxiGraphStore(NeatGraphStoreBase):
-    """OxiGraph is a class that stores the graph using OxiGraph and provides methods to read/write data it contains
-
-
-    Args:
-        graph : Instance of rdflib.Graph class for graph storage
-        base_prefix : Used as a base prefix for graph namespace, allowing querying graph data using a shortform of a URI
-        namespace : Namespace (aka URI) used to resolve any relative URI in the graph
-        prefixes : Dictionary of additional prefixes used and bounded to the graph
-    """
-
-    rdf_store_type = "oxigraph"
-
-    def __init__(
-        self,
-        graph: Graph | None = None,
-        base_prefix: str = "",  # usually empty
-        namespace: Namespace = DEFAULT_NAMESPACE,
-        prefixes: dict[str, Namespace] | None = None,
-    ):
-        prefixes = prefixes if prefixes else get_default_prefixes()
-        super().__init__(graph, base_prefix, namespace, prefixes)
-
-    def _set_graph(self) -> None:
-        logging.info("Initializing Oxigraph store")
-        local_import("pyoxigraph", "oxi")
-        import pyoxigraph
-
-        from cognite.neat.graph.stores import _oxrdflib
-
-        # Adding support for both in-memory and file-based storage
-        for i in range(4):
-            try:
-                oxstore = pyoxigraph.Store(
-                    path=(str(self.internal_storage_dir) if self.internal_storage_dir else None)
-                )  # Store (Rust object) accepts only str as path and not Path.
-                break
-            except OSError as e:
-                if "lock" in str(e) and i < 3:
-                    # lock originated from another instance of the store
-                    logging.error("Error initializing Oxigraph store: %s", e)
-                else:
-                    raise e
-        else:
-            raise Exception("Error initializing Oxigraph store")
-
-        self.graph = Graph(store=_oxrdflib.OxigraphStore(store=oxstore))
-        self.graph.default_union = True
-        self.garbage_collector()
-
-    def close(self):
-        """Closes the graph."""
-        if self.graph is not None:
-            try:
-                self.graph.store._inner.flush()  # type: ignore[attr-defined]
-                self.graph.close(True)
-            except Exception as e:
-                logging.debug("Error closing graph: %s", e)
-
-    def restart(self):
-        """Restarts the graph"""
-        self.close()
-        self.reinitialize_graph()
-        logging.info("GraphStore restarted")
-
-    def import_from_file(
-        self,
-        graph_file: Path,
-        mime_type: MIMETypes = "application/rdf+xml",
-        add_base_iri: bool = True,
-    ) -> None:
-        """Imports graph data from file.
-
-        Args:
-            graph_file : File path to file containing graph data, by default None
-            mime_type : MIME type of the file, by default "application/rdf+xml"
-            add_base_iri : Add base IRI to the graph, by default True
-        """
-        if add_base_iri:
-            self.graph.store._inner.bulk_load(  # type: ignore[attr-defined]
-                str(graph_file), mime_type, base_iri=self.namespace
-            )
-        else:
-            self.graph.store._inner.bulk_load(str(graph_file), mime_type)  # type: ignore[attr-defined]
-        self.graph.store._inner.optimize()  # type: ignore[attr-defined]
-        return None
-
-    def drop(self):
-        try:
-            self.close()
-            # Due to the specifics of Oxigraph, storage directory cannot be deleted immediately
-            # after closing the graph and creating a new one
-            if self.internal_storage_dir.exists():
-                self.storage_dirs_to_delete.append(self.internal_storage_dir)
-            self.garbage_collector()
-
-        except Exception as e:
-            logging.error(f"Error dropping graph : {e}")
-
-    def garbage_collector(self):
-        """Garbage collection of the graph store."""
-        # delete all directories in self.storage_dirs_to_delete
-        for d in self.storage_dirs_to_delete:
-            shutil.rmtree(d)
-        self.storage_dirs_to_delete = []
-
-    def __del__(self):
-        if self.graph is not None:
-            if self.graph.store is not None:
-                try:
-                    self.graph.store._inner.flush()
-                except Exception:
-                    logging.debug("Error flushing graph")
-            self.graph.close()
-        # It requires more investigation os.remove(self.internal_storage_dir / "LOCK")
-
-    def commit(self):
-        """Commits the graph."""
-        if self.graph:
-            if self.graph.store:
-                logging.info("Committing graph - flushing and optimizing")
-                self.graph.store._inner.flush()
-                self.graph.store._inner.optimize()
-            self.graph.commit()
-
-    @staticmethod
-    def drop_graph_store_storage(storage_path: Path | None) -> None:
-        """Drop graph store storage on disk.
-
-        Args:
-            storage_path : Path to storage directory
-        """
-        if storage_path and storage_path.exists():
-            for f in os.listdir(storage_path):
-                (storage_path / f).unlink()
-            logging.info("Graph store dropped.")
-        else:
-            logging.info(f"Storage path {storage_path} does not exist. Skipping drop.")