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

pgbelt/util/postgres.py

@@ -6,25 +6,44 @@ from asyncpg.exceptions import UndefinedObjectError
 
 
 async def dump_sequences(
-    pool: Pool, targeted_sequences: list[str], logger: Logger
+    pool: Pool, targeted_sequences: list[str], schema: str, logger: Logger
 ) -> dict[str, int]:
     """
     return a dictionary of sequence names mapped to their last values
     """
     logger.info("Dumping sequence values...")
-    seqs = await pool.fetch("SELECT sequence_name FROM information_schema.sequences;")
+    # Get all sequences in the schema
+    seqs = await pool.fetch(
+        f"""
+        SELECT '{schema}' || '.' || sequence_name
+        FROM information_schema.sequences
+        WHERE sequence_schema = '{schema}';
+        """
+    )
+
+    # Note: When in an exodus migration with a non-public schema, the sequence names must be prefixed with the schema name.
+    # This may not be done by the user, so we must do it here.
+    proper_sequence_names = None
+    if targeted_sequences is not None:
+        proper_sequence_names = []
+        for seq in targeted_sequences:
+            if f"{schema}." not in seq:
+                proper_sequence_names.append(f"{schema}.{seq}")
+            else:
+                proper_sequence_names.append(seq)
+    targeted_sequences = proper_sequence_names
 
     seq_vals = {}
+    final_seqs = []
+    # If we get a list of targeted sequences, we only want to dump whichever of those are found in the database and schema.
     if targeted_sequences:
-        for seq in [r[0] for r in seqs if r[0] in targeted_sequences]:
-            seq_vals[seq.strip()] = await pool.fetchval(
-                f"SELECT last_value FROM {seq};"
-            )
-    else:
-        for seq in [r[0] for r in seqs]:
-            seq_vals[seq.strip()] = await pool.fetchval(
-                f"SELECT last_value FROM {seq};"
-            )
+        final_seqs = [r[0] for r in seqs if r[0] in targeted_sequences]
+    else:  # Otherwise, we want to dump all sequences found in the schema.
+        final_seqs = [r[0] for r in seqs]
+
+    for seq in final_seqs:
+        res = await pool.fetchval(f"SELECT last_value FROM {seq};")
+        seq_vals[seq.strip()] = res
 
     logger.debug(f"Dumped sequences: {seq_vals}")
     return seq_vals
@@ -46,15 +65,20 @@ async def load_sequences(pool: Pool, seqs: dict[str, int], logger: Logger) -> No
 
 
 async def compare_data(
-    src_pool: Pool, dst_pool: Pool, query: str, tables: list[str], logger: Logger
+    src_pool: Pool,
+    dst_pool: Pool,
+    query: str,
+    tables: list[str],
+    schema: str,
+    logger: Logger,
 ) -> None:
     """
     Validate data between source and destination databases by doing the following:
-    1. Get all tables with primary keys
+    1. Get all tables with primary keys (from the source)
     2. For each of those tables, select * limit 100
     3. For each row, ensure the row in the destination is identical
     """
-    pkeys, _, pkeys_raw = await analyze_table_pkeys(src_pool, logger)
+    pkeys, _, pkeys_raw = await analyze_table_pkeys(src_pool, schema, logger)
 
     pkeys_dict = {}
     # {
@@ -77,22 +101,23 @@ async def compare_data(
         # If specific table list is defined and iterated table is not in that list, skip.
         if tables and (table not in tables):
             continue
+        full_table_name = f"{schema}.{table}"
 
-        logger.debug(f"Validating table {table}...")
+        logger.debug(f"Validating table {full_table_name}...")
         order_by_pkeys = ",".join(pkeys_dict[table])
 
         src_rows = await src_pool.fetch(
-            query.format(table=table, order_by_pkeys=order_by_pkeys)
+            query.format(table=full_table_name, order_by_pkeys=order_by_pkeys)
         )
 
         # There is a chance tables are empty...
         if len(src_rows) == 0:
             dst_rows = await dst_pool.fetch(
-                query.format(table=table, order_by_pkeys=order_by_pkeys)
+                query.format(table=full_table_name, order_by_pkeys=order_by_pkeys)
             )
             if len(dst_rows) != 0:
                 raise AssertionError(
-                    f"Table {table} has 0 rows in source but nonzero rows in target... Big problem. Please investigate."
+                    f"Table {full_table_name} has 0 rows in source but nonzero rows in target... Big problem. Please investigate."
                 )
             else:
                 continue
@@ -114,7 +139,7 @@ async def compare_data(
             src_pkeys_string = src_pkeys_string[:-1]
             pkey_vals_dict[pkey] = src_pkeys_string
 
-        dst_query = f"SELECT * FROM {table} WHERE "
+        dst_query = f"SELECT * FROM {full_table_name} WHERE "
 
         for k, v in pkey_vals_dict.items():
             dst_query = dst_query + f"{k} IN ({v}) AND "
@@ -131,7 +156,7 @@ async def compare_data(
 
         if len(src_rows) != len(dst_rows):
             raise AssertionError(
-                f'Row count of the sample taken from table "{table}" '
+                f'Row count of the sample taken from table "{full_table_name}" '
                 "does not match in source and destination!\n"
                 f"Query: {dst_query}"
             )
@@ -141,7 +166,7 @@ async def compare_data(
             if src_row != dst_row:
                 raise AssertionError(
                     "Row match failure between source and destination.\n"
-                    f"Table: {table}\n"
+                    f"Table: {full_table_name}\n"
                     f"Source Row: {src_row}\n"
                     f"Dest Row: {dst_row}"
                 )
@@ -154,7 +179,7 @@ async def compare_data(
 
 
 async def compare_100_rows(
-    src_pool: Pool, dst_pool: Pool, tables: list[str], logger: Logger
+    src_pool: Pool, dst_pool: Pool, tables: list[str], schema: str, logger: Logger
 ) -> None:
     """
     Validate data between source and destination databases by doing the following:
@@ -174,11 +199,11 @@ async def compare_100_rows(
     ORDER BY {order_by_pkeys};
     """
 
-    await compare_data(src_pool, dst_pool, query, tables, logger)
+    await compare_data(src_pool, dst_pool, query, tables, schema, logger)
 
 
 async def compare_latest_100_rows(
-    src_pool: Pool, dst_pool: Pool, tables: list[str], logger: Logger
+    src_pool: Pool, dst_pool: Pool, tables: list[str], schema: str, logger: Logger
 ) -> None:
     """
     Validate data between source and destination databases by comparing the latest row:
@@ -195,30 +220,31 @@ async def compare_latest_100_rows(
     LIMIT 100;
     """
 
-    await compare_data(src_pool, dst_pool, query, tables, logger)
+    await compare_data(src_pool, dst_pool, query, tables, schema, logger)
 
 
-async def table_empty(pool: Pool, table: str, logger: Logger) -> bool:
+async def table_empty(pool: Pool, table: str, schema: str, logger: Logger) -> bool:
     """
     return true if the table is empty
     """
     logger.info(f"Checking if table {table} is empty...")
-    result = await pool.fetch(f"SELECT * FROM {table} LIMIT 1;")
+    result = await pool.fetch(f"SELECT * FROM {schema}.{table} LIMIT 1;")
     return len(result) == 0
 
 
 async def analyze_table_pkeys(
-    pool: Pool, logger: Logger
+    pool: Pool, schema: str, logger: Logger
 ) -> tuple[list[str], list[str], Record]:
     """
-    return three lists of table names. the first element is all tables
-    with pkeys in public and the second is all tables without pkeys in public.
-    The third list is the raw rows of the primary key query with the table name,
-    constraint name, position and column name for the primary key.
+    Return three lists of table names. the first element is all tables
+    with pkeys in the config's named schema and the second is all tables
+    without pkeys in that schema. The third list is the raw rows of the
+    primary key query with the table name, constraint name, position and
+    column name for the primary key.
     """
     logger.info("Checking table primary keys...")
     pkeys_raw = await pool.fetch(
-        """
+        f"""
         SELECT kcu.table_name,
             tco.constraint_name,
             kcu.ordinal_position as position,
@@ -229,7 +255,7 @@ async def analyze_table_pkeys(
             AND kcu.constraint_schema = tco.constraint_schema
             AND kcu.constraint_name = tco.constraint_name
         WHERE tco.constraint_type = 'PRIMARY KEY'
-            AND kcu.table_schema = 'public'
+            AND kcu.table_schema = '{schema}'
         ORDER BY kcu.table_name,
                 position;
         """
@@ -237,11 +263,11 @@ async def analyze_table_pkeys(
     pkeys = [r[0] for r in pkeys_raw]
 
     all_tables = await pool.fetch(
-        """SELECT table_name
+        f"""SELECT table_name
         FROM
             information_schema.tables
         WHERE
-            table_schema = 'public'
+            table_schema = '{schema}'
             AND table_name != 'pg_stat_statements'
         ORDER BY 1;"""
     )
@@ -294,7 +320,13 @@ async def enable_login_users(pool: Pool, users: list[str], logger: Logger) -> No
 
 
 async def precheck_info(
-    pool: Pool, root_name: str, owner_name: str, logger: Logger
+    pool: Pool,
+    root_name: str,
+    owner_name: str,
+    target_tables: list[str],
+    target_sequences: list[str],
+    schema: str,
+    logger: Logger,
 ) -> dict:
     """
     Return a dictionary of information about the database used to determine
@@ -302,7 +334,7 @@ async def precheck_info(
     """
     logger.info("Checking db requirements...")
     result = {
-        "server_version": await pool.fetchval("SHOW server_version;"),
+        "server_version": await pool.fetchval("SHOW server_version"),
         "max_replication_slots": await pool.fetchval("SHOW max_replication_slots;"),
         "max_worker_processes": await pool.fetchval("SHOW max_worker_processes;"),
         "max_wal_senders": await pool.fetchval("SHOW max_wal_senders;"),
@@ -311,9 +343,13 @@ async def precheck_info(
         ),
         "tables": [],
         "sequences": [],
-        "users": [],
+        "users": {},
+        "extensions": [],
     }
 
+    # server_version shows 13.14 (Debian 13.14-1.pgdg120+2) in the output. Remove the Debian part.
+    result["server_version"] = result["server_version"].split(" ")[0]
+
     try:
         result["rds.logical_replication"] = await pool.fetchval(
             "SHOW rds.logical_replication;"
@@ -333,9 +369,12 @@ async def precheck_info(
               AND n.nspname <> 'pg_catalog'
               AND n.nspname !~ '^pg_toast'
               AND n.nspname <> 'information_schema'
-          AND pg_catalog.pg_table_is_visible(c.oid)
+              AND n.nspname <> 'pglogical'
         ORDER BY 1,2;"""
     )
+    # We filter the table list if the user has specified a list of tables to target.
+    if target_tables:
+        result["tables"] = [t for t in result["tables"] if t["Name"] in target_tables]
 
     result["sequences"] = await pool.fetch(
         """
@@ -349,10 +388,16 @@ async def precheck_info(
               AND n.nspname <> 'pg_catalog'
               AND n.nspname !~ '^pg_toast'
               AND n.nspname <> 'information_schema'
-          AND pg_catalog.pg_table_is_visible(c.oid)
+              AND n.nspname <> 'pglogical'
         ORDER BY 1,2;"""
     )
 
+    # We filter the sequence list if the user has specified a list of sequences to target.
+    if target_sequences:
+        result["sequences"] = [
+            s for s in result["sequences"] if s["Name"] in target_sequences
+        ]
+
     users = await pool.fetch(
         f"""
         SELECT r.rolname, r.rolsuper, r.rolinherit,
@@ -364,15 +409,104 @@ async def precheck_info(
                 WHERE m.member = r.oid) as memberof
         , r.rolreplication
         , r.rolbypassrls
+        , has_schema_privilege(r.rolname, '{schema}', 'CREATE') AS can_create
         FROM pg_catalog.pg_roles r
         WHERE r.rolname !~ '^pg_' AND (r.rolname = '{root_name}' OR r.rolname = '{owner_name}')
         ORDER BY 1;"""
     )
 
+    # We only care about the root and owner users.
     for u in users:
         if u[0] == root_name:
-            result["root"] = u
+            result["users"]["root"] = u
         if u[0] == owner_name:
-            result["owner"] = u
+            result["users"]["owner"] = u
+
+    result["extensions"] = await pool.fetch(
+        """
+        SELECT extname
+        FROM pg_extension
+        ORDER BY extname;
+        """
+    )
 
     return result
+
+
+# TODO: Need to add schema here when working on non-public schema support.
+async def get_dataset_size(
+    tables: list[str], schema: str, pool: Pool, logger: Logger
+) -> str:
+    """
+    Get the total disk size of a dataset (via list of tables).
+
+    This function ALWAYS expects a list of tables. If not, the calling function should handle that.
+    """
+    logger.info("Getting the targeted dataset size...")
+
+    # Tables string must be of form "'table1', 'table2', ..."
+    tables_string = ", ".join([f"'{t}'" for t in tables])
+
+    query = f"""
+    SELECT
+        sum(pg_total_relation_size(schemaname || '.' || tablename)) AS total_relation_size
+    FROM
+        pg_tables
+    WHERE
+        schemaname = '{schema}'
+    AND tablename IN ({tables_string});
+    """
+
+    # Yes it's a duplicate, but it's a pretty one. Rather let Postgres do this than Python.
+    pretty_query = f"""
+    SELECT
+        pg_size_pretty(sum(pg_total_relation_size(schemaname || '.' || tablename))) AS total_relation_size
+    FROM
+        pg_tables
+    WHERE
+        schemaname = '{schema}'
+    AND tablename IN ({tables_string});
+    """
+
+    result = {
+        "db_size": await pool.fetchval(query),
+        "db_size_pretty": await pool.fetchval(pretty_query),
+    }
+
+    return result
+
+
+async def initialization_progress(
+    tables: list[str],
+    src_schema: str,
+    dst_schema: str,
+    src_pool: Pool,
+    dst_pool: Pool,
+    src_logger: Logger,
+    dst_logger: Logger,
+) -> dict[str, str]:
+    """
+    Get the size progress of the initialization stage
+    """
+
+    src_dataset_size = await get_dataset_size(tables, src_schema, src_pool, src_logger)
+    dst_dataset_size = await get_dataset_size(tables, dst_schema, dst_pool, dst_logger)
+
+    # Eliminate None values
+    if src_dataset_size["db_size"] is None:
+        src_dataset_size["db_size"] = 0
+    if dst_dataset_size["db_size"] is None:
+        dst_dataset_size["db_size"] = 0
+
+    # Eliminate division by zero
+    if src_dataset_size["db_size"] == 0 and dst_dataset_size["db_size"] == 0:
+        progress = "0 %"
+    else:
+        progress = f"{str(round(int(dst_dataset_size['db_size'])/int(src_dataset_size['db_size'])*100 ,1))} %"
+
+    status = {
+        "src_dataset_size": src_dataset_size["db_size_pretty"] or "0 bytes",
+        "dst_dataset_size": dst_dataset_size["db_size_pretty"] or "0 bytes",
+        "progress": progress,
+    }
+    return status

{pgbelt-0.6.2.dist-info → pgbelt-0.7.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pgbelt
-Version: 0.6.2
+Version: 0.7.1
 Summary: A CLI tool used to manage Postgres data migrations from beginning to end, for a single database or a fleet, leveraging pglogical replication.
 Author: Varjitt Jeeva
 Author-email: varjitt.jeeva@autodesk.com
@@ -65,6 +65,10 @@ Install pgbelt locally:
 
 See [this doc](docs/quickstart.md)!
 
+## Playbook
+
+This playbook gets updated actively. If you have any issues, solutions could be found in [this playbook](docs/playbook.md).
+
 ## Contributing
 
 We welcome contributions! See [this doc](CONTRIBUTING.md) on how to do so, including setting up your local development environment.

pgbelt-0.7.1.dist-info/RECORD

@@ -0,0 +1,27 @@
+pgbelt/__init__.py,sha256=0FxxDo2hLiQpxjY5NS48g2XC_nRMKPJf5X0aGIxmMo4,137
+pgbelt/cmd/__init__.py,sha256=28SxbIChyvJ1mLoI8XJQPrPPwKA6852V9jtnFm8R0Nw,462
+pgbelt/cmd/convenience.py,sha256=RY762BYoyOccmhHqSOaHWIPVBzY4kU0qZautTprGps0,5857
+pgbelt/cmd/helpers.py,sha256=anMAkZjJRy7n0EwewyMNJGJ4Jx9keSNZdzWqi3ICNgI,5292
+pgbelt/cmd/login.py,sha256=kyDT755YsScg3QA4ZlkyIq3J0WO4XqMepMchufjuaDE,3043
+pgbelt/cmd/preflight.py,sha256=-78Puqqf1rCxNQEyn4bQIAORez_tsy6zJAbSuO_de9s,20676
+pgbelt/cmd/schema.py,sha256=XAAj2BH6HVoA1LxuWw2kheesGDu41L3GAd5UahNp3nE,4760
+pgbelt/cmd/setup.py,sha256=CYaRGQG_ooz0BpP5g1YVmPvt8zcT4DSIdQM6nASr9Qc,5908
+pgbelt/cmd/status.py,sha256=LTvteFwA1jcfnGJdlA2siJn5TUtlf14ffMkGVSi5gi0,4730
+pgbelt/cmd/sync.py,sha256=8TSNbn1dpkyxlEMImY7FjOrzhriSXlorqE-HBjqbUfU,9027
+pgbelt/cmd/teardown.py,sha256=TTSmhmD1bqJ8-P5YyK5ZpCsmneczO0nvd2Q2s2oyXY4,3526
+pgbelt/config/__init__.py,sha256=SXok1aZcpMYJpX_hk5cuKO33CJ5s8IESkswNN9KsVSo,35
+pgbelt/config/config.py,sha256=Kw2H-G1Evfj0TXIbh3k06gE72dZEp_wXWJ2Icq_T54c,3817
+pgbelt/config/models.py,sha256=4qm23kdNNZiQOZrejD1yXP6h0Z38LNFSyI4d0i-IdUs,5830
+pgbelt/config/remote.py,sha256=RQ_dfL5g2ChVP6jeGWnpQMVKQZK1o2m9-ZPYBMaaGj4,5297
+pgbelt/main.py,sha256=YiagBiGt8pbNlukkRxROXnQX1Tx6ax7c6riuHRCrPYU,186
+pgbelt/util/__init__.py,sha256=-6KkvVMz-yGNQfeoo4CZZrgWKXYmFd4CMyoiao8OnFE,40
+pgbelt/util/asyncfuncs.py,sha256=7i_GpBmUNNZ8RUGvU-q5nclsoaCm6Lx8jLP8usYvmZc,583
+pgbelt/util/dump.py,sha256=AwyOAd9CP014gvsl-qlo1lbnXZfxoeN4ujZWUIq7KM8,14715
+pgbelt/util/logs.py,sha256=l2jT-WKZ-33eNDw4S4W1_eE4ISo4rtDRXYLVf4QTV4Y,1699
+pgbelt/util/pglogical.py,sha256=r4I95CG4eRFsEPnN30P4yxt8HeGMzKBTEm3fecfw4ws,13611
+pgbelt/util/postgres.py,sha256=909eSLhnhAkGe_4IVt-j72I3wXHArT5V_YikJnl6gdc,18178
+pgbelt-0.7.1.dist-info/LICENSE,sha256=FQ5cFkW02dKK3LmKH8z-rwn93tWSCh7lsxfNUiWcFsg,10758
+pgbelt-0.7.1.dist-info/METADATA,sha256=uXZgpAXeisAYU3M89TVK_w18R3VaEKteXZVQgQB_Kg8,2964
+pgbelt-0.7.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+pgbelt-0.7.1.dist-info/entry_points.txt,sha256=SCz_poPjkaVnWpJ-CeytAnDzbVc6l0WalOwitIqW_3g,40
+pgbelt-0.7.1.dist-info/RECORD,,

{pgbelt-0.6.2.dist-info → pgbelt-0.7.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.8.1
+Generator: poetry-core 1.9.0
 Root-Is-Purelib: true
 Tag: py3-none-any

pgbelt-0.6.2.dist-info/RECORD

@@ -1,27 +0,0 @@
-pgbelt/__init__.py,sha256=0FxxDo2hLiQpxjY5NS48g2XC_nRMKPJf5X0aGIxmMo4,137
-pgbelt/cmd/__init__.py,sha256=28SxbIChyvJ1mLoI8XJQPrPPwKA6852V9jtnFm8R0Nw,462
-pgbelt/cmd/convenience.py,sha256=xeweZUYPKarueYZQBD59meqXzAJ16uaC72pvEgua1Ec,5849
-pgbelt/cmd/helpers.py,sha256=anMAkZjJRy7n0EwewyMNJGJ4Jx9keSNZdzWqi3ICNgI,5292
-pgbelt/cmd/login.py,sha256=kyDT755YsScg3QA4ZlkyIq3J0WO4XqMepMchufjuaDE,3043
-pgbelt/cmd/preflight.py,sha256=CwP74UPpmnZ3RqJ1zFeVj3WReoXJd_2xVwWJVfTnuyg,8213
-pgbelt/cmd/schema.py,sha256=XAAj2BH6HVoA1LxuWw2kheesGDu41L3GAd5UahNp3nE,4760
-pgbelt/cmd/setup.py,sha256=v-phLiNQYFportKHvcG0EBfarfv_F2xekeL9P_wuI5c,5252
-pgbelt/cmd/status.py,sha256=r-Zv5QiO7oQQBTewnTXOdl-IslNGE6lAxlgArAYRTWU,3190
-pgbelt/cmd/sync.py,sha256=-ffZOCTWcFXjcpRe7MFfBRdV8tjSh6uoVALvykX85pU,8375
-pgbelt/cmd/teardown.py,sha256=Yps1iejx9amW2b3kmtKjw5-ySrk_qK3LfaH6EOM7NHE,3490
-pgbelt/config/__init__.py,sha256=SXok1aZcpMYJpX_hk5cuKO33CJ5s8IESkswNN9KsVSo,35
-pgbelt/config/config.py,sha256=Kw2H-G1Evfj0TXIbh3k06gE72dZEp_wXWJ2Icq_T54c,3817
-pgbelt/config/models.py,sha256=XrTZLps-jQvOaOdf9BWV5cx6xqHFRD7TnxtQMpohR1c,5432
-pgbelt/config/remote.py,sha256=RQ_dfL5g2ChVP6jeGWnpQMVKQZK1o2m9-ZPYBMaaGj4,5297
-pgbelt/main.py,sha256=YiagBiGt8pbNlukkRxROXnQX1Tx6ax7c6riuHRCrPYU,186
-pgbelt/util/__init__.py,sha256=-6KkvVMz-yGNQfeoo4CZZrgWKXYmFd4CMyoiao8OnFE,40
-pgbelt/util/asyncfuncs.py,sha256=7i_GpBmUNNZ8RUGvU-q5nclsoaCm6Lx8jLP8usYvmZc,583
-pgbelt/util/dump.py,sha256=p6M26s3FXnSCTmee4v7X0-azYvmcVz6DkZf8J8YxU0s,14357
-pgbelt/util/logs.py,sha256=l2jT-WKZ-33eNDw4S4W1_eE4ISo4rtDRXYLVf4QTV4Y,1699
-pgbelt/util/pglogical.py,sha256=URi4L_4PWYIz0iWw7HcXNfjzvWXPjQCm71cAOqnt3Qk,13059
-pgbelt/util/postgres.py,sha256=9AgnmHdYpWVypyBNTEQEpd4COoa1JEQXoELOR8-HrzQ,13668
-pgbelt-0.6.2.dist-info/LICENSE,sha256=FQ5cFkW02dKK3LmKH8z-rwn93tWSCh7lsxfNUiWcFsg,10758
-pgbelt-0.6.2.dist-info/METADATA,sha256=7bEElDGURHPEE3cVWbF19aPGhAXNQ0_YxPRpdZK_Fc8,2826
-pgbelt-0.6.2.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-pgbelt-0.6.2.dist-info/entry_points.txt,sha256=SCz_poPjkaVnWpJ-CeytAnDzbVc6l0WalOwitIqW_3g,40
-pgbelt-0.6.2.dist-info/RECORD,,