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