[DPE-5373] Create backup action (#104)

* feat: Port VM changes

* feat: Use provided path to save snapshots

* chore: Remove local dev files

actions.yaml

@@ -35,4 +35,4 @@ resume-upgrade:
   description: Resume a rolling upgrade after asserting successful upgrade of a new revision.
 
 create-backup:
-  description: TODO. This action is only used for testing at the moment.
+  description: Create a database backup and send it to an object storage. S3 credentials are retrieved from a relation with the S3 integrator charm.

metadata.yaml

@@ -28,7 +28,7 @@ resources:
   zookeeper-image:
     type: oci-image
     description: OCI Image for Apache ZooKeeper
-    upstream-source: ghcr.io/canonical/charmed-zookeeper@sha256:07a8ec3abc0aa73c2824437bc9f792cde1d0e2da5e312606c97c21def34d1a6d
+    upstream-source: ghcr.io/canonical/charmed-zookeeper@sha256:e37016b76aef497d1b4de2332af479592e9a2e8dcabc296e6c7876e0334906ad
 
 peers:
   cluster:

pyproject.toml

@@ -46,8 +46,9 @@ jsonschema = ">=4.10"
 lightkube = "^0.15.0"
 typing-extensions = "^4.9.0"
 boto3 = "^1.34.159"
-boto3-stubs = { extras = ["s3"], version = "^1.35.8" }
+boto3-stubs = {extras = ["s3"], version = "^1.35.8"}
 httpx = "^0.27.2"
+rich = "^13.8.1"
 
 [tool.poetry.group.fmt]
 optional = true
@@ -73,7 +74,7 @@ pytest = ">=7.2"
 coverage = { extras = ["toml"], version = ">7.0" }
 jsonschema = ">=4.10"
 pytest-mock = "^3.11.1"
-ops-scenario = "^6.0.0"
+ops-scenario = "^7.0.0"
 
 [tool.poetry.group.integration]
 optional = true

requirements.txt

@@ -1,13 +1,13 @@
 anyio==4.4.0 ; python_version >= "3.10" and python_version < "4.0"
 attrs==24.2.0 ; python_version >= "3.10" and python_version < "4.0"
-boto3-stubs[s3]==1.35.8 ; python_version >= "3.10" and python_version < "4.0"
-boto3==1.35.8 ; python_version >= "3.10" and python_version < "4.0"
-botocore-stubs==1.35.8 ; python_version >= "3.10" and python_version < "4.0"
-botocore==1.35.8 ; python_version >= "3.10" and python_version < "4.0"
-certifi==2024.7.4 ; python_version >= "3.10" and python_version < "4.0"
-cffi==1.17.0 ; python_version >= "3.10" and python_version < "4.0" and platform_python_implementation != "PyPy"
-cosl==0.0.24 ; python_version >= "3.10" and python_version < "4.0"
-cryptography==43.0.0 ; python_version >= "3.10" and python_version < "4.0"
+boto3-stubs[s3]==1.35.17 ; python_version >= "3.10" and python_version < "4.0"
+boto3==1.35.17 ; python_version >= "3.10" and python_version < "4.0"
+botocore-stubs==1.35.17 ; python_version >= "3.10" and python_version < "4.0"
+botocore==1.35.17 ; python_version >= "3.10" and python_version < "4.0"
+certifi==2024.8.30 ; python_version >= "3.10" and python_version < "4.0"
+cffi==1.17.1 ; python_version >= "3.10" and python_version < "4.0" and platform_python_implementation != "PyPy"
+cosl==0.0.32 ; python_version >= "3.10" and python_version < "4.0"
+cryptography==43.0.1 ; python_version >= "3.10" and python_version < "4.0"
 exceptiongroup==1.2.2 ; python_version >= "3.10" and python_version < "3.11"
 h11==0.14.0 ; python_version >= "3.10" and python_version < "4.0"
 httpcore==1.0.5 ; python_version >= "3.10" and python_version < "4.0"
@@ -19,20 +19,24 @@ jsonschema==4.23.0 ; python_version >= "3.10" and python_version < "4.0"
 kazoo==2.9.0 ; python_version >= "3.10" and python_version < "4.0"
 lightkube-models==1.30.0.8 ; python_version >= "3.10" and python_version < "4.0"
 lightkube==0.15.4 ; python_version >= "3.10" and python_version < "4.0"
-mypy-boto3-s3==1.35.2 ; python_version >= "3.10" and python_version < "4.0"
-ops==2.16.0 ; python_version >= "3.10" and python_version < "4.0"
+markdown-it-py==3.0.0 ; python_version >= "3.10" and python_version < "4.0"
+mdurl==0.1.2 ; python_version >= "3.10" and python_version < "4.0"
+mypy-boto3-s3==1.35.16 ; python_version >= "3.10" and python_version < "4.0"
+ops==2.16.1 ; python_version >= "3.10" and python_version < "4.0"
 pure-sasl==0.6.2 ; python_version >= "3.10" and python_version < "4.0"
 pycparser==2.22 ; python_version >= "3.10" and python_version < "4.0" and platform_python_implementation != "PyPy"
 pydantic==1.10.18 ; python_version >= "3.10" and python_version < "4.0"
+pygments==2.18.0 ; python_version >= "3.10" and python_version < "4.0"
 python-dateutil==2.9.0.post0 ; python_version >= "3.10" and python_version < "4.0"
 pyyaml==6.0.2 ; python_version >= "3.10" and python_version < "4.0"
 referencing==0.35.1 ; python_version >= "3.10" and python_version < "4.0"
+rich==13.8.1 ; python_version >= "3.10" and python_version < "4.0"
 rpds-py==0.18.1 ; python_version >= "3.10" and python_version < "4.0"
 s3transfer==0.10.2 ; python_version >= "3.10" and python_version < "4.0"
 six==1.16.0 ; python_version >= "3.10" and python_version < "4.0"
 sniffio==1.3.1 ; python_version >= "3.10" and python_version < "4.0"
 tenacity==9.0.0 ; python_version >= "3.10" and python_version < "4.0"
-types-awscrt==0.21.2 ; python_version >= "3.10" and python_version < "4.0"
+types-awscrt==0.21.5 ; python_version >= "3.10" and python_version < "4.0"
 types-s3transfer==0.10.2 ; python_version >= "3.10" and python_version < "4.0"
 typing-extensions==4.12.2 ; python_version >= "3.10" and python_version < "4.0"
 urllib3==2.2.2 ; python_version >= "3.10" and python_version < "4.0"

src/core/stubs.py

@@ -16,3 +16,6 @@
         "region": str,
     },
 )
+
+
+BackupMetadata = TypedDict("BackupMetadata", {"id": str, "log-sequence-number": int, "path": str})

src/events/backup.py

@@ -32,7 +32,7 @@ def __init__(self, charm):
         super().__init__(charm, "backup")
         self.charm: "ZooKeeperCharm" = charm
         self.s3_requirer = S3Requirer(self.charm, S3_REL_NAME)
-        self.backup_manager = BackupManager(self.charm.state.cluster.s3_credentials)
+        self.backup_manager = BackupManager(self.charm.state)
 
         self.framework.observe(
             self.s3_requirer.on.credentials_changed, self._on_s3_credentials_changed
@@ -90,7 +90,6 @@ def _on_s3_credentials_gone(self, event: CredentialsGoneEvent):
         self.charm.state.cluster.update({"s3-credentials": ""})
 
     def _on_create_backup_action(self, event: ActionEvent):
-        # TODO
         failure_conditions = [
             (not self.charm.unit.is_leader(), "Action must be ran on the application leader"),
             (
@@ -110,7 +109,12 @@ def _on_create_backup_action(self, event: ActionEvent):
                 event.fail(msg)
                 return
 
-        self.backup_manager.write_test_string()
+        backup_metadata = self.backup_manager.create_backup()
+
+        output = self.backup_manager.format_backups_table(
+            [backup_metadata], title="Backup created"
+        )
+        event.log(output)
 
     def _on_list_backups_action(self, _):
         # TODO

src/literals.py

@@ -43,7 +43,7 @@
         "dependencies": {},
         "name": "zookeeper",
         "upgrade_supported": "^3.6",
-        "version": "3.8.4",
+        "version": "3.9.2",
     },
 }
 

src/managers/backup.py

@@ -5,34 +5,46 @@
 """Helpers for managing backups."""
 
 import logging
+import os
+from datetime import datetime
+from io import StringIO
 
 import boto3
+import httpx
+import yaml
 from botocore import loaders, regions
 from botocore.exceptions import ClientError
 from mypy_boto3_s3.service_resource import Bucket
+from rich.console import Console
+from rich.table import Table
 
-from core.stubs import S3ConnectionInfo
+from core.cluster import ClusterState
+from core.stubs import BackupMetadata, S3ConnectionInfo
+from literals import ADMIN_SERVER_PORT, S3_BACKUPS_PATH
 
 logger = logging.getLogger(__name__)
 
 
 class BackupManager:
     """Manager for all things backup-related."""
 
-    def __init__(self, s3_parameters: S3ConnectionInfo) -> None:
-        self.s3_parameters = s3_parameters
+    def __init__(self, state: ClusterState) -> None:
+        self.state = state
+        self.backups_path = S3_BACKUPS_PATH
 
     @property
     def bucket(self) -> Bucket:
         """S3 bucket to read from and write to."""
+        s3_parameters = self.state.cluster.s3_credentials
+        self.backups_path = s3_parameters["path"]
         s3 = boto3.resource(
             "s3",
-            aws_access_key_id=self.s3_parameters["access-key"],
-            aws_secret_access_key=self.s3_parameters["secret-key"],
-            region_name=self.s3_parameters["region"] if self.s3_parameters["region"] else None,
-            endpoint_url=self._construct_endpoint(self.s3_parameters),
+            aws_access_key_id=s3_parameters["access-key"],
+            aws_secret_access_key=s3_parameters["secret-key"],
+            region_name=s3_parameters["region"] if s3_parameters["region"] else None,
+            endpoint_url=self._construct_endpoint(s3_parameters),
         )
-        return s3.Bucket(self.s3_parameters["bucket"])
+        return s3.Bucket(s3_parameters["bucket"])
 
     def _construct_endpoint(self, s3_parameters: S3ConnectionInfo) -> str:
         """Construct the S3 service endpoint using the region.
@@ -84,6 +96,92 @@ def create_bucket(self, s3_parameters: S3ConnectionInfo) -> bool:
 
         return True
 
-    def write_test_string(self) -> None:
-        """Write content in the object storage."""
-        self.bucket.put_object(Key="test_file.txt", Body=b"test string")
+    def create_backup(self) -> BackupMetadata:
+        """Create a snapshot with ZooKeeper admin server and stream it to the object storage."""
+        zk_user = "super"
+        zk_pwd = self.state.cluster.internal_user_credentials.get("super", "")
+        date = datetime.now()
+        snapshot_name = f"{date:%Y-%m-%dT%H:%M:%SZ}"
+        snapshot_path = f"{snapshot_name}/snapshot"
+
+        # It is very likely that the file is fully loaded in memory, because the file-like interface is
+        # not seekable, and I have a strong suspicion that boto uses this to figure out if it can
+        # upload in one go or need to use a multipart request.
+        # We cannot be sure because finding this information in boto code base is time consuming.
+        # If this ever become an issue, we can find a workaround by using the 'content-length' header from
+        # the response. Or write to a temp file as a last resort.
+        with httpx.stream(
+            "GET",
+            f"http://localhost:{ADMIN_SERVER_PORT}/commands/snapshot?streaming=true",
+            headers={"Authorization": f"digest {zk_user}:{zk_pwd}"},
+        ) as response:
+
+            response_headers = response.headers
+            quorum_leader_zxid = int(response_headers["last_zxid"], base=16)
+            metadata: BackupMetadata = {
+                "id": snapshot_name,
+                "log-sequence-number": quorum_leader_zxid,
+                "path": snapshot_path,
+            }
+
+            self.bucket.put_object(
+                Key=os.path.join(self.backups_path, snapshot_name, "metadata.yaml"),
+                Body=yaml.dump(metadata, encoding="utf8"),
+            )
+
+            self.bucket.upload_fileobj(
+                _StreamingToFileSyncAdapter(response),  # type: ignore
+                os.path.join(self.backups_path, snapshot_name, "snapshot"),
+            )
+
+        return metadata
+
+    def format_backups_table(
+        self, backup_entries: list[BackupMetadata], title: str = "Backups"
+    ) -> str:
+        """Format backups metadata into a readable table."""
+        table = Table(title=title)
+
+        table.add_column("Id", no_wrap=True)
+        table.add_column("Log-sequence-number", justify="right")
+        table.add_column("path")
+
+        for meta in backup_entries:
+            table.add_row(meta["id"], str(meta["log-sequence-number"]), meta["path"])
+
+        out_f = StringIO()
+        console = Console(file=out_f, width=79)
+        console.print(table)
+
+        return out_f.getvalue()
+
+
+class _StreamingToFileSyncAdapter:
+    """Wrapper to make httpx.stream behave like a file-like object.
+
+    boto needs a .read method with an optional amount-of-bytes parameter from the file-like object.
+    Taken from https://github.com/encode/httpx/discussions/2296#discussioncomment-6781355
+    """
+
+    def __init__(self, response: httpx.Response):
+        self.streaming_source = response.iter_bytes()
+        self.buffer = b""
+        self.buffer_offset = 0
+
+    def read(self, num_bytes: int = 4096) -> bytes:
+        while len(self.buffer) - self.buffer_offset < num_bytes:
+            try:
+                chunk = next(self.streaming_source)
+                self.buffer += chunk
+            except StopIteration:
+                break
+
+        if len(self.buffer) - self.buffer_offset >= num_bytes:
+            data = self.buffer[self.buffer_offset : self.buffer_offset + num_bytes]
+            self.buffer_offset += num_bytes
+            return data
+        else:
+            data = self.buffer[self.buffer_offset :]
+            self.buffer = b""
+            self.buffer_offset = 0
+            return data

src/managers/config.py

@@ -32,6 +32,7 @@
 enforce.auth.schemes=sasl
 sessionRequireClientSASLAuth=true
 audit.enable=true
+admin.serverAddress=localhost
 """
 
 TLS_PROPERTIES = """

tests/integration/test_backup.py

@@ -103,8 +103,8 @@ async def test_relate_active_bucket_created(ops_test: OpsTest, s3_bucket):
     assert s3_bucket.meta.client.head_bucket(Bucket=s3_bucket.name)
 
 
-@pytest.mark.abort_on_fail
-async def test_write_content(ops_test: OpsTest, s3_bucket: Bucket):
+# @pytest.mark.abort_on_fail
+async def write_content(ops_test: OpsTest, s3_bucket: Bucket):
     # TODO (backup): Remove and replace with ZK snapshot write
 
     for unit in ops_test.model.applications[APP_NAME].units: