add postgres get best node query and casts for posgres support with comment

Author: taddes

tools/tokenserver/database.py

@@ -180,7 +180,9 @@
 """)
 
 
-_GET_BEST_NODE = sqltext("""\
+# MySQL: log(0) returns NULL, and NULLs sort first with ASC — zero-load
+# nodes naturally win. Original query unchanged.
+_GET_BEST_NODE_MYSQL = sqltext("""\
 select
     id, node
 from
@@ -192,7 +194,26 @@
     and downed = 0
     and backoff = 0
 order by
-    log(GREATEST(current_load, 1)) / log(capacity)
+    log(current_load) / log(capacity)
+limit 1
+""")
+
+# PostgreSQL: log(0) raises InvalidArgumentForLogarithm, so we use NULLIF
+# to convert zero to NULL explicitly. NULLS FIRST replicates MySQL's default
+# NULL-first ASC sort order, ensuring zero-load nodes are always preferred.
+_GET_BEST_NODE_POSTGRES = sqltext("""\
+select
+    id, node
+from
+    nodes
+where
+    service = :service
+    and available > 0
+    and capacity > current_load
+    and downed = 0
+    and backoff = 0
+order by
+    log(NULLIF(current_load, 0)) / log(capacity) NULLS FIRST
 limit 1
 """)
 
@@ -656,8 +677,10 @@ def add_node(self, node, capacity, **kwds):
             capacity=capacity,
             available=available,
             current_load=kwds.get("current_load", 0),
-            downed=kwds.get("downed", 0),
-            backoff=kwds.get("backoff", 0),
+            # Cast to int: optparse action="store_true" produces Python bools,
+            # which postgres rejects for INTEGER columns (MySQL coerces silently).
+            downed=int(kwds.get("downed", 0)),
+            backoff=int(kwds.get("backoff", 0)),
         )
         res.close()
 
@@ -679,6 +702,11 @@ def update_node(self, node, **kwds):
         query += """
             where service = :service and node = :node
         """
+        # Cast boolean fields to int: Python bools are rejected by postgres
+        # INTEGER columns. MySQL coerces silently; postgres does not.
+        for field in ("downed", "backoff"):
+            if field in values:
+                values[field] = int(values[field])
         values["service"] = self._get_service_id(SERVICE_NAME)
         values["node"] = node
         if kwds:
@@ -750,8 +778,16 @@ def get_best_node(self):
             # capacity.  This loop allows a maximum of five retries before
             # bailing out.
             for _ in range(5):
+                # Select the appropriate query variant — postgres requires
+                # explicit NULL handling for log(0) and NULL sort order that
+                # MySQL handles implicitly.
+                best_node_query = (
+                    _GET_BEST_NODE_POSTGRES
+                    if self.db_mode == "postgresql"
+                    else _GET_BEST_NODE_MYSQL
+                )
                 res = self._execute_sql(
-                    _GET_BEST_NODE, service=self._get_service_id(SERVICE_NAME)
+                    best_node_query, service=self._get_service_id(SERVICE_NAME)
                 )
                 row = res.fetchone()
                 res.close()