Wednesday, April 28, 2021

Java Haiku: Forever


public class Forever {
    public void forever() {
        ever = true;
        for(;ever;);
    }
}

 

Thursday, April 15, 2021

Creating AWS resources in pytest fixtures with moto

 When unit testing functions that access AWS resources, the easiest way is to use the moto library. For instance, I have to test several functions that access an SQS queue, so I proceed this way:

import boto3
from moto import mock_sqs

def create_queue():
    client = boto3.client("sqs")
    queue_url = client.create_queue(QueueName="queue")["QueueUrl"]
    sqs = boto3.resource("sqs")
    return sqs.Queue(queue_url)

@mock_sqs
def test_something()
    queue = create_queue()
    ...

Since I had this create_queue() call at the beginning of most of my test functions, I wanted to make it a fixture. So i tried this way:

import boto3
from moto import mock_sqs
from pytest import fixture

@fixture
@mock_sqs
def queue():
    client = boto3.client("sqs")
    queue_url = client.create_queue(QueueName="queue")["QueueUrl"]
    sqs = boto3.resource("sqs")
    return sqs.Queue(queue_url)

@mock_sqs
def test_something(queue)
    ...

Unfortunately, this raised an error in my test functions stating that the queue does not exist. The reason for it is that the decorator @mock_sqs on the fixture would destroy my SQS queue as soon as it would leave the queue() method.

The solution is simple: do not use the mock as a decorator, but trigger it when the fixture is initializing and destroy it when it terminates. That means using a yield within the fixture to return the queue:

import boto3
from moto import mock_sqs
from pytest import fixture

@fixture
def queue():
    mock_sqs().start()

    client = boto3.client("sqs")
    queue_url = client.create_queue(QueueName="queue")["QueueUrl"]
    sqs = boto3.resource("sqs")
    yield sqs.Queue(queue_url)
    
    mock_sqs().stop()

def test_something(queue)
    ...

Notice that we have to drop the decorator on the test function as well.

Thanks for the answers to this moto issue.

Monday, February 15, 2021

Airflow: Glue Operator looses its region name

 The other day, I tried to run an AWS Glue script from our Airflow instance. Nothing fancy, it would just convert a parquet file to CSV between two S3 buckets. Looking for an operator to use, I found that there is indeed a Glue Operator. It looks pretty easy to configure, so I tried it out:

parquet2csv_glue = AwsGlueJobOperator(
    task_id='parquet2csv_glue',
    dag=parquet2csv_dag,
    job_name='glue-parquet2csv',
    region_name='eu-west-3',
    aws_conn_id=None,
    script_args={
        "--SOURCE_PATH": input_path,
        "--TARGET_PATH": output_path,
    }
)

After triggering the DAG, it would fail, telling me that I need to set the region name. Well, I thought I did!

It probably comes from the fact that recent boto3 versions require region to be set, while it was optional in the past. As airflow is open source, I decided to have a look at the code. As it turns out, the bug was not too hard to spot. The Glue Operator creates a Glue Hook, which declares this constructor:

class AwsGlueJobHook(AwsBaseHook):
    def __init__(
        self,
        s3_bucket: Optional[str] = None,
        job_name: Optional[str] = None,
        desc: Optional[str] = None,
        concurrent_run_limitint = 1,
        script_location: Optional[str] = None,
        retry_limitint = 0,
        num_of_dpusint = 10,
        region_name: Optional[str] = None,
        iam_role_name: Optional[str] = None,
        *args,
        **kwargs,
    ):
        self.job_name = job_name
        self.desc = desc
        self.concurrent_run_limit = concurrent_run_limit
        self.script_location = script_location
        self.retry_limit = retry_limit
        self.num_of_dpus = num_of_dpus
        self.region_name = region_name
        self.s3_bucket = s3_bucket
        self.role_name = iam_role_name
        self.s3_glue_logs = 'logs/glue-logs/'
        kwargs['client_type'] = 'glue'
        super().__init__(*args, **kwargs)

The Hook derives from a base class, AwsBaseHook, that handles the common connection part for all AWS Hooks. The call to the constructor of the super class does not forward the region name. It should probably be called like this:

        super().__init__(region_name=region_name, *args, **kwargs)

I opened a bug report. But in the meanwhile, I still needed my code to work. So I found quite an ugly patch. There is probably better, but since I could see that the boto3 session was created in the AwsBaseHook class, and since our Airflow instance is running from an EC2 and I can just inherit its profile, I made this simple workaround:

import boto3
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook

# ugly patch to resolve the bug in the Glue Hook
def _get_credentials(selfregion_name):
    return boto3.session.Session(region_name="eu-west-3"), None

AwsBaseHook._get_credentials = _get_credentials

Hopefully, it won't stay here long...

Friday, December 18, 2020

Group files by folders in Python

 I sometimes need to display a list of files coming in this format:

folder1/file1

folder2/file1

folder1/file2 ...

 And I want to display it in this format:

folder1

file1

file2

folder2

 file1

 Here is my code. It uses the groupby function from the itertools library:

from itertools import groupby

def format_files_by_folder(folder,filenames):
    return folder + "\n  " + "\n  ".join([f[1for f in filenames])

def file_by_folder(file_list):
    files_and_folders = [(f.split('/')[0], '/'.join(f.split('/')[1:])) 
        for f in file_list]

    # Group by folder
    files_and_folders.sort(key=lambda f: f[0])
    files_by_folder = groupby(files_and_folders, lambda f: f[0])

    return "\n\n".join(
        [format_files_by_folder(folder, filenames) 
            for folder, filenames in files_by_folder])


Thursday, November 19, 2020

Snake Case in Terraform

 How do you convert from camel case to snake case in Terraform? How do you go from "MyProjectName" to "my-project-name"? Here is a simple solution:

locals {
# convert to snake case
snake_case_name = lower(replace(var.camel_case_name
"/(\\w)([A-Z])/""$${1}_$2"))
}

Basically, you add an underscore before each capital letters. The strange syntax with the double dollars is a workaround to a strange bug in Terraform, which considers '$1_' as meaning something. 

This line works well if you have simple cases like "MyDatabaseName". However, I had to handle some more complex cases, like "ABCProject", or "ProjectABC". In that case, you have to work a bit more.

My solution was to implement two replace functions:

  • One for a series of capital letters anywhere in my word. I insert an underscore only if my capital letter is followed by a lower case.
  • One for a series of capital letters at the end of my word. I insert an underscore before the first letter of the series.
This gives this more complex command:

locals {
  # convert to snake case
  snake_case_name lower(replace(replace(var.camel_case_name,
      # add underscore before capital letter followed by lowcase
      "/(\\w)([A-Z][a-z])/""$${1}_$2"),
      # add underscore before capital letters at the end of the word
      "/([A-Z]+)$/""_$1"))
}

Friday, June 26, 2020

Work around EC2 Termination Protection in Jenkins Pipeline

Recently, we enabled Termination Protection on our EC2 instances on our AWS cloud. That means that it is not possible to accidentally switch off an instance, you have to switch off the PRotection flag first. In our Terraform files, it is easy to put in place:
resource "aws_instance" "myinstance" {
  ...
  disable_api_termination = true
}
The problem is, when we run our Terraform scripts, and a new instance has to be created instead of an old one (like when we modify user data), it won't let us. You'd have to manually remove the flag. Since we decided that running a Terraform script for deployment means that we really want to be able to terminate our instance, we decide to allow the Jenkins pipeline to remove the flag before running the deployment scripts.
In order to do this, we wrote a small Groovy script at the beginning of our pipeline:
def removeTerminationProtection(instanceName) {
    echo "Looking for '${instanceName}'"
    def instanceId = sh (script: "aws ec2 describe-instances --region ${AWS_REGION} --filters Name=tag:Name,Values=${instanceName} --query 'Reservations[0].Instances[0].InstanceId' | xargs echo -n", returnStdout: true)
    if ('null'.equals(instanceId)) {
        echo "Instance ${instanceName} not found."
    }
    else {
        echo "Removing termination protection for '${instanceId}'"
        sh (script: "aws ec2 modify-instance-attribute --instance-id ${instanceId} --region ${AWS_REGION} --disable-api-termination Value=false || exit 1")
    }
}
It uses AWS CLI to first find our instance by its name, and then disable termination protection. The way to call it from a pipeline stage is the following:
pipeline {
    ...
    stages {
        ...
        stage('Remove Termination Protection') {
            environment {
                AWS_PROFILE = "${PROFILE_NAME}"
            }

            steps {
                script {
                    removeTerminationProtection("myinstance")
                }
            }
        }
        ...
    }
}
Since we run the AWS CLI, we must have AWS credentials, so we do it using an AWS Profile. We run our Terraform script in a later stage.

Wednesday, April 15, 2020

S3 Multipart Upload of Memory Mapped File in Java

In AWS,when uploading files bigger than 5 GB, you have no choice bu to use a Multipart Upload. In the first versions of the AWS SDK in Java, you had a TransferManager class to handle all the low level bits for you. Unfortunately, I did not find any trace of it in the latest versions.
The low level operations of the Multipart Upload expect you to divide your file into parts, each time providing a buffer to the part you want to upload. The best way to go through a big file without reading it all in memory is to use a Memory Mapped File. Here is my code:
private void uploadMultiPartFile(String bucketName, Path path, String key) throws IOException {
    S3Client s3 = S3Client.builder().region(REGION).build();
 
    CreateMultipartUploadRequest createMultipartUploadRequest = CreateMultipartUploadRequest.builder()
       .bucket(bucketName)
       .key(key)
       .build();

    CreateMultipartUploadResponse response = s3.createMultipartUpload(createMultipartUploadRequest);
    String uploadId = response.uploadId();

    long position = 0;
    long fileSize = Files.size(path);
    int part = 1;
    List<CompletedPart> completedParts = new ArrayList<>();

    try (FileChannel channel = FileChannel.open(path, StandardOpenOption.READ)) {
        while (position < fileSize) {
            long remaining = fileSize - position;
            int toRead = (int) Math.min(BUFFER_SIZE, remaining);
            MappedByteBuffer map = channel.map(MapMode.READ_ONLY, position, toRead);

            UploadPartRequest uploadPartRequest = UploadPartRequest.builder()
                .bucket(bucketName)
                .key(key)
                .uploadId(uploadId)
                .partNumber(part)
                .build();
            String etag = s3.uploadPart(uploadPartRequest, RequestBody.fromByteBuffer(map)).eTag();
            CompletedPart completed = CompletedPart.builder().partNumber(part).eTag(etag).build();
            completedParts.add(completed);
   
            position += BUFFER_SIZE;
            part++;
        }
    }  
 
    CompletedMultipartUpload completedMultipartUpload = CompletedMultipartUpload.builder().parts(completedParts).build();
    CompleteMultipartUploadRequest completeMultipartUploadRequest = CompleteMultipartUploadRequest.builder()
        .bucket(bucketName)
        .key(key)
        .uploadId(uploadId)
        .multipartUpload(completedMultipartUpload)
        .build();
    s3.completeMultipartUpload(completeMultipartUploadRequest);
}
For my case, I set the BUFFER_SIZE to 100MB. The AWS API states that buffers can be up to 5GB, but I had the bad surprise of noticing that the buffer size parameter in the Java SDK is an int instead of a long, so you cannot use anything above 2GB. When I tried to use 2GB, the process would just hang, so I guess there must be other limitations in the OS.
However, 100MB is OK for me because it allows me to upload files up to 1TB (10000 parts is the maximum allowed). But you can play with this parameter since bigger numbers allow you to upload files faster. Another optimization would be to use a parallel stream to run several uploads in parallel, but I didn't try it.
Another piece of advice if you start playing with Multipart Uploads: if your upload fails at some point, the parts already uploaded are stored in your S3, but you cannot see them, also you pay for them. So do not forget to attach a lifecycle policy for aborted Multipart Uploads.