feat: use a separate render function for limited query

Author: pik94

data_diff/databases/base.py

@@ -202,7 +202,6 @@ class BaseDialect(abc.ABC):
     TYPE_CLASSES: ClassVar[Dict[str, Type[ColType]]] = {}
 
     PLACEHOLDER_TABLE = None  # Used for Oracle
-    USE_TOP_INSTEAD_LIMIT: bool = False  # True for MsSQL or Teradata
 
     def parse_table_name(self, name: str) -> DbPath:
         "Parse the given table name into a DbPath"
@@ -472,10 +471,7 @@ def render_select(self, parent_c: Compiler, elem: Select) -> str:
         columns = ", ".join(map(compile_fn, elem.columns)) if elem.columns else "*"
         distinct = "DISTINCT " if elem.distinct else ""
         optimizer_hints = self.optimizer_hints(elem.optimizer_hints) if elem.optimizer_hints else ""
-        if elem.limit_expr is not None and self.USE_TOP_INSTEAD_LIMIT:
-            select = f"SELECT TOP {elem.limit_expr} {optimizer_hints}{distinct}{columns}"
-        else:
-            select = f"SELECT {optimizer_hints}{distinct}{columns}"
+        select = f"SELECT {optimizer_hints}{distinct}{columns}"
 
         if elem.table:
             select += " FROM " + self.compile(c, elem.table)
@@ -495,9 +491,9 @@ def render_select(self, parent_c: Compiler, elem: Select) -> str:
         if elem.order_by_exprs:
             select += " ORDER BY " + ", ".join(map(compile_fn, elem.order_by_exprs))
 
-        if elem.limit_expr is not None and not self.USE_TOP_INSTEAD_LIMIT:
+        if elem.limit_expr is not None:
             has_order_by = bool(elem.order_by_exprs)
-            select += " " + self.offset_limit(0, elem.limit_expr, has_order_by=has_order_by)
+            select = self.limit_select(select_query=select, offset=0, limit=elem.limit_expr, has_order_by=has_order_by)
 
         if parent_c.in_select:
             select = f"({select}) {c.new_unique_name()}"
@@ -609,14 +605,17 @@ def render_inserttotable(self, c: Compiler, elem: InsertToTable) -> str:
 
         return f"INSERT INTO {self.compile(c, elem.path)}{columns} {expr}"
 
-    def offset_limit(
-        self, offset: Optional[int] = None, limit: Optional[int] = None, has_order_by: Optional[bool] = None
+    def limit_select(
+        self,
+        select_query: str,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+        has_order_by: Optional[bool] = None,
     ) -> str:
-        "Provide SQL fragment for limit and offset inside a select"
         if offset:
             raise NotImplementedError("No support for OFFSET in query")
 
-        return f"LIMIT {limit}"
+        return f"SELECT * FROM ({select_query}) AS LIMITED_SELECT LIMIT {limit}"
 
     def concat(self, items: List[str]) -> str:
         "Provide SQL for concatenating a bunch of columns into a string"
@@ -1107,7 +1106,7 @@ def _query_cursor(self, c, sql_code: str) -> QueryResult:
                 return result
         except Exception as _e:
             # logger.exception(e)
-            # logger.error(f'Caused by SQL: {sql_code}')
+            # logger.error(f"Caused by SQL: {sql_code}")
             raise
 
     def _query_conn(self, conn, sql_code: Union[str, ThreadLocalInterpreter]) -> QueryResult:

data_diff/databases/mssql.py

@@ -110,8 +110,12 @@ def is_distinct_from(self, a: str, b: str) -> str:
         # See: https://stackoverflow.com/a/18684859/857383
         return f"(({a}<>{b} OR {a} IS NULL OR {b} IS NULL) AND NOT({a} IS NULL AND {b} IS NULL))"
 
-    def offset_limit(
-        self, offset: Optional[int] = None, limit: Optional[int] = None, has_order_by: Optional[bool] = None
+    def limit_select(
+        self,
+        select_query: str,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+        has_order_by: Optional[bool] = None,
     ) -> str:
         if offset:
             raise NotImplementedError("No support for OFFSET in query")
@@ -121,7 +125,7 @@ def offset_limit(
             result += "ORDER BY 1"
 
         result += f" OFFSET 0 ROWS FETCH NEXT {limit} ROWS ONLY"
-        return result
+        return f"SELECT * FROM ({select_query}) AS LIMITED_SELECT {result}"
 
     def constant_values(self, rows) -> str:
         values = ", ".join("(%s)" % ", ".join(self._constant_value(v) for v in row) for row in rows)

data_diff/databases/oracle.py

@@ -64,13 +64,17 @@ def quote(self, s: str):
     def to_string(self, s: str):
         return f"cast({s} as varchar(1024))"
 
-    def offset_limit(
-        self, offset: Optional[int] = None, limit: Optional[int] = None, has_order_by: Optional[bool] = None
+    def limit_select(
+        self,
+        select_query: str,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+        has_order_by: Optional[bool] = None,
     ) -> str:
         if offset:
             raise NotImplementedError("No support for OFFSET in query")
 
-        return f"FETCH NEXT {limit} ROWS ONLY"
+        return f"SELECT * FROM ({select_query}) FETCH NEXT {limit} ROWS ONLY"
 
     def concat(self, items: List[str]) -> str:
         joined_exprs = " || ".join(items)

tests/test_query.py

@@ -50,11 +50,16 @@ def current_database(self) -> str:
     def current_schema(self) -> str:
         return "current_schema()"
 
-    def offset_limit(
-        self, offset: Optional[int] = None, limit: Optional[int] = None, has_order_by: Optional[bool] = None
+    def limit_select(
+        self,
+        select_query: str,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+        has_order_by: Optional[bool] = None,
     ) -> str:
         x = offset and f"OFFSET {offset}", limit and f"LIMIT {limit}"
-        return " ".join(filter(None, x))
+        result = " ".join(filter(None, x))
+        return f"SELECT * FROM ({select_query}) AS LIMITED_SELECT {result}"
 
     def explain_as_text(self, query: str) -> str:
         return f"explain {query}"