Merging upstream version 11.3.3.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 24 insertions, 9 deletions

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 79501ef..4504e95 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -112,6 +112,9 @@ class Generator:
     # Whether or not to treat the division operator "/" as integer division
     INTEGER_DIVISION = True
 
+    # Whether or not MERGE ... WHEN MATCHED BY SOURCE is allowed
+    MATCHED_BY_SOURCE = True
+
     TYPE_MAPPING = {
         exp.DataType.Type.NCHAR: "CHAR",
         exp.DataType.Type.NVARCHAR: "VARCHAR",
@@ -688,7 +691,8 @@ class Generator:
             else ""
         )
         where_sql = self.sql(expression, "where")
-        sql = f"DELETE{this}{using_sql}{where_sql}"
+        returning = self.sql(expression, "returning")
+        sql = f"DELETE{this}{using_sql}{where_sql}{returning}"
         return self.prepend_ctes(expression, sql)
 
     def drop_sql(self, expression: exp.Drop) -> str:
@@ -952,8 +956,9 @@ class Generator:
             self.sql(expression, "partition") if expression.args.get("partition") else ""
         )
         expression_sql = self.sql(expression, "expression")
+        returning = self.sql(expression, "returning")
         sep = self.sep() if partition_sql else ""
-        sql = f"INSERT{alternative}{this}{exists}{partition_sql}{sep}{expression_sql}"
+        sql = f"INSERT{alternative}{this}{exists}{partition_sql}{sep}{expression_sql}{returning}"
         return self.prepend_ctes(expression, sql)
 
     def intersect_sql(self, expression: exp.Intersect) -> str:
@@ -971,6 +976,9 @@ class Generator:
     def pseudotype_sql(self, expression: exp.PseudoType) -> str:
         return expression.name.upper()
 
+    def returning_sql(self, expression: exp.Returning) -> str:
+        return f"{self.seg('RETURNING')} {self.expressions(expression, flat=True)}"
+
     def rowformatdelimitedproperty_sql(self, expression: exp.RowFormatDelimitedProperty) -> str:
         fields = expression.args.get("fields")
         fields = f" FIELDS TERMINATED BY {fields}" if fields else ""
@@ -1009,7 +1017,7 @@ class Generator:
 
         return f"{table}{system_time}{alias}{hints}{laterals}{joins}{pivots}"
 
-    def tablesample_sql(self, expression: exp.TableSample) -> str:
+    def tablesample_sql(self, expression: exp.TableSample, seed_prefix: str = "SEED") -> str:
         if self.alias_post_tablesample and expression.this.alias:
             this = self.sql(expression.this, "this")
             alias = f" AS {self.sql(expression.this, 'alias')}"
@@ -1017,7 +1025,7 @@ class Generator:
             this = self.sql(expression, "this")
             alias = ""
         method = self.sql(expression, "method")
-        method = f" {method.upper()} " if method else ""
+        method = f"{method.upper()} " if method else ""
         numerator = self.sql(expression, "bucket_numerator")
         denominator = self.sql(expression, "bucket_denominator")
         field = self.sql(expression, "bucket_field")
@@ -1029,8 +1037,9 @@ class Generator:
         rows = f"{rows} ROWS" if rows else ""
         size = self.sql(expression, "size")
         seed = self.sql(expression, "seed")
-        seed = f" SEED ({seed})" if seed else ""
-        return f"{this} TABLESAMPLE{method}({bucket}{percent}{rows}{size}){seed}{alias}"
+        seed = f" {seed_prefix} ({seed})" if seed else ""
+        kind = expression.args.get("kind", "TABLESAMPLE")
+        return f"{this} {kind} {method}({bucket}{percent}{rows}{size}){seed}{alias}"
 
     def pivot_sql(self, expression: exp.Pivot) -> str:
         this = self.sql(expression, "this")
@@ -1050,7 +1059,8 @@ class Generator:
         set_sql = self.expressions(expression, flat=True)
         from_sql = self.sql(expression, "from")
         where_sql = self.sql(expression, "where")
-        sql = f"UPDATE {this} SET {set_sql}{from_sql}{where_sql}"
+        returning = self.sql(expression, "returning")
+        sql = f"UPDATE {this} SET {set_sql}{from_sql}{where_sql}{returning}"
         return self.prepend_ctes(expression, sql)
 
     def values_sql(self, expression: exp.Values) -> str:
@@ -1297,6 +1307,7 @@ class Generator:
             self.sql(expression, "limit"),
             self.sql(expression, "offset"),
             self.sql(expression, "lock"),
+            self.sql(expression, "sample"),
             sep="",
         )
 
@@ -1956,7 +1967,11 @@ class Generator:
         return self.binary(expression, "=>")
 
     def when_sql(self, expression: exp.When) -> str:
-        this = self.sql(expression, "this")
+        matched = "MATCHED" if expression.args["matched"] else "NOT MATCHED"
+        source = " BY SOURCE" if self.MATCHED_BY_SOURCE and expression.args.get("source") else ""
+        condition = self.sql(expression, "condition")
+        condition = f" AND {condition}" if condition else ""
+
         then_expression = expression.args.get("then")
         if isinstance(then_expression, exp.Insert):
             then = f"INSERT {self.sql(then_expression, 'this')}"
@@ -1969,7 +1984,7 @@ class Generator:
                 then = f"UPDATE SET {self.expressions(then_expression, flat=True)}"
         else:
             then = self.sql(then_expression)
-        return f"WHEN {this} THEN {then}"
+        return f"WHEN {matched}{source}{condition} THEN {then}"
 
     def merge_sql(self, expression: exp.Merge) -> str:
         this = self.sql(expression, "this")