From 94ddfedbb509eeccf899b8bce778a326ad87ef68 Mon Sep 17 00:00:00 2001
From: mreid-tt <943378+mreid-tt@users.noreply.github.com>
Date: Mon, 15 Jan 2024 14:31:41 -0400
Subject: [PATCH] Amend major_version check

- Add type column to Firmware table
- Increase length of version column
- Filter by type for closest firmware when getting catalog
---
 ...add_firmware_type_and_increase_version_.py | 47 +++++++++++++++++++
 spkrepo/models.py                             |  3 +-
 spkrepo/views/nas.py                          |  2 +-
 3 files changed, 50 insertions(+), 2 deletions(-)
 create mode 100644 migrations/versions/f95855ce9471_add_firmware_type_and_increase_version_.py

diff --git a/migrations/versions/f95855ce9471_add_firmware_type_and_increase_version_.py b/migrations/versions/f95855ce9471_add_firmware_type_and_increase_version_.py
new file mode 100644
index 0000000..564122c
--- /dev/null
+++ b/migrations/versions/f95855ce9471_add_firmware_type_and_increase_version_.py
@@ -0,0 +1,47 @@
+"""Add firmware type and increase version length
+
+Revision ID: f95855ce9471
+Revises: 76d559b4e873
+Create Date: 2024-01-15 13:58:34.160242
+
+"""
+revision = "f95855ce9471"
+down_revision = "76d559b4e873"
+
+import sqlalchemy as sa
+from alembic import op
+
+
+def upgrade():
+    op.add_column("firmware", sa.Column("type", sa.Unicode(length=4)))
+    # Set type based on version
+    op.execute(
+        """
+        UPDATE firmware
+        SET type = CASE
+            WHEN version LIKE '1.%' THEN 'srm'
+            ELSE 'dsm'
+        END
+    """
+    )
+    # Modify the column to be NOT NULL after setting the values
+    op.alter_column("firmware", "type", nullable=False)
+
+    op.alter_column(
+        "firmware",
+        "version",
+        existing_type=sa.VARCHAR(length=3),
+        type_=sa.Unicode(length=4),
+        existing_nullable=False,
+    )
+
+
+def downgrade():
+    op.alter_column(
+        "firmware",
+        "version",
+        existing_type=sa.Unicode(length=4),
+        type_=sa.VARCHAR(length=3),
+        existing_nullable=False,
+    )
+    op.drop_column("firmware", "type")
diff --git a/spkrepo/models.py b/spkrepo/models.py
index 0d0c294..2cc0613 100644
--- a/spkrepo/models.py
+++ b/spkrepo/models.py
@@ -134,8 +134,9 @@ class Firmware(db.Model):
 
     # Columns
     id = db.Column(db.Integer, primary_key=True)
-    version = db.Column(db.Unicode(3), nullable=False)
+    version = db.Column(db.Unicode(4), nullable=False)
     build = db.Column(db.Integer, unique=True, nullable=False)
+    type = db.Column(db.Unicode(4), nullable=False)
 
     @classmethod
     def find(cls, build):
diff --git a/spkrepo/views/nas.py b/spkrepo/views/nas.py
index f7d9cf0..9cbeeed 100644
--- a/spkrepo/views/nas.py
+++ b/spkrepo/views/nas.py
@@ -46,7 +46,7 @@ def is_valid_language(language):
 def get_catalog(arch, build, language, beta):
     # Find the closest matching firmware for the provided build
     closest_firmware = (
-        Firmware.query.filter(Firmware.build <= build)
+        Firmware.query.filter(Firmware.build <= build, Firmware.type == "dsm")
         .order_by(Firmware.build.desc())
         .first()
     )