Python SR/MR cluster management

Author: Mitchell Elholm

.github/workflows/python-cm-integ-tests.yml

@@ -55,6 +55,8 @@ jobs:
 
     - name: Configure and run integration for cluster management
       working-directory: ./python/cluster_management
+      env:
+        IS_CI: "TRUE"
       run: |
         python3 -m venv cm-integ
         source cm-integ/bin/activate

python/cluster_management/README.md

@@ -2,29 +2,61 @@
 
 ## Overview
 
-The code examples in this topic show you how to use the AWS Python SDK with DSQL to create, read, update, and delete clusters.
+The code examples in this topic show you how to use the AWS Python SDK with DSQL
+to create, update, get, and delete single- and multi-Region clusters.
+
+Each file in the [/src](src) directory demonstrates a minimum
+working example for each operation. The example function for each operation is invoked
+in [`test_dsql_cluster_management.py`](test/test_dsql_cluster_management.py).
 
 ## Run the examples
 
+### ⚠️ Important
+
+* Running this code might result in charges to your AWS account.
+* We recommend that you grant your code least privilege. At most, grant only the
+  minimum permissions required to perform the task. For more information, see
+  [Grant least privilege](https://docs.aws.amazon.com/IAM/latest/UserGuide/best-practices.html#grant-least-privilege).
+* This code is not tested in every AWS Region. For more information, see
+  [AWS Regional Services](https://aws.amazon.com/about-aws/global-infrastructure/regional-product-services).
+
 ### Prerequisites
 
-* python version >=3.8.0 is needed
+- Python version >= 3.8.0 is installed.
+- Valid AWS credentials can be discovered by the [default provider chain](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html).
 
-### Setup test running environment 
+### Execute tests to create and delete clusters
 
 ```sh
-source setup.sh
+# Optional: Single-Region examples will execute in REGION_1. Defaults to 'us-east-1'.
+export REGION_1="us-east-1"
+
+# Optional: Multi-Region examples will create clusters in REGION_1 and REGION_2
+# with WITNESS_REGION as witness for both. Defaults to 'us-east-2' for REGION_2
+# and 'us-west-2' for WITNESS_REGION.
+export REGION_2="us-east-2"
+export WITNESS_REGION="us-west-2"
+
+python3 -m venv .venv && source .venv/bin/activate
+pip install -r requirements.txt
+
+# Will create, update, read, then delete clusters, use -s to see print statements when running tests
+pytest test/test_dsql_cluster_management.py [-s]
 ```
 
-### Run the example tests
+Test execution will take around five minutes as it waits for clusters to complete activation and deletion.
 
-```sh
-# Use the account credentials dedicated for python
-pytest test/test_example.py
+### Executing single operations
+Files in [src/](src/) each have a `main()` function that let you exercise single operations.
+
+```shell
+# Check each operation for its expected environment variables
+REGION_1="us-east-1" CLUSTER_ID_1="<your cluster id>" \
+  python src/get_cluster.py
 ```
 
 ---
 
-Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
 SPDX-License-Identifier: MIT-0

python/cluster_management/pyproject.toml

@@ -1,4 +1,4 @@
 [tool.pytest.ini_options]
 pythonpath = [
   "src"
-]
+]

python/cluster_management/requirements.txt

@@ -1,3 +1,3 @@
-pytest>=8
-botocore>=1.35.74
-boto3>=1.35.74
+boto3>=1.38.15
+botocore>=1.38.15
+pytest>=8

python/cluster_management/setup.sh

@@ -1,28 +1,70 @@
 import boto3
+import os
 
-def create_multi_region_clusters(client, linkedRegionList, witnessRegion, clusterProperties):
+
+def create_multi_region_clusters(region_1, region_2, witness_region):
     try:
-        response = client.create_multi_region_clusters(
-            linkedRegionList=linkedRegionList,
-            witnessRegion=witnessRegion,
-            clusterProperties=clusterProperties,
+        client_1 = boto3.client("dsql", region_name=region_1)
+        client_2 = boto3.client("dsql", region_name=region_2)
+
+        # We can only set the witness region for the first cluster
+        cluster_1 = client_1.create_cluster(
+            deletionProtectionEnabled=True,
+            multiRegionProperties={"witnessRegion": witness_region},
+            tags={"Name": "Python-CM-Example-Multi-Region", "Repo": "aws-samples/aurora-dsql-samples"}
+        )
+        print(f"Created {cluster_1["arn"]}")
+
+        # For the second cluster we can set witness region and designate cluster_1 as a peer
+        cluster_2 = client_2.create_cluster(
+            deletionProtectionEnabled=True,
+            multiRegionProperties={"witnessRegion": witness_region, "clusters": [cluster_1["arn"]]},
+            tags={"Name": "Python-CM-Example-Multi-Region", "Repo": "aws-samples/aurora-dsql-samples"}
         )
-        return response
+
+        print(f"Created {cluster_2["arn"]}")
+        # Now that we know the cluster_2 arn we can set it as a peer of cluster_1
+        client_1.update_cluster(
+            identifier=cluster_1["identifier"],
+            multiRegionProperties={"witnessRegion": witness_region, "clusters": [cluster_2["arn"]]}
+        )
+        print(f"Added {cluster_2["arn"]} as a peer of {cluster_1["arn"]}")
+
+        # Now that multiRegionProperties is fully defined for both clusters
+        # they'll begin the transition to ACTIVE
+        print(f"Waiting for {cluster_1["arn"]} to become ACTIVE")
+        client_1.get_waiter("cluster_active").wait(
+            identifier=cluster_1["identifier"],
+            WaiterConfig={
+                'Delay': 10,
+                'MaxAttempts': 50
+            }
+        )
+
+        print(f"Waiting for {cluster_2["arn"]} to become ACTIVE")
+        client_2.get_waiter("cluster_active").wait(
+            identifier=cluster_2["identifier"],
+            WaiterConfig={
+                'Delay': 10,
+                'MaxAttempts': 50
+            }
+        )
+
+        return cluster_1, cluster_2
+
     except:
-        print("Unable to create multi-region cluster")
+        print("Unable to create cluster")
         raise
 
+
 def main():
-    region = "us-east-1"
-    client = boto3.client("dsql", region_name=region)
-    linkedRegionList = ["us-east-1", "us-east-2"]
-    witnessRegion = "us-west-2"
-    clusterProperties = {
-        "us-east-1": {"tags": {"Name": "Foo"}},
-        "us-east-2": {"tags": {"Name": "Bar"}}
-    }
-    response = create_multi_region_clusters(client, linkedRegionList, witnessRegion, clusterProperties)
-    print("Linked Cluster Arns:", response['linkedClusterArns'])
+    region_1 = os.environ.get("REGION_1", "us-east-1")
+    region_2 = os.environ.get("REGION_2", "us-east-2")
+    witness_region = os.environ.get("WITNESS_REGION", "us-west-2")
+    (cluster_1, cluster_2) = create_multi_region_clusters(region_1, region_2, witness_region)
+    print("Created multi region clusters:")
+    print("Cluster id: " + cluster_1['arn'])
+    print("Cluster id: " + cluster_2['arn'])
 
 
 if __name__ == "__main__":

python/cluster_management/src/create_multi_region.py

@@ -1,20 +1,33 @@
 import boto3
+import os
 
-def create_cluster(client, tags, deletion_protection):
+
+def create_cluster(region):
     try:
-        response = client.create_cluster(tags=tags, deletionProtectionEnabled=deletion_protection)
-        return response
+        client = boto3.client("dsql", region_name=region)
+        tags = {"Name": "Python-CM-Example-Single-Region", "Repo": "aws-samples/aurora-dsql-samples"}
+        cluster = client.create_cluster(tags=tags, deletionProtectionEnabled=True)
+        print(f"Initiated creation of cluster: {cluster["identifier"]}")
+
+        print(f"Waiting for {cluster["arn"]} to become ACTIVE")
+        client.get_waiter("cluster_active").wait(
+            identifier=cluster["identifier"],
+            WaiterConfig={
+                'Delay': 10,
+                'MaxAttempts': 50
+            }
+        )
+
+        return cluster
     except:
         print("Unable to create cluster")
         raise
 
+
 def main():
-    region = "us-east-1"
-    client = boto3.client("dsql", region_name=region)
-    tag = {"Name": "FooBar"}
-    deletion_protection = True
-    response = create_cluster(client, tags=tag, deletion_protection=deletion_protection)
-    print("Cluster id: " + response['identifier'])
+    region = os.environ.get("REGION_1", "us-east-1")
+    response = create_cluster(region)
+    print(f"Created cluster: {response["arn"]}")
 
 
 if __name__ == "__main__":

python/cluster_management/src/create_single_region.py

@@ -1,17 +1,56 @@
 import boto3
+import os
+
+
+def delete_multi_region_clusters(region_1, cluster_id_1, region_2, cluster_id_2):
+    try:
+
+        client_1 = boto3.client("dsql", region_name=region_1)
+        client_2 = boto3.client("dsql", region_name=region_2)
+
+        client_1.delete_cluster(identifier=cluster_id_1)
+        print(f"Deleting cluster {cluster_id_1} in {region_1}")
+
+        # cluster_1 will stay in PENDING_DELETE state until cluster_2 is deleted
+
+        client_2.delete_cluster(identifier=cluster_id_2)
+        print(f"Deleting cluster {cluster_id_2} in {region_2}")
+
+        # Now that both clusters have been marked for deletion they will transition
+        # to DELETING state and finalize deletion
+        print(f"Waiting for {cluster_id_1} to finish deletion")
+        client_1.get_waiter("cluster_not_exists").wait(
+            identifier=cluster_id_1,
+            WaiterConfig={
+                'Delay': 10,
+                'MaxAttempts': 50
+            }
+        )
+
+        print(f"Waiting for {cluster_id_2} to finish deletion")
+        client_2.get_waiter("cluster_not_exists").wait(
+            identifier=cluster_id_2,
+            WaiterConfig={
+                'Delay': 10,
+                'MaxAttempts': 50
+            }
+        )
+
+    except:
+        print("Unable to delete cluster")
+        raise
 
-def delete_multi_region_clusters(linkedClusterArns, client):
-    client.delete_multi_region_clusters(linkedClusterArns=linkedClusterArns)
 
 def main():
-    region = "us-east-1"
-    client = boto3.client("dsql", region_name=region)
-    linkedClusterArns = [
-        "arn:aws:dsql:us-east-1:111111999999::cluster/foo0bar1baz2quux3quuux4",
-        "arn:aws:dsql:us-east-2:111111999999::cluster/bar0foo1baz2quux3quuux4"
-    ]
-    delete_multi_region_clusters(linkedClusterArns, client)
-    print("Deleting clusters with ARNs:", linkedClusterArns)
+    region_1 = os.environ.get("REGION_1", "us-east-1")
+    cluster_id_1 = os.environ.get("CLUSTER_ID_1")
+    assert cluster_id_1 is not None, "Must provide CLUSTER_ID_1"
+    region_2 = os.environ.get("REGION_2", "us-east-2")
+    cluster_id_2 = os.environ.get("CLUSTER_ID_2")
+    assert cluster_id_2 is not None, "Must provide CLUSTER_ID_2"
+
+    delete_multi_region_clusters(region_1, cluster_id_1, region_2, cluster_id_2)
+    print(f"Deleted {cluster_id_1} in {region_1} and {cluster_id_2} in {region_2}")
 
 
 if __name__ == "__main__":

python/cluster_management/src/delete_multi_region.py

@@ -1,18 +1,32 @@
 import boto3
+import os
 
-def delete_cluster(cluster_id, client):
-    try: 
-        return client.delete_cluster(identifier=cluster_id)
+
+def delete_cluster(region, identifier):
+    try:
+        client = boto3.client("dsql", region_name=region)
+        cluster = client.delete_cluster(identifier=identifier)
+        print(f"Initiated delete of {cluster["arn"]}")
+
+        print("Waiting for cluster to finish deletion")
+        client.get_waiter("cluster_not_exists").wait(
+            identifier=cluster["identifier"],
+            WaiterConfig={
+                'Delay': 10,
+                'MaxAttempts': 50
+            }
+        )
     except:
-        print("Unable to delete cluster " + cluster_id)
+        print("Unable to delete cluster " + identifier)
         raise
 
+
 def main():
-    region = "us-east-1"
-    client = boto3.client("dsql", region_name=region)
-    cluster_id = "foo0bar1baz2quux3quuux4"
-    response = delete_cluster(cluster_id, client)
-    print("Deleting cluster with ID: " + cluster_id +  ", Cluster Status: " + response['status'])
+    region = os.environ.get("REGION_1", "us-east-1")
+    cluster_id = os.environ.get("CLUSTER_ID_1")
+    assert cluster_id is not None, "Must provide CLUSTER_ID_1"
+    delete_cluster(region, cluster_id)
+    print(f"Deleted {cluster_id}")
 
 
 if __name__ == "__main__":