(fix): Give precedence to user path for Athena UNLOAD S3 Output Location (#1216)

Author: jaidisido

awswrangler/athena/_read.py

@@ -298,33 +298,31 @@ def _resolve_query_without_cache_unload(
     partitioned_by: Optional[List[str]],
     database: Optional[str],
     data_source: Optional[str],
-    s3_output: str,
+    s3_output: Optional[str],
     keep_files: bool,
     chunksize: Union[int, bool, None],
     categories: Optional[List[str]],
     encryption: Optional[str],
     kms_key: Optional[str],
     workgroup: Optional[str],
-    wg_config: _WorkGroupConfig,
     use_threads: Union[bool, int],
     s3_additional_kwargs: Optional[Dict[str, Any]],
     boto3_session: boto3.Session,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     query_metadata = _unload(
-        sql,
-        s3_output,
-        file_format,
-        compression,
-        field_delimiter,
-        partitioned_by,
-        workgroup,
-        wg_config,
-        database,
-        encryption,
-        kms_key,
-        boto3_session,
-        data_source,
+        sql=sql,
+        path=s3_output,
+        file_format=file_format,
+        compression=compression,
+        field_delimiter=field_delimiter,
+        partitioned_by=partitioned_by,
+        workgroup=workgroup,
+        database=database,
+        encryption=encryption,
+        kms_key=kms_key,
+        boto3_session=boto3_session,
+        data_source=data_source,
     )
     if file_format == "PARQUET":
         return _fetch_parquet_result(
@@ -351,11 +349,13 @@ def _resolve_query_without_cache_regular(
     encryption: Optional[str],
     workgroup: Optional[str],
     kms_key: Optional[str],
-    wg_config: _WorkGroupConfig,
     use_threads: Union[bool, int],
     s3_additional_kwargs: Optional[Dict[str, Any]],
     boto3_session: boto3.Session,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
+    wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
+    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
     _logger.debug("sql: %s", sql)
     query_id: str = _start_query_execution(
         sql=sql,
@@ -413,9 +413,6 @@ def _resolve_query_without_cache(
 
     Usually called by `read_sql_query` when using cache is not possible.
     """
-    wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
-    _s3_output: str = _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
     if ctas_approach is True:
         if ctas_temp_table_name is not None:
             name: str = catalog.sanitize_table_name(ctas_temp_table_name)
@@ -426,7 +423,7 @@ def _resolve_query_without_cache(
                 sql=sql,
                 database=database,
                 data_source=data_source,
-                s3_output=_s3_output,
+                s3_output=s3_output,
                 keep_files=keep_files,
                 chunksize=chunksize,
                 categories=categories,
@@ -456,14 +453,13 @@ def _resolve_query_without_cache(
             partitioned_by=unload_parameters.get("partitioned_by"),
             database=database,
             data_source=data_source,
-            s3_output=_s3_output,
+            s3_output=s3_output,
             keep_files=keep_files,
             chunksize=chunksize,
             categories=categories,
             encryption=encryption,
             kms_key=kms_key,
             workgroup=workgroup,
-            wg_config=wg_config,
             use_threads=use_threads,
             s3_additional_kwargs=s3_additional_kwargs,
             boto3_session=boto3_session,
@@ -473,14 +469,13 @@ def _resolve_query_without_cache(
         sql=sql,
         database=database,
         data_source=data_source,
-        s3_output=_s3_output,
+        s3_output=s3_output,
         keep_files=keep_files,
         chunksize=chunksize,
         categories=categories,
         encryption=encryption,
         workgroup=workgroup,
         kms_key=kms_key,
-        wg_config=wg_config,
         use_threads=use_threads,
         s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=boto3_session,
@@ -489,19 +484,26 @@ def _resolve_query_without_cache(
 
 def _unload(
     sql: str,
-    path: str,
+    path: Optional[str],
     file_format: str,
     compression: Optional[str],
     field_delimiter: Optional[str],
     partitioned_by: Optional[List[str]],
     workgroup: Optional[str],
-    wg_config: _WorkGroupConfig,
     database: Optional[str],
     encryption: Optional[str],
     kms_key: Optional[str],
     boto3_session: boto3.Session,
     data_source: Optional[str],
 ) -> _QueryMetadata:
+    wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
+    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
+
     # Set UNLOAD parameters
     unload_parameters = f"  format='{file_format}'"
     if compression:
@@ -520,7 +522,7 @@ def _unload(
             wg_config=wg_config,
             database=database,
             data_source=data_source,
-            s3_output=path,
+            s3_output=s3_output,
             encryption=encryption,
             kms_key=kms_key,
             boto3_session=boto3_session,
@@ -1177,24 +1179,22 @@ def unload(
 
     """
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
-    wg_config: _WorkGroupConfig = _get_workgroup_config(session=session, workgroup=workgroup)
     # Substitute query parameters
     if params is None:
         params = {}
     for key, value in params.items():
         sql = sql.replace(f":{key};", str(value))
     return _unload(
-        sql,
-        path,
-        file_format,
-        compression,
-        field_delimiter,
-        partitioned_by,
-        workgroup,
-        wg_config,
-        database,
-        encryption,
-        kms_key,
-        session,
-        data_source,
+        sql=sql,
+        path=path,
+        file_format=file_format,
+        compression=compression,
+        field_delimiter=field_delimiter,
+        partitioned_by=partitioned_by,
+        workgroup=workgroup,
+        database=database,
+        encryption=encryption,
+        kms_key=kms_key,
+        boto3_session=session,
+        data_source=data_source,
     )