Add rollback logic when db update fails after DNS changes

src/ddns_service/cleanup.py

@@ -4,7 +4,7 @@ import logging
 import threading
 
 from . import now_utc
-from .models import Hostname, User
+from .models import DatabaseError, Hostname, User
 from datetime import timedelta
 
 
@@ -31,13 +31,15 @@ def cleanup_expired(app, start_time=None):
         ipv6_expired = False
 
         if hostname.last_ipv4:
-            last_update = max(hostname.last_ipv4_update, start_time)
+            last_update = max(hostname.last_ipv4_update, start_time) if start_time \
+                else hostname.last_ipv4_update
             expiry_time = last_update + timedelta(seconds=hostname.expiry_ttl)
             if now > expiry_time:
                 ipv4_expired = True
 
         if hostname.last_ipv6:
-            last_update = max(hostname.last_ipv6_update, start_time)
+            last_update = max(hostname.last_ipv6_update, start_time) if start_time \
+                else hostname.last_ipv6_update
             expiry_time = last_update + timedelta(seconds=hostname.expiry_ttl)
             if now > expiry_time:
                 ipv6_expired = True
@@ -45,58 +47,92 @@ def cleanup_expired(app, start_time=None):
         if not ipv4_expired and not ipv6_expired:
             continue
 
-        ipv4_deleted = False
-        ipv6_deleted = False
+        old_ipv4 = hostname.last_ipv4
+        old_ipv6 = hostname.last_ipv6
+        ipv4_dns_deleted = False
+        ipv6_dns_deleted = False
 
-        if app.dns_service:
-            if ipv4_expired:
-                logging.info(
-                    f"Cleanup: Host expired: hostname={hostname.hostname} zone={hostname.zone} "
-                    f"ipv4={hostname.last_ipv4}"
+        if ipv4_expired:
+            logging.info(
+                f"Cleanup: Host expired: hostname={hostname.hostname} zone={hostname.zone} "
+                f"ipv4={hostname.last_ipv4}"
+            )
+            try:
+                ipv4_exists = app.dns_service.query_record(
+                    hostname.hostname, hostname.zone, "A")
+                if ipv4_exists:
+                    app.dns_service.delete_record(
+                        hostname.hostname, hostname.zone, "A")
+                    ipv4_dns_deleted = True
+                hostname.last_ipv4 = None
+            except Exception as e:
+                logging.error(f"DNS error: {e}")
+                logging.error(
+                    f"Cleanup failed: hostname={hostname.hostname} "
+                    f"zone={hostname.zone} type=A"
                 )
-                try:
-                    app.dns_service.delete_record(hostname.hostname, hostname.zone, "A")
-                    ipv4_deleted = True
-                except Exception as e:
-                    logging.error(f"DNS error: {e}")
-                    logging.error(
-                        f"Cleanup failed: hostname={hostname.hostname} "
-                        f"zone={hostname.zone} type=A"
-                    )
 
-            if ipv6_expired:
-                logging.info(
-                    f"Cleanup: Host expired: hostname={hostname.hostname} zone={hostname.zone} "
-                    f"ipv6={hostname.last_ipv6}"
+        if ipv6_expired:
+            logging.info(
+                f"Cleanup: Host expired: hostname={hostname.hostname} zone={hostname.zone} "
+                f"ipv6={hostname.last_ipv6}"
+            )
+            try:
+                ipv6_exists = app.dns_service.query_record(
+                    hostname.hostname, hostname.zone, "AAAA")
+                if ipv6_exists:
+                    app.dns_service.delete_record(
+                        hostname.hostname, hostname.zone, "AAAA")
+                    ipv6_dns_deleted = True
+                hostname.last_ipv6 = None
+            except Exception as e:
+                logging.error(f"DNS error: {e}")
+                logging.error(
+                    f"Cleanup failed: hostname={hostname.hostname} "
+                    f"zone={hostname.zone} type=AAAA"
                 )
-                try:
-                    app.dns_service.delete_record(hostname.hostname, hostname.zone, "AAAA")
-                    ipv6_deleted = True
-                except Exception as e:
-                    logging.error(f"DNS error: {e}")
-                    logging.error(
-                        f"Cleanup failed: hostname={hostname.hostname} "
-                        f"zone={hostname.zone} type=AAAA"
-                    )
 
-        if not (ipv4_deleted or ipv6_deleted):
+        if hostname.last_ipv4 == old_ipv4 and hostname.last_ipv6 == old_ipv6:
+            continue
+
+        try:
+            hostname.save()
+        except DatabaseError as e:
+            logging.error(
+                f"DB save failed after retries: hostname={hostname.hostname} "
+                f"zone={hostname.zone}: {e}"
+            )
+
+            # Rollback: re-add DNS records that were deleted
+            if ipv4_dns_deleted:
+                try:
+                    app.dns_service.update_record(
+                        hostname.hostname, hostname.zone,
+                        old_ipv4, hostname.dns_ttl)
+                except Exception as e:
+                    logging.error(f"DNS rollback failed (A): {e}")
+
+            if ipv6_dns_deleted:
+                try:
+                    app.dns_service.update_record(
+                        hostname.hostname, hostname.zone,
+                        old_ipv6, hostname.dns_ttl)
+                except Exception as e:
+                    logging.error(f"DNS rollback failed (AAAA): {e}")
+
             continue
 
         if app.email_service:
+            # Restore old IPs on in-memory model for email template
+            hostname.last_ipv4 = old_ipv4
+            hostname.last_ipv6 = old_ipv6
             app.email_service.send_expiry_notification(
                 hostname.user.email,
                 hostname,
-                ipv4_deleted,
-                ipv6_deleted
+                hostname.last_ipv4 == old_ipv4,
+                hostname.last_ipv6 == old_ipv6
             )
 
-        # Clear IP addresses only if DNS delete succeeded
-        if ipv4_deleted:
-            hostname.last_ipv4 = None
-        if ipv6_deleted:
-            hostname.last_ipv6 = None
-
-        hostname.save()
         expired_count += 1
 
     return expired_count

src/ddns_service/server.py

@@ -452,14 +452,29 @@ class DDNSRequestHandler(BaseHTTPRequestHandler):
 
         return hostname
 
+    def _rollback_dns(self, hostname, old_ip, record_type):
+        """Roll back a DNS record to its previous value."""
+        try:
+            if old_ip:
+                self.app.dns_service.update_record(
+                    hostname.hostname, hostname.zone,
+                    old_ip, hostname.dns_ttl)
+            else:
+                self.app.dns_service.delete_record(
+                    hostname.hostname, hostname.zone, record_type)
+        except Exception as e:
+            logging.error(f"DNS rollback failed ({record_type}): {e}")
+
     def _process_ip_update(self, client_ip, user, hostname, ipv4, ipv6, notify_change):
         """Process IP update for hostname."""
         now = now_utc()
 
+        old_ipv4 = hostname.last_ipv4
+        old_ipv6 = hostname.last_ipv6
         ipv4_changed = False
         ipv6_changed = False
+
         if ipv4:
-            hostname.last_ipv4_update = now
             if ipv4 != hostname.last_ipv4:
                 # Update DNS IPv4 record
                 try:
@@ -469,10 +484,8 @@ class DDNSRequestHandler(BaseHTTPRequestHandler):
                         ipv4,
                         hostname.dns_ttl
                     )
-                    hostname.last_ipv4 = ipv4
                     ipv4_changed = True
                 except Exception as e:
-                    hostname.save()
                     logging.error(f"DNS error: {e}")
                     raise DDNSError(
                         "Update failed",
@@ -484,7 +497,6 @@ class DDNSRequestHandler(BaseHTTPRequestHandler):
                     )
 
         if ipv6:
-            hostname.last_ipv6_update = now
             if ipv6 != hostname.last_ipv6:
                 # Update DNS IPv6 record
                 try:
@@ -494,11 +506,12 @@ class DDNSRequestHandler(BaseHTTPRequestHandler):
                         ipv6,
                         hostname.dns_ttl
                     )
-                    hostname.last_ipv6 = ipv6
                     ipv6_changed = True
                 except Exception as e:
-                    hostname.save()
                     logging.error(f"DNS error: {e}")
+                    # Roll back IPv4 DNS if it was changed
+                    if ipv4_changed:
+                        self._rollback_dns(hostname, old_ipv4, "A")
                     raise DDNSError(
                         "Update failed",
                         STATUS_DNSERR,
@@ -508,8 +521,32 @@ class DDNSRequestHandler(BaseHTTPRequestHandler):
                         ipv6=ipv6
                     )
 
+        if ipv4_changed:
+            hostname.last_ipv4 = ipv4
+            hostname.last_ipv4_update = now
+        if ipv6_changed:
+            hostname.last_ipv6 = ipv6
+            hostname.last_ipv6_update = now
+
         # Update database
-        hostname.save()
+        try:
+            hostname.save()
+        except DatabaseError as e:
+            logging.error(
+                f"DB save failed after retries: hostname={hostname.hostname} "
+                f"zone={hostname.zone}: {e}"
+            )
+            if ipv4_changed:
+                self._rollback_dns(hostname, old_ipv4, "A")
+            if ipv6_changed:
+                self._rollback_dns(hostname, old_ipv6, "AAAA")
+            raise DDNSError(
+                "Update failed",
+                STATUS_DNSERR,
+                client=client_ip,
+                hostname=hostname.hostname,
+                zone=hostname.zone
+            )
 
         changed_addrs = ""
         if ipv4_changed: