Repository: kayak/pypika
Pull request #472 (2020-08-06 19:05:32)

Added drop statement and improved create statement

mikeengland: Thanks for the contribution @nibrag. Would you mind adding some documentation for create/drop table as there is currently no documentation?

Aside from that the code looks good.



Patch:
diff --git a/docs/3_advanced.rst b/docs/3_advanced.rst
index b1dd51d..d6ae1e5 100644
--- a/docs/3_advanced.rst
+++ b/docs/3_advanced.rst
@@ -176,4 +176,92 @@ following the current row by omitting the ``constant_value`` parameter like ``an
 
     SELECT SUM("total") OVER(PARTITION BY "item_id" ORDER BY "day" ROWS BETWEEN 7 PRECEDING AND CURRENT ROW) FROM "t_customers"
 
+DDL
+===
 
+CREATE TABLE
+------------
+
+Simple example of creating table:
+
+.. code-block:: python
+
+    from pypika import Query, Table, Columns
+
+    columns = Columns(
+        ('id', 'integer'),
+        ('price', 'decimal(15, 9)'),
+        ('name', 'varchar(128)')
+    )
+
+    Query.create_table(Table('items')).columns(*columns).if_not_exists()
+
+.. code-block:: sql
+
+    CREATE TABLE IF NOT EXISTS items (
+        id integer,
+        price decimal(15, 9),
+        name varchar(128)
+    )
+
+Create table from query:
+
+.. code-block:: python
+
+    from pypika import Query, Table
+
+    items = Table('items')
+
+    query = Query.select(items.id, items.price, items.name).from_(items)
+
+    Query.create_table(Table('items')).as_select(query)
+
+.. code-block:: sql
+
+    CREATE TABLE items
+    AS
+        (SELECT id, price, name FROM items)
+
+TEMPORARY and UNLOGGED tables:
+
+.. code-block:: python
+
+    from pypika import Query, Table, Columns
+
+    columns = Columns(
+        ('id', 'integer'),
+        ('price', 'decimal(15, 9)'),
+        ('name', 'varchar(128)')
+    )
+
+    Query.create_table(Table('items')).columns(*columns).temporary()
+    Query.create_table(Table('items')).columns(*columns).unlogged()
+
+.. code-block:: sql
+
+    CREATE TEMPORARY TABLE items (
+        id integer,
+        price decimal(15, 9),
+        name varchar(128)
+    )
+
+    CREATE UNLOGGED TABLE items (
+        id integer,
+        price decimal(15, 9),
+        name varchar(128)
+    )
+
+DROP TABLE
+----------
+
+.. code-block:: python
+
+    from pypika import Query, Table
+
+    Query.drop_table(Table('items'))
+    Query.drop_table(Table('items')).if_exists()
+
+.. code-block:: sql
+
+    DROP TABLE items
+    DROP TABLE IF EXISTS items
diff --git a/pypika/queries.py b/pypika/queries.py
index 03a3ad1..0ebb74c 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -384,6 +384,18 @@ class Query:
         """
         return CreateQueryBuilder().create_table(table)
 
+    @classmethod
+    def drop_table(cls, table: Union[str, 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 table: An instance of a Table object or a string table name.
+
+        :return: DropQueryBuilder
+        """
+        return DropQueryBuilder().drop_table(table)
+
     @classmethod
     def into(cls, table: Union[Table, str], **kwargs: Any) -> "QueryBuilder":
         """
@@ -1637,12 +1649,14 @@ class CreateQueryBuilder:
     def __init__(self, dialect: Optional[Dialects] = None) -> None:
         self._create_table = None
         self._temporary = False
+        self._unlogged = False
         self._as_select = None
         self._columns = []
         self._period_fors = []
         self._with_system_versioning = False
         self._primary_key = None
         self._uniques = []
+        self._if_not_exists = False
         self.dialect = dialect
 
     def _set_kwargs_defaults(self, kwargs: dict) -> None:
@@ -1679,6 +1693,16 @@ class CreateQueryBuilder:
         """
         self._temporary = True
 
+    @builder
+    def unlogged(self) -> "CreateQueryBuilder":
+        """
+        Makes the table unlogged.
+
+        :return:
+            CreateQueryBuilder.
+        """
+        self._unlogged = True
+
     @builder
     def with_system_versioning(self) -> "CreateQueryBuilder":
         """
@@ -1793,6 +1817,10 @@ class CreateQueryBuilder:
 
         self._as_select = query_builder
 
+    @builder
+    def if_not_exists(self) -> "CreateQueryBuilder":
+        self._if_not_exists = True
+
     def get_sql(self, **kwargs: Any) -> str:
         """
         Gets the sql statement string.
@@ -1821,8 +1849,19 @@ class CreateQueryBuilder:
         )
 
     def _create_table_sql(self, **kwargs: Any) -> str:
-        return "CREATE {temporary}TABLE {table}".format(
-            temporary="TEMPORARY " if self._temporary else "",
+        table_type = ''
+        if self._temporary:
+            table_type = 'TEMPORARY '
+        elif self._unlogged:
+            table_type = 'UNLOGGED '
+
+        if_not_exists = ''
+        if self._if_not_exists:
+            if_not_exists = 'IF NOT EXISTS '
+
+        return "CREATE {table_type}TABLE {if_not_exists}{table}".format(
+            table_type=table_type,
+            if_not_exists=if_not_exists,
             table=self._create_table.get_sql(**kwargs),
         )
 
@@ -1872,3 +1911,60 @@ class CreateQueryBuilder:
 
     def __repr__(self) -> str:
         return self.__str__()
+
+
+class DropQueryBuilder:
+    """
+    Query builder used to build DROP queries.
+    """
+
+    QUOTE_CHAR = '"'
+    SECONDARY_QUOTE_CHAR = "'"
+    ALIAS_QUOTE_CHAR = None
+
+    def __init__(self, dialect: Optional[Dialects] = None) -> None:
+        self._drop_table = None
+        self._if_exists = None
+        self.dialect = dialect
+
+    def _set_kwargs_defaults(self, kwargs: dict) -> None:
+        kwargs.setdefault("quote_char", self.QUOTE_CHAR)
+        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_table(self, table: Union[Table, str]) -> "DropQueryBuilder":
+        if self._drop_table:
+            raise AttributeError("'Query' object already has attribute drop_table")
+
+        self._drop_table = table if isinstance(table, Table) else Table(table)
+
+    @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 '
+
+        return "DROP TABLE {if_exists}{table}".format(
+            if_exists=if_exists,
+            table=self._drop_table.get_sql(**kwargs),
+        )
+
+    def __str__(self) -> str:
+        return self.get_sql()
+
+    def __repr__(self) -> str:
+        return self.__str__()


Test patch:
diff --git a/pypika/tests/test_create.py b/pypika/tests/test_create.py
index 17b1770..f51dacd 100644
--- a/pypika/tests/test_create.py
+++ b/pypika/tests/test_create.py
@@ -75,6 +75,16 @@ class CreateTableTests(unittest.TestCase):
 
             self.assertEqual('CREATE TABLE "abc" ("a" INT,"b" VARCHAR(100)) WITH SYSTEM VERSIONING', str(q))
 
+        with self.subTest("with unlogged keyword"):
+            q = Query.create_table(self.new_table).unlogged().columns(self.foo, self.bar)
+
+            self.assertEqual('CREATE UNLOGGED TABLE "abc" ("a" INT,"b" VARCHAR(100))', str(q))
+
+        with self.subTest("with if not exists keyword"):
+            q = Query.create_table(self.new_table).if_not_exists().columns(self.foo, self.bar)
+
+            self.assertEqual('CREATE TABLE IF NOT EXISTS "abc" ("a" INT,"b" VARCHAR(100))', str(q))
+
     def test_create_table_with_select(self):
         select = Query.from_(self.existing_table).select(self.existing_table.foo, self.existing_table.bar)
 
diff --git a/pypika/tests/test_drop.py b/pypika/tests/test_drop.py
new file mode 100644
index 0000000..e887359
--- /dev/null
+++ b/pypika/tests/test_drop.py
@@ -0,0 +1,22 @@
+import unittest
+
+from pypika import (
+    Tables,
+    Columns,
+    Query,
+)
+
+
+class DropTableTests(unittest.TestCase):
+    new_table, existing_table = Tables("abc", "efg")
+    foo, bar = Columns(("a", "INT"), ("b", "VARCHAR(100)"))
+
+    def test_drop_table(self):
+        q = Query.drop_table(self.new_table)
+
+        self.assertEqual('DROP TABLE "abc"', str(q))
+
+    def test_drop_table_if_exists(self):
+        q = Query.drop_table(self.new_table).if_exists()
+
+        self.assertEqual('DROP TABLE IF EXISTS "abc"', str(q))