Retry if MySQL Server is unreachable (#104)

Ported from
canonical/mysql-router-k8s-operator#190

Author: carlcsaposs-canonical

.gitattributes

@@ -0,0 +1 @@
+*.py.jinja linguist-language=Python

pyproject.toml

@@ -93,8 +93,9 @@ select = ["E", "W", "F", "C", "N", "R", "D", "H"]
 # Ignore D403 First word of the first line should be properly capitalized (false positive on "MySQL")
 # Ignore N818 Exception should be named with an Error suffix
 # Ignore D102 Missing docstring in public method (pydocstyle doesn't look for docstrings in super class
+# Ignore W505 So that strings in comments aren't split across lines
 # https://github.com/PyCQA/pydocstyle/issues/309) TODO: add pylint check? https://github.com/PyCQA/pydocstyle/issues/309#issuecomment-1284142716
-ignore = ["W503", "E501", "D107", "D105", "D415", "D403", "N818", "D102"]
+ignore = ["W503", "E501", "D107", "D105", "D415", "D403", "N818", "D102", "W505"]
 # D100, D101, D102, D103: Ignore missing docstrings in tests
 per-file-ignores = ["tests/*:D100,D101,D102,D103,D104"]
 docstring-convention = "google"

src/abstract_charm.py

@@ -17,6 +17,7 @@
 import machine_upgrade
 import relations.database_provides
 import relations.database_requires
+import server_exceptions
 import upgrade
 import workload
 
@@ -232,37 +233,43 @@ def reconcile(self, event=None) -> None:  # noqa: C901
             f"{self._database_requires.is_relation_breaking(event)=}, "
             f"{self._upgrade.in_progress=}"
         )
-        if self._unit_lifecycle.authorized_leader:
-            if self._database_requires.is_relation_breaking(event):
-                if self._upgrade.in_progress:
-                    logger.warning(
-                        "Modifying relations during an upgrade is not supported. The charm may be in a broken, unrecoverable state. Re-deploy the charm"
+        try:
+            if self._unit_lifecycle.authorized_leader:
+                if self._database_requires.is_relation_breaking(event):
+                    if self._upgrade.in_progress:
+                        logger.warning(
+                            "Modifying relations during an upgrade is not supported. The charm may be in a broken, unrecoverable state. Re-deploy the charm"
+                        )
+                    self._database_provides.delete_all_databags()
+                elif (
+                    not self._upgrade.in_progress
+                    and isinstance(workload_, workload.AuthenticatedWorkload)
+                    and workload_.container_ready
+                ):
+                    self._database_provides.reconcile_users(
+                        event=event,
+                        router_read_write_endpoint=self._read_write_endpoint,
+                        router_read_only_endpoint=self._read_only_endpoint,
+                        shell=workload_.shell,
                     )
-                self._database_provides.delete_all_databags()
-            elif (
-                not self._upgrade.in_progress
-                and isinstance(workload_, workload.AuthenticatedWorkload)
-                and workload_.container_ready
-            ):
-                self._database_provides.reconcile_users(
-                    event=event,
-                    router_read_write_endpoint=self._read_write_endpoint,
-                    router_read_only_endpoint=self._read_only_endpoint,
-                    shell=workload_.shell,
-                )
-        if isinstance(workload_, workload.AuthenticatedWorkload) and workload_.container_ready:
-            workload_.enable(tls=self._tls_certificate_saved, unit_name=self.unit.name)
-        elif workload_.container_ready:
-            workload_.disable()
-        # Empty waiting status means we're waiting for database requires relation before starting
-        # workload
-        if not workload_.status or workload_.status == ops.WaitingStatus():
-            self._upgrade.unit_state = "healthy"
-        if self._unit_lifecycle.authorized_leader:
-            self._upgrade.reconcile_partition()
-            if not self._upgrade.in_progress:
-                self._upgrade.set_versions_in_app_databag()
-        self.set_status(event=event)
+            if isinstance(workload_, workload.AuthenticatedWorkload) and workload_.container_ready:
+                workload_.enable(tls=self._tls_certificate_saved, unit_name=self.unit.name)
+            elif workload_.container_ready:
+                workload_.disable()
+            # Empty waiting status means we're waiting for database requires relation before
+            # starting workload
+            if not workload_.status or workload_.status == ops.WaitingStatus():
+                self._upgrade.unit_state = "healthy"
+            if self._unit_lifecycle.authorized_leader:
+                self._upgrade.reconcile_partition()
+                if not self._upgrade.in_progress:
+                    self._upgrade.set_versions_in_app_databag()
+            self.set_status(event=event)
+        except server_exceptions.Error as e:
+            # If not for `unit=False`, another `server_exceptions.Error` could be thrown here
+            self.set_status(event=event, unit=False)
+            self.unit.status = e.status
+            logger.debug(f"Set unit status to {self.unit.status}")
 
     def _on_resume_upgrade_action(self, event: ops.ActionEvent) -> None:
         if not self._unit_lifecycle.authorized_leader:

src/container.py

@@ -22,6 +22,11 @@ def relative_to_container(self) -> pathlib.PurePosixPath:
         Only differs from `self` on machine charm
         """
 
+    @abc.abstractmethod
+    def open(self, mode="r") -> typing.TextIO:
+        """Open the file in read text mode."""
+        assert mode == "r"
+
     @abc.abstractmethod
     def read_text(self) -> str:
         """Open the file in text mode, read it, and close the file."""

src/mysql_shell.py renamed to src/mysql_shell/__init__.py

@@ -9,11 +9,18 @@
 import dataclasses
 import json
 import logging
+import pathlib
 import secrets
 import string
 import typing
 
+import jinja2
+
 import container
+import server_exceptions
+
+if typing.TYPE_CHECKING:
+    import relations.database_requires
 
 logger = logging.getLogger(__name__)
 
@@ -29,33 +36,56 @@ class RouterUserInformation:
     router_id: str
 
 
+class ShellDBError(Exception):
+    """`mysqlsh.DBError` raised while executing MySQL Shell script
+
+    MySQL Shell runs Python code in a separate process from the charm Python code.
+    The `mysqlsh.DBError` was caught by the shell code, serialized to JSON, and de-serialized to
+    this exception.
+    """
+
+    def __init__(self, *, message: str, code: int, traceback_message: str):
+        super().__init__(message)
+        self.code = code
+        self.traceback_message = traceback_message
+
+
 # TODO python3.10 min version: Add `(kw_only=True)`
 @dataclasses.dataclass
 class Shell:
     """MySQL Shell connected to MySQL cluster"""
 
     _container: container.Container
-    username: str
-    _password: str
-    _host: str
-    _port: str
+    _connection_info: "relations.database_requires.CompleteConnectionInformation"
+
+    @property
+    def username(self):
+        return self._connection_info.username
+
+    def _run_code(self, code: str) -> None:
+        """Connect to MySQL cluster and run Python code."""
+        template = _jinja_env.get_template("try_except_wrapper.py.jinja")
+        error_file = self._container.path("/tmp/mysqlsh_error.json")
+
+        def render(connection_info: "relations.database_requires.ConnectionInformation"):
+            return template.render(
+                username=connection_info.username,
+                password=connection_info.password,
+                host=connection_info.host,
+                port=connection_info.port,
+                code=code,
+                error_filepath=error_file.relative_to_container,
+            )
 
-    # TODO python3.10 min version: Use `list` instead of `typing.List`
-    def _run_commands(self, commands: typing.List[str]) -> str:
-        """Connect to MySQL cluster and run commands."""
         # Redact password from log
-        logged_commands = commands.copy()
-        logged_commands.insert(
-            0, f"shell.connect('{self.username}:***@{self._host}:{self._port}')"
-        )
+        logged_script = render(self._connection_info.redacted)
 
-        commands.insert(
-            0, f"shell.connect('{self.username}:{self._password}@{self._host}:{self._port}')"
-        )
-        temporary_script_file = self._container.path("/tmp/script.py")
-        temporary_script_file.write_text("\n".join(commands))
+        script = render(self._connection_info)
+        temporary_script_file = self._container.path("/tmp/mysqlsh_script.py")
+        error_file = self._container.path("/tmp/mysqlsh_error.json")
+        temporary_script_file.write_text(script)
         try:
-            output = self._container.run_mysql_shell(
+            self._container.run_mysql_shell(
                 [
                     "--no-wizard",
                     "--python",
@@ -64,21 +94,34 @@ def _run_commands(self, commands: typing.List[str]) -> str:
                 ]
             )
         except container.CalledProcessError as e:
-            logger.exception(f"Failed to run {logged_commands=}\nstderr:\n{e.stderr}\n")
+            logger.exception(
+                f"Failed to run MySQL Shell script:\n{logged_script}\n\nstderr:\n{e.stderr}\n"
+            )
             raise
         finally:
             temporary_script_file.unlink()
-        return output
+        with error_file.open("r") as file:
+            exception = json.load(file)
+        error_file.unlink()
+        try:
+            if exception:
+                raise ShellDBError(**exception)
+        except ShellDBError as e:
+            if e.code == 2003:
+                logger.exception(server_exceptions.ConnectionError.MESSAGE)
+                raise server_exceptions.ConnectionError
+            else:
+                logger.exception(
+                    f"Failed to run MySQL Shell script:\n{logged_script}\n\nMySQL client error {e.code}\nMySQL Shell traceback:\n{e.traceback_message}\n"
+                )
+                raise
 
     # TODO python3.10 min version: Use `list` instead of `typing.List`
     def _run_sql(self, sql_statements: typing.List[str]) -> None:
         """Connect to MySQL cluster and execute SQL."""
-        commands = []
-        for statement in sql_statements:
-            # Escape double quote (") characters in statement
-            statement = statement.replace('"', r"\"")
-            commands.append('session.run_sql("' + statement + '")')
-        self._run_commands(commands)
+        self._run_code(
+            _jinja_env.get_template("run_sql.py.jinja").render(statements=sql_statements)
+        )
 
     @staticmethod
     def _generate_password() -> str:
@@ -135,14 +178,17 @@ def get_mysql_router_user_for_unit(
         again.
         """
         logger.debug(f"Getting MySQL Router user for {unit_name=}")
-        rows = json.loads(
-            self._run_commands(
-                [
-                    f"result = session.run_sql(\"SELECT USER, ATTRIBUTE->>'$.router_id' FROM INFORMATION_SCHEMA.USER_ATTRIBUTES WHERE ATTRIBUTE->'$.created_by_user'='{self.username}' AND ATTRIBUTE->'$.created_by_juju_unit'='{unit_name}'\")",
-                    "print(result.fetch_all())",
-                ]
+        output_file = self._container.path("/tmp/mysqlsh_output.json")
+        self._run_code(
+            _jinja_env.get_template("get_mysql_router_user_for_unit.py.jinja").render(
+                username=self.username,
+                unit_name=unit_name,
+                output_filepath=output_file.relative_to_container,
             )
         )
+        with output_file.open("r") as file:
+            rows = json.load(file)
+        output_file.unlink()
         if not rows:
             logger.debug(f"No MySQL Router user found for {unit_name=}")
             return
@@ -159,8 +205,10 @@ def remove_router_from_cluster_metadata(self, router_id: str) -> None:
         metadata already exists for the router ID.
         """
         logger.debug(f"Removing {router_id=} from cluster metadata")
-        self._run_commands(
-            ["cluster = dba.get_cluster()", f'cluster.remove_router_metadata("{router_id}")']
+        self._run_code(
+            _jinja_env.get_template("remove_router_from_cluster_metadata.py.jinja").render(
+                router_id=router_id
+            )
         )
         logger.debug(f"Removed {router_id=} from cluster metadata")
 
@@ -177,12 +225,24 @@ def delete_user(self, username: str, *, must_exist=True) -> None:
     def is_router_in_cluster_set(self, router_id: str) -> bool:
         """Check if MySQL Router is part of InnoDB ClusterSet."""
         logger.debug(f"Checking if {router_id=} in cluster set")
-        output = json.loads(
-            self._run_commands(
-                ["cluster_set = dba.get_cluster_set()", "print(cluster_set.list_routers())"]
+        output_file = self._container.path("/tmp/mysqlsh_output.json")
+        self._run_code(
+            _jinja_env.get_template("get_routers_in_cluster_set.py.jinja").render(
+                output_filepath=output_file.relative_to_container
             )
         )
+        with output_file.open("r") as file:
+            output = json.load(file)
+        output_file.unlink()
         cluster_set_router_ids = output["routers"].keys()
         logger.debug(f"{cluster_set_router_ids=}")
         logger.debug(f"Checked if {router_id in cluster_set_router_ids=}")
         return router_id in cluster_set_router_ids
+
+
+_jinja_env = jinja2.Environment(
+    autoescape=False,
+    trim_blocks=True,
+    loader=jinja2.FileSystemLoader(pathlib.Path(__file__).parent / "templates"),
+    undefined=jinja2.StrictUndefined,
+)

src/mysql_shell/templates/get_mysql_router_user_for_unit.py.jinja

@@ -0,0 +1,12 @@
+import json
+
+result = session.run_sql(
+    "SELECT USER, ATTRIBUTE->>'$.router_id' FROM INFORMATION_SCHEMA.USER_ATTRIBUTES WHERE ATTRIBUTE->'$.created_by_user'='{{ username }}' AND ATTRIBUTE->'$.created_by_juju_unit'='{{ unit_name }}'"
+)
+rows = result.fetch_all()
+# mysqlsh objects are weird—they quack (i.e. duck typing) like standard Python objects (e.g. list,
+# dict), but do not serialize to JSON correctly.
+# Cast to str & load from JSON str before serializing
+rows = json.loads(str(rows))
+with open("{{ output_filepath }}", "w") as file:
+    json.dump(rows, file)

src/mysql_shell/templates/get_routers_in_cluster_set.py.jinja

@@ -0,0 +1,10 @@
+import json
+
+cluster_set = dba.get_cluster_set()
+routers = cluster_set.list_routers()
+# mysqlsh objects are weird—they quack (i.e. duck typing) like standard Python objects (e.g. list,
+# dict), but do not serialize to JSON correctly.
+# Cast to str & load from JSON str before serializing
+routers = json.loads(str(routers))
+with open("{{ output_filepath }}", "w") as file:
+    json.dump(routers, file)

src/mysql_shell/templates/remove_router_from_cluster_metadata.py.jinja

@@ -0,0 +1,2 @@
+cluster = dba.get_cluster()
+cluster.remove_router_metadata("{{ router_id }}")

src/mysql_shell/templates/run_sql.py.jinja

@@ -0,0 +1,3 @@
+{% for statement in statements %}
+session.run_sql("{{ statement|replace('"', '\\"') }}")
+{% endfor %}

src/mysql_shell/templates/try_except_wrapper.py.jinja

@@ -0,0 +1,17 @@
+import json
+import mysqlsh
+import traceback
+
+try:
+    shell.connect("{{ username }}:{{ password }}@{{ host }}:{{ port }}")
+    {{ code|indent(width=4) }}
+except mysqlsh.DBError as exception:
+    error = {
+        "message": str(exception),
+        "code": exception.code,
+        "traceback_message": "".join(traceback.format_exception(exception)),
+    }
+else:
+    error = None
+with open("{{ error_filepath }}", "w") as file:
+    json.dump(error, file)