Repository: tobymao/sqlglot

Issue:
Support parsing UDF definition in Databricks SQL
**Is your feature request related to a problem? Please describe.**
context: UDF definition on Databricks: https://docs.databricks.com/sql/language-manual/sql-ref-syntax-ddl-create-sql-function.html#create-python-functions

Tried to use sqlglot to parse a creation sql, but no luck:
```
from sqlglot import parse_one, expressions
def_func = """CREATE OR REPLACE FUNCTION my_catalog.my_schema.greet(s STRING)
RETURNS STRING
LANGUAGE SQL
CONTAINS SQL
RETURN s"""

parse_one(def_func)
```
`parse_one` gives a parse error:
```
ParseError: Invalid expression / Unexpected token. Line 5, Col: 6.
  CE FUNCTION my_catalog.my_schema.greet(s STRING)
RETURNS STRING
LANGUAGE SQL
CONTAINS SQL
[4mRETURN[0m s
```

**Describe the solution you'd like**
I hope sqlglot will be able to parse the FUNCTION creation sql, recognize it's about FUNCTION, and find out elements such as catalog, database, function name, etc.


Discussion:
not a priority for us but happy to accept a pr

Resolving patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2252,6 +2252,10 @@ class SortKeyProperty(Property):
     arg_types = {"this": True, "compound": False}
 
 
+class SqlReadWriteProperty(Property):
+    arg_types = {"this": True}
+
+
 class SqlSecurityProperty(Property):
     arg_types = {"definer": True}
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -97,6 +97,7 @@ class Generator:
         exp.SampleProperty: lambda self, e: f"SAMPLE BY {self.sql(e, 'this')}",
         exp.SetProperty: lambda self, e: f"{'MULTI' if e.args.get('multi') else ''}SET",
         exp.SettingsProperty: lambda self, e: f"SETTINGS{self.seg('')}{(self.expressions(e))}",
+        exp.SqlReadWriteProperty: lambda self, e: e.name,
         exp.SqlSecurityProperty: lambda self, e: f"SQL SECURITY {'DEFINER' if e.args.get('definer') else 'INVOKER'}",
         exp.StabilityProperty: lambda self, e: e.name,
         exp.TemporaryProperty: lambda self, e: f"TEMPORARY",
@@ -334,6 +335,7 @@ class Generator:
         exp.SettingsProperty: exp.Properties.Location.POST_SCHEMA,
         exp.SetProperty: exp.Properties.Location.POST_CREATE,
         exp.SortKeyProperty: exp.Properties.Location.POST_SCHEMA,
+        exp.SqlReadWriteProperty: exp.Properties.Location.POST_SCHEMA,
         exp.SqlSecurityProperty: exp.Properties.Location.POST_CREATE,
         exp.StabilityProperty: exp.Properties.Location.POST_SCHEMA,
         exp.TemporaryProperty: exp.Properties.Location.POST_CREATE,
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -681,6 +681,7 @@ class Parser(metaclass=_Parser):
             exp.CollateProperty, **kwargs
         ),
         "COMMENT": lambda self: self._parse_property_assignment(exp.SchemaCommentProperty),
+        "CONTAINS": lambda self: self._parse_contains_property(),
         "COPY": lambda self: self._parse_copy_property(),
         "DATABLOCKSIZE": lambda self, **kwargs: self._parse_datablocksize(**kwargs),
         "DEFINER": lambda self: self._parse_definer(),
@@ -711,6 +712,7 @@ class Parser(metaclass=_Parser):
         "LOG": lambda self, **kwargs: self._parse_log(**kwargs),
         "MATERIALIZED": lambda self: self.expression(exp.MaterializedProperty),
         "MERGEBLOCKRATIO": lambda self, **kwargs: self._parse_mergeblockratio(**kwargs),
+        "MODIFIES": lambda self: self._parse_modifies_property(),
         "MULTISET": lambda self: self.expression(exp.SetProperty, multi=True),
         "NO": lambda self: self._parse_no_property(),
         "ON": lambda self: self._parse_on_property(),
@@ -722,6 +724,7 @@ class Parser(metaclass=_Parser):
         "PARTITIONED_BY": lambda self: self._parse_partitioned_by(),
         "PRIMARY KEY": lambda self: self._parse_primary_key(in_props=True),
         "RANGE": lambda self: self._parse_dict_range(this="RANGE"),
+        "READS": lambda self: self._parse_reads_property(),
         "REMOTE": lambda self: self._parse_remote_with_connection(),
         "RETURNS": lambda self: self._parse_returns(),
         "ROW": lambda self: self._parse_row(),
@@ -1852,9 +1855,21 @@ def _parse_withdata(self, no: bool = False) -> exp.WithDataProperty:
 
         return self.expression(exp.WithDataProperty, no=no, statistics=statistics)
 
-    def _parse_no_property(self) -> t.Optional[exp.NoPrimaryIndexProperty]:
+    def _parse_contains_property(self) -> t.Optional[exp.SqlReadWriteProperty]:
+        if self._match_text_seq("SQL"):
+            return self.expression(exp.SqlReadWriteProperty, this="CONTAINS SQL")
+        return None
+
+    def _parse_modifies_property(self) -> t.Optional[exp.SqlReadWriteProperty]:
+        if self._match_text_seq("SQL", "DATA"):
+            return self.expression(exp.SqlReadWriteProperty, this="MODIFIES SQL DATA")
+        return None
+
+    def _parse_no_property(self) -> t.Optional[exp.Expression]:
         if self._match_text_seq("PRIMARY", "INDEX"):
             return exp.NoPrimaryIndexProperty()
+        if self._match_text_seq("SQL"):
+            return self.expression(exp.SqlReadWriteProperty, this="NO SQL")
         return None
 
     def _parse_on_property(self) -> t.Optional[exp.Expression]:
@@ -1864,6 +1879,11 @@ def _parse_on_property(self) -> t.Optional[exp.Expression]:
             return exp.OnCommitProperty(delete=True)
         return self.expression(exp.OnProperty, this=self._parse_schema(self._parse_id_var()))
 
+    def _parse_reads_property(self) -> t.Optional[exp.SqlReadWriteProperty]:
+        if self._match_text_seq("SQL", "DATA"):
+            return self.expression(exp.SqlReadWriteProperty, this="READS SQL DATA")
+        return None
+
     def _parse_distkey(self) -> exp.DistKeyProperty:
         return self.expression(exp.DistKeyProperty, this=self._parse_wrapped(self._parse_id_var))
 


Test patch:
diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql
--- a/tests/fixtures/identity.sql
+++ b/tests/fixtures/identity.sql
@@ -624,6 +624,9 @@ CREATE FUNCTION a(b INT, c VARCHAR) AS 'SELECT 1'
 CREATE FUNCTION a() LANGUAGE sql
 CREATE FUNCTION a() LANGUAGE sql RETURNS INT
 CREATE FUNCTION a.b(x INT) RETURNS INT AS RETURN x + 1
+CREATE FUNCTION a.b(x TEXT) RETURNS TEXT CONTAINS SQL AS RETURN x
+CREATE FUNCTION a.b(x TEXT) RETURNS TEXT LANGUAGE SQL MODIFIES SQL DATA AS RETURN x
+CREATE FUNCTION a.b(x TEXT) LANGUAGE SQL READS SQL DATA RETURNS TEXT AS RETURN x
 CREATE FUNCTION a.b.c()
 CREATE INDEX abc ON t(a)
 CREATE INDEX "abc" ON t(a)