Fix apply_configs decorator causing function signature to be lost (#1858)

Author: LeonLuttenberger

awswrangler/_config.py

@@ -3,7 +3,7 @@
 import inspect
 import logging
 import os
-from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, Type, Union, cast
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, Type, TypeVar, Union, cast
 
 import botocore.config
 import pandas as pd
@@ -466,7 +466,10 @@ def _inject_config_doc(doc: Optional[str], available_configs: Tuple[str, ...]) -
     return _insert_str(text=doc, token="\n    Parameters", insert=insertion)
 
 
-def apply_configs(function: Callable[..., Any]) -> Callable[..., Any]:
+FunctionType = TypeVar("FunctionType", bound=Callable[..., Any])
+
+
+def apply_configs(function: FunctionType) -> FunctionType:
     """Decorate some function with configs."""
     signature = inspect.signature(function)
     args_names: Tuple[str, ...] = tuple(signature.parameters.keys())
@@ -495,7 +498,7 @@ def wrapper(*args_raw: Any, **kwargs: Any) -> Any:
     wrapper.__doc__ = _inject_config_doc(doc=function.__doc__, available_configs=available_configs)
     wrapper.__name__ = function.__name__
     wrapper.__setattr__("__signature__", signature)  # pylint: disable=no-member
-    return wrapper
+    return wrapper  # type: ignore
 
 
 config: _Config = _Config()

awswrangler/athena/_read.py

@@ -106,6 +106,9 @@ def _fetch_parquet_result(
 
         database, temp_table_name = map(lambda x: x.replace('"', ""), temp_table_fqn.split("."))
         dtype_dict = catalog.get_table_types(database=database, table=temp_table_name, boto3_session=boto3_session)
+        if dtype_dict is None:
+            raise exceptions.ResourceDoesNotExist(f"Temp table {temp_table_fqn} not found.")
+
         df = pd.DataFrame(columns=list(dtype_dict.keys()))
         df = cast_pandas_with_athena_types(df=df, dtype=dtype_dict)
         df = _apply_query_metadata(df=df, query_metadata=query_metadata)

awswrangler/athena/_utils.py

@@ -554,15 +554,18 @@ def repair_table(
     if (database is not None) and (not database.startswith("`")):
         database = f"`{database}`"
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
-    query_id = start_query_execution(
-        sql=query,
-        database=database,
-        data_source=data_source,
-        s3_output=s3_output,
-        workgroup=workgroup,
-        encryption=encryption,
-        kms_key=kms_key,
-        boto3_session=session,
+    query_id = cast(
+        str,
+        start_query_execution(
+            sql=query,
+            database=database,
+            data_source=data_source,
+            s3_output=s3_output,
+            workgroup=workgroup,
+            encryption=encryption,
+            kms_key=kms_key,
+            boto3_session=session,
+        ),
     )
     response: Dict[str, Any] = wait_query(query_execution_id=query_id, boto3_session=session)
     return cast(str, response["Status"]["State"])
@@ -624,14 +627,17 @@ def describe_table(
     if (database is not None) and (not database.startswith("`")):
         database = f"`{database}`"
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
-    query_id = start_query_execution(
-        sql=query,
-        database=database,
-        s3_output=s3_output,
-        workgroup=workgroup,
-        encryption=encryption,
-        kms_key=kms_key,
-        boto3_session=session,
+    query_id = cast(
+        str,
+        start_query_execution(
+            sql=query,
+            database=database,
+            s3_output=s3_output,
+            workgroup=workgroup,
+            encryption=encryption,
+            kms_key=kms_key,
+            boto3_session=session,
+        ),
     )
     query_metadata: _QueryMetadata = _get_query_metadata(query_execution_id=query_id, boto3_session=session)
     raw_result = _fetch_txt_result(
@@ -643,7 +649,7 @@ def describe_table(
 @apply_configs
 def create_ctas_table(  # pylint: disable=too-many-locals
     sql: str,
-    database: str,
+    database: Optional[str] = None,
     ctas_table: Optional[str] = None,
     ctas_database: Optional[str] = None,
     s3_output: Optional[str] = None,
@@ -669,7 +675,7 @@ def create_ctas_table(  # pylint: disable=too-many-locals
     ----------
     sql : str
         SELECT SQL query.
-    database : str
+    database : Optional[str], optional
         The name of the database where the original table is stored.
     ctas_table : Optional[str], optional
         The name of the CTAS table.
@@ -756,6 +762,10 @@ def create_ctas_table(  # pylint: disable=too-many-locals
     """
     ctas_table = catalog.sanitize_table_name(ctas_table) if ctas_table else f"temp_table_{uuid.uuid4().hex}"
     ctas_database = ctas_database if ctas_database else database
+
+    if ctas_database is None:
+        raise exceptions.InvalidArgumentCombination("Either ctas_database or database must be defined.")
+
     fully_qualified_name = f'"{ctas_database}"."{ctas_table}"'
 
     wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
@@ -904,14 +914,17 @@ def show_create_table(
     if (database is not None) and (not database.startswith("`")):
         database = f"`{database}`"
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
-    query_id = start_query_execution(
-        sql=query,
-        database=database,
-        s3_output=s3_output,
-        workgroup=workgroup,
-        encryption=encryption,
-        kms_key=kms_key,
-        boto3_session=session,
+    query_id = cast(
+        str,
+        start_query_execution(
+            sql=query,
+            database=database,
+            s3_output=s3_output,
+            workgroup=workgroup,
+            encryption=encryption,
+            kms_key=kms_key,
+            boto3_session=session,
+        ),
     )
     query_metadata: _QueryMetadata = _get_query_metadata(query_execution_id=query_id, boto3_session=session)
     raw_result = _fetch_txt_result(

awswrangler/exceptions.py

@@ -115,3 +115,7 @@ class FailedQualityCheck(Exception):
 
 class AlreadyExists(Exception):
     """AlreadyExists."""
+
+
+class ResourceDoesNotExist(Exception):
+    """ResourceDoesNotExist."""

awswrangler/quicksight/_utils.py

@@ -1,7 +1,7 @@
 """Internal (private) Amazon QuickSight Utilities Module."""
 
 import logging
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, cast
 
 import boto3
 
@@ -29,7 +29,9 @@ def extract_athena_query_columns(
     data_source: Dict[str, Any] = [x for x in data_sources if x["Arn"] == data_source_arn][0]
     workgroup: str = data_source["DataSourceParameters"]["AthenaParameters"]["WorkGroup"]
     sql_wrapped: str = f"/* QuickSight */\nSELECT ds.* FROM ( {sql} ) ds LIMIT 0"
-    query_id: str = athena.start_query_execution(sql=sql_wrapped, workgroup=workgroup, boto3_session=boto3_session)
+    query_id = cast(
+        str, athena.start_query_execution(sql=sql_wrapped, workgroup=workgroup, boto3_session=boto3_session)
+    )
     athena.wait_query(query_execution_id=query_id, boto3_session=boto3_session)
     dtypes: Dict[str, str] = athena.get_query_columns_types(query_execution_id=query_id, boto3_session=boto3_session)
     return [{"Name": name, "Type": _data_types.athena2quicksight(dtype=dtype)} for name, dtype in dtypes.items()]

awswrangler/redshift.py

@@ -254,7 +254,7 @@ def _validate_parameters(
 
 
 def _redshift_types_from_path(
-    path: Optional[Union[str, List[str]]],
+    path: Union[str, List[str]],
     varchar_lengths_default: int,
     varchar_lengths: Optional[Dict[str, int]],
     parquet_infer_sampling: float,

awswrangler/s3/_download.py

@@ -1,7 +1,7 @@
 """Amazon S3 Download Module (PRIVATE)."""
 
 import logging
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional, Union, cast
 
 import boto3
 
@@ -76,7 +76,7 @@ def download(
         if isinstance(local_file, str):
             _logger.debug("Downloading local_file: %s", local_file)
             with open(file=local_file, mode="wb") as local_f:
-                local_f.write(s3_f.read())
+                local_f.write(cast(bytes, s3_f.read()))
         else:
             _logger.debug("Downloading file-like object.")
             local_file.write(s3_f.read())

awswrangler/s3/_fs.py

@@ -228,7 +228,7 @@ def __init__(
         else:
             raise RuntimeError(f"Invalid mode: {self._mode}")
 
-    def __enter__(self) -> Union["_S3ObjectBase"]:
+    def __enter__(self) -> "_S3ObjectBase":
         return self
 
     def __exit__(self, exc_type: Any, exc_value: Any, exc_traceback: Any) -> None:

awswrangler/s3/_merge_upsert_table.py

@@ -77,6 +77,9 @@ def _generate_empty_frame_for_table(
     boto3_session: Optional[boto3.Session] = None,
 ) -> pandas.DataFrame:
     type_dict = wr.catalog.get_table_types(database=database, table=table, boto3_session=boto3_session)
+    if type_dict is None:
+        raise wr.exceptions.ResourceDoesNotExist(f"Table {table} from database {database} does not exist.")
+
     empty_frame = pandas.DataFrame(columns=type_dict.keys())
     return _data_types.cast_pandas_with_athena_types(empty_frame, type_dict)
 

awswrangler/s3/_upload.py

@@ -71,7 +71,7 @@ def upload(
         if isinstance(local_file, str):
             _logger.debug("Uploading local_file: %s", local_file)
             with open(file=local_file, mode="rb") as local_f:
-                s3_f.write(local_f.read())
+                s3_f.write(local_f.read())  # type: ignore
         else:
             _logger.debug("Uploading file-like object.")
             s3_f.write(local_file.read())