tol-sdk 1.6.37__py3-none-any.whl → 1.7.0__py3-none-any.whl

tol/core/requested_fields.py

@@ -0,0 +1,189 @@
+# SPDX-FileCopyrightText: 2025 Genome Research Ltd.
+#
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+
+from . import DataSourceError, OperableDataSource
+from .operator import Relational
+from .relationship import RelationshipConfig
+
+
+class ReqFieldsTree:
+    """
+    Acts as a template for which related objects and attributes to fetch from
+    the DataSource, and for serialzing them in the response.
+    """
+
+    def __init__(
+        self,
+        object_type: str,
+        data_source: OperableDataSource,
+        requested_fields: list[str] | None = None,
+        include_all_to_ones: bool = True,
+    ) -> None:
+        self.__object_type: str = object_type
+        self.__attributes: dict[str, bool] = {}
+        self.__sub_trees: dict[str, ReqFieldsTree] = {}
+        self.__rel_conf: RelationshipConfig | None = (
+            data_source.relationship_config.get(object_type)
+            if isinstance(data_source, Relational)
+            else None
+        )
+        if requested_fields:
+            self.add_requested_tree(data_source, requested_fields)
+        elif include_all_to_ones:
+            self.add_all_to_ones(data_source)
+
+    def __eq__(self, other):
+        if isinstance(other, self.__class__):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def object_type(self) -> str:
+        return self.__object_type
+
+    @property
+    def attribute_names(self) -> list[str]:
+        return list(self.__attributes)
+
+    def add_attribute(self, name: str) -> None:
+        """
+        Add attribute `name`. Attributes are implemented as a dict rather than
+        a set to preserve order.
+        """
+        self.__attributes[name] = True
+
+    def has_attribute(self, name: str) -> bool:
+        return name in self.__attributes
+
+    def add_sub_tree(self, name: str, sub: ReqFieldsTree) -> None:
+        self.__sub_trees[name] = sub
+
+    def get_sub_tree(self, name: str) -> ReqFieldsTree | None:
+        return self.__sub_trees.get(name)
+
+    def sub_trees(self) -> Iterator[str, ReqFieldsTree]:
+        """
+        Iterator over `name, sub_tree`.
+        """
+        yield from self.__sub_trees.items()
+
+    @property
+    def is_leaf(self) -> bool:
+        """
+        This tree is a leaf if it has no sub-trees.
+        """
+        return not self.__sub_trees
+
+    @property
+    def is_stub(self) -> bool:
+        """
+        This tree is a stub if it has no sub-trees and has one attribute which
+        is "id".
+        """
+        if self.__sub_trees:
+            return False
+        return len(self.__attributes) == 1 and 'id' in self.__attributes
+
+    @property
+    def has_relationships(self) -> bool:
+        return self.__rel_conf is not None
+
+    def to_one_names(self) -> Iterable[str]:
+        return x.keys() if (x := self.__rel_conf.to_one) else ()
+
+    def to_many_names(self) -> Iterable[str]:
+        return x.keys() if (x := self.__rel_conf.to_many) else ()
+
+    def get_relationship(self, name: str) -> str:
+        """
+        Fetches the related object type from the `to_one` or `to_many` fields
+        of the attached `RelationshipConfig`.
+        """
+        if rel_conf := self.__rel_conf:
+            for attr in 'to_one', 'to_many':
+                if (cfg := getattr(rel_conf, attr)) and (type_name := cfg.get(name)):
+                    return type_name
+
+    def get_attribute_type(self, data_source, name) -> type | None:
+        return data_source.attribute_types[self.object_type].get(name)
+
+    def add_all_to_ones(self, data_source: OperableDataSource) -> None:
+        if not self.__rel_conf:
+            return
+        for name in self.to_one_names():
+            sub_type = self.get_relationship(name)
+            self.add_sub_tree(
+                name,
+                self.__class__(
+                    sub_type,
+                    data_source,
+                    include_all_to_ones=False,
+                ),
+            )
+
+    def add_requested_tree(
+        self,
+        data_source: OperableDataSource,
+        requested_fields: list[str],
+    ) -> None:
+        err_title = 'Bad Requested Fields Path Element'
+        for path_str in requested_fields:
+            tree = self
+            for name in path_str.split('.'):
+                if name == '':
+                    msg = f"Empty element in path path '{path_str}'"
+                    raise DataSourceError(title=err_title, detail=msg)
+                elif not tree:
+                    msg = f"Element '{name}' appears after an attribute name in path '{path_str}'"
+                    raise DataSourceError(title=err_title, detail=msg)
+                elif sub_tree := tree.get_sub_tree(name):
+                    # `name` is a relationship we already have so move pointer
+                    # to the sub tree.
+                    tree = sub_tree
+                elif tree.has_attribute(name):
+                    # Already have attribute `name`. Unset `tree` to trap path
+                    # elements follwing an attribute.
+                    tree = None
+                elif sub_type := tree.get_relationship(name):
+                    # New sub-tree from relationship name.
+                    sub_tree = self.__class__(
+                        sub_type,
+                        data_source,
+                        include_all_to_ones=False,
+                    )
+                    tree.add_sub_tree(name, sub_tree)
+                    tree = sub_tree
+                elif tree.get_attribute_type(data_source, name):
+                    tree.add_attribute(name)
+                    # Unset `tree` to trap path elements follwing an
+                    # attribute.
+                    tree = None
+                else:
+                    msg = (
+                        f'{name!r} in path {path_str!r} is not a known relationship'
+                        f' or attribute of {tree.object_type!r} objects'
+                    )
+                    raise DataSourceError(title=err_title, detail=msg)
+
+    def __str__(self):
+        return ','.join(self.__to_strings_iter())
+
+    def to_paths(self) -> list[str]:
+        return list(self.__to_strings_iter())
+
+    def __to_strings_iter(self, *path: list[str]) -> Iterator[str]:
+        for name in self.__attributes:
+            yield '.'.join([*path, name])
+        if self.is_leaf:
+            # We only need to return the path to self if it hasn't already
+            # appeared as a root path under an attribtue.
+            if path and not self.__attributes:
+                yield '.'.join(path)
+        else:
+            for name, tree in self.sub_trees():
+                yield from tree.__to_strings_iter(*path, name)

tol/elastic/elastic_datasource.py

@@ -215,7 +215,8 @@ class ElasticDataSource(
         page_size: Optional[int] = None,
         object_filters: Optional[DataSourceFilter] = None,
         search_after: list[str] | None = None,
-        session: Optional[OperableSession] = None
+        session: Optional[OperableSession] = None,
+        **kwargs,
     ) -> tuple[Iterable[DataObject], list[str] | None]:
 
         resp = self.__get_page_response(

tol/flows/converters/benchling_extraction_to_elastic_extraction_converter.py

@@ -18,13 +18,32 @@ class BenchlingExtractionToElasticExtractionConverter(
                 'extraction',
                 data_object.id,
                 attributes={
-                    'sample': {'id': data_object.sts_id},
-                    'species': {'id': data_object.taxon_id},
-                    'specimen': {'id': data_object.specimen_id},
-                    'tolid': {'id': data_object.programme_id},
-                    'tissue_prep': {'id': data_object.eln_tissue_prep_id},
                     **{k: v
                        for k, v in data_object.attributes.items()
                        if k not in ['sts_id', 'specimen_id', 'taxon_id',
-                                    'programme_id', 'eln_tissue_prep_id']}})
+                                    'programme_id', 'eln_tissue_prep_id']}
+                },
+                to_one={
+                    'sample': self._data_object_factory(
+                        'sample',
+                        data_object.sts_id
+                    ) if data_object.sts_id is not None else None,
+                    'species': self._data_object_factory(
+                        'species',
+                        data_object.taxon_id
+                    ) if data_object.taxon_id is not None else None,
+                    'specimen': self._data_object_factory(
+                        'specimen',
+                        data_object.specimen_id
+                    ) if data_object.specimen_id is not None else None,
+                    'tolid': self._data_object_factory(
+                        'tolid',
+                        data_object.programme_id
+                    ) if data_object.programme_id is not None else None,
+                    'tissue_prep': self._data_object_factory(
+                        'tissue_prep',
+                        data_object.eln_tissue_prep_id
+                    ) if data_object.eln_tissue_prep_id is not None else None,
+                }
+            )
             yield ret

tol/flows/converters/benchling_extraction_to_elastic_sequencing_request_converter.py

@@ -18,12 +18,33 @@ class BenchlingExtractionToElasticSequencingRequestConverter(
                 'sequencing_request',
                 data_object.id,
                 attributes={
-                    'extraction': {'id': data_object.id},
-                    'sample': {'id': data_object.sts_id},
-                    'species': {'id': data_object.taxon_id},
-                    'specimen': {'id': data_object.specimen_id},
-                    'tolid': {'id': data_object.programme_id},
-                    'tissue_prep': {'id': data_object.eln_tissue_prep_id},
                     'sequencing_platform': 'pacbio'
-                })
+                },
+                to_one={
+                    'extraction': self._data_object_factory(
+                        'extraction',
+                        data_object.id
+                    ) if data_object.id is not None else None,
+                    'sample': self._data_object_factory(
+                        'sample',
+                        data_object.sts_id
+                    ) if data_object.sts_id is not None else None,
+                    'species': self._data_object_factory(
+                        'species',
+                        data_object.taxon_id
+                    ) if data_object.taxon_id is not None else None,
+                    'specimen': self._data_object_factory(
+                        'specimen',
+                        data_object.specimen_id
+                    ) if data_object.specimen_id is not None else None,
+                    'tolid': self._data_object_factory(
+                        'tolid',
+                        data_object.programme_id
+                    ) if data_object.programme_id is not None else None,
+                    'tissue_prep': self._data_object_factory(
+                        'tissue_prep',
+                        data_object.eln_tissue_prep_id
+                    ) if data_object.eln_tissue_prep_id is not None else None,
+                }
+            )
             yield ret

tol/flows/converters/benchling_sequencing_request_to_elastic_sequencing_request_converter.py

@@ -17,22 +17,43 @@ class BenchlingSequencingRequestToElasticSequencingRequestConverter(
             extraction = None
             tissue_prep = None
             if 'extraction_id' in data_object.attributes:
-                extraction = {'id': data_object.extraction_id}
+                extraction = self._data_object_factory(
+                    'extraction',
+                    data_object.extraction_id
+                )
             if 'tissue_prep_id' in data_object.attributes:
-                tissue_prep = {'id': data_object.tissue_prep_id}
+                tissue_prep = self._data_object_factory(
+                    'tissue_prep',
+                    data_object.tissue_prep_id
+                )
             ret = self._data_object_factory(
                 'sequencing_request',
                 data_object.sanger_sample_id,
                 attributes={
-                    'sample': {'id': str(data_object.sts_id)},
-                    'specimen': {'id': str(data_object.specimen_id)},
-                    'species': {'id': str(data_object.taxon_id)},
-                    'tolid': {'id': data_object.programme_id},
-                    'extraction': extraction,
-                    'tissue_prep': tissue_prep,
                     **{k: v
                        for k, v in data_object.attributes.items()
                        if k not in ['sanger_sample_id', 'sts_id',
                                     'specimen_id', 'taxon_id', 'extraction_id',
-                                    'programme_id', 'tissue_prep_id']}})
+                                    'programme_id', 'tissue_prep_id']}
+                },
+                to_one={
+                    'sample': self._data_object_factory(
+                        'sample',
+                        data_object.sts_id
+                    ) if data_object.sts_id is not None else None,
+                    'specimen': self._data_object_factory(
+                        'specimen',
+                        data_object.specimen_id
+                    ) if data_object.specimen_id is not None else None,
+                    'species': self._data_object_factory(
+                        'species',
+                        data_object.taxon_id
+                    ) if data_object.taxon_id is not None else None,
+                    'tolid': self._data_object_factory(
+                        'tolid',
+                        data_object.programme_id
+                    ) if data_object.programme_id is not None else None,
+                    'extraction': extraction,
+                    'tissue_prep': tissue_prep,
+                })
             yield ret

tol/flows/converters/benchling_tissue_prep_to_elastic_tissue_prep_converter.py

@@ -17,15 +17,26 @@ class BenchlingTissuePrepToElasticTissuePrepConverter(
             'tissue_prep',
             data_object.eln_tissue_prep_id,
             attributes={
-                'sample': {'id': str(data_object.sts_id)},
-                'species': {'id': str(data_object.taxon_id)},
-                'tolid': {'id': data_object.programme_id},
                 **{k: v
                    for k, v in data_object.attributes.items()
                    if k not in ['eln_tissue_prep_id',
                                 'sts_id',
                                 'taxon_id',
                                 'programme_id']}
+            },
+            to_one={
+                'sample': self._data_object_factory(
+                    'sample',
+                    data_object.sts_id
+                ) if data_object.sts_id is not None else None,
+                'species': self._data_object_factory(
+                    'species',
+                    data_object.taxon_id
+                ) if data_object.taxon_id is not None else None,
+                'tolid': self._data_object_factory(
+                    'tolid',
+                    data_object.programme_id
+                ) if data_object.programme_id is not None else None,
             }
         )
         return iter([ret])

tol/flows/converters/elastic_sample_to_benchling_tissue_update_converter.py

@@ -34,7 +34,7 @@ class ElasticSampleToBenchlingTissueUpdateConverter(
                         if species.sts_taxon_group else 'NA',
                     'genome_size': str(species.sts_genome_size),
                     # 'freezer': None,
-                    'location': data_object.sts_labwhere_parentage,  # Previously shelf
+                    'location': data_object.sts_labwhere_parentage,
                     'tray': data_object.sts_labwhere_name,
                     'specimen_id': specimen.id,
                     'programme_id': data_object.sts_tolid.id,

tol/flows/converters/elastic_sample_to_elastic_sequencing_request_update_converter.py

@@ -19,7 +19,10 @@ class ElasticSampleToElasticSequencingRequestUpdateConverter(
         yield (
             None,
             {
-                'mlwh_sample': {'id': data_object.id},
+                'mlwh_sample': self._data_object_factory(
+                    'sample',
+                    data_object.id
+                ),
                 'mlwh_specimen.id': specimen.id
             }
         )

tol/flows/converters/elastic_tolid_to_elastic_genome_note_update_converter.py

@@ -19,7 +19,10 @@ class ElasticTolidToElasticGenomeNoteUpdateConverter(
         yield (
             None,
             {
-                'gn_species': {'id': species.id},
+                'gn_species': self._data_object_factory(
+                    'species',
+                    species.id
+                ),
                 'gn_tolid.id': data_object.id
             }
         )

tol/flows/converters/elastic_tolid_to_elastic_sample_update_converter.py

@@ -21,7 +21,10 @@ class ElasticTolidToElasticSampleUpdateConverter(
             yield (
                 None,
                 {
-                    'tolid_tolid': {'id': data_object.id},
+                    'tolid_tolid': self._data_object_factory(
+                        'tolid',
+                        data_object.id
+                    ),
                     'sts_species.id':
                         data_object.requested_taxonomy_id
                         if data_object.requested_taxonomy_id is not None

tol/sources/sts.py

@@ -14,13 +14,17 @@ from ..core import (
 )
 
 
-def sts(retries: int = 5, **kwargs) -> ApiDataSource:
+def sts(
+    retries: int = 5,
+    status_forcelist: list[int] | None = [429, 500, 502, 503, 504]
+) -> ApiDataSource:
     sts = create_api_datasource(
         api_url=os.getenv('STS_URL', Defaults.STS_URL)
         + os.getenv('STS_API_PATH', Defaults.STS_API_PATH),
         token=os.getenv('STS_API_KEY'),
         data_prefix=os.getenv('STS_API_DATA_PATH', Defaults.STS_API_DATA_PATH),
-        retries=retries
+        retries=retries,
+        status_forcelist=status_forcelist,
     )
     core_data_object(sts)
     return sts

tol/sql/database.py

@@ -5,18 +5,22 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Iterable, List, Optional, Type
+from collections.abc import Iterable
+from typing import Any, Dict, List, Optional, Type
 
 from sqlalchemy import distinct, func
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.orm import MappedColumn, Query, Session, joinedload
+from sqlalchemy.orm import Load, MappedColumn, Query, Session, joinedload, load_only, raiseload
 from sqlalchemy.orm.attributes import flag_modified
 
 from .filter import DatabaseFilter
 from .model import Model
 from .session import SessionFactory
 from .sort import DatabaseSorter
-from ..core import DataSourceError
+from ..core import DataSourceError, ReqFieldsTree
+
+
+SubPath = tuple[str | None, ReqFieldsTree]
 
 
 class Database(ABC):
@@ -28,7 +32,7 @@ class Database(ABC):
         tablename: str,
         instance_id: Any,
         in_session: Session,
-        requested_relationships: dict[str, str] | None = None,
+        requested_tree: ReqFieldsTree | None = None,
     ) -> Optional[Model]:
         """
         Gets a single instance by its instance-ID, or None if not found.
@@ -46,7 +50,7 @@ class Database(ABC):
         sort_by: Optional[DatabaseSorter] = None,
         offset: Optional[int] = None,
         limit: Optional[int] = None,
-        requested_relationships: dict[str, str] | None = None
+        requested_tree: ReqFieldsTree | None = None,
     ) -> Iterable[Model]:
         """
         Returns an Iterable of `Model` instances according
@@ -191,14 +195,13 @@ class DefaultDatabase(Database):
         tablename: str,
         instance_id: Any,
         in_session: Session,
-        requested_relationships: dict[str, str] | None = None,
+        requested_tree: ReqFieldsTree | None = None,
     ) -> Optional[Model]:
-
         result = self.__get_instance_by_id(
             tablename,
             instance_id,
             in_session,
-            requested_relationships,
+            requested_tree,
         )
         return result
 
@@ -210,13 +213,13 @@ class DefaultDatabase(Database):
         sort_by: Optional[DatabaseSorter] = None,
         offset: Optional[int] = None,
         limit: Optional[int] = None,
-        requested_relationships: dict[str, str] | None = None,
+        requested_tree: ReqFieldsTree | None = None,
     ) -> Iterable[Model]:
 
         _, query = self.__get_model_query(
             tablename,
             in_session,
-            requested_relationships,
+            requested_tree=requested_tree,
             filters=filters,
         )
         if filters is not None:
@@ -225,7 +228,8 @@ class DefaultDatabase(Database):
             query = filters.filter(query, tablename, self.__tablename_model_dict)
         if sort_by is not None:
             query = sort_by.sort(query, tablename, self.__tablename_model_dict, filters)
-        query = query.limit(limit).offset(offset)
+        if limit is not None and offset is not None:
+            query = query.limit(limit).offset(offset)
         results = query.all()
         return results
 
@@ -236,7 +240,7 @@ class DefaultDatabase(Database):
         filters: Optional[DatabaseFilter] = None
     ) -> int:
 
-        _, query = self.__get_model_query(tablename, in_session, None, filters=filters)
+        _, query = self.__get_model_query(tablename, in_session, filters=filters)
         if filters is not None:
             query = filters.filter(query, tablename, self.__tablename_model_dict)
         count = query.count()
@@ -346,7 +350,6 @@ class DefaultDatabase(Database):
         model, query = self.__get_model_query(
             tablename,
             in_session,
-            None,
             filters=filters,
         )
 
@@ -605,43 +608,18 @@ class DefaultDatabase(Database):
         self,
         tablename: str,
         in_session: Session,
-        requested_relationships: dict | None,
+        requested_tree: ReqFieldsTree | None = None,
         filters: DatabaseFilter | None = None,
     ) -> tuple[Type[Model], Query]:
 
         model = self.__tablename_model_dict[tablename]
-        query = in_session.query(model) if not filters else filters.get_query(in_session, model)
+        query = filters.get_query(in_session, model) if filters else in_session.query(model)
 
-        if requested_relationships:
-            query = self.__apply_requested_relationships(
-                query,
-                requested_relationships
-            )
+        if requested_tree:
+            query = self.add_options_to_query(query, tablename, requested_tree)
 
         return model, query
 
-    def __apply_requested_relationships(
-        self,
-        query: Query,
-        requested_relationships: dict[str, str]
-    ) -> Query:
-
-        tablename = requested_relationships.pop('__tablename__')
-        if not requested_relationships:
-            return query
-
-        model = self.__tablename_model_dict[tablename]
-
-        for r_name, r_dict in requested_relationships.items():
-            relationship = getattr(model, r_name)
-            query.options(
-                joinedload(relationship)
-            )
-
-            query = self.__apply_requested_relationships(query, r_dict)
-
-        return query
-
     def __commit_session(
         self,
         in_session: Session,
@@ -677,7 +655,7 @@ class DefaultDatabase(Database):
         tablename: str,
         instance_id: str,
         in_session: Session,
-        requested_relationships: dict[str, str] | None = None,
+        requested_tree: ReqFieldsTree | None = None,
     ) -> Optional[Model]:
         """
         Gets an instance by its tablename and id.
@@ -686,7 +664,7 @@ class DefaultDatabase(Database):
         model, query = self.__get_model_query(
             tablename,
             in_session,
-            requested_relationships
+            requested_tree=requested_tree,
         )
         id_column = getattr(model, model.get_id_column_name())
         result = query.filter(id_column == instance_id).one_or_none()
@@ -854,3 +832,61 @@ class DefaultDatabase(Database):
                 f'Hint - check the following tables: "{relationship_names}".'
             )
         )
+
+    def add_options_to_query(
+        self,
+        query: Query,
+        tablename: str,
+        requested_tree: ReqFieldsTree,
+    ):
+        options = self.joinedload_options(requested_tree)
+        # `raiseload(*)` acts as a trap, raising an exception if any methods
+        # on the returned objects are called which would trigger loading data
+        # from the database via another SELECT.
+        return query.options(options, raiseload('*')) if options else query
+
+    def joinedload_options(self, tree: ReqFieldsTree) -> list[Load]:
+        """
+        Returns a list of SQLAlchemy `Load` objects based on the supplied
+        `ReqFieldsTree` which specify which related tables to join into and
+        which attribute columns to select.  This list of `Load` objects can
+        then be added to a SQLAlchemy `Query` via a call to `options()`.
+        """
+        sub: SubPath = None, tree
+        return list(self.__joinedload_iter(sub))
+
+    def __joinedload_iter(self, sub: SubPath, *path: list[SubPath]):
+        path = [*path, sub]
+        tree = sub[1]
+        if tree.is_leaf:
+            if options := self.__joinedload_options_from_path(path):
+                yield options
+        else:
+            for sub in tree.sub_trees():
+                yield from self.__joinedload_iter(sub, *path)
+
+    def __joinedload_options_from_path(self, path: list[SubPath]):
+        load = None
+        prev_model = None
+        for rel_name, tree in path:
+            model = self.__tablename_model_dict[tree.object_type]
+            if prev_model:
+                # Add a joinedload for this model if `tree` isn't the root
+                relation = getattr(prev_model, rel_name)
+                load = load.joinedload(relation) if load else joinedload(relation)
+            if names := tree.attribute_names:
+                for col_name in model.get_all_foreign_key_names():
+                    # Always load any to-one ID columns where the relation
+                    # isn't being fetched so that we can create stub objects
+                    # for them.
+                    if not tree.has_attribute(col_name):
+                        names.append(col_name)
+                cols = [getattr(model, x) for x in names if x != 'id']
+                load = (
+                    load.load_only(*cols, raiseload=True)
+                    if load
+                    else load_only(*cols, raiseload=True)
+                )
+            prev_model = model
+
+        return load

tol/sql/factory.py

@@ -31,10 +31,10 @@ def __model_converter_factory(
     type_function: TypeFunction
 ) -> ConverterFactory:
 
-    return lambda do_factory, requested_fields: DefaultModelConverter(
+    return lambda do_factory, req_fields_tree: DefaultModelConverter(
         type_function,
         do_factory,
-        requested_fields=requested_fields
+        requested_tree=req_fields_tree
     )