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