feat: managed database support for hybrid setup

agent/bench.py

@@ -37,6 +37,7 @@ def __init__(self, name, server, mounts=None):
             self.directory, "sites", "common_site_config.json"
         )
         self.host = self.config.get("db_host", "localhost")
+        self.db_port = self.config.get("db_port", 3306)
         self.docker_image = self.bench_config.get("docker_image")
         self.mounts = mounts
         if not (
@@ -186,20 +187,31 @@ def docker_execute(
         )
 
     @step("New Site")
-    def bench_new_site(self, name, mariadb_root_password, admin_password):
-        site_database, temp_user, temp_password = self.create_mariadb_user(
-            name, mariadb_root_password
-        )
-        try:
-            return self.docker_execute(
-                f"bench new-site --no-mariadb-socket "
-                f"--mariadb-root-username {temp_user} "
-                f"--mariadb-root-password {temp_password} "
-                f"--admin-password {admin_password} "
-                f"--db-name {site_database} {name}"
+    def bench_new_site(self, name, mariadb_root_password, admin_password, managed_database_config=None):
+        if managed_database_config:
+            try:
+                return self.docker_execute(
+                    f"bench new-site --no-mariadb-socket "
+                    f"--mariadb-root-username {managed_database_config.get('database_root_user')} "
+                    f"--mariadb-root-password '{mariadb_root_password}' "
+                    f"--admin-password {admin_password} {name}"
+                )
+            except Exception:
+                traceback.print_exc()
+        else:
+            site_database, temp_user, temp_password = self.create_mariadb_user(
+                name, mariadb_root_password
             )
-        finally:
-            self.drop_mariadb_user(name, mariadb_root_password, site_database)
+            try:
+                return self.docker_execute(
+                    f"bench new-site --no-mariadb-socket "
+                    f"--mariadb-root-username {temp_user} "
+                    f"--mariadb-root-password {temp_password} "
+                    f"--admin-password {admin_password} "
+                    f"--db-name {site_database} {name}"
+                )
+            finally:
+                self.drop_mariadb_user(name, mariadb_root_password, site_database)
 
     @job("Create User", priority="high")
     def create_user(
@@ -379,9 +391,12 @@ def _inactive_web_sites(bench):
 
     @job("New Site", priority="high")
     def new_site(
-        self, name, config, apps, mariadb_root_password, admin_password, create_user: dict = None
+        self, name, config, apps, mariadb_root_password, admin_password, managed_database_config, create_user: dict = None
     ):
-        self.bench_new_site(name, mariadb_root_password, admin_password)
+
+        self.bench_new_site(
+            name, mariadb_root_password, admin_password, managed_database_config=managed_database_config
+        )
         site = Site(name, self)
         site.install_apps(apps)
         site.update_config(config)
@@ -409,18 +424,25 @@ def new_site_from_backup(
         public,
         private,
         skip_failing_patches,
+        managed_database_config=None,
     ):
         files = self.download_files(name, database, public, private)
-        self.bench_new_site(name, mariadb_root_password, admin_password)
+        self.bench_new_site(
+            name, mariadb_root_password, admin_password, managed_database_config=managed_database_config
+        )
         site = Site(name, self)
         site.update_config(default_config)
+        managed_database = 1 if managed_database_config else 0
+
         try:
             site.restore(
                 mariadb_root_password,
                 admin_password,
                 files["database"],
                 files["public"],
                 files["private"],
+                managed_database=managed_database,
+                mariadb_root_user=managed_database_config['database_root_user']
             )
             if site_config:
                 site_config = json.loads(site_config)

agent/site.py

@@ -3,6 +3,7 @@
 import re
 import shutil
 import time
+import traceback
 from datetime import datetime
 from typing import Dict, TYPE_CHECKING
 from shlex import quote
@@ -11,7 +12,7 @@
 
 from agent.base import AgentException, Base
 from agent.job import job, step
-from agent.utils import b2mb, get_size
+from agent.utils import b2mb, get_size, cint
 
 if TYPE_CHECKING:
     from agent.bench import Bench
@@ -46,6 +47,7 @@ def __init__(self, name: str, bench: "Bench"):
         self.user = self.config["db_name"]
         self.password = self.config["db_password"]
         self.host = self.config.get("db_host", self.bench.host)
+        self.port = cint(self.config.get("db_port", self.bench.db_port))
 
     def bench_execute(self, command, input=None):
         return self.bench.docker_execute(
@@ -111,6 +113,8 @@ def restore(
         database_file,
         public_file,
         private_file,
+        managed_database=False,
+        mariadb_root_user=None,
     ):
         sites_directory = self.bench.sites_directory
         database_file = database_file.replace(
@@ -130,23 +134,37 @@ def restore(
             f"--with-private-files {private_file} " if private_file else ""
         )
 
-        _, temp_user, temp_password = self.bench.create_mariadb_user(
-            self.name, mariadb_root_password, self.database
-        )
-        try:
-            return self.bench_execute(
-                "--force restore "
-                f"--mariadb-root-username {temp_user} "
-                f"--mariadb-root-password {temp_password} "
-                f"--admin-password {admin_password} "
-                f"{public_file_option} "
-                f"{private_file_option} "
-                f"{database_file}"
-            )
-        finally:
-            self.bench.drop_mariadb_user(
+        if managed_database:
+            try:
+                return self.bench_execute(
+                    "--force restore "
+                    f"--mariadb-root-username {mariadb_root_user} "
+                    f"--mariadb-root-password '{mariadb_root_password}' "
+                    f"--admin-password {admin_password} "
+                    f"{public_file_option} "
+                    f"{private_file_option} "
+                    f"{database_file}"
+                )
+            except Exception:
+                traceback.print_exc()
+        else:
+            _, temp_user, temp_password = self.bench.create_mariadb_user(
                 self.name, mariadb_root_password, self.database
             )
+            try:
+                return self.bench_execute(
+                    "--force restore "
+                    f"--mariadb-root-username {temp_user} "
+                    f"--mariadb-root-password {temp_password} "
+                    f"--admin-password {admin_password} "
+                    f"{public_file_option} "
+                    f"{private_file_option} "
+                    f"{database_file}"
+                )
+            finally:
+                self.bench.drop_mariadb_user(
+                    self.name, mariadb_root_password, self.database
+            )
 
     @job("Restore Site")
     def restore_job(
@@ -158,15 +176,20 @@ def restore_job(
         public,
         private,
         skip_failing_patches,
+        managed_database_config=None,
     ):
         files = self.bench.download_files(self.name, database, public, private)
+
+        managed_database = 1 if managed_database_config else 0
         try:
             self.restore(
                 mariadb_root_password,
                 admin_password,
                 files["database"],
                 files["public"],
                 files["private"],
+                managed_database=managed_database,
+                mariadb_root_user=managed_database_config['database_root_user'],
             )
         finally:
             self.bench.delete_downloaded_files(files["directory"])
@@ -191,29 +214,43 @@ def reinstall(
         self,
         mariadb_root_password,
         admin_password,
+        managed_database_config=None
     ):
-        _, temp_user, temp_password = self.bench.create_mariadb_user(
-            self.name, mariadb_root_password, self.database
-        )
-        try:
-            return self.bench_execute(
-                f"reinstall --yes "
-                f"--mariadb-root-username {temp_user} "
-                f"--mariadb-root-password {temp_password} "
-                f"--admin-password {admin_password}"
-            )
-        finally:
-            self.bench.drop_mariadb_user(
+        if managed_database_config:
+            try:
+                mariadb_root_user = managed_database_config["database_root_user"]
+                return self.bench_execute(
+                    f"reinstall --yes "
+                    f"--mariadb-root-username {mariadb_root_user} "
+                    f"--mariadb-root-password '{mariadb_root_password}' "
+                    f"--admin-password {admin_password}"
+                )
+            except Exception:
+                traceback.print_exc()
+        else:
+            _, temp_user, temp_password = self.bench.create_mariadb_user(
                 self.name, mariadb_root_password, self.database
             )
+            try:
+                return self.bench_execute(
+                    f"reinstall --yes "
+                    f"--mariadb-root-username {temp_user} "
+                    f"--mariadb-root-password {temp_password} "
+                    f"--admin-password {admin_password}"
+                )
+            finally:
+                self.bench.drop_mariadb_user(
+                    self.name, mariadb_root_password, self.database
+                )
 
     @job("Reinstall Site")
     def reinstall_job(
         self,
         mariadb_root_password,
         admin_password,
+        managed_database_config=None
     ):
-        return self.reinstall(mariadb_root_password, admin_password)
+        return self.reinstall(mariadb_root_password, admin_password, managed_database_config)
 
     @job("Install App on Site")
     def install_app_job(self, app):
@@ -272,7 +309,7 @@ def create_database_access_credentials(self, mode, mariadb_root_password):
         ]
         for query in queries:
             command = (
-                f"mysql -h {self.host} -uroot -p{mariadb_root_password}"
+                f"mysql -h {self.host} -uroot -p{mariadb_root_password} -P {self.port}"
                 f' -e "{query}"'
             )
             self.execute(command)
@@ -288,7 +325,7 @@ def revoke_database_access_credentials(self, user, mariadb_root_password):
         ]
         for query in queries:
             command = (
-                f"mysql -h {self.host} -uroot -p{mariadb_root_password}"
+                f"mysql -h {self.host} -uroot -p{mariadb_root_password} -P {self.port}"
                 f' -e "{query}"'
             )
             self.execute(command)
@@ -318,7 +355,7 @@ def restore_site_tables(self):
             output = self.execute(
                 "set -o pipefail && "
                 f"gunzip -c '{backup_file_path}' | "
-                f"mysql -h {self.host} -u {self.user} -p{self.password} "
+                f"mysql -h {self.host} -u {self.user} -p{self.password} -P {self.port} "
                 f"{self.database}",
                 executable="/bin/bash",
             )
@@ -468,6 +505,8 @@ def tablewise_backup(self):
             sort_keys=True,
         )
         data = {"tables": {}}
+        db_port = self.config.get("db_port", 3306)
+
         for table in tables:
             backup_file = os.path.join(
                 self.backup_directory, f"{table}.sql.gz"
@@ -476,7 +515,7 @@ def tablewise_backup(self):
                 "set -o pipefail && "
                 "mysqldump --single-transaction --quick --lock-tables=false "
                 f"-h {self.host} -u {self.user} -p{self.password} "
-                f"{self.database} '{table}' "
+                f"{self.database} -P {self.port} '{table}' "
                 f" | gzip > '{backup_file}'",
                 executable="/bin/bash",
             )
@@ -544,7 +583,7 @@ def restore_touched_tables(self):
                 output = self.execute(
                     "set -o pipefail && "
                     f"gunzip -c '{backup_file}' | "
-                    f"mysql -h {self.host} -u {self.user} -p{self.password} "
+                    f"mysql -h {self.host} -u {self.user} -p{self.password} -P {self.port} "
                     f"{self.database}",
                     executable="/bin/bash",
                 )
@@ -559,7 +598,7 @@ def drop_new_tables(self):
         data = {"dropped": {}}
         for table in new_tables:
             output = self.execute(
-                f"mysql -h {self.host} -u {self.user} -p{self.password} "
+                f"mysql -h {self.host} -u {self.user} -p{self.password} -P {self.port} "
                 f"{self.database} -e 'DROP TABLE `{table}`'"
             )
             data["dropped"][table] = output
@@ -641,7 +680,7 @@ def timezone(self):
         )
         try:
             timezone = self.execute(
-                f"mysql -h {self.host} -u{self.database} -p{self.password} "
+                f"mysql -h {self.host} -u{self.database} -p{self.password} -P {self.port}"
                 f'--connect-timeout 3 -sN -e "{query}"'
             )["output"].strip()
         except Exception:
@@ -652,7 +691,7 @@ def timezone(self):
     def tables(self):
         return self.execute(
             "mysql --disable-column-names -B -e 'SHOW TABLES' "
-            f"-h {self.host} -u {self.user} -p{self.password} {self.database}"
+            f"-h {self.host} -u {self.user} -p{self.password} -P {self.port} {self.database}"
         )["output"].split("\n")
 
     @property
@@ -685,7 +724,7 @@ def optimize_tables(self):
         for table in tables:
             query = f"OPTIMIZE TABLE `{table}`"
             self.execute(
-                f"mysql -sN -h {self.host} -u{self.user} -p{self.password}"
+                f"mysql -sN -h {self.host} -u{self.user} -p{self.password} -P {self.port}"
                 f" {self.database} -e '{query}'"
             )
 
@@ -761,7 +800,7 @@ def get_database_size(self):
             " GROUP BY `table_schema`"
         )
         command = (
-            f"mysql -sN -h {self.host} -u{self.user} -p{self.password}"
+            f"mysql -sN -h {self.host} -u{self.user} -p{self.password} -P {self.port}"
             f" -e '{query}'"
         )
         database_size = self.execute(command).get("output")
@@ -810,7 +849,7 @@ def get_database_free_size(self):
             " GROUP BY `table_schema`"
         )
         command = (
-            f"mysql -sN -h {self.host} -u{self.user} -p{self.password}"
+            f"mysql -sN -h {self.host} -u{self.user} -p{self.password} -P {self.port}"
             f" -e '{query}'"
         )
         database_size = self.execute(command).get("output")
@@ -831,7 +870,7 @@ def get_database_free_tables(self):
                 " OR `data_free` > 100 * 1024 * 1024)"
             )
             command = (
-                f"mysql -sN -h {self.host} -u{self.user} -p{self.password}"
+                f"mysql -sN -h {self.host} -u{self.user} -p{self.password} -P {self.port}"
                 f" -e '{query}'"
             )
             output = self.execute(command).get("output")