Add qualify

pypika/queries.py

@@ -736,6 +736,7 @@ def __init__(
         self._groupbys = []
         self._with_totals = False
         self._havings = None
+        self._qualifys = None
         self._orderbys = []
         self._joins = []
         self._unions = []
@@ -835,6 +836,7 @@ def replace_table(self, current_table: Optional[Table], new_table: Optional[Tabl
         self._prewheres = self._prewheres.replace_table(current_table, new_table) if self._prewheres else None
         self._groupbys = [groupby.replace_table(current_table, new_table) for groupby in self._groupbys]
         self._havings = self._havings.replace_table(current_table, new_table) if self._havings else None
+        self._qualifys = self._qualifys.replace_table(current_table, new_table) if self._qualifys else None
         self._orderbys = [
             (orderby[0].replace_table(current_table, new_table), orderby[1]) for orderby in self._orderbys
         ]
@@ -969,6 +971,16 @@ def having(self, criterion: Union[Term, EmptyCriterion]) -> "QueryBuilder":
         else:
             self._havings = criterion
 
+    @builder
+    def qualify(self, criterion: Union[Term, EmptyCriterion]) -> "QueryBuilder":
+        if isinstance(criterion, EmptyCriterion):
+            return
+
+        if self._qualifys:
+            self._qualifys &= criterion
+        else:
+            self._qualifys = criterion
+
     @builder
     def groupby(self, *terms: Union[str, int, Term]) -> "QueryBuilder":
         for term in terms:
@@ -1351,6 +1363,9 @@ def get_sql(self, with_alias: bool = False, subquery: bool = False, **kwargs: An
         if self._havings:
             querystring += self._having_sql(**kwargs)
 
+        if self._qualifys:
+            querystring += self._qualify_sql(**kwargs)
+
         if self._orderbys:
             querystring += self._orderby_sql(**kwargs)
 
@@ -1544,6 +1559,9 @@ def _rollup_sql(self) -> str:
     def _having_sql(self, quote_char: Optional[str] = None, **kwargs: Any) -> str:
         return " HAVING {having}".format(having=self._havings.get_sql(quote_char=quote_char, **kwargs))
 
+    def _qualify_sql(self, quote_char: Optional[str] = None, **kwargs: Any) -> str:
+        return " QUALIFY {qualify}".format(qualify=self._qualifys.get_sql(quote_char=quote_char, **kwargs))
+
     def _offset_sql(self) -> str:
         return " OFFSET {offset}".format(offset=self._offset)
 

pypika/tests/test_analytic_queries.py

@@ -608,3 +608,27 @@ def test_lead_generates_correct_sql(self):
             self.assertEqual(
                 'SELECT LEAD("date",1,\'2000-01-01\') OVER(PARTITION BY "foo" ORDER BY "date") FROM "abc"', str(q)
             )
+
+    def test_qualify(self):
+        expr = an.Rank().over(self.table_abc.foo).orderby(self.table_abc.date)
+        q = Query.from_(self.table_abc).select("*").qualify(expr > 1)
+
+        self.assertEqual(
+            "SELECT " "*" ' FROM "abc" ' 'QUALIFY RANK() OVER(PARTITION BY "foo" ORDER BY "date")>1',
+            str(q),
+        )
+
+    def test_qualify_and(self):
+        expr1 = an.Rank().over(self.table_abc.foo).orderby(self.table_abc.date)
+        expr2 = an.RowNumber().over(self.table_abc.foo).orderby(self.table_abc.date)
+        q = Query.from_(self.table_abc).select("*").qualify((expr1 > 1) & (expr2 > 2))
+
+        self.assertEqual(
+            "SELECT "
+            "*"
+            ' FROM "abc" '
+            'QUALIFY RANK() OVER(PARTITION BY "foo" ORDER BY "date")>1 '
+            "AND "
+            'ROW_NUMBER() OVER(PARTITION BY "foo" ORDER BY "date")>2',
+            str(q),
+        )

pypika/tests/test_selects.py

@@ -954,6 +954,60 @@ def test_redshift_query_uses_double_quote_chars(self):
         self.assertEqual('SELECT "foo" FROM "abc" GROUP BY "foo" HAVING "buz"=\'fiz\'', str(q))
 
 
+class QualifyTests(unittest.TestCase):
+    table_abc = Table("abc")
+
+    def test_qualify_greater_than(self):
+        q = (
+            Query.from_(self.table_abc)
+            .select(self.table_abc.foo, fn.Sum(self.table_abc.bar))
+            .groupby(self.table_abc.foo)
+            .qualify(self.table_abc.bar > 1)
+        )
+
+        self.assertEqual(
+            'SELECT "foo",SUM("bar") FROM "abc" GROUP BY "foo" QUALIFY "bar">1',
+            str(q),
+        )
+
+    def test_qualify_and(self):
+        q = (
+            Query.from_(self.table_abc)
+            .select(self.table_abc.foo, fn.Sum(self.table_abc.bar))
+            .groupby(self.table_abc.foo)
+            .qualify((self.table_abc.bar > 1) & (self.table_abc.bar < 100))
+        )
+
+        self.assertEqual(
+            'SELECT "foo",SUM("bar") FROM "abc" GROUP BY "foo" QUALIFY "bar">1 AND "bar"<100',
+            str(q),
+        )
+
+    def test_mysql_query_uses_backtick_quote_chars(self):
+        q = MySQLQuery.from_(self.table_abc).select(self.table_abc.foo).qualify(self.table_abc.buz == "fiz")
+        self.assertEqual("SELECT `foo` FROM `abc` QUALIFY `buz`='fiz'", str(q))
+
+    def test_vertica_query_uses_double_quote_chars(self):
+        q = VerticaQuery.from_(self.table_abc).select(self.table_abc.foo).qualify(self.table_abc.buz == "fiz")
+        self.assertEqual('SELECT "foo" FROM "abc" QUALIFY "buz"=\'fiz\'', str(q))
+
+    def test_mssql_query_uses_double_quote_chars(self):
+        q = MSSQLQuery.from_(self.table_abc).select(self.table_abc.foo).qualify(self.table_abc.buz == "fiz")
+        self.assertEqual('SELECT "foo" FROM "abc" QUALIFY "buz"=\'fiz\'', str(q))
+
+    def test_oracle_query_uses_no_quote_chars(self):
+        q = OracleQuery.from_(self.table_abc).select(self.table_abc.foo).qualify(self.table_abc.buz == "fiz")
+        self.assertEqual('SELECT foo FROM abc QUALIFY buz=\'fiz\'', str(q))
+
+    def test_postgres_query_uses_double_quote_chars(self):
+        q = PostgreSQLQuery.from_(self.table_abc).select(self.table_abc.foo).qualify(self.table_abc.buz == "fiz")
+        self.assertEqual('SELECT "foo" FROM "abc" QUALIFY "buz"=\'fiz\'', str(q))
+
+    def test_redshift_query_uses_double_quote_chars(self):
+        q = RedshiftQuery.from_(self.table_abc).select(self.table_abc.foo).qualify(self.table_abc.buz == "fiz")
+        self.assertEqual('SELECT "foo" FROM "abc" QUALIFY "buz"=\'fiz\'', str(q))
+
+
 class OrderByTests(unittest.TestCase):
     t = Table("abc")