refactor parseJSONQuery variants to single function (#215)

getzep · Oct 6, 2023 · 40a6cad · 40a6cad
1 parent b3b42ac
commit 40a6cad
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 50 deletions.
diff --git a/pkg/store/postgres/document_search.go b/pkg/store/postgres/document_search.go
@@ -217,7 +217,7 @@ func (dso *documentSearchOperation) applyDocsMetadataFilter(
 		if err != nil {
 			return nil, fmt.Errorf("error unmarshalling metadata %w", err)
 		}
-		qb = parseDocumentJSONQuery(qb, &jq, false)
+		qb = parseJSONQuery(qb, &jq, false, "")
 	}
 
 	query = qb.Unwrap().(*bun.SelectQuery)

diff --git a/pkg/store/postgres/search_memory.go b/pkg/store/postgres/search_memory.go
@@ -114,7 +114,7 @@ func applyMessagesMetadataFilter(
 		if err != nil {
 			return nil, store.NewStorageError("error unmarshalling metadata", err)
 		}
-		qb = parseJSONQuery(qb, &jq, false)
+		qb = parseJSONQuery(qb, &jq, false, "m")
 	}
 
 	addMessageDateFilters(&qb, metadata)

diff --git a/pkg/store/postgres/search_utils.go b/pkg/store/postgres/search_utils.go
@@ -1,65 +1,29 @@
 package postgres
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/uptrace/bun"
 )
 
-// TODO: refactor to a single function used across both document and memory search
-
 // parseJSONQuery recursively parses a JSONQuery and returns a bun.QueryBuilder.
 // TODO: fix the addition of extraneous parentheses in the query
-func parseJSONQuery(qb bun.QueryBuilder, jq *JSONQuery, isOr bool) bun.QueryBuilder {
-	if jq.JSONPath != "" {
-		path := strings.ReplaceAll(jq.JSONPath, "'", "\"")
-		if isOr {
-			qb = qb.WhereOr(
-				"jsonb_path_exists(m.metadata, ?)",
-				path,
-			)
-		} else {
-			qb = qb.Where(
-				"jsonb_path_exists(m.metadata, ?)",
-				path,
-			)
-		}
-	}
-
-	if len(jq.And) > 0 {
-		qb = qb.WhereGroup(" AND ", func(qq bun.QueryBuilder) bun.QueryBuilder {
-			for _, subQuery := range jq.And {
-				qq = parseJSONQuery(qq, subQuery, false)
-			}
-			return qq
-		})
-	}
-
-	if len(jq.Or) > 0 {
-		qb = qb.WhereGroup(" AND ", func(qq bun.QueryBuilder) bun.QueryBuilder {
-			for _, subQuery := range jq.Or {
-				qq = parseJSONQuery(qq, subQuery, true)
-			}
-			return qq
-		})
+func parseJSONQuery(qb bun.QueryBuilder, jq *JSONQuery, isOr bool, tablePrefix string) bun.QueryBuilder {
+	var tp string
+	if tablePrefix != "" {
+		tp = tablePrefix + "."
 	}
-
-	return qb
-}
-
-// parseDocumentJSONQuery recursively parses a JSONQuery and returns a bun.QueryBuilder.
-// TODO: fix the addition of extraneous parentheses in the query
-func parseDocumentJSONQuery(qb bun.QueryBuilder, jq *JSONQuery, isOr bool) bun.QueryBuilder {
 	if jq.JSONPath != "" {
 		path := strings.ReplaceAll(jq.JSONPath, "'", "\"")
 		if isOr {
 			qb = qb.WhereOr(
-				"jsonb_path_exists(metadata, ?)",
+				fmt.Sprintf("jsonb_path_exists(%smetadata, ?)", tp),
 				path,
 			)
 		} else {
 			qb = qb.Where(
-				"jsonb_path_exists(metadata, ?)",
+				"jsonb_path_exists(m.metadata, ?)",
 				path,
 			)
 		}
@@ -68,7 +32,7 @@ func parseDocumentJSONQuery(qb bun.QueryBuilder, jq *JSONQuery, isOr bool) bun.Q
 	if len(jq.And) > 0 {
 		qb = qb.WhereGroup(" AND ", func(qq bun.QueryBuilder) bun.QueryBuilder {
 			for _, subQuery := range jq.And {
-				qq = parseDocumentJSONQuery(qq, subQuery, false)
+				qq = parseJSONQuery(qq, subQuery, false, tablePrefix)
 			}
 			return qq
 		})
@@ -77,7 +41,7 @@ func parseDocumentJSONQuery(qb bun.QueryBuilder, jq *JSONQuery, isOr bool) bun.Q
 	if len(jq.Or) > 0 {
 		qb = qb.WhereGroup(" AND ", func(qq bun.QueryBuilder) bun.QueryBuilder {
 			for _, subQuery := range jq.Or {
-				qq = parseDocumentJSONQuery(qq, subQuery, true)
+				qq = parseJSONQuery(qq, subQuery, true, tablePrefix)
 			}
 			return qq
 		})

diff --git a/pkg/store/postgres/search_utils_test.go b/pkg/store/postgres/search_utils_test.go
@@ -15,21 +15,25 @@ func TestParseJSONQuery(t *testing.T) {
 		name         string
 		jsonQuery    string
 		expectedCond string
+		tablePrefix  string
 	}{
 		{
 			name:         "Test 1",
 			jsonQuery:    `{"where": {"jsonpath": "$.system.entities[*] ? (@.Label == \"DATE\")"}}`,
 			expectedCond: `WHERE (jsonb_path_exists(m.metadata, '$.system.entities[*] ? (@.Label == "DATE")'))`,
+			tablePrefix:  "m",
 		},
 		{
-			name:         "Test 2",
+			name:         "Without Prefix",
 			jsonQuery:    `{"where": {"or": [{"jsonpath": "$.system.entities[*] ? (@.Label == \"DATE\")"},{"jsonpath": "$.system.entities[*] ? (@.Label == \"ORG\")"}]}}`,
-			expectedCond: `WHERE ((jsonb_path_exists(m.metadata, '$.system.entities[*] ? (@.Label == "DATE")')) OR (jsonb_path_exists(m.metadata, '$.system.entities[*] ? (@.Label == "ORG")')))`,
+			expectedCond: `WHERE ((jsonb_path_exists(metadata, '$.system.entities[*] ? (@.Label == "DATE")')) OR (jsonb_path_exists(metadata, '$.system.entities[*] ? (@.Label == "ORG")')))`,
+			tablePrefix:  "",
 		},
 		{
 			name:         "Test 3",
 			jsonQuery:    `{"where": {"and": [{"jsonpath": "$.system.entities[*] ? (@.Label == \"DATE\")"},{"jsonpath": "$.system.entities[*] ? (@.Label == \"ORG\")"},{"or": [{"jsonpath": "$.system.entities[*] ? (@.Name == \"Iceland\")"},{"jsonpath": "$.system.entities[*] ? (@.Name == \"Canada\")"}]}]}}`,
 			expectedCond: `WHERE ((jsonb_path_exists(m.metadata, '$.system.entities[*] ? (@.Label == "DATE")')) AND (jsonb_path_exists(m.metadata, '$.system.entities[*] ? (@.Label == "ORG")')) AND ((jsonb_path_exists(m.metadata, '$.system.entities[*] ? (@.Name == "Iceland")')) OR (jsonb_path_exists(m.metadata, '$.system.entities[*] ? (@.Name == "Canada")'))))`,
+			tablePrefix:  "m",
 		},
 	}
 
@@ -50,7 +54,7 @@ func TestParseJSONQuery(t *testing.T) {
 			err = json.Unmarshal(query, &jsonQuery)
 			assert.NoError(t, err)
 
-			qb = parseJSONQuery(qb, &jsonQuery, false)
+			qb = parseJSONQuery(qb, &jsonQuery, false, tt.tablePrefix)
 
 			selectQuery := qb.Unwrap().(*bun.SelectQuery)