<p><div><div><div><div data-show-faces="false" data-height="21" data-width="100" data-layout="button_count" data-send="true" data-href="https://www.r-bloggers.com/alleviating-aws-athena-aggravation-with-asynchronous-assistance/"></div></div><div></div></div><div></div></div>
<div>
(This article was first published on <strong><a href="https://rud.is/b/2018/07/14/alleviating-aws-athena-aggravation-with-asynchronous-assistance/"> R – rud.is</a></strong>, and kindly contributed to <a rel="nofollow" href="https://www.r-bloggers.com/">R-bloggers)</a> </div>
<p>I’ve blogged about <a rel="nofollow" href="https://rud.is/b/2018/04/20/painless-odbc-dplyr-connections-to-amazon-athena-and-apache-drill-with-r-odbc/" target="_blank">how to use Amazon Athena with R</a> before and if you are a regular Athena user, you’ve likely run into a situation where you prepare a <code>dplyr</code> chain, fire off a <code>collect()</code> and then wait.</p>
<p>And, wait.</p>
<p>And, wait.</p>
<p>And, wait.</p>
<p>Queries that take significant processing time or have large result sets do not play nicely with the provided ODBC and JDBC drivers. This means “hung” R sessions and severe frustration, especially when you can login to the AWS Athena console and see that the results are <em>right there</em>!!</p>
<p>I’ve been crafting SQL by hand or using <code>sql_render()</code> by hand to avoid this (when I remember to) but finally felt sufficient frustration to craft a better way, provided you can install and run <code>rJava</code>-based code (it’s 2018 and that still is not an easy given on many systems unfortunately).</p>
<p>There are two functions below:</p>
<ul>
<li><code>collect_async()</code>, and</li>
<li><code>gather_results()</code></li>
</ul>
<p>The <code>collect_async()</code> function is designed to be used like <code>collect()</code> but uses Athena components from the AWS SDK for Java to execute the SQL query behind the <code>dplyr</code> chain <em>asynchronously</em>. The companion function <code>gather_results()</code> takes the object created by <code>collect_async()</code> and checks to see if the results are ready. If if they are, it will use the <code>aws.s3</code> package to download them. Personally, I’d just <code>aws s3 sync ...</code> from the command line vs use the <code>aws.s3</code> package but that’s not everyone’s cup of tea.</p>
<p>Once I figure out the best package API for this I’ll add it to the <code>metis</code> package. There are many AWS idiosyncrasies that need to be accounted for and I’d rather ship this current set of functions via the blog so folks can use it (and tweak it to their needs) before waiting for perfection.</p>
<p>Here’s the code:</p>
<pre><code>library(rJava)
library(awsjavasdk)
library(aws.signature)
library(aws.s3)
library(odbc)
library(tidyverse)
library(dbplyr)

#' Collect Amazon Athena query results asynchronously
#' 
#' Long running Athena queries and Athena queries with large result
#' sets can seriously stall a `dplyr` processing chain due to poorly
#' implemented ODBC and JDBC drivers. The AWS SDK for Athena has 
#' methods that support submitting a query asynchronously for &quot;batch&quot;
#' processing. All Athena resutls are stored in CSV files in S3 and it's
#' easy to use the R `aws.s3` package to grab these or perform an
#' `aws s3 sync ...` operation on the command line.
#' 
#' <a rel="nofollow" href="http://twitter.com/md" target="_blank">@md</a>
#' <a rel="nofollow" href="http://twitter.com/param" target="_blank">@param</a> obj the `dplyr` chain
#' <a rel="nofollow" href="http://twitter.com/param" target="_blank">@param</a> schema Athena schema (usually matches the `Schema` parameter to the 
#'        Simba ODBC connection)
#' <a rel="nofollow" href="http://twitter.com/param" target="_blank">@param</a> region Your AWS region. All lower case with dashes (usually matches
#'        the `AwsRegion` parameter to the Simba ODBC connection)
#' <a rel="nofollow" href="http://twitter.com/param" target="_blank">@param</a> results_bucket the S3 results bucket where query results are stored 
#'        (usually matches the `S3OutputLocation` parameter to the Simba ODBC
#'        connection)
#' <a rel="nofollow" href="http://twitter.com/return" target="_blank">@return</a> a `list` with the query execution ID and the S3 bucket. This object
#'         is designed to be passed to the companion `gather_results()` if you
#'         want to use the `aws.s3` package to retrieve the results. Otherwise,
#'         sync the file however you want using the query execution id.
#' <a rel="nofollow" href="http://twitter.com/note" target="_blank">@note</a> You may need to change up the authentication provider depending on how 
#'       you use credentials with Athena
collect_async &lt;- function(obj, schema, region, results_bucket) {

  ugly_query &lt;- as.character(sql_render(obj))

  region &lt;- toupper(region)
  region &lt;- gsub(&quot;-&quot;, &quot;_&quot;, region, fixed=TRUE)

  regions &lt;- J(&quot;com.amazonaws.regions.Regions&quot;)

  available_regions &lt;- grep(&quot;^[[:upper:][:digit:]_]+$&quot;, names(regions), value=TRUE)
  if (!region %in% available_regions) stop(&quot;Invalid region.&quot;, call.=FALSE)

  switch(
    region,
    &quot;GovCloud&quot; = regions$GovCloud,
    &quot;US_EAST_1&quot; = regions$US_EAST_1,
    &quot;US_EAST_2&quot; = regions$US_EAST_2,
    &quot;US_WEST_1&quot; = regions$US_WEST_1,
    &quot;US_WEST_2&quot; = regions$US_WEST_2,
    &quot;EU_WEST_1&quot; = regions$EU_WEST_1,
    &quot;EU_WEST_2&quot; = regions$EU_WEST_2,
    &quot;EU_WEST_3&quot; = regions$EU_WEST_3,
    &quot;EU_CENTRAL_1&quot; = regions$EU_CENTRAL_1,
    &quot;AP_SOUTH_1&quot; = regions$AP_SOUTH_1,
    &quot;AP_SOUTHEAST_1&quot; = regions$AP_SOUTHEAST_1,
    &quot;AP_SOUTHEAST_2&quot; = regions$AP_SOUTHEAST_2,
    &quot;AP_NORTHEAST_1&quot; = regions$AP_NORTHEAST_1,
    &quot;AP_NORTHEAST_2&quot; = regions$AP_NORTHEAST_2,
    &quot;SA_EAST_1&quot; = regions$SA_EAST_1,
    &quot;CN_NORTH_1&quot; = regions$CN_NORTH_1,
    &quot;CN_NORTHWEST_1&quot; = regions$CN_NORTHWEST_1,
    &quot;CA_CENTRAL_1&quot; = regions$CA_CENTRAL_1,
    &quot;DEFAULT_REGION&quot; = regions$DEFAULT_REGION
  ) -> region

  provider &lt;- J(&quot;com.amazonaws.auth.DefaultAWSCredentialsProviderChain&quot;)
  client &lt;- J(&quot;com.amazonaws.services.athena.AmazonAthenaAsyncClientBuilder&quot;)

  my_client &lt;- client$standard()
  my_client &lt;- my_client$withRegion(region)
  my_client &lt;- my_client$withCredentials(provider$getInstance())
  my_client &lt;- my_client$build()

  queryExecutionContext &lt;- .jnew(&quot;com.amazonaws.services.athena.model.QueryExecutionContext&quot;)
  context &lt;- queryExecutionContext$withDatabase(schema)
  result &lt;- .jnew(&quot;com.amazonaws.services.athena.model.ResultConfiguration&quot;)
  result$setOutputLocation(results_bucket)

  startQueryExecutionRequest &lt;- .jnew(&quot;com.amazonaws.services.athena.model.StartQueryExecutionRequest&quot;)
  startQueryExecutionRequest$setQueryString(ugly_query)
  startQueryExecutionRequest$setQueryExecutionContext(context)
  startQueryExecutionRequest$setResultConfiguration(result)

  res &lt;- my_client$startQueryExecutionAsync(startQueryExecutionRequest)

  r &lt;- res$get()
  qex_id &lt;- r$getQueryExecutionId()

  list(
    qex_id = qex_id,
    results_bucket = results_bucket
  )

}

#' Gather the results of an asynchronous query
#'
#' <a rel="nofollow" href="http://twitter.com/md" target="_blank">@md</a>
#' <a rel="nofollow" href="http://twitter.com/param" target="_blank">@param</a> async_result the result of a call to `collect_async()`
#' <a rel="nofollow" href="http://twitter.com/return" target="_blank">@return</a> a data frame (tibble) or `NULL` if the query results are not ready yet
gather_results &lt;- function(async_result) {
  if (bucket_exists(sprintf(&quot;%s/%s&quot;, async_result$results_bucket, async_result$qex_id))) {
    readr::read_csv(
      get_object(sprintf(&quot;%s/%s.csv&quot;, async_result$results_bucket, async_result$qex_id))
    )
  } else {
    message(&quot;Results are not in the designated bucket.&quot;)
    return(NULL)
  }
}</code></pre>
<p>Now, we give it a go:</p>
<pre><code># Setup the credentials you're using
use_credentials(&quot;personal&quot;)

# load the AWS Java SDK classes
awsjavasdk::load_sdk()

# necessary for Simba ODBC and the async query ops
aws_region &lt;- &quot;us-east-1&quot;
athena_schema &lt;- &quot;sampledb&quot;
athena_results_bucket &lt;- &quot;s3://aws-athena-query-results-redacted&quot;

# connect to Athena and the sample database
DBI::dbConnect(
  odbc::odbc(),
  driver = &quot;/Library/simba/athenaodbc/lib/libathenaodbc_sbu.dylib&quot;,
  Schema = athena_schema,
  AwsRegion = aws_region,
  AuthenticationType = &quot;IAM Profile&quot;,
  AwsProfile = &quot;personal&quot;,
  S3OutputLocation = athena_results_bucket
) -> con

# the sample table in the sample db/schema
elb_logs &lt;- tbl(con, &quot;elb_logs&quot;)

# create your dplyr chain. This one is small so I don't incur charges
# collect_async() MUST be the LAST item in the dplyr chain.
elb_logs %>%
  filter(requestip == &quot;253.89.30.138&quot;) %>%
  collect_async(
    schema = athena_schema,
    region = aws_region,
    results_bucket = athena_results_bucket
  ) -> async_result

async_result
## $qex_id
## [1] &quot;d5fe7754-919b-47c5-bd7d-3ccdb1a3a414&quot;
## 
## $results_bucket
## [1] &quot;s3://aws-athena-query-results-redacted&quot;

# For long queries we can wait a bit but the function will tell us if the results
# are there or not.

gather_results(async_result)
## Parsed with column specification:
## cols(
##   timestamp = col_datetime(format = &quot;&quot;),
##   elbname = col_character(),
##   requestip = col_character(),
##   requestport = col_integer(),
##   backendip = col_character(),
##   backendport = col_integer(),
##   requestprocessingtime = col_double(),
##   backendprocessingtime = col_double(),
##   clientresponsetime = col_double(),
##   elbresponsecode = col_integer(),
##   backendresponsecode = col_integer(),
##   receivedbytes = col_integer(),
##   sentbytes = col_integer(),
##   requestverb = col_character(),
##   url = col_character(),
##   protocol = col_character()
## )
## # A tibble: 1 x 16
##   timestamp           elbname requestip     requestport backendip     backendport
##   </code></pre>