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

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/solr_collection.py

@@ -0,0 +1,132 @@
+# 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,82 @@
+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)