add check for maximum databases size to migrate

Teemu Harju · Teemu Harju · commit fc94257a36fa · 2022-01-12T10:51:44.000+02:00
diff --git a/README.md b/README.md
@@ -122,6 +122,8 @@ Use `--no-replicate-extension-tables` to skip extension tables.  By default it a
 
 With `--force-method` you can specify if you wish to use either replication or dump method. Otherwise the most suitable method is chosen automatically.
 
+Using `--dbs-max-total-size` together with `--validate` you can check if the size of the source database in below some threshold.
+
 ### API example
 
 Migrating from AWS RDS to Aiven for PostgreSQL. Logical replication is enabled in source AWS RDS PostgreSQL
diff --git a/aiven_db_migrate/migrate/pgmigrate.py b/aiven_db_migrate/migrate/pgmigrate.py
@@ -389,6 +389,18 @@ def mangle_db_name(self, db_name: str) -> str:
 
 class PGSource(PGCluster):
     """Source PostgreSQL cluster"""
+    def get_size(self, *, dbname, only_tables: Optional[List[str]] = None) -> float:
+        if only_tables == []:
+            return 0
+        if only_tables is not None:
+            query = "SELECT SUM(pg_total_relation_size(tablename)) AS size FROM UNNEST(%s) AS tablename"
+            args = [only_tables]
+        else:
+            query = "SELECT pg_database_size(oid) AS size FROM pg_catalog.pg_database WHERE datname = %s"
+            args = [dbname]
+        result = self.c(query, args=args, dbname=dbname)
+        return result[0]["size"] or 0
+
     def create_publication(self, *, dbname: str, only_tables: Optional[List[str]] = None) -> str:
         mangled_name = self.mangle_db_name(dbname)
         pubname = f"aiven_db_migrate_{mangled_name}_pub"
@@ -852,6 +864,17 @@ def _check_databases(self):
                 else:
                     self.log.info("Database %r already exists in target", dbname)
 
+    def _check_database_size(self, max_size: float):
+        dbs_size = 0
+        for dbname, source_db in self.source.databases.items():
+            only_tables = self.filter_tables(db=source_db)
+            db_size = self.source.get_size(dbname=dbname, only_tables=only_tables)
+            dbs_size += db_size
+        if dbs_size > max_size:
+            raise PGMigrateValidationFailedError(
+                f"Databases do not fit to the required maximum size ({dbs_size} > {max_size})"
+            )
+
     def _check_pg_lang(self):
         source_lang = {lan["lanname"] for lan in self.source.pg_lang}
         target_lang = {lan["lanname"] for lan in self.target.pg_lang}
@@ -1131,7 +1154,7 @@ def _db_migrate(self, *, pgtask: PGMigrateTask) -> PGMigrateStatus:
         pgtask.method = PGMigrateMethod.dump
         return self._dump_data(db=pgtask.source_db)
 
-    def validate(self):
+    def validate(self, dbs_max_total_size: Optional[float] = None):
         """
         Do best effort validation whether all the bits and pieces are in place for migration to succeed.
         * Migrating to same server is not supported (doable but requires obviously different dbname)
@@ -1156,6 +1179,8 @@ def validate(self):
             # but it can be newer than the source version: source <= pgdump <= target
             self.pgbin = find_pgbin_dir(str(self.source.version), max_pgversion=str(self.target.version))
             self._check_databases()
+            if dbs_max_total_size is not None:
+                self._check_database_size(max_size=dbs_max_total_size)
             self._check_pg_lang()
             self._check_pg_ext()
         except KeyError as err:
@@ -1304,6 +1329,12 @@ def main(args=None, *, prog="pg_migrate"):
         default=None,
         help="Force the migration method to be used as either replication or dump.",
     )
+    parser.add_argument(
+        "--dbs-max-total-size",
+        type=int,
+        default=-1,
+        help="Max total size of databases to be migrated, ignored by default",
+    )
 
     args = parser.parse_args(args)
     log_format = "%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s"
@@ -1334,7 +1365,8 @@ def main(args=None, *, prog="pg_migrate"):
             raise ValueError(f"Unsupported migration method '{args.force_method}'") from e
 
     if args.validate:
-        pg_mig.validate()
+        dbs_max_total_size = None if args.dbs_max_total_size == -1 else args.dbs_max_total_size
+        pg_mig.validate(dbs_max_total_size=dbs_max_total_size)
     else:
         result: PGMigrateResult = pg_mig.migrate(force_method=method)
         print()
diff --git a/test/test_db_size_check.py b/test/test_db_size_check.py
@@ -0,0 +1,47 @@
+from aiven_db_migrate.migrate.pgmigrate import PGMigrate
+from test.conftest import PGRunner
+from test.utils import random_string
+from typing import Tuple
+
+import psycopg2
+import pytest
+
+
+def test_db_size(pg_source_and_target_replication: Tuple[PGRunner, PGRunner]):
+    source, target = pg_source_and_target_replication
+
+    db_name = random_string(6)
+    other_db_name = random_string(6)
+
+    source.create_db(dbname=db_name)
+    source.create_db(dbname=other_db_name)
+
+    pg_mig = PGMigrate(
+        source_conn_info=source.super_conn_info(),
+        target_conn_info=target.super_conn_info(),
+        verbose=True,
+    )
+
+    # Create few tables and insert some data
+    tables = [f'table_{i}' for i in range(4)]
+    for dbname in {db_name, other_db_name}:
+        with source.cursor(dbname=dbname) as c:
+            for t in tables:
+                c.execute(f"DROP TABLE IF EXISTS {t}")
+                c.execute(f"CREATE TABLE {t} (foo INT)")
+                c.execute(f"INSERT INTO {t} (foo) VALUES (1), (2), (3)")
+
+    size = pg_mig.source.get_size(dbname=db_name, only_tables=[])
+    assert size == 0
+
+    size = pg_mig.source.get_size(dbname=db_name)
+    assert size >= 0  # returns slightly different values per pg version
+
+    size = pg_mig.source.get_size(dbname=db_name, only_tables=tables)
+    assert size == 32768
+
+    size = pg_mig.source.get_size(dbname=db_name, only_tables=tables[:1])
+    assert size == 8192
+
+    with pytest.raises(psycopg2.OperationalError):
+        size = pg_mig.source.get_size(dbname="notfound")
diff --git a/test/test_migrate_checks.py b/test/test_migrate_checks.py
@@ -0,0 +1,113 @@
+# Copyright (c) 2021 Aiven, Helsinki, Finland. https://aiven.io/
+
+from aiven_db_migrate.migrate.errors import PGMigrateValidationFailedError
+from aiven_db_migrate.migrate.pgmigrate import PGMigrate
+from test.conftest import PGRunner
+from test.utils import random_string
+from typing import Tuple
+from unittest.mock import patch
+
+import pytest
+
+
+def test_dbs_max_total_size_check(pg_source_and_target: Tuple[PGRunner, PGRunner]):
+    source, target = pg_source_and_target
+    dbnames = {random_string() for _ in range(3)}
+
+    for dbname in dbnames:
+        source.create_db(dbname=dbname)
+        target.create_db(dbname=dbname)
+
+    # This DB seems to be created outside the tests
+    dbnames.add("postgres")
+
+    # Create few tables and insert some data
+    tables = [f'table_{i}' for i in range(4)]
+    for dbname in dbnames:
+        with source.cursor(dbname=dbname) as c:
+            for t in tables:
+                c.execute(f"DROP TABLE IF EXISTS {t}")
+                c.execute(f"CREATE TABLE {t} (foo INT)")
+                c.execute(f"INSERT INTO {t} (foo) VALUES (1), (2), (3)")
+
+    pg_mig = PGMigrate(
+        source_conn_info=source.conn_info(),
+        target_conn_info=target.conn_info(),
+        createdb=False,
+        verbose=True,
+    )
+
+    with patch(
+        "aiven_db_migrate.migrate.pgmigrate.PGMigrate._check_database_size", side_effect=pg_mig._check_database_size
+    ) as mock_db_size_check:
+        # DB size check is not run
+        pg_mig.validate()
+        mock_db_size_check.assert_not_called()
+
+        mock_db_size_check.reset_mock()
+
+        # DB size check with max size of zero
+        with pytest.raises(PGMigrateValidationFailedError) as e:
+            pg_mig.validate(dbs_max_total_size=0)
+        assert "Databases do not fit to the required maximum size" in str(e)
+        mock_db_size_check.assert_called_once_with(max_size=0)
+
+        mock_db_size_check.reset_mock()
+
+        # DB size check with enough size
+        pg_mig.validate(dbs_max_total_size=1073741824)
+        mock_db_size_check.assert_called_once_with(max_size=1073741824)
+
+    # Test with DB name filtering
+    pg_mig = PGMigrate(
+        source_conn_info=source.conn_info(),
+        target_conn_info=target.conn_info(),
+        createdb=False,
+        verbose=True,
+        filtered_db=",".join(dbnames),
+    )
+
+    with patch(
+        "aiven_db_migrate.migrate.pgmigrate.PGMigrate._check_database_size", side_effect=pg_mig._check_database_size
+    ) as mock_db_size_check:
+        # Should pass as all DBs are filtered out from size calculations
+        pg_mig.validate(dbs_max_total_size=0)
+        mock_db_size_check.assert_called_once_with(max_size=0)
+
+    # Test with table filtering
+
+    # Include all tables in "skip_tables"
+    pg_mig = PGMigrate(
+        source_conn_info=source.conn_info(),
+        target_conn_info=target.conn_info(),
+        createdb=False,
+        verbose=True,
+        skip_tables=tables,  # skip all tables
+    )
+
+    with patch(
+        "aiven_db_migrate.migrate.pgmigrate.PGMigrate._check_database_size", side_effect=pg_mig._check_database_size
+    ) as mock_db_size_check:
+        # Should pass as all tables are filtered out from size calculations
+        pg_mig.validate(dbs_max_total_size=0)
+        mock_db_size_check.assert_called_once_with(max_size=0)
+
+    # Only the first table is included
+    pg_mig = PGMigrate(
+        source_conn_info=source.conn_info(),
+        target_conn_info=target.conn_info(),
+        createdb=False,
+        verbose=True,
+        with_tables=tables[:1],  # include only one table
+    )
+    with patch(
+        "aiven_db_migrate.migrate.pgmigrate.PGMigrate._check_database_size", side_effect=pg_mig._check_database_size
+    ) as mock_db_size_check:
+        # This fails as one table is included in check and it should have data
+        with pytest.raises(PGMigrateValidationFailedError) as e:
+            pg_mig.validate(dbs_max_total_size=0)
+        assert "Databases do not fit to the required maximum size" in str(e)
+        mock_db_size_check.assert_called_once_with(max_size=0)
+
+        # Should easily fit
+        pg_mig.validate(dbs_max_total_size=1073741824)
