From 0de929074db26e5ca3d47ab9adfbf9bc3d055f0f Mon Sep 17 00:00:00 2001
From: Varjitt Jeeva <varjitt.jeeva@autodesk.com>
Date: Thu, 9 Nov 2023 14:17:10 -0500
Subject: [PATCH] fix: grant pglogical schema usage #278 (#318)

* fix: grant pglogical schema usage #278

* doc: comment

* doc: comment
---
 pgbelt/cmd/setup.py      | 14 ++++++++++++--
 pgbelt/util/pglogical.py | 17 ++++++++++++++---
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/pgbelt/cmd/setup.py b/pgbelt/cmd/setup.py
index 3991f7c..63754aa 100644
--- a/pgbelt/cmd/setup.py
+++ b/pgbelt/cmd/setup.py
@@ -65,7 +65,12 @@ async def setup(
     try:
         src_logger = get_logger(conf.db, conf.dc, "setup.src")
         dst_logger = get_logger(conf.db, conf.dc, "setup.dst")
-        await configure_pgl(src_root_pool, conf.src.pglogical_user.pw, src_logger)
+        await configure_pgl(
+            src_root_pool,
+            conf.src.pglogical_user.pw,
+            src_logger,
+            conf.src.owner_user.name,
+        )
         await grant_pgl(src_owner_pool, conf.tables, src_logger)
 
         schema_load_task = None
@@ -80,7 +85,12 @@ async def setup(
         if schema_load_task is not None:
             await schema_load_task
 
-        await configure_pgl(dst_root_pool, conf.dst.pglogical_user.pw, dst_logger)
+        await configure_pgl(
+            dst_root_pool,
+            conf.dst.pglogical_user.pw,
+            dst_logger,
+            conf.dst.owner_user.name,
+        )
         await grant_pgl(dst_owner_pool, conf.tables, dst_logger)
         await configure_node(dst_root_pool, "pg2", conf.dst.pglogical_dsn, dst_logger)
 
diff --git a/pgbelt/util/pglogical.py b/pgbelt/util/pglogical.py
index d54661b..9734221 100644
--- a/pgbelt/util/pglogical.py
+++ b/pgbelt/util/pglogical.py
@@ -11,10 +11,12 @@
 from asyncpg.exceptions import UniqueViolationError
 
 
-async def configure_pgl(pool: Pool, pgl_pw: str, logger: Logger) -> None:
+async def configure_pgl(
+    pool: Pool, pgl_pw: str, logger: Logger, owner_user: str
+) -> None:
     """
-    Set up the pglogical role, grant it superuser and replication, and create
-    the extension.
+    Set up the pglogical role, grant it superuser and replication, create
+    the extension and grant USAGE to its schema to the owner user.
     """
     logger.info("Creating pglogical user and extension...")
     async with pool.acquire() as conn:
@@ -54,6 +56,15 @@ async def configure_pgl(pool: Pool, pgl_pw: str, logger: Logger) -> None:
             except DuplicateObjectError:
                 logger.debug("pglogical extension already created")
 
+    # TODO: Somehow test for this working in our integration test.
+    #     We need to make the DBs have a separate schema owner role to test this.
+    async with pool.acquire() as conn:
+        async with conn.transaction():
+            await conn.execute(f"GRANT USAGE ON SCHEMA pglogical TO {owner_user};")
+            logger.debug(
+                f"GRANTed USAGE ON pglogical schema to Schema Owner {owner_user}"
+            )
+
 
 async def grant_pgl(pool: Pool, tables: list[str], logger: Logger) -> None:
     """