add jsonb type

klahnakoski · May 23, 2024 · 1ff7931 · 1ff7931
1 parent 32d1928
commit 1ff7931
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 1 deletion.
diff --git a/mo_sql_parsing/keywords.py b/mo_sql_parsing/keywords.py
@@ -126,6 +126,17 @@
 NEQ = (Literal("!=") | Literal("<>")).set_parser_name("neq")
 ASSIGN = Literal(":=").set_parser_name("assign")
 
+JSON_GET = Literal("->").set_parser_name("json_get")
+JSON_GET_TEXT = Literal("->>").set_parser_name("json_get_text")
+JSON_PATH = Literal("#>").set_parser_name("json_path")
+JSON_PATH_TEXT = Literal("#>>").set_parser_name("json_path_text")
+JSON_SUBSUMES = Literal("@>").set_parser_name("json_subsumes")
+JSON_SUBSUMED_BY = Literal("<@").set_parser_name("json_subsumed_by")
+JSON_CONTAINS = Literal("?").set_parser_name("json_contains")
+JSON_CONTAINS_ANY = Literal("?|").set_parser_name("json_contains_any")
+JSON_CONTAINS_ALL = Literal("?&").set_parser_name("json_contains_all")
+JSON_PATH_DEL = Literal("#-").set_parser_name("json_path_del")
+
 AND = keyword("and")
 APPLY = keyword("apply")
 BEGIN = keyword("begin").suppress()
@@ -289,6 +300,16 @@
     "validate_conversion": 0,
     "collate": 0,
     "concat": 1,
+    "json_get": 1.5,
+    "json_get_text": 1.5,
+    "json_path": 1.5,
+    "json_path_text": 1.5,
+    "json_subsumes": 1.5,
+    "json_subsumed_by": 1.5,
+    "json_contains": 1.5,
+    "json_contains_any": 1.5,
+    "json_contains_all": 1.5,
+    "json_path_del": 1.5,
     "mul": 2,
     "div": 1.5,
     "mod": 2,
@@ -350,6 +371,16 @@
 KNOWN_OPS = [
     COLLATE,
     CONCAT,
+    JSON_GET_TEXT
+    | JSON_GET
+    | JSON_PATH_TEXT
+    | JSON_PATH
+    | JSON_SUBSUMES
+    | JSON_SUBSUMED_BY
+    | JSON_CONTAINS_ANY
+    | JSON_CONTAINS_ALL
+    | JSON_CONTAINS
+    | JSON_PATH_DEL,
     POS,
     NEG,
     MUL | DIV | MOD,

diff --git a/mo_sql_parsing/types.py b/mo_sql_parsing/types.py
@@ -116,6 +116,7 @@
 MAP_TYPE = (keyword("map")("op") + LB + delimited_list(simple_types("params")) + RB) / to_json_call
 ARRAY_TYPE = (keyword("array")("op") + LB + simple_types("params") + RB) / to_json_call
 JSON = Group(keyword("json")("op")) / to_json_call
+JSONB = Group(keyword("jsonb")("op")) / to_json_call
 
 DATE = keyword("date")
 DATETIME = keyword("datetime")
@@ -171,6 +172,7 @@
     INT64,
     BYTEINT,
     JSON,
+    JSONB,
     NCHAR,
     NVARCHAR,
     NUMBER,

diff --git a/mo_sql_parsing/utils.py b/mo_sql_parsing/utils.py
@@ -237,6 +237,16 @@ def to_tuple_call(token, index, string):
     "COLLATE": "collate",
     ":": "get",
     "||": "concat",
+    "->": "json_get",
+    "->>": "json_get_text",
+    "#>": "json_path",
+    "#>>": "json_path_text",
+    "@>": "json_subsumes",
+    "<@": "json_subsumed_by",
+    "?": "json_contains",
+    "?|": "json_contains_any",
+    "?&": "json_contains_all",
+    "#-": "json_path_del",
     "*": "mul",
     "/": "div",
     "%": "mod",
@@ -274,7 +284,6 @@ def to_tuple_call(token, index, string):
     "not_simlilar_to": "not_similar_to",
     "or": "or",
     "and": "and",
-    "->": "lambda",
     ":=": "assign",
     "union": "union",
     "union_all": "union_all",

diff --git a/tests/test_postgres.py b/tests/test_postgres.py
@@ -529,3 +529,18 @@ def test_issue_175_extract_millennium(self):
         result = parse(sql)
         expected = {"select": {"value": {"extract": ["millennium", "date"]}}}
         self.assertEqual(result, expected)
+
+    def test_issue_239_jsonb1(self):
+        sql = """select jsonb ->> 'field_key' FROM a"""
+        result = parse(sql)
+        expected = {"from": "a", "select": {"value": {"json_get_text": ["jsonb", {"literal": "field_key"}]}}}
+        self.assertEqual(result, expected)
+
+    def test_issue_239_jsonb2(self):
+        sql = """select name::jsonb ->> 'field_key' FROM a"""
+        result = parse(sql)
+        expected = {
+            "from": "a",
+            "select": {"value": {"json_get_text": [{"cast": ["name", {"jsonb": {}}]}, {"literal": "field_key"}]}},
+        }
+        self.assertEqual(result, expected)