Skip to content

test: add KMS Encryption plugin integration tests #1349

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
JuanLeee wants to merge 1 commit into
base: feat/kms-stack-4-plugin
Choose a base branch
from
Draft

test: add KMS Encryption plugin integration tests #1349

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
248 changes: 248 additions & 0 deletions .github/workflows/kms_encryption_integration_tests.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,248 @@
name: KMS Encryption Integration Tests

# The LocalStack job needs no AWS account and no secrets, so it also runs on pull requests from forks.
# The real-key and Aurora jobs read the key ARN, its region and the roles from repository variables,
# which an administrator sets up once; CI never creates or changes KMS keys.

on:
pull_request:
paths:
- 'aws_advanced_python_wrapper/**'
- 'tests/integration/**'
- 'pyproject.toml'
- 'poetry.lock'
- '.github/workflows/kms_encryption_integration_tests.yml'
push:
branches: [ main ]
workflow_dispatch:
inputs:
lane:
description: 'Which KMS environment to run'
type: choice
options: [ localstack, real-kms, aurora ]
default: localstack

permissions:
contents: read

concurrency:
group: ${{ github.workflow }}-${{ github.ref }}-${{ inputs.lane }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}

jobs:
kms-localstack:
name: Run KMS Encryption Tests with LocalStack
if: github.event_name != 'workflow_dispatch' || inputs.lane == 'localstack'
runs-on: ubuntu-latest
timeout-minutes: 45
strategy:
fail-fast: false
matrix:
include:
- { engine: pg, skip: mysql }
- { engine: mysql, skip: pg }

steps:
- name: 'Clone repository'
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6

- name: 'Set up JDK 8'
uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5
with:
distribution: 'corretto'
java-version: 8

- name: 'Set up Python 3.12'
uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'

- name: Install poetry
shell: bash
run: |
pipx install poetry==1.8.2
poetry config virtualenvs.prefer-active-python true

- name: Install dependencies
run: poetry install

- name: 'Run the KMS tests on Docker with LocalStack'
run: |
DB_PASSWORD="$(openssl rand -hex 16)"
echo "::add-mask::${DB_PASSWORD}"
export DB_PASSWORD DB_USERNAME=test_user DB_DATABASE_NAME=test_db
./gradlew --no-parallel --no-daemon test-kms-encryption-docker \
-Dexclude-${{ matrix.skip }}-engine=true -Dexclude-${{ matrix.skip }}-driver=true

# pytest exits 0 when every test is skipped, so count the KMS test cases that ran.
- name: 'Fail when no KMS test ran'
run: |
python3 -c 'import glob,sys,xml.etree.ElementTree as E; c=[t for f in glob.glob("tests/integration/container/reports/*.xml") for t in E.parse(f).iter("testcase") if "kms_encryption" in t.get("classname","") and t.find("skipped") is None]; print("KMS tests run:", len(c)); sys.exit(0 if c else 1)'

- name: 'Archive results'
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: kms-localstack-${{ matrix.engine }}
path: ./tests/integration/container/reports
retention-days: 5

kms-real-key:
name: Run KMS Encryption Tests with a real KMS key
if: (github.event_name == 'push' && github.ref == 'refs/heads/main') || inputs.lane == 'real-kms'
runs-on: ubuntu-latest
timeout-minutes: 45
permissions:
id-token: write # This is required for requesting the JWT
contents: read # This is required for actions/checkout
strategy:
fail-fast: false
matrix:
include:
- { engine: pg, skip: mysql }
- { engine: mysql, skip: pg }

steps:
- name: 'Clone repository'
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6

- name: 'Set up JDK 8'
uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5
with:
distribution: 'corretto'
java-version: 8

- name: 'Set up Python 3.12'
uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'

- name: Install poetry
shell: bash
run: |
pipx install poetry==1.8.2
poetry config virtualenvs.prefer-active-python true

- name: Install dependencies
run: poetry install

- name: 'Configure AWS Credentials'
uses: aws-actions/configure-aws-credentials@d979d5b3a71173a29b74b5b88418bfda9437d885 # v6
with:
role-to-assume: ${{ vars.KMS_IT_ROLE_ARN }}
role-session-name: python_kms_integration_tests
role-duration-seconds: 3600
aws-region: ${{ vars.KMS_REGION }}

- name: 'Run the KMS tests on Docker with a real KMS key'
run: |
DB_PASSWORD="$(openssl rand -hex 16)"
echo "::add-mask::${DB_PASSWORD}"
export DB_PASSWORD DB_USERNAME=test_user DB_DATABASE_NAME=test_db
./gradlew --no-parallel --no-daemon test-kms-encryption-docker-kms \
-Dexclude-${{ matrix.skip }}-engine=true -Dexclude-${{ matrix.skip }}-driver=true
env:
KMS_KEY_ARN: ${{ vars.KMS_KEY_ARN }}
KMS_REGION: ${{ vars.KMS_REGION }}
KMS_DECRYPT_ONLY_ROLE_ARN: ${{ vars.KMS_DECRYPT_ONLY_ROLE_ARN }}

- name: 'Fail when no KMS test ran'
run: |
python3 -c 'import glob,sys,xml.etree.ElementTree as E; c=[t for f in glob.glob("tests/integration/container/reports/*.xml") for t in E.parse(f).iter("testcase") if "kms_encryption" in t.get("classname","") and t.find("skipped") is None]; print("KMS tests run:", len(c)); sys.exit(0 if c else 1)'

- name: 'Archive results'
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: kms-real-key-${{ matrix.engine }}
path: ./tests/integration/container/reports
retention-days: 5

kms-aurora:
name: Run KMS Encryption Tests on Aurora
if: inputs.lane == 'aurora'
runs-on: ubuntu-latest
timeout-minutes: 150
permissions:
id-token: write # This is required for requesting the JWT
contents: read # This is required for actions/checkout
strategy:
fail-fast: false
matrix:
include:
- { engine: pg, skip: mysql }
- { engine: mysql, skip: pg }
concurrency:
group: kms-aurora-${{ matrix.engine }}
cancel-in-progress: false
env:
# The harness doesn't delete its cluster when a run is cancelled, so the last step deletes this one by name.
RDS_DB_NAME: kms-it-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.engine }}

steps:
- name: 'Clone repository'
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6

- name: 'Set up JDK 8'
uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5
with:
distribution: 'corretto'
java-version: 8

- name: 'Set up Python 3.12'
uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'

- name: Install poetry
shell: bash
run: |
pipx install poetry==1.8.2
poetry config virtualenvs.prefer-active-python true

- name: Install dependencies
run: poetry install

- name: 'Configure AWS Credentials'
uses: aws-actions/configure-aws-credentials@d979d5b3a71173a29b74b5b88418bfda9437d885 # v6
with:
role-to-assume: arn:aws:iam::${{ secrets.AWS_ACCOUNT_ID }}:role/${{ secrets.AWS_DEPLOY_ROLE }}
role-session-name: python_integration_tests
role-duration-seconds: 21600
aws-region: ${{ secrets.AWS_DEFAULT_REGION }}

- name: 'Run the KMS tests on Aurora'
run: |
DB_PASSWORD="$(openssl rand -hex 16)"
echo "::add-mask::${DB_PASSWORD}"
export DB_PASSWORD DB_USERNAME=test_user DB_DATABASE_NAME=test_db
./gradlew --no-parallel --no-daemon test-kms-encryption-aurora \
-Dexclude-${{ matrix.skip }}-engine=true -Dexclude-${{ matrix.skip }}-driver=true
env:
NUM_INSTANCES: 2
RDS_DB_REGION: ${{ secrets.AWS_DEFAULT_REGION }}
KMS_KEY_ARN: ${{ vars.KMS_KEY_ARN }}
KMS_REGION: ${{ vars.KMS_REGION }}

- name: 'Fail when no KMS test ran'
run: |
python3 -c 'import glob,sys,xml.etree.ElementTree as E; c=[t for f in glob.glob("tests/integration/container/reports/*.xml") for t in E.parse(f).iter("testcase") if "kms_encryption" in t.get("classname","") and t.find("skipped") is None]; print("KMS tests run:", len(c)); sys.exit(0 if c else 1)'

- name: 'Archive results'
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: kms-aurora-${{ matrix.engine }}
path: ./tests/integration/container/reports
retention-days: 5

# Deletes the cluster named RDS_DB_NAME and its member instances, never clusters matched by a name prefix.
- name: 'Delete the cluster of this run'
if: always()
run: |
instances="$(aws rds describe-db-clusters --db-cluster-identifier "${RDS_DB_NAME}" \
--query 'DBClusters[0].DBClusterMembers[].DBInstanceIdentifier' --output text 2>/dev/null)" || exit 0
for instance in ${instances}; do
aws rds delete-db-instance --db-instance-identifier "${instance}" --skip-final-snapshot > /dev/null || true
done
aws rds delete-db-cluster --db-cluster-identifier "${RDS_DB_NAME}" --skip-final-snapshot > /dev/null || true
40 changes: 34 additions & 6 deletions docs/development-guide/IntegrationTests.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,10 @@ unset FILTER # Done testing the IAM tests, unset FILTER

## Running async integration tests

The wrapper ships async counterparts to every sync integration test file. Async tests exercise `AsyncAwsWrapperConnection` (raw) and `create_async_engine` with the wrapper's dialects (`postgresql+aws_wrapper_psycopg` — shared with sync via `get_async_dialect_cls`; `mysql+aws_wrapper_aiomysql`). They are invoked via dedicated Gradle tasks, independent of the sync tasks:
The wrapper ships an async counterpart for each sync test file at the top level of `tests/integration/container`, except the KMS Encryption Plugin's files, because the plugin is sync only.
The `sqlalchemy/` and `django/` folders have no async counterparts.
Async tests exercise `AsyncAwsWrapperConnection` (raw) and `create_async_engine` with the wrapper's dialects: `postgresql+aws_wrapper_psycopg`, shared with sync through `get_async_dialect_cls`, and `mysql+aws_wrapper_aiomysql`.
They run through their own Gradle tasks, separate from the sync tasks:

| Deployment | Engine | Sync task | Async task |
|---|---|---|---|
Expand All @@ -85,6 +88,29 @@ Async tests do not introduce any new pytest dependencies — they follow the rep
2. **Cold-start latency.** Serverless v2 can take 30–60 seconds to scale from 0 ACUs on the first connection. Warm the cluster with a trivial query (`SELECT 1`) before running timing-sensitive tests so that warmup cost isn't charged against a test's 600-second timeout.
3. **Topology-gated tests skip environmentally.** Multi-instance tests (failover, some read/write splitting assertions) need ≥2 instances to execute; they skip on single-writer topologies. Blue/Green and Global Database tests require specific engine versions / cluster configurations. Tests skipped under these conditions are environmental, not regressions — they are equivalent to the sync suite's existing behavior under the same cluster topology.

## Running the KMS encryption tests

The KMS Encryption Plugin's tests are the `test_kms_encryption_*.py` files in `tests/integration/container`, plus `sqlalchemy/test_sqlalchemy_kms_encryption.py` and `django/test_django_kms_encryption.py`.
Each environment has its own Gradle task.

| Task | Databases | KMS | AWS credentials |
|---|---|---|---|
| `test-kms-encryption-docker` | Docker | LocalStack | None |
| `test-kms-encryption-docker-kms` | Docker | A real key | Limited to the test key |
| `test-kms-encryption-aurora` | Aurora | A real key | As for the Aurora tasks |

The LocalStack task runs every module except `test_kms_encryption_real_kms.py`, which runs only on the real-key Docker task, and `test_kms_encryption_aurora.py`, which runs only on the Aurora task.
The LocalStack task needs no AWS account.
The harness gives the test container dummy credentials and passes LocalStack's address to the tests as `KMS_ENDPOINT_URL`.
The KMS tests skip outside these tasks, and fail inside them when the endpoint or the key is missing.
The real-key tasks read the key ARN from `KMS_KEY_ARN`, its region from `KMS_REGION`, and the ARN of a role that may only decrypt from `KMS_DECRYPT_ONLY_ROLE_ARN`.
Use full key ARNs, not aliases.
The tests create keys only in LocalStack, so an administrator creates the real test key, once.
A key costs $1 a month, and a key scheduled for deletion waits 7 to 30 days before AWS deletes it.
Give the test credentials only `kms:GenerateDataKey`, `kms:Decrypt` and `kms:DescribeKey` on the test key, and `sts:AssumeRole` on the decrypt-only role.
Reports keep failure output, which can include the database password, so use throwaway database passwords and keep reports private.
To run the same tests against local servers and LocalStack without the harness, use `tests/integration/kms_local/run.sh`.

## Debugging Aurora Integration Tests - Pycharm

> [!WARNING]\
Expand Down Expand Up @@ -144,13 +170,15 @@ Async tests do not introduce any new pytest dependencies — they follow the rep
| `RDS_DB_NAME` | Yes | The database identifier for your Aurora or RDS cluster. Must be a unique value to avoid conflicting with existing clusters. | `db-identifier` |
| `RDS_DB_DOMAIN` | No | The existing database connection suffix. Use this variable to run against an existing database. | `XYZ.us-east-2.rds.amazonaws.com` |
| `IAM_USER` | No | User within the database that is identified with AWSAuthenticationPlugin. This is used for AWS IAM Authentication and is optional | `example_user_name` |
| `AWS_ACCESS_KEY_ID` | Yes | An AWS access key associated with an IAM user or role with RDS permissions. | `ASIAIOSFODNN7EXAMPLE` |
| `AWS_SECRET_ACCESS_KEY` | Yes | The secret key associated with the provided AWS_ACCESS_KEY_ID. | `wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY` |
| `AWS_ACCESS_KEY_ID` | Aurora and RDS only | An AWS access key associated with an IAM user or role with RDS permissions. | `ASIAIOSFODNN7EXAMPLE` |
| `AWS_SECRET_ACCESS_KEY` | Aurora and RDS only | The secret key associated with the provided AWS_ACCESS_KEY_ID. | `wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY` |
| `AWS_SESSION_TOKEN` | No | AWS Session Token for CLI, SDK, & API access. This value is for MFA credentials only. See: [temporary AWS credentials](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_temp_use-resources.html). | `AQoDYXdzEJr...<remainder of session token>` | |
| `REUSE_RDS_DB` | Yes | Set to true if you would like to use an existing cluster for your tests. | `false` |
| `RDS_DB_REGION` | Yes | The database region. | `us-east-2` |
| `DEBUG_ENV` | No | The IDE you will be using to debug the tests, values are either `PYCHARM` or `VSCODE` | `PYCHARM` |

Docker tasks don't read the AWS credential variables, except the KMS task that uses a real key (see [Running the KMS encryption tests](#running-the-kms-encryption-tests)).

## Managing Environment Variables

After following steps 1-3 from the [Run Integration Tests](#running-aurora-integration-tests) section, your default
Expand Down Expand Up @@ -236,6 +264,6 @@ can also be useful.

## Viewing test results

The test results can be viewed in the test report generated at
`tests/integration/container/reports/integration_tests.html`. The logs can be viewed from the console where you executed
the tests or by clicking `show details` on a given test in the test report.
Each run writes its test report to `tests/integration/container/reports/<environment>.html`, with a JUnit XML report, `<environment>.xml`, beside it.
`<environment>` names the deployment, the engine, the number of instances and the Python version.
The logs are in the console where you ran the tests, and under `show details` for each test in the report.
2 changes: 2 additions & 0 deletions docs/using-the-python-wrapper/using-plugins/UsingTheKmsEncryptionPlugin.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -344,6 +344,8 @@ The trigger checks a new value against the column's current HMAC key, and rotati
Three controls decide how much the encryption protects: which KMS keys the application can use, who can write the metadata schema, and who can change the trigger. The first three steps cover them, and the fourth covers what none of them stops.

- **KMS keys.** The plugin takes a column's master key from its `key_storage` row, not from its own settings, so the allow-list and the application's IAM policy keep a changed row from pointing it at another key. List every master key the columns use in `encryption_allowed_master_key_arns`, in the form `key_storage.master_key_arn` holds. The plugin refuses any other key before calling KMS (`UnauthorizedMasterKeyError`), on reads and writes alike, and `encryption_return_unverified_data` doesn't change that. Add a new master key before rotating a column onto it, and keep the old one listed, with `kms:Decrypt` allowed, while values written under it remain. Give the application's credentials only `kms:Decrypt`, on those key ARNs, never `"Resource": "*"`, and use a separate key per environment. Run `KeyManagementUtility`, which also calls `kms:GenerateDataKey`, `kms:CreateKey`, `kms:CreateAlias` and `kms:DescribeKey`, under a separate administrative credential. Don't add a `kms:ViaService` condition to these policies. It matches only requests an AWS service makes on your behalf, and the plugin calls KMS directly, so every `Decrypt` would be denied.
Store key ARNs, not alias ARNs, in `key_storage` and in `encryption_allowed_master_key_arns`.
Repointing an alias makes every value under the old key unreadable.

```json
{"Version": "2012-10-17", "Statement": [{"Effect": "Allow", "Action": "kms:Decrypt",
Expand Down
Loading
Loading