Repository: tobymao/sqlglot

Issue:
tsql dialect replaces "TOP N" with "FETCH FIRST N ROWS ONLY", which is not supported by SQL Server (MSSQL)
import sqlglot
from sqlglot.optimizer import optimize

if __name__ == '__main__':
  sqlglot_metadata = {
    'sales': {'ReceiptDate': 'DATE', 'Qtty': 'FLOAT', 'Price': 'FLOAT'}
  }

  print(sqlglot.__version__) # 16.7.7

  experiments = [
    {
      'sql': """SELECT TOP 3 ReceiptDate, AVG(Qtty * Price) AS AverageSale FROM sales WHERE ReceiptDate >= '2023-01-01' GROUP BY ReceiptDate""",
      'read': 'tsql',
      'dialect': 'tsql'
    }
  ]

  for exp in experiments:
    expression = optimize(
      expression=sqlglot.parse_one(sql=exp['sql'], read=exp['read']),
      schema=sqlglot_metadata,
      rules=optimize_rules,
    )

    print(expression.sql(dialect=exp['dialect']))


Discussion:

Resolving patch:
diff --git a/sqlglot/dialects/tsql.py b/sqlglot/dialects/tsql.py
--- a/sqlglot/dialects/tsql.py
+++ b/sqlglot/dialects/tsql.py
@@ -467,6 +467,7 @@ def _parse_user_defined_function(
 
     class Generator(generator.Generator):
         LOCKING_READS_SUPPORTED = True
+        LIMIT_IS_TOP = True
 
         TYPE_MAPPING = {
             **generator.Generator.TYPE_MAPPING,
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -146,6 +146,9 @@ class Generator:
     # 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")
 
@@ -1385,7 +1388,7 @@ def lateral_sql(self, expression: exp.Lateral) -> str:
         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
@@ -1395,7 +1398,7 @@ def limit_sql(self, expression: exp.Limit) -> str:
             )
             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")
@@ -1552,6 +1555,9 @@ def matchrecognize_sql(self, expression: exp.MatchRecognize) -> str:
     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):
@@ -1559,6 +1565,12 @@ def query_modifiers(self, expression: exp.Expression, *sqls: str) -> str:
 
         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 []],
@@ -1569,7 +1581,7 @@ def query_modifiers(self, expression: exp.Expression, *sqls: str) -> str:
             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="",
         )
@@ -1600,6 +1612,13 @@ def select_sql(self, expression: exp.Select) -> str:
         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:
@@ -1626,7 +1645,7 @@ def select_sql(self, expression: exp.Select) -> str:
         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),
         )


Test patch:
diff --git a/tests/dialects/test_dialect.py b/tests/dialects/test_dialect.py
--- a/tests/dialects/test_dialect.py
+++ b/tests/dialects/test_dialect.py
@@ -1307,7 +1307,7 @@ def test_limit(self):
             write={
                 "sqlite": "SELECT x FROM y LIMIT 10",
                 "oracle": "SELECT x FROM y FETCH FIRST 10 ROWS ONLY",
-                "tsql": "SELECT x FROM y FETCH FIRST 10 ROWS ONLY",
+                "tsql": "SELECT TOP 10 x FROM y",
             },
         )
         self.validate_all(