From 4e59f872fcd1ca82cfc4b740986daa143849c81d Mon Sep 17 00:00:00 2001 From: Abdel Jaidi Date: Thu, 4 Mar 2021 18:06:29 +0000 Subject: [PATCH] Making datetime timezone aware in CW module and adding security token in Redshift --- awswrangler/cloudwatch.py | 17 ++++++++++++++--- tests/test_redshift.py | 5 +++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/awswrangler/cloudwatch.py b/awswrangler/cloudwatch.py index 7da49b0a9..03f632087 100644 --- a/awswrangler/cloudwatch.py +++ b/awswrangler/cloudwatch.py @@ -15,10 +15,20 @@ _QUERY_WAIT_POLLING_DELAY: float = 0.2 # SECONDS +def _validate_args( + start_timestamp: int, + end_timestamp: int, +) -> None: + if start_timestamp < 0: + raise exceptions.InvalidArgument("`start_time` cannot be a negative value.") + if start_timestamp >= end_timestamp: + raise exceptions.InvalidArgumentCombination("`start_time` must be inferior to `end_time`.") + + def start_query( query: str, log_group_names: List[str], - start_time: datetime.datetime = datetime.datetime(year=1970, month=1, day=1), + start_time: datetime.datetime = datetime.datetime(year=1970, month=1, day=1, tzinfo=datetime.timezone.utc), end_time: datetime.datetime = datetime.datetime.now(), limit: Optional[int] = None, boto3_session: Optional[boto3.Session] = None, @@ -61,6 +71,7 @@ def start_query( end_timestamp: int = int(1000 * end_time.timestamp()) _logger.debug("start_timestamp: %s", start_timestamp) _logger.debug("end_timestamp: %s", end_timestamp) + _validate_args(start_timestamp=start_timestamp, end_timestamp=end_timestamp) args: Dict[str, Any] = { "logGroupNames": log_group_names, "startTime": start_timestamp, @@ -120,7 +131,7 @@ def wait_query(query_id: str, boto3_session: Optional[boto3.Session] = None) -> def run_query( query: str, log_group_names: List[str], - start_time: datetime.datetime = datetime.datetime(year=1970, month=1, day=1), + start_time: datetime.datetime = datetime.datetime(year=1970, month=1, day=1, tzinfo=datetime.timezone.utc), end_time: datetime.datetime = datetime.datetime.now(), limit: Optional[int] = None, boto3_session: Optional[boto3.Session] = None, @@ -174,7 +185,7 @@ def run_query( def read_logs( query: str, log_group_names: List[str], - start_time: datetime.datetime = datetime.datetime(year=1970, month=1, day=1), + start_time: datetime.datetime = datetime.datetime(year=1970, month=1, day=1, tzinfo=datetime.timezone.utc), end_time: datetime.datetime = datetime.datetime.now(), limit: Optional[int] = None, boto3_session: Optional[boto3.Session] = None, diff --git a/tests/test_redshift.py b/tests/test_redshift.py index 893dac79f..defa245c1 100644 --- a/tests/test_redshift.py +++ b/tests/test_redshift.py @@ -831,6 +831,7 @@ def test_copy_unload_creds(path, redshift_table): mode="overwrite", aws_access_key_id=credentials.access_key, aws_secret_access_key=credentials.secret_key, + aws_session_token=credentials.token, ) df2 = wr.redshift.unload( sql=f"SELECT * FROM public.{redshift_table}", @@ -839,6 +840,7 @@ def test_copy_unload_creds(path, redshift_table): keep_files=False, aws_access_key_id=credentials.access_key, aws_secret_access_key=credentials.secret_key, + aws_session_token=credentials.token, ) assert df2.shape == (3, 1) wr.redshift.copy( @@ -850,6 +852,7 @@ def test_copy_unload_creds(path, redshift_table): mode="append", aws_access_key_id=credentials.access_key, aws_secret_access_key=credentials.secret_key, + aws_session_token=credentials.token, ) df2 = wr.redshift.unload( sql=f"SELECT * FROM public.{redshift_table}", @@ -858,6 +861,7 @@ def test_copy_unload_creds(path, redshift_table): keep_files=False, aws_access_key_id=credentials.access_key, aws_secret_access_key=credentials.secret_key, + aws_session_token=credentials.token, ) assert df2.shape == (6, 1) dfs = wr.redshift.unload( @@ -868,6 +872,7 @@ def test_copy_unload_creds(path, redshift_table): chunked=True, aws_access_key_id=credentials.access_key, aws_secret_access_key=credentials.secret_key, + aws_session_token=credentials.token, ) for chunk in dfs: assert len(chunk.columns) == 1