From 6324bc77a8796d0a721ffc4302852512d0d56911 Mon Sep 17 00:00:00 2001
From: nadzyah <nadzya.info@gmail.com>
Date: Thu, 10 Aug 2023 14:05:05 +0300
Subject: [PATCH] Add artefact validator

---
 backend/test_observer/data_access/models.py   | 23 +++++++++++-
 .../test_observer/data_access/repository.py   |  1 -
 backend/test_observer/data_access/setup.py    |  8 ++++-
 .../test_observer/data_access/validators.py   | 36 +++++++++++++++++++
 4 files changed, 65 insertions(+), 3 deletions(-)
 create mode 100644 backend/test_observer/data_access/validators.py

diff --git a/backend/test_observer/data_access/models.py b/backend/test_observer/data_access/models.py
index 27c5caec..56c58bdb 100644
--- a/backend/test_observer/data_access/models.py
+++ b/backend/test_observer/data_access/models.py
@@ -28,9 +28,14 @@
     Mapped,
     mapped_column,
     relationship,
+    validates,
 )
 
-from test_observer.data_access.models_enums import ArtefactStatus, TestExecutionStatus
+from test_observer.data_access.models_enums import (
+    ArtefactStatus,
+    TestExecutionStatus,
+    FamilyName,
+)
 
 
 class Base(DeclarativeBase):
@@ -92,6 +97,22 @@ class Artefact(Base):
         UniqueConstraint("name", "version", "source", name="unique_artefact"),
     )
 
+    @validates("source")
+    def validate_source(
+        self, key: str, value: dict, family_name: str | None = None  # noqa: ARG002
+    ) -> dict:
+        """Validate source field"""
+        if (
+            self.stage and self.stage.family.name == FamilyName.SNAP.value
+        ) or family_name == FamilyName.SNAP.value:
+            # Check that store key is specified
+            if "store" not in value:
+                raise ValueError("Snap artefacts should have store key in source")
+            # Check that store key has a correct value
+            if not isinstance(value["store"], str):
+                raise ValueError("Store key in source field should be a string")
+        return value
+
 
 class ArtefactBuild(Base):
     """A model to represent specific builds of artefact (e.g. arm64 revision 2)"""
diff --git a/backend/test_observer/data_access/repository.py b/backend/test_observer/data_access/repository.py
index 688909e4..8e258c5c 100644
--- a/backend/test_observer/data_access/repository.py
+++ b/backend/test_observer/data_access/repository.py
@@ -21,7 +21,6 @@
 
 
 from sqlalchemy import and_, func
-from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.orm import Session, joinedload
 from sqlalchemy.exc import IntegrityError
 
diff --git a/backend/test_observer/data_access/setup.py b/backend/test_observer/data_access/setup.py
index ae5373b9..bf7256f3 100644
--- a/backend/test_observer/data_access/setup.py
+++ b/backend/test_observer/data_access/setup.py
@@ -1,7 +1,10 @@
 from os import environ
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, event
 from sqlalchemy.orm import sessionmaker
 
+from .models import Artefact
+from .validators import validate_artefact
+
 DEFAULT_DB_URL = "postgresql+pg8000://postgres:password@test-observer-db:5432/postgres"
 DB_URL = environ.get("DB_URL", DEFAULT_DB_URL)
 
@@ -9,6 +12,9 @@
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
 
+event.listen(Artefact, "before_insert", validate_artefact)
+
+
 # Dependency
 def get_db():
     db = SessionLocal()
diff --git a/backend/test_observer/data_access/validators.py b/backend/test_observer/data_access/validators.py
new file mode 100644
index 00000000..7e7d24b4
--- /dev/null
+++ b/backend/test_observer/data_access/validators.py
@@ -0,0 +1,36 @@
+# Copyright 2023 Canonical Ltd.
+# All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+# Written by:
+#        Nadzeya Hutsko <nadzeya.hutsko@canonical.com>
+"""Functions to validate the whole models (not just specific field)"""
+
+from sqlalchemy import select
+from sqlalchemy.engine.base import Connection
+from sqlalchemy.orm import Mapper
+from .models import Stage, Family, Artefact
+
+
+def validate_artefact(
+    mapper: Mapper, connection: Connection, target: Artefact  # noqa: ARG001
+) -> None:
+    """Validate artefact before insertion"""
+    # Validate source
+    stage = select(Stage.family_id).where(Stage.id == target.stage_id)
+    family_id = connection.execute(stage).scalar()
+    stage_family = select(Family.name).where(Family.id == family_id)
+    family_name = connection.execute(stage_family).scalar()
+    target.source = target.validate_source("source", target.source, family_name)