sqlglot.parsers.postgres

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, parser
  6from sqlglot.dialects.dialect import (
  7    DialectType,
  8    binary_from_function,
  9    build_formatted_time,
 10    build_json_extract_path,
 11    build_timestamp_trunc,
 12)
 13from sqlglot.helper import is_int, seq_get
 14from sqlglot.parser import binary_range_parser
 15from sqlglot.tokens import TokenType
 16
 17
 18def _build_generate_series(args: list) -> exp.ExplodingGenerateSeries:
 19    # The goal is to convert step values like '1 day' or INTERVAL '1 day' into INTERVAL '1' day
 20    # Note: postgres allows calls with just two arguments -- the "step" argument defaults to 1
 21    step = seq_get(args, 2)
 22    if step is not None:
 23        if step.is_string:
 24            args[2] = exp.to_interval(step.this)
 25        elif isinstance(step, exp.Interval) and not step.args.get("unit"):
 26            args[2] = exp.to_interval(step.this.this)
 27
 28    return exp.ExplodingGenerateSeries.from_arg_list(args)
 29
 30
 31def _build_to_timestamp(args: list) -> exp.UnixToTime | exp.StrToTime:
 32    # TO_TIMESTAMP accepts either a single double argument or (text, text)
 33    if len(args) == 1:
 34        # https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-TABLE
 35        return exp.UnixToTime.from_arg_list(args)
 36
 37    # https://www.postgresql.org/docs/current/functions-formatting.html
 38    return build_formatted_time(exp.StrToTime, "postgres")(args)
 39
 40
 41def _build_regexp_replace(args: list, dialect: DialectType = None) -> exp.RegexpReplace:
 42    # The signature of REGEXP_REPLACE is:
 43    # regexp_replace(source, pattern, replacement [, start [, N ]] [, flags ])
 44    #
 45    # Any one of `start`, `N` and `flags` can be column references, meaning that
 46    # unless we can statically see that the last argument is a non-integer string
 47    # (eg. not '0'), then it's not possible to construct the correct AST
 48    regexp_replace = None
 49    if len(args) > 3:
 50        last = args[-1]
 51        if not is_int(last.name):
 52            if not last.type or last.is_type(exp.DType.UNKNOWN, exp.DType.NULL):
 53                from sqlglot.optimizer.annotate_types import annotate_types
 54
 55                last = annotate_types(last, dialect=dialect)
 56
 57            if last.is_type(*exp.DataType.TEXT_TYPES):
 58                regexp_replace = exp.RegexpReplace.from_arg_list(args[:-1])
 59                regexp_replace.set("modifiers", last)
 60
 61    regexp_replace = regexp_replace or exp.RegexpReplace.from_arg_list(args)
 62    regexp_replace.set("single_replace", True)
 63    return regexp_replace
 64
 65
 66def _build_levenshtein_less_equal(args: list) -> exp.Levenshtein:
 67    # Postgres has two signatures for levenshtein_less_equal function, but in both cases
 68    # max_dist is the last argument
 69    # levenshtein_less_equal(source, target, ins_cost, del_cost, sub_cost, max_d)
 70    # levenshtein_less_equal(source, target, max_d)
 71    max_dist = args.pop()
 72
 73    return exp.Levenshtein(
 74        this=seq_get(args, 0),
 75        expression=seq_get(args, 1),
 76        ins_cost=seq_get(args, 2),
 77        del_cost=seq_get(args, 3),
 78        sub_cost=seq_get(args, 4),
 79        max_dist=max_dist,
 80    )
 81
 82
 83class PostgresParser(parser.Parser):
 84    SUPPORTS_OMITTED_INTERVAL_SPAN_UNIT = True
 85
 86    PROPERTY_PARSERS = {
 87        **{k: v for k, v in parser.Parser.PROPERTY_PARSERS.items() if k != "INPUT"},
 88        "SET": lambda self: self.expression(exp.SetConfigProperty(this=self._parse_set())),
 89    }
 90
 91    PLACEHOLDER_PARSERS = {
 92        **parser.Parser.PLACEHOLDER_PARSERS,
 93        TokenType.PLACEHOLDER: lambda self: self.expression(exp.Placeholder(jdbc=True)),
 94        TokenType.MOD: lambda self: self._parse_query_parameter(),
 95    }
 96
 97    FUNCTIONS = {
 98        **parser.Parser.FUNCTIONS,
 99        "ARRAY_PREPEND": lambda args: exp.ArrayPrepend(
100            this=seq_get(args, 1), expression=seq_get(args, 0)
101        ),
102        "BIT_AND": exp.BitwiseAndAgg.from_arg_list,
103        "BIT_OR": exp.BitwiseOrAgg.from_arg_list,
104        "BIT_XOR": exp.BitwiseXorAgg.from_arg_list,
105        "VERSION": exp.CurrentVersion.from_arg_list,
106        "DATE_TRUNC": build_timestamp_trunc,
107        "DIV": lambda args: exp.cast(binary_from_function(exp.IntDiv)(args), exp.DType.DECIMAL),
108        "GENERATE_SERIES": _build_generate_series,
109        "GET_BIT": lambda args: exp.Getbit(
110            this=seq_get(args, 0), expression=seq_get(args, 1), zero_is_msb=True
111        ),
112        "JSON_EXTRACT_PATH": build_json_extract_path(exp.JSONExtract),
113        "JSON_EXTRACT_PATH_TEXT": build_json_extract_path(exp.JSONExtractScalar),
114        "LENGTH": lambda args: exp.Length(this=seq_get(args, 0), encoding=seq_get(args, 1)),
115        "MAKE_TIME": exp.TimeFromParts.from_arg_list,
116        "MAKE_TIMESTAMP": exp.TimestampFromParts.from_arg_list,
117        "NOW": exp.CurrentTimestamp.from_arg_list,
118        "REGEXP_REPLACE": _build_regexp_replace,
119        "TO_CHAR": build_formatted_time(exp.TimeToStr, "postgres"),
120        "TO_DATE": build_formatted_time(exp.StrToDate, "postgres"),
121        "TO_TIMESTAMP": _build_to_timestamp,
122        "UNNEST": exp.Explode.from_arg_list,
123        "SHA256": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(256)),
124        "SHA384": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(384)),
125        "SHA512": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(512)),
126        "LEVENSHTEIN_LESS_EQUAL": _build_levenshtein_less_equal,
127        "JSON_OBJECT_AGG": lambda args: exp.JSONObjectAgg(expressions=args),
128        "JSONB_OBJECT_AGG": exp.JSONBObjectAgg.from_arg_list,
129        "WIDTH_BUCKET": lambda args: (
130            exp.WidthBucket(this=seq_get(args, 0), threshold=seq_get(args, 1))
131            if len(args) == 2
132            else exp.WidthBucket.from_arg_list(args)
133        ),
134    }
135
136    NO_PAREN_FUNCTION_PARSERS = {
137        **parser.Parser.NO_PAREN_FUNCTION_PARSERS,
138        "VARIADIC": lambda self: self.expression(exp.Variadic(this=self._parse_bitwise())),
139    }
140
141    NO_PAREN_FUNCTIONS = {
142        **parser.Parser.NO_PAREN_FUNCTIONS,
143        TokenType.LOCALTIME: exp.Localtime,
144        TokenType.LOCALTIMESTAMP: exp.Localtimestamp,
145        TokenType.CURRENT_CATALOG: exp.CurrentCatalog,
146        TokenType.SESSION_USER: exp.SessionUser,
147        TokenType.CURRENT_SCHEMA: exp.CurrentSchema,
148    }
149
150    FUNCTION_PARSERS = {
151        **parser.Parser.FUNCTION_PARSERS,
152        "DATE_PART": lambda self: self._parse_date_part(),
153        "JSON_AGG": lambda self: self.expression(
154            exp.JSONArrayAgg(this=self._parse_lambda(), order=self._parse_order())
155        ),
156        "JSONB_EXISTS": lambda self: self._parse_jsonb_exists(),
157    }
158
159    BITWISE = {
160        **parser.Parser.BITWISE,
161        TokenType.HASH: exp.BitwiseXor,
162    }
163
164    EXPONENT = {
165        TokenType.CARET: exp.Pow,
166    }
167
168    RANGE_PARSERS = {
169        **parser.Parser.RANGE_PARSERS,
170        TokenType.DAMP: binary_range_parser(exp.ArrayOverlaps),
171        TokenType.DAT: lambda self, this: self.expression(
172            exp.MatchAgainst(this=self._parse_bitwise(), expressions=[this])
173        ),
174    }
175
176    STATEMENT_PARSERS = {
177        **parser.Parser.STATEMENT_PARSERS,
178        TokenType.END: lambda self: self._parse_commit_or_rollback(),
179    }
180
181    UNARY_PARSERS = {
182        **parser.Parser.UNARY_PARSERS,
183        # The `~` token is remapped from TILDE to RLIKE in Postgres due to the binary REGEXP LIKE operator
184        TokenType.RLIKE: lambda self: self.expression(exp.BitwiseNot(this=self._parse_unary())),
185    }
186
187    JSON_ARROWS_REQUIRE_JSON_TYPE = True
188
189    COLUMN_OPERATORS = {
190        **parser.Parser.COLUMN_OPERATORS,
191        TokenType.ARROW: lambda self, this, path: self.validate_expression(
192            build_json_extract_path(
193                exp.JSONExtract, arrow_req_json_type=self.JSON_ARROWS_REQUIRE_JSON_TYPE
194            )([this, path])
195        ),
196        TokenType.DARROW: lambda self, this, path: self.validate_expression(
197            build_json_extract_path(
198                exp.JSONExtractScalar, arrow_req_json_type=self.JSON_ARROWS_REQUIRE_JSON_TYPE
199            )([this, path])
200        ),
201    }
202
203    ARG_MODE_TOKENS: t.ClassVar = {TokenType.IN, TokenType.OUT, TokenType.INOUT, TokenType.VARIADIC}
204
205    def _parse_parameter_mode(self) -> TokenType | None:
206        """
207        Parse PostgreSQL function parameter mode (IN, OUT, INOUT, VARIADIC).
208
209        Disambiguates between mode keywords and identifiers with the same name:
210        - MODE TYPE      -> keyword is identifier (e.g., "out INT")
211        - MODE NAME TYPE -> keyword is mode (e.g., "OUT x INT")
212
213        Returns:
214            Mode token type if current token is a mode keyword, None otherwise.
215        """
216        if not self._match_set(self.ARG_MODE_TOKENS, advance=False) or not self._next:
217            return None
218
219        mode_token = self._curr
220
221        # Check Pattern 1: MODE TYPE
222        # Try parsing next token as a built-in type (not UDT)
223        # If successful, the keyword is an identifier, not a mode
224        is_followed_by_builtin_type = self._try_parse(
225            lambda: (
226                self._advance()  # type: ignore
227                or self._parse_types(check_func=False, allow_identifiers=False)
228            ),
229            retreat=True,
230        )
231        if is_followed_by_builtin_type:
232            return None  # Pattern: "out INT" -> out is parameter name
233
234        # Check Pattern 2: MODE NAME TYPE
235        # If next token is an identifier, check if there's a type after it
236        # The type can be built-in or user-defined (allow_identifiers=True)
237        if self._next.token_type not in self.ID_VAR_TOKENS:
238            return None
239
240        is_followed_by_any_type = self._try_parse(
241            lambda: (
242                self._advance(2)  # type: ignore
243                or self._parse_types(check_func=False, allow_identifiers=True)
244            ),
245            retreat=True,
246        )
247
248        if is_followed_by_any_type:
249            return mode_token.token_type  # Pattern: "OUT x INT" -> OUT is mode
250
251        return None
252
253    def _create_mode_constraint(self, param_mode: TokenType) -> exp.InOutColumnConstraint:
254        """
255        Create parameter mode constraint for function parameters.
256
257        Args:
258            param_mode: The parameter mode token (IN, OUT, INOUT, or VARIADIC).
259
260        Returns:
261            InOutColumnConstraint expression representing the parameter mode.
262        """
263        return self.expression(
264            exp.InOutColumnConstraint(
265                input_=(param_mode in {TokenType.IN, TokenType.INOUT}),
266                output=(param_mode in {TokenType.OUT, TokenType.INOUT}),
267                variadic=(param_mode == TokenType.VARIADIC),
268            )
269        )
270
271    def _parse_function_parameter(self) -> exp.Expr | None:
272        param_mode = self._parse_parameter_mode()
273
274        if param_mode:
275            self._advance()
276
277        # Parse parameter name and type
278        param_name = self._parse_id_var()
279        column_def = self._parse_column_def(this=param_name, computed_column=False)
280
281        # Attach mode as constraint
282        if param_mode and column_def:
283            constraint = self._create_mode_constraint(param_mode)
284            if not column_def.args.get("constraints"):
285                column_def.set("constraints", [])
286            column_def.args["constraints"].insert(0, constraint)
287
288        return column_def
289
290    def _parse_query_parameter(self) -> exp.Expr | None:
291        this = (
292            self._parse_wrapped(self._parse_id_var)
293            if self._match(TokenType.L_PAREN, advance=False)
294            else None
295        )
296        self._match_text_seq("S")
297        return self.expression(exp.Placeholder(this=this))
298
299    def _parse_date_part(self) -> exp.Expr:
300        part = self._parse_type()
301        self._match(TokenType.COMMA)
302        value = self._parse_bitwise()
303
304        if part and isinstance(part, (exp.Column, exp.Literal)):
305            part = exp.var(part.name)
306
307        return self.expression(exp.Extract(this=part, expression=value))
308
309    def _parse_unique_key(self) -> exp.Expr | None:
310        return None
311
312    def _parse_jsonb_exists(self) -> exp.JSONBExists:
313        return self.expression(
314            exp.JSONBExists(
315                this=self._parse_bitwise(),
316                path=self._match(TokenType.COMMA)
317                and self.dialect.to_json_path(self._parse_bitwise()),
318            )
319        )
320
321    def _parse_generated_as_identity(
322        self,
323    ) -> (
324        exp.GeneratedAsIdentityColumnConstraint
325        | exp.ComputedColumnConstraint
326        | exp.GeneratedAsRowColumnConstraint
327    ):
328        this = super()._parse_generated_as_identity()
329
330        if self._match_text_seq("STORED"):
331            this = self.expression(exp.ComputedColumnConstraint(this=this.expression))
332
333        return this
334
335    def _parse_user_defined_type(self, identifier: exp.Identifier) -> exp.Expr | None:
336        udt_type: exp.Identifier | exp.Dot = identifier
337
338        while self._match(TokenType.DOT):
339            part = self._parse_id_var()
340            if part:
341                udt_type = exp.Dot(this=udt_type, expression=part)
342
343        return exp.DataType.build(udt_type, udt=True)