feat(cli) - add replication role and replica identity checks (#489)

* Add replication role and replica identity checks

* Fix to show all DB results at once and formatted some msgs.

* Show the per DB results in sorted order.

Author: trivikram29

src/dma/collector/query_managers.py

@@ -284,6 +284,7 @@ def get_collection_queries(self) -> set[str]:
             "collection_postgres_schema_objects",
             "collection_postgres_settings",
             "collection_postgres_source_details",
+            "collection_postgres_replication_role",
         }
 
     def get_per_db_collection_queries(self) -> set[str]:
@@ -301,6 +302,7 @@ def get_per_db_collection_queries(self) -> set[str]:
             "collection_postgres_user_views_without_privilege",
             "collection_postgres_user_sequences_without_privilege",
             "collection_postgres_tables_with_no_primary_key",
+            "collection_postgres_tables_with_primary_key_replica_identity",
         }
 
     def get_collection_filenames(self) -> dict[str, str]:
@@ -329,6 +331,8 @@ def get_collection_filenames(self) -> dict[str, str]:
             "collection_postgres_source_details": "postgres_source_details",
             "collection_postgres_pglogical_provider_node": "postgres_pglogical_details",
             "collection_postgres_tables_with_no_primary_key": "postgres_table_details",
+            "collection_postgres_tables_with_primary_key_replica_identity": "postgres_table_details",
+            "collection_postgres_replication_role": "collection_privileges",
         }
 
 

src/dma/collector/sql/canonical/ddls-postgres-collection.sql

@@ -885,6 +885,24 @@ create or replace table collection_postgres_tables_with_no_primary_key(
     database_name VARCHAR
   );
 
+create or replace table collection_postgres_tables_with_primary_key_replica_identity(
+    pkey VARCHAR,
+    dma_source_id VARCHAR,
+    dma_manual_id VARCHAR,
+    nspname VARCHAR,
+    relname VARCHAR,
+    database_name VARCHAR
+  );
+
+create or replace table collection_postgres_replication_role(
+    pkey VARCHAR,
+    dma_source_id VARCHAR,
+    dma_manual_id VARCHAR,
+    rolname VARCHAR,
+    rolreplication VARCHAR,
+    database_name VARCHAR
+  );
+
 create or replace table collection_postgres_pglogical_provider_node(
     pkey VARCHAR,
     dma_source_id VARCHAR,

src/dma/collector/sql/sources/postgres/collection-privileges.sql

@@ -137,3 +137,16 @@ select :PKEY as pkey,
   src.relname as rel_name,
   current_database() as database_name
 from src;
+
+-- name: collection-postgres-replication-role
+with src as (
+  SELECT rolname, rolreplication FROM pg_catalog.pg_roles
+  WHERE rolname IN (SELECT CURRENT_USER)
+)
+select :PKEY as pkey,
+  :DMA_SOURCE_ID as dma_source_id,
+  :DMA_MANUAL_ID as dma_manual_id,
+  src.rolname,
+  src.rolreplication,
+  current_database() as database_name
+from src;

src/dma/collector/sql/sources/postgres/collection-table_details.sql

@@ -467,3 +467,30 @@ select :PKEY as pkey,
   src.relname,
   current_database() as database_name
 from src;
+
+-- name: collection-postgres-tables-with-primary-key-replica-identity
+with src as (
+SELECT nr.nspname, r.relname
+	FROM pg_namespace nr,
+		pg_class r,
+		pg_namespace nc,
+	  pg_constraint c
+	WHERE
+			nr.oid = r.relnamespace
+			AND r.oid = c.conrelid
+			AND nc.oid = c.connamespace
+			AND c.contype = 'p'::"char"
+			AND r.relkind = 'r'::"char"
+			AND r.relpersistence = 'p'::"char"
+			AND r.relreplident IN ('f'::"char", 'n'::"char")
+			AND NOT pg_catalog.pg_is_other_temp_schema(nr.oid)
+			AND nr.nspname not in ('pg_catalog', 'information_schema', 'pglogical', 'pglogical_origin')
+			and nr.nspname not like 'pg\_%%'
+)
+select :PKEY as pkey,
+  :DMA_SOURCE_ID as dma_source_id,
+  :DMA_MANUAL_ID as dma_manual_id,
+  src.nspname,
+  src.relname,
+  current_database() as database_name
+from src;

src/dma/collector/workflows/readiness_check/_postgres/main.py

@@ -49,6 +49,8 @@
 PRIVILEGES: Final = "PRIVILEGES"
 PGLOGICAL_NODE_ALREADY_EXISTS: Final = "PGLOGICAL_NODE_ALREADY_EXISTS"
 TABLES_WITH_NO_PK: Final = "TABLES_WITH_NO_PK"
+UNSUPPORTED_TABLES_WITH_REPLICA_IDENTITY: Final = "UNSUPPORTED_TABLES_WITH_REPLICA_IDENTITY"
+REPLICATION_ROLE: Final = "REPLICATION_ROLE"
 CLOUDSQL: Final = "CLOUDSQL"
 ALLOYDB: Final = "ALLOYDB"
 CloudSQL_SUPER_ROLE: Final = "cloudsqladmin"
@@ -126,21 +128,23 @@ def execute(self) -> None:
         self._check_max_replication_slots()
         self._check_max_wal_senders_replication_slots()
         self._check_max_worker_processes()
-        self._check_extensions()
         self._check_fdw()
 
         # Per DB Checks.
+        self._check_extensions()
+        self._check_replication_role()
         # db_check_results stores the verification results for all DBs.
         for config in self.rule_config:
             db_check_results: dict[str, dict[str, list]] = {}
-            for db in self.get_all_dbs():
+            for db in sorted(self.get_all_dbs()):
                 is_pglogical_installed = self._check_pglogical_installed(db, db_check_results)
                 if is_pglogical_installed:
                     privilege_check_passed = self._check_privileges(db, db_check_results)
                     if not privilege_check_passed:
-                        break
+                        continue
                     self._check_if_node_exists(db, db_check_results)
                 self._check_tables_without_pk(db, db_check_results)
+                self._check_tables_replica_identity(db, db_check_results)
             self._save_results(config.db_variant, db_check_results)
 
     def _save_results(self, db_variant: PostgresVariants, db_check_results: dict[str, dict[str, list]]) -> None:
@@ -157,9 +161,9 @@ def _save_results(self, db_variant: PostgresVariants, db_check_results: dict[str
 
                 if rule == PRIVILEGES:
                     if severity == ACTION_REQUIRED:
-                        output_str = ";".join(result[ACTION_REQUIRED])
+                        output_str = ";\n\n".join(result[ACTION_REQUIRED])
                     elif severity == PASS:
-                        output_str = ";".join(result[PASS])
+                        output_str = ";\n".join(result[PASS])
                     else:
                         continue
 
@@ -179,6 +183,12 @@ def _save_results(self, db_variant: PostgresVariants, db_check_results: dict[str
                     else:
                         continue
 
+                if rule == UNSUPPORTED_TABLES_WITH_REPLICA_IDENTITY:
+                    if severity == ACTION_REQUIRED:
+                        output_str = f"Source has table(s) with both primary key and replica identity FULL or NOTHING. Please remove replica identity or change it to DEFAULT to migrate: {';'.join(result[ACTION_REQUIRED])}"
+                    else:
+                        continue
+
                 if len(result[severity]) > 0:
                     self.save_rule_result(
                         db_variant,
@@ -225,6 +235,17 @@ def _check_tables_without_pk(self, db_name: str, db_check_results: dict[str, dic
         if tables:
             db_check_results[rule_code][WARNING].append(f"In database {db_name}, {tables} don't have primary keys")
 
+    def _check_tables_replica_identity(self, db_name: str, db_check_results: dict[str, dict[str, list]]) -> None:
+        rule_code = UNSUPPORTED_TABLES_WITH_REPLICA_IDENTITY
+        result = self.local_db.sql(
+            "select CONCAT(nspname, '.', relname) from collection_postgres_tables_with_primary_key_replica_identity where database_name = $db_name",
+            params={"db_name": db_name},
+        ).fetchall()
+        tables = ", ".join(row[0] for row in result)
+        init_results_dict(db_check_results, rule_code)
+        if tables:
+            db_check_results[rule_code][ACTION_REQUIRED].append(f"{tables} in database {db_name}")
+
     def _check_version(self) -> None:
         rule_code = "DATABASE_VERSION"
         self.console.print(f"version: {self.db_version}")
@@ -357,6 +378,29 @@ def check_user_obj_privileges(self, db_name: str) -> list[str]:
         errors.extend(f"user doesn't have SELECT privilege on sequence {row[0]}.{row[1]}" for row in rows)
         return errors
 
+    def _check_replication_role(self) -> None:
+        if self._is_rds():
+            return
+        rule_code = REPLICATION_ROLE
+        result = self.local_db.sql("SELECT rolreplication FROM collection_postgres_replication_role").fetchone()
+        if result is None:
+            return
+        for c in self.rule_config:
+            if result[0] == "false":
+                self.save_rule_result(
+                    c.db_variant,
+                    rule_code,
+                    ACTION_REQUIRED,
+                    "user does not have rolreplication role.",
+                )
+            else:
+                self.save_rule_result(
+                    c.db_variant,
+                    rule_code,
+                    PASS,
+                    "user has rolreplication role.",
+                )
+
     def _check_privileges(self, db_name: str, db_check_results: dict[str, dict[str, list]]) -> bool:
         rule_code = PRIVILEGES
         errors = self._check_pglogical_privileges(db_name)