Athena: start_query_execution() should support WorkGroups (getmoto#8390)

pinzon · Dec 12, 2024 · 57d0c0e · 57d0c0e
1 parent b10edaa
commit 57d0c0e
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/moto/athena/responses.py b/moto/athena/responses.py
@@ -89,7 +89,7 @@ def get_query_execution(self) -> str:
                     "QueryPlanningTimeInMillis": 0,
                     "ServiceProcessingTimeInMillis": 0,
                 },
-                "WorkGroup": execution.workgroup,
+                "WorkGroup": execution.workgroup.name if execution.workgroup else None,
             }
         }
         if execution.execution_parameters is not None:

diff --git a/tests/test_athena/test_athena.py b/tests/test_athena/test_athena.py
@@ -237,6 +237,29 @@ def test_stop_query_execution():
     assert details["Status"]["State"] == "CANCELLED"
 
 
+@mock_aws
+def test_start_execution_with_workgroup():
+    client = boto3.client("athena", region_name="us-east-1")
+
+    client.create_work_group(
+        Name="myworkgroup",
+        Description="Test work group",
+        Configuration={
+            "ResultConfiguration": {"OutputLocation": "s3://bucket-name/prefix/"}
+        },
+    )
+
+    exec_id = client.start_query_execution(
+        QueryString="SELECT stuff",
+        QueryExecutionContext={"Database": "database"},
+        ResultConfiguration={"OutputLocation": "s3://bucket-name/prefix/"},
+        WorkGroup="myworkgroup",
+    )["QueryExecutionId"]
+
+    execution = client.get_query_execution(QueryExecutionId=exec_id)["QueryExecution"]
+    assert execution["WorkGroup"] == "myworkgroup"
+
+
 @mock_aws
 def test_create_named_query():
     client = boto3.client("athena", region_name="us-east-1")