From f9f9392b93134ba9dc07082155bee2b79e5bfc8d Mon Sep 17 00:00:00 2001
From: Alex Vandiver <alexmv@zulip.com>
Date: Tue, 5 Mar 2024 16:16:31 +0000
Subject: [PATCH] zulip_tools: Unify get_config and get_config_bool using
 overloads.

---
 scripts/lib/email-mirror-postfix  |  4 ++--
 scripts/lib/upgrade-zulip-stage-2 |  5 ++---
 scripts/lib/zulip_tools.py        | 28 +++++++++++++++-------------
 scripts/restart-server            |  4 ++--
 4 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/scripts/lib/email-mirror-postfix b/scripts/lib/email-mirror-postfix
index 8ccb9b991b..49a51db126 100755
--- a/scripts/lib/email-mirror-postfix
+++ b/scripts/lib/email-mirror-postfix
@@ -54,7 +54,7 @@ from urllib.parse import urlencode, urljoin, urlsplit
 from urllib.request import Request, urlopen
 
 sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
-from scripts.lib.zulip_tools import get_config_bool, get_config_file
+from scripts.lib.zulip_tools import get_config, get_config_file
 
 parser = argparse.ArgumentParser()
 
@@ -133,7 +133,7 @@ def send_email_mirror(
 
     if not urlsplit(host).scheme:
         config_file = get_config_file()
-        http_only = get_config_bool(config_file, "application_server", "http_only", False)
+        http_only = get_config(config_file, "application_server", "http_only", False)
         scheme = "http://" if http_only else "https://"
         host = scheme + host
 
diff --git a/scripts/lib/upgrade-zulip-stage-2 b/scripts/lib/upgrade-zulip-stage-2
index fee5d30a60..a31b979539 100755
--- a/scripts/lib/upgrade-zulip-stage-2
+++ b/scripts/lib/upgrade-zulip-stage-2
@@ -28,7 +28,6 @@ from scripts.lib.zulip_tools import (
     DEPLOYMENTS_DIR,
     assert_running_as_root,
     get_config,
-    get_config_bool,
     get_config_file,
     get_zulip_pwent,
     listening_publicly,
@@ -221,12 +220,12 @@ if not (args.skip_restart or args.skip_puppet):
     # We need to temporarily hold pgroonga, if installed -- upgrading
     # it without running the appropriate upgrade SQL can cause
     # PostgreSQL to crash
-    if get_config_bool(config_file, "machine", "pgroonga"):
+    if get_config(config_file, "machine", "pgroonga", False):
         subprocess.check_call(["apt-mark", "hold", "postgresql-*-pgdg-pgroonga"])
     logging.info("Upgrading system packages...")
     subprocess.check_call(["apt-get", "update"])
     subprocess.check_call(["apt-get", "-y", "--with-new-pkgs", "upgrade"])
-    if get_config_bool(config_file, "machine", "pgroonga"):
+    if get_config(config_file, "machine", "pgroonga", False):
         subprocess.check_call(["apt-mark", "unhold", "postgresql-*-pgdg-pgroonga"])
 
 # To bootstrap zulip-puppet-apply, we need to install the system yaml
diff --git a/scripts/lib/zulip_tools.py b/scripts/lib/zulip_tools.py
index f884a56f91..5189738385 100755
--- a/scripts/lib/zulip_tools.py
+++ b/scripts/lib/zulip_tools.py
@@ -16,7 +16,7 @@ import sys
 import time
 import uuid
 from datetime import datetime, timedelta
-from typing import IO, Any, Dict, List, Sequence, Set
+from typing import IO, Any, Dict, List, Sequence, Set, Union, overload
 from urllib.parse import SplitResult
 
 DEPLOYMENTS_DIR = "/home/zulip/deployments"
@@ -583,24 +583,26 @@ def assert_running_as_root(strip_lib_from_paths: bool = False) -> None:
         sys.exit(1)
 
 
+@overload
+def get_config(
+    config_file: configparser.RawConfigParser, section: str, key: str, default_value: str = ""
+) -> str: ...
+@overload
+def get_config(
+    config_file: configparser.RawConfigParser, section: str, key: str, default_value: bool
+) -> bool: ...
 def get_config(
     config_file: configparser.RawConfigParser,
     section: str,
     key: str,
-    default_value: str = "",
-) -> str:
-    if config_file.has_option(section, key):
-        return config_file.get(section, key)
-    return default_value
-
-
-def get_config_bool(
-    config_file: configparser.RawConfigParser, section: str, key: str, default_value: bool = False
-) -> bool:
+    default_value: Union[str, bool] = "",
+) -> Union[str, bool]:
     if config_file.has_option(section, key):
         val = config_file.get(section, key)
-        # This list is parallel to puppet/zulip/lib/puppet/functions/zulipconf.rb
-        return val.lower() in ["1", "y", "t", "true", "yes", "enable", "enabled"]
+        if isinstance(default_value, bool):
+            # This list is parallel to puppet/zulip/lib/puppet/functions/zulipconf.rb
+            return val.lower() in ["1", "y", "t", "true", "yes", "enable", "enabled"]
+        return val
     return default_value
 
 
diff --git a/scripts/restart-server b/scripts/restart-server
index ca8e069a2a..211286cb5b 100755
--- a/scripts/restart-server
+++ b/scripts/restart-server
@@ -18,7 +18,7 @@ from scripts.lib.zulip_tools import (
     ENDC,
     OKGREEN,
     WARNING,
-    get_config_bool,
+    get_config,
     get_config_file,
     get_tornado_ports,
     has_application_server,
@@ -185,7 +185,7 @@ if has_application_server():
     if (
         action == "restart"
         and not args.less_graceful
-        and get_config_bool(config_file, "application_server", "rolling_restart")
+        and get_config(config_file, "application_server", "rolling_restart", False)
         and os.path.exists("/home/zulip/deployments/uwsgi-control")
     ):
         # See if it's currently running