From 8ce29cd039c60f9bf8221dfc227041becd30fde7 Mon Sep 17 00:00:00 2001
From: Tanner Rogalsky <tanner@tannerrogalsky.com>
Date: Tue, 7 Mar 2023 15:49:56 +0000
Subject: [PATCH] Add binary operators for Postgres' JSON traversal.

---
 src/backend/postgres/query.rs |  2 ++
 src/extension/postgres/mod.rs |  4 ++++
 tests/postgres/query.rs       | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 40 insertions(+)

diff --git a/src/backend/postgres/query.rs b/src/backend/postgres/query.rs
index 4a51ac081..e452a1352 100644
--- a/src/backend/postgres/query.rs
+++ b/src/backend/postgres/query.rs
@@ -53,6 +53,8 @@ impl QueryBuilder for PostgresQueryBuilder {
                     PgBinOper::SimilarityDistance => "<->",
                     PgBinOper::WordSimilarityDistance => "<<->",
                     PgBinOper::StrictWordSimilarityDistance => "<<<->",
+                    PgBinOper::GetJsonField => "->",
+                    PgBinOper::CastJsonField => "->>",
                 }
             )
             .unwrap(),
diff --git a/src/extension/postgres/mod.rs b/src/extension/postgres/mod.rs
index f84241686..54d7942dd 100644
--- a/src/extension/postgres/mod.rs
+++ b/src/extension/postgres/mod.rs
@@ -25,6 +25,10 @@ pub enum PgBinOper {
     SimilarityDistance,
     WordSimilarityDistance,
     StrictWordSimilarityDistance,
+    /// `->`. Retrieves JSON field as JSON value.
+    GetJsonField,
+    /// `->>`. Retrieves JSON field and casts it to an appropriate SQL type.
+    CastJsonField,
 }
 
 impl From<PgBinOper> for BinOper {
diff --git a/tests/postgres/query.rs b/tests/postgres/query.rs
index 174635ad2..352177db5 100644
--- a/tests/postgres/query.rs
+++ b/tests/postgres/query.rs
@@ -1782,3 +1782,37 @@ fn sub_query_with_fn() {
         r#"SELECT ARRAY((SELECT * FROM "character"))"#
     );
 }
+
+#[test]
+fn get_json_field_bin_oper() {
+    assert_eq!(
+        Query::select()
+            .column(Char::Character)
+            .from(Char::Table)
+            .and_where(
+                Expr::col(Char::Character).binary(PgBinOper::GetJsonField, Expr::val("test"))
+            )
+            .build(PostgresQueryBuilder),
+        (
+            r#"SELECT "character" FROM "character" WHERE "character" -> $1"#.to_owned(),
+            Values(vec!["test".into()])
+        )
+    );
+}
+
+#[test]
+fn cast_json_field_bin_oper() {
+    assert_eq!(
+        Query::select()
+            .column(Char::Character)
+            .from(Char::Table)
+            .and_where(
+                Expr::col(Char::Character).binary(PgBinOper::CastJsonField, Expr::val("test"))
+            )
+            .build(PostgresQueryBuilder),
+        (
+            r#"SELECT "character" FROM "character" WHERE "character" ->> $1"#.to_owned(),
+            Values(vec!["test".into()])
+        )
+    );
+}