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

linkml_store/api/stores/duckdb/duckdb_database.py

@@ -1,5 +1,6 @@
 import json
 import logging
+from pathlib import Path
 from typing import Optional
 
 import pandas as pd
@@ -22,6 +23,7 @@ TYPE_MAP = {
     "DATE": "date",
     "DOUBLE": "float",
     "INTEGER": "integer",
+    "JSON": "Any",
 }
 
 
@@ -33,9 +35,13 @@ class DuckDBDatabase(Database):
     _engine: sqlalchemy.Engine = None
     collection_class = DuckDBCollection
 
-    def __init__(self, handle: Optional[str] = None, **kwargs):
+    def __init__(self, handle: Optional[str] = None, recreate_if_exists: bool = False, **kwargs):
         if handle is None:
             handle = "duckdb:///:memory:"
+        if recreate_if_exists:
+            path = Path(handle.replace("duckdb:///", ""))
+            if path.exists():
+                path.unlink()
         super().__init__(handle=handle, **kwargs)
 
     @property
@@ -69,16 +75,19 @@ class DuckDBDatabase(Database):
                 if qr.num_rows == 0:
                     logger.debug(f"Table {query.from_table} not created yet")
                     return QueryResult(query=query, num_rows=0, rows=[])
-            sv = self._schema_view
+            if not query.from_table.startswith("information_schema"):
+                sv = self.schema_view
+            else:
+                sv = None
             if sv:
                 cd = None
                 for c in self._collections.values():
-                    if c.name == query.from_table:
+                    if c.name == query.from_table or c.metadata.alias == query.from_table:
                         cd = c.class_definition()
                         break
                 if cd:
-                    for att in cd.attributes.values():
-                        if att.inlined:
+                    for att in sv.class_induced_slots(cd.name):
+                        if att.inlined or att.inlined_as_list:
                             json_encoded_cols.append(att.name)
         with self.engine.connect() as conn:
             count_query_str = text(query_to_sql(query, count=True))
@@ -107,7 +116,10 @@ class DuckDBDatabase(Database):
 
     def init_collections(self):
         # TODO: unify schema introspection
-        schema = introspect_schema(self.engine)
+        if not self.schema_view:
+            schema = introspect_schema(self.engine)
+        else:
+            schema = self.schema_view.schema
         table_names = schema.classes.keys()
         if self._collections is None:
             self._collections = {}
@@ -119,7 +131,7 @@ class DuckDBDatabase(Database):
     def induce_schema_view(self) -> SchemaView:
         # TODO: unify schema introspection
         # TODO: handle case where schema is provided in advance
-        logger.info(f"Inducing schema view for {self.metadata.handle}")
+        logger.info(f"Inducing schema view for {self.metadata.handle} // {self}")
         sb = SchemaBuilder()
         schema = sb.schema
         query = Query(from_table="information_schema.tables", where_clause={"table_type": "BASE TABLE"})
@@ -144,8 +156,10 @@ class DuckDBDatabase(Database):
             sd = SlotDefinition(
                 row["column_name"], required=row["is_nullable"] == "NO", multivalued=multivalued, range=rng
             )
+            if dt == "JSON":
+                sd.inlined_as_list = True
             sb.schema.classes[tbl_name].attributes[sd.name] = sd
-            logger.info(f"Introspected slot: {tbl_name}.{sd.name}: {sd.range}")
+            logger.info(f"Introspected slot: {tbl_name}.{sd.name}: {sd.range} FROM {dt}")
         sb.add_defaults()
         for cls_name in schema.classes:
             if cls_name in self.metadata.collections:

linkml_store/api/stores/duckdb/mappings.py

@@ -3,5 +3,6 @@ import sqlalchemy as sqla
 TMAP = {
     "string": sqla.String,
     "integer": sqla.Integer,
+    "float": sqla.Float,
     "linkml:Any": sqla.JSON,
 }

linkml_store/api/stores/filesystem/__init__.py

@@ -0,0 +1,16 @@
+"""
+Adapter for DuckDB embedded database.
+
+Handles have the form:
+
+ - ``duckdb:///<path>`` for a file-based database
+ - ``duckdb:///:memory:`` for an in-memory database
+"""
+
+from linkml_store.api.stores.duckdb.duckdb_collection import DuckDBCollection
+from linkml_store.api.stores.duckdb.duckdb_database import DuckDBDatabase
+
+__all__ = [
+    "DuckDBCollection",
+    "DuckDBDatabase",
+]

linkml_store/api/stores/filesystem/filesystem_collection.py

@@ -0,0 +1,142 @@
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+import sqlalchemy as sqla
+from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
+from sqlalchemy import Column, Table, delete, insert, inspect, text
+from sqlalchemy.sql.ddl import CreateTable
+
+from linkml_store.api import Collection
+from linkml_store.api.collection import DEFAULT_FACET_LIMIT, OBJECT
+from linkml_store.api.queries import Query
+from linkml_store.api.stores.duckdb.mappings import TMAP
+from linkml_store.utils.sql_utils import facet_count_sql
+
+logger = logging.getLogger(__name__)
+
+
+class FileSystemCollection(Collection):
+    _table_created: bool = None
+
+    def insert(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)
+        self._create_table(cd)
+        table = self._sqla_table(cd)
+        logger.info(f"Inserting into: {self.alias} // T={table.name}")
+        engine = self.parent.engine
+        col_names = [c.name for c in table.columns]
+        objs = [{k: obj.get(k, None) for k in col_names} for obj in objs]
+        with engine.connect() as conn:
+            with conn.begin():
+                conn.execute(insert(table), objs)
+            conn.commit()
+
+    def delete(self, objs: Union[OBJECT, List[OBJECT]], **kwargs) -> Optional[int]:
+        if not isinstance(objs, list):
+            objs = [objs]
+        cd = self.class_definition()
+        if not cd:
+            cd = self.induce_class_definition_from_objects(objs)
+        table = self._sqla_table(cd)
+        engine = self.parent.engine
+        with engine.connect() as conn:
+            for obj in objs:
+                conditions = [table.c[k] == v for k, v in obj.items() if k in cd.attributes]
+                stmt = delete(table).where(*conditions)
+                stmt = stmt.compile(engine)
+                conn.execute(stmt)
+                conn.commit()
+        return
+
+    def delete_where(self, where: Optional[Dict[str, Any]] = None, missing_ok=True, **kwargs) -> Optional[int]:
+        logger.info(f"Deleting from {self.target_class_name} where: {where}")
+        if where is None:
+            where = {}
+        cd = self.class_definition()
+        if not cd:
+            logger.info(f"No class definition found for {self.target_class_name}, assuming not prepopulated")
+            return 0
+        table = self._sqla_table(cd)
+        engine = self.parent.engine
+        inspector = inspect(engine)
+        table_exists = table.name in inspector.get_table_names()
+        if not table_exists:
+            logger.info(f"Table {table.name} does not exist, assuming no data")
+            return 0
+        with engine.connect() as conn:
+            conditions = [table.c[k] == v for k, v in where.items()]
+            stmt = delete(table).where(*conditions)
+            stmt = stmt.compile(engine)
+            result = conn.execute(stmt)
+            deleted_rows_count = result.rowcount
+            if deleted_rows_count == 0 and not missing_ok:
+                raise ValueError(f"No rows found for {where}")
+            conn.commit()
+            return deleted_rows_count if deleted_rows_count > -1 else None
+
+    def query_facets(
+        self, where: Dict = None, facet_columns: List[str] = None, facet_limit=DEFAULT_FACET_LIMIT, **kwargs
+    ) -> Dict[str, Dict[str, int]]:
+        results = {}
+        cd = self.class_definition()
+        with self.parent.engine.connect() as conn:
+            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]
+                facet_query = self._create_query(where_clause=where)
+                facet_query_str = facet_count_sql(facet_query, col, multivalued=sd.multivalued)
+                logger.debug(f"Facet query: {facet_query_str}")
+                rows = list(conn.execute(text(facet_query_str)))
+                results[col] = rows
+            return results
+
+    def _sqla_table(self, cd: ClassDefinition) -> Table:
+        schema_view = self.parent.schema_view
+        metadata_obj = sqla.MetaData()
+        cols = []
+        for att in schema_view.class_induced_slots(cd.name):
+            typ = TMAP.get(att.range, sqla.String)
+            if att.inlined:
+                typ = sqla.JSON
+            if att.multivalued:
+                typ = sqla.ARRAY(typ, dimensions=1)
+            if att.array:
+                typ = sqla.ARRAY(typ, dimensions=1)
+            col = Column(att.name, typ)
+            cols.append(col)
+        t = Table(self.alias, metadata_obj, *cols)
+        return t
+
+    def _create_table(self, cd: ClassDefinition):
+        if self._table_created or self.metadata.is_prepopulated:
+            logger.info(f"Already have table for: {cd.name}")
+            return
+        query = Query(
+            from_table="information_schema.tables", where_clause={"table_type": "BASE TABLE", "table_name": self.alias}
+        )
+        qr = self.parent.query(query)
+        if qr.num_rows > 0:
+            logger.info(f"Table already exists for {cd.name}")
+            self._table_created = True
+            self.metadata.is_prepopulated = True
+            return
+        logger.info(f"Creating table for {cd.name}")
+        t = self._sqla_table(cd)
+        ct = CreateTable(t)
+        ddl = str(ct.compile(self.parent.engine))
+        with self.parent.engine.connect() as conn:
+            conn.execute(text(ddl))
+            conn.commit()
+        self._table_created = True
+        self.metadata.is_prepopulated = True

linkml_store/api/stores/filesystem/filesystem_database.py

@@ -0,0 +1,36 @@
+import logging
+from typing import Optional
+
+from linkml_store.api import Collection, Database
+from linkml_store.api.config import CollectionConfig
+from linkml_store.api.stores.duckdb import DuckDBDatabase
+from linkml_store.api.stores.filesystem.filesystem_collection import FileSystemCollection
+
+logger = logging.getLogger(__name__)
+
+
+class FileSystemDatabase(Database):
+    collection_class = FileSystemCollection
+    wrapped_database: Database = None
+
+    def __init__(self, handle: Optional[str] = None, recreate_if_exists: bool = False, **kwargs):
+        self.wrapped_database = DuckDBDatabase("duckdb:///:memory:")
+        super().__init__(handle=handle, **kwargs)
+
+    def commit(self, **kwargs):
+        # TODO: sync
+        pass
+
+    def close(self, **kwargs):
+        self.wrapped_database.close()
+
+    def create_collection(
+        self,
+        name: str,
+        alias: Optional[str] = None,
+        metadata: Optional[CollectionConfig] = None,
+        recreate_if_exists=False,
+        **kwargs,
+    ) -> Collection:
+        wd = self.wrapped_database
+        wd.create_collection()

linkml_store/api/stores/hdf5/__init__.py

@@ -0,0 +1,7 @@
+"""
+Adapter for HDF5 file storage.
+
+.. warning::
+
+    Experimental support for HDF5 storage.
+"""

linkml_store/api/stores/hdf5/hdf5_collection.py

@@ -46,7 +46,7 @@ class HDF5Collection(Collection):
         return 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}")
+        logger.info(f"Deleting from {self.target_class_name} where: {where}")
         if where is None:
             where = {}
         results = self.query(Query(where_clause=where)).rows

linkml_store/api/stores/mongodb/__init__.py

@@ -0,0 +1,25 @@
+"""
+Adapter for MongoDB document store.
+
+Handles have the form: ``mongodb://<host>:<port>/<database>``
+
+To use this, you must have the `pymongo` extra installed.
+
+.. code-block:: bash
+
+    pip install linkml-store[mongodb]
+
+or
+
+.. code-block:: bash
+
+    pip install linkml-store[all]
+"""
+
+from linkml_store.api.stores.mongodb.mongodb_collection import MongoDBCollection
+from linkml_store.api.stores.mongodb.mongodb_database import MongoDBDatabase
+
+__all__ = [
+    "MongoDBCollection",
+    "MongoDBDatabase",
+]

linkml_store/api/stores/mongodb/mongodb_collection.py

@@ -1,4 +1,5 @@
 import logging
+from copy import copy
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 from linkml_runtime.linkml_model import SlotDefinition
@@ -12,6 +13,14 @@ logger = logging.getLogger(__name__)
 
 
 class MongoDBCollection(Collection):
+    """
+    Adapter for collections in a MongoDB database.
+
+    .. note::
+
+        You should not use or manipulate this class directly.
+        Instead, use the general :class:`linkml_store.api.Collection`
+    """
 
     @property
     def mongo_collection(self) -> MongoCollection:
@@ -31,7 +40,12 @@ class MongoDBCollection(Collection):
         else:
             cursor = self.mongo_collection.find(mongo_filter)
 
-        rows = list(cursor)
+        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)
@@ -56,24 +70,31 @@ class MongoDBCollection(Collection):
             if isinstance(col, tuple):
                 sd = SlotDefinition(name="PLACEHOLDER")
             else:
-                sd = cd.attributes[col]
-
-            if sd.multivalued:
+                if col in cd.attributes:
+                    sd = cd.attributes[col]
+                else:
+                    logger.info(f"No schema metadata for {col}")
+                    sd = SlotDefinition(name=col)
+            group = {"$group": {"_id": f"${col}", "count": {"$sum": 1}}}
+            if isinstance(col, tuple):
+                q = {k.replace(".", ""): f"${k}" for k in col}
+                group["$group"]["_id"] = q
+            if sd and sd.multivalued:
                 facet_pipeline = [
                     {"$match": where} if where else {"$match": {}},
                     {"$unwind": f"${col}"},
-                    {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
+                    group,
                     {"$sort": {"count": -1}},
                     {"$limit": facet_limit},
                 ]
             else:
                 facet_pipeline = [
                     {"$match": where} if where else {"$match": {}},
-                    {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
+                    group,
                     {"$sort": {"count": -1}},
                     {"$limit": facet_limit},
                 ]
-
+            logger.info(f"Facet pipeline: {facet_pipeline}")
             facet_results = list(self.mongo_collection.aggregate(facet_pipeline))
             results[col] = [(result["_id"], result["count"]) for result in facet_results]
 
@@ -92,7 +113,7 @@ class MongoDBCollection(Collection):
         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}")
+        logger.info(f"Deleting from {self.target_class_name} where: {where}")
         if where is None:
             where = {}
         result = self.mongo_collection.delete_many(where)

linkml_store/api/stores/solr/__init__.py

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

linkml_store/api/stores/solr/solr_collection.py

@@ -5,6 +5,7 @@ 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
@@ -119,7 +120,7 @@ class SolrCollection(Collection):
         conditions = []
         if self.parent.metadata.collection_type_slot:
             where_clause = copy(where_clause)
-            where_clause[self.parent.metadata.collection_type_slot] = self._alias
+            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]

linkml_store/api/stores/solr/solr_database.py

@@ -2,6 +2,7 @@ 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

linkml_store/cli.py

@@ -11,12 +11,19 @@ from pydantic import BaseModel
 from linkml_store import Client
 from linkml_store.api import Collection, Database
 from linkml_store.api.queries import Query
+from linkml_store.index import get_indexer
 from linkml_store.index.implementations.simple_indexer import SimpleIndexer
 from linkml_store.index.indexer import Indexer
-from linkml_store.utils.format_utils import Format, load_objects, render_output
+from linkml_store.utils.format_utils import Format, guess_format, load_objects, render_output
 from linkml_store.utils.object_utils import object_path_update
 
-index_type_option = click.option("--index-type", "-t")
+index_type_option = click.option(
+    "--index-type",
+    "-t",
+    default="simple",
+    show_default=True,
+    help="Type of index to create. Values: simple, llm",
+)
 
 logger = logging.getLogger(__name__)
 
@@ -70,6 +77,9 @@ class ContextSettings(BaseModel):
 format_choice = click.Choice([f.value for f in Format])
 
 
+include_internal_option = click.option("--include-internal/--no-include-internal", default=False, show_default=True)
+
+
 @click.group()
 @click.option("--database", "-d", help="Database name")
 @click.option("--collection", "-c", help="Collection name")
@@ -89,6 +99,15 @@ def cli(ctx, verbose: int, quiet: bool, stacktrace: bool, database, collection,
     if not stacktrace:
         sys.tracebacklimit = 0
     logger = logging.getLogger()
+    # Set handler for the root logger to output to the console
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
+
+    # Clear existing handlers to avoid duplicate messages if function runs multiple times
+    logger.handlers = []
+
+    # Add the newly created console handler to the logger
+    logger.addHandler(console_handler)
     if verbose >= 2:
         logger.setLevel(logging.DEBUG)
     elif verbose == 1:
@@ -193,6 +212,35 @@ def store(ctx, files, object, format):
             click.echo(f"Inserted {len(objects)} objects from {object_str} into collection '{db.name}'.")
 
 
+@cli.command(name="import")
+@click.argument("files", type=click.Path(exists=True), nargs=-1)
+@click.option("--format", "-f", help="Input format")
+@click.pass_context
+def import_database(ctx, files, format):
+    """Imports a database from a dump."""
+    settings = ctx.obj["settings"]
+    db = settings.database
+    if not files and not object:
+        files = ["-"]
+    for file_path in files:
+        db.import_database(file_path, source_format=format)
+
+
+@cli.command()
+@click.option("--output-type", "-O", type=format_choice, default="json", help="Output format")
+@click.option("--output", "-o", required=True, type=click.Path(), help="Output file path")
+@click.pass_context
+def export(ctx, output_type, output):
+    """Exports a database to a dump."""
+    settings = ctx.obj["settings"]
+    db = settings.database
+    if output_type is None:
+        output_type = guess_format(output)
+    if output_type is None:
+        raise ValueError(f"Output format must be specified can't be inferred from {output}.")
+    db.export_database(output, target_format=output_type)
+
+
 @cli.command()
 @click.option("--where", "-w", type=click.STRING, help="WHERE clause for the query")
 @click.option("--limit", "-l", type=click.INT, help="Maximum number of results to return")
@@ -216,9 +264,10 @@ def query(ctx, where, limit, output_type, output):
 
 @cli.command()
 @click.pass_context
-def list_collections(ctx):
+@include_internal_option
+def list_collections(ctx, **kwargs):
     db = ctx.obj["settings"].database
-    for collection in db.list_collections():
+    for collection in db.list_collections(**kwargs):
         click.echo(collection.name)
         click.echo(render_output(collection.metadata))
 
@@ -254,7 +303,7 @@ def fq(ctx, where, limit, columns, output_type, output):
 
     def _untuple(key):
         if isinstance(key, tuple):
-            return "+".join(key)
+            return "+".join([str(x) for x in key])
         return key
 
     count_dict = {}
@@ -279,8 +328,10 @@ def _get_index(index_type=None, **kwargs) -> Indexer:
 
 @cli.command()
 @index_type_option
+@click.option("--cached-embeddings-database", "-E", help="Path to the database where embeddings are cached")
+@click.option("--text-template", "-T", help="Template for text embeddings")
 @click.pass_context
-def index(ctx, index_type):
+def index(ctx, index_type, **kwargs):
     """
     Create an index over a collection.
 
@@ -289,7 +340,7 @@ def index(ctx, index_type):
     :return:
     """
     collection = ctx.obj["settings"].collection
-    ix = _get_index(index_type)
+    ix = get_indexer(index_type, **kwargs)
     collection.attach_indexer(ix)
 
 
@@ -322,14 +373,17 @@ def schema(ctx, output_type, output):
 @click.option("--limit", "-l", type=click.INT, help="Maximum number of search results")
 @click.option("--output-type", "-O", type=format_choice, default="json", help="Output format")
 @click.option("--output", "-o", type=click.Path(), help="Output file path")
+@click.option(
+    "--auto-index/--no-auto-index", default=False, show_default=True, help="Automatically index the collection"
+)
 @index_type_option
 @click.pass_context
-def search(ctx, search_term, where, limit, index_type, output_type, output):
+def search(ctx, search_term, where, limit, index_type, output_type, output, auto_index):
     """Search objects in the specified collection."""
     collection = ctx.obj["settings"].collection
-    ix = _get_index(index_type)
+    ix = get_indexer(index_type)
     logger.info(f"Attaching index to collection {collection.name}: {ix.model_dump()}")
-    collection.attach_indexer(ix, auto_index=False)
+    collection.attach_indexer(ix, auto_index=auto_index)
     result = collection.search(search_term, where=where, limit=limit)
     output_data = render_output([{"score": row[0], **row[1]} for row in result.ranked_rows], output_type)
     if output:

linkml_store/index/__init__.py

@@ -22,7 +22,7 @@ def get_indexer_class(name: str) -> Type[Indexer]:
     return INDEXER_CLASSES[name]
 
 
-def get_indexer(name: str, *args, **kwargs) -> Indexer:
+def get_indexer(name: str, **kwargs) -> Indexer:
     """
     Get an indexer by name.
 
@@ -30,4 +30,8 @@ def get_indexer(name: str, *args, **kwargs) -> Indexer:
     :param kwargs: additional arguments to pass to the indexer
     :return: the indexer
     """
-    return get_indexer_class(name)(*args, **kwargs)
+    kwargs = {k: v for k, v in kwargs.items() if v is not None}
+    cls = get_indexer_class(name)
+    kwargs["name"] = name
+    indexer = cls(**kwargs)
+    return indexer