diff options
Diffstat (limited to 'sqlglot/generator.py')
| -rw-r--r-- | sqlglot/generator.py | 91 |
1 files changed, 66 insertions, 25 deletions
diff --git a/sqlglot/generator.py b/sqlglot/generator.py index a41af12..1ce2aaa 100644 --- a/sqlglot/generator.py +++ b/sqlglot/generator.py @@ -140,9 +140,21 @@ class Generator: # Whether or not table hints should be generated TABLE_HINTS = True + # Whether or not query hints should be generated + QUERY_HINTS = True + + # What kind of separator to use for query hints + QUERY_HINT_SEP = ", " + # Whether or not comparing against booleans (e.g. x IS TRUE) is supported IS_BOOL_ALLOWED = True + # Whether or not to include the "SET" keyword in the "INSERT ... ON DUPLICATE KEY UPDATE" statement + DUPLICATE_KEY_UPDATE_WITH_SET = True + + # Whether or not to generate the limit as TOP <value> instead of LIMIT <value> + LIMIT_IS_TOP = False + # https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax SELECT_KINDS: t.Tuple[str, ...] = ("STRUCT", "VALUE") @@ -268,6 +280,7 @@ class Generator: STRICT_STRING_CONCAT = False NORMALIZE_FUNCTIONS: bool | str = "upper" NULL_ORDERING = "nulls_are_small" + ESCAPE_LINE_BREAK = False can_identify: t.Callable[[str, str | bool], bool] @@ -286,8 +299,6 @@ class Generator: HEX_END: t.Optional[str] = None BYTE_START: t.Optional[str] = None BYTE_END: t.Optional[str] = None - RAW_START: t.Optional[str] = None - RAW_END: t.Optional[str] = None __slots__ = ( "pretty", @@ -486,7 +497,10 @@ class Generator: return expression if key: - return self.sql(expression.args.get(key)) + value = expression.args.get(key) + if value: + return self.sql(value) + return "" if self._cache is not None: expression_id = hash(expression) @@ -779,10 +793,7 @@ class Generator: return this def rawstring_sql(self, expression: exp.RawString) -> str: - string = expression.this - if self.RAW_START: - return f"{self.RAW_START}{self.escape_str(expression.this)}{self.RAW_END}" - string = self.escape_str(string.replace("\\", "\\\\")) + string = self.escape_str(expression.this.replace("\\", "\\\\")) return f"{self.QUOTE_START}{string}{self.QUOTE_END}" def datatypesize_sql(self, expression: exp.DataTypeSize) -> str: @@ -818,15 +829,14 @@ class Generator: def delete_sql(self, expression: exp.Delete) -> str: this = self.sql(expression, "this") this = f" FROM {this}" if this else "" - using_sql = ( - f" USING {self.expressions(expression, key='using', sep=', USING ')}" - if expression.args.get("using") - else "" - ) - where_sql = self.sql(expression, "where") + using = self.sql(expression, "using") + using = f" USING {using}" if using else "" + where = self.sql(expression, "where") returning = self.sql(expression, "returning") limit = self.sql(expression, "limit") - sql = f"DELETE{this}{using_sql}{where_sql}{returning}{limit}" + tables = self.expressions(expression, key="tables") + tables = f" {tables}" if tables else "" + sql = f"DELETE{tables}{this}{using}{where}{returning}{limit}" return self.prepend_ctes(expression, sql) def drop_sql(self, expression: exp.Drop) -> str: @@ -867,9 +877,11 @@ class Generator: return f"{this} FILTER({where})" def hint_sql(self, expression: exp.Hint) -> str: - if self.sql(expression, "this"): + if not self.QUERY_HINTS: self.unsupported("Hints are not supported") - return "" + return "" + + return f" /*+ {self.expressions(expression, sep=self.QUERY_HINT_SEP).strip()} */" def index_sql(self, expression: exp.Index) -> str: unique = "UNIQUE " if expression.args.get("unique") else "" @@ -1109,6 +1121,8 @@ class Generator: alternative = expression.args.get("alternative") alternative = f" OR {alternative}" if alternative else "" + ignore = " IGNORE" if expression.args.get("ignore") else "" + this = f"{this} {self.sql(expression, 'this')}" exists = " IF EXISTS" if expression.args.get("exists") else "" @@ -1120,7 +1134,7 @@ class Generator: expression_sql = f"{self.sep()}{self.sql(expression, 'expression')}" conflict = self.sql(expression, "conflict") returning = self.sql(expression, "returning") - sql = f"INSERT{alternative}{this}{exists}{partition_sql}{where}{expression_sql}{conflict}{returning}" + sql = f"INSERT{alternative}{ignore}{this}{exists}{partition_sql}{where}{expression_sql}{conflict}{returning}" return self.prepend_ctes(expression, sql) def intersect_sql(self, expression: exp.Intersect) -> str: @@ -1147,8 +1161,9 @@ class Generator: do = "" if expression.args.get("duplicate") else " DO " nothing = "NOTHING" if expression.args.get("nothing") else "" expressions = self.expressions(expression, flat=True) + set_keyword = "SET " if self.DUPLICATE_KEY_UPDATE_WITH_SET else "" if expressions: - expressions = f"UPDATE SET {expressions}" + expressions = f"UPDATE {set_keyword}{expressions}" return f"{self.seg(conflict)} {constraint}{key}{do}{nothing}{expressions}" def returning_sql(self, expression: exp.Returning) -> str: @@ -1195,7 +1210,7 @@ class Generator: hints = f" {hints}" if hints and self.TABLE_HINTS else "" pivots = self.expressions(expression, key="pivots", sep=" ", flat=True) pivots = f" {pivots}" if pivots else "" - joins = self.expressions(expression, key="joins", sep="") + joins = self.expressions(expression, key="joins", sep="", skip_first=True) laterals = self.expressions(expression, key="laterals", sep="") system_time = expression.args.get("system_time") system_time = f" {self.sql(expression, 'system_time')}" if system_time else "" @@ -1287,6 +1302,10 @@ class Generator: def group_sql(self, expression: exp.Group) -> str: group_by = self.op_expressions("GROUP BY", expression) + + if expression.args.get("all"): + return f"{group_by} ALL" + grouping_sets = self.expressions(expression, key="grouping_sets", indent=False) grouping_sets = ( f"{self.seg('GROUPING SETS')} {self.wrap(grouping_sets)}" if grouping_sets else "" @@ -1379,7 +1398,7 @@ class Generator: alias = f" AS {alias}" if alias else "" return f"LATERAL {this}{alias}" - def limit_sql(self, expression: exp.Limit) -> str: + def limit_sql(self, expression: exp.Limit, top: bool = False) -> str: this = self.sql(expression, "this") args = ", ".join( sql @@ -1389,7 +1408,7 @@ class Generator: ) if sql ) - return f"{this}{self.seg('LIMIT')} {args}" + return f"{this}{self.seg('TOP' if top else 'LIMIT')} {args}" def offset_sql(self, expression: exp.Offset) -> str: this = self.sql(expression, "this") @@ -1441,7 +1460,9 @@ class Generator: def escape_str(self, text: str) -> str: text = text.replace(self.QUOTE_END, self._escaped_quote_end) - if self.pretty: + if self.ESCAPE_LINE_BREAK: + text = text.replace("\n", "\\n") + elif self.pretty: text = text.replace("\n", self.SENTINEL_LINE_BREAK) return text @@ -1544,6 +1565,9 @@ class Generator: def query_modifiers(self, expression: exp.Expression, *sqls: str) -> str: limit: t.Optional[exp.Fetch | exp.Limit] = expression.args.get("limit") + # If the limit is generated as TOP, we need to ensure it's not generated twice + with_offset_limit_modifiers = not isinstance(limit, exp.Limit) or not self.LIMIT_IS_TOP + if self.LIMIT_FETCH == "LIMIT" and isinstance(limit, exp.Fetch): limit = exp.Limit(expression=limit.args.get("count")) elif self.LIMIT_FETCH == "FETCH" and isinstance(limit, exp.Limit): @@ -1551,6 +1575,12 @@ class Generator: fetch = isinstance(limit, exp.Fetch) + offset_limit_modifiers = ( + self.offset_limit_modifiers(expression, fetch, limit) + if with_offset_limit_modifiers + else [] + ) + return csv( *sqls, *[self.sql(join) for join in expression.args.get("joins") or []], @@ -1561,7 +1591,7 @@ class Generator: self.sql(expression, "having"), *self.after_having_modifiers(expression), self.sql(expression, "order"), - *self.offset_limit_modifiers(expression, fetch, limit), + *offset_limit_modifiers, *self.after_limit_modifiers(expression), sep="", ) @@ -1580,6 +1610,9 @@ class Generator: self.seg("WINDOW ") + self.expressions(expression, key="windows", flat=True) if expression.args.get("windows") else "", + self.sql(expression, "distribute"), + self.sql(expression, "sort"), + self.sql(expression, "cluster"), ] def after_limit_modifiers(self, expression: exp.Expression) -> t.List[str]: @@ -1592,6 +1625,13 @@ class Generator: distinct = self.sql(expression, "distinct") distinct = f" {distinct}" if distinct else "" kind = self.sql(expression, "kind").upper() + limit = expression.args.get("limit") + top = ( + self.limit_sql(limit, top=True) + if isinstance(limit, exp.Limit) and self.LIMIT_IS_TOP + else "" + ) + expressions = self.expressions(expression) if kind: @@ -1618,7 +1658,7 @@ class Generator: expressions = f"{self.sep()}{expressions}" if expressions else expressions sql = self.query_modifiers( expression, - f"SELECT{hint}{distinct}{kind}{expressions}", + f"SELECT{top}{hint}{distinct}{kind}{expressions}", self.sql(expression, "into", comment=False), self.sql(expression, "from", comment=False), ) @@ -2288,6 +2328,7 @@ class Generator: sqls: t.Optional[t.List[str]] = None, flat: bool = False, indent: bool = True, + skip_first: bool = False, sep: str = ", ", prefix: str = "", ) -> str: @@ -2321,7 +2362,7 @@ class Generator: result_sqls.append(f"{prefix}{sql}{comments}{sep if i + 1 < num_sqls else ''}") result_sql = "\n".join(result_sqls) if self.pretty else "".join(result_sqls) - return self.indent(result_sql, skip_first=False) if indent else result_sql + return self.indent(result_sql, skip_first=skip_first) if indent else result_sql def op_expressions(self, op: str, expression: exp.Expression, flat: bool = False) -> str: flat = flat or isinstance(expression.parent, exp.Properties) |