Repository: kayak/pypika
Pull request #601 (2021-05-22 23:16:05)

Improvement on TOP statement support for `mssql`.
Added support for WITH TIES and PERCENT statements when using MSSQL dialect.


Patch:
diff --git a/README.rst b/README.rst
index 6d7a4aa..c66d176 100644
--- a/README.rst
+++ b/README.rst
@@ -862,25 +862,25 @@ MySQL
 
     customers = Table('customers')
 
-    q = Query.into(customers)\
-        .insert(1, 'Jane', 'Doe', 'jane@example.com')\
+    q = MySQLQuery.into(customers) \
+        .insert(1, 'Jane', 'Doe', 'jane@example.com') \
         .on_duplicate_key_ignore())
 
 .. code-block:: sql
 
-    INSERT INTO customers VALUES (1,'Jane','Doe','jane@example.com') ON DUPLICATE KEY IGNORE
+    INSERT INTO `customers` VALUES (1,'Jane','Doe','jane@example.com') ON DUPLICATE KEY IGNORE
 
 .. code-block:: python
 
     customers = Table('customers')
 
-    q = Query.into(customers)\
-        .insert(1, 'Jane', 'Doe', 'jane@example.com')\
+    q = MySQLQuery.into(customers) \
+        .insert(1, 'Jane', 'Doe', 'jane@example.com') \
         .on_duplicate_key_update(customers.email, Values(customers.email))
 
 .. code-block:: sql
 
-    INSERT INTO customers VALUES (1,'Jane','Doe','jane@example.com') ON DUPLICATE KEY UPDATE `email`=VALUES(`email`)
+    INSERT INTO `customers` VALUES (1,'Jane','Doe','jane@example.com') ON DUPLICATE KEY UPDATE `email`=VALUES(`email`)
 
 ``.on_duplicate_key_update`` works similar to ``.set`` for updating rows, additionally it provides the ``Values``
 wrapper to update to the value specified in the ``INSERT`` clause.
@@ -892,22 +892,22 @@ PostgreSQL
 
     customers = Table('customers')
 
-    q = Query.into(customers)\
-        .insert(1, 'Jane', 'Doe', 'jane@example.com')\
-        .on_conflict(customers.email)
+    q = PostgreSQLQuery.into(customers) \
+        .insert(1, 'Jane', 'Doe', 'jane@example.com') \
+        .on_conflict(customers.email) \
         .do_nothing()
 
 .. code-block:: sql
 
-    INSERT INTO "abc" VALUES (1,'Jane','Doe','jane@example.com') ON CONFLICT ("email") DO NOTHING
+    INSERT INTO "customers" VALUES (1,'Jane','Doe','jane@example.com') ON CONFLICT ("email") DO NOTHING
 
 .. code-block:: python
 
     customers = Table('customers')
 
-    q = Query.into(customers)\
-        .insert(1, 'Jane', 'Doe', 'jane@example.com')\
-        .on_conflict(customers.email)
+    q = PostgreSQLQuery.into(customers) \
+        .insert(1, 'Jane', 'Doe', 'jane@example.com') \
+        .on_conflict(customers.email) \
         .do_update(customers.email, 'bob@example.com')
 
 .. code-block:: sql
diff --git a/pypika/dialects.py b/pypika/dialects.py
index ec5a4e4..4f42465 100644
--- a/pypika/dialects.py
+++ b/pypika/dialects.py
@@ -607,15 +607,14 @@ class MSSQLQueryBuilder(QueryBuilder):
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(dialect=Dialects.MSSQL, **kwargs)
-        self._top = None
+        self._top: Union[int, None] = None
+        self._top_with_ties: bool = False
+        self._top_percent: bool = False
 
     @builder
-    def top(self, value: Union[str, int]) -> "MSSQLQueryBuilder":
+    def top(self, value: Union[str, int], percent: bool = False, with_ties: bool = False) -> "MSSQLQueryBuilder":
         """
         Implements support for simple TOP clauses.
-
-        Does not include support for PERCENT or WITH TIES.
-
         https://docs.microsoft.com/en-us/sql/t-sql/queries/top-transact-sql?view=sql-server-2017
         """
         try:
@@ -623,6 +622,11 @@ class MSSQLQueryBuilder(QueryBuilder):
         except ValueError:
             raise QueryException("TOP value must be an integer")
 
+        if percent and not (0 <= int(value) <= 100):
+            raise QueryException("TOP value must be between 0 and 100 when `percent`" " is specified")
+        self._top_percent: bool = percent
+        self._top_with_ties: bool = with_ties
+
     @builder
     def fetch_next(self, limit: int) -> "MSSQLQueryBuilder":
         # Overridden to provide a more domain-specific API for T-SQL users
@@ -652,10 +656,15 @@ class MSSQLQueryBuilder(QueryBuilder):
         return super().get_sql(*args, **kwargs)
 
     def _top_sql(self) -> str:
+        _top_statement: str = ""
         if self._top:
-            return "TOP ({}) ".format(self._top)
-        else:
-            return ""
+            _top_statement = f"TOP ({self._top}) "
+            if self._top_percent:
+                _top_statement = f"{_top_statement}PERCENT "
+            if self._top_with_ties:
+                _top_statement = f"{_top_statement}WITH TIES "
+
+        return _top_statement
 
     def _select_sql(self, **kwargs: Any) -> str:
         return "SELECT {distinct}{top}{select}".format(
@@ -705,10 +714,10 @@ class ClickHouseQueryBuilder(QueryBuilder):
 
 
 class SQLLiteValueWrapper(ValueWrapper):
-    def get_value_sql(self, *args: Any, **kwargs: Any) -> str:
+    def get_value_sql(self, **kwargs: Any) -> str:
         if isinstance(self.value, bool):
             return "1" if self.value else "0"
-        return super().get_value_sql(*args, **kwargs)
+        return super().get_value_sql(**kwargs)
 
 
 class SQLLiteQuery(Query):
@@ -718,8 +727,22 @@ class SQLLiteQuery(Query):
 
     @classmethod
     def _builder(cls, **kwargs: Any) -> "SQLLiteQueryBuilder":
-        return SQLLiteQueryBuilder(dialect=Dialects.SQLLITE, wrapper_cls=SQLLiteValueWrapper, **kwargs)
+        return SQLLiteQueryBuilder(**kwargs)
 
 
 class SQLLiteQueryBuilder(QueryBuilder):
     QUERY_CLS = SQLLiteQuery
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(dialect=Dialects.SQLLITE, wrapper_cls=SQLLiteValueWrapper, **kwargs)
+        self._insert_or_replace = False
+
+    @builder
+    def insert_or_replace(self, *terms: Any) -> "SQLLiteQueryBuilder":
+        self._apply_terms(*terms)
+        self._replace = True
+        self._insert_or_replace = True
+
+    def _replace_sql(self, **kwargs: Any) -> str:
+        prefix = "INSERT OR " if self._insert_or_replace else ""
+        return prefix + super()._replace_sql(**kwargs)
diff --git a/pypika/queries.py b/pypika/queries.py
index 36397eb..3ea70ab 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -837,24 +837,12 @@ class QueryBuilder(Selectable, Term):
 
     @builder
     def insert(self, *terms: Any) -> "QueryBuilder":
-        if self._insert_table is None:
-            raise AttributeError("'Query' object has no attribute '%s'" % "insert")
-
-        if not terms:
-            return
-        else:
-            self._validate_terms_and_append(*terms)
+        self._apply_terms(*terms)
         self._replace = False
 
     @builder
     def replace(self, *terms: Any) -> "QueryBuilder":
-        if self._insert_table is None:
-            raise AttributeError("'Query' object has no attribute '%s'" % "insert")
-
-        if not terms:
-            return
-        else:
-            self._validate_terms_and_append(*terms)
+        self._apply_terms(*terms)
         self._replace = True
 
     @builder
@@ -1140,11 +1128,17 @@ class QueryBuilder(Selectable, Term):
         subquery.alias = "sq%d" % self._subquery_count
         self._subquery_count += 1
 
-    def _validate_terms_and_append(self, *terms: Any) -> None:
+    def _apply_terms(self, *terms: Any) -> None:
         """
         Handy function for INSERT and REPLACE statements in order to check if
         terms are introduced and how append them to `self._values`
         """
+        if self._insert_table is None:
+            raise AttributeError("'Query' object has no attribute '%s'" % "insert")
+
+        if not terms:
+            return
+
         if not isinstance(terms[0], (list, tuple, set)):
             terms = [terms]
 


Test patch:
diff --git a/pypika/tests/dialects/test_mssql.py b/pypika/tests/dialects/test_mssql.py
index 1c7f78b..f8940b0 100644
--- a/pypika/tests/dialects/test_mssql.py
+++ b/pypika/tests/dialects/test_mssql.py
@@ -27,6 +27,23 @@ class SelectTests(unittest.TestCase):
 
         self.assertEqual('SELECT TOP (10) "def" FROM "abc"', str(q))
 
+    def test_top_with_ties_select(self):
+        q = MSSQLQuery.from_("abc").select("def").top(10, with_ties=True)
+
+        self.assertEqual('SELECT TOP (10) WITH TIES "def" FROM "abc"', str(q))
+
+    def test_top_percent_select(self):
+        q = MSSQLQuery.from_("abc").select("def").top(10, percent=True)
+
+        self.assertEqual('SELECT TOP (10) PERCENT "def" FROM "abc"', str(q))
+
+    def test_top_percent_invalid_range(self):
+        for invalid_top in [-1, 101]:
+            with self.assertRaisesRegex(
+                QueryException, "TOP value must be between 0 and 100 when `percent` is specified"
+            ):
+                MSSQLQuery.from_("abc").select("def").top(invalid_top, percent=True)
+
     def test_top_select_non_int(self):
         with self.assertRaisesRegex(QueryException, "TOP value must be an integer"):
             MSSQLQuery.from_("abc").select("def").top("a")
diff --git a/pypika/tests/dialects/test_sqllite.py b/pypika/tests/dialects/test_sqllite.py
index b3b269a..104a723 100644
--- a/pypika/tests/dialects/test_sqllite.py
+++ b/pypika/tests/dialects/test_sqllite.py
@@ -16,3 +16,20 @@ class SelectTests(unittest.TestCase):
         q = SQLLiteQuery.from_("abc").select(False)
 
         self.assertEqual('SELECT 0 FROM "abc"', str(q))
+
+
+class ReplaceTests(unittest.TestCase):
+    def test_normal_replace(self):
+        query = SQLLiteQuery.into("abc").replace("v1", "v2", "v3")
+        expected_output = "REPLACE INTO \"abc\" VALUES ('v1','v2','v3')"
+        self.assertEqual(expected_output, str(query))
+
+    def test_replace_subquery(self):
+        query = SQLLiteQuery.into("abc").replace(SQLLiteQuery.from_("def").select("f1", "f2"))
+        expected_output = 'REPLACE INTO "abc" VALUES ((SELECT "f1","f2" FROM "def"))'
+        self.assertEqual(expected_output, str(query))
+
+    def test_insert_or_replace(self):
+        query = SQLLiteQuery.into("abc").insert_or_replace("v1", "v2", "v3")
+        expected_output = "INSERT OR REPLACE INTO \"abc\" VALUES ('v1','v2','v3')"
+        self.assertEqual(expected_output, str(query))