Repository: tobymao/sqlglot

Issue:
Support for Databricks Unity Catalog Python UDF's
Unity Catalog allows to do somethings like this:

```python

import sqlglot

print(
    sqlglot.parse_one(
        """ CREATE OR REPLACE FUNCTION my_catalog.my_schema.mixed_process(data STRING)
RETURNS STRING
LANGUAGE PYTHON
ENVIRONMENT (
  dependencies = '["simplejson==3.19.3", "/Volumes/my_catalog/my_schema/my_volume/packages/custom_package-1.0.0.whl", "https://my-bucket.s3.amazonaws.com/packages/special_package-2.0.0.whl?Expires=2043167927&Signature=abcd"]'
  environment_version = 'None'
)
AS $$
import simplejson as json
import custom_package
return json.dumps(custom_package.process(data))
$$""",
        read="databricks",
    ).sql("databricks")
)


```

Sqlglot cannot parse this currently. Of course I do not expect it to parse Python, that would probably just be a string

See also [docs](https://docs.databricks.com/aws/en/udf/python)

Discussion:

Solution patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index f61ab9cb9d..8cd66f59b7 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2869,6 +2869,10 @@ class LanguageProperty(Property):
     arg_types = {"this": True}
 
 
+class EnviromentProperty(Property):
+    arg_types = {"expressions": True}
+
+
 # spark ddl
 class ClusteredByProperty(Property):
     arg_types = {"expressions": True, "sorted_by": False, "buckets": True}
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 11e6a45b57..cccb92ae16 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -139,6 +139,7 @@ class Generator(metaclass=_Generator):
         exp.DynamicProperty: lambda *_: "DYNAMIC",
         exp.EmptyProperty: lambda *_: "EMPTY",
         exp.EncodeColumnConstraint: lambda self, e: f"ENCODE {self.sql(e, 'this')}",
+        exp.EnviromentProperty: lambda self, e: f"ENVIRONMENT ({self.expressions(e, flat=True)})",
         exp.EphemeralColumnConstraint: lambda self,
         e: f"EPHEMERAL{(' ' + self.sql(e, 'this')) if e.this else ''}",
         exp.ExcludeColumnConstraint: lambda self, e: f"EXCLUDE {self.sql(e, 'this').lstrip()}",
@@ -543,6 +544,7 @@ class Generator(metaclass=_Generator):
         exp.EmptyProperty: exp.Properties.Location.POST_SCHEMA,
         exp.EncodeProperty: exp.Properties.Location.POST_EXPRESSION,
         exp.EngineProperty: exp.Properties.Location.POST_SCHEMA,
+        exp.EnviromentProperty: exp.Properties.Location.POST_SCHEMA,
         exp.ExecuteAsProperty: exp.Properties.Location.POST_SCHEMA,
         exp.ExternalProperty: exp.Properties.Location.POST_CREATE,
         exp.FallbackProperty: exp.Properties.Location.POST_NAME,
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index bfa843a215..ef4f67b4f1 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -963,6 +963,9 @@ class Parser(metaclass=_Parser):
         "DISTSTYLE": lambda self: self._parse_property_assignment(exp.DistStyleProperty),
         "EMPTY": lambda self: self.expression(exp.EmptyProperty),
         "ENGINE": lambda self: self._parse_property_assignment(exp.EngineProperty),
+        "ENVIRONMENT": lambda self: self.expression(
+            exp.EnviromentProperty, expressions=self._parse_wrapped_csv(self._parse_assignment)
+        ),
         "EXECUTE": lambda self: self._parse_property_assignment(exp.ExecuteAsProperty),
         "EXTERNAL": lambda self: self.expression(exp.ExternalProperty),
         "FALLBACK": lambda self, **kwargs: self._parse_fallback(**kwargs),


Test patch:
diff --git a/tests/dialects/test_databricks.py b/tests/dialects/test_databricks.py
index b0b921d7f3..723767c0cd 100644
--- a/tests/dialects/test_databricks.py
+++ b/tests/dialects/test_databricks.py
@@ -338,3 +338,8 @@ def test_analyze(self):
         self.validate_identity(
             "ANALYZE TABLE ctlg.db.tbl PARTITION(foo = 'foo', bar = 'bar') COMPUTE STATISTICS NOSCAN"
         )
+
+    def test_udf_environment_property(self):
+        self.validate_identity(
+            """CREATE FUNCTION a() ENVIRONMENT (dependencies = '["foo1==1", "foo2==2"]', environment_version = 'None')"""
+        )