pgbelt 0.6.2__py3-none-any.whl → 0.7.1__py3-none-any.whl

pgbelt/cmd/setup.py

@@ -29,15 +29,23 @@ async def _dump_and_load_schema(
 async def _setup_src_node(
     conf: DbupgradeConfig, src_root_pool: Pool, src_logger: Logger
 ) -> None:
+    """
+    Configure the pglogical node and replication set on the Source database.
+    """
+
     await configure_node(src_root_pool, "pg1", conf.src.pglogical_dsn, src_logger)
     async with create_pool(conf.src.pglogical_uri, min_size=1) as src_pglogical_pool:
-        pkey_tables, _, _ = await analyze_table_pkeys(src_pglogical_pool, src_logger)
+        pkey_tables, _, _ = await analyze_table_pkeys(
+            src_pglogical_pool, conf.schema_name, src_logger
+        )
 
     pglogical_tables = pkey_tables
     if conf.tables:
         pglogical_tables = [t for t in pkey_tables if t in conf.tables]
 
-    await configure_replication_set(src_root_pool, pglogical_tables, src_logger)
+    await configure_replication_set(
+        src_root_pool, pglogical_tables, conf.schema_name, src_logger
+    )
 
 
 @run_with_configs
@@ -65,33 +73,40 @@ async def setup(
     try:
         src_logger = get_logger(conf.db, conf.dc, "setup.src")
         dst_logger = get_logger(conf.db, conf.dc, "setup.dst")
+
+        # Configure Source for pglogical (before we can configure the plugin)
         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)
+        await grant_pgl(src_owner_pool, conf.tables, conf.schema_name, src_logger)
 
+        # Load schema into destination
         schema_load_task = None
         if schema:
             schema_load_task = create_task(
                 _dump_and_load_schema(conf, src_logger, dst_logger)
             )
 
+        # Configure Pglogical plugin on Source
         src_node_task = create_task(_setup_src_node(conf, src_root_pool, src_logger))
 
-        # we need to wait for the schema to exist in the target before setting up pglogical there
+        # We need to wait for the schema to exist in the target before setting up pglogical there
         if schema_load_task is not None:
             await schema_load_task
 
+        # Configure Destination for pglogical (before we can configure the plugin)
         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 grant_pgl(dst_owner_pool, conf.tables, conf.schema_name, dst_logger)
+
+        # Also configure the node on the destination... of itself. #TODO: This is a bit weird, confirm if this is necessary.
         await configure_node(dst_root_pool, "pg2", conf.dst.pglogical_dsn, dst_logger)
 
         # The source node must be set up before we create a subscription
@@ -123,14 +138,18 @@ async def setup_back_replication(config_future: Awaitable[DbupgradeConfig]) -> N
 
     try:
         src_logger = get_logger(conf.db, conf.dc, "setup.src")
-        pkeys, _, _ = await analyze_table_pkeys(src_pglogical_pool, src_logger)
+        pkeys, _, _ = await analyze_table_pkeys(
+            src_pglogical_pool, conf.schema_name, src_logger
+        )
         dst_logger = get_logger(conf.db, conf.dc, "setup.src")
 
         pglogical_tables = pkeys
         if conf.tables:
             pglogical_tables = [t for t in pkeys if t in conf.tables]
 
-        await configure_replication_set(dst_root_pool, pglogical_tables, dst_logger)
+        await configure_replication_set(
+            dst_root_pool, pglogical_tables, conf.schema_name, dst_logger
+        )
         await configure_subscription(
             src_root_pool, "pg2_pg1", conf.dst.pglogical_dsn, src_logger
         )

pgbelt/cmd/status.py

@@ -7,6 +7,8 @@ from pgbelt.config.models import DbupgradeConfig
 from pgbelt.util import get_logger
 from pgbelt.util.pglogical import dst_status
 from pgbelt.util.pglogical import src_status
+from pgbelt.util.postgres import initialization_progress
+from pgbelt.util.postgres import analyze_table_pkeys
 from tabulate import tabulate
 from typer import echo
 from typer import style
@@ -22,6 +24,9 @@ async def _print_status_table(results: list[dict[str, str]]) -> list[list[str]]:
             style("flush_lag", "yellow"),
             style("write_lag", "yellow"),
             style("replay_lag", "yellow"),
+            style("src_dataset_size", "yellow"),
+            style("dst_dataset_size", "yellow"),
+            style("progress", "yellow"),
         ]
     ]
 
@@ -41,6 +46,9 @@ async def _print_status_table(results: list[dict[str, str]]) -> list[list[str]]:
                 style(r["flush_lag"], "green" if r["flush_lag"] == "0" else "red"),
                 style(r["write_lag"], "green" if r["write_lag"] == "0" else "red"),
                 style(r["replay_lag"], "green" if r["replay_lag"] == "0" else "red"),
+                style(r["src_dataset_size"], "green"),
+                style(r["dst_dataset_size"], "green"),
+                style(r["progress"], "green"),
             ]
         )
 
@@ -77,14 +85,42 @@ async def status(conf_future: Awaitable[DbupgradeConfig]) -> dict[str, str]:
     )
     src_pool, dst_pool = pools
 
+    # Get the list of targeted tables by first getting all tables, then filtering whatever is in the config.
+    pkey_tables, non_pkey_tables, _ = await analyze_table_pkeys(
+        src_pool, conf.schema_name, src_logger
+    )
+    all_tables = pkey_tables + non_pkey_tables
+    target_tables = all_tables
+    if conf.tables:
+        target_tables = [t for t in all_tables if t in conf.tables]
+
+    if not target_tables:
+        raise ValueError(
+            f"Targeted tables not found in the source database. Please check your config's schema and tables. DB: {conf.db} DC: {conf.dc}, SCHEMA: {conf.schema_name} TABLES: {conf.tables}."
+        )
+
     try:
         result = await gather(
             src_status(src_pool, src_logger),
             dst_status(dst_pool, dst_logger),
+            initialization_progress(
+                target_tables,
+                conf.schema_name,
+                conf.schema_name,
+                src_pool,
+                dst_pool,
+                src_logger,
+                dst_logger,
+            ),
         )
 
         result[0].update(result[1])
         result[0]["db"] = conf.db
+        if result[0]["pg1_pg2"] == "replicating":
+            result[2]["src_dataset_size"] = "n/a"
+            result[2]["dst_dataset_size"] = "n/a"
+            result[2]["progress"] = "n/a"
+        result[0].update(result[2])
         return result[0]
     finally:
         await gather(*[p.close() for p in pools])

pgbelt/cmd/sync.py

@@ -22,12 +22,14 @@ from typer import Option
 
 async def _sync_sequences(
     targeted_sequences: list[str],
+    schema: str,
     src_pool: Pool,
     dst_pool: Pool,
     src_logger: Logger,
     dst_logger: Logger,
 ) -> None:
-    seq_vals = await dump_sequences(src_pool, targeted_sequences, src_logger)
+
+    seq_vals = await dump_sequences(src_pool, targeted_sequences, schema, src_logger)
     await load_sequences(dst_pool, seq_vals, dst_logger)
 
 
@@ -47,7 +49,7 @@ async def sync_sequences(config_future: Awaitable[DbupgradeConfig]) -> None:
         src_logger = get_logger(conf.db, conf.dc, "sync.src")
         dst_logger = get_logger(conf.db, conf.dc, "sync.dst")
         await _sync_sequences(
-            conf.sequences, src_pool, dst_pool, src_logger, dst_logger
+            conf.sequences, conf.schema_name, src_pool, dst_pool, src_logger, dst_logger
         )
     finally:
         await gather(*[p.close() for p in pools])
@@ -72,7 +74,7 @@ async def dump_tables(
         tables = tables.split(",")
     else:
         async with create_pool(conf.src.pglogical_uri, min_size=1) as src_pool:
-            _, tables, _ = await analyze_table_pkeys(src_pool, logger)
+            _, tables, _ = await analyze_table_pkeys(src_pool, conf.schema_name, logger)
 
         if conf.tables:
             tables = [t for t in tables if t in conf.tables]
@@ -127,7 +129,9 @@ async def sync_tables(
         dump_tables = tables.split(",")
     else:
         async with create_pool(conf.src.pglogical_uri, min_size=1) as src_pool:
-            _, dump_tables, _ = await analyze_table_pkeys(src_pool, src_logger)
+            _, dump_tables, _ = await analyze_table_pkeys(
+                src_pool, conf.schema_name, src_logger
+            )
 
         if conf.tables:
             dump_tables = [t for t in dump_tables if t in conf.tables]
@@ -167,8 +171,10 @@ async def validate_data(config_future: Awaitable[DbupgradeConfig]) -> None:
     try:
         logger = get_logger(conf.db, conf.dc, "sync")
         await gather(
-            compare_100_rows(src_pool, dst_pool, conf.tables, logger),
-            compare_latest_100_rows(src_pool, dst_pool, conf.tables, logger),
+            compare_100_rows(src_pool, dst_pool, conf.tables, conf.schema_name, logger),
+            compare_latest_100_rows(
+                src_pool, dst_pool, conf.tables, conf.schema_name, logger
+            ),
         )
     finally:
         await gather(*[p.close() for p in pools])
@@ -177,7 +183,7 @@ async def validate_data(config_future: Awaitable[DbupgradeConfig]) -> None:
 async def _dump_and_load_all_tables(
     conf: DbupgradeConfig, src_pool: Pool, src_logger: Logger, dst_logger: Logger
 ) -> None:
-    _, tables, _ = await analyze_table_pkeys(src_pool, src_logger)
+    _, tables, _ = await analyze_table_pkeys(src_pool, conf.schema_name, src_logger)
     if conf.tables:
         tables = [t for t in tables if t in conf.tables]
     await dump_source_tables(conf, tables, src_logger)
@@ -185,7 +191,9 @@ async def _dump_and_load_all_tables(
 
 
 @run_with_configs
-async def sync(config_future: Awaitable[DbupgradeConfig]) -> None:
+async def sync(
+    config_future: Awaitable[DbupgradeConfig], no_schema: bool = False
+) -> None:
     """
     Sync and validate all data that is not replicated with pglogical. This includes all
     tables without primary keys and all sequences. Also loads any previously omitted
@@ -210,7 +218,12 @@ async def sync(config_future: Awaitable[DbupgradeConfig]) -> None:
 
         await gather(
             _sync_sequences(
-                conf.sequences, src_pool, dst_root_pool, src_logger, dst_logger
+                conf.sequences,
+                conf.schema_name,
+                src_pool,
+                dst_root_pool,
+                src_logger,
+                dst_logger,
             ),
             _dump_and_load_all_tables(conf, src_pool, src_logger, dst_logger),
         )
@@ -218,15 +231,27 @@ async def sync(config_future: Awaitable[DbupgradeConfig]) -> None:
         # Creating indexes should run before validations and ANALYZE, but after all the data exists
         # in the destination database.
 
-        await gather(
-            apply_target_constraints(conf, dst_logger),
-            create_target_indexes(conf, dst_logger, during_sync=True),
-        )
+        # Do not load NOT VALID constraints or create INDEXes for exodus-style migrations
+        if not no_schema:
+            await gather(
+                apply_target_constraints(conf, dst_logger),
+                create_target_indexes(conf, dst_logger, during_sync=True),
+            )
 
         await gather(
-            compare_100_rows(src_pool, dst_owner_pool, conf.tables, validation_logger),
+            compare_100_rows(
+                src_pool,
+                dst_owner_pool,
+                conf.tables,
+                conf.schema_name,
+                validation_logger,
+            ),
             compare_latest_100_rows(
-                src_pool, dst_owner_pool, conf.tables, validation_logger
+                src_pool,
+                dst_owner_pool,
+                conf.tables,
+                conf.schema_name,
+                validation_logger,
             ),
             run_analyze(dst_owner_pool, dst_logger),
         )

pgbelt/cmd/teardown.py

@@ -80,8 +80,8 @@ async def teardown(
             teardown_node(dst_root_pool, "pg2", dst_logger),
         )
         await gather(
-            revoke_pgl(src_root_pool, conf.tables, src_logger),
-            revoke_pgl(dst_root_pool, conf.tables, dst_logger),
+            revoke_pgl(src_root_pool, conf.tables, conf.schema_name, src_logger),
+            revoke_pgl(dst_root_pool, conf.tables, conf.schema_name, dst_logger),
         )
 
         if full:

pgbelt/config/models.py

@@ -49,7 +49,7 @@ class DbConfig(BaseModel):
     db: str The dbname to operate on. If you want to migrate multiple dbs in a single instance set up a separate config.
     port: str The port to connect to.
     root_user: User A superuser. Usually the postgres user.
-    owner_user: User A user who owns all the data in the public schema or who has equivalent permissions. # noqa: RST301
+    owner_user: User A user who owns all the data in the your specified schema or who has equivalent permissions. # noqa: RST301
                      This user will end up owning all the data if this is describing the target instance.
     pglogical_user: User A user for use with pglogical. Will be created if it does not exist.
     other_users: list[User] A list of other users whose passwords we might not know.
@@ -105,6 +105,9 @@ class DbupgradeConfig(BaseModel):
     dc: str A name used to identify the environment this database pair is in. Used in cli commands.
     src: DbConfig The database we are moving data out of.
     dst: DbConfig The database we are moving data into.
+    tables: Optional[list[str]] A list of tables to replicate. If not provided all tables in the named schema will be replicated.
+    sequences: Optional[list[str]] A list of sequences to replicate. If not provided all sequences in the named schema will be replicated.
+    schema_name: Optional[str] The schema to operate on. Defaults to "public".
     """
 
     db: str
@@ -113,6 +116,7 @@ class DbupgradeConfig(BaseModel):
     dst: Optional[DbConfig] = None
     tables: Optional[list[str]] = None
     sequences: Optional[list[str]] = None
+    schema_name: Optional[str] = "public"
 
     _not_empty = validator("db", "dc", allow_reuse=True)(not_empty)
 

pgbelt/util/dump.py

@@ -93,7 +93,7 @@ async def dump_source_tables(
                 [
                     "pg_dump",
                     "--data-only",
-                    f"--table={table}",
+                    f"--table={config.schema_name}.{table}",
                     "-Fc",
                     "-f",
                     table_file(config.db, config.dc, table),
@@ -125,7 +125,7 @@ async def load_dumped_tables(
     async with create_pool(config.dst.root_uri, min_size=1) as pool:
         to_load = []
         for t in tables:
-            if await table_empty(pool, t, logger):
+            if await table_empty(pool, t, config.schema_name, logger):
                 to_load.append(table_file(config.db, config.dc, t))
             else:
                 logger.warning(
@@ -164,10 +164,14 @@ async def dump_source_schema(config: DbupgradeConfig, logger: Logger) -> None:
         "--schema-only",
         "--no-owner",
         "-n",
-        "public",
+        config.schema_name,
         config.src.pglogical_dsn,
     ]
 
+    # TODO: We should exclude the creation of a schema in the schema dump and load, and made that the responsibility of the user.
+    # Confirm if the CREATE SCHEMA statement is included in the schema dump, and if yes, exclude it.
+    # This will reveal itself in the integration test.
+
     out = await _execute_subprocess(command, "Retrieved source schema", logger)
 
     commands_raw = _parse_dump_commands(out.decode("utf-8"))
@@ -239,7 +243,7 @@ async def dump_dst_not_valid_constraints(
         "--schema-only",
         "--no-owner",
         "-n",
-        "public",
+        config.schema_name,
         config.dst.pglogical_dsn,
     ]
 
@@ -345,7 +349,7 @@ async def dump_dst_create_index(config: DbupgradeConfig, logger: Logger) -> None
         "--schema-only",
         "--no-owner",
         "-n",
-        "public",
+        config.schema_name,
         config.dst.pglogical_dsn,
     ]
 

pgbelt/util/pglogical.py

@@ -66,7 +66,7 @@ async def configure_pgl(
             )
 
 
-async def grant_pgl(pool: Pool, tables: list[str], logger: Logger) -> None:
+async def grant_pgl(pool: Pool, tables: list[str], schema: str, logger: Logger) -> None:
     """
     Grant pglogical access to the data
 
@@ -78,21 +78,22 @@ async def grant_pgl(pool: Pool, tables: list[str], logger: Logger) -> None:
     async with pool.acquire() as conn:
         async with conn.transaction():
             if tables:
+                tables_with_schema = [f"{schema}.{table}" for table in tables]
                 await conn.execute(
-                    f"GRANT ALL ON TABLE {','.join(tables)} TO pglogical;"
+                    f"GRANT ALL ON TABLE {','.join(tables_with_schema)} TO pglogical;"
                 )
             else:
                 await conn.execute(
-                    "GRANT ALL ON ALL TABLES IN SCHEMA public TO pglogical;"
+                    f"GRANT ALL ON ALL TABLES IN SCHEMA {schema} TO pglogical;"
                 )
             await conn.execute(
-                "GRANT ALL ON ALL SEQUENCES IN SCHEMA public TO pglogical;"
+                f"GRANT ALL ON ALL SEQUENCES IN SCHEMA {schema} TO pglogical;"
             )
             logger.debug("pglogical data grants complete")
 
 
 async def configure_replication_set(
-    pool: Pool, tables: list[str], logger: Logger
+    pool: Pool, tables: list[str], schema: str, logger: Logger
 ) -> None:
     """
     Add each table in the given list to the default replication set
@@ -105,17 +106,23 @@ async def configure_replication_set(
         except Exception as e:
             logger.debug(f"Could not create replication set 'pgbelt': {e}")
 
-    logger.info(f"Configuring 'pgbelt' replication set with tables: {tables}")
+    logger.info(
+        f"Configuring 'pgbelt' replication set with tables from schema {schema}: {tables}"
+    )
     for table in tables:
         async with pool.acquire() as conn:
             async with conn.transaction():
                 try:
                     await conn.execute(
-                        f"SELECT pglogical.replication_set_add_table('pgbelt', '\"{table}\"');"
+                        f"SELECT pglogical.replication_set_add_table('pgbelt', '\"{schema}\".\"{table}\"');"
+                    )
+                    logger.debug(
+                        f"Table '{table}' added to 'pgbelt' replication set from schema {schema}"
                     )
-                    logger.debug(f"Table '{table}' added to 'pgbelt' replication set")
                 except UniqueViolationError:
-                    logger.debug(f"Table '{table}' already in 'pgbelt' replication set")
+                    logger.debug(
+                        f"Table '{table}' already in 'pgbelt' replication set from schema {schema}"
+                    )
 
 
 async def configure_node(pool: Pool, name: str, dsn: str, logger: Logger) -> None:
@@ -202,7 +209,7 @@ async def teardown_replication_set(pool: Pool, logger: Logger) -> None:
     """
     Tear down the replication_set
     """
-    logger.info("Dropping replication set 'default'...")
+    logger.info("Dropping replication set 'pgbelt'...")
     async with pool.acquire() as conn:
         async with conn.transaction():
             try:
@@ -213,10 +220,16 @@ async def teardown_replication_set(pool: Pool, logger: Logger) -> None:
                 UndefinedFunctionError,
                 InternalServerError,
             ):
-                logger.debug("Replication set 'default' does not exist")
+                logger.debug("Replication set 'pgbelt' does not exist")
+            except ObjectNotInPrerequisiteStateError:
+                logger.debug(
+                    "pglogical node was already dropped, so we can't drop the replication set. This is okay, keep going."
+                )
 
 
-async def revoke_pgl(pool: Pool, tables: list[str], logger: Logger) -> None:
+async def revoke_pgl(
+    pool: Pool, tables: list[str], schema: str, logger: Logger
+) -> None:
     """
     Revoke data access permissions from pglogical, and drop the pglogical role
     """
@@ -225,10 +238,10 @@ async def revoke_pgl(pool: Pool, tables: list[str], logger: Logger) -> None:
         async with conn.transaction():
             try:
                 await conn.execute(
-                    "REVOKE ALL ON ALL TABLES IN SCHEMA public FROM pglogical;"
+                    f"REVOKE ALL ON ALL TABLES IN SCHEMA {schema} FROM pglogical;"
                 )
                 await conn.execute(
-                    "REVOKE ALL ON ALL SEQUENCES IN SCHEMA public FROM pglogical;"
+                    f"REVOKE ALL ON ALL SEQUENCES IN SCHEMA {schema} FROM pglogical;"
                 )
                 logger.debug("Data access permissions revoked")
             except UndefinedObjectError as e: