From ed3e30ea7581a67b83cdf0f097324a338da92e2f Mon Sep 17 00:00:00 2001
From: Donkie <daniel.cf.hultgren@gmail.com>
Date: Tue, 9 May 2023 20:03:19 +0200
Subject: [PATCH] Improved compatibility with postgres

---
 pyproject.toml                |  2 +-
 spoolman/database/database.py |  8 ++++-
 spoolman/database/spool.py    | 60 ++++++++++++++++++++++++-----------
 3 files changed, 50 insertions(+), 20 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 189af35..7f6caf8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,7 @@ dependencies = {file = ["requirements.txt"]}
 
 [tool.ruff]
 select = ["ALL"]
-ignore = ["A003", "D101", "D104", "D406", "D407", "S104", "TRY201", "TRY003", "EM101", "EM102"]
+ignore = ["A003", "D101", "D104", "D406", "D407", "S104", "TRY201", "TRY003", "EM101", "EM102", "DTZ003"]
 line-length = 120
 target-version = "py39"
 
diff --git a/spoolman/database/database.py b/spoolman/database/database.py
index b8e59c9..2394586 100644
--- a/spoolman/database/database.py
+++ b/spoolman/database/database.py
@@ -20,7 +20,13 @@ class Database:
 
     def connect(self: "Database") -> None:
         """Connect to the database."""
-        self.engine = create_async_engine(self.connection_url, connect_args={"check_same_thread": False})
+        connect_args = {}
+        if self.connection_url.drivername == "sqlite":
+            connect_args["check_same_thread"] = False
+        elif self.connection_url.drivername == "postgresql":
+            connect_args["options"] = "-c timezone=utc"
+
+        self.engine = create_async_engine(self.connection_url, connect_args=connect_args)
         self.session_maker = async_sessionmaker(self.engine, autocommit=False, autoflush=True)
 
     async def create_tables(self: "Database") -> None:
diff --git a/spoolman/database/spool.py b/spoolman/database/spool.py
index 1af9bef..5d1803c 100644
--- a/spoolman/database/spool.py
+++ b/spoolman/database/spool.py
@@ -3,7 +3,7 @@
 from datetime import datetime, timezone
 from typing import Optional
 
-from sqlalchemy import select
+import sqlalchemy
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import contains_eager, joinedload
 
@@ -48,12 +48,11 @@ async def create(
     return db_item
 
 
-async def get_by_id(db: AsyncSession, spool_id: int, with_for_update: Optional[bool] = None) -> models.Spool:
+async def get_by_id(db: AsyncSession, spool_id: int) -> models.Spool:
     """Get a spool object from the database by the unique ID."""
     spool = await db.get(
         models.Spool,
         spool_id,
-        with_for_update=with_for_update,  # type: ignore  # noqa: PGH003
         options=[joinedload("*")],  # Load all nested objects as well
     )
     if spool is None:
@@ -74,7 +73,7 @@ async def find(
 ) -> list[models.Spool]:
     """Find a list of spool objects by search criteria."""
     stmt = (
-        select(models.Spool)
+        sqlalchemy.select(models.Spool)
         .join(models.Spool.filament)
         .join(models.Filament.vendor)
         .options(contains_eager(models.Spool.filament).contains_eager(models.Filament.vendor))
@@ -125,6 +124,21 @@ async def delete(db: AsyncSession, spool_id: int) -> None:
     await db.delete(spool)
 
 
+async def use_weight_safe(db: AsyncSession, spool_id: int, weight: float) -> None:
+    """Consume filament from a spool by weight in a way that is safe against race conditions.
+
+    Args:
+        db (AsyncSession): Database session
+        spool_id (int): Spool ID
+        weight (float): Filament weight to consume, in grams
+    """
+    await db.execute(
+        sqlalchemy.update(models.Spool)
+        .where(models.Spool.id == spool_id)
+        .values(used_weight=models.Spool.used_weight + weight),
+    )
+
+
 # TODO: Make unit tests for race conditions on these
 async def use_weight(db: AsyncSession, spool_id: int, weight: float) -> models.Spool:
     """Consume filament from a spool by weight.
@@ -140,12 +154,13 @@ async def use_weight(db: AsyncSession, spool_id: int, weight: float) -> models.S
     Returns:
         models.Spool: Updated spool object
     """
-    spool = await get_by_id(db, spool_id, with_for_update=True)
-    spool.used_weight += weight
+    await use_weight_safe(db, spool_id, weight)
+
+    spool = await get_by_id(db, spool_id)
 
     if spool.first_used is None:
-        spool.first_used = datetime.now(tz=timezone.utc)
-    spool.last_used = datetime.now(tz=timezone.utc)
+        spool.first_used = datetime.utcnow()
+    spool.last_used = datetime.utcnow()
 
     await db.flush()
     return spool
@@ -165,19 +180,28 @@ async def use_length(db: AsyncSession, spool_id: int, length: float) -> models.S
     Returns:
         models.Spool: Updated spool object
     """
-    spool = await get_by_id(db, spool_id, with_for_update=True)
-
-    filament = spool.filament
-
-    spool.used_weight += weight_from_length(
-        length=length,
-        radius=filament.diameter / 2,
-        density=filament.density,
+    # Get filament diameter and density
+    result = await db.execute(
+        sqlalchemy.select(models.Filament.diameter, models.Filament.density)
+        .join(models.Spool, models.Spool.filament_id == models.Filament.id)
+        .where(models.Spool.id == spool_id),
     )
+    filament_info = result.one()
+
+    # Calculate and use weight
+    weight = weight_from_length(
+        length=length,
+        radius=filament_info[0] / 2,
+        density=filament_info[1],
+    )
+    await use_weight_safe(db, spool_id, weight)
+
+    # Get spool with new weight and update first_used and last_used
+    spool = await get_by_id(db, spool_id)
 
     if spool.first_used is None:
-        spool.first_used = datetime.now(tz=timezone.utc)
-    spool.last_used = datetime.now(tz=timezone.utc)
+        spool.first_used = datetime.utcnow()
+    spool.last_used = datetime.utcnow()
 
     await db.flush()
     return spool