jsonb/
keypath.rs

1// Copyright 2023 Datafuse Labs.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::borrow::Cow;
16use std::fmt::Display;
17use std::fmt::Formatter;
18
19use nom::branch::alt;
20use nom::character::complete::char;
21use nom::character::complete::i32;
22use nom::character::complete::multispace0;
23use nom::combinator::map;
24use nom::multi::separated_list1;
25use nom::sequence::delimited;
26use nom::sequence::preceded;
27use nom::sequence::terminated;
28use nom::IResult;
29use nom::Parser;
30
31use crate::jsonpath::raw_string;
32use crate::jsonpath::string;
33use crate::Error;
34
35/// Represents a set of key path chains.
36/// Compatible with PostgreSQL extracts JSON sub-object paths syntax.
37#[derive(Debug, Clone, Eq, PartialEq, Hash)]
38pub struct KeyPaths<'a> {
39    pub paths: Vec<KeyPath<'a>>,
40}
41
42/// Represents a valid key path.
43#[derive(Debug, Clone, Eq, PartialEq, Hash)]
44pub enum KeyPath<'a> {
45    /// represents the index of an Array, allow negative indexing.
46    Index(i32),
47    /// represents the quoted field name of an Object.
48    QuotedName(Cow<'a, str>),
49    /// represents the field name of an Object.
50    Name(Cow<'a, str>),
51}
52
53impl Display for KeyPaths<'_> {
54    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
55        write!(f, "{{")?;
56        for (i, path) in self.paths.iter().enumerate() {
57            if i > 0 {
58                write!(f, ",")?;
59            }
60            write!(f, "{path}")?;
61        }
62        write!(f, "}}")?;
63        Ok(())
64    }
65}
66
67impl Display for KeyPath<'_> {
68    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
69        match self {
70            KeyPath::Index(idx) => {
71                write!(f, "{idx}")?;
72            }
73            KeyPath::QuotedName(name) => {
74                write!(f, "\"{name}\"")?;
75            }
76            KeyPath::Name(name) => {
77                write!(f, "{name}")?;
78            }
79        }
80        Ok(())
81    }
82}
83
84/// Parsing the input string to key paths.
85pub fn parse_key_paths(input: &[u8]) -> Result<KeyPaths<'_>, Error> {
86    match key_paths(input) {
87        Ok((rest, paths)) => {
88            if !rest.is_empty() {
89                return Err(Error::InvalidKeyPath);
90            }
91            let key_paths = KeyPaths { paths };
92            Ok(key_paths)
93        }
94        Err(nom::Err::Error(_) | nom::Err::Failure(_)) => Err(Error::InvalidKeyPath),
95        Err(nom::Err::Incomplete(_)) => unreachable!(),
96    }
97}
98
99fn key_path(input: &[u8]) -> IResult<&[u8], KeyPath<'_>> {
100    alt((
101        map(i32, KeyPath::Index),
102        map(string, KeyPath::QuotedName),
103        map(raw_string, KeyPath::Name),
104    ))
105    .parse(input)
106}
107
108fn key_paths(input: &[u8]) -> IResult<&[u8], Vec<KeyPath<'_>>> {
109    alt((
110        delimited(
111            preceded(multispace0, char('{')),
112            separated_list1(char(','), delimited(multispace0, key_path, multispace0)),
113            terminated(char('}'), multispace0),
114        ),
115        map(
116            delimited(
117                preceded(multispace0, char('{')),
118                multispace0,
119                terminated(char('}'), multispace0),
120            ),
121            |_| vec![],
122        ),
123    ))
124    .parse(input)
125}