linkml-store 0.0.0__py3-none-any.whl → 0.1.7__py3-none-any.whl

linkml_store/api/config.py

@@ -0,0 +1,101 @@
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class CollectionConfig(BaseModel):
+    name: Optional[str] = Field(
+        default=None,
+        description="An optional name for the collection",
+    )
+    alias: Optional[str] = Field(
+        default=None,
+        description="An optional alias for the collection",
+    )
+    type: Optional[str] = Field(
+        default=None,
+        description="The type of object in the collection. TODO; use this instead of name",
+    )
+    metadata: Optional[Dict] = Field(
+        default=None,
+        description="Optional metadata for the collection",
+    )
+    attributes: Optional[Dict[str, Dict]] = Field(
+        default=None,
+        description="Optional attributes for the collection, following LinkML schema",
+    )
+    indexers: Optional[Dict[str, Dict]] = Field(
+        default=None,
+        description="Optional configuration for indexers",
+    )
+    hidden: Optional[bool] = Field(
+        default=False,
+        description="Whether the collection is hidden",
+    )
+    is_prepopulated: Optional[bool] = Field(
+        default=False,
+        description="Whether the collection is prepopulated",
+    )
+
+
+class DatabaseConfig(BaseModel):
+    handle: str = Field(
+        default="duckdb:///:memory:",
+        description="The database handle, e.g., 'duckdb:///:memory:' or 'mongodb://localhost:27017'",
+    )
+    alias: Optional[str] = Field(
+        default=None,
+        description="An optional alias for the database",
+    )
+    schema_location: Optional[str] = Field(
+        default=None,
+        description="The location of the schema file, either a path on disk or URL",
+    )
+    schema_dict: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="The LinkML schema as a dictionary",
+    )
+    collections: Dict[str, CollectionConfig] = Field(
+        default={},
+        description="A dictionary of collection configurations",
+    )
+    recreate_if_exists: bool = Field(
+        default=False,
+        description="Whether to recreate the database if it already exists",
+    )
+    collection_type_slot: Optional[str] = Field(
+        default=None,
+        description=(
+            "For databases that combine multiple collections into a single space, this field"
+            "specifies the field that contains the collection type. An example of this is a Solr"
+            "index that does not use cores for collections, and instead uses a single global"
+            "document space; if this has a field 'document_type', then this field should be set"
+        ),
+    )
+    searchable_slots: Optional[List[str]] = Field(
+        default=None,
+        description="Optional configuration for search fields",
+    )
+    ensure_referential_integrity: bool = Field(
+        default=False,
+        description="Whether to ensure referential integrity",
+    )
+
+
+class ClientConfig(BaseModel):
+    handle: Optional[str] = Field(
+        default=None,
+        description="The client handle",
+    )
+    databases: Dict[str, DatabaseConfig] = Field(
+        default={},
+        description="A dictionary of database configurations",
+    )
+    schema_path: Optional[str] = Field(
+        default=None,
+        description="The path to the LinkML schema file",
+    )
+    base_dir: Optional[str] = Field(
+        default=None,
+        description="The base directory for the client",
+    )

linkml_store/api/database.py

@@ -1,15 +1,28 @@
+import logging
 from abc import ABC
-from dataclasses import dataclass
-from typing import Dict, Optional, Sequence
+from collections import defaultdict
+from copy import copy
+from pathlib import Path
+from typing import TYPE_CHECKING, ClassVar, Dict, Iterator, Optional, Sequence, Type, Union
+
+try:
+    from linkml.validator.report import Severity, ValidationResult
+except ImportError:
+    ValidationResult = None
 
 from linkml_runtime import SchemaView
+from linkml_runtime.linkml_model import ClassDefinition, SchemaDefinition
 
 from linkml_store.api.collection import Collection
-from linkml_store.api.metadata import MetaData
+from linkml_store.api.config import CollectionConfig, DatabaseConfig
 from linkml_store.api.queries import Query, QueryResult
 
+if TYPE_CHECKING:
+    from linkml_store.api.client import Client
+
+logger = logging.getLogger(__name__)
+
 
-@dataclass
 class Database(ABC):
     """
     A Database provides access to named collections of data.
@@ -27,7 +40,7 @@ class Database(ABC):
     >>> db.get_collection("Person") == collection
     True
     >>> objs = [{"id": "P1", "name": "John", "age_in_years": 30}, {"id": "P2", "name": "Alice", "age_in_years": 25}]
-    >>> collection.add(objs)
+    >>> collection.insert(objs)
     >>> qr = collection.find()
     >>> len(qr.rows)
     2
@@ -43,10 +56,82 @@ class Database(ABC):
 
     """
 
-    handle: Optional[str] = None
-    recreate_if_exists: Optional[bool] = False
     _schema_view: Optional[SchemaView] = None
     _collections: Optional[Dict[str, Collection]] = None
+    parent: Optional["Client"] = None
+    metadata: Optional[DatabaseConfig] = None
+    collection_class: ClassVar[Optional[Type[Collection]]] = None
+
+    def __init__(self, handle: Optional[str] = None, metadata: Optional[DatabaseConfig] = None, **kwargs):
+        if metadata:
+            self.metadata = metadata
+        else:
+            self.metadata = DatabaseConfig(handle=handle, **kwargs)
+        if handle is not None and self.metadata.handle is not None and handle != self.metadata.handle:
+            raise ValueError(f"Handle mismatch: {handle} != {self.metadata.handle}")
+        self._initialize_schema()
+        self._initialize_collections()
+
+    def _initialize_schema(self, **kwargs):
+        db_config = self.metadata
+        if db_config.schema_location:
+            schema_location = db_config.schema_location.format(base_dir=self.parent.metadata.base_dir)
+            logger.info(f"Loading schema from: {schema_location}")
+            self.load_schema_view(schema_location)
+        if db_config.schema_dict:
+            schema_dict = copy(db_config.schema_dict)
+            if "id" not in schema_dict:
+                schema_dict["id"] = "tmp"
+            if "name" not in schema_dict:
+                schema_dict["name"] = "tmp"
+            self.set_schema_view(SchemaView(SchemaDefinition(**schema_dict)))
+
+    def from_config(self, db_config: DatabaseConfig, **kwargs):
+        """
+        Initialize a database from a configuration.
+
+        TODO: DEPRECATE
+
+        :param db_config: database configuration
+        :param kwargs: additional arguments
+        """
+        self.metadata = db_config
+        self._initialize_schema()
+        self._initialize_collections()
+        return self
+
+    def _initialize_collections(self):
+        for name, collection_config in self.metadata.collections.items():
+            alias = collection_config.alias
+            typ = collection_config.type
+            # if typ and alias is None:
+            #    alias = name
+            # if typ is None:
+            #    typ = name
+            # collection = self.create_collection(
+            #    typ, alias=alias, metadata=collection_config.metadata
+            # )
+            if False and typ is not None:
+                if not alias:
+                    alias = name
+                name = typ
+            if not collection_config.name:
+                collection_config.name = name
+            _collection = self.create_collection(name, alias=alias, metadata=collection_config)
+            if collection_config.attributes:
+                sv = self.schema_view
+                cd = ClassDefinition(name, attributes=collection_config.attributes)
+                sv.schema.classes[cd.name] = cd
+                sv.set_modified()
+                # assert collection.class_definition() is not None
+
+    @property
+    def recreate_if_exists(self) -> bool:
+        return self.metadata.recreate_if_exists
+
+    @property
+    def handle(self) -> str:
+        return self.metadata.handle
 
     def store(self, obj: Dict[str, str], **kwargs):
         """
@@ -55,13 +140,28 @@ class Database(ABC):
         :param obj: object to store
         :param kwargs: additional arguments
         """
+        sv = self.schema_view
+        roots = [c for c in sv.all_classes().values() if c.tree_root]
+        root = roots[0] if roots else None
         for k, v in obj.items():
+            if root:
+                slot = sv.induced_slot(k, root.name)
+                if not slot:
+                    raise ValueError(f"Cannot determine type for {k}")
+            else:
+                slot = None
+            if isinstance(v, dict):
+                logger.debug(f"Coercing dict to list: {v}")
+                v = [v]
             if not isinstance(v, list):
                 continue
             if not v:
                 continue
-            collection = self.get_collection(k, create_if_not_exists=True)
-            collection.add(v)
+            if slot:
+                collection = self.get_collection(slot.range, create_if_not_exists=True)
+            else:
+                collection = self.get_collection(k, create_if_not_exists=True)
+            collection.replace(v)
 
     def commit(self, **kwargs):
         """
@@ -75,8 +175,17 @@ class Database(ABC):
         """
         raise NotImplementedError()
 
+    @property
+    def _collection_class(self) -> Type[Collection]:
+        raise NotImplementedError()
+
     def create_collection(
-        self, name: str, alias: Optional[str] = None, metadata: Optional[MetaData] = None, **kwargs
+        self,
+        name: str,
+        alias: Optional[str] = None,
+        metadata: Optional[CollectionConfig] = None,
+        recreate_if_exists=False,
+        **kwargs,
     ) -> Collection:
         """
         Create a new collection
@@ -91,11 +200,28 @@ class Database(ABC):
         :param name: name of the collection
         :param alias: alias for the collection
         :param metadata: metadata for the collection
+        :param recreate_if_exists: recreate the collection if it already exists
         :param kwargs: additional arguments
         """
-        raise NotImplementedError()
-
-    def list_collections(self) -> Sequence[Collection]:
+        if not name:
+            raise ValueError(f"Collection name must be provided: alias: {alias} metadata: {metadata}")
+        collection_cls = self.collection_class
+        collection = collection_cls(name=name, alias=alias, parent=self, metadata=metadata)
+        if metadata and metadata.attributes:
+            sv = self.schema_view
+            schema = sv.schema
+            cd = ClassDefinition(name=metadata.type, attributes=metadata.attributes)
+            schema.classes[cd.name] = cd
+        if not self._collections:
+            self._collections = {}
+        if not alias:
+            alias = name
+        self._collections[alias] = collection
+        if recreate_if_exists:
+            collection.delete_where({}, missing_ok=True)
+        return collection
+
+    def list_collections(self, include_internal=False) -> Sequence[Collection]:
         """
         List all collections.
 
@@ -112,10 +238,32 @@ class Database(ABC):
         >>> [c.name for c in collections]
         ['Person', 'Product']
 
+        :param include_internal: include internal collections
+        :return: list of collections
         """
         if not self._collections:
             self.init_collections()
-        return list(self._collections.values())
+        return [c for c in self._collections.values() if include_internal or not c.is_internal]
+
+    def list_collection_names(self, **kwargs) -> Sequence[str]:
+        """
+        List all collection names.
+
+        Examples
+        --------
+        >>> from linkml_store.api.client import Client
+        >>> client = Client()
+        >>> db = client.attach_database("duckdb", alias="test")
+        >>> c1 = db.create_collection("Person")
+        >>> c2 = db.create_collection("Product")
+        >>> collection_names = db.list_collection_names()
+        >>> len(collection_names)
+        2
+        >>> collection_names
+        ['Person', 'Product']
+
+        """
+        return [c.name for c in self.list_collections(**kwargs)]
 
     def get_collection(self, name: str, create_if_not_exists=True, **kwargs) -> "Collection":
         """
@@ -140,7 +288,7 @@ class Database(ABC):
         """
         if not self._collections:
             self.init_collections()
-        if name not in self._collections:
+        if name not in self._collections.keys():
             if create_if_not_exists:
                 self._collections[name] = self.create_collection(name)
             else:
@@ -167,7 +315,7 @@ class Database(ABC):
         >>> client = Client()
         >>> db = client.attach_database("duckdb", alias="test")
         >>> collection = db.create_collection("Person")
-        >>> collection.add([{"id": "P1", "name": "John"}, {"id": "P2", "name": "Alice"}])
+        >>> collection.insert([{"id": "P1", "name": "John"}, {"id": "P2", "name": "Alice"}])
         >>> query = Query(from_table="Person", where_clause={"name": "John"})
         >>> result = db.query(query)
         >>> len(result.rows)
@@ -187,12 +335,64 @@ class Database(ABC):
         """
         Return a schema view for the named collection
         """
+        if not self._schema_view:
+            self._initialize_schema()
         if not self._schema_view:
             self._schema_view = self.induce_schema_view()
         return self._schema_view
 
-    def set_schema_view(self, schema_view: SchemaView):
+    def set_schema_view(self, schema_view: Union[str, Path, SchemaView]):
+        """
+        Set the schema view for the database.
+
+        :param schema_view:
+        :return:
+        """
+        if isinstance(schema_view, Path):
+            schema_view = str(schema_view)
+        if isinstance(schema_view, str):
+            schema_view = SchemaView(schema_view)
         self._schema_view = schema_view
+        if not self._collections:
+            return
+        # align with induced schema
+        roots = [c for c in schema_view.all_classes().values() if c.tree_root]
+        if len(roots) == 0:
+            all_ranges = set()
+            for cn in schema_view.all_classes():
+                for slot in schema_view.class_induced_slots(cn):
+                    if slot.range:
+                        all_ranges.add(slot.range)
+            roots = [
+                c
+                for c in schema_view.all_classes().values()
+                if not all_ranges.intersection(schema_view.class_ancestors(c.name, reflexive=True))
+            ]
+        if len(roots) == 1:
+            root = roots[0]
+            for slot in schema_view.class_induced_slots(root.name):
+                inlined = slot.inlined or slot.inlined_as_list
+                if inlined and slot.range:
+                    if slot.name in self._collections:
+                        coll = self._collections[slot.name]
+                        if not coll.metadata.type:
+                            coll.metadata.type = slot.range
+
+    def load_schema_view(self, path: Union[str, Path]):
+        """
+        Load a schema view from a file.
+
+        >>> from linkml_store.api.client import Client
+        >>> client = Client()
+        >>> db = client.attach_database("duckdb", alias="test")
+        >>> db.load_schema_view("tests/input/countries/countries.linkml.yaml")
+
+        :param path:
+        :return:
+        """
+        if isinstance(path, Path):
+            path = str(path)
+        self.set_schema_view(SchemaView(path))
 
     def induce_schema_view(self) -> SchemaView:
         """
@@ -203,13 +403,77 @@ class Database(ABC):
         >>> client = Client()
         >>> db = client.attach_database("duckdb", alias="test")
         >>> collection = db.create_collection("Person")
-        >>> collection.add([{"id": "P1", "name": "John", "age_in_years": 25},
+        >>> collection.insert([{"id": "P1", "name": "John", "age_in_years": 25},
         ...                 {"id": "P2", "name": "Alice", "age_in_years": 25}])
         >>> schema_view = db.induce_schema_view()
         >>> cd = schema_view.get_class("Person")
         >>> cd.attributes["id"].range
         'string'
+        >>> cd.attributes["age_in_years"].range
+        'integer'
 
         :return: A schema view
         """
         raise NotImplementedError()
+
+    def iter_validate_database(self, **kwargs) -> Iterator["ValidationResult"]:
+        """
+        Validate the contents of the database.
+
+        :param kwargs:
+        :return: iterator over validation results
+        """
+        for collection in self.list_collections():
+            yield from collection.iter_validate_collection(**kwargs)
+        if self.metadata.ensure_referential_integrity:
+            yield from self._validate_referential_integrity(**kwargs)
+
+    def _validate_referential_integrity(self, **kwargs) -> Iterator["ValidationResult"]:
+        """
+        Validate referential integrity of the database.
+
+        :param kwargs:
+        :return: iterator over validation results
+        """
+        sv = self.schema_view
+        cmap = defaultdict(list)
+        for collection in self.list_collections():
+            if not collection.target_class_name:
+                raise ValueError(f"Collection {collection.name} has no target class")
+            cmap[collection.target_class_name].append(collection)
+        for collection in self.list_collections():
+            cd = collection.class_definition()
+            induced_slots = sv.class_induced_slots(cd.name)
+            slot_map = {s.name: s for s in induced_slots}
+            # rmap = {s.name: s.range for s in induced_slots}
+            sr_to_coll = {s.name: cmap.get(s.range, []) for s in induced_slots if s.range}
+            for obj in collection.find_iter():
+                for k, v in obj.items():
+                    if k not in sr_to_coll:
+                        continue
+                    ref_colls = sr_to_coll[k]
+                    if not ref_colls:
+                        continue
+                    if not isinstance(v, (str, int)):
+                        continue
+                    slot = slot_map[k]
+                    found = False
+                    for ref_coll in ref_colls:
+                        ref_obj = ref_coll.get_one(v)
+                        if ref_obj:
+                            found = True
+                            break
+                    if not found:
+                        yield ValidationResult(
+                            type="ReferentialIntegrity",
+                            severity=Severity.ERROR,
+                            message=f"Referential integrity error: {slot.range} not found",
+                            instantiates=slot.range,
+                            instance=v,
+                        )
+
+    def drop(self, **kwargs):
+        """
+        Drop the database and all collections
+        """
+        raise NotImplementedError()

linkml_store/api/queries.py

@@ -22,7 +22,7 @@ class Query(BaseModel):
 
     from_table: Optional[str]
     select_cols: Optional[List[str]] = None
-    where_clause: Optional[Union[str, List[str], Dict[str, str]]] = None
+    where_clause: Optional[Union[str, List[str], Dict[str, Any]]] = None
     sort_by: Optional[List[str]] = None
     limit: Optional[int] = None
     offset: Optional[int] = None
@@ -30,12 +30,21 @@ class Query(BaseModel):
     facet_slots: Optional[List[str]] = None
 
 
+class FacetCountResult(BaseModel):
+    """
+    A facet count result
+    """
+
+    as_dict: Dict[FACET_GROUP, List[Tuple[FACET_GROUP, int]]]
+
+
 class QueryResult(BaseModel):
     """
     A query result
     """
 
     query: Optional[Query] = None
+    search_term: Optional[str] = None
     num_rows: int
     offset: Optional[int] = 0
     rows: Optional[List[Dict[str, Any]]] = None
@@ -45,6 +54,8 @@ class QueryResult(BaseModel):
 
     @property
     def rows_dataframe(self) -> pd.DataFrame:
+        if self.ranked_rows is not None:
+            self._rows_dataframe = pd.DataFrame([{"score": score, **row} for score, row in self.ranked_rows])
         if self._rows_dataframe is None and self.rows:
             self._rows_dataframe = pd.DataFrame(self.rows)
         return self._rows_dataframe

linkml_store/api/stores/chromadb/__init__.py

@@ -0,0 +1,3 @@
+"""
+Support for ChromaDB is experimental.
+"""

linkml_store/api/stores/chromadb/chromadb_collection.py

@@ -0,0 +1,121 @@
+"""
+ChromaDB Collection
+"""
+
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from chromadb.api.models.Collection import Collection as ChromaCollection
+from linkml_runtime.linkml_model import SlotDefinition
+
+from linkml_store.api import Collection
+from linkml_store.api.collection import DEFAULT_FACET_LIMIT, OBJECT
+from linkml_store.api.queries import Query, QueryResult
+from linkml_store.index import Indexer
+
+logger = logging.getLogger(__name__)
+
+
+class ChromaDBCollection(Collection):
+    """
+    A wrapper for ChromaDB collections.
+    """
+
+    @property
+    def native_collection(self) -> ChromaCollection:
+        return self.parent.client.get_collection(self.name)
+
+    def insert(self, objs: Union[OBJECT, List[OBJECT]], **kwargs):
+        if not isinstance(objs, list):
+            objs = [objs]
+
+        documents = []
+        metadatas = []
+        ids = []
+        indexer = Indexer()
+
+        for obj in objs:
+            obj_id = self.object_identifier(obj)
+            ids.append(obj_id)
+            doc_text = indexer.object_to_text(obj)
+            documents.append(doc_text)
+            # TODO: handle nesting
+            metadata = {k: v for k, v in obj.items()}
+            metadatas.append(metadata)
+
+        self.native_collection.add(
+            documents=documents,
+            metadatas=metadatas,
+            ids=ids,
+        )
+
+    def delete(self, objs: Union[OBJECT, List[OBJECT]], **kwargs) -> int:
+        if not isinstance(objs, list):
+            objs = [objs]
+        ids = [obj["id"] for obj in objs]
+        self.native_collection.delete(ids=ids)
+        return len(ids)
+
+    def delete_where(self, where: Optional[Dict[str, Any]] = None, missing_ok=True, **kwargs) -> int:
+        logger.info(f"Deleting from {self.target_class_name} where: {where}")
+        if where is None:
+            where = {}
+        results = self.native_collection.get(where=where)
+        ids = [result["id"] for result in results]
+        self.native_collection.delete(ids=ids)
+        return len(ids)
+
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        chroma_filter = self._build_chroma_filter(query.where_clause)
+        if query.limit:
+            results = self.native_collection.get(where=chroma_filter, limit=query.limit)
+        else:
+            results = self.native_collection.get(where=chroma_filter)
+
+        count = len(results)
+        return QueryResult(query=query, num_rows=count, rows=results)
+
+    def query_facets(
+        self, where: Dict = None, facet_columns: List[str] = None, facet_limit=DEFAULT_FACET_LIMIT, **kwargs
+    ) -> Dict[str, List[Tuple[Any, int]]]:
+        results = {}
+        cd = self.class_definition()
+        if not facet_columns:
+            facet_columns = list(self.class_definition().attributes.keys())
+
+        for col in facet_columns:
+            logger.debug(f"Faceting on {col}")
+            if isinstance(col, tuple):
+                sd = SlotDefinition(name="PLACEHOLDER")
+            else:
+                sd = cd.attributes[col]
+
+            if sd.multivalued:
+                facet_results = self.native_collection.aggregate(
+                    aggregation=[
+                        {"$match": where} if where else {"$match": {}},
+                        {"$unwind": f"${col}"},
+                        {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
+                        {"$sort": {"count": -1}},
+                        {"$limit": facet_limit},
+                    ]
+                )
+            else:
+                facet_results = self.native_collection.aggregate(
+                    aggregation=[
+                        {"$match": where} if where else {"$match": {}},
+                        {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
+                        {"$sort": {"count": -1}},
+                        {"$limit": facet_limit},
+                    ]
+                )
+
+            results[col] = [(result["_id"], result["count"]) for result in facet_results]
+
+        return results
+
+    def _build_chroma_filter(self, where_clause: Dict[str, Any]) -> Dict[str, Any]:
+        chroma_filter = {}
+        for field, value in where_clause.items():
+            chroma_filter[field] = value
+        return chroma_filter