Add AWS STS module.

Author: igorborgest

awswrangler/__init__.py

@@ -7,8 +7,7 @@
 
 import logging as _logging
 
-from awswrangler import athena, catalog, cloudwatch, db, emr, exceptions, quicksight, s3  # noqa
+from awswrangler import athena, catalog, cloudwatch, db, emr, exceptions, quicksight, s3, sts  # noqa
 from awswrangler.__metadata__ import __description__, __license__, __title__, __version__  # noqa
-from awswrangler._utils import get_account_id  # noqa
 
 _logging.getLogger("awswrangler").addHandler(_logging.NullHandler())

awswrangler/_utils.py

@@ -207,12 +207,6 @@ def get_directory(path: str) -> str:
     return path.rsplit(sep="/", maxsplit=1)[0] + "/"
 
 
-def get_account_id(boto3_session: Optional[boto3.Session] = None) -> str:
-    """Get Account ID."""
-    session: boto3.Session = ensure_session(session=boto3_session)
-    return client(service_name="sts", session=session).get_caller_identity().get("Account")
-
-
 def get_region_from_subnet(subnet_id: str, boto3_session: Optional[boto3.Session] = None) -> str:  # pragma: no cover
     """Extract region from Subnet ID."""
     session: boto3.Session = ensure_session(session=boto3_session)

awswrangler/athena.py

@@ -13,7 +13,7 @@
 import pandas as pd  # type: ignore
 import pyarrow as pa  # type: ignore
 
-from awswrangler import _data_types, _utils, catalog, exceptions, s3
+from awswrangler import _data_types, _utils, catalog, exceptions, s3, sts
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -72,7 +72,7 @@ def create_athena_bucket(boto3_session: Optional[boto3.Session] = None) -> str:
 
     """
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
-    account_id: str = _utils.get_account_id(boto3_session=session)
+    account_id: str = sts.get_account_id(boto3_session=session)
     region_name: str = str(session.region_name).lower()
     s3_output = f"s3://aws-athena-query-results-{account_id}-{region_name}/"
     s3_resource = session.resource("s3")
@@ -501,7 +501,7 @@ def read_sql_query(  # pylint: disable=too-many-branches,too-many-locals,too-man
 
     if cache_info["has_valid_cache"] is True:
         _logger.debug("Valid cache found. Retrieving...")
-        cache_result: Union[pd.DataFrame, Iterator[pd.DataFrame]] = None
+        cache_result: Union[pd.DataFrame, Iterator[pd.DataFrame], None] = None
         try:
             cache_result = _resolve_query_with_cache(
                 cache_info=cache_info,
@@ -1020,7 +1020,7 @@ def _check_for_cached_results(
             if (current_timestamp - query_info["Status"]["CompletionDateTime"]).total_seconds() > max_cache_seconds:
                 break
 
-            comparison_query: Optional[str] = None
+            comparison_query: Optional[str]
             if query_info["StatementType"] == "DDL" and query_info["Query"].startswith("CREATE TABLE"):
                 parsed_query: Optional[str] = _parse_select_query_from_possible_ctas(query_info["Query"])
                 if parsed_query is not None:

awswrangler/emr.py

@@ -7,7 +7,7 @@
 
 import boto3  # type: ignore
 
-from awswrangler import _utils, exceptions
+from awswrangler import _utils, exceptions, sts
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -48,7 +48,7 @@ def _get_default_logging_path(
     """
     if account_id is None:
         boto3_session = _utils.ensure_session(session=boto3_session)
-        _account_id: str = _utils.get_account_id(boto3_session=boto3_session)
+        _account_id: str = sts.get_account_id(boto3_session=boto3_session)
     else:
         _account_id = account_id
     if (region is None) and (subnet_id is not None):
@@ -61,7 +61,7 @@ def _get_default_logging_path(
 
 
 def _build_cluster_args(**pars):  # pylint: disable=too-many-branches,too-many-statements
-    account_id: str = _utils.get_account_id(boto3_session=pars["boto3_session"])
+    account_id: str = sts.get_account_id(boto3_session=pars["boto3_session"])
     region: str = _utils.get_region_from_session(boto3_session=pars["boto3_session"])
 
     # S3 Logging path
@@ -846,6 +846,7 @@ def build_step(
     Examples
     --------
     >>> import awswrangler as wr
+    >>> steps = []
     >>> for cmd in ['echo "Hello"', "ls -la"]:
     ...     steps.append(wr.emr.build_step(name=cmd, command=cmd))
     >>> wr.emr.submit_steps(cluster_id="cluster-id", steps=steps)

awswrangler/quicksight/_cancel.py

@@ -5,7 +5,7 @@
 
 import boto3  # type: ignore
 
-from awswrangler import _utils, exceptions
+from awswrangler import _utils, exceptions, sts
 from awswrangler.quicksight._get_list import get_dataset_id
 
 _logger: logging.Logger = logging.getLogger(__name__)
@@ -51,7 +51,7 @@ def cancel_ingestion(
         raise exceptions.InvalidArgument("You must pass a not None name or dataset_id argument.")
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     if (dataset_id is None) and (dataset_name is not None):
         dataset_id = get_dataset_id(name=dataset_name, account_id=account_id, boto3_session=session)
     client: boto3.client = _utils.client(service_name="quicksight", session=session)

awswrangler/quicksight/_create.py

@@ -6,7 +6,7 @@
 
 import boto3  # type: ignore
 
-from awswrangler import _utils, exceptions
+from awswrangler import _utils, exceptions, sts
 from awswrangler.quicksight._get_list import get_data_source_arn, get_dataset_id
 from awswrangler.quicksight._utils import extract_athena_query_columns, extract_athena_table_columns
 
@@ -157,7 +157,7 @@ def create_athena_data_source(
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     client: boto3.client = _utils.client(service_name="quicksight", session=session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     args: Dict[str, Any] = {
         "AwsAccountId": account_id,
         "DataSourceId": name,
@@ -282,7 +282,7 @@ def create_athena_dataset(
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     client: boto3.client = _utils.client(service_name="quicksight", session=session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     if (data_source_arn is None) and (data_source_name is not None):
         data_source_arn = get_data_source_arn(name=data_source_name, account_id=account_id, boto3_session=session)
     if sql is not None:
@@ -379,7 +379,7 @@ def create_ingestion(
     """
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     if (dataset_name is None) and (dataset_id is None):
         raise exceptions.InvalidArgument("You must pass a not None dataset_name or dataset_id argument.")
     if (dataset_id is None) and (dataset_name is not None):

awswrangler/quicksight/_delete.py

@@ -5,7 +5,7 @@
 
 import boto3  # type: ignore
 
-from awswrangler import _utils, exceptions
+from awswrangler import _utils, exceptions, sts
 from awswrangler.quicksight._get_list import (
     get_dashboard_id,
     get_data_source_id,
@@ -25,7 +25,7 @@ def _delete(
 ) -> None:
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     client: boto3.client = _utils.client(service_name="quicksight", session=session)
     func: Callable = getattr(client, func_name)
     func(AwsAccountId=account_id, **kwargs)
@@ -253,7 +253,7 @@ def delete_all_dashboards(account_id: Optional[str] = None, boto3_session: Optio
     """
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     for dashboard in list_dashboards(account_id=account_id, boto3_session=session):
         delete_dashboard(dashboard_id=dashboard["DashboardId"], account_id=account_id, boto3_session=session)
 
@@ -280,7 +280,7 @@ def delete_all_datasets(account_id: Optional[str] = None, boto3_session: Optiona
     """
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     for dataset in list_datasets(account_id=account_id, boto3_session=session):
         delete_dataset(dataset_id=dataset["DataSetId"], account_id=account_id, boto3_session=session)
 
@@ -307,7 +307,7 @@ def delete_all_data_sources(account_id: Optional[str] = None, boto3_session: Opt
     """
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     for data_source in list_data_sources(account_id=account_id, boto3_session=session):
         delete_data_source(data_source_id=data_source["DataSourceId"], account_id=account_id, boto3_session=session)
 
@@ -334,6 +334,6 @@ def delete_all_templates(account_id: Optional[str] = None, boto3_session: Option
     """
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     for template in list_templates(account_id=account_id, boto3_session=session):
         delete_template(template_id=template["TemplateId"], account_id=account_id, boto3_session=session)

awswrangler/quicksight/_describe.py

@@ -5,7 +5,7 @@
 
 import boto3  # type: ignore
 
-from awswrangler import _utils, exceptions
+from awswrangler import _utils, exceptions, sts
 from awswrangler.quicksight._get_list import get_dashboard_id, get_data_source_id, get_dataset_id
 
 _logger: logging.Logger = logging.getLogger(__name__)
@@ -48,7 +48,7 @@ def describe_dashboard(
         raise exceptions.InvalidArgument("You must pass a not None name or dashboard_id argument.")
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     if (dashboard_id is None) and (name is not None):
         dashboard_id = get_dashboard_id(name=name, account_id=account_id, boto3_session=session)
     client: boto3.client = _utils.client(service_name="quicksight", session=session)
@@ -92,7 +92,7 @@ def describe_data_source(
         raise exceptions.InvalidArgument("You must pass a not None name or data_source_id argument.")
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     if (data_source_id is None) and (name is not None):
         data_source_id = get_data_source_id(name=name, account_id=account_id, boto3_session=session)
     client: boto3.client = _utils.client(service_name="quicksight", session=session)
@@ -136,7 +136,7 @@ def describe_data_source_permissions(
         raise exceptions.InvalidArgument("You must pass a not None name or data_source_id argument.")
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     if (data_source_id is None) and (name is not None):
         data_source_id = get_data_source_id(name=name, account_id=account_id, boto3_session=session)
     client: boto3.client = _utils.client(service_name="quicksight", session=session)
@@ -180,7 +180,7 @@ def describe_dataset(
         raise exceptions.InvalidArgument("You must pass a not None name or dataset_id argument.")
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     if (dataset_id is None) and (name is not None):
         dataset_id = get_dataset_id(name=name, account_id=account_id, boto3_session=session)
     client: boto3.client = _utils.client(service_name="quicksight", session=session)
@@ -227,7 +227,7 @@ def describe_ingestion(
         raise exceptions.InvalidArgument("You must pass a not None name or dataset_id argument.")
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     if (dataset_id is None) and (dataset_name is not None):
         dataset_id = get_dataset_id(name=dataset_name, account_id=account_id, boto3_session=session)
     client: boto3.client = _utils.client(service_name="quicksight", session=session)

awswrangler/quicksight/_get_list.py

@@ -9,7 +9,7 @@
 
 import boto3  # type: ignore
 
-from awswrangler import _utils, exceptions
+from awswrangler import _utils, exceptions, sts
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -23,7 +23,7 @@ def _list(
 ) -> List[Dict[str, Any]]:
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     client: boto3.client = _utils.client(service_name="quicksight", session=session)
     func: Callable = getattr(client, func_name)
     response = func(AwsAccountId=account_id, **kwargs)
@@ -408,7 +408,7 @@ def list_ingestions(
         raise exceptions.InvalidArgument("You must pass a not None name or dataset_id argument.")
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     if account_id is None:
-        account_id = _utils.get_account_id(boto3_session=session)
+        account_id = sts.get_account_id(boto3_session=session)
     if (dataset_id is None) and (dataset_name is not None):
         dataset_id = get_dataset_id(name=dataset_name, account_id=account_id, boto3_session=session)
     return _list(

tests/test_moto.py

@@ -2,7 +2,7 @@
 
 import boto3
 import botocore
-import mock
+from unittest import mock
 import moto
 import pandas as pd
 import pytest