From b56c5314f8882923907ac56c295e8579d20ea8c9 Mon Sep 17 00:00:00 2001
From: omicron <omicron.me@protonmail.com>
Date: Sat, 28 Oct 2023 04:03:51 +0200
Subject: [PATCH] make d2warehouse.db ready for the flask webui

 - Add a function to automatically initialize the db with the schema if
   it isn't yet initialized
 - Remove the DROP TABLE queries from the schema in case the above^ goes
   wrong somehow
 - Make Item.write_to_db and Item.load_from_db take a db argument so
   that multiple threads can be supported

This also changes the schema to never drop existing tables, just in case
something goes wrong in the logic.
---
 d2warehouse/db.py      |  8 +++++
 d2warehouse/item.py    | 66 ++++++++++++++++++++++++++++++++++++++----
 d2warehouse/schema.sql |  8 ++---
 3 files changed, 73 insertions(+), 9 deletions(-)

diff --git a/d2warehouse/db.py b/d2warehouse/db.py
index 9708622..437fec8 100644
--- a/d2warehouse/db.py
+++ b/d2warehouse/db.py
@@ -1,6 +1,7 @@
 import atexit
 import os
 import sqlite3
+from flask import g
 
 
 _schema_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "schema.sql")
@@ -32,6 +33,13 @@ def close_db() -> None:
         _db = None
 
 
+def init_db() -> None:
+    db = get_db()
+    count = db.execute("SELECT count(*) FROM sqlite_master").fetchone()[0]
+    if count == 0:
+        create_db()
+
+
 def create_db() -> None:
     db = get_db()
     with open(_schema_path, encoding="utf-8") as f:
diff --git a/d2warehouse/item.py b/d2warehouse/item.py
index bc759b0..4c6ae19 100644
--- a/d2warehouse/item.py
+++ b/d2warehouse/item.py
@@ -17,7 +17,7 @@
 import json
 import os
 import re
-from typing import Optional
+from typing import Optional, assert_never
 from bitarray import bitarray
 from bitarray.util import int2ba
 from dataclasses import dataclass
@@ -122,6 +122,60 @@ class Item:
     quantity: int | None = None
     stats: list[Stat] | None = None
 
+    @property
+    def basename(self) -> str:
+        return lookup_basetype(self.code)["name"]
+
+    @property
+    def name(self) -> str:
+        match self.quality:
+            case Quality.LOW:
+                return f"{self.low_quality} {self.basename}"
+            case Quality.NORMAL | None:
+                return self.basename
+            case Quality.HIGH:
+                return f"Superior {self.basename}"
+            case Quality.MAGIC:
+                return f"<prefix> {self.basename} <suffix>"
+            case Quality.SET:
+                assert self.set_id is not None
+                return lookup_set_item(self.set_id)["name"]
+            case Quality.RARE:
+                # FIXME
+                return "<rare> <name>"
+            case Quality.UNIQUE:
+                assert self.unique_id is not None
+                return lookup_unique(self.unique_id)["name"]
+            case Quality.CRAFTED:
+                # FIXME
+                return "<crafted> <name>"
+            case _ as unreachable:
+                assert_never(unreachable)
+
+    @property
+    def color(self) -> str:
+        if self.is_runeword:
+            return "runeword"
+        match self.quality:
+            case Quality.LOW:
+                return "low"
+            case Quality.NORMAL | None:
+                return "normal"
+            case Quality.HIGH:
+                return "normal"
+            case Quality.MAGIC:
+                return "magic"
+            case Quality.SET:
+                return "set"
+            case Quality.RARE:
+                return "rare"
+            case Quality.UNIQUE:
+                return "unique"
+            case Quality.CRAFTED:
+                return "crafted"
+            case _ as unreachable:
+                assert_never(unreachable)
+
     def print(self, indent=5, with_raw=False):
         properties = []
         base_name = lookup_basetype(self.code)["name"]
@@ -204,7 +258,9 @@ class Item:
         base = lookup_basetype(self.code)
         return base["width"], base["height"]
 
-    def write_to_db(self, socketed_into=None, commit=True) -> int:
+    def write_to_db(self, socketed_into=None, commit=True, db=None) -> int:
+        if db is None:
+            db = get_db()
         name = lookup_basetype(self.code)["name"]
         # FIXME: handle magic & rare names
         if self.is_runeword:
@@ -218,7 +274,6 @@ class Item:
             lookup_set_item(self.set_id)["set"] if self.quality == Quality.SET else None
         )
 
-        db = get_db()
         cur = db.cursor()
         cur.execute(
             """INSERT INTO item (itembase_name, socketed_into, raw_data, raw_version,
@@ -305,8 +360,9 @@ class Item:
 
         return item_id
 
-    def load_from_db(id: int) -> "Item":
-        db = get_db()
+    def load_from_db(id: int, db=None) -> "Item":
+        if db is None:
+            db = get_db()
         row = db.execute(
             """SELECT raw_data, raw_version, is_identified, is_socketed,
             is_beginner, is_simple, is_ethereal, is_personalized, is_runeword, code,
diff --git a/d2warehouse/schema.sql b/d2warehouse/schema.sql
index 7ba639a..5c7ae95 100644
--- a/d2warehouse/schema.sql
+++ b/d2warehouse/schema.sql
@@ -1,7 +1,7 @@
-DROP TABLE IF EXISTS item_stat;
-DROP TABLE IF EXISTS item_affix;
-DROP TABLE IF EXISTS item_extra;
-DROP TABLE IF EXISTS item;
+--DROP TABLE IF EXISTS item_stat;
+--DROP TABLE IF EXISTS item_affix;
+--DROP TABLE IF EXISTS item_extra;
+--DROP TABLE IF EXISTS item;
 
 CREATE TABLE item (
     id INTEGER PRIMARY KEY,