From 826d5ce6264d19d9c2d6cf02cffb3fd13e925fbb Mon Sep 17 00:00:00 2001
From: Christoph Ladurner <christoph.ladurner@tugraz.at>
Date: Mon, 16 Dec 2024 12:59:33 +0100
Subject: [PATCH] model: add UTCDateTime

---
 invenio_banners/records/models.py | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/invenio_banners/records/models.py b/invenio_banners/records/models.py
index c9238d1..8988edb 100644
--- a/invenio_banners/records/models.py
+++ b/invenio_banners/records/models.py
@@ -8,18 +8,37 @@
 
 """Models."""
 
-from datetime import datetime
+from datetime import datetime, timezone
 
 import sqlalchemy as sa
 from flask import current_app
 from invenio_db import db
 from sqlalchemy import or_
 from sqlalchemy.sql import text
+from sqlalchemy.types import DateTime, TypeDecorator
 from sqlalchemy_utils.models import Timestamp
 
 from ..services.errors import BannerNotExistsError
 
 
+class UTCDateTime(TypeDecorator):
+    """Custom UTC datetime type."""
+
+    impl = DateTime
+
+    def process_bind_param(self, value, dialect):
+        """Process value storing into database."""
+        if isinstance(value, datetime):
+            return value.replace(tzinfo=None)
+        return value
+
+    def process_result_value(self, value, dialect):
+        """Process value retrieving from database."""
+        if isinstance(value, datetime):
+            return value.replace(tzinfo=None)
+        return value
+
+
 class BannerModel(db.Model, Timestamp):
     """Defines a message to show to users."""
 
@@ -36,10 +55,12 @@ class BannerModel(db.Model, Timestamp):
     category = db.Column(db.String(20), nullable=False)
     """Category of the message, for styling messages per category."""
 
-    start_datetime = db.Column(db.DateTime, nullable=False, default=datetime.utcnow)
+    start_datetime = db.Column(
+        UTCDateTime, nullable=False, default=lambda: datetime.now(timezone.utc)
+    )
     """Start date and time (UTC), can be immediate or delayed."""
 
-    end_datetime = db.Column(db.DateTime, nullable=True)
+    end_datetime = db.Column(UTCDateTime, nullable=True)
     """End date and time (UTC), must be after `start` or forever if null."""
 
     active = db.Column(db.Boolean(name="active"), nullable=False, default=True)