sqlglot.dialects.postgres

  1from __future__ import annotations
  2
  3from sqlglot import exp, generator, parser, tokens
  4from sqlglot.dialects.dialect import (
  5    Dialect,
  6    arrow_json_extract_scalar_sql,
  7    arrow_json_extract_sql,
  8    format_time_lambda,
  9    min_or_least,
 10    no_paren_current_date_sql,
 11    no_tablesample_sql,
 12    no_trycast_sql,
 13    rename_func,
 14    str_position_sql,
 15    trim_sql,
 16)
 17from sqlglot.helper import seq_get
 18from sqlglot.tokens import TokenType
 19from sqlglot.transforms import delegate, preprocess
 20
 21DATE_DIFF_FACTOR = {
 22    "MICROSECOND": " * 1000000",
 23    "MILLISECOND": " * 1000",
 24    "SECOND": "",
 25    "MINUTE": " / 60",
 26    "HOUR": " / 3600",
 27    "DAY": " / 86400",
 28}
 29
 30
 31def _date_add_sql(kind):
 32    def func(self, expression):
 33        from sqlglot.optimizer.simplify import simplify
 34
 35        this = self.sql(expression, "this")
 36        unit = self.sql(expression, "unit")
 37        expression = simplify(expression.args["expression"])
 38
 39        if not isinstance(expression, exp.Literal):
 40            self.unsupported("Cannot add non literal")
 41
 42        expression = expression.copy()
 43        expression.args["is_string"] = True
 44        return f"{this} {kind} {self.sql(exp.Interval(this=expression, unit=unit))}"
 45
 46    return func
 47
 48
 49def _date_diff_sql(self, expression):
 50    unit = expression.text("unit").upper()
 51    factor = DATE_DIFF_FACTOR.get(unit)
 52
 53    end = f"CAST({expression.this} AS TIMESTAMP)"
 54    start = f"CAST({expression.expression} AS TIMESTAMP)"
 55
 56    if factor is not None:
 57        return f"CAST(EXTRACT(epoch FROM {end} - {start}){factor} AS BIGINT)"
 58
 59    age = f"AGE({end}, {start})"
 60
 61    if unit == "WEEK":
 62        unit = f"EXTRACT(year FROM {age}) * 48 + EXTRACT(month FROM {age}) * 4 + EXTRACT(day FROM {age}) / 7"
 63    elif unit == "MONTH":
 64        unit = f"EXTRACT(year FROM {age}) * 12 + EXTRACT(month FROM {age})"
 65    elif unit == "QUARTER":
 66        unit = f"EXTRACT(year FROM {age}) * 4 + EXTRACT(month FROM {age}) / 3"
 67    elif unit == "YEAR":
 68        unit = f"EXTRACT(year FROM {age})"
 69    else:
 70        unit = age
 71
 72    return f"CAST({unit} AS BIGINT)"
 73
 74
 75def _substring_sql(self, expression):
 76    this = self.sql(expression, "this")
 77    start = self.sql(expression, "start")
 78    length = self.sql(expression, "length")
 79
 80    from_part = f" FROM {start}" if start else ""
 81    for_part = f" FOR {length}" if length else ""
 82
 83    return f"SUBSTRING({this}{from_part}{for_part})"
 84
 85
 86def _string_agg_sql(self, expression):
 87    expression = expression.copy()
 88    separator = expression.args.get("separator") or exp.Literal.string(",")
 89
 90    order = ""
 91    this = expression.this
 92    if isinstance(this, exp.Order):
 93        if this.this:
 94            this = this.this
 95            this.pop()
 96        order = self.sql(expression.this)  # Order has a leading space
 97
 98    return f"STRING_AGG({self.format_args(this, separator)}{order})"
 99
100
101def _datatype_sql(self, expression):
102    if expression.this == exp.DataType.Type.ARRAY:
103        return f"{self.expressions(expression, flat=True)}[]"
104    return self.datatype_sql(expression)
105
106
107def _auto_increment_to_serial(expression):
108    auto = expression.find(exp.AutoIncrementColumnConstraint)
109
110    if auto:
111        expression = expression.copy()
112        expression.args["constraints"].remove(auto.parent)
113        kind = expression.args["kind"]
114
115        if kind.this == exp.DataType.Type.INT:
116            kind.replace(exp.DataType(this=exp.DataType.Type.SERIAL))
117        elif kind.this == exp.DataType.Type.SMALLINT:
118            kind.replace(exp.DataType(this=exp.DataType.Type.SMALLSERIAL))
119        elif kind.this == exp.DataType.Type.BIGINT:
120            kind.replace(exp.DataType(this=exp.DataType.Type.BIGSERIAL))
121
122    return expression
123
124
125def _serial_to_generated(expression):
126    kind = expression.args["kind"]
127
128    if kind.this == exp.DataType.Type.SERIAL:
129        data_type = exp.DataType(this=exp.DataType.Type.INT)
130    elif kind.this == exp.DataType.Type.SMALLSERIAL:
131        data_type = exp.DataType(this=exp.DataType.Type.SMALLINT)
132    elif kind.this == exp.DataType.Type.BIGSERIAL:
133        data_type = exp.DataType(this=exp.DataType.Type.BIGINT)
134    else:
135        data_type = None
136
137    if data_type:
138        expression = expression.copy()
139        expression.args["kind"].replace(data_type)
140        constraints = expression.args["constraints"]
141        generated = exp.ColumnConstraint(kind=exp.GeneratedAsIdentityColumnConstraint(this=False))
142        notnull = exp.ColumnConstraint(kind=exp.NotNullColumnConstraint())
143        if notnull not in constraints:
144            constraints.insert(0, notnull)
145        if generated not in constraints:
146            constraints.insert(0, generated)
147
148    return expression
149
150
151def _generate_series(args):
152    # The goal is to convert step values like '1 day' or INTERVAL '1 day' into INTERVAL '1' day
153    step = seq_get(args, 2)
154
155    if step is None:
156        # Postgres allows calls with just two arguments -- the "step" argument defaults to 1
157        return exp.GenerateSeries.from_arg_list(args)
158
159    if step.is_string:
160        args[2] = exp.to_interval(step.this)
161    elif isinstance(step, exp.Interval) and not step.args.get("unit"):
162        args[2] = exp.to_interval(step.this.this)
163
164    return exp.GenerateSeries.from_arg_list(args)
165
166
167def _to_timestamp(args):
168    # TO_TIMESTAMP accepts either a single double argument or (text, text)
169    if len(args) == 1:
170        # https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-TABLE
171        return exp.UnixToTime.from_arg_list(args)
172    # https://www.postgresql.org/docs/current/functions-formatting.html
173    return format_time_lambda(exp.StrToTime, "postgres")(args)
174
175
176class Postgres(Dialect):
177    null_ordering = "nulls_are_large"
178    time_format = "'YYYY-MM-DD HH24:MI:SS'"
179    time_mapping = {
180        "AM": "%p",
181        "PM": "%p",
182        "D": "%u",  # 1-based day of week
183        "DD": "%d",  # day of month
184        "DDD": "%j",  # zero padded day of year
185        "FMDD": "%-d",  # - is no leading zero for Python; same for FM in postgres
186        "FMDDD": "%-j",  # day of year
187        "FMHH12": "%-I",  # 9
188        "FMHH24": "%-H",  # 9
189        "FMMI": "%-M",  # Minute
190        "FMMM": "%-m",  # 1
191        "FMSS": "%-S",  # Second
192        "HH12": "%I",  # 09
193        "HH24": "%H",  # 09
194        "MI": "%M",  # zero padded minute
195        "MM": "%m",  # 01
196        "OF": "%z",  # utc offset
197        "SS": "%S",  # zero padded second
198        "TMDay": "%A",  # TM is locale dependent
199        "TMDy": "%a",
200        "TMMon": "%b",  # Sep
201        "TMMonth": "%B",  # September
202        "TZ": "%Z",  # uppercase timezone name
203        "US": "%f",  # zero padded microsecond
204        "WW": "%U",  # 1-based week of year
205        "YY": "%y",  # 15
206        "YYYY": "%Y",  # 2015
207    }
208
209    class Tokenizer(tokens.Tokenizer):
210        QUOTES = ["'", "$$"]
211
212        BIT_STRINGS = [("b'", "'"), ("B'", "'")]
213        HEX_STRINGS = [("x'", "'"), ("X'", "'")]
214        BYTE_STRINGS = [("e'", "'"), ("E'", "'")]
215
216        KEYWORDS = {
217            **tokens.Tokenizer.KEYWORDS,
218            "~~": TokenType.LIKE,
219            "~~*": TokenType.ILIKE,
220            "~*": TokenType.IRLIKE,
221            "~": TokenType.RLIKE,
222            "BEGIN": TokenType.COMMAND,
223            "BEGIN TRANSACTION": TokenType.BEGIN,
224            "BIGSERIAL": TokenType.BIGSERIAL,
225            "CHARACTER VARYING": TokenType.VARCHAR,
226            "DECLARE": TokenType.COMMAND,
227            "DO": TokenType.COMMAND,
228            "HSTORE": TokenType.HSTORE,
229            "JSONB": TokenType.JSONB,
230            "REFRESH": TokenType.COMMAND,
231            "REINDEX": TokenType.COMMAND,
232            "RESET": TokenType.COMMAND,
233            "RETURNING": TokenType.RETURNING,
234            "REVOKE": TokenType.COMMAND,
235            "SERIAL": TokenType.SERIAL,
236            "SMALLSERIAL": TokenType.SMALLSERIAL,
237            "TEMP": TokenType.TEMPORARY,
238            "UUID": TokenType.UUID,
239            "CSTRING": TokenType.PSEUDO_TYPE,
240        }
241
242        SINGLE_TOKENS = {
243            **tokens.Tokenizer.SINGLE_TOKENS,
244            "$": TokenType.PARAMETER,
245        }
246
247    class Parser(parser.Parser):
248        STRICT_CAST = False
249
250        FUNCTIONS = {
251            **parser.Parser.FUNCTIONS,  # type: ignore
252            "NOW": exp.CurrentTimestamp.from_arg_list,
253            "TO_TIMESTAMP": _to_timestamp,
254            "TO_CHAR": format_time_lambda(exp.TimeToStr, "postgres"),
255            "GENERATE_SERIES": _generate_series,
256        }
257
258        BITWISE = {
259            **parser.Parser.BITWISE,  # type: ignore
260            TokenType.HASH: exp.BitwiseXor,
261        }
262
263        FACTOR = {**parser.Parser.FACTOR, TokenType.CARET: exp.Pow}
264
265    class Generator(generator.Generator):
266        LOCKING_READS_SUPPORTED = True
267        PARAMETER_TOKEN = "$"
268
269        TYPE_MAPPING = {
270            **generator.Generator.TYPE_MAPPING,  # type: ignore
271            exp.DataType.Type.TINYINT: "SMALLINT",
272            exp.DataType.Type.FLOAT: "REAL",
273            exp.DataType.Type.DOUBLE: "DOUBLE PRECISION",
274            exp.DataType.Type.BINARY: "BYTEA",
275            exp.DataType.Type.VARBINARY: "BYTEA",
276            exp.DataType.Type.DATETIME: "TIMESTAMP",
277        }
278
279        TRANSFORMS = {
280            **generator.Generator.TRANSFORMS,  # type: ignore
281            exp.BitwiseXor: lambda self, e: self.binary(e, "#"),
282            exp.ColumnDef: preprocess(
283                [
284                    _auto_increment_to_serial,
285                    _serial_to_generated,
286                ],
287                delegate("columndef_sql"),
288            ),
289            exp.JSONExtract: arrow_json_extract_sql,
290            exp.JSONExtractScalar: arrow_json_extract_scalar_sql,
291            exp.JSONBExtract: lambda self, e: self.binary(e, "#>"),
292            exp.JSONBExtractScalar: lambda self, e: self.binary(e, "#>>"),
293            exp.JSONBContains: lambda self, e: self.binary(e, "?"),
294            exp.Pow: lambda self, e: self.binary(e, "^"),
295            exp.CurrentDate: no_paren_current_date_sql,
296            exp.CurrentTimestamp: lambda *_: "CURRENT_TIMESTAMP",
297            exp.DateAdd: _date_add_sql("+"),
298            exp.DateSub: _date_add_sql("-"),
299            exp.DateDiff: _date_diff_sql,
300            exp.LogicalOr: rename_func("BOOL_OR"),
301            exp.Min: min_or_least,
302            exp.RegexpLike: lambda self, e: self.binary(e, "~"),
303            exp.RegexpILike: lambda self, e: self.binary(e, "~*"),
304            exp.StrPosition: str_position_sql,
305            exp.StrToTime: lambda self, e: f"TO_TIMESTAMP({self.sql(e, 'this')}, {self.format_time(e)})",
306            exp.Substring: _substring_sql,
307            exp.TimeStrToTime: lambda self, e: f"CAST({self.sql(e, 'this')} AS TIMESTAMP)",
308            exp.TimeToStr: lambda self, e: f"TO_CHAR({self.sql(e, 'this')}, {self.format_time(e)})",
309            exp.TableSample: no_tablesample_sql,
310            exp.Trim: trim_sql,
311            exp.TryCast: no_trycast_sql,
312            exp.UnixToTime: lambda self, e: f"TO_TIMESTAMP({self.sql(e, 'this')})",
313            exp.DataType: _datatype_sql,
314            exp.GroupConcat: _string_agg_sql,
315            exp.Array: lambda self, e: f"{self.normalize_func('ARRAY')}({self.sql(e.expressions[0])})"
316            if isinstance(seq_get(e.expressions, 0), exp.Select)
317            else f"{self.normalize_func('ARRAY')}[{self.expressions(e, flat=True)}]",
318        }
319
320        PROPERTIES_LOCATION = {
321            **generator.Generator.PROPERTIES_LOCATION,  # type: ignore
322            exp.TransientProperty: exp.Properties.Location.UNSUPPORTED,
323        }