From 573d7ead6cf2ff8e8382a0674453354be37419f1 Mon Sep 17 00:00:00 2001
From: minhtrannhat <minh@minhtrannhat.com>
Date: Fri, 1 Mar 2024 01:33:21 -0500
Subject: [PATCH] fix(db): resolve pg role issues

- Seperate posgres user for development and testing
---
 backend/development.env                    |  2 +-
 backend/pyproject.toml                     |  2 +-
 backend/src/neo_neo_todo/migrations/0.py   |  3 ++-
 backend/src/neo_neo_todo/utils/database.py | 18 +++++++++++++++---
 backend/testing.env                        |  2 +-
 5 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/backend/development.env b/backend/development.env
index df2bd3d..670ad37 100644
--- a/backend/development.env
+++ b/backend/development.env
@@ -1,4 +1,4 @@
 UVICORN_PORT="5050" 
 TODO_DEBUG=true
 TODO_SECRET_KEY="secret key"
-TODO_DB_DATABASE_URL="postgresql://todo:todo@0.0.0.0:5432/todo"
+TODO_DB_DATABASE_URL="postgresql://todo_dev:todo@0.0.0.0:5432/todo_dev"
diff --git a/backend/pyproject.toml b/backend/pyproject.toml
index d1ed66a..db3adc0 100644
--- a/backend/pyproject.toml
+++ b/backend/pyproject.toml
@@ -66,4 +66,4 @@ dev = {cmd = "uvicorn neo_neo_todo.main:app --reload", env_file = "development.e
 recreate-db-base = "python3 src/neo_neo_todo/utils/database.py"
 migration-base = "python3 src/neo_neo_todo/migrations/0.py"
 recreate-db = {composite = ["recreate-db-base", "migration-base"], env_file = "development.env"}
-test = {composite = ["recreate-db-base", "pytest tests/"], env_file = "testing.env"}
+test = {composite = ["recreate-db-base", "migration-base", "pytest tests/"], env_file = "testing.env"}
diff --git a/backend/src/neo_neo_todo/migrations/0.py b/backend/src/neo_neo_todo/migrations/0.py
index a75cab6..60c0ef3 100644
--- a/backend/src/neo_neo_todo/migrations/0.py
+++ b/backend/src/neo_neo_todo/migrations/0.py
@@ -21,7 +21,8 @@ def migrate() -> None:
     port = url_parts.port
 
     print(
-        f"The database name is {dbname}, the username is {user_name}, the password is {password}, the host is {host}, the port is {port}"
+        f"The database name is {dbname}, the username is {user_name}, "
+        f"the password is {password}, the host is {host}, the port is {port}"
     )
 
     with psycopg.connect(
diff --git a/backend/src/neo_neo_todo/utils/database.py b/backend/src/neo_neo_todo/utils/database.py
index 601f08f..d4d2c78 100644
--- a/backend/src/neo_neo_todo/utils/database.py
+++ b/backend/src/neo_neo_todo/utils/database.py
@@ -21,7 +21,9 @@ def recreate_db() -> None:
     port = url_parts.port
 
     print(
-        f"The database name is {dbname}, the username is {user_name}, the password is {password}, the host is {host}, the port is {port}"
+        "Running database recreation script, "
+        f"The database name is {dbname}, the username is {user_name}, "
+        f"the password is {password}, the host is {host}, the port is {port}"
     )
 
     with psycopg.connect(
@@ -31,8 +33,17 @@ def recreate_db() -> None:
             conn._set_autocommit(True)
 
             cur.execute(
-                sql.SQL("DROP OWNED BY {}").format(
-                    sql.Identifier(user_name)  # type: ignore
+                sql.SQL(
+                    """
+                    DO $$
+                    BEGIN
+                        IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = '{}') THEN
+                            DROP OWNED BY {};
+                        END IF;
+                    END $$;
+                    """
+                ).format(
+                    sql.Identifier(user_name), sql.Identifier(user_name)  # type: ignore
                 )
             )
 
@@ -42,6 +53,7 @@ def recreate_db() -> None:
                     dbname=sql.Identifier(dbname),
                 )
             )
+
             # Drop the user if it exists
             cur.execute(
                 sql.SQL("DROP USER IF EXISTS {}").format(
diff --git a/backend/testing.env b/backend/testing.env
index e04212b..2a293e8 100644
--- a/backend/testing.env
+++ b/backend/testing.env
@@ -2,4 +2,4 @@ UVICORN_PORT="5050"
 TODO_DEBUG=true
 TODO_SECRET_KEY="secret key" 
 TODO_TESTING=true
-TODO_DB_DATABASE_URL="postgresql://todo:todo@0.0.0.0:5432/todo_test"
+TODO_DB_DATABASE_URL="postgresql://todo_test:todo@0.0.0.0:5432/todo_test"