diff --git a/.github/workflows/baseshift.yml b/.github/workflows/baseshift.yml new file mode 100644 index 0000000..26065ef --- /dev/null +++ b/.github/workflows/baseshift.yml @@ -0,0 +1,111 @@ +name: LocalStack Baseshift Extension Tests + +on: + schedule: + - cron: '0 2 * * 1-5' + pull_request: + branches: + - main + paths: + - .github/workflows/baseshift.yml + - 'baseshift/**' + push: + branches: + - main + paths: + - .github/workflows/baseshift.yml + - 'baseshift/**' + workflow_dispatch: + +env: + LOCALSTACK_DISABLE_EVENTS: "1" + LOCALSTACK_AUTH_TOKEN: ${{ secrets.LOCALSTACK_AUTH_TOKEN }} + +jobs: + integration-tests: + name: Run Baseshift Extension Tests (${{ matrix.emulator }}, ${{ matrix.tag }}) + runs-on: ubuntu-latest + permissions: + contents: read + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + # the Snowflake emulator image also provides the AWS services, and is used for the end-to-end demo + emulator: + - aws + - snowflake + tag: + - latest + - dev + env: + LOCALSTACK_EMULATOR: ${{ matrix.emulator }} + LOCALSTACK_TAG: ${{ matrix.tag }} + steps: + - name: Checkout + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + + - name: Install lstk CLI + # version is pinned; there is no lockfile for global CLI tools + run: | # zizmor: ignore[adhoc-packages] + npm install -g @localstack/lstk@1.2.0 + lstk --version + + - name: Build extension + run: | + cd baseshift + make install + make lint + make dist + + - name: Start LocalStack with extension + run: | + cd baseshift + + # Baseshift clone images are private to each customer, hence we use a plain + # Postgres image as a stand-in (clones also accept connections without a password) + docker pull postgres:17 + + # mount the built extension into the container, so LocalStack can install it at startup + CONFIG_FILE="$(lstk config path)" + mkdir -p "$(dirname "$CONFIG_FILE")" + cat > "$CONFIG_FILE" < /dev/tcp/localhost/5432) 2>/dev/null; do sleep 2; done' + + - name: Run integration tests + run: | + cd baseshift + make test + + - name: Run end-to-end demo + if: matrix.emulator == 'snowflake' + run: | + cd baseshift + make demo + + - name: Print logs + if: always() + run: | + lstk logs --verbose + lstk stop diff --git a/README.md b/README.md index a138a81..b41d8a6 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,7 @@ You can install the respective extension by calling `localstack extensions insta | Extension | Install name | Version | Support status | |----------------------------------------------------------------------------------------------------| ------------ |---------| -------------- | | [AWS Proxy](https://github.com/localstack/localstack-extensions/tree/main/aws-proxy) | localstack-extension-aws-proxy | 0.2.1 | Experimental | +| [Baseshift](https://github.com/localstack/localstack-extensions/tree/main/baseshift) | localstack-baseshift | 0.1.0 | Experimental | | [Diagnosis Viewer](https://github.com/localstack/localstack-extensions/tree/main/diagnosis-viewer) | localstack-extension-diagnosis-viewer | 0.1.0 | Stable | | [Hello World](https://github.com/localstack/localstack-extensions/tree/main/hello-world) | localstack-extension-hello-world | 0.1.0 | Stable | | [httpbin](https://github.com/localstack/localstack-extensions/tree/main/httpbin) | localstack-extension-httpbin | 0.1.0 | Stable | diff --git a/baseshift/.gitignore b/baseshift/.gitignore new file mode 100644 index 0000000..1602d05 --- /dev/null +++ b/baseshift/.gitignore @@ -0,0 +1,8 @@ +.venv +dist +build +**/*.egg-info +.eggs +.pytest_cache +__pycache__ +demo/baseshift-selfhosted/.env diff --git a/baseshift/Makefile b/baseshift/Makefile new file mode 100644 index 0000000..c7371dd --- /dev/null +++ b/baseshift/Makefile @@ -0,0 +1,51 @@ +VENV_BIN = python3 -m venv +VENV_DIR ?= .venv +VENV_ACTIVATE = $(VENV_DIR)/bin/activate +VENV_RUN = . $(VENV_ACTIVATE) +TEST_PATH ?= tests + +usage: ## Shows usage for this Makefile + @cat Makefile | grep -E '^[a-zA-Z_-]+:.*?## .*$$' | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-15s\033[0m %s\n", $$1, $$2}' + +venv: $(VENV_ACTIVATE) + +$(VENV_ACTIVATE): pyproject.toml + test -d .venv || $(VENV_BIN) .venv + $(VENV_RUN); pip install --upgrade pip setuptools plux + $(VENV_RUN); pip install -e .[dev] + touch $(VENV_DIR)/bin/activate + +clean: + rm -rf .venv/ + rm -rf build/ + rm -rf .eggs/ + rm -rf *.egg-info/ + +install: venv ## Install dependencies + $(VENV_RUN); python -m plux entrypoints + +dist: venv ## Create distribution + $(VENV_RUN); python -m build + +publish: clean-dist venv dist ## Publish extension to pypi + $(VENV_RUN); pip install --upgrade twine; twine upload dist/* + +entrypoints: venv # Generate plugin entrypoints for Python package + $(VENV_RUN); python -m plux entrypoints + +format: ## Run ruff to format the whole codebase + $(VENV_RUN); python -m ruff format .; python -m ruff check --output-format=full --fix . + +lint: ## Run ruff to lint the codebase + $(VENV_RUN); python -m ruff check --output-format=full . + +test: ## Run integration tests (requires LocalStack running with the Extension installed) + $(VENV_RUN); pytest $(PYTEST_ARGS) $(TEST_PATH) + +demo: venv ## Run the end-to-end demo (requires LocalStack with the Snowflake emulator and the Extension installed) + $(VENV_RUN); pip install -q -r demo/requirements.txt; python demo/demo.py all + +clean-dist: clean + rm -rf dist/ + +.PHONY: clean clean-dist demo dist install publish usage venv format lint test diff --git a/baseshift/README.md b/baseshift/README.md new file mode 100644 index 0000000..3cb117f --- /dev/null +++ b/baseshift/README.md @@ -0,0 +1,216 @@ +# Baseshift on LocalStack + +This repo contains a [LocalStack Extension](https://github.com/localstack/localstack-extensions) that runs a [Baseshift](https://baseshift.com) database clone next to LocalStack. + +Baseshift creates masked, writable clones of your production database (PostgreSQL or MySQL) for development and testing. With this extension, your application code running in LocalStack (e.g., Lambda functions or ECS tasks) can work against a realistic, anonymized copy of your data, all on your local machine. + +For PostgreSQL clones, the database is available through the LocalStack gateway (`localhost.localstack.cloud:4566`), as well as on the regular host port `5432`. + +## Architecture + +```mermaid +flowchart LR + subgraph source["1 · Production data"] + rds[("LocalStack RDS
PostgreSQL with PII")] + end + + subgraph baseshift["2 · Baseshift (self-hosted)"] + direction TB + cloud["Baseshift Cloud
control plane"] + connector["Connector
masking policy"] + repserver["Replication server
masked replica, snapshots"] + cloud -.- repserver + connector -- "masked data" --> repserver + end + + subgraph registry["3 · Snapshot registry"] + ecr[("LocalStack ECR
Docker snapshot images")] + end + + subgraph clones["4 · Local clones (this extension)"] + direction TB + extension["Baseshift extension
clones API"] + clone1[("Clone 'default'
:5432")] + clone2[("Clone 'pr-123'
:15432")] + extension -- "start / stop" --> clone1 + extension -- "start / stop" --> clone2 + end + + gateway["LocalStack gateway :4566
PostgreSQL routing"] + + subgraph consumers["5 · Consumers of masked data"] + direction TB + app["App code
Lambda, ECS, ..."] + dev["Developer / CI
psql, IDE, tests"] + subgraph analytics["Analytics pipeline"] + direction LR + elt["ELT job"] -- "CSV" --> s3[("LocalStack S3")] -- "COPY INTO" --> snowflake[("LocalStack
Snowflake")] + end + end + + rds -- "replicate" --> connector + repserver -- "push" --> ecr + ecr -- "pull" --> extension + clone1 --> gateway + gateway -- "SQL" --> app + gateway -- "SQL" --> dev + gateway -- "SQL" --> elt + clone2 -. "SQL via host port" .-> dev +``` + +1. **Production data**: the source database, e.g. PostgreSQL in (LocalStack) RDS, containing PII. +2. **Baseshift** replicates the source database via its connector, which applies the masking policy, and the replication server maintains a masked replica and creates snapshots. The components are managed via the Baseshift Cloud control plane. (In the [demo](demo/), a stand-in script currently takes the place of this step.) +3. The replication server publishes Docker **snapshot images** to a registry, e.g. AWS ECR, or the LocalStack ECR registry for a fully local setup. +4. This **extension** starts clones from these images as containers next to LocalStack: a default clone at startup (`BASESHIFT_IMAGE`), and further clones on demand via the clones API. Each clone is a writable copy of the masked database. +5. **Consumers** work with the masked data only: app code running in LocalStack and developer tools connect through the LocalStack gateway (the first PostgreSQL clone is detected by its protocol handshake on port 4566), or to the host port of a clone. An analytics pipeline can extract data from a clone, and load it into Snowflake (see the [demo](demo/)). + +## Prerequisites + +- Docker +- LocalStack Pro (free trial available) +- [`lstk`](https://docs.localstack.cloud/aws/developer-tools/running-localstack/lstk/) CLI (`npm install -g @localstack/lstk`) +- A Baseshift account with a Dub, and a Docker snapshot (clone image) of it +- AWS CLI, for pulling the clone image from ECR +- `make` + +## Getting the clone image + +Baseshift clone images are private to your organization, and are stored in an AWS ECR repository. In the Baseshift dashboard, select your Dub and click **Get docker snapshot** to get the commands for authenticating to ECR and pulling the image, for example: + +```bash +aws ecr get-login-password --region | \ + docker login --username AWS --password-stdin .dkr.ecr..amazonaws.com +docker pull .dkr.ecr..amazonaws.com/:latest +``` + +The image needs to be pulled on the host **before** starting LocalStack, as the extension does not authenticate against ECR. Note that ECR logins expire after 12 hours. + +See the [Baseshift docs on local clones](https://docs.baseshift.com/docs/clones/local-clones) for more details. + +## Install from GitHub repository + +`lstk` does not have commands for managing extensions, but LocalStack can install the extension at startup via the `EXTENSION_AUTO_INSTALL` config variable: + +```bash +LOCALSTACK_EXTENSION_AUTO_INSTALL="git+https://github.com/localstack/localstack-extensions.git#egg=localstack-baseshift&subdirectory=baseshift" \ +LOCALSTACK_BASESHIFT_IMAGE=.dkr.ecr..amazonaws.com/:latest \ + lstk start +``` + +Alternatively, if you are using the legacy `localstack` CLI: + +```bash +localstack extensions install "git+https://github.com/localstack/localstack-extensions.git#egg=localstack-baseshift&subdirectory=baseshift" +``` + +## Install local development version + +To install the extension into LocalStack in developer mode, you will need Python 3.11, and create a virtual environment in the extensions project. + +In the newly generated project, simply run + +```bash +make install +``` + +Developer mode currently requires the legacy `localstack` CLI, as `lstk` does not support extensions yet. To enable the extension for LocalStack, run + +```bash +localstack extensions dev enable . +``` + +You can then start LocalStack with `EXTENSION_DEV_MODE=1` to load all enabled extensions (the `localstack` CLI mounts the extension sources into the container): + +```bash +EXTENSION_DEV_MODE=1 LOCALSTACK_BASESHIFT_IMAGE= localstack start +``` + +## Usage + +### Default clone + +Start LocalStack with `BASESHIFT_IMAGE` pointing to your clone image, and (if one was configured when creating the Dub) the encryption password: + +```bash +LOCALSTACK_BASESHIFT_IMAGE=.dkr.ecr..amazonaws.com/:latest \ +LOCALSTACK_BASESHIFT_ENCRYPTION_PASSWORD= \ + lstk start +``` + +The clone (named `default`) is started in the background once LocalStack is ready. Connect to it, for example with `psql`. Local clones do not require a password; use one of the database users replicated from your source database: + +```bash +# through the LocalStack gateway +psql -h localhost.localstack.cloud -p 4566 -U + +# or directly on the host port +psql -h localhost -p 5432 -U +``` + +From inside LocalStack (e.g., Lambda functions), connect to `localhost.localstack.cloud:4566`. + +### Clones API + +Clones can also be started and stopped on demand (e.g., one clone per pull request or coding agent), via the API at `http://baseshift.localhost.localstack.cloud:4566/clones`: + +```bash +# start a clone +curl -X POST http://baseshift.localhost.localstack.cloud:4566/clones \ + -d '{"name": "pr-123", "image": ""}' + +# list clones, with their status and endpoints +curl http://baseshift.localhost.localstack.cloud:4566/clones + +# get / stop a clone +curl http://baseshift.localhost.localstack.cloud:4566/clones/pr-123 +curl -X DELETE http://baseshift.localhost.localstack.cloud:4566/clones/pr-123 +``` + +The request accepts `name` (lowercase letters, digits, and dashes), `image`, and optionally `dbType` (`postgres` or `mysql`) and `env` (additional environment variables for the clone container). Clones start asynchronously - their `status` changes from `starting` to `running` (or `failed`, with an `error` message). + +Each clone gets its own host port. The first PostgreSQL clone gets port `5432` and is also available through the LocalStack gateway; further clones get ports from `15432` upwards (see the `endpoints` in the API response). + +### Clone images in LocalStack ECR + +Clone images can also be served by the LocalStack ECR registry, e.g., to share snapshot images within a team, or to emulate the full Baseshift flow locally: + +```bash +lstk aws ecr create-repository --repository-name baseshift/my-dub +docker tag 000000000000.dkr.ecr.us-east-1.localhost.localstack.cloud:4566/baseshift/my-dub:latest +docker push 000000000000.dkr.ecr.us-east-1.localhost.localstack.cloud:4566/baseshift/my-dub:latest + +curl -X POST http://baseshift.localhost.localstack.cloud:4566/clones \ + -d '{"name": "my-dub", "image": "000000000000.dkr.ecr.us-east-1.localhost.localstack.cloud:4566/baseshift/my-dub:latest"}' +``` + +The LocalStack ECR registry does not require a `docker login`. + +### MySQL clones + +Set `BASESHIFT_DB_TYPE=mysql` (or `"dbType": "mysql"` in the API) for MySQL clones. MySQL clones are only available on their host port (`3306` for the first one): unlike PostgreSQL, the MySQL protocol starts with the server sending a greeting, so MySQL connections cannot be told apart from other traffic on the gateway port `4566`. + +### Environment Variables + +- `BASESHIFT_IMAGE`: Image of the default clone to start once LocalStack is ready (optional, clones can also be started via the API) +- `BASESHIFT_DB_TYPE`: Database engine of the default clone, `postgres` (default) or `mysql` +- `BASESHIFT_ENCRYPTION_PASSWORD`: Encryption password defined when the Dub was created (passed to the clones as `PASSWORD`) +- `BASESHIFT_CLONE_`: Passed to the clone containers as ``, for the [advanced clone options](https://docs.baseshift.com/docs/clones/local-clones#advanced-configuration), e.g. `BASESHIFT_CLONE_BACKUP_SCHEDULE`, `BASESHIFT_CLONE_MAX_BACKUPS`, or `BASESHIFT_CLONE_SPACE_USAGE_MIN_PERCENT` + +Note: When starting LocalStack via `lstk` (or the `localstack` CLI), prefix environment variables with `LOCALSTACK_` to forward them to the container, e.g. `LOCALSTACK_BASESHIFT_IMAGE`. + +### Limitations + +- PostgreSQL connections on the gateway port are detected by their protocol handshake. Running this extension together with another extension that serves PostgreSQL on the gateway (e.g., ParadeDB) is not supported. +- The clone ports are published on the host, so they must not be in use by another database. + +## Demo + +The [`demo/`](demo/) directory contains an end-to-end demo of masked production data flowing through a local pipeline: a "production" database in LocalStack RDS, a masked snapshot image in LocalStack ECR, a clone started via this extension, and an ELT pipeline into the LocalStack Snowflake emulator. + +## Change Log + +- `0.1.0`: Initial release of the extension + +## License + +This project is licensed under the Apache License, Version 2.0. diff --git a/baseshift/demo/README.md b/baseshift/demo/README.md new file mode 100644 index 0000000..cf33025 --- /dev/null +++ b/baseshift/demo/README.md @@ -0,0 +1,72 @@ +# Baseshift on LocalStack - Demo + +This demo shows how masked production data flows through a fully local development setup: + +``` + "production" DB masked snapshot clone ELT pipeline warehouse + LocalStack RDS ---> LocalStack ECR ---> Baseshift ---> S3 (CSV export) ---> Snowflake + (PostgreSQL + PII) (Docker image) extension emulator +``` + +At the end, the same customer record is shown at every stage of the pipeline - raw in the source database, and masked in the clone and in Snowflake: + +``` +column source (RDS) Baseshift clone Snowflake +-------- ---------------------------- ---------------------------- ---------------------------- +name Jane Roe Customer 9231c1 Customer 9231c1 +email jane.roe@gmail.com user-ae9cacc9@example.com user-ae9cacc9@example.com +ssn 078-05-1120 ***-**-1120 ***-**-1120 +phone +1 415 555 0142 +1 555 3021670 +1 555 3021670 +address 221B Baker Street, London 101 Masked Street 101 Masked Street + +No raw PII values reached the clone or Snowflake. +``` + +## Prerequisites + +- Docker, and the [`lstk`](https://docs.localstack.cloud/aws/developer-tools/running-localstack/lstk/) CLI +- A LocalStack license that includes the Snowflake emulator +- Python 3, with the demo dependencies: `pip install -r demo/requirements.txt` + +The demo uses the LocalStack Snowflake emulator image (`lstk` emulator type `snowflake`), which also provides the AWS services used here (RDS, ECR, S3), so everything runs in a single LocalStack container. + +## Running the demo + +Start LocalStack with the Snowflake emulator and the extension installed: + +```bash +lstk start --type snowflake # with LOCALSTACK_EXTENSION_AUTO_INSTALL set, see the extension README +``` + +Then run all steps of the demo (from the `baseshift` directory): + +```bash +make demo +``` + +Or run the steps one by one, to walk through the flow: + +```bash +python demo/demo.py source # create the "production" RDS database, and seed it with PII +python demo/demo.py snapshot # create a masked snapshot image, and push it to LocalStack ECR +python demo/demo.py clone # start a clone from the snapshot image, via the extension API +python demo/demo.py pipeline # export the clone data to S3, and load it into Snowflake +python demo/demo.py compare # show the same record at every stage +``` + +Set `LOCALSTACK_HOST` if LocalStack runs on a different port (default: `localhost.localstack.cloud:4566`). + +## What is emulated + +The `snapshot` step is currently a **stand-in** for the Baseshift replication server: it reads the source database, applies a simple deterministic masking policy, and builds a PostgreSQL image with the masked data. Its purpose is to illustrate the flow - the actual Baseshift components provide PII scanning, masking policies, subsetting, and continuous snapshots via logical replication. + +The next step is to replace the stand-in with the actual Baseshift self-hosted components, running against LocalStack: + +- The [`baseshift-selfhosted/`](baseshift-selfhosted/) directory contains Helm values for LocalStack (`values-localstack.yaml`, validated against the Baseshift chart schema), and a `docker-compose.yml` translated from the rendered chart, as a lightweight alternative to Kubernetes. +- The source database would be the RDS instance created by the `source` step (which already enables logical replication), and Docker snapshots would be pushed to the LocalStack ECR registry (as a plain Docker v2 registry). +- This requires access to the Baseshift replication server image, and a Dub server enrolled in Baseshift Cloud (the control plane). + +## Notes + +- LocalStack RDS does not apply the `rds.logical_replication` parameter (yet) - the `source` step sets `wal_level = logical` directly, and reboots the instance. +- The ELT step is a plain script for simplicity; in a real setup, it would typically be a scheduled job (e.g., a Lambda function or an Airflow DAG) running against the clone. diff --git a/baseshift/demo/baseshift-selfhosted/.env.example b/baseshift/demo/baseshift-selfhosted/.env.example new file mode 100644 index 0000000..267faad --- /dev/null +++ b/baseshift/demo/baseshift-selfhosted/.env.example @@ -0,0 +1,8 @@ +# IDs and tokens for the self-hosted Dub, provided by Baseshift (Dub setup wizard) +BASESHIFT_ORG_ID= +BASESHIFT_SERVER_ID= +BASESHIFT_ENROLLMENT_TOKEN= +BASESHIFT_AGENT_ID= +BASESHIFT_AGENT_TOKEN= +# encryption password for the Docker snapshot images (letters and numbers only) +BASESHIFT_ENCRYPTION_PASSWORD=localstack123 diff --git a/baseshift/demo/baseshift-selfhosted/docker-compose.yml b/baseshift/demo/baseshift-selfhosted/docker-compose.yml new file mode 100644 index 0000000..41b90e5 --- /dev/null +++ b/baseshift/demo/baseshift-selfhosted/docker-compose.yml @@ -0,0 +1,86 @@ +# Baseshift self-hosted components (replication server + connector) as plain containers, as a +# lightweight alternative to the Kubernetes deployment. Translated from the rendered `baseshift-dub` +# Helm chart (bundled in the `baseshift` CLI) with `values-localstack.yaml`. +# +# Untested so far: requires access to the (private) replication server image, and the IDs/tokens +# from Baseshift in `.env` (see `.env.example`). +# +# docker login registry.gitlab.com # with the credentials provided by Baseshift +# cp .env.example .env # and fill in the values +# docker compose up -d + +x-localstack-hosts: &localstack-hosts + # LocalStack hostnames resolve to 127.0.0.1 - inside these containers they need to point to the host + # (where LocalStack publishes its ports). Note: host-gateway reaches ports bound to 127.0.0.1 on + # Docker Desktop, but not on Linux Docker engines. + - "localhost.localstack.cloud:host-gateway" + - "000000000000.dkr.ecr.us-east-1.localhost.localstack.cloud:host-gateway" + +services: + repserver: + image: registry.gitlab.com/dubhub/core/multi-ubuntu2404/repserver:master + hostname: repserver + restart: unless-stopped + shm_size: 512m + extra_hosts: *localstack-hosts + environment: + SERVERID: ${BASESHIFT_SERVER_ID:?} + ORGID: ${BASESHIFT_ORG_ID:?} + GATEWAY: aws-gw-elb.dubhub.io:50052 + ENVIRONMENT: production + ENROLLMENT_TOKEN: ${BASESHIFT_ENROLLMENT_TOKEN:?} + CERT_DIR: /dubhub/certs + EXTRAFLAGS: -proxytimeout=45s + SERVERNAME: localstack-demo + # the Kubernetes service maps port 15056 to the container's SSH port 22 + SERVERURI: repserver:22 + PASSWORD: ${BASESHIFT_ENCRYPTION_PASSWORD:-localstack123} + PROXYID: ${BASESHIFT_AGENT_ID:?} + PROXYTOKEN: ${BASESHIFT_AGENT_TOKEN:?} + TDS_PORT_RANGE_START: "32768" + TDS_PORT_RANGE_END: "32812" + REPSERVERTAG: master + EXTERNAL_HOSTNAME: localhost + # Docker snapshots are pushed to the LocalStack ECR registry (plain Docker v2 registry over HTTP) + DOCKER_SNAPSHOT_REGISTRY_ENABLED: "true" + DOCKER_SNAPSHOT_REGISTRY_PROVIDER: harbor + DOCKER_SNAPSHOT_REGISTRY: 000000000000.dkr.ecr.us-east-1.localhost.localstack.cloud + DOCKER_SNAPSHOT_REGISTRY_PORT: "4566" + DOCKER_SNAPSHOT_REGISTRY_PROJECT: baseshift + DOCKER_SNAPSHOT_REGISTRY_USE_TLS: "false" + DOCKER_SNAPSHOT_REGISTRY_AUTH_TYPE: basic + DOCKER_SNAPSHOT_REGISTRY_USERNAME: AWS + DOCKER_SNAPSHOT_REGISTRY_PASSWORD: test + # same mounts as the Helm chart (its init container also pre-creates /dubhub/logs/repserver + # and makes the directories world-writable - may be required here as well) + volumes: + - repserver-targets:/dubhub/targets + - repserver-logs:/dubhub/logs + - repserver-certs:/dubhub/certs + ports: + # server clones + - "127.0.0.1:32768-32812:32768-32812" + + agent: + image: public.ecr.aws/x5e2w4i3/proxy:master + restart: unless-stopped + depends_on: + - repserver + extra_hosts: *localstack-hosts + environment: + T: ${BASESHIFT_AGENT_TOKEN:?} + TOKENSERVER: repserver:15057 + GRPCSERVER: repserver:50055 + USETLS: "false" + AWS_REGION: us-east-1 + AWS_DEFAULT_REGION: us-east-1 + PROXY_SOURCE_CONNECT_TIMEOUT_SEC: "8" + PROXY_SOURCE_CONNECT_MAX_DURATION_SEC: "30" + PROXY_SOURCE_CONNECT_RETRY_DELAY_SEC: "3" + REQUESTED_RESTART_EXIT_CODE: "124" + EXIT_CONTAINER_ON_REQUESTED_RESTART: "true" + +volumes: + repserver-targets: + repserver-logs: + repserver-certs: diff --git a/baseshift/demo/baseshift-selfhosted/values-localstack.yaml b/baseshift/demo/baseshift-selfhosted/values-localstack.yaml new file mode 100644 index 0000000..5805735 --- /dev/null +++ b/baseshift/demo/baseshift-selfhosted/values-localstack.yaml @@ -0,0 +1,60 @@ +# Helm values for running the Baseshift self-hosted components (connector + replication server) +# against LocalStack: the source database runs in LocalStack RDS, and Docker snapshots are pushed +# to the LocalStack ECR registry. +# +# Values in <...> are provided by Baseshift (see the Dub setup wizard, or `baseshift dub create +# --deployment self-hosted`). Untested so far - requires access to the replication server image. + +orgId: "" +serverId: "" +enrollmentToken: "" +# encryption password for the Docker snapshot images (letters and numbers only) +password: "localstack123" + +# LocalStack ECR is a plain Docker v2 registry, hence we use the generic ("harbor") provider over +# HTTP, instead of the "aws_ecr" provider (which would call the real AWS ECR APIs) +dockerSnapshotRegistry: + enabled: true + provider: harbor + registry: "000000000000.dkr.ecr.us-east-1.localhost.localstack.cloud" + port: "4566" + project: "baseshift" + useTls: false + auth: + type: basic + # LocalStack ECR does not enforce authentication + username: "AWS" + password: "test" + +repserver: + image: + repository: registry.gitlab.com/dubhub/core/multi-ubuntu2404/repserver + tag: master + # a small demo database needs far less than the production defaults + resources: + requests: + cpu: "1" + memory: "2Gi" + limits: + cpu: "2" + memory: "4Gi" + storage: + size: "10Gi" + +agent: + id: "" + token: "" + resources: + requests: + cpu: "250m" + memory: "512Mi" + limits: + cpu: "1" + memory: "2Gi" + +# pull credentials for registry.gitlab.com (replication server image), provided by Baseshift +imageCredentials: + enabled: true + registry: registry.gitlab.com + username: "" + token: "" diff --git a/baseshift/demo/demo.py b/baseshift/demo/demo.py new file mode 100644 index 0000000..2c95b03 --- /dev/null +++ b/baseshift/demo/demo.py @@ -0,0 +1,458 @@ +""" +Baseshift on LocalStack - end-to-end demo. + +Walks through the flow of masked production data, fully local: + + source - "production" PostgreSQL database in LocalStack RDS, seeded with PII + snapshot - create a masked Docker snapshot of the database, pushed to LocalStack ECR + clone - start a clone from the snapshot image, via the extension's clones API + pipeline - export the clone data to S3, and load it into Snowflake (LocalStack Snowflake emulator) + compare - show the same record in the source, the clone, and Snowflake side by side + +Note: the `snapshot` step is a stand-in for the Baseshift replication server (connector + masking + +snapshot image creation), until the demo can run the Baseshift self-hosted components. Its masking +rules are deliberately simple, and only serve to illustrate the flow. + +Requires LocalStack with the Snowflake emulator (e.g., `lstk` with `type = "snowflake"`) and the +Baseshift extension installed. +""" + +import argparse +import csv +import hashlib +import io +import os +import subprocess +import sys +import tempfile +import time + +import boto3 +import psycopg2 +import requests +import snowflake.connector + +LOCALSTACK_HOST = os.environ.get("LOCALSTACK_HOST", "localhost.localstack.cloud:4566") +ENDPOINT_URL = f"http://{LOCALSTACK_HOST}" +GATEWAY_HOST, GATEWAY_PORT = LOCALSTACK_HOST.split(":") +CLONES_API = f"http://baseshift.{LOCALSTACK_HOST}/clones" + +DB_INSTANCE = "prod-db" +DB_NAME = "shop" +DB_USER = "admin" +DB_PASSWORD = "prod-secret-123" +SNAPSHOT_REPO = "baseshift/shop-dub" +CLONE_NAME = "shop-dev" +EXPORT_BUCKET = "shop-analytics-export" +CUSTOMER_ID = 1 + +CUSTOMERS = [ + ( + "Jane Roe", + "jane.roe@gmail.com", + "078-05-1120", + "+1 415 555 0142", + "221B Baker Street, London", + ), + ( + "Max Mustermann", + "max.mustermann@web.de", + "219-09-9999", + "+49 30 1234567", + "Unter den Linden 1, Berlin", + ), + ( + "Priya Patel", + "priya.patel@outlook.com", + "457-55-5462", + "+1 212 555 0199", + "350 5th Ave, New York", + ), + ( + "Kenji Sato", + "kenji.sato@yahoo.co.jp", + "123-45-6789", + "+81 3 1234 5678", + "1-1 Chiyoda, Tokyo", + ), +] +ORDERS = [(1, 129.99), (1, 42.50), (2, 18.00), (3, 999.00), (3, 12.75), (4, 64.20)] + + +def aws_client(service: str): + return boto3.client( + service, + endpoint_url=ENDPOINT_URL, + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + + +def log(message: str): + print(f"==> {message}", flush=True) + + +# ---------- source: "production" database in RDS ---------- + + +def get_source_endpoint() -> tuple[str, int]: + instance = aws_client("rds").describe_db_instances(DBInstanceIdentifier=DB_INSTANCE) + endpoint = instance["DBInstances"][0]["Endpoint"] + return endpoint["Address"], endpoint["Port"] + + +def connect_source(): + host, port = get_source_endpoint() + return psycopg2.connect( + host=host, port=port, user=DB_USER, password=DB_PASSWORD, dbname=DB_NAME + ) + + +def wait_for_db_instance(): + rds = aws_client("rds") + for _ in range(100): + status = rds.describe_db_instances(DBInstanceIdentifier=DB_INSTANCE)[ + "DBInstances" + ][0] + if status["DBInstanceStatus"] == "available": + return + time.sleep(3) + raise TimeoutError(f"RDS instance {DB_INSTANCE} did not become available") + + +def with_db(connect_fn, work_fn, autocommit: bool = False, retries: int = 30): + """ + Run `work_fn(cursor)` on a fresh connection, retrying on connection errors. Databases may drop + connections while (re)starting (e.g., right after RDS reports an instance as available). Without + autocommit, the work runs in a single transaction, which is rolled back if the connection drops. + """ + for attempt in range(retries): + try: + conn = connect_fn() + try: + conn.autocommit = autocommit + with conn.cursor() as cursor: + result = work_fn(cursor) + if not autocommit: + conn.commit() + return result + finally: + conn.close() + except psycopg2.OperationalError: + if attempt == retries - 1: + raise + time.sleep(2) + + +def cmd_source(_args): + rds = aws_client("rds") + log(f"Creating 'production' RDS PostgreSQL instance {DB_INSTANCE}") + rds.create_db_instance( + DBInstanceIdentifier=DB_INSTANCE, + Engine="postgres", + EngineVersion="17", + DBInstanceClass="db.t3.micro", + MasterUsername=DB_USER, + MasterUserPassword=DB_PASSWORD, + DBName=DB_NAME, + AllocatedStorage=20, + ) + wait_for_db_instance() + + # Baseshift uses logical replication for continuous snapshots. On AWS, this is enabled via the + # `rds.logical_replication` parameter - LocalStack RDS does not apply it (yet), hence we set the + # WAL level directly (the RDS master user is a superuser in LocalStack) and reboot the instance. + log("Enabling logical replication (wal_level=logical)") + with_db( + connect_source, + lambda cursor: cursor.execute("ALTER SYSTEM SET wal_level = 'logical'"), + autocommit=True, + ) + rds.reboot_db_instance(DBInstanceIdentifier=DB_INSTANCE) + time.sleep(3) + wait_for_db_instance() + + def _check_wal_level(cursor): + cursor.execute("SHOW wal_level") + wal_level = cursor.fetchone()[0] + if wal_level != "logical": + # the instance may not have restarted yet - retry + raise psycopg2.OperationalError(f"wal_level is still {wal_level}") + return wal_level + + print(f" wal_level: {with_db(connect_source, _check_wal_level)}") + + log("Seeding customer data (with PII) and orders") + + def _seed(cursor): + cursor.execute( + """ + CREATE TABLE customers ( + id SERIAL PRIMARY KEY, name TEXT, email TEXT, ssn TEXT, phone TEXT, address TEXT + ); + CREATE TABLE orders ( + id SERIAL PRIMARY KEY, customer_id INT REFERENCES customers(id), amount NUMERIC(10, 2) + ); + """ + ) + cursor.executemany( + "INSERT INTO customers (name, email, ssn, phone, address) VALUES (%s, %s, %s, %s, %s)", + CUSTOMERS, + ) + cursor.executemany( + "INSERT INTO orders (customer_id, amount) VALUES (%s, %s)", ORDERS + ) + + with_db(connect_source, _seed) + host, port = get_source_endpoint() + print(f" source database: postgresql://{DB_USER}@{host}:{port}/{DB_NAME}") + + +# ---------- snapshot: masked Docker snapshot in ECR (stand-in for the Baseshift replication server) ---------- + + +def mask_value(column: str, value: str) -> str: + """Deterministic masking - the same input always yields the same output, preserving joins.""" + digest = hashlib.sha256(value.encode()).hexdigest() + if column == "name": + return f"Customer {digest[:6]}" + if column == "email": + return f"user-{digest[:8]}@example.com" + if column == "ssn": + return f"***-**-{value[-4:]}" + if column == "phone": + return f"+1 555 {int(digest[:6], 16) % 10000000:07d}" + if column == "address": + return f"{int(digest[:4], 16) % 999 + 1} Masked Street" + return value + + +def cmd_snapshot(_args): + log( + "Reading source data and applying masking policy (name, email, ssn, phone, address)" + ) + + def _read(cursor): + cursor.execute( + "SELECT id, name, email, ssn, phone, address FROM customers ORDER BY id" + ) + customers = cursor.fetchall() + cursor.execute("SELECT id, customer_id, amount FROM orders ORDER BY id") + return customers, cursor.fetchall() + + customers, orders = with_db(connect_source, _read) + columns = ["name", "email", "ssn", "phone", "address"] + masked = [ + ( + row[0], + *[mask_value(col, val) for col, val in zip(columns, row[1:], strict=True)], + ) + for row in customers + ] + + def sql_literal(value) -> str: + return "'" + str(value).replace("'", "''") + "'" + + init_sql = [ + "CREATE TABLE customers (id INT PRIMARY KEY, name TEXT, email TEXT, ssn TEXT, phone TEXT, address TEXT);", + "CREATE TABLE orders (id INT PRIMARY KEY, customer_id INT REFERENCES customers(id), amount NUMERIC(10, 2));", + ] + init_sql += [ + f"INSERT INTO customers VALUES ({', '.join(map(sql_literal, row))});" + for row in masked + ] + init_sql += [ + f"INSERT INTO orders VALUES ({', '.join(map(sql_literal, row))});" + for row in orders + ] + + ecr = aws_client("ecr") + try: + repo = ecr.create_repository(repositoryName=SNAPSHOT_REPO)["repository"] + except ecr.exceptions.RepositoryAlreadyExistsException: + repo = ecr.describe_repositories(repositoryNames=[SNAPSHOT_REPO])[ + "repositories" + ][0] + image = f"{repo['repositoryUri']}:latest" + + log(f"Building snapshot image {image}") + with tempfile.TemporaryDirectory() as build_dir: + with open(os.path.join(build_dir, "init.sql"), "w") as f: + f.write("\n".join(init_sql) + "\n") + with open(os.path.join(build_dir, "Dockerfile"), "w") as f: + f.write( + "FROM postgres:17\n" + # like Baseshift local clones, the snapshot accepts connections without a password + f"ENV POSTGRES_HOST_AUTH_METHOD=trust POSTGRES_DB={DB_NAME}\n" + "COPY init.sql /docker-entrypoint-initdb.d/\n" + ) + subprocess.run(["docker", "build", "-q", "-t", image, build_dir], check=True) + log("Pushing snapshot image to LocalStack ECR") + subprocess.run(["docker", "push", "-q", image], check=True) + # remove the local tag, so that the clone is actually pulled from the registry + subprocess.run(["docker", "rmi", image], check=True, capture_output=True) + print(f" snapshot image: {image}") + + +# ---------- clone: start a clone from the snapshot, via the extension ---------- + + +def get_clone() -> dict: + return requests.get(f"{CLONES_API}/{CLONE_NAME}").json() + + +def connect_clone(): + clone = get_clone() + host, port = clone["endpoints"][-1].split(":") + return psycopg2.connect(host=host, port=int(port), user="postgres", dbname=DB_NAME) + + +def cmd_clone(_args): + image = aws_client("ecr").describe_repositories(repositoryNames=[SNAPSHOT_REPO])[ + "repositories" + ][0] + image = f"{image['repositoryUri']}:latest" + log(f"Starting clone {CLONE_NAME} from {image}") + response = requests.post(CLONES_API, json={"name": CLONE_NAME, "image": image}) + if response.status_code == 409: + print(f" clone {CLONE_NAME} already exists") + else: + response.raise_for_status() + clone = get_clone() + for _ in range(90): + clone = get_clone() + if clone["status"] in ("running", "failed"): + break + time.sleep(2) + print( + f" clone status: {clone['status']}, endpoints: {', '.join(clone['endpoints'])}" + ) + if clone["status"] != "running": + sys.exit(f"Clone failed to start: {clone.get('error')}") + + +# ---------- pipeline: clone -> S3 -> Snowflake ---------- + + +def connect_snowflake(): + return snowflake.connector.connect( + user="test", + password="test", + account="test", + database="test", + schema="public", + host=f"snowflake.{GATEWAY_HOST}", + port=int(GATEWAY_PORT), + protocol="http", + ) + + +def cmd_pipeline(_args): + s3 = aws_client("s3") + try: + s3.create_bucket(Bucket=EXPORT_BUCKET) + except s3.exceptions.BucketAlreadyOwnedByYou: + pass + + log("Exporting tables from the clone to S3 (CSV)") + + def _export(cursor): + for table in ("customers", "orders"): + cursor.execute(f"SELECT * FROM {table} ORDER BY id") + buffer = io.StringIO() + writer = csv.writer(buffer) + writer.writerow([col.name for col in cursor.description]) + writer.writerows(cursor.fetchall()) + s3.put_object( + Bucket=EXPORT_BUCKET, Key=f"{table}/data.csv", Body=buffer.getvalue() + ) + print(f" s3://{EXPORT_BUCKET}/{table}/data.csv") + + with_db(connect_clone, _export) + + log("Loading the data into Snowflake (S3 stage + COPY INTO)") + conn = connect_snowflake() + cursor = conn.cursor() + statements = [ + "CREATE OR REPLACE TABLE customers (id INT, name TEXT, email TEXT, ssn TEXT, phone TEXT, address TEXT)", + "CREATE OR REPLACE TABLE orders (id INT, customer_id INT, amount NUMBER(10, 2))", + ] + for table in ("customers", "orders"): + statements += [ + f"CREATE OR REPLACE STAGE {table}_stage URL='s3://{EXPORT_BUCKET}/{table}/' " + "CREDENTIALS=(AWS_KEY_ID='test' AWS_SECRET_KEY='test')", + f"COPY INTO {table} FROM @{table}_stage FILE_FORMAT=(TYPE=CSV SKIP_HEADER=1)", + ] + for statement in statements: + cursor.execute(statement) + cursor.execute( + "SELECT c.id, c.name, COUNT(o.id), SUM(o.amount) FROM customers c " + "JOIN orders o ON o.customer_id = c.id GROUP BY c.id, c.name ORDER BY c.id" + ) + print(" revenue per customer (Snowflake):") + for row in cursor.fetchall(): + print(f" {row[0]} {row[1]:<18} orders={row[2]} total={row[3]}") + conn.close() + + +# ---------- compare: the same record at every stage ---------- + + +def cmd_compare(_args): + query = f"SELECT name, email, ssn, phone, address FROM customers WHERE id = {CUSTOMER_ID}" + + def _fetch(cursor): + cursor.execute(query) + return cursor.fetchone() + + source = with_db(connect_source, _fetch) + clone = with_db(connect_clone, _fetch) + sf_conn = connect_snowflake() + snowflake_row = sf_conn.cursor().execute(query).fetchone() + sf_conn.close() + + columns = ["name", "email", "ssn", "phone", "address"] + widths = [8, 28, 28, 28] + header = ["column", "source (RDS)", "Baseshift clone", "Snowflake"] + print(f"\nCustomer {CUSTOMER_ID} at every stage of the pipeline:\n") + print(" ".join(h.ljust(w) for h, w in zip(header, widths, strict=True))) + print(" ".join("-" * w for w in widths)) + for i, column in enumerate(columns): + values = [column, source[i], clone[i], snowflake_row[i]] + print( + " ".join(str(v)[:w].ljust(w) for v, w in zip(values, widths, strict=True)) + ) + leaked = [ + c for i, c in enumerate(columns) if source[i] in (clone[i], snowflake_row[i]) + ] + print() + if leaked: + sys.exit(f"PII leaked into downstream stages: {leaked}") + print("No raw PII values reached the clone or Snowflake.") + + +def cmd_all(args): + for step in (cmd_source, cmd_snapshot, cmd_clone, cmd_pipeline, cmd_compare): + step(args) + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + commands = { + "source": cmd_source, + "snapshot": cmd_snapshot, + "clone": cmd_clone, + "pipeline": cmd_pipeline, + "compare": cmd_compare, + "all": cmd_all, + } + parser.add_argument("command", choices=commands) + args = parser.parse_args() + commands[args.command](args) + + +if __name__ == "__main__": + main() diff --git a/baseshift/demo/requirements.txt b/baseshift/demo/requirements.txt new file mode 100644 index 0000000..a32ee97 --- /dev/null +++ b/baseshift/demo/requirements.txt @@ -0,0 +1,4 @@ +boto3 +psycopg2-binary +requests +snowflake-connector-python diff --git a/baseshift/localstack_baseshift/__init__.py b/baseshift/localstack_baseshift/__init__.py new file mode 100644 index 0000000..95f1bb1 --- /dev/null +++ b/baseshift/localstack_baseshift/__init__.py @@ -0,0 +1 @@ +name = "localstack_baseshift" diff --git a/baseshift/localstack_baseshift/extension.py b/baseshift/localstack_baseshift/extension.py new file mode 100644 index 0000000..00d7f9e --- /dev/null +++ b/baseshift/localstack_baseshift/extension.py @@ -0,0 +1,303 @@ +import logging +import os +import re +import socket +import threading +from dataclasses import dataclass, field + +from localstack import config +from localstack.extensions.api import Extension, http +from localstack.http import Request, Response, route +from localstack.utils.container_utils.container_client import ( + NoSuchImage, + PortMappings, +) +from localstack.utils.docker_utils import DOCKER_CLIENT +from localstack.utils.net import get_addressable_container_host +from localstack.utils.sync import retry +from localstack_extensions.utils.tcp_protocol_router import ( + patch_gateway_for_tcp_routing, + register_tcp_extension, + unregister_tcp_extension, +) +from rolo.routing import RuleAdapter, WithHost + +LOG = logging.getLogger(__name__) + +# Image of the Baseshift clone (Docker snapshot) to start when LocalStack is ready, e.g. +# ".dkr.ecr..amazonaws.com/:latest", or an image in the LocalStack ECR registry +ENV_BASESHIFT_IMAGE = "BASESHIFT_IMAGE" +# Database engine of the clone: "postgres" (default) or "mysql" +ENV_BASESHIFT_DB_TYPE = "BASESHIFT_DB_TYPE" +# Encryption password defined when the Dub was created (passed to clones as PASSWORD) +ENV_BASESHIFT_ENCRYPTION_PASSWORD = "BASESHIFT_ENCRYPTION_PASSWORD" +# Any variable with this prefix is passed to the clone containers with the prefix stripped, +# e.g. BASESHIFT_CLONE_BACKUP_SCHEDULE -> BACKUP_SCHEDULE +ENV_CLONE_PREFIX = "BASESHIFT_CLONE_" + +API_HOST = "baseshift." +DEFAULT_CLONE_NAME = "default" +DB_PORTS = {"postgres": 5432, "mysql": 3306} +# host ports for additional clones, if the default port of the engine is already taken +EXTRA_PORT_RANGE = range(15432, 15532) +CLONE_NAME_REGEX = re.compile(r"^[a-z0-9]([a-z0-9-]{0,38}[a-z0-9])?$") +# clones may need a while to start up from the snapshot +STARTUP_RETRIES = 120 +STARTUP_SLEEP = 2 + + +def is_postgres_handshake(data: bytes) -> bool: + """ + Identify PostgreSQL connections by protocol handshake: the first message is either + an SSL request (code 80877103) or a startup message with protocol version 3.0. + """ + if len(data) < 8: + return False + return data[4:8] in (b"\x04\xd2\x16\x2f", b"\x00\x03\x00\x00") + + +@dataclass +class Clone: + name: str + image: str + db_type: str + port: int + env_vars: dict[str, str] = field(default_factory=dict) + status: str = "starting" + error: str | None = None + + @property + def container_name(self) -> str: + return f"ls-baseshift-clone-{self.name}" + + @property + def gateway_routed(self) -> bool: + # only PostgreSQL can be detected on the shared gateway port (MySQL is a server-first + # protocol), and only the clone on the default port is routed through the gateway + return self.db_type == "postgres" and self.port == DB_PORTS["postgres"] + + def to_dict(self) -> dict: + result = { + "name": self.name, + "image": self.image, + "dbType": self.db_type, + "status": self.status, + "hostPort": self.port, + "endpoints": [f"localhost:{self.port}"], + } + if self.gateway_routed: + # the gateway endpoint as seen by clients (e.g., a custom port configured via LOCALSTACK_HOST) + result["endpoints"].insert(0, config.LOCALSTACK_HOST.host_and_port()) + if self.error: + result["error"] = self.error + return result + + +class ClonesApi: + """ + HTTP API to manage clones, similar to the `baseshift server-clone` CLI commands. + Note: handlers receive the matched host placeholders (e.g., `domain`) as keyword arguments. + """ + + def __init__(self, extension: "BaseshiftExtension"): + self.extension = extension + + @route("/clones", methods=["GET"]) + def list_clones(self, request: Request, **kwargs): + return {"clones": [clone.to_dict() for clone in self.extension.list_clones()]} + + @route("/clones", methods=["POST"]) + def create_clone(self, request: Request, **kwargs): + try: + payload = request.get_json(force=True) or {} + clone = self.extension.start_clone( + name=payload.get("name", ""), + image=payload.get("image", ""), + db_type=payload.get("dbType", "postgres"), + env_vars=payload.get("env") or {}, + ) + except ValueError as e: + return Response.for_json({"error": str(e)}, status=400) + except KeyError as e: + return Response.for_json({"error": str(e.args[0])}, status=409) + return Response.for_json(clone.to_dict(), status=202) + + @route("/clones/", methods=["GET"]) + def get_clone(self, request: Request, name: str, **kwargs): + if clone := self.extension.get_clone(name): + return clone.to_dict() + return Response.for_json({"error": f"Clone {name} not found"}, status=404) + + @route("/clones/", methods=["DELETE"]) + def delete_clone(self, request: Request, name: str, **kwargs): + if self.extension.stop_clone(name): + return Response(status=204) + return Response.for_json({"error": f"Clone {name} not found"}, status=404) + + +class BaseshiftExtension(Extension): + name = "localstack-baseshift" + + def __init__(self): + self.default_image = os.getenv(ENV_BASESHIFT_IMAGE, "").strip() + self.default_db_type = self._validate_db_type( + os.getenv(ENV_BASESHIFT_DB_TYPE, "postgres") + ) + self.base_env_vars = { + key.removeprefix(ENV_CLONE_PREFIX): value + for key, value in os.environ.items() + if key.startswith(ENV_CLONE_PREFIX) and key != ENV_CLONE_PREFIX + } + if password := os.getenv(ENV_BASESHIFT_ENCRYPTION_PASSWORD): + self.base_env_vars["PASSWORD"] = password + self.container_host = get_addressable_container_host() + self._clones: dict[str, Clone] = {} + self._lock = threading.RLock() + + def update_gateway_routes(self, router: http.Router[http.RouteHandler]): + router.add(WithHost(API_HOST, [RuleAdapter(ClonesApi(self))])) + patch_gateway_for_tcp_routing() + + def on_platform_ready(self): + url = f"http://baseshift.{config.LOCALSTACK_HOST.host_and_port()}/clones" + if not self.default_image: + LOG.info( + "Baseshift extension ready. %s is not set - start clones via the API: %s", + ENV_BASESHIFT_IMAGE, + url, + ) + return + # start the clone after LocalStack is ready - the image may be served by the LocalStack ECR registry + self.start_clone(DEFAULT_CLONE_NAME, self.default_image, self.default_db_type) + LOG.info( + "Baseshift extension ready, starting default clone. Manage clones via: %s", + url, + ) + + def on_platform_shutdown(self): + for clone in self.list_clones(): + self._remove_container(clone) + + # clone management + + def list_clones(self) -> list[Clone]: + with self._lock: + return list(self._clones.values()) + + def get_clone(self, name: str) -> Clone | None: + with self._lock: + return self._clones.get(name) + + def start_clone( + self, + name: str, + image: str, + db_type: str = "postgres", + env_vars: dict | None = None, + ) -> Clone: + if not CLONE_NAME_REGEX.match(name or ""): + raise ValueError( + "Clone name must consist of lowercase letters, digits, and dashes (max. 40 chars)" + ) + if not image: + raise ValueError("Clone image must be specified") + db_type = self._validate_db_type(db_type) + with self._lock: + if name in self._clones: + raise KeyError(f"Clone {name} already exists") + clone = Clone( + name=name, + image=image, + db_type=db_type, + port=self._allocate_port(db_type), + env_vars={ + **self.base_env_vars, + **{k: str(v) for k, v in env_vars.items()}, + } + if env_vars + else dict(self.base_env_vars), + ) + self._clones[name] = clone + threading.Thread(target=self._run_clone, args=(clone,), daemon=True).start() + return clone + + def stop_clone(self, name: str) -> bool: + with self._lock: + clone = self._clones.pop(name, None) + if not clone: + return False + self._remove_container(clone) + return True + + def _run_clone(self, clone: Clone): + try: + self._pull_image(clone.image) + ports = PortMappings() + ports.add(clone.port, DB_PORTS[clone.db_type]) + DOCKER_CLIENT.run_container( + clone.image, + name=clone.container_name, + detach=True, + remove=True, + ports=ports, + env_vars=clone.env_vars or None, + ) + + def _check_port(): + with socket.create_connection( + (self.container_host, clone.port), timeout=2 + ): + pass + + retry(_check_port, retries=STARTUP_RETRIES, sleep=STARTUP_SLEEP) + if clone.gateway_routed: + register_tcp_extension( + extension_name=self._tcp_route_name(clone), + matcher=is_postgres_handshake, + backend_host=self.container_host, + backend_port=clone.port, + ) + clone.status = "running" + LOG.info( + "Baseshift clone %s is running: %s", + clone.name, + clone.to_dict()["endpoints"], + ) + except Exception as e: + LOG.warning("Failed to start Baseshift clone %s: %s", clone.name, e) + clone.status = "failed" + clone.error = str(e) + self._remove_container(clone) + + def _pull_image(self, image: str): + try: + DOCKER_CLIENT.inspect_image(image, pull=False) + except NoSuchImage: + LOG.info("Pulling Baseshift clone image %s", image) + DOCKER_CLIENT.pull_image(image) + + def _remove_container(self, clone: Clone): + if clone.gateway_routed: + unregister_tcp_extension(self._tcp_route_name(clone)) + DOCKER_CLIENT.remove_container( + clone.container_name, force=True, check_existence=False + ) + + def _allocate_port(self, db_type: str) -> int: + used_ports = {clone.port for clone in self._clones.values()} + for port in [DB_PORTS[db_type], *EXTRA_PORT_RANGE]: + if port not in used_ports: + return port + raise ValueError("No free port available for another clone") + + def _tcp_route_name(self, clone: Clone) -> str: + return f"{self.name}-{clone.name}" + + @staticmethod + def _validate_db_type(db_type: str) -> str: + db_type = (db_type or "").strip().lower() + if db_type not in DB_PORTS: + raise ValueError( + f"Database type must be one of {sorted(DB_PORTS)}, got: {db_type}" + ) + return db_type diff --git a/baseshift/pyproject.toml b/baseshift/pyproject.toml new file mode 100644 index 0000000..0ba7b41 --- /dev/null +++ b/baseshift/pyproject.toml @@ -0,0 +1,41 @@ +[build-system] +requires = ["setuptools", "wheel", "plux>=1.3.1"] +build-backend = "setuptools.build_meta" + +[project] +name = "localstack-baseshift" +version = "0.1.0" +description = "Baseshift Extension for LocalStack" +readme = {file = "README.md", content-type = "text/markdown; charset=UTF-8"} +requires-python = ">=3.9" +authors = [ + { name = "LocalStack team"} +] +keywords = ["LocalStack", "Baseshift", "PostgreSQL", "MySQL"] +classifiers = [] +dependencies = [ + "localstack-extensions-utils" +] + +[project.urls] +Homepage = "https://github.com/localstack/localstack-extensions/tree/main/baseshift" + +[project.optional-dependencies] +dev = [ + "boto3", + "build", + "jsonpatch", + "localstack-core", + "psycopg2-binary", + "requests", + "pytest", + "rolo", + "ruff", + "twisted" +] + +[tool.setuptools] +packages = ["localstack_baseshift"] + +[project.entry-points."localstack.extensions"] +localstack-baseshift = "localstack_baseshift.extension:BaseshiftExtension" diff --git a/baseshift/tests/__init__.py b/baseshift/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/baseshift/tests/test_extension.py b/baseshift/tests/test_extension.py new file mode 100644 index 0000000..ec2cfce --- /dev/null +++ b/baseshift/tests/test_extension.py @@ -0,0 +1,166 @@ +import os + +import boto3 +import psycopg2 +import pytest +import requests +from localstack.utils.docker_utils import DOCKER_CLIENT +from localstack.utils.strings import short_uid +from localstack.utils.sync import retry + +# Note: these tests connect to databases served by Baseshift clone containers. In CI, a plain +# Postgres image is used as a stand-in for a clone image (clone images are private to each Baseshift +# customer), and LocalStack is started with: +# BASESHIFT_IMAGE=postgres:17 BASESHIFT_CLONE_POSTGRES_HOST_AUTH_METHOD=trust +# To run the tests against a real clone, override the user/database via the env variables below. + +USER = os.environ.get("BASESHIFT_TEST_USER", "postgres") +DATABASE = os.environ.get("BASESHIFT_TEST_DATABASE", "postgres") +STAND_IN_IMAGE = "postgres:17" + +GATEWAY_PORT = int(os.environ.get("GATEWAY_PORT", "4566")) +GATEWAY_ENDPOINT = ("localhost.localstack.cloud", GATEWAY_PORT) +HOST_PORT_ENDPOINT = ("localhost", 5432) +API_URL = f"http://baseshift.localhost.localstack.cloud:{GATEWAY_PORT}/clones" + + +def connect(host: str, port: int): + def _connect(): + return psycopg2.connect(host=host, port=port, user=USER, dbname=DATABASE) + + return retry(_connect, retries=15, sleep=2.0) + + +def wait_for_clone_status(name: str, status: str = "running") -> dict: + def _check(): + clone = requests.get(f"{API_URL}/{name}").json() + assert clone["status"] == status, clone + return clone + + return retry(_check, retries=60, sleep=2) + + +@pytest.fixture +def ecr_image(): + """Push the stand-in image into the LocalStack ECR registry, emulating a Docker snapshot of a Dub.""" + ecr = boto3.client( + "ecr", + endpoint_url=f"http://localhost.localstack.cloud:{GATEWAY_PORT}", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + repo_name = f"dub-snapshots/dub-{short_uid()}" + repo_uri = ecr.create_repository(repositoryName=repo_name)["repository"][ + "repositoryUri" + ] + image = f"{repo_uri}:latest" + DOCKER_CLIENT.tag_image(STAND_IN_IMAGE, image) + DOCKER_CLIENT.push_image(image) + # remove the local tag, to make sure the extension pulls the image from the registry + DOCKER_CLIENT.remove_image(image) + yield image + ecr.delete_repository(repositoryName=repo_name, force=True) + + +@pytest.mark.parametrize( + "endpoint", [GATEWAY_ENDPOINT, HOST_PORT_ENDPOINT], ids=["gateway", "host-port"] +) +def test_query_default_clone(endpoint): + with connect(*endpoint) as conn, conn.cursor() as cursor: + cursor.execute("SELECT version()") + assert "PostgreSQL" in cursor.fetchone()[0] + + +def test_default_clone_is_writable(): + table = f"test_{short_uid()}" + with connect(*GATEWAY_ENDPOINT) as conn, conn.cursor() as cursor: + cursor.execute(f"CREATE TABLE {table} (id SERIAL PRIMARY KEY, name TEXT)") + cursor.execute( + f"INSERT INTO {table} (name) VALUES (%s), (%s)", ("alice", "bob") + ) + cursor.execute(f"SELECT name FROM {table} ORDER BY id") + assert [row[0] for row in cursor.fetchall()] == ["alice", "bob"] + cursor.execute(f"DROP TABLE {table}") + + +def test_concurrent_connections(): + connections = [connect(*GATEWAY_ENDPOINT) for _ in range(3)] + try: + for i, conn in enumerate(connections): + with conn.cursor() as cursor: + cursor.execute("SELECT %s", (i,)) + assert cursor.fetchone()[0] == i + finally: + for conn in connections: + conn.close() + + +def test_list_clones(): + clones = requests.get(API_URL).json()["clones"] + default = next(clone for clone in clones if clone["name"] == "default") + assert default["status"] == "running" + assert default["hostPort"] == 5432 + assert default["image"] == STAND_IN_IMAGE + + +def test_start_clone_from_ecr_image(ecr_image): + name = f"pr-{short_uid()}" + response = requests.post( + API_URL, + json={ + "name": name, + "image": ecr_image, + "env": {"POSTGRES_HOST_AUTH_METHOD": "trust"}, + }, + ) + assert response.status_code == 202 + try: + clone = wait_for_clone_status(name) + # the default clone owns the gateway and the default port, hence a separate host port is used + assert clone["hostPort"] != 5432 + assert clone["endpoints"] == [f"localhost:{clone['hostPort']}"] + + with connect("localhost", clone["hostPort"]) as conn, conn.cursor() as cursor: + cursor.execute("SELECT 1") + assert cursor.fetchone()[0] == 1 + finally: + assert requests.delete(f"{API_URL}/{name}").status_code == 204 + + assert requests.get(f"{API_URL}/{name}").status_code == 404 + assert not DOCKER_CLIENT.is_container_running(f"ls-baseshift-clone-{name}") + + +def test_start_clone_with_invalid_image(): + name = f"broken-{short_uid()}" + response = requests.post( + API_URL, json={"name": name, "image": "localhost:1/does-not-exist"} + ) + assert response.status_code == 202 + try: + clone = wait_for_clone_status(name, "failed") + assert clone["error"] + finally: + requests.delete(f"{API_URL}/{name}") + + +def test_api_errors(): + # invalid name + response = requests.post( + API_URL, json={"name": "Invalid_Name", "image": STAND_IN_IMAGE} + ) + assert response.status_code == 400 + # missing image + response = requests.post(API_URL, json={"name": "no-image"}) + assert response.status_code == 400 + # invalid database type + response = requests.post( + API_URL, json={"name": "mongo", "image": STAND_IN_IMAGE, "dbType": "mongodb"} + ) + assert response.status_code == 400 + # duplicate name + response = requests.post(API_URL, json={"name": "default", "image": STAND_IN_IMAGE}) + assert response.status_code == 409 + # unknown clone + assert requests.get(f"{API_URL}/unknown").status_code == 404 + assert requests.delete(f"{API_URL}/unknown").status_code == 404