Improve AWS STS module.

Author: igorborgest

README.md

@@ -101,10 +101,11 @@ wr.db.to_sql(df, engine, schema="test", name="my_table")
   - [Amazon S3](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#amazon-s3)
   - [AWS Glue Catalog](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#aws-glue-catalog)
   - [Amazon Athena](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#amazon-athena)
-  - [Databases (Redshift, PostgreSQL, MySQL)](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#databases-redshift-postgresql-mysql)
-  - [EMR Cluster](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#emr-cluster)
-  - [CloudWatch Logs](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#cloudwatch-logs)
-  - [QuickSight](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#quicksight)
+  - [Databases (Amazon Redshift, PostgreSQL, MySQL)](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#databases-amazon-redshift-postgresql-mysql)
+  - [Amazon EMR](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#amazon-emr)
+  - [Amazon CloudWatch Logs](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#amazon-cloudwatch-logs)
+  - [Amazon QuickSight](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#amazon-quicksight)
+  - [AWS STS](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#aws-sts)
 - [**License**](https://github.com/awslabs/aws-data-wrangler/blob/master/LICENSE)
 - [**Contributing**](https://github.com/awslabs/aws-data-wrangler/blob/master/CONTRIBUTING.md)
 - [**Legacy Docs** (pre-1.0.0)](https://aws-data-wrangler.readthedocs.io/en/legacy/)

awswrangler/sts.py

@@ -0,0 +1,80 @@
+"""STS module."""
+
+import logging
+from typing import Optional
+
+import boto3  # type: ignore
+
+from awswrangler import _utils
+
+_logger: logging.Logger = logging.getLogger(__name__)
+
+
+def get_account_id(boto3_session: Optional[boto3.Session] = None) -> str:
+    """Get Account ID.
+
+    Parameters
+    ----------
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    str
+        Account ID.
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> account_id = wr.sts.get_account_id()
+
+    """
+    session: boto3.Session = _utils.ensure_session(session=boto3_session)
+    return _utils.client(service_name="sts", session=session).get_caller_identity().get("Account")
+
+
+def get_current_identity_arn(boto3_session: Optional[boto3.Session] = None) -> str:
+    """Get current user/role ARN.
+
+    Parameters
+    ----------
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    str
+        User/role ARN.
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> arn = wr.sts.get_current_identity_arn()
+
+    """
+    session: boto3.Session = _utils.ensure_session(session=boto3_session)
+    return _utils.client(service_name="sts", session=session).get_caller_identity().get("Arn")
+
+
+def get_current_identity_name(boto3_session: Optional[boto3.Session] = None) -> str:
+    """Get current user/role name.
+
+    Parameters
+    ----------
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    str
+        User/role name.
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> name = wr.sts.get_current_identity_name()
+
+    """
+    arn: str = get_current_identity_arn(boto3_session=boto3_session)
+    name: str = arn.rpartition("/")[-1]
+    return name

docs/source/api.rst

@@ -87,8 +87,8 @@ Amazon Athena
     stop_query_execution
     wait_query
 
-Databases (Redshift, PostgreSQL, MySQL)
----------------------------------------
+Databases (Amazon Redshift, PostgreSQL, MySQL)
+----------------------------------------------
 
 .. currentmodule:: awswrangler.db
 
@@ -106,8 +106,8 @@ Databases (Redshift, PostgreSQL, MySQL)
     unload_redshift_to_files
     write_redshift_copy_manifest
 
-EMR
----
+Amazon EMR
+----------
 
 .. currentmodule:: awswrangler.emr
 
@@ -125,8 +125,8 @@ EMR
     submit_steps
     terminate_cluster
 
-CloudWatch Logs
----------------
+Amazon CloudWatch Logs
+----------------------
 
 .. currentmodule:: awswrangler.cloudwatch
 
@@ -184,3 +184,15 @@ Amazon QuickSight
     list_templates
     list_users
     list_user_groups
+
+AWS STS
+-----------------
+
+.. currentmodule:: awswrangler.sts
+
+.. autosummary::
+    :toctree: stubs
+
+    get_account_id
+    get_current_identity_arn
+    get_current_identity_name

tests/test_moto.py

@@ -1,8 +1,8 @@
+from unittest import mock
 from unittest.mock import ANY
 
 import boto3
 import botocore
-from unittest import mock
 import moto
 import pandas as pd
 import pytest