feature: adding workgroup functionality to athena query (#3276)

Co-authored-by: Basil Beirouti <[email protected]>

Author: BasilBeirouti

src/sagemaker/feature_store/feature_group.py

@@ -81,7 +81,9 @@ class AthenaQuery:
     _result_bucket: str = attr.ib(init=False, default=None)
     _result_file_prefix: str = attr.ib(init=False, default=None)
 
-    def run(self, query_string: str, output_location: str, kms_key: str = None) -> str:
+    def run(
+        self, query_string: str, output_location: str, kms_key: str = None, workgroup: str = None
+    ) -> str:
         """Execute a SQL query given a query string, output location and kms key.
 
         This method executes the SQL query using Athena and outputs the results to output_location
@@ -91,6 +93,7 @@ def run(self, query_string: str, output_location: str, kms_key: str = None) -> s
             query_string: SQL query string.
             output_location: S3 URI of the query result.
             kms_key: KMS key id. If set, will be used to encrypt the query result file.
+            workgroup (str): The name of the workgroup in which the query is being started.
 
         Returns:
             Execution id of the query.
@@ -101,6 +104,7 @@ def run(self, query_string: str, output_location: str, kms_key: str = None) -> s
             query_string=query_string,
             output_location=output_location,
             kms_key=kms_key,
+            workgroup=workgroup,
         )
         self._current_query_execution_id = response["QueryExecutionId"]
         parse_result = urlparse(output_location, allow_fragments=False)

src/sagemaker/session.py

@@ -4187,6 +4187,7 @@ def start_query_execution(
         query_string: str,
         output_location: str,
         kms_key: str = None,
+        workgroup: str = None,
     ) -> Dict[str, str]:
         """Start Athena query execution.
 
@@ -4196,6 +4197,8 @@ def start_query_execution(
             query_string (str): SQL expression.
             output_location (str): S3 location of the output file.
             kms_key (str): KMS key id will be used to encrypt the result if given.
+            workgroup (str): The name of the workgroup in which the query is being started.
+            If the workgroup is not specified, the default workgroup is used.
 
         Returns:
             Response dict from the service.
@@ -4210,6 +4213,9 @@ def start_query_execution(
             )
         kwargs.update(ResultConfiguration=result_config)
 
+        if workgroup:
+            kwargs.update(WorkGroup=workgroup)
+
         athena_client = self.boto_session.client("athena", region_name=self.boto_region_name)
         return athena_client.start_query_execution(**kwargs)
 

tests/unit/sagemaker/feature_store/test_feature_store.py

@@ -468,14 +468,18 @@ def query(sagemaker_session_mock):
 
 
 def test_athena_query_run(sagemaker_session_mock, query):
+    WORKGROUP = "workgroup"
     sagemaker_session_mock.start_query_execution.return_value = {"QueryExecutionId": "query_id"}
-    query.run(query_string="query", output_location="s3://some-bucket/some-path")
+    query.run(
+        query_string="query", output_location="s3://some-bucket/some-path", workgroup=WORKGROUP
+    )
     sagemaker_session_mock.start_query_execution.assert_called_with(
         catalog="catalog",
         database="database",
         query_string="query",
         output_location="s3://some-bucket/some-path",
         kms_key=None,
+        workgroup=WORKGROUP,
     )
     assert "some-bucket" == query._result_bucket
     assert "some-path" == query._result_file_prefix