Reset/Remove Static Athena Query Results

[robwhiteston]

Before I get to the main question, a little context about what I am trying to do. At my company I have set up a AWS MWAA environment and I am using motoserver as a way to run integration tests against my DAG's locally, and AWS Glue and S3 its been working FANTASTIC. One important bit of information: I am not using the moto mocks in my pytest code. My tests are executing outside of the container and calling the airflow REST API with Airflow's api client.

I am setting our environment up to leverage the newish Athena SQL Connection provided by the the Amazon airflow provider. I read through how to set up static query results with the moto Athena backend and have successfully gotten my simple integration test to run green. Right now I have 3 pytest fixtures which each load a different static result into the static query results for Athena. All three basically look like this (I figured posting all three would create clutter):

@pytest.fixture
def med_claims_zero_row_count_result():
    expected_results = {
        "results": [
                {
                    "rows": [{"Data": [{"VarCharValue": "0"}]}],
                    "column_info": [{
                        "CatalogName": "AWSDataCatalog",
                        "SchemaName": "data_local_db_claims",
                        "TableName": "ah_medical_claims",
                        "Name": "col_0",
                        "Type": "string",
                        "Precision": 1,
                        "Scale": 0,
                        "Nullable": "NOT_NULL",
                        "CaseSensitive": True,
                    }]
                }
            ]
        }
    resp = requests.post(
        url="http://127.0.0.1:6000//moto-api/static/athena/query-results",
        json=expected_results
    )

    return resp

Here is the integration test:

def test_execute_example_athena_sqlcheck_dag(dag_id, 
                                    med_claims_zero_row_count_result, #seed query results fixture 1
                                    med_claims_table_check_count_is_zero, #seed query results fixture 2
                                    sql_check_operator_matching_counts, #seed query results fixture 3
                                    airflow_api_client, 
                                    dag_run_client, 
                                    athena_dag_run_id_prefix ):
    query_result_resp = med_claims_zero_row_count_result
    print(query_result_resp)
    table_check_result_resp = med_claims_table_check_count_is_zero
    sql_check_result_resp = sql_check_operator_matching_counts
    dag_run_api, dag_run = pytest.helpers.build_dag_run_request(airflow_api_client, dag_id, athena_dag_run_id_prefix)
    resp = dag_run_api.post_dag_run(dag_id, dag_run)
    pprint(resp)
    pytest.helpers.post_dag_run_and_poll_status(dag_run_client, dag_id, dag_run)

    dag_run_result = dag_run_api.get_dag_run(dag_id=dag_id, dag_run_id=dag_run.dag_run_id)
    assert dag_run_result['state'].value == 'success'

The Athena SQL Connection doesn't actually use the boto client, rather, it uses PyAthena for the DB API 2.0 support. So, there isn't anyway for me to predict what QueryExecutionId will be assigned to each of the three seeded query results. That is fine, so long as I can delete the static query results between test method executions. Is there a way for me to delete the static query results from the motoserver running inside my containerized environment? I don't want to call moto's reset endpoint because when the environment starts up I create a bunch of AWS resources in moto which other tests depend on.

It doesn't seem like there is a DELETE endpoint for the static query results resource in the api, so if thats the best way to do it for my situation I'll gladly work on a PR to add it. I feel like there must be a way to do this and I am just misunderstanding something.

Thanks!!

Oh, and FWIW, the DAG under test is just a simple dag meant to illustrate how to use the SQLExecuteQueryOperator, the SQLTableCheckOperator, and the SQLCheckOperator. I am also happy to post my docker-compose.yml file that shows how my environment is set up if that is helpful.

[bblommers]

Hi @robwhiteston! To answer your question directly: No, there is currently no way to delete the static query results.

However: The query results are stored in a queue, and are tied to the query execution ID's as they come in. So if you have query_exec_id1, it will take the first response from the queue of possible query results and return that. If you request the results for query_exec_id1 for a second time, it will return the first response again.
When requesting the results for a different query, say query_exec_id2, it will take the next item from the queue, and return that instead.

So as long as you use three different query execution ID's, you should not have to worry about clashing results.

Does that make sense, or am I misunderstanding the problem?

[robwhiteston]

It makes sense, and is also sort of what I feared. The way I am using moto for this project, just as a stub for AWS, I don't have a way to explicitly set the QueryExecutionId, or to pass in a particular QueryExecutionId when retrieving the static query results. My integration test is running the airflow DAG, and in that DAG I have three separate SQL operators making queries against Athena using an Athena connection type that is based on PyAthena. I have found that as long as I create the static query results in the order in which those results would be retrieved from Athena, my test works. Once I add a second test to the mix and run the test more than once, things start breaking.

I could probably find some way to cobble together a complicated patching scheme, but I think the simplest thing would be to have a DELETE endpoint for the static query results so that I can remove the static query results when a test method finishes executing. Or more accurately, I could have a fixture that exists only to clean up everything:

@pytest.fixture
def clean_up_static_query_results():
  do_cleanup = True 
  #I am a pytest noob, but I think I need to yield something to ensure this runs after everything is done
  yield do_cleanup
  requests.delete(
        url="http://127.0.0.1:6000//moto-api/static/athena/query-results",
    )
  

I am happy to put up a PR for this delete endpoint!

[bblommers]

It does make sense to have a DELETE variant, so a PR would be very welcome @robwhiteston!

FYI: The request currently comes in via this method:

moto/moto/moto_api/_internal/responses.py

Line 144 in 792f7a3

def set_athena_result(

There is no matching on the HTTP method yet, so the same method is executed for GET/POST/PATCH/DELETE/etc. So introducing the delete variant would always involve an if/else here:

if request.method == "POST": ... 
elif request.method == "DELETE":...

Let me know if you need any other pointers!