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

linkml_store/api/database.py

@@ -1,15 +1,27 @@
+import logging
 from abc import ABC
-from dataclasses import dataclass
-from typing import Dict, Optional, Sequence
+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 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 +39,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 +55,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 +139,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 +174,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 +199,29 @@ 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_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,13 +335,37 @@ 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):
+        """
+        Set the schema view for the database.
+
+        :param schema_view:
+        :return:
+        """
         self._schema_view = schema_view
 
+    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:
         """
         Induce a schema view from a schema definition.
@@ -203,13 +375,31 @@ 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)
+
+    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/chromadb_collection.py

@@ -0,0 +1,114 @@
+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):
+
+    @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

linkml_store/api/stores/chromadb/chromadb_database.py

@@ -0,0 +1,89 @@
+# chromadb_database.py
+
+import logging
+from typing import Optional
+
+import chromadb
+from chromadb.config import Settings
+from linkml_runtime import SchemaView
+from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
+from linkml_runtime.utils.schema_builder import SchemaBuilder
+
+from linkml_store.api import Database
+from linkml_store.api.queries import Query, QueryResult
+from linkml_store.api.stores.chromadb.chromadb_collection import ChromaDBCollection
+
+logger = logging.getLogger(__name__)
+
+
+class ChromaDBDatabase(Database):
+    _client: chromadb.Client = None
+    collection_class = ChromaDBCollection
+
+    def __init__(self, handle: Optional[str] = None, **kwargs):
+        if handle is None:
+            handle = ".chromadb"
+        super().__init__(handle=handle, **kwargs)
+
+    @property
+    def client(self) -> chromadb.Client:
+        if self._client is None:
+            self._client = chromadb.Client(
+                Settings(
+                    chroma_db_impl="duckdb+parquet",
+                    persist_directory=self.handle,
+                )
+            )
+        return self._client
+
+    def commit(self, **kwargs):
+        pass
+
+    def close(self, **kwargs):
+        if self._client:
+            self._client.close()
+
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        if query.from_table:
+            collection = self.get_collection(query.from_table)
+            return collection.query(query, **kwargs)
+
+    def init_collections(self):
+        if self._collections is None:
+            self._collections = {}
+
+        for collection_name in self.client.list_collections():
+            if collection_name not in self._collections:
+                collection = ChromaDBCollection(name=collection_name, parent=self)
+                self._collections[collection_name] = collection
+
+    def induce_schema_view(self) -> SchemaView:
+        logger.info(f"Inducing schema view for {self.handle}")
+        sb = SchemaBuilder()
+        schema = sb.schema
+
+        for collection_name in self.client.list_collections():
+            sb.add_class(collection_name)
+            chroma_collection = self.client.get_collection(collection_name)
+            sample_doc = chroma_collection.peek(1)
+            if sample_doc:
+                for field, value in sample_doc[0].items():
+                    if field == "_id":
+                        continue
+                    sd = SlotDefinition(field)
+                    if isinstance(value, list):
+                        sd.multivalued = True
+                    if isinstance(value, dict):
+                        sd.inlined = True
+                    sb.schema.classes[collection_name].attributes[sd.name] = sd
+
+        sb.add_defaults()
+        for cls_name in schema.classes:
+            if cls_name in self.metadata.collections:
+                collection_metadata = self.metadata.collections[cls_name]
+                if collection_metadata.attributes:
+                    del schema.classes[cls_name]
+                    cls = ClassDefinition(name=collection_metadata.type, attributes=collection_metadata.attributes)
+                    schema.classes[cls.name] = cls
+
+        return SchemaView(schema)