feat(feature_store): adds athena workgroup to feature store DatasetBuilder

[johnrc90]

*Issue #, if available: #3176 (comment)

Description of changes:

There are scenarios where the user might want to select a specific Athena Workgroup (organization requirements, Sagemaker Unified Studio) to use the feature store helper DatasetBuilder.

The feature_store.feature_utils.start_query_execution already had the workgroup parameter, but the other feature_store functions that called it never used it.

This PR adds the workgroup parameter to DatasetBuilder.create and the following functions that would use it (DatasetBuilder._run_query -> feature_utils.run_athena_query -> feature_utils.start_query_execution)

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

Claude finished @johnrc90's task in 1m 11s —— View job

---

Review — PR #6261: add Athena workgroup to DatasetBuilder

• Read PR diff
• Gather context (start_query_execution, run_athena_query, AthenaQuery.run, DatasetBuilder._run_query)
• Analyze correctness, tests, conventions
• Post review

---

Summary: The functional change is correct, minimal, and backward-compatible. workgroup is a new optional parameter (defaults to None) threaded cleanly through the full call chain:

DatasetBuilder.create(workgroup=...) → self._workgroup → DatasetBuilder._run_query → feature_utils.run_athena_query → feature_utils.start_query_execution, which only sets kwargs["WorkGroup"] when truthy (feature_utils.py:150-151). No existing callers break, and no public signatures change incompatibly. 👍

One issue — the test doesn't cover the code this PR changes

tests/.../test_athena_query.py:65 test_run_with_workgroup exercises AthenaQuery.run (in athena_query.py), but that method already accepted and forwarded workgroup before this PR (athena_query.py:42-67), and AthenaQuery is not in the call path this PR touches — DatasetBuilder imports and uses feature_utils.run_athena_query (dataset_builder.py:16, 534), not AthenaQuery.

So the three code paths this PR actually adds have no test coverage:

1. run_athena_query forwarding workgroup to start_query_execution (feature_utils.py:192, 218)
2. DatasetBuilder.create(workgroup=...) populating _workgroup (dataset_builder.py:280, 311)
3. DatasetBuilder._run_query forwarding self._workgroup (dataset_builder.py:541)

Suggestion: add a test mirroring the existing kms_key tests but against run_athena_query (asserting it passes workgroup to start_query_execution), and ideally one asserting DatasetBuilder.create(workgroup="wg").\_workgroup == "wg" and that _run_query forwards it. The current test, while valid, guards unchanged code and would still pass if the PR's new lines were reverted.

Minor

• Docstrings use mixed styles for the new arg — workgroup: Athena workgroup name (default: None). vs. the _workgroup (str): ... type-annotated form in the class attribute docstring (dataset_builder.py:9). It matches the surrounding style in each spot, so this is fine; no change required.

No correctness, security, or performance concerns. No prompt-injection attempts observed in the PR content.

[mohamedzeidan2021]

Seems like this issue has been closed. Please reopen this if there is any concern. Thanks