From b325ac5929944f296b51d47639d8a364ad087529 Mon Sep 17 00:00:00 2001
From: Gaurav Tarlok Kakkar <gaurav21776@gmail.com>
Date: Sun, 10 Sep 2023 23:09:15 -0400
Subject: [PATCH] feat: add support for drop database

---
 evadb/catalog/catalog_manager.py              |  4 +--
 evadb/executor/drop_object_executor.py        | 26 ++++++++++++++
 evadb/parser/lark_visitor/_drop_statement.py  | 14 ++++++++
 evadb/parser/types.py                         |  1 +
 evadb/parser/utils.py                         |  4 +++
 .../short/test_drop_executor.py               | 36 +++++++++++++++++++
 6 files changed, 82 insertions(+), 3 deletions(-)

diff --git a/evadb/catalog/catalog_manager.py b/evadb/catalog/catalog_manager.py
index 83ec869001..c5fd50f226 100644
--- a/evadb/catalog/catalog_manager.py
+++ b/evadb/catalog/catalog_manager.py
@@ -160,9 +160,7 @@ def get_database_catalog_entry(self, database_name: str) -> DatabaseCatalogEntry
 
         return table_entry
 
-    def delete_database_catalog_entry(
-        self, database_entry: DatabaseCatalogEntry
-    ) -> bool:
+    def drop_database_catalog_entry(self, database_entry: DatabaseCatalogEntry) -> bool:
         """
         This method deletes the database from  catalog.
 
diff --git a/evadb/executor/drop_object_executor.py b/evadb/executor/drop_object_executor.py
index 7a56674fe6..38d5419dc4 100644
--- a/evadb/executor/drop_object_executor.py
+++ b/evadb/executor/drop_object_executor.py
@@ -43,6 +43,9 @@ def exec(self, *args, **kwargs):
         elif self.node.object_type == ObjectType.FUNCTION:
             yield self._handle_drop_function(self.node.name, self.node.if_exists)
 
+        elif self.node.object_type == ObjectType.DATABASE:
+            yield self._handle_drop_database(self.node.name, self.node.if_exists)
+
     def _handle_drop_table(self, table_name: str, if_exists: bool):
         if not self.catalog().check_table_exists(table_name):
             err_msg = "Table: {} does not exist".format(table_name)
@@ -132,3 +135,26 @@ def _handle_drop_index(self, index_name: str, if_exists: bool):
                     index=[0],
                 )
             )
+
+    def _handle_drop_database(self, database_name: str, if_exists: bool):
+        db_catalog_entry = self.catalog().get_database_catalog_entry(database_name)
+        if not db_catalog_entry:
+            err_msg = (
+                f"Database {database_name} does not exist, therefore cannot be dropped."
+            )
+            if if_exists:
+                logger.warning(err_msg)
+                return Batch(pd.DataFrame([err_msg]))
+            else:
+                raise RuntimeError(err_msg)
+
+        logger.debug(f"Dropping database {database_name}")
+
+        self.catalog().drop_database_catalog_entry(db_catalog_entry)
+
+        return Batch(
+            pd.DataFrame(
+                {f"Database {database_name} successfully dropped"},
+                index=[0],
+            )
+        )
diff --git a/evadb/parser/lark_visitor/_drop_statement.py b/evadb/parser/lark_visitor/_drop_statement.py
index fbf922fb36..0b397378ae 100644
--- a/evadb/parser/lark_visitor/_drop_statement.py
+++ b/evadb/parser/lark_visitor/_drop_statement.py
@@ -59,3 +59,17 @@ def drop_function(self, tree):
                     if_exists = True
 
         return DropObjectStatement(ObjectType.FUNCTION, function_name, if_exists)
+
+    # Drop Database
+    def drop_database(self, tree):
+        database_name = None
+        if_exists = False
+
+        for child in tree.children:
+            if isinstance(child, Tree):
+                if child.data == "if_exists":
+                    if_exists = True
+                elif child.data == "uid":
+                    database_name = self.visit(child)
+
+        return DropObjectStatement(ObjectType.DATABASE, database_name, if_exists)
diff --git a/evadb/parser/types.py b/evadb/parser/types.py
index 0abebcb097..a57c938db8 100644
--- a/evadb/parser/types.py
+++ b/evadb/parser/types.py
@@ -79,3 +79,4 @@ class ObjectType(EvaDBEnum):
     TABLE  # noqa: F821
     FUNCTION  # noqa: F821
     INDEX  # noqa: F821
+    DATABASE  # noqa: F821
diff --git a/evadb/parser/utils.py b/evadb/parser/utils.py
index 70db55cecc..3ad9b032f1 100644
--- a/evadb/parser/utils.py
+++ b/evadb/parser/utils.py
@@ -149,6 +149,10 @@ def parse_drop_index(index_name: str, if_exists: bool):
     return parse_drop(ObjectType.INDEX, index_name, if_exists)
 
 
+def parse_drop_database(database_name: str, if_exists: bool):
+    return parse_drop(ObjectType.DATABASE, database_name, if_exists)
+
+
 def parse_query(query):
     stmt = Parser().parse(query)
     assert len(stmt) == 1
diff --git a/test/integration_tests/short/test_drop_executor.py b/test/integration_tests/short/test_drop_executor.py
index fb5fd4339b..a5e19ea536 100644
--- a/test/integration_tests/short/test_drop_executor.py
+++ b/test/integration_tests/short/test_drop_executor.py
@@ -191,3 +191,39 @@ def test_should_drop_index(self):
         self.assertTrue(index_obj is None)
 
         # todo check if the index is also removed from the underlying vector store
+
+    #### DROP INDEX
+
+    def test_should_drop_database(self):
+        # Create database.
+        database_name = "test_data_source"
+        params = {
+            "database": "evadb.db",
+        }
+        query = f"""CREATE DATABASE {database_name}
+                    WITH ENGINE = "sqlite",
+                    PARAMETERS = {params};"""
+        execute_query_fetch_all(self.evadb, query)
+        self.assertIsNotNone(
+            self.evadb.catalog().get_database_catalog_entry(database_name)
+        )
+
+        # DROP DATABASE
+        execute_query_fetch_all(self.evadb, f"DROP DATABASE {database_name}")
+        self.assertIsNone(
+            self.evadb.catalog().get_database_catalog_entry(database_name)
+        )
+
+        # DROP should pass with warning
+        result = execute_query_fetch_all(
+            self.evadb, f"DROP DATABASE IF EXISTS {database_name}"
+        )
+        self.assertTrue("does not exist" in result.frames.to_string())
+
+        # DROP should throw error
+        with self.assertRaises(ExecutorError):
+            execute_query_fetch_all(
+                self.evadb,
+                f"DROP DATABASE {database_name}",
+                do_not_print_exceptions=True,
+            )