Repository: kayak/pypika
Pull request #529 (2020-12-08 16:35:10)

Reference query class from query builders
* Added QUERY_CLS class variable which references the relevant Query class. 
* Renamed PostgreQueryBuilder to PostgreSQLQueryBuilder
* Renamed SnowFlakeQueryBuilder to SnowflakeQueryBuilder

coveralls: 
[![Coverage Status](https://coveralls.io/builds/35541921/badge)](https://coveralls.io/builds/35541921)

Coverage increased (+0.01%) to 98.209% when pulling **8a320ca2c849aadbe5ec63530db581d3eeb0abdf on mikeengland:reference-query-cls-from-querybuilders** into **a3f1d025822ed34408ad93c8e8d9f3f794378f27 on kayak:master**.


Patch:
diff --git a/pypika/dialects.py b/pypika/dialects.py
index 991df2b..3aa847c 100644
--- a/pypika/dialects.py
+++ b/pypika/dialects.py
@@ -15,27 +15,43 @@ from pypika.utils import (
 )
 
 
-class SnowFlakeQueryBuilder(QueryBuilder):
+class SnowflakeQuery(Query):
+    """
+    Defines a query class for use with Snowflake.
+    """
+
+    @classmethod
+    def _builder(cls, **kwargs: Any) -> "SnowflakeQueryBuilder":
+        return SnowflakeQueryBuilder(**kwargs)
+
+
+class SnowflakeQueryBuilder(QueryBuilder):
     QUOTE_CHAR = None
     ALIAS_QUOTE_CHAR = '"'
     QUERY_ALIAS_QUOTE_CHAR = ''
+    QUERY_CLS = SnowflakeQuery
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(dialect=Dialects.SNOWFLAKE, **kwargs)
 
 
-class SnowflakeQuery(Query):
+class MySQLQuery(Query):
     """
-    Defines a query class for use with Snowflake.
+    Defines a query class for use with MySQL.
     """
 
     @classmethod
-    def _builder(cls, **kwargs: Any) -> SnowFlakeQueryBuilder:
-        return SnowFlakeQueryBuilder(**kwargs)
+    def _builder(cls, **kwargs: Any) -> "MySQLQueryBuilder":
+        return MySQLQueryBuilder(**kwargs)
+
+    @classmethod
+    def load(cls, fp: str) -> "MySQLLoadQueryBuilder":
+        return MySQLLoadQueryBuilder().load(fp)
 
 
 class MySQLQueryBuilder(QueryBuilder):
     QUOTE_CHAR = "`"
+    QUERY_CLS = MySQLQuery
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(dialect=Dialects.MYSQL, wrap_set_operation_queries=False, **kwargs)
@@ -108,16 +124,18 @@ class MySQLQueryBuilder(QueryBuilder):
 
 
 class MySQLLoadQueryBuilder:
+    QUERY_CLS = MySQLQuery
+
     def __init__(self) -> None:
         self._load_file = None
         self._into_table = None
 
     @builder
-    def load(self, fp: str) -> "MySQLQueryBuilder":
+    def load(self, fp: str) -> "MySQLLoadQueryBuilder":
         self._load_file = fp
 
     @builder
-    def into(self, table: Union[str, Table]) -> "MySQLQueryBuilder":
+    def into(self, table: Union[str, Table]) -> "MySQLLoadQueryBuilder":
         self._into_table = table if isinstance(table, Table) else Table(table)
 
     def get_sql(self, *args: Any, **kwargs: Any) -> str:
@@ -142,21 +160,27 @@ class MySQLLoadQueryBuilder:
         return self.get_sql()
 
 
-class MySQLQuery(Query):
+class VerticaQuery(Query):
     """
-    Defines a query class for use with MySQL.
+    Defines a query class for use with Vertica.
     """
 
     @classmethod
-    def _builder(cls, **kwargs: Any) -> "MySQLQueryBuilder":
-        return MySQLQueryBuilder(**kwargs)
+    def _builder(cls, **kwargs) -> "VerticaQueryBuilder":
+        return VerticaQueryBuilder(**kwargs)
 
     @classmethod
-    def load(cls, fp: str) -> "MySQLQueryBuilder":
-        return MySQLLoadQueryBuilder().load(fp)
+    def from_file(cls, fp: str) -> "VerticaCopyQueryBuilder":
+        return VerticaCopyQueryBuilder().from_file(fp)
+
+    @classmethod
+    def create_table(cls, table: Union[str, Table]) -> "VerticaCreateQueryBuilder":
+        return VerticaCreateQueryBuilder().create_table(table)
 
 
 class VerticaQueryBuilder(QueryBuilder):
+    QUERY_CLS = VerticaQuery
+
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(dialect=Dialects.VERTICA, **kwargs)
         self._hint = None
@@ -175,6 +199,8 @@ class VerticaQueryBuilder(QueryBuilder):
 
 
 class VerticaCreateQueryBuilder(CreateQueryBuilder):
+    QUERY_CLS = VerticaQuery
+
     def __init__(self) -> None:
         super().__init__(dialect=Dialects.VERTICA)
         self._local = False
@@ -217,6 +243,8 @@ class VerticaCreateQueryBuilder(CreateQueryBuilder):
 
 
 class VerticaCopyQueryBuilder:
+    QUERY_CLS = VerticaQuery
+
     def __init__(self) -> None:
         self._copy_table = None
         self._from_file = None
@@ -251,26 +279,19 @@ class VerticaCopyQueryBuilder:
         return self.get_sql()
 
 
-class VerticaQuery(Query):
+class OracleQuery(Query):
     """
-    Defines a query class for use with Vertica.
+    Defines a query class for use with Oracle.
     """
 
     @classmethod
-    def _builder(cls, **kwargs) -> VerticaQueryBuilder:
-        return VerticaQueryBuilder(**kwargs)
-
-    @classmethod
-    def from_file(cls, fp: str) -> VerticaCopyQueryBuilder:
-        return VerticaCopyQueryBuilder().from_file(fp)
-
-    @classmethod
-    def create_table(cls, table: Union[str, Table]) -> VerticaCreateQueryBuilder:
-        return VerticaCreateQueryBuilder().create_table(table)
+    def _builder(cls, **kwargs: Any) -> "OracleQueryBuilder":
+        return OracleQueryBuilder(**kwargs)
 
 
 class OracleQueryBuilder(QueryBuilder):
     QUOTE_CHAR = None
+    QUERY_CLS = OracleQuery
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(dialect=Dialects.ORACLE, **kwargs)
@@ -282,20 +303,21 @@ class OracleQueryBuilder(QueryBuilder):
         return super().get_sql(*args, **kwargs)
 
 
-class OracleQuery(Query):
+class PostgreSQLQuery(Query):
     """
-    Defines a query class for use with Oracle.
+    Defines a query class for use with PostgreSQL.
     """
 
     @classmethod
-    def _builder(cls, **kwargs: Any) -> OracleQueryBuilder:
-        return OracleQueryBuilder(**kwargs)
+    def _builder(cls, **kwargs) -> "PostgreSQLQueryBuilder":
+        return PostgreSQLQueryBuilder(**kwargs)
 
 
-class PostgreQueryBuilder(QueryBuilder):
+class PostgreSQLQueryBuilder(QueryBuilder):
     ALIAS_QUOTE_CHAR = '"'
+    QUERY_CLS = PostgreSQLQuery
 
-    def __init__(self, **kwargs: str) -> None:
+    def __init__(self, **kwargs: Any) -> None:
         super().__init__(dialect=Dialects.POSTGRESQL, **kwargs)
         self._returns = []
         self._return_star = False
@@ -309,14 +331,14 @@ class PostgreQueryBuilder(QueryBuilder):
 
         self._distinct_on = []
 
-    def __copy__(self) -> "PostgreQueryBuilder":
+    def __copy__(self) -> "PostgreSQLQueryBuilder":
         newone = super().__copy__()
         newone._returns = copy(self._returns)
         newone._on_conflict_do_updates = copy(self._on_conflict_do_updates)
         return newone
 
     @builder
-    def distinct_on(self, *fields: Union[str, Term]) -> "PostgreQueryBuilder":
+    def distinct_on(self, *fields: Union[str, Term]) -> "PostgreSQLQueryBuilder":
         for field in fields:
             if isinstance(field, str):
                 self._distinct_on.append(Field(field))
@@ -324,7 +346,7 @@ class PostgreQueryBuilder(QueryBuilder):
                 self._distinct_on.append(field)
 
     @builder
-    def on_conflict(self, *target_fields: Union[str, Term]) -> "PostgreQueryBuilder":
+    def on_conflict(self, *target_fields: Union[str, Term]) -> "PostgreSQLQueryBuilder":
         if not self._insert_table:
             raise QueryException("On conflict only applies to insert query")
 
@@ -337,13 +359,15 @@ class PostgreQueryBuilder(QueryBuilder):
                 self._on_conflict_fields.append(target_field)
 
     @builder
-    def do_nothing(self) -> "PostgreQueryBuilder":
+    def do_nothing(self) -> "PostgreSQLQueryBuilder":
         if len(self._on_conflict_do_updates) > 0:
             raise QueryException("Can not have two conflict handlers")
         self._on_conflict_do_nothing = True
 
     @builder
-    def do_update(self, update_field: Union[str, Field], update_value: Optional[Any] = None) -> "PostgreQueryBuilder":
+    def do_update(
+        self, update_field: Union[str, Field], update_value: Optional[Any] = None
+    ) -> "PostgreSQLQueryBuilder":
         if self._on_conflict_do_nothing:
             raise QueryException("Can not have two conflict handlers")
 
@@ -360,7 +384,7 @@ class PostgreQueryBuilder(QueryBuilder):
             self._on_conflict_do_updates.append((field, None))
 
     @builder
-    def where(self, criterion: Criterion) -> "PostgreQueryBuilder":
+    def where(self, criterion: Criterion) -> "PostgreSQLQueryBuilder":
         if not self._on_conflict:
             return super().where(criterion)
 
@@ -444,7 +468,7 @@ class PostgreQueryBuilder(QueryBuilder):
         return ''
 
     @builder
-    def returning(self, *terms: Any) -> "PostgreQueryBuilder":
+    def returning(self, *terms: Any) -> "PostgreSQLQueryBuilder":
         for term in terms:
             if isinstance(term, Field):
                 self._return_field(term)
@@ -507,7 +531,7 @@ class PostgreQueryBuilder(QueryBuilder):
     def get_sql(self, with_alias: bool = False, subquery: bool = False, **kwargs: Any) -> str:
         self._set_kwargs_defaults(kwargs)
 
-        querystring = super(PostgreQueryBuilder, self).get_sql(with_alias, subquery, **kwargs)
+        querystring = super(PostgreSQLQueryBuilder, self).get_sql(with_alias, subquery, **kwargs)
 
         querystring += self._on_conflict_sql(**kwargs)
         querystring += self._on_conflict_action_sql(**kwargs)
@@ -518,27 +542,33 @@ class PostgreQueryBuilder(QueryBuilder):
         return querystring
 
 
-class PostgreSQLQuery(Query):
+class RedshiftQuery(Query):
     """
-    Defines a query class for use with PostgreSQL.
+    Defines a query class for use with Amazon Redshift.
     """
 
     @classmethod
-    def _builder(cls, **kwargs) -> PostgreQueryBuilder:
-        return PostgreQueryBuilder(**kwargs)
+    def _builder(cls, **kwargs: Any) -> "RedShiftQueryBuilder":
+        return RedShiftQueryBuilder(dialect=Dialects.REDSHIFT, **kwargs)
 
 
-class RedshiftQuery(Query):
+class RedShiftQueryBuilder(QueryBuilder):
+    QUERY_CLS = RedshiftQuery
+
+
+class MSSQLQuery(Query):
     """
-    Defines a query class for use with Amazon Redshift.
+    Defines a query class for use with Microsoft SQL Server.
     """
 
     @classmethod
-    def _builder(cls, **kwargs: Any) -> "QueryBuilder":
-        return QueryBuilder(dialect=Dialects.REDSHIFT, **kwargs)
+    def _builder(cls, **kwargs: Any) -> "MSSQLQueryBuilder":
+        return MSSQLQueryBuilder(**kwargs)
 
 
 class MSSQLQueryBuilder(QueryBuilder):
+    QUERY_CLS = MSSQLQuery
+
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(dialect=Dialects.MSSQL, **kwargs)
         self._top = None
@@ -599,17 +629,21 @@ class MSSQLQueryBuilder(QueryBuilder):
         )
 
 
-class MSSQLQuery(Query):
+class ClickHouseQuery(Query):
     """
-    Defines a query class for use with Microsoft SQL Server.
+    Defines a query class for use with Yandex ClickHouse.
     """
 
     @classmethod
-    def _builder(cls, **kwargs: Any) -> MSSQLQueryBuilder:
-        return MSSQLQueryBuilder(**kwargs)
+    def _builder(cls, **kwargs: Any) -> "ClickHouseQueryBuilder":
+        return ClickHouseQueryBuilder(
+            dialect=Dialects.CLICKHOUSE, wrap_set_operation_queries=False, as_keyword=True, **kwargs
+        )
 
 
 class ClickHouseQueryBuilder(QueryBuilder):
+    QUERY_CLS = ClickHouseQuery
+
     @staticmethod
     def _delete_sql(**kwargs: Any) -> str:
         return 'ALTER TABLE'
@@ -634,18 +668,6 @@ class ClickHouseQueryBuilder(QueryBuilder):
         )
 
 
-class ClickHouseQuery(Query):
-    """
-    Defines a query class for use with Yandex ClickHouse.
-    """
-
-    @classmethod
-    def _builder(cls, **kwargs: Any) -> QueryBuilder:
-        return ClickHouseQueryBuilder(
-            dialect=Dialects.CLICKHOUSE, wrap_set_operation_queries=False, as_keyword=True, **kwargs
-        )
-
-
 class SQLLiteValueWrapper(ValueWrapper):
     def get_value_sql(self, *args: Any, **kwargs: Any) -> str:
         if isinstance(self.value, bool):
@@ -659,5 +681,9 @@ class SQLLiteQuery(Query):
     """
 
     @classmethod
-    def _builder(cls, **kwargs: Any) -> QueryBuilder:
-        return QueryBuilder(dialect=Dialects.SQLLITE, wrapper_cls=SQLLiteValueWrapper, **kwargs)
+    def _builder(cls, **kwargs: Any) -> "SQLLiteQueryBuilder":
+        return SQLLiteQueryBuilder(dialect=Dialects.SQLLITE, wrapper_cls=SQLLiteValueWrapper, **kwargs)
+
+
+class SQLLiteQueryBuilder(QueryBuilder):
+    QUERY_CLS = SQLLiteQuery
diff --git a/pypika/queries.py b/pypika/queries.py
index 9bc9dc5..9abd624 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -642,6 +642,7 @@ class QueryBuilder(Selectable, Term):
     SECONDARY_QUOTE_CHAR = "'"
     ALIAS_QUOTE_CHAR = None
     QUERY_ALIAS_QUOTE_CHAR = None
+    QUERY_CLS = Query
 
     def __init__(
         self,
@@ -1648,6 +1649,7 @@ class CreateQueryBuilder:
     QUOTE_CHAR = '"'
     SECONDARY_QUOTE_CHAR = "'"
     ALIAS_QUOTE_CHAR = None
+    QUERY_CLS = Query
 
     def __init__(self, dialect: Optional[Dialects] = None) -> None:
         self._create_table = None
@@ -1924,6 +1926,7 @@ class DropQueryBuilder:
     QUOTE_CHAR = '"'
     SECONDARY_QUOTE_CHAR = "'"
     ALIAS_QUOTE_CHAR = None
+    QUERY_CLS = Query
 
     def __init__(self, dialect: Optional[Dialects] = None) -> None:
         self._drop_table = None


Test patch:
diff --git a/pypika/tests/test_query.py b/pypika/tests/test_query.py
index 0cb6bc0..460b7d7 100644
--- a/pypika/tests/test_query.py
+++ b/pypika/tests/test_query.py
@@ -1,6 +1,30 @@
 import unittest
 
 from pypika import Case, Query, Tables, Tuple, functions
+from pypika.dialects import (
+    ClickHouseQuery,
+    ClickHouseQueryBuilder,
+    MSSQLQuery,
+    MSSQLQueryBuilder,
+    MySQLLoadQueryBuilder,
+    MySQLQuery,
+    MySQLQueryBuilder,
+    OracleQuery,
+    OracleQueryBuilder,
+    PostgreSQLQuery,
+    PostgreSQLQueryBuilder,
+    RedShiftQueryBuilder,
+    RedshiftQuery,
+    SQLLiteQuery,
+    SQLLiteQueryBuilder,
+    SnowflakeQuery,
+    SnowflakeQueryBuilder,
+    VerticaCopyQueryBuilder,
+    VerticaCreateQueryBuilder,
+    VerticaQuery,
+    VerticaQueryBuilder,
+)
+from pypika.queries import CreateQueryBuilder, DropQueryBuilder, QueryBuilder
 
 
 class QueryTablesTests(unittest.TestCase):
@@ -132,3 +156,51 @@ class QueryTablesTests(unittest.TestCase):
 
         self.assertTrue(q.is_joined(self.table_b))
         self.assertFalse(q.is_joined(self.table_c))
+
+
+class QueryBuilderTests(unittest.TestCase):
+    def test_query_builders_have_reference_to_correct_query_class(self):
+        with self.subTest('QueryBuilder'):
+            self.assertEqual(Query, QueryBuilder.QUERY_CLS)
+
+        with self.subTest('DropQueryBuilder'):
+            self.assertEqual(Query, DropQueryBuilder.QUERY_CLS)
+
+        with self.subTest('CreateQueryBuilder'):
+            self.assertEqual(Query, CreateQueryBuilder.QUERY_CLS)
+
+        with self.subTest('MySQLQueryBuilder'):
+            self.assertEqual(MySQLQuery, MySQLQueryBuilder.QUERY_CLS)
+
+        with self.subTest('MySQLLoadQueryBuilder'):
+            self.assertEqual(MySQLQuery, MySQLLoadQueryBuilder.QUERY_CLS)
+
+        with self.subTest('VerticaQueryBuilder'):
+            self.assertEqual(VerticaQuery, VerticaQueryBuilder.QUERY_CLS)
+
+        with self.subTest('VerticaCreateQueryBuilder'):
+            self.assertEqual(VerticaQuery, VerticaCreateQueryBuilder.QUERY_CLS)
+
+        with self.subTest('VerticaCopyQueryBuilder'):
+            self.assertEqual(VerticaQuery, VerticaCopyQueryBuilder.QUERY_CLS)
+
+        with self.subTest('PostgreSQLQueryBuilder'):
+            self.assertEqual(PostgreSQLQuery, PostgreSQLQueryBuilder.QUERY_CLS)
+
+        with self.subTest('MSSQLQueryBuilder'):
+            self.assertEqual(MSSQLQuery, MSSQLQueryBuilder.QUERY_CLS)
+
+        with self.subTest('SnowflakeQueryBuilder'):
+            self.assertEqual(SnowflakeQuery, SnowflakeQueryBuilder.QUERY_CLS)
+
+        with self.subTest('ClickHouseQueryBuilder'):
+            self.assertEqual(ClickHouseQuery, ClickHouseQueryBuilder.QUERY_CLS)
+
+        with self.subTest('RedShiftQueryBuilder'):
+            self.assertEqual(RedshiftQuery, RedShiftQueryBuilder.QUERY_CLS)
+
+        with self.subTest('SQLLiteQueryBuilder'):
+            self.assertEqual(SQLLiteQuery, SQLLiteQueryBuilder.QUERY_CLS)
+
+        with self.subTest('OracleQueryBuilder'):
+            self.assertEqual(OracleQuery, OracleQueryBuilder.QUERY_CLS)