linkml-store 0.3.0__py3-none-any.whl

linkml_store/api/stores/solr/solr_collection.py

@@ -0,0 +1,224 @@
+# solr_collection.py
+
+import logging
+from copy import copy
+from typing import Any, Dict, List, Optional, Union, Tuple
+
+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.alias}"
+        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")
+        if not qfs:
+            qfs = list(self.class_definition().attributes.keys())
+        solr_query = self._build_solr_query(where, search_term=query, extra={"defType": index_name, "qf": qfs})
+        logger.info(f"Querying Solr collection {self.alias} 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.alias} with query: {solr_query}")
+
+        response = requests.get(f"{self._collection_base}/select", params=solr_query)
+        response.raise_for_status()
+
+        data = response.json()
+        logger.debug(f"Response: {data}")
+        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[Union[str, Tuple[str, ...]]] = None,
+        facet_limit=DEFAULT_FACET_LIMIT,
+        facet_min_count: int = 1,
+        **kwargs,
+    ) -> Dict[Union[str, Tuple[str, ...]], List[Tuple[Any, int]]]:
+        """
+        Query facet counts for fields or field combinations.
+        
+        :param where: Filter conditions
+        :param facet_columns: List of fields to facet on. Elements can be:
+                            - Simple strings for single field facets
+                            - Tuples of strings for field combinations (pivot facets)
+        :param facet_limit: Maximum number of facet values to return
+        :param facet_min_count: Minimum count for facet values to be included
+        :return: Dictionary mapping fields or field tuples to lists of (value, count) tuples
+        """
+        solr_query = self._build_solr_query(where)
+        
+        # Separate single fields and tuple fields
+        single_fields = []
+        tuple_fields = []
+        
+        if facet_columns:
+            for field in facet_columns:
+                if isinstance(field, str):
+                    single_fields.append(field)
+                elif isinstance(field, tuple):
+                    tuple_fields.append(field)
+        
+        # Process regular facets
+        results = {}
+        if single_fields:
+            solr_query["facet"] = "true"
+            solr_query["facet.field"] = single_fields
+            solr_query["facet.limit"] = facet_limit
+            solr_query["facet.mincount"] = facet_min_count
+            
+            logger.info(f"Querying Solr collection {self.alias} 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"]
+            
+            for facet_field, counts in facet_counts.items():
+                results[facet_field] = list(zip(counts[::2], counts[1::2]))
+        
+        # Process pivot facets for tuple fields
+        if tuple_fields:
+            # TODO: Add a warning if Solr < 4.0, when this was introduced
+            for field_tuple in tuple_fields:
+                # Create a query for this specific field tuple
+                pivot_query = self._build_solr_query(where)
+                pivot_query["facet"] = "true"
+                
+                # Create pivot facet
+                field_str = ','.join(field_tuple)
+                pivot_query["facet.pivot"] = field_str
+                pivot_query["facet.pivot.mincount"] = facet_min_count
+                pivot_query["facet.limit"] = facet_limit
+                
+                logger.info(f"Querying Solr collection {self.alias} for pivot facets with query: {pivot_query}")
+                response = requests.get(f"{self._collection_base}/select", params=pivot_query)
+                response.raise_for_status()
+                
+                data = response.json()
+                pivot_facets = data.get("facet_counts", {}).get("facet_pivot", {})
+                
+                # Process pivot facets into the same format as MongoDB results
+                field_str = ','.join(field_tuple)
+                pivot_data = pivot_facets.get(field_str, [])
+                
+                # Build a list of tuples (field values, count)
+                pivot_results = []
+                self._process_pivot_facets(pivot_data, [], pivot_results, field_tuple)
+                
+                results[field_tuple] = pivot_results
+        
+        return results
+        
+    def _process_pivot_facets(self, pivot_data, current_values, results, field_tuple):
+        """
+        Recursively process pivot facet results to extract combinations of field values.
+        
+        :param pivot_data: The pivot facet data from Solr
+        :param current_values: The current path of values in the recursion
+        :param results: The result list to populate
+        :param field_tuple: The original field tuple for reference
+        """
+        for item in pivot_data:
+            # Add the current field value
+            value = item.get("value")
+            count = item.get("count", 0)
+            
+            # Update the current path with this value
+            values = current_values + [value]
+            
+            # If we have all the fields from the tuple, add a result
+            if len(values) == len(field_tuple):
+                # Create a tuple of values corresponding to the field tuple
+                results.append((tuple(values), count))
+            
+            # Process child pivot fields recursively
+            pivot = item.get("pivot", [])
+            if pivot and len(values) < len(field_tuple):
+                self._process_pivot_facets(pivot, values, results, field_tuple)
+
+    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)

linkml_store/api/types.py

@@ -0,0 +1,4 @@
+from typing import TypeVar
+
+DatabaseType = TypeVar("DatabaseType", bound="Database")  # noqa: F821
+CollectionType = TypeVar("CollectionType", bound="Collection")  # noqa: F821