Skip to content

Commit

Permalink
Merge cfe5b02 into 40d7885
Browse files Browse the repository at this point in the history
  • Loading branch information
gleb-kov committed Jun 17, 2021
2 parents 40d7885 + cfe5b02 commit d24b185
Show file tree
Hide file tree
Showing 4 changed files with 199 additions and 33 deletions.
66 changes: 59 additions & 7 deletions pypika/dialects.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,17 +5,15 @@
from pypika.enums import Dialects
from pypika.queries import (
CreateQueryBuilder,
Query,
QueryBuilder,
Database,
DropQueryBuilder,
Selectable,
Table,
DropQueryBuilder,
Query,
QueryBuilder,
)
from pypika.terms import ArithmeticExpression, Criterion, EmptyCriterion, Field, Function, Star, Term, ValueWrapper
from pypika.utils import (
QueryException,
builder,
)
from pypika.utils import QueryException, builder, format_quotes


class SnowflakeQuery(Query):
Expand Down Expand Up @@ -685,6 +683,30 @@ def _builder(cls, **kwargs: Any) -> "ClickHouseQueryBuilder":
dialect=Dialects.CLICKHOUSE, wrap_set_operation_queries=False, as_keyword=True, **kwargs
)

@classmethod
def drop_database(self, database: Union[Database, str]) -> "ClickHouseDropQueryBuilder":
return ClickHouseDropQueryBuilder().drop_database(database)

@classmethod
def drop_table(self, table: Union[Table, str]) -> "ClickHouseDropQueryBuilder":
return ClickHouseDropQueryBuilder().drop_table(table)

@classmethod
def drop_dictionary(self, dictionary: str) -> "ClickHouseDropQueryBuilder":
return ClickHouseDropQueryBuilder().drop_dictionary(dictionary)

@classmethod
def drop_quota(self, quota: str) -> "ClickHouseDropQueryBuilder":
return ClickHouseDropQueryBuilder().drop_quota(quota)

@classmethod
def drop_user(self, user: str) -> "ClickHouseDropQueryBuilder":
return ClickHouseDropQueryBuilder().drop_user(user)

@classmethod
def drop_view(self, view: str) -> "ClickHouseDropQueryBuilder":
return ClickHouseDropQueryBuilder().drop_view(view)


class ClickHouseQueryBuilder(QueryBuilder):
QUERY_CLS = ClickHouseQuery
Expand Down Expand Up @@ -713,6 +735,36 @@ def _set_sql(self, **kwargs: Any) -> str:
)


class ClickHouseDropQueryBuilder(DropQueryBuilder):
QUERY_CLS = ClickHouseQuery

def __init__(self):
super().__init__(dialect=Dialects.CLICKHOUSE)
self._cluster_name = None

@builder
def drop_dictionary(self, dictionary: str) -> "ClickHouseDropQueryBuilder":
super()._set_target('DICTIONARY', dictionary)

@builder
def drop_quota(self, quota: str) -> "ClickHouseDropQueryBuilder":
super()._set_target('QUOTA', quota)

@builder
def on_cluster(self, cluster: str) -> "ClickHouseDropQueryBuilder":
if self._cluster_name:
raise AttributeError("'DropQuery' object already has attribute cluster_name")
self._cluster_name = cluster

def get_sql(self, **kwargs: Any) -> str:
query = super().get_sql(**kwargs)

if self._drop_target_kind != "DICTIONARY" and self._cluster_name is not None:
query += " ON CLUSTER " + format_quotes(self._cluster_name, super().QUOTE_CHAR)

return query


class SQLLiteValueWrapper(ValueWrapper):
def get_value_sql(self, **kwargs: Any) -> str:
if isinstance(self.value, bool):
Expand Down
91 changes: 71 additions & 20 deletions pypika/queries.py
Original file line number Diff line number Diff line change
Expand Up @@ -384,6 +384,18 @@ def create_table(cls, table: Union[str, Table]) -> "CreateQueryBuilder":
"""
return CreateQueryBuilder().create_table(table)

@classmethod
def drop_database(cls, database: Union[Database, Table]) -> "DropQueryBuilder":
"""
Query builder entry point. Initializes query building and sets the table name to be dropped. When using this
function, the query becomes a DROP statement.
:param database: An instance of a Database object or a string database name.
:return: DropQueryBuilder
"""
return DropQueryBuilder().drop_database(database)

@classmethod
def drop_table(cls, table: Union[str, Table]) -> "DropQueryBuilder":
"""
Expand All @@ -396,6 +408,30 @@ def drop_table(cls, table: Union[str, Table]) -> "DropQueryBuilder":
"""
return DropQueryBuilder().drop_table(table)

@classmethod
def drop_user(cls, user: str) -> "DropQueryBuilder":
"""
Query builder entry point. Initializes query building and sets the table name to be dropped. When using this
function, the query becomes a DROP statement.
:param user: String user name.
:return: DropQueryBuilder
"""
return DropQueryBuilder().drop_user(user)

@classmethod
def drop_view(cls, view: str) -> "DropQueryBuilder":
"""
Query builder entry point. Initializes query building and sets the table name to be dropped. When using this
function, the query becomes a DROP statement.
:param view: String view name.
:return: DropQueryBuilder
"""
return DropQueryBuilder().drop_view(view)

@classmethod
def into(cls, table: Union[Table, str], **kwargs: Any) -> "QueryBuilder":
"""
Expand Down Expand Up @@ -1938,7 +1974,8 @@ class DropQueryBuilder:
QUERY_CLS = Query

def __init__(self, dialect: Optional[Dialects] = None) -> None:
self._drop_table = None
self._drop_target_kind = None
self._drop_target: Union[Database, Table, str] = ""
self._if_exists = None
self.dialect = dialect

Expand All @@ -1947,35 +1984,49 @@ def _set_kwargs_defaults(self, kwargs: dict) -> None:
kwargs.setdefault("secondary_quote_char", self.SECONDARY_QUOTE_CHAR)
kwargs.setdefault("dialect", self.dialect)

def get_sql(self, **kwargs: Any) -> str:
self._set_kwargs_defaults(kwargs)

if not self._drop_table:
return ""

querystring = self._drop_table_sql(**kwargs)

return querystring
@builder
def drop_database(self, database: Union[Database, str]) -> "DropQueryBuilder":
target = database if isinstance(database, Database) else Database(database)
self._set_target('DATABASE', target)

@builder
def drop_table(self, table: Union[Table, str]) -> "DropQueryBuilder":
if self._drop_table:
raise AttributeError("'Query' object already has attribute drop_table")
target = table if isinstance(table, Table) else Table(table)
self._set_target('TABLE', target)

@builder
def drop_user(self, user: str) -> "DropQueryBuilder":
self._set_target('USER', user)

self._drop_table = table if isinstance(table, Table) else Table(table)
@builder
def drop_view(self, view: str) -> "DropQueryBuilder":
self._set_target('VIEW', view)

@builder
def if_exists(self) -> "DropQueryBuilder":
self._if_exists = True

def _drop_table_sql(self, **kwargs: Any) -> str:
if_exists = ''
if self._if_exists:
if_exists = 'IF EXISTS '
def _set_target(self, kind: str, target: Union[Database, Table, str]) -> None:
if self._drop_target:
raise AttributeError("'DropQuery' object already has attribute drop_target")
self._drop_target_kind = kind
self._drop_target = target

def get_sql(self, **kwargs: Any) -> str:
self._set_kwargs_defaults(kwargs)

if_exists = 'IF EXISTS ' if self._if_exists else ''
target_name: str = ""

if isinstance(self._drop_target, Database):
target_name = self._drop_target.get_sql(**kwargs)
elif isinstance(self._drop_target, Table):
target_name = self._drop_target.get_sql(**kwargs)
else:
target_name = format_quotes(self._drop_target, self.QUOTE_CHAR)

return "DROP TABLE {if_exists}{table}".format(
if_exists=if_exists,
table=self._drop_table.get_sql(**kwargs),
return "DROP {kind} {if_exists}{name}".format(
kind=self._drop_target_kind, if_exists=if_exists, name=target_name
)

def __str__(self) -> str:
Expand Down
43 changes: 43 additions & 0 deletions pypika/tests/dialects/test_clickhouse.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from pypika import (
ClickHouseQuery,
Database,
Table,
)

Expand Down Expand Up @@ -41,3 +42,45 @@ def test_update(self):
q = ClickHouseQuery.update(self.table_abc).where(self.table_abc.foo == 0).set("foo", "bar")

self.assertEqual('ALTER TABLE "abc" UPDATE "foo"=\'bar\' WHERE "foo"=0', str(q))


class ClickHouseDropQuery(TestCase):
table_abc = Table("abc")
database_xyz = Database("mydb")
cluster_name = "mycluster"

def test_drop_database(self):
q1 = ClickHouseQuery.drop_database(self.database_xyz)
q2 = ClickHouseQuery.drop_database(self.database_xyz).on_cluster(self.cluster_name)
q3 = ClickHouseQuery.drop_database(self.database_xyz).if_exists().on_cluster(self.cluster_name)

self.assertEqual('DROP DATABASE "mydb"', str(q1))
self.assertEqual('DROP DATABASE "mydb" ON CLUSTER "mycluster"', str(q2))
self.assertEqual('DROP DATABASE IF EXISTS "mydb" ON CLUSTER "mycluster"', str(q3))

def test_drop_table(self):
q1 = ClickHouseQuery.drop_table(self.table_abc)
q2 = ClickHouseQuery.drop_table(self.table_abc).on_cluster(self.cluster_name)
q3 = ClickHouseQuery.drop_table(self.table_abc).if_exists().on_cluster(self.cluster_name)

self.assertEqual('DROP TABLE "abc"', str(q1))
self.assertEqual('DROP TABLE "abc" ON CLUSTER "mycluster"', str(q2))
self.assertEqual('DROP TABLE IF EXISTS "abc" ON CLUSTER "mycluster"', str(q3))

def test_drop_dictionary(self):
q1 = ClickHouseQuery.drop_dictionary("dict")
q2 = ClickHouseQuery.drop_dictionary("dict").on_cluster(self.cluster_name)
q3 = ClickHouseQuery.drop_dictionary("dict").if_exists().on_cluster(self.cluster_name)

self.assertEqual('DROP DICTIONARY "dict"', str(q1))
self.assertEqual('DROP DICTIONARY "dict"', str(q2)) # NO CLUSTER
self.assertEqual('DROP DICTIONARY IF EXISTS "dict"', str(q3)) # NO CLUSTER

def test_drop_other(self):
q1 = ClickHouseQuery.drop_quota("myquota")
q2 = ClickHouseQuery.drop_user("myuser")
q3 = ClickHouseQuery.drop_view("myview")

self.assertEqual('DROP QUOTA "myquota"', str(q1))
self.assertEqual('DROP USER "myuser"', str(q2))
self.assertEqual('DROP VIEW "myview"', str(q3))
32 changes: 26 additions & 6 deletions pypika/tests/test_drop.py
Original file line number Diff line number Diff line change
@@ -1,22 +1,42 @@
import unittest

from pypika import (
Tables,
Columns,
Database,
Query,
Tables,
)


class DropTableTests(unittest.TestCase):
database_xyz = Database("mydb")
new_table, existing_table = Tables("abc", "efg")
foo, bar = Columns(("a", "INT"), ("b", "VARCHAR(100)"))

def test_drop_database(self):
q1 = Query.drop_database(self.database_xyz)
q2 = Query.drop_database(self.database_xyz).if_exists()

self.assertEqual('DROP DATABASE "mydb"', str(q1))
self.assertEqual('DROP DATABASE IF EXISTS "mydb"', str(q2))

def test_drop_table(self):
q = Query.drop_table(self.new_table)
q1 = Query.drop_table(self.new_table)
q2 = Query.drop_table(self.new_table).if_exists()

self.assertEqual('DROP TABLE "abc"', str(q1))
self.assertEqual('DROP TABLE IF EXISTS "abc"', str(q2))

def test_drop_user(self):
q1 = Query.drop_user("myuser")
q2 = Query.drop_user("myuser").if_exists()

self.assertEqual('DROP TABLE "abc"', str(q))
self.assertEqual('DROP USER "myuser"', str(q1))
self.assertEqual('DROP USER IF EXISTS "myuser"', str(q2))

def test_drop_table_if_exists(self):
q = Query.drop_table(self.new_table).if_exists()
def test_drop_view(self):
q1 = Query.drop_view("myview")
q2 = Query.drop_view("myview").if_exists()

self.assertEqual('DROP TABLE IF EXISTS "abc"', str(q))
self.assertEqual('DROP VIEW "myview"', str(q1))
self.assertEqual('DROP VIEW IF EXISTS "myview"', str(q2))

0 comments on commit d24b185

Please sign in to comment.