Fix error with writing prices to database

src/helpers/database.py

@@ -1,8 +1,9 @@
 from datetime import datetime, timezone
 
+import psycopg.errors
+import sqlalchemy
 from hexbytes import HexBytes
-import psycopg
-from sqlalchemy import text, insert, Table, Column, Integer, LargeBinary, MetaData
+from sqlalchemy import Table, MetaData, insert, text, update
 from sqlalchemy.engine import Engine
 
 from src.helpers.config import check_db_connection, logger
@@ -167,28 +168,60 @@ def write_transaction_tokens(
 
     def write_prices_new(self, prices: list[tuple[str, int, float, str]]) -> None:
         """Write prices to database."""
-        query = (
-            "INSERT INTO prices (token_address, time, price, source) "
-            "VALUES (:token_address, :time, :price, :source);"
-        )
-        for token_address, time, price, source in prices:
-            try:
-                self.execute_and_commit(
-                    query,
-                    {
-                        "token_address": bytes.fromhex(token_address[2:]),
-                        "time": datetime.fromtimestamp(time, tz=timezone.utc),
-                        "price": price,
-                        "source": source,
-                    },
-                )
-            except Exception as err:
-                pass
-            # except psycopg.errors.NumericValueOutOfRange:
-            #     logger.info(
-            #         f"Error while writing price data. token: {token_address}, "
-            #         f"time: {time}, price: {price}, source: {source}"
-            #     )
+        try:
+            prices_table = Table("prices", MetaData(), autoload_with=self.engine)
+
+            with self.engine.connect() as conn:
+                with conn.begin():
+                    for token_address, time, price, source in prices:
+                        savepoint = conn.begin_nested()
+                        # First try to insert
+                        try:
+                            insert_stmt = insert(prices_table).values(
+                                token_address=bytes.fromhex(token_address[2:]),
+                                time=datetime.fromtimestamp(time, tz=timezone.utc),
+                                price=price,
+                                source=source,
+                            )
+                            conn.execute(insert_stmt)
+                        except (
+                            sqlalchemy.exc.IntegrityError
+                        ):  # update if insertion failed
+                            logger.info(
+                                f"Error while inserting price data. token: {token_address}, "
+                                f"time: {time}, price: {price}, source: {source}\n"
+                                "Updating price instead."
+                            )
+                            savepoint.rollback()
+                            update_stmt = (
+                                update(prices_table)
+                                .where(
+                                    prices_table.c.token_address
+                                    == bytes.fromhex(token_address[2:])
+                                )
+                                .where(
+                                    prices_table.c.time
+                                    == datetime.fromtimestamp(time, tz=timezone.utc)
+                                )
+                                .where(prices_table.c.source == source)
+                                .values(price=price)
+                            )
+                            res = conn.execute(update_stmt)
+                            if res.rowcount == 0:
+                                logger.error(
+                                    f"Update failed, no matching record found. token: "
+                                    f"{token_address}, time: {time}, price: {price},"
+                                    f"source: {source}\n"
+                                    "This indicates a faulty database state."
+                                )
+
+                        except psycopg.errors.NumericValueOutOfRange:
+                            logger.info(
+                                f"Value out of range for price data. token: {token_address}, "
+                                f"time: {time}, price: {price}, source: {source}"
+                            )
+        except Exception as e:
+            logger.error(f"Error writing prices: {e}")
 
     def get_latest_transaction(self) -> str | None:
         """Get latest transaction hash.

tests/unit/test_database.py

@@ -1,5 +1,6 @@
 from datetime import datetime, timezone
 
+import pytest
 from hexbytes import HexBytes
 from sqlalchemy import create_engine, text
 
@@ -75,13 +76,21 @@ def tests_write_prices():
     token_prices = [
         (
             "0xA0B86991C6218B36C1D19D4A2E9EB0CE3606EB48",
-            int(datetime.fromisoformat("2024-10-10 16:48:47.000000").timestamp()),
+            int(
+                datetime.fromisoformat("2024-10-10 16:48:47.000000")
+                .replace(tzinfo=timezone.utc)
+                .timestamp()
+            ),
             0.000420454193230350,
             "coingecko",
         ),
         (
             "0x68BBED6A47194EFF1CF514B50EA91895597FC91E",
-            int(datetime.fromisoformat("2024-10-10 16:49:47.000000").timestamp()),
+            int(
+                datetime.fromisoformat("2024-10-10 16:49:47.000000")
+                .replace(tzinfo=timezone.utc)
+                .timestamp()
+            ),
             0.000000050569218629,
             "moralis",
         ),
@@ -99,11 +108,74 @@ def tests_write_prices():
         ).all()
     for i, (token_address, time, price, source) in enumerate(token_prices):
         assert HexBytes(res[i][0]) == HexBytes(token_address)
-        assert res[i][1].timestamp() == time
+        assert res[i][1].replace(tzinfo=timezone.utc).timestamp() == time
         assert float(res[i][2]) == price
         assert res[i][3] == source
 
 
+def tests_write_prices_duplicates():
+    engine = create_engine(
+        f"postgresql+psycopg://postgres:postgres@localhost:5432/mainnet"
+    )
+    db = Database(engine, "mainnet")
+    token_address = "0xA0B86991C6218B36C1D19D4A2E9EB0CE3606EB48"
+    time = int(
+        datetime.fromisoformat("2024-10-10 16:48:47.000000")
+        .replace(tzinfo=timezone.utc)
+        .timestamp()
+    )
+    source = "coingecko"
+    price_1 = 0.000420454193230350
+    price_2 = price_1 + 0.0001
+    # truncate table
+    with engine.connect() as conn:
+        conn.execute(text("TRUNCATE prices"))
+        conn.commit()
+    # write data twice
+    db.write_prices_new([(token_address, time, price_1, source)])
+    db.write_prices_new([(token_address, time, price_2, source)])
+    # read data
+    with engine.connect() as conn:
+        res = conn.execute(
+            text("SELECT token_address, time, price, source FROM prices")
+        ).all()
+    assert HexBytes(res[0][0]) == HexBytes(token_address)
+    assert res[0][1].replace(tzinfo=timezone.utc).timestamp() == time
+    assert float(res[0][2]) == pytest.approx(price_2)
+    assert res[0][3] == source
+
+
+def tests_write_prices_large_value():
+    """Test that writing large prices does not crash.
+
+    The expected behavior is to not write the price to the database and log an error.
+    """
+    engine = create_engine(
+        f"postgresql+psycopg://postgres:postgres@localhost:5432/mainnet"
+    )
+    db = Database(engine, "mainnet")
+    token_address = "0xA0B86991C6218B36C1D19D4A2E9EB0CE3606EB48"
+    time = int(
+        datetime.fromisoformat("2024-10-10 16:48:47.000000")
+        .replace(tzinfo=timezone.utc)
+        .timestamp()
+    )
+    source = "coingecko"
+    price_1 = 1e100
+    # truncate table
+    with engine.connect() as conn:
+        conn.execute(text("TRUNCATE prices"))
+        conn.commit()
+    # write data twice
+    db.write_prices_new([(token_address, time, price_1, source)])
+    # read data
+    with engine.connect() as conn:
+        res = conn.execute(
+            text("SELECT token_address, time, price, source FROM prices")
+        ).all()
+    assert len(res) == 0
+
+
 def test_get_latest_transaction():
     # import has to happen after patching environment variable
     from src.helpers.database import Database