Feat(oracle): add support for JSON_ARRAYAGG (#2189)

tobymao · Sep 11, 2023 · 0c536bd · 0c536bd
1 parent d192515
commit 0c536bd
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 4 deletions.
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -4494,6 +4494,18 @@ class JSONObject(Func):
     }
 
 
+# https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/JSON_ARRAYAGG.html
+class JSONArrayAgg(Func):
+    arg_types = {
+        "this": True,
+        "format_json": False,
+        "order": False,
+        "null_handling": False,
+        "return_type": False,
+        "strict": False,
+    }
+
+
 class OpenJSONColumnDef(Expression):
     arg_types = {"this": True, "kind": True, "path": False, "as_json": False}
 

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -2042,6 +2042,21 @@ def jsonobject_sql(self, expression: exp.JSONObject) -> str:
             suffix=f"{null_handling}{unique_keys}{return_type}{format_json}{encoding})",
         )
 
+    def jsonarrayagg_sql(self, expression: exp.JSONArrayAgg) -> str:
+        this = self.sql(expression, "this")
+        format_json = " FORMAT JSON" if expression.args.get("format_json") else ""
+        order = self.sql(expression, "order")
+        null_handling = expression.args.get("null_handling")
+        null_handling = f" {null_handling}" if null_handling else ""
+        return_type = self.sql(expression, "return_type")
+        return_type = f" RETURNING {return_type}" if return_type else ""
+        strict = " STRICT" if expression.args.get("strict") else ""
+        return self.func(
+            "JSON_ARRAYAGG",
+            this,
+            suffix=f"{format_json}{order}{null_handling}{return_type}{strict})",
+        )
+
     def openjsoncolumndef_sql(self, expression: exp.OpenJSONColumnDef) -> str:
         this = self.sql(expression, "this")
         kind = self.sql(expression, "kind")

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -769,6 +769,7 @@ class Parser(metaclass=_Parser):
         "CONVERT": lambda self: self._parse_convert(self.STRICT_CAST),
         "DECODE": lambda self: self._parse_decode(),
         "EXTRACT": lambda self: self._parse_extract(),
+        "JSON_ARRAYAGG": lambda self: self._parse_json_arrayagg(),
         "JSON_OBJECT": lambda self: self._parse_json_object(),
         "LOG": lambda self: self._parse_logarithm(),
         "MATCH": lambda self: self._parse_match_against(),
@@ -4183,16 +4184,21 @@ def _parse_json_key_value(self) -> t.Optional[exp.JSONKeyValue]:
             return None
         return self.expression(exp.JSONKeyValue, this=key, expression=value)
 
-    def _parse_json_object(self) -> exp.JSONObject:
-        star = self._parse_star()
-        expressions = [star] if star else self._parse_csv(self._parse_json_key_value)
-
+    def _parse_null_handling(self) -> t.Optional[str]:
+        # Parses Oracle's {NULL|ABSENT} ON NULL syntax
         null_handling = None
         if self._match_text_seq("NULL", "ON", "NULL"):
             null_handling = "NULL ON NULL"
         elif self._match_text_seq("ABSENT", "ON", "NULL"):
             null_handling = "ABSENT ON NULL"
 
+        return null_handling
+
+    def _parse_json_object(self) -> exp.JSONObject:
+        star = self._parse_star()
+        expressions = [star] if star else self._parse_csv(self._parse_json_key_value)
+        null_handling = self._parse_null_handling()
+
         unique_keys = None
         if self._match_text_seq("WITH", "UNIQUE"):
             unique_keys = True
@@ -4215,6 +4221,17 @@ def _parse_json_object(self) -> exp.JSONObject:
             encoding=encoding,
         )
 
+    def _parse_json_arrayagg(self) -> exp.JSONArrayAgg:
+        return self.expression(
+            exp.JSONArrayAgg,
+            this=self._parse_bitwise(),
+            format_json=self._match_text_seq("FORMAT", "JSON"),
+            order=self._parse_order(),
+            null_handling=self._parse_null_handling(),
+            return_type=self._match_text_seq("RETURNING") and self._parse_type(),
+            strict=self._match_text_seq("STRICT"),
+        )
+
     def _parse_logarithm(self) -> exp.Func:
         # Default argument order is base, expression
         args = self._parse_csv(self._parse_range)

diff --git a/tests/dialects/test_oracle.py b/tests/dialects/test_oracle.py
@@ -25,6 +25,9 @@ def test_oracle(self):
         self.validate_identity("SELECT * FROM table_name@dblink_name.database_link_domain")
         self.validate_identity("SELECT * FROM table_name SAMPLE (25) s")
         self.validate_identity("SELECT * FROM V$SESSION")
+        self.validate_identity(
+            "SELECT JSON_ARRAYAGG(FOO() FORMAT JSON ORDER BY bar NULL ON NULL RETURNING CLOB STRICT)"
+        )
         self.validate_identity(
             "SELECT COUNT(1) INTO V_Temp FROM TABLE(CAST(somelist AS data_list)) WHERE col LIKE '%contact'"
         )