From e2bcd61ace10e9ab9dfa70f14215bbe01c42551f Mon Sep 17 00:00:00 2001
From: Gregoo <versatile.mailbox@gmail.com>
Date: Fri, 24 Jan 2025 10:03:53 +0100
Subject: [PATCH] Take a more generic approach at counting all the tables in
 the database file

---
 bricktracker/sql.py                | 19 +++++++------
 bricktracker/sql/schema/tables.sql |  1 +
 bricktracker/sql_counter.py        | 43 +++++++++++++++++++++++++++---
 templates/admin/database.html      |  8 +++---
 4 files changed, 53 insertions(+), 18 deletions(-)
 create mode 100644 bricktracker/sql/schema/tables.sql

diff --git a/bricktracker/sql.py b/bricktracker/sql.py
index 77d5466..000e19e 100644
--- a/bricktracker/sql.py
+++ b/bricktracker/sql.py
@@ -20,14 +20,6 @@ G_ENVIRONMENT: Final[str] = 'database_environment'
 G_DEFER: Final[str] = 'database_defer'
 G_STATS: Final[str] = 'database_stats'
 
-COUNTERS: Final[list[BrickCounter]] = [
-    BrickCounter('Sets', 'sets', icon='hashtag'),
-    BrickCounter('Minifigures', 'minifigures', icon='group-line'),
-    BrickCounter('Parts', 'inventory', icon='shapes-line'),
-    BrickCounter('Missing', 'missing', icon='error-warning-line'),
-    BrickCounter('Wishlist', 'wishlist', icon='gift-line'),
-]
-
 
 # SQLite3 client with our extra features
 class BrickSQL(object):
@@ -123,7 +115,12 @@ class BrickSQL(object):
 
     # Count the database records
     def count_records(self) -> list[BrickCounter]:
-        for counter in COUNTERS:
+        counters: list[BrickCounter] = []
+
+        # Get all tables
+        for table in self.fetchall('schema/tables'):
+            counter = BrickCounter(table['name'])
+
             # Failsafe this one
             try:
                 record = self.fetchone('schema/count', table=counter.table)
@@ -133,7 +130,9 @@ class BrickSQL(object):
             except Exception:
                 pass
 
-        return COUNTERS
+            counters.append(counter)
+
+        return counters
 
     # Defer a call to execute
     def defer(self, query: str, parameters: dict[str, Any], /):
diff --git a/bricktracker/sql/schema/tables.sql b/bricktracker/sql/schema/tables.sql
new file mode 100644
index 0000000..44ac37a
--- /dev/null
+++ b/bricktracker/sql/schema/tables.sql
@@ -0,0 +1 @@
+SELECT "name" FROM "sqlite_master" WHERE type='table' ORDER BY "name" ASC
\ No newline at end of file
diff --git a/bricktracker/sql_counter.py b/bricktracker/sql_counter.py
index 44c9733..d104269 100644
--- a/bricktracker/sql_counter.py
+++ b/bricktracker/sql_counter.py
@@ -1,11 +1,46 @@
+from typing import Tuple
+
+# Some table aliases to make it look cleaner (id: (name, icon))
+ALIASES: dict[str, Tuple[str, str]] = {
+    'bricktracker_set_checkboxes': ('Checkboxes', 'checkbox-line'),
+    'bricktracker_set_statuses': ('Bricktracker sets status', 'checkbox-line'),
+    'bricktracker_sets': ('Bricktracker sets', 'hashtag'),
+    'bricktracker_wishes': ('Bricktracker wishes', 'gift-line'),
+    'inventory': ('Parts', 'shapes-line'),
+    'minifigures': ('Minifigures', 'group-line'),
+    'missing': ('Missing', 'error-warning-line'),
+    'rebrickable_sets': ('Rebrickable sets', 'hashtag'),
+    'sets': ('Sets', 'hashtag'),
+    'sets_old': ('Sets (legacy)', 'hashtag'),
+    'wishlist': ('Wishlist', 'gift-line'),
+    'wishlist_old': ('Wishlist (legacy)', 'gift-line'),
+}
+
+
 class BrickCounter(object):
     name: str
     table: str
     icon: str
     count: int
 
-    def __init__(self, name: str, table: str, /, *, icon: str = ''):
-        self.name = name
+    def __init__(
+        self,
+        table: str,
+        /,
+        *,
+        name: str | None = None,
+        icon: str = 'question-line'
+    ):
         self.table = table
-        self.icon = icon
-        self.count = 0
+
+        # Check if there is an alias
+        if table in ALIASES:
+            self.name = ALIASES[table][0]
+            self.icon = ALIASES[table][1]
+        else:
+            if name is None:
+                self.name = table
+            else:
+                self.name = name
+
+            self.icon = icon
diff --git a/templates/admin/database.html b/templates/admin/database.html
index 9e397ef..b52ebab 100644
--- a/templates/admin/database.html
+++ b/templates/admin/database.html
@@ -22,13 +22,13 @@
   <div class="d-flex justify-content-start">
     <ul class="list-group me-2">
       {% for counter in database_counters %}
-      {% if not (loop.index % 5) %}
-        </ul>
-        <ul class="list-group">
-      {% endif %}
         <li class="list-group-item d-flex justify-content-between align-items-start">
           <span><i class="ri-{{ counter.icon }}"></i> {{ counter.name }}</span> <span class="badge text-bg-primary rounded-pill ms-2">{{ counter.count }}</span>
         </li>
+        {% if not (loop.index % 5) %}
+          </ul>
+          <ul class="list-group me-2">
+        {% endif %}
       {% endfor %}
     </ul>
   </div>