s3_output: where the query can use the workgroup setting, only set s3_output if passed by user. For queries that require a location, keep previous impl of s3_output either by user or by workgroup setting

Author: inbar-beehero

awswrangler/athena/_read.py

@@ -25,6 +25,7 @@
     _get_default_workgroup_config,
     _get_query_metadata,
     _get_s3_output,
+    _get_workgroup_config,
     _QueryMetadata,
     _start_query_execution,
     _WorkGroupConfig,
@@ -432,8 +433,7 @@ def _resolve_query_without_cache_regular(
     client_request_token: str | None = None,
 ) -> pd.DataFrame | Iterator[pd.DataFrame]:
     wg_config: _WorkGroupConfig = _get_default_workgroup_config()
-    s3_output = _get_s3_output(s3_output=s3_output, wg_config=wg_config, boto3_session=boto3_session)
-    s3_output = s3_output[:-1] if s3_output[-1] == "/" else s3_output
+    s3_output = s3_output[:-1] if s3_output and s3_output[-1] == "/" else s3_output
     _logger.debug("Executing sql: %s", sql)
     query_id: str = _start_query_execution(
         sql=sql,
@@ -597,13 +597,13 @@ def _unload(
     athena_query_wait_polling_delay: float,
     execution_params: list[str] | None,
 ) -> _QueryMetadata:
-    wg_config: _WorkGroupConfig = _get_default_workgroup_config()
+    wg_config: _WorkGroupConfig = _get_workgroup_config(workgroup=workgroup, session=boto3_session)
     s3_output: str = _get_s3_output(s3_output=path, wg_config=wg_config, boto3_session=boto3_session)
-    s3_output = s3_output[:-1] if s3_output[-1] == "/" else s3_output
-    # Athena does not enforce a Query Result Location for UNLOAD. Thus, the workgroup output location
-    # is only used if no path is supplied.
-    if not path:
-        path = s3_output
+    s3_output = s3_output[:-1] if s3_output and s3_output[-1] == "/" else s3_output
+    if not s3_output:
+        raise exceptions.InvalidArgumentValue(
+            "Output S3 location is required for UNLOAD, either as the path argument or as a workgroup configuration"
+        )
 
     # Set UNLOAD parameters
     unload_parameters = f"  format='{file_format}'"
@@ -614,7 +614,7 @@ def _unload(
     if partitioned_by:
         unload_parameters += f"  , partitioned_by=ARRAY{partitioned_by}"
 
-    sql = f"UNLOAD ({sql}) TO '{path}' WITH ({unload_parameters})"
+    sql = f"UNLOAD ({sql}) TO '{s3_output}' WITH ({unload_parameters})"
     _logger.debug("Executing unload query: %s", sql)
     try:
         query_id: str = _start_query_execution(

awswrangler/athena/_utils.py

@@ -93,17 +93,20 @@ def _start_query_execution(
     args: dict[str, Any] = {"QueryString": sql}
 
     # s3_output
-    args["ResultConfiguration"] = {
-        "OutputLocation": _get_s3_output(s3_output=s3_output, wg_config=wg_config, boto3_session=boto3_session)
-    }
+    if s3_output:
+        args["ResultConfiguration"] = {"OutputLocation": s3_output}
 
     # encryption
     if wg_config.enforced is True:
+        if "ResultConfiguration" not in args:
+            args["ResultConfiguration"] = {}
         if wg_config.encryption is not None:
             args["ResultConfiguration"]["EncryptionConfiguration"] = {"EncryptionOption": wg_config.encryption}
             if wg_config.kms_key is not None:
                 args["ResultConfiguration"]["EncryptionConfiguration"]["KmsKey"] = wg_config.kms_key
     elif encryption is not None:
+        if "ResultConfiguration" not in args:
+            args["ResultConfiguration"] = {}
         args["ResultConfiguration"]["EncryptionConfiguration"] = {"EncryptionOption": encryption}
         if kms_key is not None:
             args["ResultConfiguration"]["EncryptionConfiguration"]["KmsKey"] = kms_key
@@ -146,6 +149,29 @@ def _get_default_workgroup_config() -> _WorkGroupConfig:
     return wg_config
 
 
+def _get_workgroup_config(session: boto3.Session | None = None, workgroup: str = "primary") -> _WorkGroupConfig:
+    enforced: bool
+    wg_s3_output: str | None
+    wg_encryption: str | None
+    wg_kms_key: str | None
+
+    enforced, wg_s3_output, wg_encryption, wg_kms_key = False, None, None, None
+    if workgroup is not None:
+        res = get_work_group(workgroup=workgroup, boto3_session=session)
+        enforced = res["WorkGroup"]["Configuration"]["EnforceWorkGroupConfiguration"]
+        config: dict[str, Any] = res["WorkGroup"]["Configuration"].get("ResultConfiguration")
+        if config is not None:
+            wg_s3_output = config.get("OutputLocation")
+            encrypt_config: dict[str, str] | None = config.get("EncryptionConfiguration")
+            wg_encryption = None if encrypt_config is None else encrypt_config.get("EncryptionOption")
+            wg_kms_key = None if encrypt_config is None else encrypt_config.get("KmsKey")
+    wg_config: _WorkGroupConfig = _WorkGroupConfig(
+        enforced=enforced, s3_output=wg_s3_output, encryption=wg_encryption, kms_key=wg_kms_key
+    )
+    _logger.debug("Workgroup config:\n%s", wg_config)
+    return wg_config
+
+
 def _fetch_txt_result(
     query_metadata: _QueryMetadata,
     keep_files: bool,
@@ -767,8 +793,7 @@ def create_ctas_table(
     fully_qualified_name = f'"{ctas_database}"."{ctas_table}"'
 
     wg_config: _WorkGroupConfig = _get_default_workgroup_config()
-    s3_output = _get_s3_output(s3_output=s3_output, wg_config=wg_config, boto3_session=boto3_session)
-    s3_output = s3_output[:-1] if s3_output[-1] == "/" else s3_output
+    s3_output = s3_output[:-1] if s3_output and s3_output[-1] == "/" else s3_output
     # If the workgroup enforces an external location, then it overrides the user supplied argument
     external_location_str: str = (
         f"    external_location = '{s3_output}/{ctas_table}',\n" if (not wg_config.enforced) and (s3_output) else ""