fix(athena): only replace commas instead of all special characters in column names (#11138)

Author: cpcloud

ibis/backends/athena/tests/test_client.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+import ibis
 from ibis.backends.tests.errors import PyAthenaOperationalError
 from ibis.util import gen_name
 
@@ -28,3 +29,14 @@ def test_create_and_drop_database(con):
     # drop it again (should fail)
     with pytest.raises(PyAthenaOperationalError):
         con.drop_database(name)
+
+
+def test_column_name_with_slash(con):
+    table = ibis.memtable({"inventarnr_/_mde_dummy": [1, 2, 3]})
+
+    renamed_table = con.execute(
+        table.select("inventarnr_/_mde_dummy")
+        .rename({"dummy": "inventarnr_/_mde_dummy"})
+        .dummy
+    )
+    assert set(renamed_table.values) == {1, 2, 3}

ibis/backends/sql/compilers/athena.py

@@ -1,16 +1,12 @@
 from __future__ import annotations
 
-import re
-
 from sqlglot.dialects import Athena
 
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
 from ibis.backends.sql.compilers.trino import TrinoCompiler
 from ibis.backends.sql.datatypes import AthenaType
 
-_NAME_REGEX = re.compile(r'[^!"$()*,./;?@[\\\]^`{}~\n]+')
-
 
 class AthenaCompiler(TrinoCompiler):
     __slots__ = ()
@@ -35,7 +31,7 @@ class AthenaCompiler(TrinoCompiler):
 
     @staticmethod
     def _gen_valid_name(name: str) -> str:
-        return "_".join(map(str.strip, _NAME_REGEX.findall(name))) or "tmp"
+        return name.replace(",", ";")
 
     def visit_Cast(self, op, *, arg, to):
         from_ = op.arg.dtype

ibis/backends/tests/snapshots/test_sql/test_isin_bug/athena/out.sql

@@ -5,5 +5,5 @@ SELECT
     FROM "t" AS "t0"
     WHERE
       "t0"."x" > 2
-  ) AS "InSubquery_x"
+  ) AS "InSubquery(x)"
 FROM "t" AS "t0"

ibis/backends/tests/snapshots/test_sql/test_to_sql_default_backend/athena/to_sql.sql

@@ -1,5 +1,5 @@
 SELECT
-  COUNT(*) AS "CountStar"
+  COUNT(*) AS "CountStar()"
 FROM (
   SELECT
     *

ibis/backends/tests/test_temporal.py

@@ -1765,13 +1765,8 @@ def test_timestamp_extract_milliseconds_with_big_value(con):
     assert result == 333
 
 
-@pytest.mark.notimpl(
-    ["oracle"],
-    raises=OracleDatabaseError,
-    reason="ORA-00932",
-)
+@pytest.mark.notimpl(["oracle"], raises=OracleDatabaseError, reason="ORA-00932")
 @pytest.mark.notimpl(["exasol"], raises=ExaQueryError)
-@pytest.mark.notimpl(["athena"], raises=PyAthenaOperationalError)
 def test_integer_cast_to_timestamp_column(backend, alltypes, df):
     expr = alltypes.int_col.cast("timestamp")
     expected = pd.to_datetime(df.int_col, unit="s").rename(expr.get_name())
@@ -1781,7 +1776,6 @@ def test_integer_cast_to_timestamp_column(backend, alltypes, df):
 
 @pytest.mark.notimpl(["exasol"], raises=ExaQueryError)
 @pytest.mark.notimpl(["oracle"], raises=OracleDatabaseError)
-@pytest.mark.notimpl(["athena"], raises=PyAthenaOperationalError)
 def test_integer_cast_to_timestamp_scalar(alltypes, df):
     expr = alltypes.int_col.min().cast("timestamp")
     result = expr.execute()