linkml-store 0.0.0__py3-none-any.whl

linkml_store/api/database.py

@@ -0,0 +1,215 @@
+from abc import ABC
+from dataclasses import dataclass
+from typing import Dict, Optional, Sequence
+
+from linkml_runtime import SchemaView
+
+from linkml_store.api.collection import Collection
+from linkml_store.api.metadata import MetaData
+from linkml_store.api.queries import Query, QueryResult
+
+
+@dataclass
+class Database(ABC):
+    """
+    A Database provides access to named collections of data.
+
+    Examples
+    --------
+    >>> from linkml_store.api.client import Client
+    >>> client = Client()
+    >>> db = client.attach_database("duckdb", alias="test")
+    >>> db.handle
+    'duckdb:///:memory:'
+    >>> collection = db.create_collection("Person")
+    >>> len(db.list_collections())
+    1
+    >>> 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)
+    >>> qr = collection.find()
+    >>> len(qr.rows)
+    2
+    >>> qr.rows[0]["id"]
+    'P1'
+    >>> qr.rows[1]["name"]
+    'Alice'
+    >>> qr = collection.find({"name": "John"})
+    >>> len(qr.rows)
+    1
+    >>> qr.rows[0]["name"]
+    'John'
+
+    """
+
+    handle: Optional[str] = None
+    recreate_if_exists: Optional[bool] = False
+    _schema_view: Optional[SchemaView] = None
+    _collections: Optional[Dict[str, Collection]] = None
+
+    def store(self, obj: Dict[str, str], **kwargs):
+        """
+        Store an object in the database
+
+        :param obj: object to store
+        :param kwargs: additional arguments
+        """
+        for k, v in obj.items():
+            if not isinstance(v, list):
+                continue
+            if not v:
+                continue
+            collection = self.get_collection(k, create_if_not_exists=True)
+            collection.add(v)
+
+    def commit(self, **kwargs):
+        """
+        Commit any pending changes to the database
+        """
+        raise NotImplementedError()
+
+    def close(self, **kwargs):
+        """
+        Close the database and all connection objects
+        """
+        raise NotImplementedError()
+
+    def create_collection(
+        self, name: str, alias: Optional[str] = None, metadata: Optional[MetaData] = None, **kwargs
+    ) -> Collection:
+        """
+        Create a new collection
+
+        >>> from linkml_store.api.client import Client
+        >>> client = Client()
+        >>> db = client.attach_database("duckdb", alias="test")
+        >>> collection = db.create_collection("Person")
+        >>> collection.name
+        'Person'
+
+        :param name: name of the collection
+        :param alias: alias for the collection
+        :param metadata: metadata for the collection
+        :param kwargs: additional arguments
+        """
+        raise NotImplementedError()
+
+    def list_collections(self) -> Sequence[Collection]:
+        """
+        List all collections.
+
+        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")
+        >>> collections = db.list_collections()
+        >>> len(collections)
+        2
+        >>> [c.name for c in collections]
+        ['Person', 'Product']
+
+        """
+        if not self._collections:
+            self.init_collections()
+        return list(self._collections.values())
+
+    def get_collection(self, name: str, create_if_not_exists=True, **kwargs) -> "Collection":
+        """
+        Get a named collection.
+
+        Examples
+        --------
+        >>> from linkml_store.api.client import Client
+        >>> client = Client()
+        >>> db = client.attach_database("duckdb", alias="test")
+        >>> collection = db.create_collection("Person")
+        >>> db.get_collection("Person") == collection
+        True
+        >>> db.get_collection("NonExistent", create_if_not_exists=False)
+        Traceback (most recent call last):
+            ...
+        KeyError: 'Collection NonExistent does not exist'
+
+        :param name: name of the collection
+        :param create_if_not_exists: create the collection if it does not exist
+
+        """
+        if not self._collections:
+            self.init_collections()
+        if name not in self._collections:
+            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 init_collections(self):
+        """
+        Initialize collections.
+
+        Not typically called directly: consider making hidden
+        :return:
+        """
+        raise NotImplementedError
+
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        """
+        Run a query against the database.
+
+        Examples
+        --------
+        >>> from linkml_store.api.client import Client
+        >>> from linkml_store.api.queries import Query
+        >>> client = Client()
+        >>> db = client.attach_database("duckdb", alias="test")
+        >>> collection = db.create_collection("Person")
+        >>> collection.add([{"id": "P1", "name": "John"}, {"id": "P2", "name": "Alice"}])
+        >>> query = Query(from_table="Person", where_clause={"name": "John"})
+        >>> result = db.query(query)
+        >>> len(result.rows)
+        1
+        >>> result.rows[0]["id"]
+        'P1'
+
+        :param query:
+        :param kwargs:
+        :return:
+
+        """
+        raise NotImplementedError
+
+    @property
+    def schema_view(self) -> SchemaView:
+        """
+        Return a schema view for the named collection
+        """
+        if not self._schema_view:
+            self._schema_view = self.induce_schema_view()
+        return self._schema_view
+
+    def set_schema_view(self, schema_view: SchemaView):
+        self._schema_view = schema_view
+
+    def induce_schema_view(self) -> SchemaView:
+        """
+        Induce a schema view from a schema definition.
+
+        >>> from linkml_store.api.client import Client
+        >>> from linkml_store.api.queries import Query
+        >>> client = Client()
+        >>> db = client.attach_database("duckdb", alias="test")
+        >>> collection = db.create_collection("Person")
+        >>> collection.add([{"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'
+
+        :return: A schema view
+        """
+        raise NotImplementedError()

linkml_store/api/metadata.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class MetaData(BaseModel):
+    pass

linkml_store/api/queries.py

@@ -0,0 +1,56 @@
+from collections import namedtuple
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import pandas as pd
+from pydantic import BaseModel
+
+# defined a named tuple called between with two values (start, end):
+# This is used in the Query class to represent a range of values
+# This is used in the Query class to represent a range of values
+Between = namedtuple("Between", "min max")
+
+FACET_GROUP_ATOM = Union[str, int, float, Between]
+FACET_GROUP = Union[FACET_GROUP_ATOM, Tuple[FACET_GROUP_ATOM, ...]]
+
+
+class Query(BaseModel):
+    """
+    A query object.
+
+    - In SQL this would be a SQL query string
+    """
+
+    from_table: Optional[str]
+    select_cols: Optional[List[str]] = None
+    where_clause: Optional[Union[str, List[str], Dict[str, str]]] = None
+    sort_by: Optional[List[str]] = None
+    limit: Optional[int] = None
+    offset: Optional[int] = None
+    include_facet_counts: bool = False
+    facet_slots: Optional[List[str]] = None
+
+
+class QueryResult(BaseModel):
+    """
+    A query result
+    """
+
+    query: Optional[Query] = None
+    num_rows: int
+    offset: Optional[int] = 0
+    rows: Optional[List[Dict[str, Any]]] = None
+    ranked_rows: Optional[List[Tuple[float, Dict[str, Any]]]] = None
+    _rows_dataframe: Optional[pd.DataFrame] = None
+    facet_counts: Optional[Dict[str, List[Tuple[FACET_GROUP, int]]]] = None
+
+    @property
+    def rows_dataframe(self) -> pd.DataFrame:
+        if self._rows_dataframe is None and self.rows:
+            self._rows_dataframe = pd.DataFrame(self.rows)
+        return self._rows_dataframe
+
+    def set_rows(self, rows: pd.DataFrame):
+        self._rows_dataframe = rows
+
+    class Config:
+        arbitrary_types_allowed = True

linkml_store/api/stores/duckdb/duckdb_collection.py

@@ -0,0 +1,109 @@
+from dataclasses import dataclass
+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, text
+from sqlalchemy.sql.ddl import CreateTable
+
+from linkml_store.api import Collection
+from linkml_store.api.collection import OBJECT
+from linkml_store.api.stores.duckdb.mappings import TMAP
+from linkml_store.utils.sql_utils import facet_count_sql
+
+
+@dataclass
+class DuckDBCollection(Collection):
+    _table_created: bool = None
+
+    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)
+        self._create_table(cd)
+        table = self._sqla_table(cd)
+        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) -> 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 len(objs)
+
+    def delete_where(self, where: Optional[Dict[str, Any]] = None, **kwargs) -> int:
+        cd = self.class_definition()
+        table = self._sqla_table(cd)
+        engine = self.parent.engine
+        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)
+            conn.execute(stmt)
+            conn.commit()
+        return 0
+
+    def query_facets(self, where: Dict = None, facet_columns: List[str] = None) -> 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:
+                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)
+                rows = list(conn.execute(text(facet_query_str)))
+                results[col] = rows
+            return results
+
+    def _sqla_table(self, cd: ClassDefinition) -> Table:
+        metadata_obj = sqla.MetaData()
+        cols = []
+        for att in cd.attributes.values():
+            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.name, metadata_obj, *cols)
+        return t
+
+    def _create_table(self, cd: ClassDefinition):
+        if self._table_created:
+            return
+        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()
+        if not self._table_created:
+            self._table_created = True

linkml_store/api/stores/duckdb/duckdb_database.py

@@ -0,0 +1,166 @@
+import json
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+import pandas as pd
+import sqlalchemy
+from duckdb import DuckDBPyConnection
+from linkml_runtime import SchemaView
+from linkml_runtime.linkml_model import SlotDefinition
+from linkml_runtime.utils.schema_builder import SchemaBuilder
+from sqlalchemy import text
+
+from linkml_store.api import Database
+from linkml_store.api.queries import Query, QueryResult
+from linkml_store.api.stores.duckdb.duckdb_collection import DuckDBCollection
+from linkml_store.utils.sql_utils import introspect_schema, query_to_sql
+
+TYPE_MAP = {
+    "VARCHAR": "string",
+    "BIGINT": "integer",
+    "BOOLEAN": "boolean",
+}
+
+
+logger = logging.getLogger(__name__)
+
+
+def run_query(con: DuckDBPyConnection, query: Query, **kwargs):
+    """
+    Run a query and return the result.
+
+    >>> import duckdb
+    >>> con = duckdb.connect("db/mgi.db")
+    >>> query = Query(from_table="gaf_association", limit=5)
+    >>> result = run_query(con, query)
+    >>> print(result.num_rows)
+    532233
+
+    :param con:
+    :param query:
+    :return:
+    """
+    count_query_str = query_to_sql(query, count=True)
+    num_rows = con.execute(count_query_str).fetchall()[0][0]
+    logger.debug(f"num_rows: {num_rows}")
+    query_str = query_to_sql(query, **kwargs)
+    logger.debug(f"query_str: {query_str}")
+    rows = con.execute(query_str).fetchdf()
+    qr = QueryResult(query=query, num_rows=num_rows)
+    qr.set_rows(rows)
+    return qr
+
+
+@dataclass
+class DuckDBDatabase(Database):
+    _connection: DuckDBPyConnection = None
+    _engine: sqlalchemy.Engine = None
+
+    def __post_init__(self):
+        if not self.handle:
+            self.handle = "duckdb:///:memory:"
+
+    @property
+    def engine(self) -> sqlalchemy.Engine:
+        if not self._engine:
+            handle = self.handle
+            if not handle.startswith("duckdb://") and not handle.startswith(":"):
+                handle = f"duckdb://{handle}"
+            self._engine = sqlalchemy.create_engine(handle)
+        return self._engine
+
+    def commit(self, **kwargs):
+        with self.engine.connect() as conn:
+            conn.commit()
+
+    def close(self, **kwargs):
+        self.engine.dispose()
+
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        json_encoded_cols = []
+        if query.from_table:
+            sv = self._schema_view
+            if sv:
+                cd = None
+                for c in self._collections.values():
+                    if c.name == query.from_table:
+                        cd = c.class_definition()
+                        break
+                if cd:
+                    for att in cd.attributes.values():
+                        if att.inlined:
+                            json_encoded_cols.append(att.name)
+        with self.engine.connect() as conn:
+            count_query_str = text(query_to_sql(query, count=True))
+            num_rows = list(conn.execute(count_query_str))[0][0]
+            logger.debug(f"num_rows: {num_rows}")
+            query_str = query_to_sql(query, **kwargs)  # include offset, limit
+            logger.debug(f"query_str: {query_str}")
+            rows = list(conn.execute(text(query_str)).mappings())
+            qr = QueryResult(query=query, num_rows=num_rows, rows=rows)
+            if json_encoded_cols:
+                for row in qr.rows:
+                    for col in json_encoded_cols:
+                        if row[col]:
+                            if isinstance(row[col], list):
+                                for i in range(len(row[col])):
+                                    row[col][i] = json.loads(row[col][i])
+                            else:
+                                row[col] = json.loads(row[col])
+            qr.set_rows(pd.DataFrame(rows))
+            facet_columns = query.facet_slots
+            if query.include_facet_counts and not facet_columns:
+                raise ValueError("Facet counts requested but no facet columns specified")
+            if facet_columns:
+                raise NotImplementedError
+            return qr
+
+    def init_collections(self):
+        # TODO: unify schema introspection
+        schema = introspect_schema(self.engine)
+        table_names = schema.classes.keys()
+        if self._collections is None:
+            self._collections = {}
+        for table_name in table_names:
+            if table_name not in self._collections:
+                collection = DuckDBCollection(name=table_name, parent=self)
+                self._collections[table_name] = collection
+
+    def create_collection(self, name: str, alias: Optional[str] = None, **kwargs) -> DuckDBCollection:
+        collection = DuckDBCollection(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:
+        # TODO: unify schema introspection
+        sb = SchemaBuilder()
+        schema = sb.schema
+        query = Query(from_table="information_schema.tables", where_clause={"table_type": "BASE TABLE"})
+        qr = self.query(query)
+        if qr.num_rows:
+            table_names = [row["table_name"] for row in qr.rows]
+            for tbl in table_names:
+                sb.add_class(tbl)
+        query = Query(from_table="information_schema.columns", sort_by=["ordinal_position"])
+        for row in self.query(query, limit=-1).rows:
+            tbl_name = row["table_name"]
+            if tbl_name not in sb.schema.classes:
+                continue
+            dt = row["data_type"]
+            if dt.endswith("[]"):
+                dt = dt[0:-2]
+                multivalued = True
+            else:
+                multivalued = False
+            rng = TYPE_MAP.get(dt, "string")
+            sd = SlotDefinition(
+                row["column_name"], required=row["is_nullable"] == "NO", multivalued=multivalued, range=rng
+            )
+            sb.schema.classes[tbl_name].attributes[sd.name] = sd
+        sb.add_defaults()
+        return SchemaView(schema)

linkml_store/api/stores/duckdb/mappings.py

@@ -0,0 +1,7 @@
+import sqlalchemy as sqla
+
+TMAP = {
+    "string": sqla.String,
+    "integer": sqla.Integer,
+    "linkml:Any": sqla.JSON,
+}

linkml_store/api/stores/mongodb/mongodb_collection.py

@@ -0,0 +1,56 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+from linkml_store.api import Collection
+from linkml_store.api.collection import OBJECT
+
+
+@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)
+
+    def delete(self, objs: Union[OBJECT, List[OBJECT]], **kwargs) -> int:
+        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
+
+    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_facets(self, where: Dict = None, facet_columns: List[str] = None) -> Dict[str, Dict[str, int]]:
+        results = {}
+        _cd = self.class_definition()
+        collection = self.parent.database[self.name]
+        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]
+        return results