Repository: tobymao/sqlglot

Issue:
Support User Defined Functions on Athena Dialect
It looks like sqlglot is not able to parse [AWS Athena's user defined functions syntax](https://docs.aws.amazon.com/athena/latest/ug/querying-udf.html):

```py
from sqlglot import parse
from sqlglot.dialects import Trino

parse("""
    USING EXTERNAL FUNCTION some_function(input VARBINARY)
        RETURNS VARCHAR
        LAMBDA 'some-name'
    SELECT
        some_function(1)
""", dialect=Trino)
```

Exception:

```
sqlglot.errors.ParseError: Invalid expression / Unexpected token. Line 2, Col: 9.
  
    USING EXTERNAL FUNCTION some_function(input VARBINARY)
        RETURNS VARCHAR
        LAMBDA 'some-name'
```

We are using `Trino` dialect since sqlglot does not have a dedicated one for Athena, as far as I understand, but Athena is based off Trino, so this dialect works otherwise perfectly for our codebase :slightly_smiling_face: 

Am I missing something? Does it need a dedicated dialect for Athena?



Discussion:
Hey @sbrandtb, thanks for the report.

AWS Athena is not supported right now, so this behavior is expected as it's not valid _Trino_ syntax, IIUC. Closing this as not planned for now since we don't plan to add Athena in the near future.
@georgesittas Given that Athena is basically Trino, what speaks against deriving a new Athena dialect, starting with UDF syntax? I am willing to help.
We actually just discussed this internally and decided to implement the `Athena` dialect and have this be parsed into a `Command` for now. Should be a good enough compromise to start with. Appreciate your willingness to help here.
So should have a fix for this particular issue soon. Feel free to contribute improvements to Athena if you find more discrepancies in the future.

Resolving patch:
diff --git a/sqlglot/dialects/__init__.py b/sqlglot/dialects/__init__.py
--- a/sqlglot/dialects/__init__.py
+++ b/sqlglot/dialects/__init__.py
@@ -61,6 +61,7 @@ class Generator(Generator):
 ----
 """
 
+from sqlglot.dialects.athena import Athena
 from sqlglot.dialects.bigquery import BigQuery
 from sqlglot.dialects.clickhouse import ClickHouse
 from sqlglot.dialects.databricks import Databricks
diff --git a/sqlglot/dialects/athena.py b/sqlglot/dialects/athena.py
new file mode 100644
--- /dev/null
+++ b/sqlglot/dialects/athena.py
@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from sqlglot.dialects.trino import Trino
+from sqlglot.tokens import TokenType
+
+
+class Athena(Trino):
+    class Parser(Trino.Parser):
+        STATEMENT_PARSERS = {
+            **Trino.Parser.STATEMENT_PARSERS,
+            TokenType.USING: lambda self: self._parse_as_command(self._prev),
+        }
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -31,6 +31,7 @@ class Dialects(str, Enum):
 
     DIALECT = ""
 
+    ATHENA = "athena"
     BIGQUERY = "bigquery"
     CLICKHOUSE = "clickhouse"
     DATABRICKS = "databricks"


Test patch:
diff --git a/tests/dialects/test_athena.py b/tests/dialects/test_athena.py
new file mode 100644
--- /dev/null
+++ b/tests/dialects/test_athena.py
@@ -0,0 +1,16 @@
+from tests.dialects.test_dialect import Validator
+
+
+class TestAthena(Validator):
+    dialect = "athena"
+    maxDiff = None
+
+    def test_athena(self):
+        self.validate_identity(
+            """USING EXTERNAL FUNCTION some_function(input VARBINARY)
+            RETURNS VARCHAR
+                LAMBDA 'some-name'
+            SELECT
+            some_function(1)""",
+            check_command_warning=True,
+        )