deriva-ml 1.17.10__py3-none-any.whl → 1.17.11__py3-none-any.whl

deriva_ml/model/catalog.py

@@ -12,10 +12,17 @@ from collections import Counter, defaultdict
 from graphlib import CycleError, TopologicalSorter
 from typing import Any, Callable, Final, Iterable, NewType, TypeAlias
 
-from deriva.core.ermrest_catalog import ErmrestCatalog
-
-# Deriva imports
-from deriva.core.ermrest_model import Column, FindAssociationResult, Model, Schema, Table
+# Deriva imports - use importlib to avoid shadowing by local 'deriva.py' files
+import importlib
+_ermrest_catalog = importlib.import_module("deriva.core.ermrest_catalog")
+_ermrest_model = importlib.import_module("deriva.core.ermrest_model")
+
+ErmrestCatalog = _ermrest_catalog.ErmrestCatalog
+Column = _ermrest_model.Column
+FindAssociationResult = _ermrest_model.FindAssociationResult
+Model = _ermrest_model.Model
+Schema = _ermrest_model.Schema
+Table = _ermrest_model.Table
 
 # Third-party imports
 from pydantic import ConfigDict, validate_call
@@ -23,14 +30,16 @@ from pydantic import ConfigDict, validate_call
 from deriva_ml.core.definitions import (
     ML_SCHEMA,
     RID,
+    SYSTEM_SCHEMAS,
     DerivaAssetColumns,
     TableDefinition,
+    get_domain_schemas,
+    is_system_schema,
 )
 from deriva_ml.core.exceptions import DerivaMLException, DerivaMLTableTypeError
 
 # Local imports
 from deriva_ml.feature import Feature
-from deriva_ml.protocols.dataset import DatasetLike
 
 try:
     from icecream import ic
@@ -61,12 +70,12 @@ class DerivaModel:
     This class provides a number of DerivaML specific methods that augment the interface in the deriva model class.
 
     Attributes:
-        domain_schema: Schema name for domain-specific tables and relationships.
         model: ERMRest model for the catalog.
-        catalog: ERMRest catalog for the model
-        hostname: ERMRest catalog for the model
-        ml_schema: The ML schema for the catalog.
-        domain_schema: The domain schema for the catalog.
+        catalog: ERMRest catalog for the model.
+        hostname: Hostname of the ERMRest server.
+        ml_schema: The ML schema name for the catalog.
+        domain_schemas: Frozenset of all domain schema names in the catalog.
+        default_schema: The default schema for table creation operations.
 
     """
 
@@ -74,17 +83,22 @@ class DerivaModel:
         self,
         model: Model,
         ml_schema: str = ML_SCHEMA,
-        domain_schema: str | None = None,
+        domain_schemas: set[str] | None = None,
+        default_schema: str | None = None,
     ):
-        """Create and initialize a DerivaML instance.
+        """Create and initialize a DerivaModel instance.
 
-        This method will connect to a catalog, and initialize local configuration for the ML execution.
+        This method will connect to a catalog and initialize schema configuration.
         This class is intended to be used as a base class on which domain-specific interfaces are built.
 
         Args:
             model: The ERMRest model for the catalog.
             ml_schema: The ML schema name.
-            domain_schema: The domain schema name.
+            domain_schemas: Optional explicit set of domain schema names. If None,
+                auto-detects all non-system schemas.
+            default_schema: The default schema for table creation operations. If None
+                and there is exactly one domain schema, that schema is used as default.
+                If there are multiple domain schemas, default_schema must be specified.
         """
         self.model = model
         self.configuration = None
@@ -92,27 +106,182 @@ class DerivaModel:
         self.hostname = self.catalog.deriva_server.server if isinstance(self.catalog, ErmrestCatalog) else "localhost"
 
         self.ml_schema = ml_schema
-        builtin_schemas = ("public", self.ml_schema, "www", "WWW")
-        if domain_schema:
-            self.domain_schema = domain_schema
+        self._system_schemas = frozenset(SYSTEM_SCHEMAS | {ml_schema})
+
+        # Determine domain schemas
+        if domain_schemas is not None:
+            self.domain_schemas = frozenset(domain_schemas)
         else:
-            if len(user_schemas := {k for k in self.model.schemas.keys()} - set(builtin_schemas)) == 1:
-                self.domain_schema = user_schemas.pop()
-            else:
-                raise DerivaMLException(f"Ambiguous domain schema: {user_schemas}")
+            # Auto-detect all domain schemas
+            self.domain_schemas = get_domain_schemas(self.model.schemas.keys(), ml_schema)
+
+        # Determine default schema for table creation
+        if default_schema is not None:
+            if default_schema not in self.domain_schemas:
+                raise DerivaMLException(
+                    f"default_schema '{default_schema}' is not in domain_schemas: {self.domain_schemas}"
+                )
+            self.default_schema = default_schema
+        elif len(self.domain_schemas) == 1:
+            # Single domain schema - use it as default
+            self.default_schema = next(iter(self.domain_schemas))
+        elif len(self.domain_schemas) == 0:
+            # No domain schemas - default_schema will be None
+            self.default_schema = None
+        else:
+            # Multiple domain schemas, no explicit default
+            self.default_schema = None
+
+    def is_system_schema(self, schema_name: str) -> bool:
+        """Check if a schema is a system or ML schema.
+
+        Args:
+            schema_name: Name of the schema to check.
+
+        Returns:
+            True if the schema is a system or ML schema.
+        """
+        return is_system_schema(schema_name, self.ml_schema)
+
+    def is_domain_schema(self, schema_name: str) -> bool:
+        """Check if a schema is a domain schema.
+
+        Args:
+            schema_name: Name of the schema to check.
+
+        Returns:
+            True if the schema is a domain schema.
+        """
+        return schema_name in self.domain_schemas
+
+    def _require_default_schema(self) -> str:
+        """Get default schema, raising an error if not set.
+
+        Returns:
+            The default schema name.
+
+        Raises:
+            DerivaMLException: If default_schema is not set.
+        """
+        if self.default_schema is None:
+            raise DerivaMLException(
+                f"No default_schema set. With multiple domain schemas {self.domain_schemas}, "
+                "you must either specify a default_schema when creating DerivaML or "
+                "pass an explicit schema parameter to this method."
+            )
+        return self.default_schema
 
     def refresh_model(self) -> None:
         self.model = self.catalog.getCatalogModel()
 
-    @property
-    def schemas(self) -> dict[str, Schema]:
-        return self.model.schemas
-
     @property
     def chaise_config(self) -> dict[str, Any]:
         """Return the chaise configuration."""
         return self.model.chaise_config
 
+    def get_schema_description(self, include_system_columns: bool = False) -> dict[str, Any]:
+        """Return a JSON description of the catalog schema structure.
+
+        Provides a structured representation of the domain and ML schemas including
+        tables, columns, foreign keys, and relationships. Useful for understanding
+        the data model structure programmatically.
+
+        Args:
+            include_system_columns: If True, include RID, RCT, RMT, RCB, RMB columns.
+                Default False to reduce output size.
+
+        Returns:
+            Dictionary with schema structure:
+            {
+                "domain_schemas": ["schema_name1", "schema_name2"],
+                "default_schema": "schema_name1",
+                "ml_schema": "deriva-ml",
+                "schemas": {
+                    "schema_name": {
+                        "tables": {
+                            "TableName": {
+                                "comment": "description",
+                                "is_vocabulary": bool,
+                                "is_asset": bool,
+                                "is_association": bool,
+                                "columns": [...],
+                                "foreign_keys": [...],
+                                "features": [...]
+                            }
+                        }
+                    }
+                }
+            }
+        """
+        system_columns = {"RID", "RCT", "RMT", "RCB", "RMB"}
+        result = {
+            "domain_schemas": sorted(self.domain_schemas),
+            "default_schema": self.default_schema,
+            "ml_schema": self.ml_schema,
+            "schemas": {},
+        }
+
+        # Include all domain schemas and the ML schema
+        for schema_name in [*self.domain_schemas, self.ml_schema]:
+            schema = self.model.schemas.get(schema_name)
+            if not schema:
+                continue
+
+            schema_info = {"tables": {}}
+
+            for table_name, table in schema.tables.items():
+                # Get columns
+                columns = []
+                for col in table.columns:
+                    if not include_system_columns and col.name in system_columns:
+                        continue
+                    columns.append({
+                        "name": col.name,
+                        "type": str(col.type.typename),
+                        "nullok": col.nullok,
+                        "comment": col.comment or "",
+                    })
+
+                # Get foreign keys
+                foreign_keys = []
+                for fk in table.foreign_keys:
+                    fk_cols = [c.name for c in fk.foreign_key_columns]
+                    ref_cols = [c.name for c in fk.referenced_columns]
+                    foreign_keys.append({
+                        "columns": fk_cols,
+                        "referenced_table": f"{fk.pk_table.schema.name}.{fk.pk_table.name}",
+                        "referenced_columns": ref_cols,
+                    })
+
+                # Get features if this is a domain table
+                features = []
+                if self.is_domain_schema(schema_name):
+                    try:
+                        for f in self.find_features(table):
+                            features.append({
+                                "name": f.feature_name,
+                                "feature_table": f.feature_table.name,
+                            })
+                    except Exception:
+                        pass  # Table may not support features
+
+                table_info = {
+                    "comment": table.comment or "",
+                    "is_vocabulary": self.is_vocabulary(table),
+                    "is_asset": self.is_asset(table),
+                    "is_association": bool(self.is_association(table)),
+                    "columns": columns,
+                    "foreign_keys": foreign_keys,
+                }
+                if features:
+                    table_info["features"] = features
+
+                schema_info["tables"][table_name] = table_info
+
+            result["schemas"][schema_name] = schema_info
+
+        return result
+
     def __getattr__(self, name: str) -> Any:
         # Called only if `name` is not found in Manager.  Delegate attributes to model class.
         return getattr(self.model, name)
@@ -120,20 +289,28 @@ class DerivaModel:
     def name_to_table(self, table: TableInput) -> Table:
         """Return the table object corresponding to the given table name.
 
-        If the table name appears in more than one schema, return the first one you find.
+        Searches domain schemas first (in sorted order), then ML schema, then WWW.
+        If the table name appears in more than one schema, returns the first match.
 
         Args:
           table: A ERMRest table object or a string that is the name of the table.
 
         Returns:
           Table object.
+
+        Raises:
+          DerivaMLException: If the table doesn't exist in any searchable schema.
         """
         if isinstance(table, Table):
             return table
-        if table in (s := self.model.schemas[self.domain_schema].tables):
-            return s[table]
-        for s in [self.model.schemas[sname] for sname in [self.domain_schema, self.ml_schema, "WWW"]]:
-            if table in s.tables.keys():
+
+        # Search domain schemas (sorted for deterministic order), then ML schema, then WWW
+        search_order = [*sorted(self.domain_schemas), self.ml_schema, "WWW"]
+        for sname in search_order:
+            if sname not in self.model.schemas:
+                continue
+            s = self.model.schemas[sname]
+            if table in s.tables:
                 return s.tables[table]
         raise DerivaMLException(f"The table {table} doesn't exist.")
 
@@ -220,21 +397,28 @@ class DerivaModel:
         return [t for s in self.model.schemas.values() for t in s.tables.values() if self.is_asset(t)]
 
     def find_vocabularies(self) -> list[Table]:
-        """Return a list of all the controlled vocabulary tables in the domain schema."""
-        return [t for s in self.model.schemas.values() for t in s.tables.values() if self.is_vocabulary(t)]
+        """Return a list of all controlled vocabulary tables in domain and ML schemas."""
+        tables = []
+        for schema_name in [*self.domain_schemas, self.ml_schema]:
+            schema = self.model.schemas.get(schema_name)
+            if schema:
+                tables.extend(t for t in schema.tables.values() if self.is_vocabulary(t))
+        return tables
 
     @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
-    def find_features(self, table: TableInput) -> Iterable[Feature]:
-        """List the names of the features in the specified table.
+    def find_features(self, table: TableInput | None = None) -> Iterable[Feature]:
+        """List features in the catalog.
+
+        If a table is specified, returns only features for that table.
+        If no table is specified, returns all features across all tables in the catalog.
 
         Args:
-            table: The table to find features for.
-            table: Table | str:
+            table: Optional table to find features for. If None, returns all features
+                in the catalog.
 
         Returns:
-            An iterable of FeatureResult instances that describe the current features in the table.
+            An iterable of Feature instances describing the features.
         """
-        table = self.name_to_table(table)
 
         def is_feature(a: FindAssociationResult) -> bool:
             """Check if association represents a feature.
@@ -250,9 +434,24 @@ class DerivaModel:
                 a.self_fkey.foreign_key_columns[0].name,
             }.issubset({c.name for c in a.table.columns})
 
-        return [
-            Feature(a, self) for a in table.find_associations(min_arity=3, max_arity=3, pure=False) if is_feature(a)
-        ]
+        def find_table_features(t: Table) -> list[Feature]:
+            """Find all features for a single table."""
+            return [
+                Feature(a, self) for a in t.find_associations(min_arity=3, max_arity=3, pure=False) if is_feature(a)
+            ]
+
+        if table is not None:
+            # Find features for a specific table
+            return find_table_features(self.name_to_table(table))
+        else:
+            # Find all features across all domain and ML schema tables
+            features: list[Feature] = []
+            for schema_name in [*self.domain_schemas, self.ml_schema]:
+                schema = self.model.schemas.get(schema_name)
+                if schema:
+                    for t in schema.tables.values():
+                        features.extend(find_table_features(t))
+            return features
 
     def lookup_feature(self, table: TableInput, feature_name: str) -> Feature:
         """Lookup the named feature associated with the provided table.
@@ -290,6 +489,20 @@ class DerivaModel:
         else:
             self.model.apply()
 
+    def is_dataset_rid(self, rid: RID, deleted: bool = False) -> bool:
+        """Check if a given RID is a dataset RID."""
+        try:
+            rid_info = self.model.catalog.resolve_rid(rid, self.model)
+        except KeyError as _e:
+            raise DerivaMLException(f"Invalid RID {rid}")
+        if rid_info.table.name != "Dataset":
+            return False
+        elif deleted:
+            # Got a dataset rid. Now check to see if its deleted or not.
+            return True
+        else:
+            return not list(rid_info.datapath.entities().fetch())[0]["Deleted"]
+
     def list_dataset_element_types(self) -> list[Table]:
         """
         Lists the data types of elements contained within a dataset.
@@ -307,15 +520,14 @@ class DerivaModel:
 
         dataset_table = self.name_to_table("Dataset")
 
-        def domain_table(table: Table) -> bool:
-            return table.schema.name == self.domain_schema or table.name == dataset_table.name
+        def is_domain_or_dataset_table(table: Table) -> bool:
+            return self.is_domain_schema(table.schema.name) or table.name == dataset_table.name
 
-        return [t for a in dataset_table.find_associations() if domain_table(t := a.other_fkeys.pop().pk_table)]
+        return [t for a in dataset_table.find_associations() if is_domain_or_dataset_table(t := a.other_fkeys.pop().pk_table)]
 
-    def _prepare_wide_table(self,
-                            dataset,
-                            dataset_rid: RID,
-                            include_tables: list[str]) -> tuple[dict[str, Any], list[tuple]]:
+    def _prepare_wide_table(
+        self, dataset, dataset_rid: RID, include_tables: list[str]
+    ) -> tuple[dict[str, Any], list[tuple]]:
         """
         Generates details of a wide table from the model
 
@@ -344,11 +556,6 @@ class DerivaModel:
         for p in table_paths:
             paths_by_element[p[2].name].append(p)
 
-        # Get the names of all of the tables that can be dataset elements.
-        dataset_element_tables = {
-            e.name for e in self.list_dataset_element_types() if e.schema.name == self.domain_schema
-        }
-
         skip_columns = {"RCT", "RMT", "RCB", "RMB"}
         element_tables = {}
         for element_table, paths in paths_by_element.items():
@@ -446,9 +653,11 @@ class DerivaModel:
 
         def find_arcs(table: Table) -> set[Table]:
             """Given a path through the model, return the FKs that link the tables"""
+            # Valid schemas for traversal: all domain schemas + ML schema
+            valid_schemas = self.domain_schemas | {self.ml_schema}
             arc_list = [fk.pk_table for fk in table.foreign_keys] + [fk.table for fk in table.referenced_by]
-            arc_list = [t for t in arc_list if t.schema.name in {self.domain_schema, self.ml_schema}]
-            domain_tables = [t for t in arc_list if t.schema.name == self.domain_schema]
+            arc_list = [t for t in arc_list if t.schema.name in valid_schemas]
+            domain_tables = [t for t in arc_list if self.is_domain_schema(t.schema.name)]
             if multiple_columns := [c for c, cnt in Counter(domain_tables).items() if cnt > 1]:
                 raise DerivaMLException(f"Ambiguous relationship in {table.name} {multiple_columns}")
             return set(arc_list)
@@ -466,7 +675,8 @@ class DerivaModel:
             return paths
 
         for child in find_arcs(root):
-            if child.name in {"Dataset_Execution", "Dataset_Dataset", "Execution"}:
+            #        if child.name in {"Dataset_Execution", "Dataset_Dataset", "Execution"}:
+            if child.name in {"Dataset_Dataset", "Execution"}:
                 continue
             if child == parent:
                 # Don't loop back via referred_by
@@ -479,7 +689,23 @@ class DerivaModel:
             paths.extend(self._schema_to_paths(child, path))
         return paths
 
-    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
-    def create_table(self, table_def: TableDefinition) -> Table:
-        """Create a new table from TableDefinition."""
-        return self.model.schemas[self.domain_schema].create_table(table_def.model_dump())
+    def create_table(self, table_def: TableDefinition, schema: str | None = None) -> Table:
+        """Create a new table from TableDefinition.
+
+        Args:
+            table_def: Table definition (dataclass or dict).
+            schema: Schema to create the table in. If None, uses default_schema.
+
+        Returns:
+            The newly created Table.
+
+        Raises:
+            DerivaMLException: If no schema specified and default_schema is not set.
+
+        Note: @validate_call removed because TableDefinition is now a dataclass from
+        deriva.core.typed and Pydantic validation doesn't work well with dataclass fields.
+        """
+        schema = schema or self._require_default_schema()
+        # Handle both TableDefinition (dataclass with to_dict) and plain dicts
+        table_dict = table_def.to_dict() if hasattr(table_def, 'to_dict') else table_def
+        return self.model.schemas[schema].create_table(table_dict)