From e80ad226c7d631c66866e06a677a975bd535737f Mon Sep 17 00:00:00 2001
From: Bhavika Tekwani <4955119+bhavika@users.noreply.github.com>
Date: Mon, 10 Feb 2020 11:44:54 -0500
Subject: [PATCH] [AIRFLOW-6767] Correct name for default Athena workgroup
 (#7394)

cherry-picked from b7aa778b3
---
 airflow/contrib/hooks/aws_athena_hook.py            | 4 ++--
 airflow/contrib/operators/aws_athena_operator.py    | 2 +-
 tests/contrib/operators/test_aws_athena_operator.py | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/airflow/contrib/hooks/aws_athena_hook.py b/airflow/contrib/hooks/aws_athena_hook.py
index 7f17601116b9d2..ce95cb1aa8b60d 100644
--- a/airflow/contrib/hooks/aws_athena_hook.py
+++ b/airflow/contrib/hooks/aws_athena_hook.py
@@ -54,7 +54,7 @@ def get_conn(self):
         return self.conn
 
     def run_query(self, query, query_context, result_configuration, client_request_token=None,
-                  workgroup='default'):
+                  workgroup='primary'):
         """
         Run Presto query on athena with provided config and return submitted query_execution_id
 
@@ -66,7 +66,7 @@ def run_query(self, query, query_context, result_configuration, client_request_t
         :type result_configuration: dict
         :param client_request_token: Unique token created by user to avoid multiple executions of same query
         :type client_request_token: str
-        :param workgroup: Athena workgroup name, when not specified, will be 'default'
+        :param workgroup: Athena workgroup name, when not specified, will be 'primary'
         :type workgroup: str
         :return: str
         """
diff --git a/airflow/contrib/operators/aws_athena_operator.py b/airflow/contrib/operators/aws_athena_operator.py
index e2928741b9e1f7..0e08ffda00fd07 100644
--- a/airflow/contrib/operators/aws_athena_operator.py
+++ b/airflow/contrib/operators/aws_athena_operator.py
@@ -53,7 +53,7 @@ class AWSAthenaOperator(BaseOperator):
     @apply_defaults
     def __init__(self, query, database, output_location, aws_conn_id='aws_default', client_request_token=None,
                  query_execution_context=None, result_configuration=None, sleep_time=30, max_tries=None,
-                 workgroup='default',
+                 workgroup='primary',
                  *args, **kwargs):
         super(AWSAthenaOperator, self).__init__(*args, **kwargs)
         self.query = query
diff --git a/tests/contrib/operators/test_aws_athena_operator.py b/tests/contrib/operators/test_aws_athena_operator.py
index 9d5f894dedec04..3d1b912a65dcd6 100644
--- a/tests/contrib/operators/test_aws_athena_operator.py
+++ b/tests/contrib/operators/test_aws_athena_operator.py
@@ -37,7 +37,7 @@
     'database': 'TEST_DATABASE',
     'outputLocation': 's3://test_s3_bucket/',
     'client_request_token': 'eac427d0-1c6d-4dfb-96aa-2835d3ac6595',
-    'workgroup': 'default'
+    'workgroup': 'primary'
 }
 
 query_context = {