linkml-store 0.3.0__py3-none-any.whl

linkml_store/api/stores/duckdb/__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/duckdb/duckdb_collection.py

@@ -0,0 +1,339 @@
+import logging
+from typing import Any, Dict, List, Optional, Tuple, 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, QueryResult
+from linkml_store.api.stores.duckdb.mappings import TMAP
+from linkml_store.utils.sql_utils import facet_count_sql
+
+logger = logging.getLogger(__name__)
+
+
+class DuckDBCollection(Collection):
+    _table_created: bool = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def insert(self, objs: Union[OBJECT, List[OBJECT]], **kwargs):
+        logger.debug(f"Inserting {len(objs)}")
+        if not isinstance(objs, list):
+            objs = [objs]
+        if not objs:
+            return
+        cd = self.class_definition()
+        if not cd:
+            logger.debug(f"No class definition defined for {self.alias} {self.target_class_name}; will induce")
+            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]
+        bad_objs = [obj for obj in objs if not isinstance(obj, dict)]
+        if bad_objs:
+            logger.error(f"Bad objects: {bad_objs}")
+        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()
+        self._post_insert_hook(objs)
+
+    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 or not cd.attributes:
+            cd = self.induce_class_definition_from_objects(objs)
+        assert cd.attributes
+        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()
+        self._post_delete_hook()
+        return None
+
+    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()
+            self._post_delete_hook()
+            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[Union[str, Tuple[str, ...]], List[Tuple[Any, int]]]:
+        if facet_limit is None:
+            facet_limit = DEFAULT_FACET_LIMIT
+        results = {}
+        cd = self.class_definition()
+        with self.parent.engine.connect() as conn:
+            if not facet_columns:
+                if not cd:
+                    raise ValueError(f"No class definition found for {self.target_class_name}")
+                facet_columns = list(cd.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] = [tuple(row) for row in 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 or att.inlined_as_list:
+                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 _check_if_initialized(self) -> bool:
+        # if self._initialized:
+        #    return True
+        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:
+            return True
+        return False
+
+    def group_by(
+        self,
+        group_by_fields: List[str],
+        inlined_field="objects",
+        agg_map: Optional[Dict[str, str]] = None,
+        where: Optional[Dict] = None,
+        **kwargs,
+    ) -> QueryResult:
+        """
+        Group objects in the collection by specified fields using SQLAlchemy.
+
+        This implementation leverages DuckDB's SQL capabilities for more efficient grouping.
+
+        :param group_by_fields: List of fields to group by
+        :param inlined_field: Field name to store aggregated objects
+        :param agg_map: Dictionary mapping aggregation types to fields
+        :param where: Filter conditions
+        :param kwargs: Additional arguments
+        :return: Query result containing grouped data
+        """
+        if isinstance(group_by_fields, str):
+            group_by_fields = [group_by_fields]
+
+        cd = self.class_definition()
+        if not cd:
+            logger.debug(f"No class definition defined for {self.alias} {self.target_class_name}")
+            return super().group_by(group_by_fields, inlined_field, agg_map, where, **kwargs)
+
+        # Check if the table exists
+        if not self.parent._table_exists(self.alias):
+            logger.debug(f"Table {self.alias} doesn't exist, falling back to parent implementation")
+            return super().group_by(group_by_fields, inlined_field, agg_map, where, **kwargs)
+
+        # Get table definition
+        table = self._sqla_table(cd)
+        engine = self.parent.engine
+
+        # Create a SQLAlchemy select statement for groups
+        from sqlalchemy import select
+
+        group_cols = [table.c[field] for field in group_by_fields if field in table.columns.keys()]
+
+        if not group_cols:
+            logger.warning(f"None of the group_by fields {group_by_fields} found in table columns")
+            return super().group_by(group_by_fields, inlined_field, agg_map, where, **kwargs)
+
+        stmt = select(*group_cols).distinct()
+
+        # Add where conditions if specified
+        if where:
+            conditions = []
+            for k, v in where.items():
+                if k in table.columns.keys():
+                    # Handle different operator types (dict values for operators)
+                    if isinstance(v, dict):
+                        for op, val in v.items():
+                            if op == "$gt":
+                                conditions.append(table.c[k] > val)
+                            elif op == "$gte":
+                                conditions.append(table.c[k] >= val)
+                            elif op == "$lt":
+                                conditions.append(table.c[k] < val)
+                            elif op == "$lte":
+                                conditions.append(table.c[k] <= val)
+                            elif op == "$ne":
+                                conditions.append(table.c[k] != val)
+                            elif op == "$in":
+                                conditions.append(table.c[k].in_(val))
+                            else:
+                                # Default to equality for unknown operators
+                                logger.warning(f"Unknown operator {op}, using equality")
+                                conditions.append(table.c[k] == val)
+                    else:
+                        # Direct equality comparison
+                        conditions.append(table.c[k] == v)
+
+            if conditions:
+                for condition in conditions:
+                    stmt = stmt.where(condition)
+
+        results = []
+        try:
+            with engine.connect() as conn:
+                # Get all distinct groups
+                group_result = conn.execute(stmt)
+                group_rows = list(group_result)
+
+                # For each group, get all objects
+                for group_row in group_rows:
+                    # Build conditions for this group
+                    group_conditions = []
+                    group_dict = {}
+
+                    for i, field in enumerate(group_by_fields):
+                        if field in table.columns.keys():
+                            value = group_row[i]
+                            group_dict[field] = value
+                            if value is None:
+                                group_conditions.append(table.c[field].is_(None))
+                            else:
+                                group_conditions.append(table.c[field] == value)
+
+                    # Get all rows for this group
+                    row_stmt = select(*table.columns)
+                    for condition in group_conditions:
+                        row_stmt = row_stmt.where(condition)
+
+                    # Add original where conditions
+                    if where:
+                        for k, v in where.items():
+                            if k in table.columns.keys():
+                                # Handle different operator types for the row query as well
+                                if isinstance(v, dict):
+                                    for op, val in v.items():
+                                        if op == "$gt":
+                                            row_stmt = row_stmt.where(table.c[k] > val)
+                                        elif op == "$gte":
+                                            row_stmt = row_stmt.where(table.c[k] >= val)
+                                        elif op == "$lt":
+                                            row_stmt = row_stmt.where(table.c[k] < val)
+                                        elif op == "$lte":
+                                            row_stmt = row_stmt.where(table.c[k] <= val)
+                                        elif op == "$ne":
+                                            row_stmt = row_stmt.where(table.c[k] != val)
+                                        elif op == "$in":
+                                            row_stmt = row_stmt.where(table.c[k].in_(val))
+                                        else:
+                                            # Default to equality for unknown operators
+                                            row_stmt = row_stmt.where(table.c[k] == val)
+                                else:
+                                    # Direct equality comparison
+                                    row_stmt = row_stmt.where(table.c[k] == v)
+
+                    row_result = conn.execute(row_stmt)
+                    rows = list(row_result)
+
+                    # Convert rows to dictionaries
+                    objects = []
+                    for row in rows:
+                        obj = {}
+                        for i, col in enumerate(row._fields):
+                            obj[col] = row[i]
+                        objects.append(obj)
+
+                    # Apply agg_map to filter fields if specified
+                    if agg_map and "list" in agg_map:
+                        list_fields = agg_map["list"]
+                        if list_fields:
+                            objects = [{k: obj.get(k) for k in list_fields if k in obj} for obj in objects]
+
+                    # Create the result object
+                    result_obj = group_dict.copy()
+                    result_obj[inlined_field] = objects
+                    results.append(result_obj)
+
+                return QueryResult(num_rows=len(results), rows=results)
+        except Exception as e:
+            logger.warning(f"Error in DuckDB group_by: {e}")
+            # Fall back to parent implementation
+            return super().group_by(group_by_fields, inlined_field, agg_map, where, **kwargs)
+
+    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
+        if self.parent._table_exists(self.alias):
+            logger.info(f"Table already exists for {cd.name}")
+            self._table_created = True
+            self._initialized = True
+            self.metadata.is_prepopulated = True
+            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._initialized = 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._initialized = True
+        self.metadata.is_prepopulated = True

linkml_store/api/stores/duckdb/duckdb_database.py

@@ -0,0 +1,283 @@
+import json
+import logging
+from pathlib import Path
+from typing import List, Optional, Union
+
+import pandas as pd
+import sqlalchemy
+from linkml_runtime import SchemaView
+from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
+from linkml_runtime.utils.schema_builder import SchemaBuilder
+from sqlalchemy import NullPool, 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.format_utils import Format
+from linkml_store.utils.sql_utils import introspect_schema, query_to_sql
+
+TYPE_MAP = {
+    "VARCHAR": "string",
+    "BIGINT": "integer",
+    "BOOLEAN": "boolean",
+    "DATE": "date",
+    "DOUBLE": "float",
+    "INTEGER": "integer",
+    "JSON": "Any",
+}
+
+MEMORY_HANDLE = "duckdb:///:memory:"
+
+
+logger = logging.getLogger(__name__)
+
+
+class DuckDBDatabase(Database):
+    """
+    An adapter for DuckDB databases.
+
+    Note that this adapter does not make use of a LinkML relational model transformation and
+    SQL Alchemy ORM layer. Instead, it attempts to map each collection (which is of type
+    some LinkML class) to a *single* DuckDB table. New tables are not created for nested references,
+    and linking tables are not created for many-to-many relationships.
+
+    Instead the native DuckDB ARRAY type is used to store multivalued attributes, and DuckDB JSON
+    types are used for nested inlined objects.
+    """
+
+    # _connection: DuckDBPyConnection = None
+    _engine: sqlalchemy.Engine = None
+    collection_class = DuckDBCollection
+
+    def __init__(self, handle: Optional[str] = None, recreate_if_exists: bool = False, **kwargs):
+        if handle is None:
+            handle = MEMORY_HANDLE
+        if recreate_if_exists:
+            path = Path(handle.replace("duckdb:///", ""))
+            if path.exists():
+                path.unlink()
+        super().__init__(handle=handle, **kwargs)
+
+    @property
+    def engine(self) -> sqlalchemy.Engine:
+        if not self._engine:
+            handle = self.handle
+            if not handle.startswith("duckdb://") and not handle.startswith(":") and "://" not in handle:
+                handle = f"duckdb:///{handle}"
+            if ":memory:" not in handle:
+                # TODO: investigate this; duckdb appears to be prematurely caching
+                self._engine = sqlalchemy.create_engine(handle, poolclass=NullPool)
+            else:
+                self._engine = sqlalchemy.create_engine(handle)
+        return self._engine
+
+    @property
+    def _is_sqlite(self) -> bool:
+        return self.handle and self.handle.startswith("sqlite:")
+
+    def commit(self, **kwargs):
+        with self.engine.connect() as conn:
+            conn.commit()
+
+    def close(self, **kwargs):
+        self.engine.dispose()
+
+    def drop(self, missing_ok=True, **kwargs):
+        self.close()
+        if self.handle == MEMORY_HANDLE:
+            return
+        path = Path(self.handle.replace("duckdb:///", ""))
+        if path.exists():
+            path.unlink()
+        else:
+            if not missing_ok:
+                raise FileNotFoundError(f"Database file not found: {path}")
+
+    def _table_exists(self, table: str) -> bool:
+        if self._is_sqlite:
+            if table == "sqlite_master":
+                return True
+            meta_query = Query(
+                from_table="sqlite_master",
+                where_clause={
+                    # "type": "table",
+                    "name": table,
+                },
+            )
+        else:
+            if table.startswith("information_schema"):
+                return True
+            meta_query = Query(
+                from_table="information_schema.tables",
+                where_clause={
+                    "table_type": "BASE TABLE",
+                    "table_name": table,
+                },
+            )
+
+        qr = self.query(meta_query)
+        if qr.num_rows == 0:
+            logger.debug(f"Table {self.alias} not created yet")
+            return False
+        return True
+
+    def _json_encoded_cols(self, table_name: str) -> Optional[List[str]]:
+        json_encoded_cols = []
+        if table_name:
+            if table_name.startswith("information_schema") or table_name.startswith("sqlite"):
+                return []
+            sv = self.schema_view
+            if sv:
+                cd = None
+                for c in self._collections.values():
+                    if c.alias == table_name or c.target_class_name == table_name:
+                        cd = c.class_definition()
+                        break
+                if cd:
+                    for att in sv.class_induced_slots(cd.name):
+                        if att.inlined or att.inlined_as_list:
+                            json_encoded_cols.append(att.name)
+        return json_encoded_cols
+
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        if not self._table_exists(query.from_table):
+            return QueryResult(query=query, num_rows=0, rows=[])
+        json_encoded_cols = self._json_encoded_cols(query.from_table)
+
+        with self.engine.connect() as conn:
+            count_query_str = text(query_to_sql(query, count=True))
+            logger.debug(f"count_query_str: {count_query_str}")
+            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])):
+                                    try:
+                                        parsed_val = json.loads(row[col][i])
+                                    except json.JSONDecodeError as e:
+                                        logger.error(f"Failed to parse col {col}[{i}] == {row[col][i]}")
+                                        raise e
+                                    row[col][i] = parsed_val
+                            elif isinstance(row[col], dict):
+                                pass
+                            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
+
+    @property
+    def supports_sql(self) -> bool:
+        """Return True - DuckDB supports raw SQL queries."""
+        return True
+
+    def execute_sql(self, sql: str, **kwargs) -> QueryResult:
+        """
+        Execute a raw SQL query against the DuckDB database.
+
+        :param sql: SQL query string
+        :param kwargs: Additional arguments
+        :return: QueryResult containing the results
+        """
+        with self.engine.connect() as conn:
+            result = conn.execute(text(sql))
+            rows = [dict(row._mapping) for row in result]
+            return QueryResult(num_rows=len(rows), rows=rows)
+
+    def init_collections(self):
+        # TODO: unify schema introspection
+        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 = {}
+        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 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} // {self}")
+        sb = SchemaBuilder()
+        schema = sb.schema
+        logger.info(f"Checking if {self.metadata.handle} is sqlite: {self._is_sqlite}")
+        if self._is_sqlite:
+            return SchemaView(schema)
+        query = Query(from_table="information_schema.tables", where_clause={"table_type": "BASE TABLE"})
+        qr = self.query(query)
+        logger.info(f"Found {qr.num_rows} information_schema.tables // {qr.rows}")
+        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
+            )
+            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} FROM {dt}")
+        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)
+
+    def export_database(self, location: str, target_format: Optional[Union[str, Format]] = None, **kwargs):
+        if target_format == "duckdb" or target_format == Format.SQLDUMP_DUCKDB:
+            path = Path(location)
+            if path.exists():
+                if path.is_file():
+                    path.unlink()
+            with self.engine.connect() as conn:
+                sql = text(f"EXPORT DATABASE '{location}'")
+                conn.execute(sql)
+        else:
+            super().export_database(location, target_format=target_format, **kwargs)
+
+    def import_database(self, location: str, source_format: Optional[str] = None, **kwargs):
+        """
+        Import a database from a file or location.
+
+        :param location: location of the file
+        :param source_format: source format
+        :param kwargs: additional arguments
+        """
+        if source_format == Format.SQLDUMP_DUCKDB.value or source_format == Format.SQLDUMP_DUCKDB:
+            with self.engine.connect() as conn:
+                sql = text(f"IMPORT DATABASE '{location}'")
+                conn.execute(sql)
+                conn.commit()
+        else:
+            super().import_database(location, source_format=source_format, **kwargs)

linkml_store/api/stores/duckdb/mappings.py

@@ -0,0 +1,8 @@
+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,15 @@
+"""
+Adapter for FileSystem wrapper
+
+Handles have the form:
+
+ - ``file:<path>`` for a local file
+"""
+
+from linkml_store.api.stores.filesystem.filesystem_collection import FileSystemCollection
+from linkml_store.api.stores.filesystem.filesystem_database import FileSystemDatabase
+
+__all__ = [
+    "FileSystemCollection",
+    "FileSystemDatabase",
+]