Mock AWS Athena for your ETL tests | by Stefano Francavilla | Towards Data Science

I spent some time last weekend thinking about strategies to improve our ETL pipeline unit and integration tests.
Just to quickly spice up an idea, a simplified component of the ETL process we have in place is something like the following:
where a series of AWS Eventbridge rules process and transform the data injected from different sources and in different formats; the data is then ready to be reviewed in Athena before being saved in our Postgres database.
The main goal of this exercise is to simulate the interaction with our datalake (Athena) in a local environment, somehow aligned with what we have in production.
First of all, what is Athena:
📖Docs
Amazon Athena is an interactive query service that makes it easy to analyze data in Amazon S3 using standard SQL. Athena is serverless, so there is no infrastructure to manage, and you pay only for the queries that you run.
Athena is easy to use. Simply point to your data in Amazon S3, define the schema, and start querying using standard SQL. Most results are delivered within seconds. With Athena, there’s no need for complex ETL jobs to prepare your data for analysis. This makes it easy for anyone with SQL skills to quickly analyze large-scale datasets.
Athena is out-of-the-box integrated with AWS Glue Data Catalog, allowing you to create a unified metadata repository across various services, crawl data sources to discover schemas and populate your Catalog with new and modified table and partition definitions, and maintain schema versioning.
Put simply, by pointing Athena to a specific structured or semi-structured data stored in S3, it makes really easy to quickly analyze results by using SQL queries. When you run a query, Athena will parallelize it across hundreds/thousands of cores and deliver the result in seconds.
Under the hood and at its core, Athena is a fusion of Hadoop/Hive for its Data Description Language (DDL) and Presto, which allows full standard SQL support and works with a variety of standard data formats, including CSV, JSON and Apache Parquet. We’re going to describe them in detail later in this guide.
Some of the ingredients needed for this recipe to work are the following Docker images:
moto server
hive/hadoop
presto
We’re going to cover all of them in details in the following sections.
📝Note
This guide assumes that you’re familiar with Docker, AWS S3 and Athena services
Below I’m going to give you an overview of the docker images that will be used in this tutorial and their role in the whole process.
At the end of this section, we’ll have a full docker compose file ready to be used for our tests.
Let’s introduce each one of them with some examples.
Moto is an amazing library to mock different AWS services both by using decorators for your specific tests or as a stand alone server. At the time of writing this article and at the best of my knowledge, it does not support the specific use case for Athena described here.
For this guide, we’re going to use Moto as a stand alone server to serve the S3 service.
Let’s start this exciting journey by filling up our empty docker-container with the first service.
I’ve create a new project folder called mothena which contains only a docker-compose file as you can seen in the following image:
The moto-server service is self-explanatory and the only thing to take into account is that we’ve exposed the port 5000, which is the default one used by Moto server.
Now, we’re ready to launch our first service from our terminal (from our mothena directory) by typing the following command:
docker-compose up
As you can see from Figure 2, the Moto standalone server is now running on localhost:5000.
Now, let’s try to create a bucket and make sure it works as expected. I’m using a jupyter notebook with the following snippet:
The output result suggests us that something happened and the bucket named mothena has been created. I’ve used as region_name=”eu-west-1” but, feel free to input any region name you prefer.
The most important piece to note here is that boto3 client accepts a parameter called endpoint_url which basically refers to our local endpoint. You can read more here.
But, it’s not finished yet, as a proof that everything worked correctly, let’s have a look at the container logs in the terminal:
Figure 4 shows us that a PUT request has been executed with HTTP status code 200.
Everything works! Now, we can stop our container (type Ctrl+C in the terminal window where the container runs) and proceed with the next section.
As said at the beginning of this guide, under the hood, Athena runs Hadoop/Hive for its DDL. So, we need to include Hadoop/Hive in the picture too.
Lets briefly describe what Hadoop and Hive are (taking the definition from their official websites):
📖Docs
The Apache Hadoop software library is a framework that allows for the distributed processing of large data sets across clusters of computers using simple programming models. It is designed to scale up from single servers to thousands of machines, each offering local computation and storage. Rather than rely on hardware to deliver high-availability, the library itself is designed to detect and handle failures at the application layer, so delivering a highly-available service on top of a cluster of computers, each of which may be prone to failures.
The Apache Hive data warehouse software facilitates reading, writing, and managing large datasets residing in distributed storage using SQL. Structure can be projected onto data already in storage. A command line tool and JDBC driver are provided to connect users to Hive.
In short, Hadoop is an open source framework responsible to easily store and manage large datasets (from gigabytes to petabytes) of data across a cluster of servers: this means that instead of using one computer, Hadoop allows the use of multiple computers (clusters) to analyze massive datasets in parallel and to quickly execute distributed processes against huge amounts of data. It’s the foundation layer on top of which other services and applications can be built.
Hive is indeed an application that runs over the Hadoop framework for processing/querying this distributed data using a SQL like interface.
This means that, when we create a table in Hadoop, we can specify the location where the data (e.g. .csv or .parquet files) resides and, in our case, it will be S3.
We’re going to use this repo for the Hadoop/Hive image and more specifically, the hdp3.1-hive.
Let’s first clone the project by running the following command in the terminal:
git clone https://github.com/prestosql/docker-images.git
Now, before building the image, there is something we need to take into account: we want the Hadoop cluster to “communicate” with the moto server.
In doing so, we need to add some properties in one specific configuration files:
prestodev/hdp3.1-hive/files/etc/hadoop/conf/core-site.xml - contains configuration information for the Hadoop daemon, including S3 specs such as endpoint, access key, access secrect, etc..
Core-site
In the below gist, you can see the properties we’re going to add in core-site.xml; all of them are related to S3 and for more information, here you can find an extensive description of all the mentioned properties.
Let’s review the ones we’ve added as there are some aspects that may require a bit of explanation:
access key and access secret can be any non-null string we want. Moto doesn’t mind. As defined in the gist, I’m going to define them as ENV variables in the docker compose file (com.amazonaws.auth.EnvironmentVariableCredentialsProvider)
we’re using s3a protocol and we’ll see later what this implies to create a table
the endpoint refers to the pre-defined moto hostname, which is moto-server and has been defined in our docker-compose file
📝Note
As you may know, by default Docker Compose sets up a single network for our app. Each service joins the default network and is both reachable and discoverable by other containers on that network by its hostname
path style access set to True is particularly useful for s3a storage as it removes the need to set up DNS for virtual hosting. So in this way, we can have the following url style for s3: https://s3.Region.amazonaws.com/bucket-name/key-name , which for Moto means http://localhost:5000/bucket-name/key-name
let’s disable ssl, not required
Now that everything is set up, we can build the docker image and running it from our docker-compose. A new service will finally be added 🤓.
Let’s follow the github README and build the image by typing in the terminal the following command:
make prestodev/hdp3.1-hive
This would take few minutes and the end result will be a successfully tagged image like what you see in Figure 6.
It’s time to test if our Hadoop cluster works and, more importantly, works with moto.
Let’s then add a new service in our docker-compose file like in Figure 7.
You can see from the newly added service that we’ve included the access key and secret as ENV variables as defined in our core-site.xml.
One question you may have is why exposing port 1000? The main reason in to allow communication with Hive from our script (via localhost), something we’re going to explore later in this guide.
You can also have a look at the Docker file prestodev/hdp3.1-hive/Dockerfile and notice that there are 2 ports open for Hive: 9083 (used to work with Presto and it’s already available inside the containers’ network) and 10000.
Ok, with 2 services in our docker-compose.yml, we can now run again from our terminal:
docker-compose up
And as you can see from Figure 9, they are running as expected.
Now, let’s do something interesting here and imagine a scenario where we need to explore a subset (just few lines) of data from the new york residential properties sales transaction dataset (.parquet format, but any other compatible format is fine):
I want to upload a file on S3 with a specific partition
I can query the file by using Hive and save the result in a dataframe
For this exercise and for all the ones you’ll see next, I’m going to use sqlalchemy with pyhive dialects installed. You can read more here in order to set them up.
Having Hadoop and moto ready, you can run the following script and see what happens.
The above code is pretty simple, basically we’re:
Uploading the sample data on moto
Creating a table and update metadata with MSCK REPAIR. I’ve also included the partitioned column downloaded_at following the specific S3 key structure. I thought this would make the exercise as much completed as possible for different use cases
Retrieving results and save them in a dataframe. See Figure 10
As reference, in Figure 11 you can also see what’s going on in our containers, including all the GET requests.
And that’s it for this section. But, in order to fully align with Athena, we need to add the last tile in the puzzle which is Presto; so, let’s stop docker-compose (CRTL-C) again and go straight to the next section.
What is Presto?
📖Docs
From prestosql: Presto is a highly parallel and distributed query engine, that is built from the ground up for efficient, low latency analytics.
Presto (I’m using prestosql) supports ANSI SQL semantics and allows querying data where it lives, including Hive, Cassandra, relational databases or even proprietary data stores. Indeed, Presto has different connectors and Hive connector is one of them.
First, we need to clone Presto docker and add a new catalog. We’ve introduced the concept of connector before, what’s a catalog?
📖Docs
A Presto catalog contains schemas and references a data source via a connector. For example, you can configure a JMX catalog to provide access to JMX information via the JMX connector. When you run a SQL statement in Presto, you are running it against one or more catalogs. Other examples of catalogs include the hive catalog to connect to a Hive data source.
When addressing a table in Presto, the fully-qualified table name is always rooted in a catalog. For example, a fully-qualified table name of hive.test_data.test refers to the test table in the test_data schema in the hive catalog.
Catalogs are defined in properties files stored in the Presto configuration directory.
Hive already defines a hive catalog, we just need to reference it by adding our hive.properties in our presto catalog folder as you can see in Figure 12.
The structure of the properties is pretty straightforward and as defined here we’re:
Using hive-hadoop2 as connector
Defining which Hadoop endpoint Presto should connect to (do you remember docker-compose, hostname, same network?) using the thrift protocol
Defining path style here again to true
Defining S3 endpoint to point to our local moto service
And that’s it! Let’s move to the Presto directory and build the image from the terminal with the following command:
Docker build . -t prestosql-docker
Once finished, we can go back to our mothena docker-compose.yml and add the new service as follows:
As you can see, Presto includes aws access secret and key as ENV variables and exposes port 8080.
We’re ready to change our previous code snippet and include Presto to retrieve our data instead of Hive.
Let’s now run back again the mothena docker-compose and update our code snippet as follows:
As you can see, we’ve added the Presto engine which connects to the hive catalog and uses it to retrieve the pandas dataframe.
The result will be the same dataframe you’ve see in Figure 10.
Wow, I have learned a lot writing this guide and I hope you too 🎉.
Now you have locally the fully (almost :)) functional Athena service that you may use in your script or unit/integration tests on CodeBuild for example.
With all the services available on localhost, other than boto3 and pandas, you can use any other library that accepts sqlalchemy connection/engine or endpoint overwriting such as pyarrow, s3fs, some awswrangler functions, etc.
Here you can find the source code where I’ve pushed missing images on Docker Hub.
Please, feel free to share your feedback in the comments, including improvements on this guide.