Merge pull request #28 from aiven/tsharju-db-size-check

add check for maximum databases size to migrate

#28

Author: kmichel-aiven

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

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"
@@ -766,7 +778,7 @@ def _convert_table_names(self, tables: Optional[List[str]]) -> Set[PGTable]:
                 )
         return ret
 
-    def filter_tables(self, db: PGDatabase) -> List[str]:
+    def filter_tables(self, db: PGDatabase) -> Optional[List[str]]:
         """
             Given a database, it will attempt to return a list of tables that should be data dumped / replicated
             based on the skip table list, with table list and the replicate extensions flag
@@ -779,7 +791,7 @@ def filter_tables(self, db: PGDatabase) -> List[str]:
             "Filtering tables for db %r, and skip tables %r and with tables %r", db, self.skip_tables, self.with_tables
         )
         if not self.skip_tables and not self.with_tables and self.replicate_extensions:
-            return []
+            return None
         if not db.tables:
             return []
         ret: Set[PGTable] = set()
@@ -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}
@@ -1053,7 +1076,8 @@ def _dump_data(self, *, db: PGDatabase) -> PGMigrateStatus:
             "--data-only",
             self.source.conn_str(dbname=dbname),
         ]
-        pg_dump_cmd.extend([f"--table={w}" for w in self.filter_tables(db)])
+        tables = self.filter_tables(db) or []
+        pg_dump_cmd.extend([f"--table={w}" for w in tables])
         subtask: PGSubTask = self._pg_dump_pipe_psql(
             pg_dump_cmd=pg_dump_cmd, target_conn_str=self.target.conn_str(dbname=dbname)
         )
@@ -1076,7 +1100,8 @@ def _db_replication(self, *, db: PGDatabase) -> PGMigrateStatus:
         dbname = db.dbname
         pubname = slotname = subname = None
         try:
-            pubname = self.source.create_publication(dbname=dbname, only_tables=self.filter_tables(db))
+            tables = self.filter_tables(db) or []
+            pubname = self.source.create_publication(dbname=dbname, only_tables=tables)
             slotname = self.source.create_replication_slot(dbname=dbname)
             subname = self.target.create_subscription(
                 conn_str=self.source.conn_str(dbname=dbname), pubname=pubname, slotname=slotname, dbname=dbname
@@ -1129,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)
@@ -1154,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:
@@ -1302,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"
@@ -1332,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()

test/conftest.py

@@ -4,6 +4,7 @@
 
 from aiven_db_migrate.migrate.pgutils import find_pgbin_dir
 from contextlib import contextmanager
+from copy import copy
 from datetime import datetime
 from distutils.version import LooseVersion
 from pathlib import Path
@@ -423,7 +424,9 @@ def generate_fixtures():
         assert version in SUPPORTED_PG_VERSIONS, f"Supported pg versions are: {SUPPORTED_PG_VERSIONS}"
         pg_target_versions.append(version)
     else:
-        pg_target_versions = SUPPORTED_PG_VERSIONS
+        # We do not support PG 9.5 as target
+        pg_target_versions = copy(SUPPORTED_PG_VERSIONS)
+        pg_target_versions.remove("9.5")
 
     for source in pg_source_versions:
         name_prefix = "pg{}".format(source.replace(".", ""))

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")

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)