AWS Step function And Athena : can I configure query string from inputpath?

Viewed 1160

in my step function, I would like to execute an Athena query. I am able to define a step and execute a query successfully. However, I would like to pass some parameters as input and use them in the query string. For example.

Let's say, my query string is:

select * from <Data Source>.<database>.<tablename> where partition_0 = '2021';

I want to be able to pass the year as a input json to the step function, something like:

{
"YYYY": 2021
}

Is it possible to insert the input "YYYY" in the query string? If so, how?

Sample step function configuration:

{
  "Comment": "Start athena exececution",
  "StartAt": "athena",
  "States": {
    "athena": {
      "Type": "Task",
      "InputPath": "$",
      "Resource": "arn:aws:states:::athena:startQueryExecution.sync",
      "Parameters": {
        "QueryString": "select * from mycatalog.mydatabase.mytable where partition_0 = '2021'",
        "WorkGroup": "primary",
        "ResultConfiguration": {
          "OutputLocation": "s3://mys3bucket"
        }
      },
      "Next": "Pass"
    },
    "Pass": {
      "Comment": "A Pass state passes its input to its output, without performing work. Pass states are useful when constructing and debugging state machines.",
      "Type": "Pass",
      "End": true
    }
  }
}
2 Answers

Use a Pass State to interpolate the input YYYY with the query string:

"QueryPassTask": {
  "Type": "Pass",
  "ResultPath": "$.athena",
  "Parameters": {
    "query.$": "States.Format('select * from mycatalog.mydatabase.mytable where partition_0 = \\'{}\\'', $.YYYY)"
  },

Pass task output is:

{
  "YYYY": 2021,
  "athena": {
    "query": "select * from mycatalog.mydatabase.mytable where partition_0 = '2021'"
  }
}

Next, provide the query string to the Athena task. Don't forget the '.$' suffix on the key. This tells Step Functions that the key's value contains a substitution.

 "Parameters": {
        "QueryString.$": "$.athena.query",

In addition to fedenov's answer: I think you don't necessarily need an intermediate Pass state. The Pass state also does not allow you to use the ResultSelector (that might be annoying). You can do the string formatting directly in the Athena Query task.

Below is an example.

{
    "athena": {
        "catalog": "my_catalog",
        "database": "my_database",
        "outputLocation": "S3://my_bucket...",
        "queryString": "SELECT ... {}",
        "year": "2022"
    }
}
    "queryData": {
      "Type": "Task",
      "Resource": "arn:aws:states:::athena:startQueryExecution.sync",
      "Parameters": {
        "QueryExecutionContext": {
          "Catalog.$": "$.athena.catalog",
          "Database.$": "$.athena.database"
        },
        "ResultConfiguration": {
          "OutputLocation.$": "$.athena.outputLocation"
        },
        "QueryString.$": "States.Format($.athena.queryString, $.athena.year)",
        "WorkGroup": "primary"
      },
      "ResultPath": "$.query",
      "Next": "nextStep"
    },
Related