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

linkml_store/api/stores/mongodb/mongodb_collection.py

@@ -1,56 +1,108 @@
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+import logging
+from copy import copy
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from linkml_runtime.linkml_model import SlotDefinition
+from pymongo.collection import Collection as MongoCollection
 
 from linkml_store.api import Collection
-from linkml_store.api.collection import OBJECT
+from linkml_store.api.collection import DEFAULT_FACET_LIMIT, OBJECT
+from linkml_store.api.queries import Query, QueryResult
+
+logger = logging.getLogger(__name__)
 
 
-@dataclass
 class MongoDBCollection(Collection):
-    """
-    A wrapper around a MongoDB collection
-    """
 
-    def add(self, objs: Union[OBJECT, List[OBJECT]], **kwargs):
-        if not isinstance(objs, list):
-            objs = [objs]
-        if not objs:
-            return
-        cd = self.class_definition()
-        if not cd:
-            cd = self.induce_class_definition_from_objects(objs)
-        collection = self.parent.database[self.name]
-        collection.insert_many(objs)
+    @property
+    def mongo_collection(self) -> MongoCollection:
+        if not self.name:
+            raise ValueError("Collection name not set")
+        return self.parent.native_db[self.name]
 
-    def delete(self, objs: Union[OBJECT, List[OBJECT]], **kwargs) -> int:
+    def insert(self, objs: Union[OBJECT, List[OBJECT]], **kwargs):
         if not isinstance(objs, list):
             objs = [objs]
-        cd = self.class_definition()
-        if not cd:
-            cd = self.induce_class_definition_from_objects(objs)
-        collection = self.parent.database[self.name]
-        deleted_count = 0
-        for obj in objs:
-            result = collection.delete_one(obj)
-            deleted_count += result.deleted_count
-        return deleted_count
+        self.mongo_collection.insert_many(objs)
 
-    def delete_where(self, where: Optional[Dict[str, Any]] = None, **kwargs) -> int:
-        collection = self.parent.database[self.name]
-        result = collection.delete_many(where)
-        return result.deleted_count
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        mongo_filter = self._build_mongo_filter(query.where_clause)
+        if query.limit:
+            cursor = self.mongo_collection.find(mongo_filter).limit(query.limit)
+        else:
+            cursor = self.mongo_collection.find(mongo_filter)
+
+        def _as_row(row: dict):
+            row = copy(row)
+            del row["_id"]
+            return row
+
+        rows = [_as_row(row) for row in cursor]
+        count = self.mongo_collection.count_documents(mongo_filter)
+
+        return QueryResult(query=query, num_rows=count, rows=rows)
+
+    def _build_mongo_filter(self, where_clause: Dict[str, Any]) -> Dict[str, Any]:
+        mongo_filter = {}
+        if where_clause:
+            for field, value in where_clause.items():
+                mongo_filter[field] = value
+        return mongo_filter
 
-    def query_facets(self, where: Dict = None, facet_columns: List[str] = None) -> Dict[str, Dict[str, int]]:
+    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()
-        collection = self.parent.database[self.name]
+        cd = self.class_definition()
         if not facet_columns:
             facet_columns = list(self.class_definition().attributes.keys())
+
         for col in facet_columns:
-            facet_pipeline = [
-                {"$match": where} if where else {"$match": {}},
-                {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
-            ]
-            facet_results = list(collection.aggregate(facet_pipeline))
-            results[col] = [(row["_id"], row["count"]) for row in facet_results]
+            logger.debug(f"Faceting on {col}")
+            if isinstance(col, tuple):
+                sd = SlotDefinition(name="PLACEHOLDER")
+            else:
+                sd = cd.attributes[col]
+
+            if sd.multivalued:
+                facet_pipeline = [
+                    {"$match": where} if where else {"$match": {}},
+                    {"$unwind": f"${col}"},
+                    {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
+                    {"$sort": {"count": -1}},
+                    {"$limit": facet_limit},
+                ]
+            else:
+                facet_pipeline = [
+                    {"$match": where} if where else {"$match": {}},
+                    {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
+                    {"$sort": {"count": -1}},
+                    {"$limit": facet_limit},
+                ]
+
+            facet_results = list(self.mongo_collection.aggregate(facet_pipeline))
+            results[col] = [(result["_id"], result["count"]) for result in facet_results]
+
         return results
+
+    def delete(self, objs: Union[OBJECT, List[OBJECT]], **kwargs) -> int:
+        if not isinstance(objs, list):
+            objs = [objs]
+        filter_conditions = []
+        for obj in objs:
+            filter_condition = {}
+            for key, value in obj.items():
+                filter_condition[key] = value
+            filter_conditions.append(filter_condition)
+        result = self.mongo_collection.delete_many({"$or": filter_conditions})
+        return result.deleted_count
+
+    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 = {}
+        result = self.mongo_collection.delete_many(where)
+        deleted_rows_count = result.deleted_count
+        if deleted_rows_count == 0 and not missing_ok:
+            raise ValueError(f"No rows found for {where}")
+        return deleted_rows_count

linkml_store/api/stores/mongodb/mongodb_database.py

@@ -1,112 +1,103 @@
-from dataclasses import dataclass
+# mongodb_database.py
+
+import logging
 from typing import Optional
 
 from linkml_runtime import SchemaView
-from linkml_runtime.linkml_model import SlotDefinition
+from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
 from linkml_runtime.utils.schema_builder import SchemaBuilder
 from pymongo import MongoClient
+from pymongo.database import Database as NativeDatabase
 
 from linkml_store.api import Database
 from linkml_store.api.queries import Query, QueryResult
 from linkml_store.api.stores.mongodb.mongodb_collection import MongoDBCollection
 
+logger = logging.getLogger(__name__)
+
 
-@dataclass
 class MongoDBDatabase(Database):
     """
-    A wrapper around a MongoDB database
+    An adapter for MongoDB databases.
+
+    The LinkML-Store Database abstraction combines mongodb Client and Database.
     """
 
-    _client: MongoClient = None
-    _database = None
+    _native_client: MongoClient = None
+    _native_db = None
+    collection_class = MongoDBCollection
 
-    def __post_init__(self):
-        if not self.handle:
-            self.handle = "mongodb://localhost:27017"
+    def __init__(self, handle: Optional[str] = None, **kwargs):
+        if handle is None:
+            handle = "mongodb://localhost:27017"
+        super().__init__(handle=handle, **kwargs)
 
     @property
-    def client(self) -> MongoClient:
-        if not self._client:
-            self._client = MongoClient(self.handle)
-        return self._client
+    def native_client(self) -> MongoClient:
+        if self._native_client is None:
+            self._native_client = MongoClient(self.handle)
+        return self._native_client
 
     @property
-    def database(self):
-        if not self._database:
-            db_name = self.handle.split("/")[-1]
-            self._database = self.client[db_name]
-        return self._database
+    def native_db(self) -> NativeDatabase:
+        if self._native_db is None:
+            alias = self.metadata.alias
+            if not alias:
+                alias = "default"
+            self._native_db = self.native_client[alias]
+        return self._native_db
 
     def commit(self, **kwargs):
         pass
 
     def close(self, **kwargs):
-        self.client.close()
+        if self._native_client:
+            self._native_client.close()
+
+    def drop(self, **kwargs):
+        self.native_client.drop_database(self.metadata.alias)
 
     def query(self, query: Query, **kwargs) -> QueryResult:
-        collection = self.database[query.from_table]
-        where_clause = query.where_clause or {}
-        cursor = collection.find(where_clause)
-        if query.limit:
-            cursor = cursor.limit(query.limit)
-        if query.offset:
-            cursor = cursor.skip(query.offset)
-        if query.sort_by:
-            sort_key = [(col, 1) for col in query.sort_by]
-            cursor = cursor.sort(sort_key)
-        rows = list(cursor)
-        num_rows = len(rows)
-        qr = QueryResult(query=query, num_rows=num_rows, rows=rows)
-        return qr
+        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.database.list_collection_names():
+
+        for collection_name in self.native_db.list_collection_names():
             if collection_name not in self._collections:
                 collection = MongoDBCollection(name=collection_name, parent=self)
                 self._collections[collection_name] = collection
 
-    def create_collection(self, name: str, alias: Optional[str] = None, **kwargs) -> MongoDBCollection:
-        collection = MongoDBCollection(name=name, parent=self)
-        if not self._collections:
-            self._collections = {}
-        if not alias:
-            alias = name
-        self._collections[alias] = collection
-        return collection
-
     def induce_schema_view(self) -> SchemaView:
+        logger.info(f"Inducing schema view for {self.handle}")
         sb = SchemaBuilder()
         schema = sb.schema
-        collection_names = self.database.list_collection_names()
-        for collection_name in collection_names:
+
+        for collection_name in self.native_db.list_collection_names():
             sb.add_class(collection_name)
-            collection = self.database[collection_name]
-            sample_doc = collection.find_one()
+            mongo_collection = self.native_db[collection_name]
+            sample_doc = mongo_collection.find_one()
             if sample_doc:
-                for key, value in sample_doc.items():
-                    if key == "_id":
+                for field, value in sample_doc.items():
+                    if field == "_id":
                         continue
+                    sd = SlotDefinition(field)
                     if isinstance(value, list):
-                        multivalued = True
-                        if value:
-                            value = value[0]
-                        else:
-                            value = None
-                    else:
-                        multivalued = False
-                    if isinstance(value, str):
-                        rng = "string"
-                    elif isinstance(value, int):
-                        rng = "integer"
-                    elif isinstance(value, float):
-                        rng = "float"
-                    elif isinstance(value, bool):
-                        rng = "boolean"
-                    else:
-                        rng = "string"
-                    sd = SlotDefinition(key, range=rng, multivalued=multivalued)
+                        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)

linkml_store/api/stores/solr/__init__.py

@@ -0,0 +1,3 @@
+"""
+Wrapper for Solr endpoints.
+"""

linkml_store/api/stores/solr/solr_collection.py

@@ -0,0 +1,133 @@
+# solr_collection.py
+
+import logging
+from copy import copy
+from typing import Any, Dict, List, Optional, Union
+
+import requests
+
+from linkml_store.api import Collection
+from linkml_store.api.collection import DEFAULT_FACET_LIMIT
+from linkml_store.api.queries import Query, QueryResult
+
+logger = logging.getLogger(__name__)
+
+
+class SolrCollection(Collection):
+
+    @property
+    def _collection_base(self) -> str:
+        if self.parent.use_cores:
+            base_url = f"{self.parent.base_url}/{self.name}"
+        else:
+            base_url = self.parent.base_url
+        return base_url
+
+    def search(
+        self,
+        query: str,
+        where: Optional[Any] = None,
+        index_name: Optional[str] = None,
+        limit: Optional[int] = None,
+        **kwargs,
+    ) -> QueryResult:
+        if index_name is None:
+            index_name = "edismax"
+        qfs = self.parent.metadata.searchable_slots
+        if not qfs:
+            raise ValueError("No searchable slots configured for Solr collection")
+        solr_query = self._build_solr_query(where, search_term=query, extra={"defType": index_name, "qf": qfs})
+        logger.info(f"Querying Solr collection {self.name} with query: {solr_query}")
+
+        response = requests.get(f"{self._collection_base}/select", params=solr_query)
+        response.raise_for_status()
+
+        data = response.json()
+        num_rows = data["response"]["numFound"]
+        rows = data["response"]["docs"]
+        ranked_rows = [(1.0, row) for row in rows]
+        return QueryResult(query=where, search_term=query, num_rows=num_rows, rows=rows, ranked_rows=ranked_rows)
+
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        solr_query = self._build_solr_query(query)
+        logger.info(f"Querying Solr collection {self.name} with query: {solr_query}")
+
+        response = requests.get(f"{self._collection_base}/select", params=solr_query)
+        response.raise_for_status()
+
+        data = response.json()
+        num_rows = data["response"]["numFound"]
+        rows = data["response"]["docs"]
+
+        return QueryResult(query=query, num_rows=num_rows, rows=rows)
+
+    def query_facets(
+        self, where: Optional[Dict] = None, facet_columns: List[str] = None, facet_limit=DEFAULT_FACET_LIMIT, **kwargs
+    ) -> Dict[str, Dict[str, int]]:
+        solr_query = self._build_solr_query(where)
+        solr_query["facet"] = "true"
+        solr_query["facet.field"] = facet_columns
+        solr_query["facet.limit"] = facet_limit
+
+        logger.info(f"Querying Solr collection {self.name} for facets with query: {solr_query}")
+
+        response = requests.get(f"{self._collection_base}/select", params=solr_query)
+        response.raise_for_status()
+
+        data = response.json()
+        facet_counts = data["facet_counts"]["facet_fields"]
+
+        results = {}
+        for facet_field, counts in facet_counts.items():
+            results[facet_field] = list(zip(counts[::2], counts[1::2]))
+
+        return results
+
+    def _build_solr_query(
+        self, query: Union[Query, Dict], search_term="*:*", extra: Optional[Dict] = None
+    ) -> Dict[str, Any]:
+        solr_query = {}
+        if query is None:
+            query = {}
+
+        if isinstance(query, Query):
+            where = query.where_clause
+            solr_query["fq"] = self._build_solr_where_clause(where)
+
+            if query.select_cols:
+                solr_query["fl"] = ",".join(query.select_cols)
+
+            if query.limit:
+                solr_query["rows"] = query.limit
+
+            if query.offset:
+                solr_query["start"] = query.offset
+
+        elif isinstance(query, dict):
+            solr_query["fq"] = self._build_solr_where_clause(query)
+
+        solr_query["wt"] = "json"
+        if "q" not in solr_query:
+            solr_query["q"] = search_term
+        if extra:
+            solr_query.update(extra)
+        logger.info(f"Built Solr query: {solr_query}")
+        return solr_query
+
+    def _build_solr_where_clause(self, where_clause: Dict) -> str:
+        if where_clause is None:
+            where_clause = {}
+        conditions = []
+        if self.parent.metadata.collection_type_slot:
+            where_clause = copy(where_clause)
+            where_clause[self.parent.metadata.collection_type_slot] = self.alias
+        for field, value in where_clause.items():
+            if not isinstance(value, (list, tuple)):
+                value = [value]
+            value = [f'"{v}"' if isinstance(v, str) else str(v) for v in value]
+            if len(value) > 1:
+                conditions.append(f"{field}:({' '.join(value)})")
+            else:
+                conditions.append(f"{field}:{value[0]}")
+
+        return " AND ".join(conditions)

linkml_store/api/stores/solr/solr_database.py

@@ -0,0 +1,83 @@
+import logging
+from typing import Optional
+
+import requests
+
+from linkml_store.api import Collection, Database
+from linkml_store.api.config import CollectionConfig
+from linkml_store.api.queries import Query, QueryResult
+from linkml_store.api.stores.solr.solr_collection import SolrCollection
+
+logger = logging.getLogger(__name__)
+
+
+class SolrDatabase(Database):
+    base_url: str
+    collection_class = SolrCollection
+    use_cores: bool = False
+
+    def __init__(self, handle: Optional[str] = None, **kwargs):
+        if handle.startswith("solr:"):
+            self.base_url = handle.replace("solr:", "")
+        else:
+            self.base_url = handle
+        super().__init__(handle=handle, **kwargs)
+
+    def get_collection(self, name: str, create_if_not_exists=True, **kwargs) -> "Collection":
+        if not self._collections:
+            self.init_collections()
+
+        if name not in self._collections.keys():
+            if create_if_not_exists:
+                self._collections[name] = self.create_collection(name)
+            else:
+                raise KeyError(f"Collection {name} does not exist")
+
+        return self._collections[name]
+
+    def create_collection(
+        self, name: str, alias: Optional[str] = None, metadata: Optional[CollectionConfig] = None, **kwargs
+    ) -> 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 not self._collections:
+            self._collections = {}
+
+        if not alias:
+            alias = name
+
+        self._collections[alias] = collection
+        return collection
+
+    def init_collections(self):
+        if self._collections is None:
+            self._collections = {}
+        if self.metadata.collection_type_slot:
+            response = requests.get(
+                f"{self.base_url}/select",
+                params={
+                    "q": "*:*",
+                    "wt": "json",
+                    "rows": 0,
+                    "facet": "true",
+                    "facet.field": self.metadata.collection_type_slot,
+                    "facet.limit": -1,
+                },
+            )
+            response.raise_for_status()
+            data = response.json()
+            coll_names = data["facet_counts"]["facet_fields"][self.metadata.collection_type_slot]
+            coll_names = coll_names[::2]
+            for coll_name in coll_names:
+                self.create_collection(coll_name)
+        else:
+            self.create_collection("default")
+
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        collection_name = query.from_table
+        collection = self.get_collection(collection_name)
+        return collection.query(query, **kwargs)