Refactor

src/test/regress/citus_tests/common.py

@@ -753,6 +753,7 @@ class Postgres(QueryRunner):
         self.subscriptions = set()
         self.publications = set()
         self.replication_slots = set()
+        self.databases = set()
         self.schemas = set()
         self.users = set()
 
@@ -993,6 +994,10 @@ class Postgres(QueryRunner):
             args = sql.SQL("")
         self.sql(sql.SQL("CREATE USER {} {}").format(sql.Identifier(name), args))
 
+    def create_database(self, name):
+        self.databases.add(name)
+        self.sql(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name)))
+
     def create_schema(self, name):
         self.schemas.add(name)
         self.sql(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(name)))
@@ -1020,6 +1025,12 @@ class Postgres(QueryRunner):
         for user in self.users:
             self.sql(sql.SQL("DROP USER IF EXISTS {}").format(sql.Identifier(user)))
 
+    def cleanup_databases(self):
+        for database in self.databases:
+            self.sql(
+                sql.SQL("DROP DATABASE IF EXISTS {}").format(sql.Identifier(database))
+            )
+
     def cleanup_schemas(self):
         for schema in self.schemas:
             self.sql(

src/test/regress/citus_tests/test/test_maintenancedeamon.py

@@ -3,8 +3,9 @@
 # is launched. This should happen even if there is no query run
 # in main_db yet.
 def test_set_maindb(coord):
+    coord.create_database("mymaindb")
+
     with coord.cur() as cur1:
-        cur1.execute("CREATE DATABASE mymaindb;")
         cur1.execute("ALTER SYSTEM SET citus.main_db='mymaindb'")
         cur1.execute("SELECT pg_reload_conf();")
         coord.restart()
@@ -25,4 +26,4 @@ def test_set_maindb(coord):
             == 1
         )
 
-        coord.sql("DROP DATABASE mymaindb;")
+    coord.cleanup_databases()