diff --git a/README.md b/README.md index ac01ece..3a9afc4 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ Install `django-pgclone` with: After this, add `pgclone` to the `INSTALLED_APPS` setting of your Django project. -**Note** Install the AWS CLI to enable the S3 storage backend. Use `pip install awscli` or follow the [installation guide here](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html). +**Note** Install the [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) to enable the S3 storage backend. Alternatively, install the optional S3 extra (`pip install django-pgclone[s3]`) and set `PGCLONE_S3_BACKEND = "boto3"`. ## Contributing Guide diff --git a/docs/installation.md b/docs/installation.md index 8721e52..a599491 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -8,4 +8,4 @@ After this, add `pgclone` to the `INSTALLED_APPS` setting of your Django project !!! note - Install the AWS CLI to enable the S3 storage backend. Use `pip install awscli` or follow the [installation guide here](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html). + Install the [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) to enable the S3 storage backend. Alternatively, install the optional S3 extra (`pip install django-pgclone[s3]`) and set `PGCLONE_S3_BACKEND = "boto3"`. diff --git a/docs/settings.md b/docs/settings.md index eee93ed..e7218bc 100644 --- a/docs/settings.md +++ b/docs/settings.md @@ -80,7 +80,7 @@ The hooks to run by default for restores before swapping happens. ## PGCLONE_S3_CONFIG -The environment variable overrides when using the AWS CLI. Only applicable when using the S3 storage backend. +The AWS credentials and region configuration for the S3 storage backend. Applies to both the boto3 and AWS CLI backends. For example: @@ -95,6 +95,20 @@ PGCLONE_S3_CONFIG = { **Default**: `{}` +## PGCLONE_S3_BACKEND + +The S3 backend to use. Must be `"boto3"` or `"awscli"`. + +When unset, the AWS CLI backend is used. + +For example: + +```python +PGCLONE_S3_BACKEND = "boto3" +``` + +**Default**: `"awscli"` + ## PGCLONE_S3_ENDPOINT_URL The S3 endpoint url to send requests to if using a non-standard AWS endpoint or an S3 service other than AWS (such as DigitalOcean Spaces or self-hosting an endpoint directly within your private VPC). Only applicable when using the S3 storage backend. diff --git a/docs/storage.md b/docs/storage.md index 0804a3b..235d958 100644 --- a/docs/storage.md +++ b/docs/storage.md @@ -8,15 +8,27 @@ S3 storage is enabled by configuring a path that starts with `s3://`. A bucket a When using S3, dumps and restores are streamed, reducing the memory consumption required for large databases. -In order to use the S3 storage backend, one must additionally install the [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html). Earlier versions of the CLI can be installed with `pip install awscli`. +### AWS CLI (default) + +Install the [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html). Earlier versions of the CLI can be installed with `pip install awscli`. + +This is the default backend and requires no additional configuration. !!! warning Installing the AWS CLI with pip can cause dependency issues in projects that depend on later versions of colorama or docutils. It is recommended to manually install the CLI using the [installation instructions](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) if the pip version doesn't work. +### boto3 (alternative) + +Install the optional S3 extra: + + pip install django-pgclone[s3] + +This installs `boto3`. No external binaries are required. Set `settings.PGCLONE_S3_BACKEND = "boto3"` to use it. + ## Configuring the S3 backend -The AWS CLI can be configured by environment variables. Inject custom environment variables by configuring `settings.PGCLONE_S3_CONFIG`. Here we override the AWS credentials and region: +S3 credentials and region can be configured with `settings.PGCLONE_S3_CONFIG`. Here we override the AWS credentials and region: ```python @@ -29,10 +41,20 @@ PGCLONE_S3_CONFIG = { } ``` -See [this guide](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html) for all environment variables that can be used with the AWS CLI. +When using the boto3 backend, unset keys fall through to boto3's default credential chain (environment variables, IAM roles, profiles, etc.). When using the AWS CLI backend, these values are passed as environment variables to the `aws` subprocess. -If using a non-standard AWS endpoint url or a non-AWS S3 provider, the endpoint url must be specified. Unfortunately, AWS CLI does not provide an environmental variable for this purpose. Instead, use the `settings.PGCLONE_S3_ENDPOINT_URL` setting, which will override the AWS CLI commands with the `--endpoint-url` option. `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` must still be specified in `settings.PGCLONE_S3_CONFIG`. +See [this guide](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html) for all environment variables that can be used with S3. + +If using a non-standard AWS endpoint url or a non-AWS S3 provider, the endpoint url must be specified with `settings.PGCLONE_S3_ENDPOINT_URL`. `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` must still be specified in `settings.PGCLONE_S3_CONFIG` when not using IAM roles or other automatic credential providers. ```python PGCLONE_S3_ENDPOINT_URL = "https://endpoint.example.com" ``` + +## S3 backend selection + +By default, `django-pgclone` uses the AWS CLI backend. Override this behavior with `settings.PGCLONE_S3_BACKEND`: + +```python +PGCLONE_S3_BACKEND = "boto3" # or "awscli" +``` diff --git a/pgclone/dump_cmd.py b/pgclone/dump_cmd.py index 7fcaa93..5d0b225 100644 --- a/pgclone/dump_cmd.py +++ b/pgclone/dump_cmd.py @@ -53,13 +53,12 @@ def _dump( # Note - do note format {db_dump_url} with an `f` string. # It will be formatted later when running the command. pg_dump_cmd_fmt = "pg_dump -Fc --no-acl --no-owner {db_dump_url} " + exclude_args - pg_dump_cmd_fmt += " " + storage_client.pg_dump(file_path) anon_pg_dump_cmd = pg_dump_cmd_fmt.format(db_dump_url="") logging.success_msg(f"Creating DB copy with cmd: {anon_pg_dump_cmd}") pg_dump_cmd = pg_dump_cmd_fmt.format(db_dump_url=db.url(dump_db)) - run.shell(pg_dump_cmd, env=storage_client.env, pipefail=True) + storage_client.run_pg_dump(pg_dump_cmd, file_path) logging.success_msg(f'Database "{database}" successfully dumped to "{dump_key}"') diff --git a/pgclone/restore_cmd.py b/pgclone/restore_cmd.py index e431850..6d19701 100644 --- a/pgclone/restore_cmd.py +++ b/pgclone/restore_cmd.py @@ -110,13 +110,12 @@ def _remote_restore( logging.success_msg(f'Running pg_restore on "{dump_key}"') pg_restore_cmd = f"pg_restore --verbose --no-acl --no-owner -d {db.url(temp_db)}" - pg_restore_cmd = storage_client.pg_restore(file_path) + " " + pg_restore_cmd # When restoring, we need to ignore errors because there are certain # errors we cannot get around when pg restoring some DBs (like Aurora). # In the future, we may parse the output of the pg_restore command to see # if an unexpected error happened. - run.shell(pg_restore_cmd, env=storage_client.env, ignore_errors=True) + storage_client.run_pg_restore(pg_restore_cmd, file_path) return dump_key diff --git a/pgclone/settings.py b/pgclone/settings.py index c56ae49..e6d5e61 100644 --- a/pgclone/settings.py +++ b/pgclone/settings.py @@ -26,6 +26,18 @@ def s3_endpoint_url() -> str | None: return getattr(settings, "PGCLONE_S3_ENDPOINT_URL", None) +def s3_backend() -> str: + backend = getattr(settings, "PGCLONE_S3_BACKEND", None) + if backend is not None: + if backend not in ("boto3", "awscli"): + raise exceptions.RuntimeError( + 'Invalid PGCLONE_S3_BACKEND setting. Must be "boto3" or "awscli".' + ) + return backend + + return "awscli" + + def storage_location() -> str: location = getattr(settings, "PGCLONE_STORAGE_LOCATION", ".pgclone") if not location.endswith("/"): # pragma: no cover diff --git a/pgclone/storage.py b/pgclone/storage.py index 1b81e5a..36601d4 100644 --- a/pgclone/storage.py +++ b/pgclone/storage.py @@ -1,27 +1,66 @@ from __future__ import annotations import abc +import importlib.util import os import pathlib import subprocess +import threading from typing import Any -from pgclone import exceptions, settings +from pgclone import exceptions, logging, run, settings +S3_BACKEND_BOTO3 = "boto3" +S3_BACKEND_AWSCLI = "awscli" -def validate_s3_support() -> None: # pragma: no cover - """Verify that pgclone has been installed with the S3 extras""" + +def _is_boto3_importable() -> bool: + return importlib.util.find_spec("boto3") is not None + + +def _is_awscli_available() -> bool: which_aws = subprocess.run( "which aws", shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE ) - if which_aws.returncode != 0: + return which_aws.returncode == 0 + + +def validate_s3_support(backend: str) -> None: # pragma: no cover + """Verify that the configured S3 backend is available.""" + if backend == S3_BACKEND_BOTO3: + if not _is_boto3_importable(): + raise exceptions.RuntimeError( + "You must install boto3 to use the boto3 S3 backend." + ' Run "pip install django-pgclone[s3]".' + ) + elif backend == S3_BACKEND_AWSCLI: + if not _is_awscli_available(): + raise exceptions.RuntimeError( + "You must install the AWS command line tool in order to enable S3 support." + ' Run "pip install awscli" or follow these instructions -' + " https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html" + ) + else: raise exceptions.RuntimeError( - "You must install the AWS command line tool in order to enable S3 support." - ' Run "pip install awscli" or follow these instructions -' - " https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html" + 'Invalid PGCLONE_S3_BACKEND setting. Must be "boto3" or "awscli".' ) +def _parse_s3_path(s3_path: str) -> tuple[str, str]: + if not s3_path.startswith("s3://"): + raise ValueError(f'Invalid S3 path "{s3_path}".') + + path = s3_path[5:] + bucket, _, key = path.partition("/") + return bucket, key + + +def _log_stream(stream: Any) -> None: + logger = logging.get_logger() + for line in iter(stream.readline, b""): + logger.info(line.decode("utf-8").rstrip()) + + class _Storage(abc.ABC): def __init__(self, storage_location: str) -> None: # Ensure the storage location always has a slash appended @@ -48,15 +87,23 @@ def pg_restore(self, file_path: str) -> str: """Given a file path, generates the CLI fragment to prepend to pg_restore""" pass + def run_pg_dump(self, pg_dump_cmd: str, file_path: str) -> None: + cmd = pg_dump_cmd + " " + self.pg_dump(file_path) + run.shell(cmd, env=self.env, pipefail=True) + + def run_pg_restore(self, pg_restore_cmd: str, file_path: str) -> None: + cmd = self.pg_restore(file_path) + " " + pg_restore_cmd + run.shell(cmd, env=self.env, ignore_errors=True) + @abc.abstractmethod def ls(self, prefix: str | None = None) -> list[str]: """Given a prefix, returns a list of dump keys""" pass -class S3(_Storage): +class S3Awscli(_Storage): def __init__(self, storage_location: str): - validate_s3_support() + validate_s3_support(S3_BACKEND_AWSCLI) self.s3_endpoint_url = ( f" --endpoint-url {settings.s3_endpoint_url()}" if settings.s3_endpoint_url() is not None @@ -80,7 +127,17 @@ def ls(self, prefix: str | None = None) -> list[str]: # pragma: no cover return [self.dump_key(path) for path in abs_paths] def get_env(self): - return settings.s3_config() + # Since AWS CLI v2.23, uploads default to sending an additional CRC32 + # integrity checksum using "aws-chunked" streaming trailers. Many + # S3-compatible providers reject this with a "XAmzContentSHA256Mismatch" + # error, so restore the pre-2.23 behavior of only checksumming when + # required. This mirrors the boto3 backend's client config and is fully + # supported by AWS S3 as well. + return { + "AWS_REQUEST_CHECKSUM_CALCULATION": "when_required", + "AWS_RESPONSE_CHECKSUM_VALIDATION": "when_required", + **settings.s3_config(), + } def pg_dump(self, file_path: str) -> str: return f"| aws s3 cp - {file_path}{self.s3_endpoint_url}" @@ -89,6 +146,146 @@ def pg_restore(self, file_path: str) -> str: return f"aws s3 cp {file_path} -{self.s3_endpoint_url} |" +S3 = S3Awscli + + +class S3Boto3(_Storage): + def __init__(self, storage_location: str): + validate_s3_support(S3_BACKEND_BOTO3) + self._s3_client: Any = None + super().__init__(storage_location) + + @property + def s3_client(self) -> Any: + if self._s3_client is None: + from boto3.session import Session + + session_kwargs = self._boto3_session_kwargs() + client_kwargs = self._boto3_client_kwargs() + session = Session(**session_kwargs) + self._s3_client = session.client("s3", **client_kwargs) + return self._s3_client + + def _boto3_session_kwargs(self) -> dict[str, Any]: + config = settings.s3_config() + kwargs: dict[str, Any] = {} + key_mapping = { + "AWS_ACCESS_KEY_ID": "aws_access_key_id", + "AWS_SECRET_ACCESS_KEY": "aws_secret_access_key", + "AWS_SESSION_TOKEN": "aws_session_token", + "AWS_DEFAULT_REGION": "region_name", + } + for env_key, boto_key in key_mapping.items(): + if env_key in config and config[env_key] is not None: + kwargs[boto_key] = config[env_key] + return kwargs + + def _boto3_client_kwargs(self) -> dict[str, Any]: + from botocore.config import Config + + kwargs: dict[str, Any] = {} + endpoint_url = settings.s3_endpoint_url() + if endpoint_url is not None and isinstance(endpoint_url, str): + kwargs["endpoint_url"] = endpoint_url + + # Since botocore 1.36, uploads default to sending an additional CRC32 + # integrity checksum using "aws-chunked" streaming trailers. Many + # S3-compatible providers reject this with a "XAmzContentSHA256Mismatch" + # error, so restore the pre-1.36 behavior of only checksumming when + # required. This is fully supported by AWS S3 as well. + kwargs["config"] = Config( + request_checksum_calculation="when_required", + response_checksum_validation="when_required", + ) + return kwargs + + def pg_dump(self, file_path: str) -> str: + raise NotImplementedError + + def pg_restore(self, file_path: str) -> str: + raise NotImplementedError + + def run_pg_dump(self, pg_dump_cmd: str, file_path: str) -> None: + bucket, key = _parse_s3_path(file_path) + process = subprocess.Popen( + pg_dump_cmd, + shell=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + if process.stdout is None or process.stderr is None: + raise AssertionError + + stderr_thread = threading.Thread( + target=_log_stream, + args=(process.stderr,), + daemon=True, + ) + stderr_thread.start() + + try: + self.s3_client.upload_fileobj(process.stdout, bucket, key) + except Exception as exc: + process.kill() + process.wait() + stderr_thread.join() + raise exceptions.RuntimeError(f"Error uploading dump to S3: {exc}") from exc + finally: + process.stdout.close() + + process.wait() + stderr_thread.join() + + if process.returncode: + try: + self.s3_client.delete_object(Bucket=bucket, Key=key) + except Exception: + pass + raise exceptions.RuntimeError("Error running command.") + + def run_pg_restore(self, pg_restore_cmd: str, file_path: str) -> None: + bucket, key = _parse_s3_path(file_path) + process = subprocess.Popen( + pg_restore_cmd, + shell=True, + stdin=subprocess.PIPE, + stderr=subprocess.STDOUT, + stdout=subprocess.PIPE, + ) + if process.stdin is None or process.stdout is None: + raise AssertionError + + output_thread = threading.Thread( + target=_log_stream, + args=(process.stdout,), + daemon=True, + ) + output_thread.start() + + try: + self.s3_client.download_fileobj(bucket, key, process.stdin) + except Exception as exc: + process.kill() + process.wait() + output_thread.join() + raise exceptions.RuntimeError(f"Error downloading dump from S3: {exc}") from exc + finally: + process.stdin.close() + + process.wait() + output_thread.join() + + def ls(self, prefix: str | None = None) -> list[str]: + bucket, base_prefix = _parse_s3_path(self.storage_location) + full_prefix = f"{base_prefix}{prefix or ''}" + paginator = self.s3_client.get_paginator("list_objects_v2") + abs_paths = [] + for page in paginator.paginate(Bucket=bucket, Prefix=full_prefix): + for obj in page.get("Contents", []): + abs_paths.append(f"s3://{bucket}/{obj['Key']}") + return [self.dump_key(path) for path in abs_paths] + + class Local(_Storage): def ls(self, prefix: str | None = None) -> list[str]: abs_paths = [ @@ -113,6 +310,10 @@ def pg_restore(self, file_path: str) -> str: def client(storage_location: str) -> _Storage: if storage_location.startswith("s3://"): # pragma: no cover - return S3(storage_location) + backend = settings.s3_backend() + validate_s3_support(backend) + if backend == S3_BACKEND_BOTO3: + return S3Boto3(storage_location) + return S3Awscli(storage_location) else: return Local(storage_location) diff --git a/pgclone/tests/test_settings.py b/pgclone/tests/test_settings.py index fce3be2..acfaae4 100644 --- a/pgclone/tests/test_settings.py +++ b/pgclone/tests/test_settings.py @@ -10,6 +10,24 @@ def test_s3_config(settings): assert pgclone_settings.s3_config() == {"AWS_ACCESS_KEY_ID": "access_key"} +def test_s3_backend_explicit(settings): + settings.PGCLONE_S3_BACKEND = "boto3" + assert pgclone_settings.s3_backend() == "boto3" + settings.PGCLONE_S3_BACKEND = "awscli" + assert pgclone_settings.s3_backend() == "awscli" + + +def test_s3_backend_invalid(settings): + settings.PGCLONE_S3_BACKEND = "invalid" + with pytest.raises(RuntimeError): + pgclone_settings.s3_backend() + + +def test_s3_backend_defaults_to_awscli(settings): + delattr(settings, "PGCLONE_S3_BACKEND") + assert pgclone_settings.s3_backend() == "awscli" + + def test_conn_db(settings): delattr(settings, "PGCLONE_CONN_DB") settings.DATABASES = {"default": {"NAME": "hello"}} diff --git a/pgclone/tests/test_storage.py b/pgclone/tests/test_storage.py index 58a1b26..ea9dbbb 100644 --- a/pgclone/tests/test_storage.py +++ b/pgclone/tests/test_storage.py @@ -1,6 +1,9 @@ +import io +from unittest import mock + import pytest -from pgclone import storage +from pgclone import exceptions, storage @pytest.fixture(autouse=True) @@ -16,7 +19,9 @@ def test_s3_env(settings): "AWS_DEFAULT_REGION": "region", } - assert storage.S3("bucket").env == { + assert storage.S3Awscli("s3://bucket/").env == { + "AWS_REQUEST_CHECKSUM_CALCULATION": "when_required", + "AWS_RESPONSE_CHECKSUM_VALIDATION": "when_required", "AWS_ACCESS_KEY_ID": "access_key", "AWS_SECRET_ACCESS_KEY": "secret_access_key", "AWS_DEFAULT_REGION": "region", @@ -24,28 +29,328 @@ def test_s3_env(settings): delattr(settings, "PGCLONE_S3_CONFIG") - assert storage.S3("bucket").env == {} + assert storage.S3Awscli("s3://bucket/").env == { + "AWS_REQUEST_CHECKSUM_CALCULATION": "when_required", + "AWS_RESPONSE_CHECKSUM_VALIDATION": "when_required", + } def test_s3_pg_dump(): - assert storage.S3("bucket").pg_dump("file_path") == "| aws s3 cp - file_path" + assert storage.S3Awscli("s3://bucket/").pg_dump("s3://bucket/file_path") == ( + "| aws s3 cp - s3://bucket/file_path" + ) def test_s3_pg_restore(): - assert storage.S3("bucket").pg_restore("file_path") == "aws s3 cp file_path - |" + assert storage.S3Awscli("s3://bucket/").pg_restore("s3://bucket/file_path") == ( + "aws s3 cp s3://bucket/file_path - |" + ) def test_s3_pg_dump_with_endpoint_url(settings): settings.PGCLONE_S3_ENDPOINT_URL = "https://endpoint.example.com" assert ( - storage.S3("bucket").pg_dump("file_path") - == "| aws s3 cp - file_path --endpoint-url https://endpoint.example.com" + storage.S3Awscli("s3://bucket/").pg_dump("s3://bucket/file_path") + == "| aws s3 cp - s3://bucket/file_path --endpoint-url https://endpoint.example.com" ) def test_s3_pg_restore_with_endpoint_url(settings): settings.PGCLONE_S3_ENDPOINT_URL = "https://endpoint.example.com" assert ( - storage.S3("bucket").pg_restore("file_path") - == "aws s3 cp file_path - --endpoint-url https://endpoint.example.com |" + storage.S3Awscli("s3://bucket/").pg_restore("s3://bucket/file_path") + == "aws s3 cp s3://bucket/file_path - --endpoint-url https://endpoint.example.com |" + ) + + +def test_s3_alias(): + assert storage.S3 is storage.S3Awscli + + +def test_local_run_pg_dump(mocker): + shell = mocker.patch("pgclone.storage.run.shell", autospec=True) + local = storage.Local("/tmp/pgclone/") + local.run_pg_dump("pg_dump cmd", "/tmp/pgclone/file.dump") + shell.assert_called_once_with( + "pg_dump cmd > /tmp/pgclone/file.dump", + env={}, + pipefail=True, + ) + + +def test_local_run_pg_restore(mocker): + shell = mocker.patch("pgclone.storage.run.shell", autospec=True) + local = storage.Local("/tmp/pgclone/") + local.run_pg_restore("pg_restore cmd", "/tmp/pgclone/file.dump") + shell.assert_called_once_with( + "cat /tmp/pgclone/file.dump | pg_restore cmd", + env={}, + ignore_errors=True, + ) + + +def test_s3_awscli_run_pg_dump(mocker): + shell = mocker.patch("pgclone.storage.run.shell", autospec=True) + s3 = storage.S3Awscli("s3://bucket/") + s3.run_pg_dump("pg_dump cmd", "s3://bucket/file_path") + shell.assert_called_once_with( + "pg_dump cmd | aws s3 cp - s3://bucket/file_path", + env={ + "AWS_REQUEST_CHECKSUM_CALCULATION": "when_required", + "AWS_RESPONSE_CHECKSUM_VALIDATION": "when_required", + }, + pipefail=True, ) + + +def test_s3_awscli_run_pg_restore(mocker): + shell = mocker.patch("pgclone.storage.run.shell", autospec=True) + s3 = storage.S3Awscli("s3://bucket/") + s3.run_pg_restore("pg_restore cmd", "s3://bucket/file_path") + shell.assert_called_once_with( + "aws s3 cp s3://bucket/file_path - | pg_restore cmd", + env={ + "AWS_REQUEST_CHECKSUM_CALCULATION": "when_required", + "AWS_RESPONSE_CHECKSUM_VALIDATION": "when_required", + }, + ignore_errors=True, + ) + + +def test_client_defaults_to_awscli_backend(mocker, settings): + delattr(settings, "PGCLONE_S3_BACKEND") + mocker.patch("importlib.util.find_spec", return_value=mock.Mock()) + client = storage.client("s3://bucket/") + assert isinstance(client, storage.S3Awscli) + + +def test_client_uses_boto3_backend(mocker): + mocker.patch("pgclone.settings.s3_backend", return_value="boto3") + client = storage.client("s3://bucket/") + assert isinstance(client, storage.S3Boto3) + + +def test_client_uses_awscli_backend(mocker): + mocker.patch("pgclone.settings.s3_backend", return_value="awscli") + client = storage.client("s3://bucket/") + assert isinstance(client, storage.S3Awscli) + + +def test_client_uses_local_backend(): + client = storage.client("/tmp/pgclone/") + assert isinstance(client, storage.Local) + + +def test_parse_s3_path(): + assert storage._parse_s3_path("s3://bucket/key/path") == ("bucket", "key/path") + with pytest.raises(ValueError): + storage._parse_s3_path("/local/path") + + +@pytest.fixture +def boto3_client(mocker): + mock_boto3 = mocker.patch("boto3.session.Session", autospec=True) + mock_session = mock_boto3.return_value + mock_client = mock_session.client.return_value + return mock_client + + +def test_s3_boto3_session_kwargs(settings, mocker): + settings.PGCLONE_S3_CONFIG = { + "AWS_ACCESS_KEY_ID": "access_key", + "AWS_SECRET_ACCESS_KEY": "secret_access_key", + "AWS_SESSION_TOKEN": "session_token", + "AWS_DEFAULT_REGION": "us-east-1", + } + mock_session_cls = mocker.patch("boto3.session.Session", autospec=True) + s3 = storage.S3Boto3("s3://bucket/") + assert s3.s3_client is mock_session_cls.return_value.client.return_value + mock_session_cls.assert_called_once_with( + aws_access_key_id="access_key", + aws_secret_access_key="secret_access_key", + aws_session_token="session_token", + region_name="us-east-1", + ) + mock_session_cls.return_value.client.assert_called_once_with("s3", config=mock.ANY) + + +def test_s3_boto3_client_kwargs(settings, mocker): + settings.PGCLONE_S3_ENDPOINT_URL = "https://endpoint.example.com" + mock_session_cls = mocker.patch("boto3.session.Session", autospec=True) + s3 = storage.S3Boto3("s3://bucket/") + client = s3.s3_client + mock_session_cls.return_value.client.assert_called_once_with( + "s3", + endpoint_url="https://endpoint.example.com", + config=mock.ANY, + ) + assert client is mock_session_cls.return_value.client.return_value + + +def test_s3_boto3_ls(boto3_client): + paginator = boto3_client.get_paginator.return_value + paginator.paginate.return_value = [ + { + "Contents": [ + {"Key": "prefix/instance/db/config/2024-01-01-00-00-00-000000.dump"}, + {"Key": "prefix/other.dump"}, + ] + } + ] + s3 = storage.S3Boto3("s3://bucket/prefix/") + dump_keys = s3.ls() + paginator.paginate.assert_called_once_with(Bucket="bucket", Prefix="prefix/") + assert dump_keys == [ + "instance/db/config/2024-01-01-00-00-00-000000.dump", + "other.dump", + ] + + +def test_s3_boto3_ls_with_prefix(boto3_client): + paginator = boto3_client.get_paginator.return_value + paginator.paginate.return_value = [{"Contents": []}] + s3 = storage.S3Boto3("s3://bucket/prefix/") + s3.ls(prefix="instance/") + paginator.paginate.assert_called_once_with(Bucket="bucket", Prefix="prefix/instance/") + + +def test_s3_boto3_run_pg_dump_success(mocker, boto3_client): + process = mock.Mock() + process.stdout = io.BytesIO(b"dump-data") + process.stderr = io.BytesIO(b"") + process.returncode = 0 + process.wait.return_value = 0 + mocker.patch("pgclone.storage.subprocess.Popen", return_value=process) + + s3 = storage.S3Boto3("s3://bucket/") + s3.run_pg_dump("pg_dump cmd", "s3://bucket/key.dump") + + boto3_client.upload_fileobj.assert_called_once_with(process.stdout, "bucket", "key.dump") + + +def test_s3_boto3_run_pg_dump_upload_failure(mocker, boto3_client): + process = mock.Mock() + process.stdout = io.BytesIO(b"dump-data") + process.stderr = io.BytesIO(b"") + mocker.patch("pgclone.storage.subprocess.Popen", return_value=process) + boto3_client.upload_fileobj.side_effect = Exception("upload failed") + + s3 = storage.S3Boto3("s3://bucket/") + with pytest.raises(exceptions.RuntimeError, match="Error uploading dump to S3"): + s3.run_pg_dump("pg_dump cmd", "s3://bucket/key.dump") + + process.kill.assert_called_once() + + +def test_s3_boto3_run_pg_dump_process_failure(mocker, boto3_client): + process = mock.Mock() + process.stdout = io.BytesIO(b"dump-data") + process.stderr = io.BytesIO(b"") + process.returncode = 1 + process.wait.return_value = 1 + mocker.patch("pgclone.storage.subprocess.Popen", return_value=process) + + s3 = storage.S3Boto3("s3://bucket/") + with pytest.raises(exceptions.RuntimeError, match="Error running command"): + s3.run_pg_dump("pg_dump cmd", "s3://bucket/key.dump") + + +def test_s3_boto3_run_pg_restore_success(mocker, boto3_client): + process = mock.Mock() + process.stdin = mock.Mock() + process.stdout = io.BytesIO(b"") + process.returncode = 0 + process.wait.return_value = 0 + mocker.patch("pgclone.storage.subprocess.Popen", return_value=process) + + s3 = storage.S3Boto3("s3://bucket/") + s3.run_pg_restore("pg_restore cmd", "s3://bucket/key.dump") + + boto3_client.download_fileobj.assert_called_once_with("bucket", "key.dump", process.stdin) + + +def test_s3_boto3_run_pg_restore_download_failure(mocker, boto3_client): + process = mock.Mock() + process.stdin = mock.Mock() + process.stdout = io.BytesIO(b"") + mocker.patch("pgclone.storage.subprocess.Popen", return_value=process) + boto3_client.download_fileobj.side_effect = Exception("download failed") + + s3 = storage.S3Boto3("s3://bucket/") + with pytest.raises(exceptions.RuntimeError, match="Error downloading dump from S3"): + s3.run_pg_restore("pg_restore cmd", "s3://bucket/key.dump") + + process.kill.assert_called_once() + + +def test_s3_boto3_run_pg_restore_ignores_pg_restore_errors(mocker, boto3_client): + process = mock.Mock() + process.stdin = mock.Mock() + process.stdout = io.BytesIO(b"") + process.returncode = 1 + process.wait.return_value = 1 + mocker.patch("pgclone.storage.subprocess.Popen", return_value=process) + + s3 = storage.S3Boto3("s3://bucket/") + s3.run_pg_restore("pg_restore cmd", "s3://bucket/key.dump") + + +def test_log_stream(mocker): + logger = mocker.patch("pgclone.storage.logging.get_logger", autospec=True) + stream = io.BytesIO(b"line1\nline2\n") + storage._log_stream(stream) + assert logger.return_value.info.call_count == 2 + + +def test_is_boto3_importable(mocker): + mocker.patch("importlib.util.find_spec", return_value=mock.Mock()) + assert storage._is_boto3_importable() is True + mocker.patch("importlib.util.find_spec", return_value=None) + assert storage._is_boto3_importable() is False + + +def test_is_awscli_available(mocker): + mocker.patch( + "pgclone.storage.subprocess.run", + return_value=mock.Mock(returncode=0), + ) + assert storage._is_awscli_available() is True + mocker.patch( + "pgclone.storage.subprocess.run", + return_value=mock.Mock(returncode=1), + ) + assert storage._is_awscli_available() is False + + +def test_s3_boto3_client_cached(mocker): + mock_session_cls = mocker.patch("boto3.session.Session", autospec=True) + s3 = storage.S3Boto3("s3://bucket/") + first_client = s3.s3_client + second_client = s3.s3_client + mock_session_cls.assert_called_once() + assert first_client is second_client + + +def test_local_ls(tmp_path): + dump_dir = tmp_path / "pgclone" / "instance" / "db" / "config" + dump_dir.mkdir(parents=True) + dump_file = dump_dir / "2024-01-01-00-00-00-000000.dump" + dump_file.write_text("dump", encoding="utf-8") + + local = storage.Local(str(tmp_path / "pgclone") + "/") + dump_keys = local.ls() + assert dump_keys == ["instance/db/config/2024-01-01-00-00-00-000000.dump"] + + assert local.ls(prefix="instance/") == ["instance/db/config/2024-01-01-00-00-00-000000.dump"] + + +def test_local_pg_dump(tmp_path): + local = storage.Local(str(tmp_path / "pgclone") + "/") + file_path = str(tmp_path / "pgclone" / "file.dump") + assert local.pg_dump(file_path) == f"> {file_path}" + assert (tmp_path / "pgclone").exists() + + +def test_local_pg_restore(): + local = storage.Local("/tmp/pgclone/") + assert local.pg_restore("/tmp/pgclone/file.dump") == "cat /tmp/pgclone/file.dump |" diff --git a/poetry.lock b/poetry.lock index a0c6d1b..2d74361 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.4.1 and should not be changed by hand. [[package]] name = "anyio" @@ -177,6 +177,48 @@ d = ["aiohttp (>=3.10)"] jupyter = ["ipython (>=7.8.0)", "tokenize-rt (>=3.2.0)"] uvloop = ["uvloop (>=0.15.2)"] +[[package]] +name = "boto3" +version = "1.43.41" +description = "The AWS SDK for Python" +optional = false +python-versions = ">=3.10" +groups = ["main", "dev"] +files = [ + {file = "boto3-1.43.41-py3-none-any.whl", hash = "sha256:f48f862d2720ea9203ed2d842d436b8eb2d459ea31654a7ad7c0756fdf36c6b2"}, + {file = "boto3-1.43.41.tar.gz", hash = "sha256:0f56811f13677bfb4542daa0cce8532c95d9afd27b4ba7b681af36a0568624ad"}, +] +markers = {main = "extra == \"s3\""} + +[package.dependencies] +botocore = ">=1.43.41,<1.44.0" +jmespath = ">=0.7.1,<2.0.0" +s3transfer = ">=0.19.0,<0.20.0" + +[package.extras] +crt = ["botocore[crt] (>=1.21.0,<2.0a0)"] + +[[package]] +name = "botocore" +version = "1.43.41" +description = "Low-level, data-driven core of boto 3." +optional = false +python-versions = ">=3.10" +groups = ["main", "dev"] +files = [ + {file = "botocore-1.43.41-py3-none-any.whl", hash = "sha256:0cc6e79b30a2a98374f16a31cd9c7a9106a51b60650bd8c34cc8223f58ae6b8d"}, + {file = "botocore-1.43.41.tar.gz", hash = "sha256:27627d79af0df7dcb7ecf78d8d3d1310da09a5e9460be30bf759f1c2ed095ee8"}, +] +markers = {main = "extra == \"s3\""} + +[package.dependencies] +jmespath = ">=0.7.1,<2.0.0" +python-dateutil = ">=2.1,<3.0.0" +urllib3 = ">=1.25.4,<2.2.0 || >2.2.0,<3" + +[package.extras] +crt = ["awscrt (==0.32.2)"] + [[package]] name = "build" version = "1.3.0" @@ -1267,6 +1309,19 @@ files = [ arrow = "*" jinja2 = "*" +[[package]] +name = "jmespath" +version = "1.1.0" +description = "JSON Matching Expressions" +optional = false +python-versions = ">=3.9" +groups = ["main", "dev"] +files = [ + {file = "jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64"}, + {file = "jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d"}, +] +markers = {main = "extra == \"s3\""} + [[package]] name = "keyring" version = "25.7.0" @@ -2147,11 +2202,12 @@ version = "2.9.0.post0" description = "Extensions to the standard Python datetime module" optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3"}, {file = "python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427"}, ] +markers = {main = "extra == \"s3\""} [package.dependencies] six = ">=1.5" @@ -2453,6 +2509,25 @@ files = [ {file = "ruff-0.14.7.tar.gz", hash = "sha256:3417deb75d23bd14a722b57b0a1435561db65f0ad97435b4cf9f85ffcef34ae5"}, ] +[[package]] +name = "s3transfer" +version = "0.19.0" +description = "An Amazon S3 Transfer Manager" +optional = false +python-versions = ">=3.10" +groups = ["main", "dev"] +files = [ + {file = "s3transfer-0.19.0-py3-none-any.whl", hash = "sha256:777cc2415536f1debadb5c2ef7779275d0fc0fe0e042411cdd6caebeb2685262"}, + {file = "s3transfer-0.19.0.tar.gz", hash = "sha256:ce436931687addc4c1712d52d40b32f53e88315723f107ffa20ba82b05a0f685"}, +] +markers = {main = "extra == \"s3\""} + +[package.dependencies] +botocore = ">=1.37.4,<2.0a0" + +[package.extras] +crt = ["botocore[crt] (>=1.37.4,<2.0a0)"] + [[package]] name = "secretstorage" version = "3.5.0" @@ -2509,11 +2584,12 @@ version = "1.17.0" description = "Python 2 and 3 compatibility utilities" optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274"}, {file = "six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81"}, ] +markers = {main = "extra == \"s3\""} [[package]] name = "sqlparse" @@ -2690,6 +2766,7 @@ files = [ {file = "typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548"}, {file = "typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466"}, ] +markers = {main = "python_version < \"3.11\""} [[package]] name = "tzdata" @@ -2710,11 +2787,12 @@ version = "2.5.0" description = "HTTP library with thread-safe connection pooling, file post, and more." optional = false python-versions = ">=3.9" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc"}, {file = "urllib3-2.5.0.tar.gz", hash = "sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760"}, ] +markers = {main = "extra == \"s3\""} [package.extras] brotli = ["brotli (>=1.0.9) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; platform_python_implementation != \"CPython\""] @@ -2985,9 +3063,12 @@ files = [ ] [package.extras] -cffi = ["cffi (>=1.17,<2.0) ; platform_python_implementation != \"PyPy\" and python_version < \"3.14\"", "cffi (>=2.0.0b) ; platform_python_implementation != \"PyPy\" and python_version >= \"3.14\""] +cffi = ["cffi (>=1.17,<2.0) ; platform_python_implementation != \"PyPy\" and python_version < \"3.14\"", "cffi (>=2.0.0b0) ; platform_python_implementation != \"PyPy\" and python_version >= \"3.14\""] + +[extras] +s3 = ["boto3"] [metadata] lock-version = "2.1" python-versions = ">=3.10.0,<4" -content-hash = "2be4a2084a8f906dc15f812ad4e80e653fbe361d213f93bd983fc65f06c56b4f" +content-hash = "bb5566e79da68d1175d0f9f4158e394a389cb5979ff9ed015928ba6da12abd40" diff --git a/pyproject.toml b/pyproject.toml index ed9ef43..af0aed8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,10 @@ documentation = "https://django-pgclone.readthedocs.io" [tool.poetry.dependencies] python = ">=3.10.0,<4" django = ">=4.2" +boto3 = { version = ">=1.26", optional = true } + +[tool.poetry.extras] +s3 = ["boto3"] [tool.poetry.dev-dependencies] pytest = "9.0.1" @@ -82,6 +86,9 @@ psycopg2-binary = "2.9.11" pytest-django = "4.11.1" django-dynamic-fixture = "4.0.1" +[tool.poetry.group.dev.dependencies] +boto3 = ">=1.26" + [tool.pytest.ini_options] xfail_strict = true testpaths = "pgclone/tests"