Making datetime timezone aware in CW module and adding security token in Redshift (#581)

Author: jaidisido

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,

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