Improved compatibility with postgres

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:

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