Merge branch 'main' into release-3.0.0

Author: kukushking

awswrangler/redshift.py

@@ -434,6 +434,7 @@ def connect(
     timeout: Optional[int] = None,
     max_prepared_statements: int = 1000,
     tcp_keepalive: bool = True,
+    **kwargs: Any,
 ) -> redshift_connector.Connection:
     """Return a redshift_connector connection from a Glue Catalog or Secret Manager.
 
@@ -455,23 +456,23 @@ def connect(
 
     Parameters
     ----------
-    connection : Optional[str]
+    connection : str, optional
         Glue Catalog Connection name.
     secret_id : Optional[str]:
         Specifies the secret containing the connection details that you want to retrieve.
         You can specify either the Amazon Resource Name (ARN) or the friendly name of the secret.
     catalog_id : str, optional
         The ID of the Data Catalog.
         If none is provided, the AWS account ID is used by default.
-    dbname : Optional[str]
+    dbname : str, optional
         Optional database name to overwrite the stored one.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
     ssl : bool
         This governs SSL encryption for TCP/IP sockets.
         This parameter is forward to redshift_connector.
         https://github.com/aws/amazon-redshift-python-driver
-    timeout : Optional[int]
+    timeout : int, optional
         This is the time in seconds before the connection to the server will time out.
         The default is None which means no timeout.
         This parameter is forward to redshift_connector.
@@ -483,6 +484,9 @@ def connect(
         If True then use TCP keepalive. The default is True.
         This parameter is forward to redshift_connector.
         https://github.com/aws/amazon-redshift-python-driver
+    **kwargs : Any
+        Forwarded to redshift_connector.connect.
+        e.g. is_serverless=True, serverless_acct_id='...', serverless_work_group='...'
 
     Returns
     -------
@@ -527,6 +531,7 @@ def connect(
         timeout=timeout,
         max_prepared_statements=max_prepared_statements,
         tcp_keepalive=tcp_keepalive,
+        **kwargs,
     )
 
 
@@ -542,6 +547,7 @@ def connect_temp(
     timeout: Optional[int] = None,
     max_prepared_statements: int = 1000,
     tcp_keepalive: bool = True,
+    **kwargs: Any,
 ) -> redshift_connector.Connection:
     """Return a redshift_connector temporary connection (No password required).
 
@@ -571,7 +577,7 @@ def connect_temp(
         This governs SSL encryption for TCP/IP sockets.
         This parameter is forward to redshift_connector.
         https://github.com/aws/amazon-redshift-python-driver
-    timeout : Optional[int]
+    timeout : int, optional
         This is the time in seconds before the connection to the server will time out.
         The default is None which means no timeout.
         This parameter is forward to redshift_connector.
@@ -583,6 +589,9 @@ def connect_temp(
         If True then use TCP keepalive. The default is True.
         This parameter is forward to redshift_connector.
         https://github.com/aws/amazon-redshift-python-driver
+    **kwargs : Any
+        Forwarded to redshift_connector.connect.
+        e.g. is_serverless=True, serverless_acct_id='...', serverless_work_group='...'
 
     Returns
     -------
@@ -623,6 +632,7 @@ def connect_temp(
         max_prepared_statements=max_prepared_statements,
         tcp_keepalive=tcp_keepalive,
         db_groups=db_groups,
+        **kwargs,
     )
 
 

awswrangler/s3/_list.pyi

@@ -0,0 +1,79 @@
+import datetime
+from typing import Any, Dict, Iterator, List, Literal, Optional, Sequence, Union, overload
+
+import boto3
+
+def _path2list(
+    path: Union[str, Sequence[str]],
+    boto3_session: boto3.Session,
+    s3_additional_kwargs: Optional[Dict[str, Any]] = ...,
+    last_modified_begin: Optional[datetime.datetime] = ...,
+    last_modified_end: Optional[datetime.datetime] = ...,
+    suffix: Union[str, List[str], None] = ...,
+    ignore_suffix: Union[str, List[str], None] = ...,
+    ignore_empty: bool = ...,
+) -> List[str]: ...
+def _prefix_cleanup(prefix: str) -> str: ...
+def does_object_exist(
+    path: str,
+    s3_additional_kwargs: Optional[Dict[str, Any]] = ...,
+    boto3_session: Optional[boto3.Session] = ...,
+    version_id: Optional[str] = ...,
+) -> bool: ...
+def list_buckets(boto3_session: Optional[boto3.Session] = ...) -> List[str]: ...
+@overload
+def list_directories(
+    path: str,
+    chunked: Literal[False],
+    s3_additional_kwargs: Union[Dict[str, Any], Dict[str, str], None] = ...,
+    boto3_session: Optional[boto3.Session] = ...,
+) -> List[str]: ...
+@overload
+def list_directories(
+    path: str,
+    *,
+    boto3_session: Optional[boto3.Session] = ...,
+    s3_additional_kwargs: Union[Dict[str, Any], Dict[str, str], None] = ...,
+) -> List[str]: ...
+@overload
+def list_directories(
+    path: str,
+    chunked: Literal[True],
+    s3_additional_kwargs: Union[Dict[str, Any], Dict[str, str], None] = ...,
+    boto3_session: Optional[boto3.Session] = ...,
+) -> Iterator[List[str]]: ...
+@overload
+def list_directories(
+    path: str,
+    chunked: bool,
+    s3_additional_kwargs: Union[Dict[str, Any], Dict[str, str], None] = ...,
+    boto3_session: Optional[boto3.Session] = ...,
+) -> Union[List[str], Iterator[List[str]]]: ...
+@overload
+def list_objects(
+    path: str,
+    chunked: Literal[False],
+    s3_additional_kwargs: Union[Dict[str, Any], Dict[str, str], None] = ...,
+    boto3_session: Optional[boto3.Session] = ...,
+) -> List[str]: ...
+@overload
+def list_objects(
+    path: str,
+    *,
+    s3_additional_kwargs: Union[Dict[str, Any], Dict[str, str], None] = ...,
+    boto3_session: Optional[boto3.Session] = ...,
+) -> List[str]: ...
+@overload
+def list_objects(
+    path: str,
+    chunked: Literal[True],
+    s3_additional_kwargs: Union[Dict[str, Any], Dict[str, str], None] = ...,
+    boto3_session: Optional[boto3.Session] = ...,
+) -> Iterator[List[str]]: ...
+@overload
+def list_objects(
+    path: str,
+    chunked: bool,
+    s3_additional_kwargs: Union[Dict[str, Any], Dict[str, str], None] = ...,
+    boto3_session: Optional[boto3.Session] = ...,
+) -> Union[List[str], Iterator[List[str]]]: ...

tests/test_catalog.py

@@ -36,14 +36,17 @@ def test_create_table(path: str, glue_database: str, glue_table: str, table_type
     assert wr.catalog.does_table_exist(database=glue_database, table=glue_table) is True
 
 
-@pytest.mark.parametrize("table_type", ["EXTERNAL_TABLE", "GOVERNED"])
-@pytest.mark.parametrize("start_transaction", [True, False])
+@pytest.mark.parametrize(
+    ("table_type", "start_transaction"),
+    [
+        ("EXTERNAL_TABLE", False),
+        ("EXTERNAL_TABLE", True),
+        pytest.param("GOVERNED", False, marks=pytest.mark.xfail(reason="TransactionCommitInProgressException")),
+    ],
+)
 def test_catalog(
     path: str, glue_database: str, glue_table: str, table_type: Optional[str], start_transaction: bool, account_id: str
 ) -> None:
-    if table_type != "GOVERNED" and start_transaction:
-        pytest.skip()
-
     transaction_id = wr.lakeformation.start_transaction() if table_type == "GOVERNED" else None
     wr.catalog.create_parquet_table(
         database=glue_database,

tests/test_redshift.py

@@ -28,7 +28,7 @@ def redshift_con():
 
 
 def test_connection():
-    wr.redshift.connect("aws-data-wrangler-redshift", timeout=10).close()
+    wr.redshift.connect("aws-data-wrangler-redshift", timeout=10, force_lowercase=False).close()
 
 
 def test_read_sql_query_simple(databases_parameters):