Fix(athena): DDL fixes (#4132)

* Fix(athena): DDL fixes

* PR feedback

Author: erindru

sqlglot/dialects/athena.py

@@ -20,12 +20,40 @@ def _generate_as_hive(expression: exp.Expression) -> bool:
         else:
             return expression.kind != "VIEW"  # CREATE VIEW is never Hive but CREATE SCHEMA etc is
 
-    elif isinstance(expression, exp.Alter) or isinstance(expression, exp.Drop):
-        return True  # all ALTER and DROP statements are Hive
+    # https://docs.aws.amazon.com/athena/latest/ug/ddl-reference.html
+    elif isinstance(expression, (exp.Alter, exp.Drop, exp.Describe)):
+        if isinstance(expression, exp.Drop) and expression.kind == "VIEW":
+            # DROP VIEW is Trino (I guess because CREATE VIEW is)
+            return False
+
+        # Everything else is Hive
+        return True
 
     return False
 
 
+def _location_property_sql(self: Athena.Generator, e: exp.LocationProperty):
+    # If table_type='iceberg', the LocationProperty is called 'location'
+    # Otherwise, it's called 'external_location'
+    # ref: https://docs.aws.amazon.com/athena/latest/ug/create-table-as.html
+
+    prop_name = "external_location"
+
+    if isinstance(e.parent, exp.Properties):
+        table_type_property = next(
+            (
+                p
+                for p in e.parent.expressions
+                if isinstance(p, exp.Property) and p.name == "table_type"
+            ),
+            None,
+        )
+        if table_type_property and table_type_property.text("value") == "iceberg":
+            prop_name = "location"
+
+    return f"{prop_name}={self.sql(e, 'this')}"
+
+
 class Athena(Trino):
     """
     Over the years, it looks like AWS has taken various execution engines, bolted on AWS-specific modifications and then
@@ -48,7 +76,7 @@ class Athena(Trino):
     Trino:
       - Uses double quotes to quote identifiers
       - Used for DDL operations that involve SELECT queries, eg:
-        - CREATE VIEW
+        - CREATE VIEW / DROP VIEW
         - CREATE TABLE... AS SELECT
       - Used for DML operations
         - SELECT, INSERT, UPDATE, DELETE, MERGE
@@ -79,27 +107,40 @@ class Parser(Trino.Parser):
             TokenType.USING: lambda self: self._parse_as_command(self._prev),
         }
 
+    class _HiveGenerator(Hive.Generator):
+        def alter_sql(self, expression: exp.Alter) -> str:
+            # package any ALTER TABLE ADD actions into a Schema object
+            # so it gets generated as `ALTER TABLE .. ADD COLUMNS(...)`
+            # instead of `ALTER TABLE ... ADD COLUMN` which is invalid syntax on Athena
+            if isinstance(expression, exp.Alter) and expression.kind == "TABLE":
+                if expression.actions and isinstance(expression.actions[0], exp.ColumnDef):
+                    new_actions = exp.Schema(expressions=expression.actions)
+                    expression.set("actions", [new_actions])
+
+            return super().alter_sql(expression)
+
     class Generator(Trino.Generator):
         """
         Generate queries for the Athena Trino execution engine
         """
 
-        TYPE_MAPPING = {
-            **Trino.Generator.TYPE_MAPPING,
-            exp.DataType.Type.TEXT: "STRING",
+        PROPERTIES_LOCATION = {
+            **Trino.Generator.PROPERTIES_LOCATION,
+            exp.LocationProperty: exp.Properties.Location.POST_WITH,
         }
 
         TRANSFORMS = {
             **Trino.Generator.TRANSFORMS,
-            exp.FileFormatProperty: lambda self, e: f"'FORMAT'={self.sql(e, 'this')}",
+            exp.FileFormatProperty: lambda self, e: f"format={self.sql(e, 'this')}",
+            exp.LocationProperty: _location_property_sql,
         }
 
         def __init__(self, *args, **kwargs):
             super().__init__(*args, **kwargs)
 
             hive_kwargs = {**kwargs, "dialect": "hive"}
 
-            self._hive_generator = Hive.Generator(*args, **hive_kwargs)
+            self._hive_generator = Athena._HiveGenerator(*args, **hive_kwargs)
 
         def generate(self, expression: exp.Expression, copy: bool = True) -> str:
             if _generate_as_hive(expression):

sqlglot/expressions.py

@@ -4398,6 +4398,15 @@ class Alter(Expression):
         "not_valid": False,
     }
 
+    @property
+    def kind(self) -> t.Optional[str]:
+        kind = self.args.get("kind")
+        return kind and kind.upper()
+
+    @property
+    def actions(self) -> t.List[Expression]:
+        return self.args.get("actions") or []
+
 
 class AddConstraint(Expression):
     arg_types = {"expressions": True}

tests/dialects/test_athena.py

@@ -1,3 +1,4 @@
+from sqlglot import exp
 from tests.dialects.test_dialect import Validator
 
 
@@ -68,6 +69,23 @@ def test_ddl(self):
             "CREATE TABLE foo AS WITH foo AS (SELECT a, b FROM bar) SELECT * FROM foo"
         )
 
+        # ALTER TABLE ADD COLUMN not supported, it needs to be generated as ALTER TABLE ADD COLUMNS
+        self.validate_identity(
+            "ALTER TABLE `foo`.`bar` ADD COLUMN `end_ts` BIGINT",
+            write_sql="ALTER TABLE `foo`.`bar` ADD COLUMNS (`end_ts` BIGINT)",
+        )
+
+    def test_dml(self):
+        self.validate_all(
+            "SELECT CAST(ds AS VARCHAR) AS ds FROM (VALUES ('2022-01-01')) AS t(ds)",
+            read={"": "SELECT CAST(ds AS STRING) AS ds FROM (VALUES ('2022-01-01')) AS t(ds)"},
+            write={
+                "hive": "SELECT CAST(ds AS STRING) AS ds FROM (VALUES ('2022-01-01')) AS t(ds)",
+                "trino": "SELECT CAST(ds AS VARCHAR) AS ds FROM (VALUES ('2022-01-01')) AS t(ds)",
+                "athena": "SELECT CAST(ds AS VARCHAR) AS ds FROM (VALUES ('2022-01-01')) AS t(ds)",
+            },
+        )
+
     def test_ddl_quoting(self):
         self.validate_identity("CREATE SCHEMA `foo`")
         self.validate_identity("CREATE SCHEMA foo")
@@ -111,6 +129,10 @@ def test_ddl_quoting(self):
             'CREATE VIEW `foo` AS SELECT "id" FROM `tbl`',
             write_sql='CREATE VIEW "foo" AS SELECT "id" FROM "tbl"',
         )
+        self.validate_identity(
+            "DROP VIEW IF EXISTS `foo`.`bar`",
+            write_sql='DROP VIEW IF EXISTS "foo"."bar"',
+        )
 
         self.validate_identity(
             'ALTER TABLE "foo" ADD COLUMNS ("id" STRING)',
@@ -128,6 +150,8 @@ def test_ddl_quoting(self):
             write_sql='CREATE TABLE "foo" AS WITH "foo" AS (SELECT "a", "b" FROM "bar") SELECT * FROM "foo"',
         )
 
+        self.validate_identity("DESCRIBE foo.bar", write_sql="DESCRIBE `foo`.`bar`", identify=True)
+
     def test_dml_quoting(self):
         self.validate_identity("SELECT a AS foo FROM tbl")
         self.validate_identity('SELECT "a" AS "foo" FROM "tbl"')
@@ -167,3 +191,39 @@ def test_dml_quoting(self):
             write_sql='WITH "foo" AS (SELECT "a", "b" FROM "bar") SELECT * FROM "foo"',
             identify=True,
         )
+
+    def test_ctas(self):
+        # Hive tables use 'external_location' to specify the table location, Iceberg tables use 'location' to specify the table location
+        # The 'table_type' property is used to determine if it's a Hive or an Iceberg table
+        # ref: https://docs.aws.amazon.com/athena/latest/ug/create-table-as.html#ctas-table-properties
+        ctas_hive = exp.Create(
+            this=exp.to_table("foo.bar"),
+            kind="TABLE",
+            properties=exp.Properties(
+                expressions=[
+                    exp.FileFormatProperty(this=exp.Literal.string("parquet")),
+                    exp.LocationProperty(this=exp.Literal.string("s3://foo")),
+                ]
+            ),
+            expression=exp.select("1"),
+        )
+        self.assertEqual(
+            ctas_hive.sql(dialect=self.dialect, identify=True),
+            "CREATE TABLE \"foo\".\"bar\" WITH (format='parquet', external_location='s3://foo') AS SELECT 1",
+        )
+
+        ctas_iceberg = exp.Create(
+            this=exp.to_table("foo.bar"),
+            kind="TABLE",
+            properties=exp.Properties(
+                expressions=[
+                    exp.Property(this=exp.var("table_type"), value=exp.Literal.string("iceberg")),
+                    exp.LocationProperty(this=exp.Literal.string("s3://foo")),
+                ]
+            ),
+            expression=exp.select("1"),
+        )
+        self.assertEqual(
+            ctas_iceberg.sql(dialect=self.dialect, identify=True),
+            "CREATE TABLE \"foo\".\"bar\" WITH (table_type='iceberg', location='s3://foo') AS SELECT 1",
+        )