def path2list(
path: object,
boto3_session: boto3.Session,
last_modified_begin: Optional[datetime.datetime] = None,
last_modified_end: Optional[datetime.datetime] = None,
suffix: str = None,
) -> List[str]:
"""Convert Amazon S3 path to list of objects."""
if isinstance(path, str): # prefix
paths: List[str] = list_objects(
path=path,
suffix=suffix,
boto3_session=boto3_session,
last_modified_begin=last_modified_begin,
last_modified_end=last_modified_end,
)
elif isinstance(path, list):
if last_modified_begin or last_modified_end:
raise exceptions.InvalidArgumentCombination(
"Specify a list of files or (last_modified_begin and last_modified_end)"
) # pragma: no cover
paths = path if suffix is None else [
x for x in path if x.endswith(suffix)
]
else:
raise exceptions.InvalidArgumentType(
f"{type(path)} is not a valid path type. Please, use str or List[str]."
)
return paths
def _path2list(
path: Union[str, Sequence[str]],
boto3_session: boto3.Session,
last_modified_begin: Optional[datetime.datetime] = None,
last_modified_end: Optional[datetime.datetime] = None,
suffix: Union[str, List[str], None] = None,
ignore_suffix: Union[str, List[str], None] = None,
) -> List[str]:
"""Convert Amazon S3 path to list of objects."""
_suffix: Optional[List[str]] = [suffix] if isinstance(suffix, str) else suffix
_ignore_suffix: Optional[List[str]] = [ignore_suffix] if isinstance(ignore_suffix, str) else ignore_suffix
if isinstance(path, str): # prefix
paths: List[str] = list_objects(
path=path,
suffix=_suffix,
ignore_suffix=_ignore_suffix,
boto3_session=boto3_session,
last_modified_begin=last_modified_begin,
last_modified_end=last_modified_end,
)
elif isinstance(path, list):
if last_modified_begin or last_modified_end:
raise exceptions.InvalidArgumentCombination(
"Specify a list of files or (last_modified_begin and last_modified_end)"
)
paths = path if _suffix is None else [x for x in path if x.endswith(tuple(_suffix))]
paths = path if _ignore_suffix is None else [x for x in paths if x.endswith(tuple(_ignore_suffix)) is False]
else:
raise exceptions.InvalidArgumentType(f"{type(path)} is not a valid path type. Please, use str or List[str].")
return paths
def read_sql_query(
sql: str,
database: str,
ctas_approach: bool = True,
categories: Optional[List[str]] = None,
chunksize: Optional[Union[int, bool]] = None,
s3_output: Optional[str] = None,
workgroup: Optional[str] = None,
encryption: Optional[str] = None,
kms_key: Optional[str] = None,
keep_files: bool = True,
ctas_temp_table_name: Optional[str] = None,
use_threads: bool = True,
boto3_session: Optional[boto3.Session] = None,
max_cache_seconds: int = 0,
max_cache_query_inspections: int = 50,
max_remote_cache_entries: int = 50,
max_local_cache_entries: int = 100,
data_source: Optional[str] = None,
params: Optional[Dict[str, Any]] = None,
s3_additional_kwargs: Optional[Dict[str, Any]] = None,
) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
"""Execute any SQL query on AWS Athena and return the results as a Pandas DataFrame.
**Related tutorial:**
- `Amazon Athena <https://aws-data-wrangler.readthedocs.io/en/2.5.0/
tutorials/006%20-%20Amazon%20Athena.html>`_
- `Athena Cache <https://aws-data-wrangler.readthedocs.io/en/2.5.0/
tutorials/019%20-%20Athena%20Cache.html>`_
- `Global Configurations <https://aws-data-wrangler.readthedocs.io/en/2.5.0/
tutorials/021%20-%20Global%20Configurations.html>`_
**There are two approaches to be defined through ctas_approach parameter:**
**1** - ctas_approach=True (Default):
Wrap the query with a CTAS and then reads the table data as parquet directly from s3.
PROS:
- Faster for mid and big result sizes.
- Can handle some level of nested types.
CONS:
- Requires create/delete table permissions on Glue.
- Does not support timestamp with time zone
- Does not support columns with repeated names.
- Does not support columns with undefined data types.
- A temporary table will be created and then deleted immediately.
- Does not support custom data_source/catalog_id.
**2** - ctas_approach=False:
Does a regular query on Athena and parse the regular CSV result on s3.
PROS:
- Faster for small result sizes (less latency).
- Does not require create/delete table permissions on Glue
- Supports timestamp with time zone.
- Support custom data_source/catalog_id.
CONS:
- Slower for big results (But stills faster than other libraries that uses the regular Athena's API)
- Does not handle nested types at all.
Note
----
The resulting DataFrame (or every DataFrame in the returned Iterator for chunked queries) have a
`query_metadata` attribute, which brings the query result metadata returned by
`Boto3/Athena <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services
/athena.html#Athena.Client.get_query_execution>`_ .
For a practical example check out the
`related tutorial <https://aws-data-wrangler.readthedocs.io/en/2.5.0/
tutorials/024%20-%20Athena%20Query%20Metadata.html>`_!
Note
----
Valid encryption modes: [None, 'SSE_S3', 'SSE_KMS'].
`P.S. 'CSE_KMS' is not supported.`
Note
----
Create the default Athena bucket if it doesn't exist and s3_output is None.
(E.g. s3://aws-athena-query-results-ACCOUNT-REGION/)
Note
----
`chunksize` argument (Memory Friendly) (i.e batching):
Return an Iterable of DataFrames instead of a regular DataFrame.
There are two batching strategies:
- If **chunksize=True**, a new DataFrame will be returned for each file in the query result.
- If **chunksize=INTEGER**, Wrangler will iterate on the data by number of rows igual the received INTEGER.
`P.S.` `chunksize=True` is faster and uses less memory while `chunksize=INTEGER` is more precise
in number of rows for each Dataframe.
`P.P.S.` If `ctas_approach=False` and `chunksize=True`, you will always receive an interador with a
single DataFrame because regular Athena queries only produces a single output file.
Note
----
In case of `use_threads=True` the number of threads
that will be spawned will be gotten from os.cpu_count().
Parameters
----------
sql : str
SQL query.
database : str
AWS Glue/Athena database name - It is only the origin database from where the query will be launched.
You can still using and mixing several databases writing the full table name within the sql
(e.g. `database.table`).
ctas_approach: bool
Wraps the query using a CTAS, and read the resulted parquet data on S3.
If false, read the regular CSV on S3.
categories: List[str], optional
List of columns names that should be returned as pandas.Categorical.
Recommended for memory restricted environments.
chunksize : Union[int, bool], optional
If passed will split the data in a Iterable of DataFrames (Memory friendly).
If `True` wrangler will iterate on the data by files in the most efficient way without guarantee of chunksize.
If an `INTEGER` is passed Wrangler will iterate on the data by number of rows igual the received INTEGER.
s3_output : str, optional
Amazon S3 path.
workgroup : str, optional
Athena workgroup.
encryption : str, optional
Valid values: [None, 'SSE_S3', 'SSE_KMS']. Notice: 'CSE_KMS' is not supported.
kms_key : str, optional
For SSE-KMS, this is the KMS key ARN or ID.
keep_files : bool
Should Wrangler delete or keep the staging files produced by Athena?
ctas_temp_table_name : str, optional
The name of the temporary table and also the directory name on S3 where the CTAS result is stored.
If None, it will use the follow random pattern: `f"temp_table_{uuid.uuid4().hex()}"`.
On S3 this directory will be under under the pattern: `f"{s3_output}/{ctas_temp_table_name}/"`.
use_threads : bool
True to enable concurrent requests, False to disable multiple threads.
If enabled os.cpu_count() will be used as the max number of threads.
boto3_session : boto3.Session(), optional
Boto3 Session. The default boto3 session will be used if boto3_session receive None.
max_cache_seconds : int
Wrangler can look up in Athena's history if this query has been run before.
If so, and its completion time is less than `max_cache_seconds` before now, wrangler
skips query execution and just returns the same results as last time.
If cached results are valid, wrangler ignores the `ctas_approach`, `s3_output`, `encryption`, `kms_key`,
`keep_files` and `ctas_temp_table_name` params.
If reading cached data fails for any reason, execution falls back to the usual query run path.
max_cache_query_inspections : int
Max number of queries that will be inspected from the history to try to find some result to reuse.
The bigger the number of inspection, the bigger will be the latency for not cached queries.
Only takes effect if max_cache_seconds > 0.
max_remote_cache_entries : int
Max number of queries that will be retrieved from AWS for cache inspection.
The bigger the number of inspection, the bigger will be the latency for not cached queries.
Only takes effect if max_cache_seconds > 0 and default value is 50.
max_local_cache_entries : int
Max number of queries for which metadata will be cached locally. This will reduce the latency and also
enables keeping more than `max_remote_cache_entries` available for the cache. This value should not be
smaller than max_remote_cache_entries.
Only takes effect if max_cache_seconds > 0 and default value is 100.
data_source : str, optional
Data Source / Catalog name. If None, 'AwsDataCatalog' will be used by default.
params: Dict[str, any], optional
Dict of parameters that will be used for constructing the SQL query. Only named parameters are supported.
The dict needs to contain the information in the form {'name': 'value'} and the SQL query needs to contain
`:name;`.
s3_additional_kwargs : Optional[Dict[str, Any]]
Forward to botocore requests. Valid parameters: "RequestPayer", "ExpectedBucketOwner".
e.g. s3_additional_kwargs={'RequestPayer': 'requester'}
Returns
-------
Union[pd.DataFrame, Iterator[pd.DataFrame]]
Pandas DataFrame or Generator of Pandas DataFrames if chunksize is passed.
Examples
--------
>>> import awswrangler as wr
>>> df = wr.athena.read_sql_query(sql="...", database="...")
>>> scanned_bytes = df.query_metadata["Statistics"]["DataScannedInBytes"]
>>> import awswrangler as wr
>>> df = wr.athena.read_sql_query(
... sql="SELECT * FROM my_table WHERE name=:name;",
... params={"name": "filtered_name"}
... )
"""
if ctas_approach and data_source not in (None, "AwsDataCatalog"):
raise exceptions.InvalidArgumentCombination(
"Queries with ctas_approach=True (default) does not support "
"data_source values different than None and 'AwsDataCatalog'. "
"Please check the related tutorial for more details "
"(https://github.com/awslabs/aws-data-wrangler/blob/main/"
"tutorials/006%20-%20Amazon%20Athena.ipynb)")
chunksize = sys.maxsize if ctas_approach is False and chunksize is True else chunksize
session: boto3.Session = _utils.ensure_session(session=boto3_session)
if params is None:
params = {}
for key, value in params.items():
sql = sql.replace(f":{key};", str(value))
if max_remote_cache_entries > max_local_cache_entries:
max_remote_cache_entries = max_local_cache_entries
_cache_manager.max_cache_size = max_local_cache_entries
cache_info: _CacheInfo = _check_for_cached_results(
sql=sql,
boto3_session=session,
workgroup=workgroup,
max_cache_seconds=max_cache_seconds,
max_cache_query_inspections=max_cache_query_inspections,
max_remote_cache_entries=max_remote_cache_entries,
)
_logger.debug("cache_info:\n%s", cache_info)
if cache_info.has_valid_cache is True:
_logger.debug("Valid cache found. Retrieving...")
try:
return _resolve_query_with_cache(
cache_info=cache_info,
categories=categories,
chunksize=chunksize,
use_threads=use_threads,
session=session,
s3_additional_kwargs=s3_additional_kwargs,
)
except Exception as e: # pylint: disable=broad-except
_logger.error(
e
) # if there is anything wrong with the cache, just fallback to the usual path
_logger.debug("Corrupted cache. Continuing to execute query...")
return _resolve_query_without_cache(
sql=sql,
database=database,
data_source=data_source,
ctas_approach=ctas_approach,
categories=categories,
chunksize=chunksize,
s3_output=s3_output,
workgroup=workgroup,
encryption=encryption,
kms_key=kms_key,
keep_files=keep_files,
ctas_temp_table_name=ctas_temp_table_name,
use_threads=use_threads,
s3_additional_kwargs=s3_additional_kwargs,
boto3_session=session,
)
def get_tables(
catalog_id: Optional[str] = None,
database: Optional[str] = None,
transaction_id: Optional[str] = None,
name_contains: Optional[str] = None,
name_prefix: Optional[str] = None,
name_suffix: Optional[str] = None,
boto3_session: Optional[boto3.Session] = None,
) -> Iterator[Dict[str, Any]]:
"""Get an iterator of tables.
Note
----
Please, does not filter using name_contains and name_prefix/name_suffix at the same time.
Only name_prefix and name_suffix can be combined together.
Parameters
----------
catalog_id : str, optional
The ID of the Data Catalog from which to retrieve Databases.
If none is provided, the AWS account ID is used by default.
database : str, optional
Database name.
transaction_id: str, optional
The ID of the transaction (i.e. used with GOVERNED tables).
name_contains : str, optional
Select by a specific string on table name
name_prefix : str, optional
Select by a specific prefix on table name
name_suffix : str, optional
Select by a specific suffix on table name
boto3_session : boto3.Session(), optional
Boto3 Session. The default boto3 session will be used if boto3_session receive None.
Returns
-------
Iterator[Dict[str, Any]]
Iterator of tables.
Examples
--------
>>> import awswrangler as wr
>>> tables = wr.catalog.get_tables()
"""
client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
paginator = client_glue.get_paginator("get_tables")
args: Dict[str, str] = {}
if (name_prefix is not None) and (name_suffix is not None) and (name_contains is not None):
raise exceptions.InvalidArgumentCombination(
"Please, does not filter using name_contains and "
"name_prefix/name_suffix at the same time. Only "
"name_prefix and name_suffix can be combined together."
)
if (name_prefix is not None) and (name_suffix is not None):
args["Expression"] = f"{name_prefix}*{name_suffix}"
elif name_contains is not None:
args["Expression"] = f"*{name_contains}*"
elif name_prefix is not None:
args["Expression"] = f"{name_prefix}*"
elif name_suffix is not None:
args["Expression"] = f"*{name_suffix}"
if database is not None:
dbs: List[str] = [database]
else:
dbs = [x["Name"] for x in get_databases(catalog_id=catalog_id)]
for db in dbs:
args["DatabaseName"] = db
response_iterator = paginator.paginate(
**_catalog_id(catalog_id=catalog_id, **_transaction_id(transaction_id=transaction_id, **args))
)
try:
for page in response_iterator:
for tbl in page["TableList"]:
yield tbl
except client_glue.exceptions.EntityNotFoundException:
continue
