Merge remote-tracking branch 'upstream/main' into fix/#2409

Author: bujjibabukatta

awswrangler/dynamodb/_read.py

@@ -456,7 +456,7 @@ def _convert_condition_base_to_expression(
 @_utils.validate_distributed_kwargs(
     unsupported_kwargs=["boto3_session", "dtype_backend"],
 )
-def read_items(  # noqa: PLR0912
+def read_items(  # noqa: PLR0912, PLR0915
     table_name: str,
     index_name: str | None = None,
     partition_values: Sequence[Any] | None = None,
@@ -475,6 +475,7 @@ def read_items(  # noqa: PLR0912
     use_threads: bool | int = True,
     boto3_session: boto3.Session | None = None,
     pyarrow_additional_kwargs: dict[str, Any] | None = None,
+    key_schema: list[dict[str, str]] | None = None,
 ) -> pd.DataFrame | Iterator[pd.DataFrame] | _ItemsListType | Iterator[_ItemsListType]:
     """Read items from given DynamoDB table.
 
@@ -551,6 +552,10 @@ def read_items(  # noqa: PLR0912
         Forwarded to `to_pandas` method converting from PyArrow tables to Pandas DataFrame.
         Valid values include "split_blocks", "self_destruct", "ignore_metadata".
         e.g. pyarrow_additional_kwargs={'split_blocks': True}.
+    key_schema
+        Key schema of the table (e.g. `[{"AttributeName": "key", "KeyType": "HASH"}]`).
+        If provided, the library will bypass the `DescribeTable` API call, which can
+        reduce network latency and prevent API throttling. Defaults to None.
 
     Raises
     ------
@@ -657,7 +662,10 @@ def read_items(  # noqa: PLR0912
     # Extract key schema
     dynamodb_client = _utils.client(service_name="dynamodb", session=boto3_session)
     serializer = TypeSerializer()
-    table_key_schema = dynamodb_client.describe_table(TableName=table_name)["Table"]["KeySchema"]
+    if key_schema:
+        table_key_schema = key_schema
+    else:
+        table_key_schema = dynamodb_client.describe_table(TableName=table_name)["Table"]["KeySchema"]
 
     # Detect sort key, if any
     if len(table_key_schema) == 1:

awswrangler/dynamodb/_read.pyi

@@ -68,6 +68,7 @@ def read_items(
     use_threads: bool | int = ...,
     boto3_session: boto3.Session | None = ...,
     pyarrow_additional_kwargs: dict[str, Any] | None = ...,
+    key_schema: list[dict[str, str]] | None = ...,
 ) -> pd.DataFrame: ...
 @overload
 def read_items(
@@ -90,6 +91,7 @@ def read_items(
     use_threads: bool | int = ...,
     boto3_session: boto3.Session | None = ...,
     pyarrow_additional_kwargs: dict[str, Any] | None = ...,
+    key_schema: list[dict[str, str]] | None = ...,
 ) -> Iterator[pd.DataFrame]: ...
 @overload
 def read_items(
@@ -112,6 +114,7 @@ def read_items(
     use_threads: bool | int = ...,
     boto3_session: boto3.Session | None = ...,
     pyarrow_additional_kwargs: dict[str, Any] | None = ...,
+    key_schema: list[dict[str, str]] | None = ...,
 ) -> _ItemsListType: ...
 @overload
 def read_items(
@@ -134,6 +137,7 @@ def read_items(
     use_threads: bool | int = ...,
     boto3_session: boto3.Session | None = ...,
     pyarrow_additional_kwargs: dict[str, Any] | None = ...,
+    key_schema: list[dict[str, str]] | None = ...,
 ) -> Iterator[_ItemsListType]: ...
 @overload
 def read_items(
@@ -156,6 +160,7 @@ def read_items(
     use_threads: bool | int = ...,
     boto3_session: boto3.Session | None = ...,
     pyarrow_additional_kwargs: dict[str, Any] | None = ...,
+    key_schema: list[dict[str, str]] | None = ...,
 ) -> pd.DataFrame | _ItemsListType: ...
 @overload
 def read_items(
@@ -178,6 +183,7 @@ def read_items(
     use_threads: bool | int = ...,
     boto3_session: boto3.Session | None = ...,
     pyarrow_additional_kwargs: dict[str, Any] | None = ...,
+    key_schema: list[dict[str, str]] | None = ...,
 ) -> Iterator[pd.DataFrame] | Iterator[_ItemsListType]: ...
 @overload
 def read_items(
@@ -200,6 +206,7 @@ def read_items(
     use_threads: bool | int = ...,
     boto3_session: boto3.Session | None = ...,
     pyarrow_additional_kwargs: dict[str, Any] | None = ...,
+    key_schema: list[dict[str, str]] | None = ...,
 ) -> pd.DataFrame | Iterator[pd.DataFrame]: ...
 @overload
 def read_items(
@@ -222,6 +229,7 @@ def read_items(
     use_threads: bool | int = ...,
     boto3_session: boto3.Session | None = ...,
     pyarrow_additional_kwargs: dict[str, Any] | None = ...,
+    key_schema: list[dict[str, str]] | None = ...,
 ) -> _ItemsListType | Iterator[_ItemsListType]: ...
 @overload
 def read_items(
@@ -244,4 +252,5 @@ def read_items(
     use_threads: bool | int = ...,
     boto3_session: boto3.Session | None = ...,
     pyarrow_additional_kwargs: dict[str, Any] | None = ...,
+    key_schema: list[dict[str, str]] | None = ...,
 ) -> pd.DataFrame | Iterator[pd.DataFrame] | _ItemsListType | Iterator[_ItemsListType]: ...

awswrangler/neptune/_neptune.py

@@ -170,7 +170,7 @@ def to_property_graph(
     ... )
     """
     # check if ~id and ~label column exist and if not throw error
-    g = gremlin.traversal().withGraph(gremlin.Graph())
+    g = gremlin.Graph().traversal()
     is_edge_df = False
     is_update_df = True
     if "~id" in df.columns:
@@ -203,6 +203,24 @@ def to_property_graph(
     return _run_gremlin_insert(client, g)
 
 
+# SPARQL 1.1 IRIREF grammar: '<' ([^<>"{}|^`\]-[#x00-#x20])* '>'
+# A cell value spliced between '<' and '>' must contain only the characters allowed
+# inside the IRIREF token. Anything else can close the token and inject arbitrary
+# SPARQL UPDATE syntax (DELETE / DROP / LOAD / ...).
+_IRIREF_INNER_RE = re.compile(r"^[^\x00-\x20<>\"{}|^`\\]*$")
+
+
+def _validate_iriref_cell(value: Any, column: str, row_index: int) -> str:
+    text = str(value)
+    if not _IRIREF_INNER_RE.match(text):
+        raise exceptions.InvalidArgumentValue(
+            f"Value in column {column!r} at row index {row_index} is not a valid IRI: "
+            f"{text!r}. Cells written by `to_rdf_graph` must conform to the SPARQL "
+            'IRIREF grammar (no whitespace, control characters, or any of <>"{}|^`\\).'
+        )
+    return text
+
+
 @_utils.check_optional_dependency(sparql, "SPARQLWrapper")
 def to_rdf_graph(
     client: NeptuneClient,
@@ -267,14 +285,18 @@ def to_rdf_graph(
     query = ""
     # Loop through items in the DF
     for i, (_, row) in enumerate(df.iterrows()):
+        subject = _validate_iriref_cell(row[subject_column], subject_column, i)
+        predicate = _validate_iriref_cell(row[predicate_column], predicate_column, i)
+        obj = _validate_iriref_cell(row[object_column], object_column, i)
         # build up a query
         if is_quads:
-            insert = f"""INSERT DATA {{ GRAPH <{row[graph_column]}> {{<{row[subject_column]}>
-                    <{str(row[predicate_column])}> <{row[object_column]}> . }} }}; """
+            graph = _validate_iriref_cell(row[graph_column], graph_column, i)
+            insert = f"""INSERT DATA {{ GRAPH <{graph}> {{<{subject}>
+                    <{predicate}> <{obj}> . }} }}; """
             query = query + insert
         else:
-            insert = f"""INSERT DATA {{ <{row[subject_column]}> <{str(row[predicate_column])}>
-                    <{row[object_column]}> . }}; """
+            insert = f"""INSERT DATA {{ <{subject}> <{predicate}>
+                    <{obj}> . }}; """
             query = query + insert
         # run the query
         if i > 0 and i % batch_size == 0:

awswrangler/redshift/_write.py

@@ -574,6 +574,7 @@ def copy(  # noqa: PLR0913
     precombine_key: str | None = None,
     use_column_names: bool = False,
     add_new_columns: bool = False,
+    pyarrow_additional_kwargs: dict[str, str] | None = None,
 ) -> None:
     """Load Pandas DataFrame as a Table on Amazon Redshift using parquet files on S3 as stage.
 
@@ -687,6 +688,9 @@ def copy(  # noqa: PLR0913
         inserted into the database columns `col1` and `col3`.
     add_new_columns
         If True, it automatically adds the new DataFrame columns into the target table.
+    pyarrow_additional_kwargs
+        Forwarded to pyarrow.
+        e.g. pyarrow_additional_kwargs={'coerce_timestamps': 'us', 'allow_truncated_timestamps': False}
 
     Examples
     --------
@@ -715,6 +719,7 @@ def copy(  # noqa: PLR0913
         s3.to_parquet(
             df=df,
             path=path,
+            pyarrow_additional_kwargs=pyarrow_additional_kwargs or {},
             index=index,
             dataset=True,
             mode="append",

tests/unit/test_moto.py

@@ -760,14 +760,3 @@ def test_extract_ctas_manifest_paths_cross_bucket_raises(moto_s3_client: "S3Clie
 
     with pytest.raises(InvalidArgumentValue, match="unexpected bucket"):
         _extract_ctas_manifest_paths(path=f"s3://bucket/{manifest_key}")
-
-
-def test_csv_pandas_mode_append(moto_s3_client: "S3Client") -> None:
-    path = "s3://bucket/test_append.csv"
-    df1 = pd.DataFrame({"col": [1, 2, 3]})
-    df2 = pd.DataFrame({"col": [4, 5, 6]})
-    wr.s3.to_csv(df=df1, path=path, index=False)
-    wr.s3.to_csv(df=df2, path=path, index=False, pandas_mode="a", header=False)
-    result = wr.s3.read_csv(path=path)
-    assert len(result) == 6
-    assert list(result["col"]) == [1, 2, 3, 4, 5, 6]