sqlglot.dialects.postgres

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, generator, parser, tokens, transforms
  6from sqlglot.dialects.dialect import (
  7    DATE_ADD_OR_SUB,
  8    Dialect,
  9    any_value_to_max_sql,
 10    bool_xor_sql,
 11    datestrtodate_sql,
 12    format_time_lambda,
 13    json_extract_segments,
 14    json_path_key_only_name,
 15    max_or_greatest,
 16    merge_without_target_sql,
 17    min_or_least,
 18    no_last_day_sql,
 19    no_map_from_entries_sql,
 20    no_paren_current_date_sql,
 21    no_pivot_sql,
 22    no_trycast_sql,
 23    parse_json_extract_path,
 24    parse_timestamp_trunc,
 25    rename_func,
 26    str_position_sql,
 27    struct_extract_sql,
 28    timestamptrunc_sql,
 29    timestrtotime_sql,
 30    trim_sql,
 31    ts_or_ds_add_cast,
 32)
 33from sqlglot.helper import seq_get
 34from sqlglot.parser import binary_range_parser
 35from sqlglot.tokens import TokenType
 36
 37DATE_DIFF_FACTOR = {
 38    "MICROSECOND": " * 1000000",
 39    "MILLISECOND": " * 1000",
 40    "SECOND": "",
 41    "MINUTE": " / 60",
 42    "HOUR": " / 3600",
 43    "DAY": " / 86400",
 44}
 45
 46
 47def _date_add_sql(kind: str) -> t.Callable[[Postgres.Generator, DATE_ADD_OR_SUB], str]:
 48    def func(self: Postgres.Generator, expression: DATE_ADD_OR_SUB) -> str:
 49        if isinstance(expression, exp.TsOrDsAdd):
 50            expression = ts_or_ds_add_cast(expression)
 51
 52        this = self.sql(expression, "this")
 53        unit = expression.args.get("unit")
 54
 55        expression = self._simplify_unless_literal(expression.expression)
 56        if not isinstance(expression, exp.Literal):
 57            self.unsupported("Cannot add non literal")
 58
 59        expression.args["is_string"] = True
 60        return f"{this} {kind} {self.sql(exp.Interval(this=expression, unit=unit))}"
 61
 62    return func
 63
 64
 65def _date_diff_sql(self: Postgres.Generator, expression: exp.DateDiff) -> str:
 66    unit = expression.text("unit").upper()
 67    factor = DATE_DIFF_FACTOR.get(unit)
 68
 69    end = f"CAST({self.sql(expression, 'this')} AS TIMESTAMP)"
 70    start = f"CAST({self.sql(expression, 'expression')} AS TIMESTAMP)"
 71
 72    if factor is not None:
 73        return f"CAST(EXTRACT(epoch FROM {end} - {start}){factor} AS BIGINT)"
 74
 75    age = f"AGE({end}, {start})"
 76
 77    if unit == "WEEK":
 78        unit = f"EXTRACT(days FROM ({end} - {start})) / 7"
 79    elif unit == "MONTH":
 80        unit = f"EXTRACT(year FROM {age}) * 12 + EXTRACT(month FROM {age})"
 81    elif unit == "QUARTER":
 82        unit = f"EXTRACT(year FROM {age}) * 4 + EXTRACT(month FROM {age}) / 3"
 83    elif unit == "YEAR":
 84        unit = f"EXTRACT(year FROM {age})"
 85    else:
 86        unit = age
 87
 88    return f"CAST({unit} AS BIGINT)"
 89
 90
 91def _substring_sql(self: Postgres.Generator, expression: exp.Substring) -> str:
 92    this = self.sql(expression, "this")
 93    start = self.sql(expression, "start")
 94    length = self.sql(expression, "length")
 95
 96    from_part = f" FROM {start}" if start else ""
 97    for_part = f" FOR {length}" if length else ""
 98
 99    return f"SUBSTRING({this}{from_part}{for_part})"
100
101
102def _string_agg_sql(self: Postgres.Generator, expression: exp.GroupConcat) -> str:
103    separator = expression.args.get("separator") or exp.Literal.string(",")
104
105    order = ""
106    this = expression.this
107    if isinstance(this, exp.Order):
108        if this.this:
109            this = this.this.pop()
110        order = self.sql(expression.this)  # Order has a leading space
111
112    return f"STRING_AGG({self.format_args(this, separator)}{order})"
113
114
115def _datatype_sql(self: Postgres.Generator, expression: exp.DataType) -> str:
116    if expression.is_type("array"):
117        return f"{self.expressions(expression, flat=True)}[]" if expression.expressions else "ARRAY"
118    return self.datatype_sql(expression)
119
120
121def _auto_increment_to_serial(expression: exp.Expression) -> exp.Expression:
122    auto = expression.find(exp.AutoIncrementColumnConstraint)
123
124    if auto:
125        expression.args["constraints"].remove(auto.parent)
126        kind = expression.args["kind"]
127
128        if kind.this == exp.DataType.Type.INT:
129            kind.replace(exp.DataType(this=exp.DataType.Type.SERIAL))
130        elif kind.this == exp.DataType.Type.SMALLINT:
131            kind.replace(exp.DataType(this=exp.DataType.Type.SMALLSERIAL))
132        elif kind.this == exp.DataType.Type.BIGINT:
133            kind.replace(exp.DataType(this=exp.DataType.Type.BIGSERIAL))
134
135    return expression
136
137
138def _serial_to_generated(expression: exp.Expression) -> exp.Expression:
139    kind = expression.args.get("kind")
140    if not kind:
141        return expression
142
143    if kind.this == exp.DataType.Type.SERIAL:
144        data_type = exp.DataType(this=exp.DataType.Type.INT)
145    elif kind.this == exp.DataType.Type.SMALLSERIAL:
146        data_type = exp.DataType(this=exp.DataType.Type.SMALLINT)
147    elif kind.this == exp.DataType.Type.BIGSERIAL:
148        data_type = exp.DataType(this=exp.DataType.Type.BIGINT)
149    else:
150        data_type = None
151
152    if data_type:
153        expression.args["kind"].replace(data_type)
154        constraints = expression.args["constraints"]
155        generated = exp.ColumnConstraint(kind=exp.GeneratedAsIdentityColumnConstraint(this=False))
156        notnull = exp.ColumnConstraint(kind=exp.NotNullColumnConstraint())
157
158        if notnull not in constraints:
159            constraints.insert(0, notnull)
160        if generated not in constraints:
161            constraints.insert(0, generated)
162
163    return expression
164
165
166def _generate_series(args: t.List) -> exp.Expression:
167    # The goal is to convert step values like '1 day' or INTERVAL '1 day' into INTERVAL '1' day
168    step = seq_get(args, 2)
169
170    if step is None:
171        # Postgres allows calls with just two arguments -- the "step" argument defaults to 1
172        return exp.GenerateSeries.from_arg_list(args)
173
174    if step.is_string:
175        args[2] = exp.to_interval(step.this)
176    elif isinstance(step, exp.Interval) and not step.args.get("unit"):
177        args[2] = exp.to_interval(step.this.this)
178
179    return exp.GenerateSeries.from_arg_list(args)
180
181
182def _to_timestamp(args: t.List) -> exp.Expression:
183    # TO_TIMESTAMP accepts either a single double argument or (text, text)
184    if len(args) == 1:
185        # https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-TABLE
186        return exp.UnixToTime.from_arg_list(args)
187
188    # https://www.postgresql.org/docs/current/functions-formatting.html
189    return format_time_lambda(exp.StrToTime, "postgres")(args)
190
191
192class Postgres(Dialect):
193    INDEX_OFFSET = 1
194    TYPED_DIVISION = True
195    CONCAT_COALESCE = True
196    NULL_ORDERING = "nulls_are_large"
197    TIME_FORMAT = "'YYYY-MM-DD HH24:MI:SS'"
198
199    TIME_MAPPING = {
200        "AM": "%p",
201        "PM": "%p",
202        "D": "%u",  # 1-based day of week
203        "DD": "%d",  # day of month
204        "DDD": "%j",  # zero padded day of year
205        "FMDD": "%-d",  # - is no leading zero for Python; same for FM in postgres
206        "FMDDD": "%-j",  # day of year
207        "FMHH12": "%-I",  # 9
208        "FMHH24": "%-H",  # 9
209        "FMMI": "%-M",  # Minute
210        "FMMM": "%-m",  # 1
211        "FMSS": "%-S",  # Second
212        "HH12": "%I",  # 09
213        "HH24": "%H",  # 09
214        "MI": "%M",  # zero padded minute
215        "MM": "%m",  # 01
216        "OF": "%z",  # utc offset
217        "SS": "%S",  # zero padded second
218        "TMDay": "%A",  # TM is locale dependent
219        "TMDy": "%a",
220        "TMMon": "%b",  # Sep
221        "TMMonth": "%B",  # September
222        "TZ": "%Z",  # uppercase timezone name
223        "US": "%f",  # zero padded microsecond
224        "WW": "%U",  # 1-based week of year
225        "YY": "%y",  # 15
226        "YYYY": "%Y",  # 2015
227    }
228
229    class Tokenizer(tokens.Tokenizer):
230        BIT_STRINGS = [("b'", "'"), ("B'", "'")]
231        HEX_STRINGS = [("x'", "'"), ("X'", "'")]
232        BYTE_STRINGS = [("e'", "'"), ("E'", "'")]
233        HEREDOC_STRINGS = ["$"]
234
235        HEREDOC_TAG_IS_IDENTIFIER = True
236        HEREDOC_STRING_ALTERNATIVE = TokenType.PARAMETER
237
238        KEYWORDS = {
239            **tokens.Tokenizer.KEYWORDS,
240            "~~": TokenType.LIKE,
241            "~~*": TokenType.ILIKE,
242            "~*": TokenType.IRLIKE,
243            "~": TokenType.RLIKE,
244            "@@": TokenType.DAT,
245            "@>": TokenType.AT_GT,
246            "<@": TokenType.LT_AT,
247            "|/": TokenType.PIPE_SLASH,
248            "||/": TokenType.DPIPE_SLASH,
249            "BEGIN": TokenType.COMMAND,
250            "BEGIN TRANSACTION": TokenType.BEGIN,
251            "BIGSERIAL": TokenType.BIGSERIAL,
252            "CHARACTER VARYING": TokenType.VARCHAR,
253            "CONSTRAINT TRIGGER": TokenType.COMMAND,
254            "DECLARE": TokenType.COMMAND,
255            "DO": TokenType.COMMAND,
256            "EXEC": TokenType.COMMAND,
257            "HSTORE": TokenType.HSTORE,
258            "JSONB": TokenType.JSONB,
259            "MONEY": TokenType.MONEY,
260            "REFRESH": TokenType.COMMAND,
261            "REINDEX": TokenType.COMMAND,
262            "RESET": TokenType.COMMAND,
263            "REVOKE": TokenType.COMMAND,
264            "SERIAL": TokenType.SERIAL,
265            "SMALLSERIAL": TokenType.SMALLSERIAL,
266            "TEMP": TokenType.TEMPORARY,
267            "CSTRING": TokenType.PSEUDO_TYPE,
268            "OID": TokenType.OBJECT_IDENTIFIER,
269            "OPERATOR": TokenType.OPERATOR,
270            "REGCLASS": TokenType.OBJECT_IDENTIFIER,
271            "REGCOLLATION": TokenType.OBJECT_IDENTIFIER,
272            "REGCONFIG": TokenType.OBJECT_IDENTIFIER,
273            "REGDICTIONARY": TokenType.OBJECT_IDENTIFIER,
274            "REGNAMESPACE": TokenType.OBJECT_IDENTIFIER,
275            "REGOPER": TokenType.OBJECT_IDENTIFIER,
276            "REGOPERATOR": TokenType.OBJECT_IDENTIFIER,
277            "REGPROC": TokenType.OBJECT_IDENTIFIER,
278            "REGPROCEDURE": TokenType.OBJECT_IDENTIFIER,
279            "REGROLE": TokenType.OBJECT_IDENTIFIER,
280            "REGTYPE": TokenType.OBJECT_IDENTIFIER,
281        }
282
283        SINGLE_TOKENS = {
284            **tokens.Tokenizer.SINGLE_TOKENS,
285            "$": TokenType.HEREDOC_STRING,
286        }
287
288        VAR_SINGLE_TOKENS = {"$"}
289
290    class Parser(parser.Parser):
291        PROPERTY_PARSERS = {
292            **parser.Parser.PROPERTY_PARSERS,
293            "SET": lambda self: self.expression(exp.SetConfigProperty, this=self._parse_set()),
294        }
295        PROPERTY_PARSERS.pop("INPUT", None)
296
297        FUNCTIONS = {
298            **parser.Parser.FUNCTIONS,
299            "DATE_TRUNC": parse_timestamp_trunc,
300            "GENERATE_SERIES": _generate_series,
301            "JSON_EXTRACT_PATH": parse_json_extract_path(exp.JSONExtract),
302            "JSON_EXTRACT_PATH_TEXT": parse_json_extract_path(exp.JSONExtractScalar),
303            "MAKE_TIME": exp.TimeFromParts.from_arg_list,
304            "MAKE_TIMESTAMP": exp.TimestampFromParts.from_arg_list,
305            "NOW": exp.CurrentTimestamp.from_arg_list,
306            "TO_CHAR": format_time_lambda(exp.TimeToStr, "postgres"),
307            "TO_TIMESTAMP": _to_timestamp,
308            "UNNEST": exp.Explode.from_arg_list,
309        }
310
311        FUNCTION_PARSERS = {
312            **parser.Parser.FUNCTION_PARSERS,
313            "DATE_PART": lambda self: self._parse_date_part(),
314        }
315
316        BITWISE = {
317            **parser.Parser.BITWISE,
318            TokenType.HASH: exp.BitwiseXor,
319        }
320
321        EXPONENT = {
322            TokenType.CARET: exp.Pow,
323        }
324
325        RANGE_PARSERS = {
326            **parser.Parser.RANGE_PARSERS,
327            TokenType.AT_GT: binary_range_parser(exp.ArrayContains),
328            TokenType.DAMP: binary_range_parser(exp.ArrayOverlaps),
329            TokenType.DAT: lambda self, this: self.expression(
330                exp.MatchAgainst, this=self._parse_bitwise(), expressions=[this]
331            ),
332            TokenType.LT_AT: binary_range_parser(exp.ArrayContained),
333            TokenType.OPERATOR: lambda self, this: self._parse_operator(this),
334        }
335
336        STATEMENT_PARSERS = {
337            **parser.Parser.STATEMENT_PARSERS,
338            TokenType.END: lambda self: self._parse_commit_or_rollback(),
339        }
340
341        def _parse_operator(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.Expression]:
342            while True:
343                if not self._match(TokenType.L_PAREN):
344                    break
345
346                op = ""
347                while self._curr and not self._match(TokenType.R_PAREN):
348                    op += self._curr.text
349                    self._advance()
350
351                this = self.expression(
352                    exp.Operator,
353                    comments=self._prev_comments,
354                    this=this,
355                    operator=op,
356                    expression=self._parse_bitwise(),
357                )
358
359                if not self._match(TokenType.OPERATOR):
360                    break
361
362            return this
363
364        def _parse_date_part(self) -> exp.Expression:
365            part = self._parse_type()
366            self._match(TokenType.COMMA)
367            value = self._parse_bitwise()
368
369            if part and part.is_string:
370                part = exp.var(part.name)
371
372            return self.expression(exp.Extract, this=part, expression=value)
373
374    class Generator(generator.Generator):
375        SINGLE_STRING_INTERVAL = True
376        RENAME_TABLE_WITH_DB = False
377        LOCKING_READS_SUPPORTED = True
378        JOIN_HINTS = False
379        TABLE_HINTS = False
380        QUERY_HINTS = False
381        NVL2_SUPPORTED = False
382        PARAMETER_TOKEN = "$"
383        TABLESAMPLE_SIZE_IS_ROWS = False
384        TABLESAMPLE_SEED_KEYWORD = "REPEATABLE"
385        SUPPORTS_SELECT_INTO = True
386        JSON_TYPE_REQUIRED_FOR_EXTRACTION = True
387        SUPPORTS_UNLOGGED_TABLES = True
388        LIKE_PROPERTY_INSIDE_SCHEMA = True
389        MULTI_ARG_DISTINCT = False
390
391        SUPPORTED_JSON_PATH_PARTS = {
392            exp.JSONPathKey,
393            exp.JSONPathRoot,
394            exp.JSONPathSubscript,
395        }
396
397        TYPE_MAPPING = {
398            **generator.Generator.TYPE_MAPPING,
399            exp.DataType.Type.TINYINT: "SMALLINT",
400            exp.DataType.Type.FLOAT: "REAL",
401            exp.DataType.Type.DOUBLE: "DOUBLE PRECISION",
402            exp.DataType.Type.BINARY: "BYTEA",
403            exp.DataType.Type.VARBINARY: "BYTEA",
404            exp.DataType.Type.DATETIME: "TIMESTAMP",
405        }
406
407        TRANSFORMS = {
408            **generator.Generator.TRANSFORMS,
409            exp.AnyValue: any_value_to_max_sql,
410            exp.Array: lambda self, e: (
411                f"{self.normalize_func('ARRAY')}({self.sql(e.expressions[0])})"
412                if isinstance(seq_get(e.expressions, 0), exp.Select)
413                else f"{self.normalize_func('ARRAY')}[{self.expressions(e, flat=True)}]"
414            ),
415            exp.ArrayConcat: rename_func("ARRAY_CAT"),
416            exp.ArrayContained: lambda self, e: self.binary(e, "<@"),
417            exp.ArrayContains: lambda self, e: self.binary(e, "@>"),
418            exp.ArrayOverlaps: lambda self, e: self.binary(e, "&&"),
419            exp.BitwiseXor: lambda self, e: self.binary(e, "#"),
420            exp.ColumnDef: transforms.preprocess([_auto_increment_to_serial, _serial_to_generated]),
421            exp.CurrentDate: no_paren_current_date_sql,
422            exp.CurrentTimestamp: lambda *_: "CURRENT_TIMESTAMP",
423            exp.CurrentUser: lambda *_: "CURRENT_USER",
424            exp.DateAdd: _date_add_sql("+"),
425            exp.DateDiff: _date_diff_sql,
426            exp.DateStrToDate: datestrtodate_sql,
427            exp.DataType: _datatype_sql,
428            exp.DateSub: _date_add_sql("-"),
429            exp.Explode: rename_func("UNNEST"),
430            exp.GroupConcat: _string_agg_sql,
431            exp.JSONExtract: json_extract_segments("JSON_EXTRACT_PATH"),
432            exp.JSONExtractScalar: json_extract_segments("JSON_EXTRACT_PATH_TEXT"),
433            exp.JSONBExtract: lambda self, e: self.binary(e, "#>"),
434            exp.JSONBExtractScalar: lambda self, e: self.binary(e, "#>>"),
435            exp.JSONBContains: lambda self, e: self.binary(e, "?"),
436            exp.JSONPathKey: json_path_key_only_name,
437            exp.JSONPathRoot: lambda *_: "",
438            exp.JSONPathSubscript: lambda self, e: self.json_path_part(e.this),
439            exp.LastDay: no_last_day_sql,
440            exp.LogicalOr: rename_func("BOOL_OR"),
441            exp.LogicalAnd: rename_func("BOOL_AND"),
442            exp.Max: max_or_greatest,
443            exp.MapFromEntries: no_map_from_entries_sql,
444            exp.Min: min_or_least,
445            exp.Merge: merge_without_target_sql,
446            exp.PartitionedByProperty: lambda self, e: f"PARTITION BY {self.sql(e, 'this')}",
447            exp.PercentileCont: transforms.preprocess(
448                [transforms.add_within_group_for_percentiles]
449            ),
450            exp.PercentileDisc: transforms.preprocess(
451                [transforms.add_within_group_for_percentiles]
452            ),
453            exp.Pivot: no_pivot_sql,
454            exp.Pow: lambda self, e: self.binary(e, "^"),
455            exp.Rand: rename_func("RANDOM"),
456            exp.RegexpLike: lambda self, e: self.binary(e, "~"),
457            exp.RegexpILike: lambda self, e: self.binary(e, "~*"),
458            exp.Select: transforms.preprocess(
459                [
460                    transforms.eliminate_semi_and_anti_joins,
461                    transforms.eliminate_qualify,
462                ]
463            ),
464            exp.StrPosition: str_position_sql,
465            exp.StrToTime: lambda self,
466            e: f"TO_TIMESTAMP({self.sql(e, 'this')}, {self.format_time(e)})",
467            exp.StructExtract: struct_extract_sql,
468            exp.Substring: _substring_sql,
469            exp.TimeFromParts: rename_func("MAKE_TIME"),
470            exp.TimestampFromParts: rename_func("MAKE_TIMESTAMP"),
471            exp.TimestampTrunc: timestamptrunc_sql,
472            exp.TimeStrToTime: timestrtotime_sql,
473            exp.TimeToStr: lambda self, e: f"TO_CHAR({self.sql(e, 'this')}, {self.format_time(e)})",
474            exp.ToChar: lambda self, e: self.function_fallback_sql(e),
475            exp.Trim: trim_sql,
476            exp.TryCast: no_trycast_sql,
477            exp.TsOrDsAdd: _date_add_sql("+"),
478            exp.TsOrDsDiff: _date_diff_sql,
479            exp.UnixToTime: lambda self, e: f"TO_TIMESTAMP({self.sql(e, 'this')})",
480            exp.VariancePop: rename_func("VAR_POP"),
481            exp.Variance: rename_func("VAR_SAMP"),
482            exp.Xor: bool_xor_sql,
483        }
484
485        PROPERTIES_LOCATION = {
486            **generator.Generator.PROPERTIES_LOCATION,
487            exp.PartitionedByProperty: exp.Properties.Location.POST_SCHEMA,
488            exp.TransientProperty: exp.Properties.Location.UNSUPPORTED,
489            exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
490        }
491
492        def bracket_sql(self, expression: exp.Bracket) -> str:
493            """Forms like ARRAY[1, 2, 3][3] aren't allowed; we need to wrap the ARRAY."""
494            if isinstance(expression.this, exp.Array):
495                expression.set("this", exp.paren(expression.this, copy=False))
496
497            return super().bracket_sql(expression)
498
499        def matchagainst_sql(self, expression: exp.MatchAgainst) -> str:
500            this = self.sql(expression, "this")
501            expressions = [f"{self.sql(e)} @@ {this}" for e in expression.expressions]
502            sql = " OR ".join(expressions)
503            return f"({sql})" if len(expressions) > 1 else sql