modified: catalog_zone_manager.py

modified:   management/commands/dns-export-zone.py
modified:   management/commands/dns-export-zones.py
modified:   management/commands/dns-transfer-endpoint.py
modified:   notify_handler.py
modified:   request_handler.py
Guard all db requests against broken database connections.

netbox_dns_bridge/management/commands/dns-export-zone.py

@@ -1,8 +1,3 @@
-import os
-import dns.zone
-import dns.rdatatype
-import dns.rdataclass
-import dns.exception
 import netbox_dns.models
 from netbox_dns_bridge.utils import export_bind_zone_file
 from django.core.management.base import BaseCommand, CommandError
@@ -24,10 +19,6 @@ class Command(BaseCommand):
         )
 
     def handle(self, *args, **options):
-        # if len(options) < 2:
-        #    print("export-zone <zone name> <file path>")
-        #    sys.exit(1)
-
         # Load parameters
         view_name = options["view"]
         zone_name = options["zone"]
@@ -48,13 +39,11 @@ class Command(BaseCommand):
 
             if nb_zone:
                 export_bind_zone_file(nb_zone, file_path=file_path)
+                self.stdout.write(
+                    self.style.SUCCESS(f"Zone '{zone_name}' exported to '{file_path}'")
+                )
             else:
-                print("Zone not found in Netbox. Aborting")
-                sys.exit(1)
+                self.stderr.write("Zone not found in Netbox. Aborting")
 
         except Exception as e:
             raise CommandError(f"Failed to export zone: {e}")
-
-        self.stdout.write(
-            self.style.SUCCESS(f"Zone '{zone_name}' exported to '{file_path}'")
-        )

netbox_dns_bridge/management/commands/dns-export-zones.py

@@ -1,8 +1,4 @@
 import os
-import dns.zone
-import dns.rdatatype
-import dns.rdataclass
-import dns.exception
 import netbox_dns.models
 from netbox_dns_bridge.utils import export_bind_zone_file
 from django.core.management.base import BaseCommand, CommandError

netbox_dns_bridge/management/commands/dns-transfer-endpoint.py

@@ -9,36 +9,24 @@ import dns.rdataclass
 import dns.rdtypes
 import dns.exception
 import dns.renderer
-import logging
 
 from django.core.management.base import BaseCommand
 from django.conf import settings
+from django.db import close_old_connections, OperationalError
 from netbox_dns.models import View
 from netbox_dns_bridge import catalog_zone_manager as catzm
 from netbox_dns_bridge.request_handler import UDPRequestHandler, TCPRequestHandler
 from netbox_dns_bridge.dns_server import UDPDNSServer, TCPDNSServer
-from netbox_dns_bridge.models import IntegerKeyValueSetting
 from netbox_dns_bridge.logger import get_logger
 
-logger = get_logger(__name__)
+LOGGER = get_logger(__name__)
+SETTINGS = settings.PLUGINS_CONFIG.get("netbox_dns_bridge", {})
 
 
 class Command(BaseCommand):
     help = "Run a minimal AXFR DNS server using data from NetBox DNS plugin"
 
-    def load_settings(self):
-        self.settings = settings.PLUGINS_CONFIG.get("netbox_dns_bridge", None)
-        if not self.settings:
-            raise RuntimeError(
-                "Command failed to initialize due to missing settings. Terminating Netbox."
-            )
-
-        self.tsig_keys = self.settings.get("tsig_keys", None)
-        if not self.tsig_keys:
-            raise RuntimeError("tsig_keys variable not set in plugin settings.")
-
-    # Load TSIG keys and map them to views
-    def load_tsig_key_settings(self):
+    def _load_tsig_key_settings(self):
         self.keyring = {}
         self.tsig_view_map = {}
 
@@ -48,7 +36,7 @@ class Command(BaseCommand):
             algorithm_str = data.get("algorithm", "hmac-sha256")
 
             if not raw_key_name or not secret:
-                logger.error(
+                LOGGER.error(
                     f"Skipping TSIG key for view {view_name}: missing keyname or secret."
                 )
                 continue
@@ -56,10 +44,12 @@ class Command(BaseCommand):
             try:
                 nb_view = View.objects.get(name=view_name)
             except View.DoesNotExist:
-                logger.error(
+                LOGGER.error(
                     f"Skipping TSIG key {raw_key_name}: view '{view_name}' not found."
                 )
                 continue
+            except OperationalError:
+                close_old_connections()
 
             # Normalize key name to absolute DNS name
             key_name_obj = dns.name.from_text(raw_key_name, origin=None).canonicalize()
@@ -71,11 +61,11 @@ class Command(BaseCommand):
                 name=key_name_obj, secret=secret, algorithm=algorithm_str
             )
             self.tsig_view_map[key_name_str] = nb_view
-            logger.info(f"Loaded TSIG key {key_name_str} for view {nb_view.name}")
+            LOGGER.info(f"Loaded TSIG key {key_name_str} for view {nb_view.name}")
 
         if not self.keyring:
             msg = "No TSIG keys found in database."
-            logger.critical(msg)
+            LOGGER.critical(msg)
             raise RuntimeError(msg)
 
     def add_arguments(self, parser):
@@ -92,9 +82,11 @@ class Command(BaseCommand):
         address = options["address"]
         catzm.init()
 
-        # Initialize settings
-        self.load_settings()
-        self.load_tsig_key_settings()
+        self.tsig_keys = SETTINGS.get("tsig_keys", None)
+        if not self.tsig_keys:
+            raise RuntimeError("tsig_keys variable not set in plugin settings.")
+
+        self._load_tsig_key_settings()
 
         udp_server = UDPDNSServer(
             (address, port), UDPRequestHandler, self.keyring, self.tsig_view_map
@@ -104,15 +96,15 @@ class Command(BaseCommand):
             (address, port), TCPRequestHandler, self.keyring, self.tsig_view_map
         )
 
-        def run_udp_server(server):
-            logger.info(f"Query endpoint listening on {address} udp/{port}")
+        def _run_udp_server(server):
+            LOGGER.info(f"Query endpoint listening on {address} udp/{port}")
             server.serve_forever()
 
         udp_thread = threading.Thread(
-            target=run_udp_server, args=(udp_server,), daemon=True
+            target=_run_udp_server, args=(udp_server,), daemon=True
         )
 
         udp_thread.start()
 
-        logger.info(f"Query endpoint listening on {address} tcp/{port}")
+        LOGGER.info(f"Query endpoint listening on {address} tcp/{port}")
         tcp_server.serve_forever()