Fix an issue in rebuilding stash data for socketed items

- 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

@@ -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:

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,68 @@ 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 raw(self):
+        parts = [self.raw_data]
+        if self.sockets:
+            for item in self.sockets:
+                if item:
+                    parts.append(item.raw_data)
+        return b"".join(parts)
+
     def print(self, indent=5, with_raw=False):
         properties = []
         base_name = lookup_basetype(self.code)["name"]
@@ -204,7 +266,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 +282,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 +368,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,

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,

d2warehouse/stash.py

@@ -44,7 +44,7 @@ class StashTab:
 
     def raw(self) -> bytes:
         """Get the computed raw representation of the stash"""
-        item_raw = b"".join(item.raw_data for item in self.items)
+        item_raw = b"".join(item.raw() for item in self.items)
         raw_length = len(item_raw) + 0x44
         return (
             STASH_TAB_MAGIC

d2warehouse/tests/test_stash.py

@@ -84,3 +84,18 @@ class StashTest(unittest.TestCase):
         self.assertEqual(len(new_stash.tabs[0].items), 3)
         self.assertEqual(len(new_stash.tabs[1].items), 1)
         self.assertEqual(len(new_stash.tabs[2].items), 25)
+
+    def test_gemmed_raw(self):
+        data = bytes.fromhex(
+            "55aa55aa0100000063000000a40100006200000000000000000000000000"
+            "000000000000000000000000000000000000000000000000000000000000"
+            "000000004a4d010010088000050094a459a496629918020a484890ff1000"
+            "a0003500e07c6f0355aa55aa0100000063000000391b0000440000000000"
+            "000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000004a4d000055aa55aa01000000630000003905"
+            "000044000000000000000000000000000000000000000000000000000000"
+            "00000000000000000000000000000000000000004a4d0000"
+        )
+        stash = parse_stash(data)
+        rebuilt = stash.raw()
+        self.assertEqual(data, rebuilt)