fix(s3): fix import order and remove unused cast import

Author: bujjibabukatta

awswrangler/athena/_write_iceberg.py

@@ -24,6 +24,21 @@
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
+def _escape_athena_string_literal(value: Any) -> str:
+    # Used for caller-supplied values spliced inside SQL string literals, e.g.
+    # COMMENT '<value>' or TBLPROPERTIES ('<key>'='<value>').
+    #
+    # The wrapping single quotes around the splice site are *delimiters* — they tell
+    # Athena "this is a string literal" and are part of the DDL grammar, not escaping.
+    # A naive splice f"'{value}'" lets a value containing ' close the literal and
+    # append arbitrary DDL (e.g. LOCATION '...').
+    #
+    # Athena/Trino string literals have exactly one escape mechanism: a single quote
+    # inside the literal must be doubled. Pre-doubling caller quotes here means the
+    # whole value parses as one literal regardless of what it contains.
+    return str(value).replace("'", "''")
+
+
 def _create_iceberg_table(
     df: pd.DataFrame,
     database: str,
@@ -50,13 +65,19 @@ def _create_iceberg_table(
         [
             f"{k} {v}"
             if (columns_comments is None or columns_comments.get(k) is None)
-            else f"{k} {v} COMMENT '{columns_comments[k]}'"
+            else f"{k} {v} COMMENT '{_escape_athena_string_literal(columns_comments[k])}'"
             for k, v in columns_types.items()
         ]
     )
     partition_cols_str: str = f"PARTITIONED BY ({', '.join([col for col in partition_cols])})" if partition_cols else ""
     table_properties_str: str = (
-        ", " + ", ".join([f"'{key}'='{value}'" for key, value in additional_table_properties.items()])
+        ", "
+        + ", ".join(
+            [
+                f"'{_escape_athena_string_literal(key)}'='{_escape_athena_string_literal(value)}'"
+                for key, value in additional_table_properties.items()
+            ]
+        )
         if additional_table_properties
         else ""
     )

awswrangler/neptune/_gremlin_init.py

@@ -8,9 +8,15 @@
     from gremlin_python.process.anonymous_traversal import traversal
     from gremlin_python.process.graph_traversal import GraphTraversalSource, __
     from gremlin_python.process.translator import Translator
-    from gremlin_python.process.traversal import Cardinality, T
+    from gremlin_python.process.traversal import Cardinality, T, TraversalStrategies
     from gremlin_python.structure.graph import Edge, Graph, Path, Property, Vertex, VertexProperty
 
+    def local_traversal_source() -> "GraphTraversalSource":
+        # gremlinpython 3.8 removed Graph().traversal() and traversal().withGraph(Graph()).
+        # We only need a traversal source to build bytecode that the Translator turns into
+        # a query string — no remote connection or strategies are needed.
+        return GraphTraversalSource(Graph(), TraversalStrategies())
+
     __all__ = [
         "__",
         "Cardinality",
@@ -22,6 +28,8 @@
         "Property",
         "T",
         "Translator",
+        "TraversalStrategies",
+        "local_traversal_source",
         "traversal",
         "Vertex",
         "VertexProperty",

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.Graph().traversal()
+    g = gremlin.local_traversal_source()
     is_edge_df = False
     is_update_df = True
     if "~id" in df.columns:
@@ -198,7 +198,7 @@ def to_property_graph(
         if index > 0 and index + 1 % batch_size == 0:
             res = _run_gremlin_insert(client, g)
             if res:
-                g = gremlin.Graph().traversal()
+                g = gremlin.local_traversal_source()
 
     return _run_gremlin_insert(client, g)
 

awswrangler/s3/_write_text.py

@@ -3,10 +3,10 @@
 from __future__ import annotations
 
 import csv
-import logging
 import io
+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

tests/unit/test_moto.py

@@ -804,6 +804,48 @@ def mock_make_api_call(self, operation_name, kwarg):
         assert describe_table_calls == 1
 
 
+def test_create_iceberg_table_escapes_single_quotes_in_columns_comments() -> None:
+    # Single quotes in caller-supplied columns_comments / additional_table_properties
+    # values must be doubled so they cannot terminate the surrounding 'literal' and
+    # change the structure of the generated DDL.
+    from awswrangler.athena import _write_iceberg
+
+    captured: list[str] = []
+
+    def fake_start(*, sql: str, **_) -> str:
+        captured.append(sql)
+        return "qid"
+
+    df = pd.DataFrame({"id": pd.Series(dtype="int64"), "user_name": pd.Series(dtype="string")})
+    wg_config = mock.MagicMock()
+    wg_config.enforce_workgroup_location = False
+
+    with mock.patch.object(_write_iceberg, "_start_query_execution", side_effect=fake_start), mock.patch.object(
+        _write_iceberg, "wait_query"
+    ):
+        _write_iceberg._create_iceberg_table(
+            df=df,
+            database="db",
+            table="t",
+            path="s3://intended/output/",
+            wg_config=wg_config,
+            partition_cols=None,
+            additional_table_properties={"prop": "val') LOCATION 's3://other/' --"},
+            index=False,
+            boto3_session=mock.MagicMock(),
+            columns_comments={"user_name": "') LOCATION 's3://other/' TBLPROPERTIES ('x'='y"},
+        )
+
+    sql = captured[0]
+    # Quotes were doubled in both splices, so unescaped caller content stays inside the
+    # COMMENT / TBLPROPERTIES string literals and does not open a new DDL clause.
+    assert "COMMENT ''') LOCATION ''s3://other/'' TBLPROPERTIES (''x''=''y'" in sql
+    assert "'prop'='val'') LOCATION ''s3://other/'' --'" in sql
+    # The intended LOCATION (un-doubled quotes) is the only top-level clause.
+    assert "LOCATION 's3://intended/output/'" in sql
+    assert "LOCATION 's3://other/'" not in sql
+
+
 def test_csv_pandas_mode_append(moto_s3_client: "S3Client") -> None:
     path = "s3://bucket/test_append.csv"
     df1 = pd.DataFrame({"col": [1, 2, 3]})