1 files changed, 51 insertions, 3 deletions
diff --git a/uploader/request_checks.py b/uploader/request_checks.py
index f1d8027..84935f9 100644
--- a/uploader/request_checks.py
+++ b/uploader/request_checks.py
@@ -2,14 +2,20 @@
 
 These are useful for reusability, and hence maintainability of the code.
 """
+import logging
+
+from typing import Callable
 from functools import wraps
 
-from gn_libs.mysqldb import database_connection
+from gn_libs.mysqldb import Connection, database_connection
 from flask import flash, url_for, redirect, current_app as app
 
 from uploader.species.models import species_by_id
 from uploader.population.models import population_by_species_and_id
 
+logger = logging.getLogger(__name__)
+
+
 def with_species(redirect_uri: str):
     """Ensure the species actually exists."""
     def __decorator__(function):
@@ -28,7 +34,7 @@ def with_species(redirect_uri: str):
                               "alert-danger")
                         return redirect(url_for(redirect_uri))
             except ValueError as _verr:
-                app.logger.debug(
+                logger.debug(
                     "Exception converting value to integer: %s",
                     kwargs.get("species_id"),
                     exc_info=True)
@@ -63,7 +69,7 @@ def with_population(species_redirect_uri: str, redirect_uri: str):
                               "alert-danger")
                         return select_population_uri
             except ValueError as _verr:
-                app.logger.debug(
+                logger.debug(
                     "Exception converting value to integer: %s",
                     kwargs.get("population_id"),
                     exc_info=True)
@@ -73,3 +79,45 @@ def with_population(species_redirect_uri: str, redirect_uri: str):
             return function(**{**kwargs, "population": population})
         return __with_population__
     return __decorator__
+
+
+def with_dataset(
+        species_redirect_uri: str,
+        population_redirect_uri: str,
+        redirect_uri: str,
+        dataset_by_id: Callable[
+            [Connection, int, int, int],
+            dict]
+):
+    """Ensure the dataset actually exists."""
+    def __decorator__(func):
+        @wraps(func)
+        @with_population(species_redirect_uri, population_redirect_uri)
+        def __with_dataset__(**kwargs):
+            try:
+                _spcid = int(kwargs["species_id"])
+                _popid = int(kwargs["population_id"])
+                _dsetid = int(kwargs.get("dataset_id"))
+                select_dataset_uri = redirect(url_for(
+                    redirect_uri, species_id=_spcid, population_id=_popid))
+                if not bool(_dsetid):
+                    flash("You need to select a valid 'dataset_id' value.",
+                          "alert-danger")
+                    return select_dataset_uri
+                with database_connection(app.config["SQL_URI"]) as conn:
+                    dataset = dataset_by_id(conn, _spcid, _popid, _dsetid)
+                    if not bool(dataset):
+                        flash("You must select a valid dataset.",
+                              "alert-danger")
+                        return select_dataset_uri
+            except ValueError as _verr:
+                logger.debug(
+                    "Exception converting 'dataset_id' to integer: %s",
+                    kwargs.get("dataset_id"),
+                    exc_info=True)
+                flash("Expected 'dataset_id' value to be an integer."
+                      "alert-danger")
+                return select_dataset_uri
+            return func(**{**kwargs, "dataset": dataset})
+        return __with_dataset__
+    return __decorator__