Glue pythonshell script to test connectivity to DocumentDB in VPC with TLS, without internet access

test_documentdb.py

### Glue pythonshell script to test connectivity to DocumentDB in VPC with TLS, without internet access

# make sure to configure connections to provision an ENI in VPC for network access:
# https://docs.aws.amazon.com/glue/latest/dg/glue-connections.html

### install packages from script to avoid pip calling PyPI for deps
# considering wheels and CA bundle are available at S3 bucket

import boto3
import subprocess
import os

bucket_name = "test"
whl_file_keys = ["dnspython-2.4.2-py3-none-any.whl", "pymongo-4.5.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"]

# https://truststore.pki.rds.amazonaws.com/global/global-bundle.pem
pem_key = "global-bundle.pem"

s3 = boto3.client("s3")

for whl_file_key in whl_file_keys:
    s3.download_file(bucket_name, whl_file_key, os.path.basename(whl_file_key))

subprocess.run(["pip", "install", "--no-index", "--find-links=.", *whl_file_keys], check=True)


# Download CA bundle for TLS connectivity with cluster:
s3.download_file(bucket_name, pem_key, os.path.basename(pem_key))


### main script
import sys
import datetime
import random
import string

from pymongo import MongoClient

print(f"Trying to connect to DocumentDB in VPC")

uri = 'mongodb://mongo:password@xxx.cluster-x.ap-northeast-1.docdb.amazonaws' \
'.com:27017/?tls=true&tlsCAFile=global-bundle.pem&replicaSet=rs0&readPreference=secondaryPreferred' \
'&retryWrites=false'
client = MongoClient(uri)

print(f"Server info: {client.server_info()}")
print(f"Databases available: {client.list_database_names()}")

## Test reads

db = client.get_database("test")
print(f"Collections in database {db.name}: {db.list_collection_names()}")

collection = db.get_collection('test_collection')
print(f"Documents in collection {collection.name}:")
for doc in db.test_collection.find():
    print(doc)

## Test writes

new_doc = {
datetime.datetime.now().strftime("%Y%m%d_%H%M%S"):
''.join(random.choice(string.ascii_uppercase) for _ in range(10))
}
print(f"Inserting document: {new_doc}")
collection.insert_one(new_doc)

print(f"Collections in database {db.name}: {db.list_collection_names()}")

print(f"Documents in collection {collection.name}:")
for doc in collection.find():
    print(doc)

client.close()