From cccd6523826339e75067411744498a3eaa6d2616 Mon Sep 17 00:00:00 2001
From: Jonathan Carvalho <49275672+johnrc90@users.noreply.github.com>
Date: Fri, 11 Sep 2026 10:32:21 -0300
Subject: [PATCH] feature: adds athena workgroup to feature store
 DatasetBuilder

---
 .../mlops/feature_store/dataset_builder.py         |  6 ++++++
 .../sagemaker/mlops/feature_store/feature_utils.py |  3 +++
 .../mlops/feature_store/test_athena_query.py       | 14 ++++++++++++++
 3 files changed, 23 insertions(+)

diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py
index bdac896ba7..416d6110ec 100644
--- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py
+++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py
@@ -214,6 +214,7 @@ class DatasetBuilder:
             (default: None).
         _kms_key_id (str): A KMS key id. If set, will be used to encrypt the result file
             (default: None).
+        _workgroup (str): Athena workgroup name in which the query will run (default: None).
         _point_in_time_accurate_join (bool): A boolean representing if point-in-time join
             is applied to the resulting dataframe when calling "to_dataframe".
             When set to True, users can retrieve data using "row-level time travel"
@@ -250,6 +251,7 @@ class DatasetBuilder:
     _event_time_identifier_feature_name: str = None
     _included_feature_names: List[str] = None
     _kms_key_id: str = None
+    _workgroup: str = None
     _event_time_identifier_feature_type: FeatureTypeEnum = None
 
     _point_in_time_accurate_join: bool = field(default=False, init=False)
@@ -274,6 +276,7 @@ def create(
         event_time_identifier_feature_name: str = None,
         included_feature_names: List[str] = None,
         kms_key_id: str = None,
+        workgroup: str = None,
         register_as_dataset: bool = False,
     ) -> "DatasetBuilder":
         """Create a DatasetBuilder for generating a Dataset.
@@ -286,6 +289,7 @@ def create(
             event_time_identifier_feature_name: Required if base is DataFrame.
             included_feature_names: Features to include in output.
             kms_key_id: KMS key for encryption.
+            workgroup: Athena workgroup name (default: None).
 
         Returns:
             DatasetBuilder instance.
@@ -304,6 +308,7 @@ def create(
             _event_time_identifier_feature_name=event_time_identifier_feature_name,
             _included_feature_names=included_feature_names,
             _kms_key_id=kms_key_id,
+            _workgroup=workgroup,
             _register_as_dataset=register_as_dataset,
         )
 
@@ -533,6 +538,7 @@ def _run_query(self, query_string: str, catalog: str, database: str) -> Dict[str
             query_string=query_string,
             output_location=self._output_path,
             kms_key=self._kms_key_id,
+            workgroup=self._workgroup,
         )
 
     def _create_temp_table(self, temp_table_name: str, s3_folder: str):
diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py
index 8c6d9b2615..ebc58c8e52 100644
--- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py
+++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py
@@ -182,6 +182,7 @@ def run_athena_query(
     query_string: str,
     output_location: str,
     kms_key: str = None,
+    workgroup: str = None,
 ) -> Dict[str, Any]:
     """Execute Athena query, wait for completion, and return result.
 
@@ -192,6 +193,7 @@ def run_athena_query(
         query_string: SQL query string.
         output_location: S3 URI for query results.
         kms_key: KMS key for encryption (default: None).
+        workgroup: Athena workgroup name (default: None).
 
     Returns:
         Query execution result dict.
@@ -206,6 +208,7 @@ def run_athena_query(
         query_string=query_string,
         output_location=output_location,
         kms_key=kms_key,
+        workgroup=workgroup,
     )
     query_id = response["QueryExecutionId"]
     wait_for_athena_query(session, query_id)
diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_athena_query.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_athena_query.py
index 5085ef9613..955970effa 100644
--- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_athena_query.py
+++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_athena_query.py
@@ -61,6 +61,20 @@ def test_run_with_kms_key(self, mock_start, athena_query):
         call_kwargs = mock_start.call_args[1]
         assert call_kwargs["kms_key"] == "arn:aws:kms:us-west-2:123:key/abc"
 
+    @patch("sagemaker.mlops.feature_store.athena_query.start_query_execution")
+    def test_run_with_workgroup(self, mock_start, athena_query):
+        mock_start.return_value = {"QueryExecutionId": "query-123"}
+
+        athena_query.run(
+            query_string="SELECT * FROM table",
+            output_location="s3://bucket/output",
+            workgroup="workgroup1",
+        )
+
+        mock_start.assert_called_once()
+        call_kwargs = mock_start.call_args[1]
+        assert call_kwargs["workgroup"] == "workgroup1"
+
     @patch("sagemaker.mlops.feature_store.athena_query.wait_for_athena_query")
     def test_wait_calls_helper(self, mock_wait, athena_query):
         athena_query._current_query_execution_id = "query-123"