fix: Use regex and cross segment (#448)

formancehq · Jun 21, 2023 · d996e0c · d996e0c
1 parent 87493ec
commit d996e0c
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 4 deletions.
diff --git a/pkg/storage/sqlstorage/accounts.go b/pkg/storage/sqlstorage/accounts.go
@@ -39,7 +39,11 @@ func (s *Store) buildAccountsQuery(p ledger.AccountsQuery) (*sqlbuilder.SelectBu
 		switch s.Schema().Flavor() {
 		case sqlbuilder.PostgreSQL:
 			src := strings.Split(address, ":")
-			sb.Where(fmt.Sprintf("jsonb_array_length(address_json) = %d", len(src)))
+			if address[len(address)-2:] != ".*" {
+				sb.Where(fmt.Sprintf("jsonb_array_length(address_json) = %d", len(src)))
+			} else {
+				src[len(src)-1] = src[len(src)-1][:len(src[len(src)-1])-2]
+			}
 
 			for i, segment := range src {
 				if segment == ".*" || segment == "*" || segment == "" {
@@ -49,6 +53,7 @@ func (s *Store) buildAccountsQuery(p ledger.AccountsQuery) (*sqlbuilder.SelectBu
 				operator := "=="
 				if !accountNameRegex.MatchString(segment) {
 					operator = "like_regex"
+					segment = strings.ReplaceAll(segment, "\\", "\\\\")
 				}
 
 				arg := sb.Args.Add(segment)

diff --git a/pkg/storage/sqlstorage/balances.go b/pkg/storage/sqlstorage/balances.go
@@ -25,7 +25,11 @@ func (s *Store) GetBalancesAggregated(ctx context.Context, q ledger.BalancesQuer
 		switch s.Schema().Flavor() {
 		case sqlbuilder.PostgreSQL:
 			src := strings.Split(q.Filters.AddressRegexp, ":")
-			sb.Where(fmt.Sprintf("jsonb_array_length(account_json) = %d", len(src)))
+			if q.Filters.AddressRegexp[len(q.Filters.AddressRegexp)-2:] != ".*" {
+				sb.Where(fmt.Sprintf("jsonb_array_length(account_json) = %d", len(src)))
+			} else {
+				src[len(src)-1] = src[len(src)-1][:len(src[len(src)-1])-2]
+			}
 
 			for i, segment := range src {
 				if segment == ".*" || segment == "*" || segment == "" {
@@ -35,6 +39,7 @@ func (s *Store) GetBalancesAggregated(ctx context.Context, q ledger.BalancesQuer
 				operator := "=="
 				if !accountNameRegex.MatchString(segment) {
 					operator = "like_regex"
+					segment = strings.ReplaceAll(segment, "\\", "\\\\")
 				}
 
 				arg := sb.Args.Add(segment)
@@ -115,14 +120,18 @@ func (s *Store) GetBalances(ctx context.Context, q ledger.BalancesQuery) (api.Cu
 		switch s.Schema().Flavor() {
 		case sqlbuilder.PostgreSQL:
 			src := strings.Split(q.Filters.AddressRegexp, ":")
-			sb.Where(fmt.Sprintf("jsonb_array_length(account_json) = %d", len(src)))
+			if q.Filters.AddressRegexp[len(q.Filters.AddressRegexp)-2:] != ".*" {
+				sb.Where(fmt.Sprintf("jsonb_array_length(account_json) = %d", len(src)))
+			} else {
+				src[len(src)-1] = src[len(src)-1][:len(src[len(src)-1])-2]
+			}
 
 			for i, segment := range src {
 				if segment == ".*" || segment == "*" || segment == "" {
 					continue
 				}
 
-				arg := sb.Args.Add(segment)
+				arg := sb.Args.Add(strings.ReplaceAll(segment, "\\", "\\\\"))
 				sb.Where(fmt.Sprintf("account_json @@ ('$[%d] like_regex \"' || %s::text || '\"')::jsonpath", i, arg))
 			}
 		case sqlbuilder.SQLite: