fix: opensearch session (#3170)

* fix: opensearch session

* fix: use session factory

* fix: types

* fix: types

* fix: doc8

Author: kukushking

awswrangler/_distributed.py

@@ -9,7 +9,7 @@
 from enum import Enum, unique
 from functools import wraps
 from importlib import reload
-from typing import Any, Callable, Literal, TypeVar, cast
+from typing import Any, Callable, Literal, TypeVar
 
 EngineLiteral = Literal["python", "ray"]
 MemoryFormatLiteral = Literal["pandas", "modin"]
@@ -112,7 +112,7 @@ def wrapper(*args: Any, **kw: dict[str, Any]) -> Any:
     def register(cls, name: EngineLiteral | None = None) -> None:
         """Register the distribution engine dispatch methods."""
         with cls._lock:
-            engine_name = cast(EngineLiteral, name or cls.get().value)
+            engine_name = name or cls.get().value
             cls.set(engine_name)
             cls._registry.clear()
 
@@ -125,7 +125,7 @@ def register(cls, name: EngineLiteral | None = None) -> None:
     def initialize(cls, name: EngineLiteral | None = None) -> None:
         """Initialize the distribution engine."""
         with cls._lock:
-            engine_name = cast(EngineLiteral, name or cls.get_installed().value)
+            engine_name = name or cls.get_installed().value
             if engine_name == EngineEnum.RAY.value:
                 from awswrangler.distributed.ray import initialize_ray
 
@@ -136,7 +136,7 @@ def initialize(cls, name: EngineLiteral | None = None) -> None:
     def is_initialized(cls, name: EngineLiteral | None = None) -> bool:
         """Check if the distribution engine is initialized."""
         with cls._lock:
-            engine_name = cast(EngineLiteral, name or cls.get_installed().value)
+            engine_name = name or cls.get_installed().value
 
             return False if not cls._initialized_engine else cls._initialized_engine.value == engine_name
 

awswrangler/athena/_write_iceberg.py

@@ -6,7 +6,7 @@
 import re
 import typing
 import uuid
-from typing import Any, Dict, Literal, TypedDict, cast
+from typing import Any, Dict, Literal, TypedDict
 
 import boto3
 import pandas as pd
@@ -501,10 +501,7 @@ def to_iceberg(  # noqa: PLR0913
         merge_condition=merge_condition,
     )
 
-    glue_table_settings = cast(
-        GlueTableSettings,
-        glue_table_settings if glue_table_settings else {},
-    )
+    glue_table_settings = glue_table_settings if glue_table_settings else {}
 
     try:
         # Create Iceberg table if it doesn't exist

awswrangler/catalog/_utils.py

@@ -196,7 +196,7 @@ def sanitize_dataframe_columns_names(df: pd.DataFrame, handle_duplicate_columns:
     """
     df.columns = [sanitize_column_name(x) for x in df.columns]
     df.index.names = [None if x is None else sanitize_column_name(x) for x in df.index.names]
-    if df.columns.duplicated().any():  # type: ignore[attr-defined]
+    if df.columns.duplicated().any():
         if handle_duplicate_columns == "warn":
             warnings.warn(
                 "Duplicate columns were detected, consider using `handle_duplicate_columns='[drop|rename]'`",

awswrangler/opensearch/_write.py

@@ -280,7 +280,7 @@ def index_json(
     path: str,
     index: str,
     doc_type: str | None = None,
-    boto3_session: boto3.Session | None = boto3.Session(),
+    boto3_session: boto3.Session | None = None,
     json_path: str | None = None,
     use_threads: bool | int = False,
     **kwargs: Any,
@@ -334,12 +334,9 @@ def index_json(
     """
     _logger.debug("indexing %s from %s", index, path)
 
-    if boto3_session is None:
-        raise ValueError("boto3_session cannot be None")
-
     if path.startswith("s3://"):
         bucket, key = parse_path(path)
-        s3 = boto3_session.client("s3")
+        s3 = _utils.client(service_name="s3", session=boto3_session)
         obj = s3.get_object(Bucket=bucket, Key=key)
         body = obj["Body"].read()
         lines = body.splitlines()

awswrangler/quicksight/_get_list.py

@@ -28,7 +28,7 @@ def _list(
     client = _utils.client(service_name="quicksight", session=boto3_session)
     func: Callable[..., dict[str, Any]] = getattr(client, func_name)
     response: dict[str, Any] = func(AwsAccountId=account_id, **kwargs)
-    next_token: str = response.get("NextToken", None)
+    next_token: str | None = response.get("NextToken", None)
     result: list[dict[str, Any]] = response[attr_name]
     while next_token is not None:
         response = func(AwsAccountId=account_id, NextToken=next_token, **kwargs)

awswrangler/s3/_write_orc.py

@@ -5,7 +5,7 @@
 import logging
 import math
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Callable, Iterator, Literal, cast
+from typing import TYPE_CHECKING, Any, Callable, Iterator, Literal
 
 import boto3
 import pandas as pd
@@ -620,10 +620,7 @@ def to_orc(
     }
 
     """
-    glue_table_settings = cast(
-        GlueTableSettings,
-        glue_table_settings if glue_table_settings else {},
-    )
+    glue_table_settings = glue_table_settings if glue_table_settings else {}
 
     table_type = glue_table_settings.get("table_type")
     description = glue_table_settings.get("description")

awswrangler/s3/_write_parquet.py

@@ -5,7 +5,7 @@
 import logging
 import math
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Callable, Iterator, Literal, cast
+from typing import TYPE_CHECKING, Any, Callable, Iterator, Literal
 
 import boto3
 import pandas as pd
@@ -678,10 +678,7 @@ def to_parquet(
     }
 
     """
-    glue_table_settings = cast(
-        GlueTableSettings,
-        glue_table_settings if glue_table_settings else {},
-    )
+    glue_table_settings = glue_table_settings if glue_table_settings else {}
 
     table_type = glue_table_settings.get("table_type")
     description = glue_table_settings.get("description")

awswrangler/s3/_write_text.py

@@ -5,7 +5,7 @@
 import csv
 import logging
 import uuid
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal
 
 import boto3
 import pandas as pd
@@ -435,10 +435,7 @@ def to_csv(  # noqa: PLR0912,PLR0915
             "e.g. wr.s3.to_csv(df, path, sep='|', na_rep='NULL', decimal=',', compression='gzip')"
         )
 
-    glue_table_settings = cast(
-        GlueTableSettings,
-        glue_table_settings if glue_table_settings else {},
-    )
+    glue_table_settings = glue_table_settings if glue_table_settings else {}
 
     table_type = glue_table_settings.get("table_type")
     description = glue_table_settings.get("description")
@@ -885,10 +882,7 @@ def to_json(  # noqa: PLR0912,PLR0915
             "e.g. wr.s3.to_json(df, path, lines=True, date_format='iso')"
         )
 
-    glue_table_settings = cast(
-        GlueTableSettings,
-        glue_table_settings if glue_table_settings else {},
-    )
+    glue_table_settings = glue_table_settings if glue_table_settings else {}
 
     table_type = glue_table_settings.get("table_type")
     description = glue_table_settings.get("description")

docs/source/layers.rst

@@ -607,4 +607,4 @@ Version 3.12.1
 | cn-northwest-1 | 3.9    | x86_64| arn:aws-cn:lambda:cn-northwest-1:406640652441:layer:AWSSDKPandas-Python39:19      |
 +----------------+--------+-------+-----------------------------------------------------------------------------------+
 | cn-northwest-1 | 3.9    | arm64 | arn:aws-cn:lambda:cn-northwest-1:406640652441:layer:AWSSDKPandas-Python39-Arm64:3 |
-+----------------+--------+-------+-----------------------------------------------------------------------------------+
++----------------+--------+-------+-----------------------------------------------------------------------------------+