From 21cc186105b3ad7d8988bb86fded3caa1f42a5d6 Mon Sep 17 00:00:00 2001
From: Adam Birds <adam.birds@adbwebdesigns.co.uk>
Date: Sun, 11 Apr 2021 17:59:48 +0000
Subject: [PATCH] installer: Add run_psql_as_postgres function zulip_tools.py.

Add a helper `run_psql_as_postgres` function in
`scripts/lib/zulip_tools.py`.  This is preparatory refactoring for the
work to add custom database and user names.
---
 scripts/lib/upgrade-zulip-stage-2 | 10 +++-------
 scripts/lib/zulip_tools.py        | 20 ++++++++++++++++++++
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/scripts/lib/upgrade-zulip-stage-2 b/scripts/lib/upgrade-zulip-stage-2
index 7969e54852..91903e6bd8 100755
--- a/scripts/lib/upgrade-zulip-stage-2
+++ b/scripts/lib/upgrade-zulip-stage-2
@@ -32,6 +32,7 @@ from scripts.lib.zulip_tools import (
     get_config_file,
     get_tornado_ports,
     parse_os_release,
+    run_psql_as_postgres,
     su_to_zulip,
 )
 
@@ -137,13 +138,8 @@ if os.path.exists("/etc/init.d/postgresql"):
 
 if glob.glob("/usr/share/postgresql/*/extension/tsearch_extras.control"):
     # Remove legacy tsearch_extras package references
-    subprocess.check_call(
-        [
-            "su",
-            "postgres",
-            "-c",
-            'psql -v ON_ERROR_STOP=1 zulip -c "DROP EXTENSION IF EXISTS tsearch_extras;"',
-        ]
+    run_psql_as_postgres(
+        sql_query="DROP EXTENSION IF EXISTS tsearch_extras;",
     )
     subprocess.check_call(["apt-get", "remove", "-y", "postgresql-*-tsearch-extras"])
 
diff --git a/scripts/lib/zulip_tools.py b/scripts/lib/zulip_tools.py
index b096b91b3d..9cb75886b4 100755
--- a/scripts/lib/zulip_tools.py
+++ b/scripts/lib/zulip_tools.py
@@ -580,6 +580,26 @@ def get_deploy_options(config_file: configparser.RawConfigParser) -> List[str]:
     return get_config(config_file, "deployment", "deploy_options", "").strip().split()
 
 
+def run_psql_as_postgres(
+    sql_query: str,
+) -> None:
+    subcmd = " ".join(
+        map(
+            shlex.quote,
+            [
+                "psql",
+                "-v",
+                "ON_ERROR_STOP=1",
+                # TODO: Stop hardcoding the database name.
+                "zulip",
+                "-c",
+                sql_query,
+            ],
+        )
+    )
+    subprocess.check_call(["su", "postgres", "-c", subcmd])
+
+
 def get_tornado_ports(config_file: configparser.RawConfigParser) -> List[int]:
     ports = []
     if config_file.has_section("tornado_sharding"):