Always use a default workgroup config, instead of retreiving a workgroup config from AWS API

Author: inbar-beehero

awswrangler/athena/_executions.py

@@ -22,7 +22,7 @@
  _QUERY_FINAL_STATES,
  _QUERY_WAIT_POLLING_DELAY,
  _apply_formatter,
- _get_workgroup_config,
+ _get_default_workgroup_config,
  _start_query_execution,
  _WorkGroupConfig,
 )
@@ -149,7 +149,7 @@ def start_query_execution(
  query_execution_id = cache_info.query_execution_id
  _logger.debug("Valid cache found. Retrieving...")
  else:
- wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
+ wg_config: _WorkGroupConfig = _get_default_workgroup_config()
  query_execution_id = _start_query_execution(
  sql=sql,
  wg_config=wg_config,

awswrangler/athena/_read.py

@@ -24,7 +24,7 @@
  _empty_dataframe_response,
  _get_query_metadata,
  _get_s3_output,
- _get_workgroup_config,
+ _get_default_workgroup_config,
  _QueryMetadata,
  _start_query_execution,
  _WorkGroupConfig,
@@ -431,7 +431,7 @@ def _resolve_query_without_cache_regular(
  dtype_backend: Literal["numpy_nullable", "pyarrow"] = "numpy_nullable",
  client_request_token: str | None = None,
 ) -> pd.DataFrame | Iterator[pd.DataFrame]:
- wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
+ 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
  _logger.debug("Executing sql: %s", sql)
@@ -597,7 +597,7 @@ def _unload(
  athena_query_wait_polling_delay: float,
  execution_params: list[str] | None,
 ) -> _QueryMetadata:
- wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
+ wg_config: _WorkGroupConfig = _get_default_workgroup_config()
  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

awswrangler/athena/_utils.py

@@ -140,26 +140,9 @@ def _start_query_execution(
  return response["QueryExecutionId"]
 
 
-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)
+def _get_default_workgroup_config() -> _WorkGroupConfig:
+ wg_config: _WorkGroupConfig = _WorkGroupConfig(enforced=False, s3_output=None, encryption=None, kms_key=None)
+ _logger.debug("Default workgroup config:\n%s", wg_config)
  return wg_config
 
 
@@ -783,7 +766,7 @@ def create_ctas_table(
 
  fully_qualified_name = f'"{ctas_database}"."{ctas_table}"'
 
- wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
+ 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
  # If the workgroup enforces an external location, then it overrides the user supplied argument

awswrangler/athena/_write_iceberg.py

@@ -15,7 +15,7 @@
 from awswrangler._config import apply_configs
 from awswrangler.athena._executions import wait_query
 from awswrangler.athena._utils import (
- _get_workgroup_config,
+ _get_default_workgroup_config,
  _start_query_execution,
  _WorkGroupConfig,
 )
@@ -303,7 +303,7 @@ def _merge_iceberg(
  None
 
  """
- wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
+ wg_config: _WorkGroupConfig = _get_default_workgroup_config()
 
  sql_statement: str
  if merge_cols:
@@ -488,7 +488,7 @@ def to_iceberg( # noqa: PLR0913
  ... )
 
  """
- wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
+ wg_config: _WorkGroupConfig = _get_default_workgroup_config()
  temp_table: str = f"temp_table_{uuid.uuid4().hex}"
 
  _validate_args(
@@ -743,7 +743,7 @@ def delete_from_iceberg_table(
  if not merge_cols:
  raise exceptions.InvalidArgumentValue("Merge columns must be specified.")
 
- wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
+ wg_config: _WorkGroupConfig = _get_default_workgroup_config()
  temp_table: str = f"temp_table_{uuid.uuid4().hex}"
 
  if not temp_path and not wg_config.s3_output: