feat(postgres): Support for IS JSON predicate (#3971)

* feat(postgres): Support for IS JSON predicate

* PR Feedback 1

sqlglot/expressions.py

@@ -5441,6 +5441,11 @@ class IsInf(Func):
  _sql_names = ["IS_INF", "ISINF"]
 
 
+# https://www.postgresql.org/docs/current/functions-json.html
+class JSON(Expression):
+ arg_types = {"this": False, "with": False, "unique": False}
+
+
 class JSONPath(Expression):
  arg_types = {"expressions": True}
 

sqlglot/generator.py

@@ -4136,3 +4136,20 @@ def converttimezone_sql(self, expression: exp.ConvertTimezone) -> str:
  expr = exp.AtTimeZone(this=timestamp, zone=target_tz)
 
  return self.sql(expr)
+
+ def json_sql(self, expression: exp.JSON) -> str:
+ this = self.sql(expression, "this")
+ this = f" {this}" if this else ""
+
+ _with = expression.args.get("with")
+
+ if _with is None:
+ with_sql = ""
+ elif not _with:
+ with_sql = " WITHOUT"
+ else:
+ with_sql = " WITH"
+
+ unique_sql = " UNIQUE KEYS" if expression.args.get("unique") else ""
+
+ return f"JSON{this}{with_sql}{unique_sql}"

sqlglot/parser.py

@@ -1252,6 +1252,8 @@ class Parser(metaclass=_Parser):
 
  COPY_INTO_VARLEN_OPTIONS = {"FILE_FORMAT", "COPY_OPTIONS", "FORMAT_OPTIONS", "CREDENTIAL"}
 
+ IS_JSON_PREDICATE_KIND = {"VALUE", "SCALAR", "ARRAY", "OBJECT"}
+
  STRICT_CAST = True
 
  PREFIXED_PIVOT_COLUMNS = False
@@ -4288,10 +4290,26 @@ def _parse_is(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.Expressi
  klass = exp.NullSafeEQ if negate else exp.NullSafeNEQ
  return self.expression(klass, this=this, expression=self._parse_bitwise())
 
- expression = self._parse_null() or self._parse_boolean()
- if not expression:
- self._retreat(index)
- return None
+ if self._match(TokenType.JSON):
+ kind = self._match_texts(self.IS_JSON_PREDICATE_KIND) and self._prev.text.upper()
+
+ if self._match_text_seq("WITH"):
+ _with = True
+ elif self._match_text_seq("WITHOUT"):
+ _with = False
+ else:
+ _with = None
+
+ unique = self._match(TokenType.UNIQUE)
+ self._match_text_seq("KEYS")
+ expression: t.Optional[exp.Expression] = self.expression(
+ exp.JSON, **{"this": kind, "with": _with, "unique": unique}
+ )
+ else:
+ expression = self._parse_null() or self._parse_boolean()
+ if not expression:
+ self._retreat(index)
+ return None
 
  this = self.expression(exp.Is, this=this, expression=expression)
  return self.expression(exp.Not, this=this) if negate else this

tests/dialects/test_postgres.py

@@ -788,6 +788,13 @@ def test_postgres(self):
  },
  )
 
+ self.validate_identity(
+ 'SELECT js, js IS JSON AS "json?", js IS JSON VALUE AS "scalar?", js IS JSON SCALAR AS "scalar?", js IS JSON OBJECT AS "object?", js IS JSON ARRAY AS "array?" FROM t'
+ )
+ self.validate_identity(
+ 'SELECT js, js IS JSON ARRAY WITH UNIQUE KEYS AS "array w. UK?", js IS JSON ARRAY WITHOUT UNIQUE KEYS AS "array w/o UK?", js IS JSON ARRAY UNIQUE KEYS AS "array w UK 2?" FROM t'
+ )
+
  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(