fix(postgres): Fix exp.WidthBucket required args (#6621)

Author: VaggelisD

sqlglot/dialects/postgres.py

@@ -447,6 +447,11 @@ class Parser(parser.Parser):
             "LEVENSHTEIN_LESS_EQUAL": _build_levenshtein_less_equal,
             "JSON_OBJECT_AGG": lambda args: exp.JSONObjectAgg(expressions=args),
             "JSONB_OBJECT_AGG": exp.JSONBObjectAgg.from_arg_list,
+            "WIDTH_BUCKET": lambda args: exp.WidthBucket(
+                this=seq_get(args, 0), threshold=seq_get(args, 1)
+            )
+            if len(args) == 2
+            else exp.WidthBucket.from_arg_list(args),
         }
 
         NO_PAREN_FUNCTIONS = {

sqlglot/expressions.py

@@ -8279,7 +8279,13 @@ class Skewness(AggFunc):
 
 
 class WidthBucket(Func):
-    arg_types = {"this": True, "min_value": True, "max_value": True, "num_buckets": True}
+    arg_types = {
+        "this": True,
+        "min_value": False,
+        "max_value": False,
+        "num_buckets": False,
+        "threshold": False,
+    }
 
 
 class CovarSamp(Binary, AggFunc):

tests/dialects/test_postgres.py

@@ -1016,6 +1016,12 @@ def test_postgres(self):
             "SELECT SLOPE(CAST('(4,4)' AS POINT), CAST('(0,0)' AS POINT))",
         )
 
+        width_bucket = self.validate_identity("WIDTH_BUCKET(10, ARRAY[5, 15])")
+        self.assertIsNotNone(width_bucket.args.get("threshold"))
+
+        width_bucket = self.validate_identity("WIDTH_BUCKET(10, 5, 15, 25)")
+        self.assertIsNone(width_bucket.args.get("threshold"))
+
     def test_ddl(self):
         # Checks that user-defined types are parsed into DataType instead of Identifier
         self.parse_one("CREATE TABLE t (a udt)").this.expressions[0].args["kind"].assert_is(