db-condenser 1.0.0__py3-none-any.whl

db_condenser/config_reader.py

@@ -0,0 +1,191 @@
+import json
+import sys
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Literal
+
+
+@dataclass
+class InitialTarget:
+    table: str
+    percent: float | None = None
+    where: str | None = None
+
+    def __post_init__(self):
+        # Exactly one of where/percent must be set
+        if (self.where is None) == (self.percent is None):
+            raise ValueError(
+                "Initial Target must specify exactly one of 'where' or 'percent'"
+            )
+
+
+class DbType(str, Enum):
+    POSTGRES = "postgres"
+    MYSQL = "mysql"
+
+
+@dataclass
+class DbConnectInfo:
+    user_name: str
+    host: str
+    db_name: str
+    port: int
+    ssl_mode: str | None = None
+    # No password will prompt user
+    password: str | None = None
+
+
+@dataclass
+class UpstreamFilter:
+    condition: str
+    table: str | None = None
+    column: str | None = None
+
+    def __post_init__(self):
+        # Exactly one of table/column must be set
+        if (self.table is None) == (self.column is None):
+            raise ValueError(
+                "Upstream filters must specify exactly one of 'table' or 'column'"
+            )
+
+
+@dataclass
+class DependencyBreak:
+    fk_table: str
+    target_table: str
+    preserve_fk_opportunistically: bool = False
+
+
+@dataclass
+class FkAugmentation:
+    fk_table: str
+    fk_columns: list[str]
+    target_table: str
+    target_columns: list[str]
+
+    def __post_init__(self):
+        if len(self.fk_columns) != len(self.target_columns):
+            raise ValueError("fk_columns and target_columns must be the same length")
+
+
+@dataclass
+class Config:
+    db_type: DbType
+    initial_targets: list[InitialTarget]
+    source_db_connection_info: DbConnectInfo
+    destination_db_connection_info: DbConnectInfo
+    keep_disconnected_tables: bool = False
+    upstream_filters: list[UpstreamFilter] = field(default_factory=list)
+    excluded_tables: list[str] = field(default_factory=list)
+    passthrough_tables: list[str] = field(default_factory=list)
+    dependency_breaks: list[DependencyBreak] = field(default_factory=list)
+    fk_augmentation: list[FkAugmentation] = field(default_factory=list)
+    max_rows_per_table: int | Literal["ALL"] | None = None
+    use_temp_tables: bool = False
+    use_copy_protocol: bool = False
+    pre_constraint_sql: list[str] = field(default_factory=list)
+    post_subset_sql: list[str] = field(default_factory=list)
+
+    @property
+    def dependency_break_set(self) -> set[tuple[str, str]]:
+        return {(b.fk_table, b.target_table) for b in self.dependency_breaks}
+
+    @property
+    def preserve_fk_opportunistically(self) -> set[tuple[str, str]]:
+        return {
+            (b.fk_table, b.target_table)
+            for b in self.dependency_breaks
+            if b.preserve_fk_opportunistically
+        }
+
+    @property
+    def initial_target_tables(self) -> list[str]:
+        return [target.table for target in self.initial_targets]
+
+
+config: Config | None = None
+
+
+def _raw_dict_to_config(raw_config: dict) -> Config:
+    initial_targets = []
+    db_type = DbType(raw_config["db_type"].lower())
+
+    initial_targets = [
+        InitialTarget(**target) for target in raw_config["initial_targets"]
+    ]
+
+    source_db = DbConnectInfo(**raw_config["source_db_connection_info"])
+    dest_db = DbConnectInfo(**raw_config["destination_db_connection_info"])
+
+    upstream_filters = [
+        UpstreamFilter(**table) for table in raw_config.get("upstream_filters", [])
+    ]
+
+    excluded_tables = [table for table in raw_config.get("excluded_tables", [])]
+    passthrough_tables = list(
+        set([table for table in raw_config.get("passthrough_tables", [])])
+    )
+    dependency_breaks = [
+        DependencyBreak(**relation)
+        for relation in raw_config.get("dependency_breaks", [])
+    ]
+    fk_augmentation = []
+    for fka in raw_config.get("fk_augmentation", []):
+        if "fk_schema" in fka:
+            fka = {
+                "fk_table": fka["fk_schema"] + "." + fka["fk_table"],
+                "fk_columns": fka["fk_columns"],
+                "target_table": fka["target_schema"] + "." + fka["target_table"],
+                "target_columns": fka["target_columns"],
+            }
+        fk_augmentation.append(FkAugmentation(**fka))
+
+    pre_constraint_sql = [sql for sql in raw_config.get("pre_constraint_sql", [])]
+    post_subset_sql = [sql for sql in raw_config.get("post_subset_sql", [])]
+    max_rows_per_table = raw_config.get("max_rows_per_table", None)
+    use_temp_tables = bool(raw_config.get("use_temp_tables", False))
+    use_copy_protocol = bool(raw_config.get("use_copy_protocol", False))
+    return Config(
+        db_type=db_type,
+        initial_targets=initial_targets,
+        source_db_connection_info=source_db,
+        destination_db_connection_info=dest_db,
+        keep_disconnected_tables=bool(
+            raw_config.get("keep_disconnected_tables", False)
+        ),
+        upstream_filters=upstream_filters,
+        excluded_tables=excluded_tables,
+        passthrough_tables=passthrough_tables,
+        dependency_breaks=dependency_breaks,
+        fk_augmentation=fk_augmentation,
+        max_rows_per_table=max_rows_per_table,
+        use_temp_tables=use_temp_tables,
+        use_copy_protocol=use_copy_protocol,
+        pre_constraint_sql=pre_constraint_sql,
+        post_subset_sql=post_subset_sql,
+    )
+
+
+def initialize(file_like=None):
+    global config
+    if config:
+        print("WARNING: Attempted to initialize configuration twice.", file=sys.stderr)
+
+    if not file_like:
+        with open("config.json", "r") as fp:
+            raw_config = json.load(fp)
+    else:
+        raw_config = json.load(file_like)
+
+    config = _raw_dict_to_config(raw_config)
+
+
+def get_config() -> Config:
+    if config is None:
+        raise RuntimeError("Config not initialized — call initialize() first")
+    return config
+
+
+def reset_config():
+    global config
+    config = None

db_condenser/data_masking.py

@@ -0,0 +1,52 @@
+from typing import Any, Final
+
+from faker import Faker
+
+fake = Faker()
+
+
+class DataMasking:
+    @staticmethod
+    def null_out(_: str):
+        return None
+
+    @staticmethod
+    def mask_numbers(value: Any) -> str | None:
+        """
+        Mask certain strings that may contain a mixture of letters,
+        normal characters, whitespaces, or special characters
+        """
+        if value is None:
+            return None
+        str_value = str(value)
+        return "".join(
+            str(fake.random_digit()) if c.isdigit() else c for c in str_value
+        )
+
+    @staticmethod
+    def mask_characters(value: Any) -> str | None:
+        """
+        Mask certain strings that may contain a mixture of letters,
+        normal characters, whitespaces, or special characters
+        """
+        if value is None:
+            return None
+        str_value = str(value)
+        return "".join(fake.random_letter() if c.isalpha() else c for c in str_value)
+
+    @staticmethod
+    def mask_email(email: Any) -> str | None:
+        if email is None:
+            return None
+        s = str(email).split("@")
+        if len(s) < 2:
+            return fake.email()
+        return f"{fake.user_name()}@{s[1]}"
+
+
+DATA_MASKING_MAPPER: Final = {
+    "null_out": DataMasking.null_out,
+    "mask_numbers": DataMasking.mask_numbers,
+    "mask_characters": DataMasking.mask_characters,
+    "mask_email": DataMasking.mask_email,
+}

db_condenser/database_helper.py

@@ -0,0 +1,13 @@
+from db_condenser.config_reader import DbType, get_config
+
+
+def get_specific_helper():
+    config = get_config()
+    if config.db_type == DbType.POSTGRES:
+        from db_condenser import psql_database_helper
+
+        return psql_database_helper
+    else:
+        from db_condenser import mysql_database_helper
+
+        return mysql_database_helper

db_condenser/db_connect.py

@@ -0,0 +1,126 @@
+import getpass
+import sys
+import time
+from datetime import datetime
+
+import mysql.connector
+import psycopg
+
+from db_condenser.config_reader import DbConnectInfo, DbType
+
+
+class DbConnection:
+    def __init__(self, connection):
+        self.connection = connection
+
+    def commit(self):
+        self.connection.commit()
+
+    def close(self):
+        self.connection.close()
+
+
+class LoggingCursor:
+    def __init__(self, cursor, verbose=False):
+        self.inner_cursor = cursor
+        self._verbose = verbose
+
+    def execute(self, query, params=None):
+        start_time = time.time()
+        if self._verbose:
+            print("Beginning query @ {}:\n\t{}".format(str(datetime.now()), query))
+            sys.stdout.flush()
+        retval = self.inner_cursor.execute(query, params)
+        if self._verbose:
+            print("\tQuery completed in {}s".format(time.time() - start_time))
+            sys.stdout.flush()
+        return retval
+
+    def __getattr__(self, name):
+        return self.inner_cursor.__getattribute__(name)
+
+    def __exit__(self, a, b, c):
+        return self.inner_cursor.__exit__(a, b, c)
+
+    def __enter__(self):
+        return LoggingCursor(self.inner_cursor.__enter__(), self._verbose)
+
+
+# small wrapper to the connection class that gives us a common interface to the cursor()
+# method across MySQL and Postgres. This one is for Postgres
+class PsqlConnection(DbConnection):
+    def __init__(self, connect, read_repeatable, verbose=False):
+        connection_args = dict(
+            dbname=connect.db_name,
+            user=connect.user,
+            password=connect.password,
+            host=connect.host,
+            port=connect.port,
+        )
+
+        if connect.ssl_mode:
+            connection_args["sslmode"] = connect.ssl_mode
+
+        DbConnection.__init__(self, psycopg.connect(**connection_args))
+        self._verbose = verbose
+        if read_repeatable:
+            self.connection.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
+
+    def cursor(self, name=None, withhold=False):
+        return LoggingCursor(
+            self.connection.cursor(name=name, withhold=withhold), self._verbose
+        )
+
+
+# small wrapper to the connection class that gives us a common interface to the cursor()
+# method across MySQL and Postgres. This one is for MySQL
+class MySqlConnection(DbConnection):
+    def __init__(self, connect, read_repeatable, verbose=False):
+        DbConnection.__init__(
+            self,
+            mysql.connector.connect(
+                host=connect.host,
+                port=connect.port,
+                user=connect.user,
+                password=connect.password,
+                database=connect.db_name,
+            ),
+        )
+
+        self.db_name = connect.db_name
+        self._verbose = verbose
+
+        if read_repeatable:
+            self.connection.start_transaction(isolation_level="REPEATABLE READ")
+
+    def cursor(self, name=None, withhold=False):
+        return LoggingCursor(self.connection.cursor(), self._verbose)
+
+
+class DbConnect:
+    def __init__(self, db_type: DbType, connection_info: DbConnectInfo, verbose=False):
+        if connection_info.password is None:
+            connection_info.password = getpass.getpass(
+                "Enter password for {0} on host {1}: ".format(
+                    connection_info.user_name, connection_info.host
+                )
+            )
+
+        self.user = connection_info.user_name
+        self.password = connection_info.password
+        self.host = connection_info.host
+        self.port = connection_info.port
+        self.db_name = connection_info.db_name
+        self.ssl_mode = connection_info.ssl_mode
+        self.__db_type = db_type
+        self._verbose = verbose
+
+    def get_db_connection(
+        self, read_repeatable=False
+    ) -> PsqlConnection | MySqlConnection:
+        if self.__db_type == DbType.POSTGRES:
+            return PsqlConnection(self, read_repeatable, self._verbose)
+        elif self.__db_type == DbType.MYSQL:
+            return MySqlConnection(self, read_repeatable, self._verbose)
+        else:
+            raise ValueError("unknown db_type " + self.__db_type)

db_condenser/direct_subset.py

@@ -0,0 +1,128 @@
+import argparse
+import sys
+import time
+
+from db_condenser import config_reader, database_helper, result_tabulator
+from db_condenser.config_reader import DbConnectInfo, DbType
+from db_condenser.db_connect import DbConnect, MySqlConnection, PsqlConnection
+from db_condenser.mysql_database_creator import MySqlDatabaseCreator
+from db_condenser.psql_database_creator import PsqlDatabaseCreator
+from db_condenser.subset import Subset
+from db_condenser.subset_utils import print_progress
+
+
+def db_creator(
+    db_type: DbType, source: DbConnect, dest: DbConnect
+) -> PsqlDatabaseCreator | MySqlDatabaseCreator:
+    if db_type == DbType.POSTGRES:
+        return PsqlDatabaseCreator(source, dest, False)
+    elif db_type == DbType.MYSQL:
+        return MySqlDatabaseCreator(source, dest)
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser(description="Database Condenser")
+    parser.add_argument("--stdin", action="store_true", help="Read config from stdin")
+    parser.add_argument(
+        "-y", "--yes", action="store_true", help="Skip destination confirmation prompt"
+    )
+    parser.add_argument(
+        "--no-constraints", action="store_true", help="Skip adding constraints"
+    )
+    parser.add_argument(
+        "-v", "--verbose", action="store_true", help="Log every query with timing"
+    )
+    return parser.parse_args()
+
+
+def _confirm_destination(dest_info: DbConnectInfo):
+    print(
+        f"\nDestination: {dest_info.host}:{dest_info.port}/{dest_info.db_name}"
+        f" (user: {dest_info.user_name})"
+    )
+    response = input("Proceed with subsetting into this destination? [y/N] ")
+    if response.lower() not in ("y", "yes"):
+        print("Aborted.")
+        sys.exit(1)
+
+
+def main():
+    args = _parse_args()
+
+    if args.stdin:
+        config_reader.initialize(sys.stdin)
+    else:
+        config_reader.initialize()
+
+    config = config_reader.get_config()
+
+    db_type = config.db_type
+    source_dbc = DbConnect(
+        db_type, config.source_db_connection_info, verbose=args.verbose
+    )
+
+    dest_info = config.destination_db_connection_info
+    if not args.yes and dest_info.host not in ("localhost", "127.0.0.1"):
+        _confirm_destination(dest_info)
+
+    destination_dbc = DbConnect(db_type, dest_info, verbose=args.verbose)
+
+    database = db_creator(db_type, source_dbc, destination_dbc)
+    database.teardown()
+    database.create()
+
+    # Get list of tables to operate on
+    db_helper = database_helper.get_specific_helper()
+    all_tables = db_helper.list_all_tables(source_dbc)
+    all_tables = [x for x in all_tables if x not in config.excluded_tables]
+
+    subsetter = Subset(source_dbc, destination_dbc, all_tables)
+
+    try:
+        subsetter.prep_temp_dbs()
+        subsetter.run_middle_out()
+
+        print("Beginning pre constraint SQL calls")
+        start_time = time.time()
+        for idx, sql in enumerate(config.pre_constraint_sql):
+            print_progress(sql, idx + 1, len(config.pre_constraint_sql))
+            db_helper.run_query(sql, destination_dbc.get_db_connection())
+        print(
+            "Completed pre constraint SQL calls in {}s".format(time.time() - start_time)
+        )
+
+        print("Adding database constraints")
+        if not args.no_constraints:
+            database.add_constraints()
+
+        print("Beginning post subset SQL calls")
+        start_time = time.time()
+        for idx, sql in enumerate(config.post_subset_sql):
+            print_progress(sql, idx + 1, len(config.post_subset_sql))
+            db_helper.run_query(sql, destination_dbc.get_db_connection())
+        print("Completed post subset SQL calls in {}s".format(time.time() - start_time))
+
+        print("Resetting sequence numbering")
+        all_tables_no_pg = [table for table in all_tables if "pgbench" not in table]
+        dest_conn = destination_dbc.get_db_connection()
+        if db_type == DbType.POSTGRES:
+            assert isinstance(dest_conn, PsqlConnection)
+            db_helper.update_sequence_numbering(dest_conn, all_tables_no_pg)
+        elif db_type == DbType.MYSQL:
+            # TODO update sequencing for mysql
+            assert isinstance(dest_conn, MySqlConnection)
+            # db_helper.update_sequence_numbering(
+            #    dest_conn, all_tables_no_pg
+            # )
+
+        result_tabulator.tabulate(source_dbc, destination_dbc, all_tables)
+    except KeyboardInterrupt:
+        print("\nInterrupted — closing connections...")
+        raise
+    finally:
+        subsetter.unprep_temp_dbs()
+        subsetter.close_connections()
+
+
+if __name__ == "__main__":
+    main()

db_condenser/mysql_database_creator.py

@@ -0,0 +1,128 @@
+import os
+import subprocess
+
+from db_condenser import config_reader, db_connect
+
+
+class MySqlDatabaseCreator:
+    def __init__(self, source_connect, destination_connect):
+        self.__source_connect = source_connect
+        self.__destination_connect = destination_connect
+
+    def create(self):
+        cur_path = os.getcwd()
+
+        mysql_bin_path = get_mysql_bin_path()
+        if mysql_bin_path != "":
+            os.chdir(mysql_bin_path)
+
+        ca = connection_args(self.__source_connect)
+        args = (
+            ["mysqldump", "--no-data", "--routines"]
+            + ca
+            + [self.__source_connect.db_name]
+        )
+        result = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        if result.returncode != 0:
+            raise Exception(
+                "Capturing schema failed. Details:\n{}".format(result.stderr)
+            )
+        commands_to_create_schema = result.stdout
+
+        ca = connection_args(self.__destination_connect)
+        args = (
+            ["mysql"]
+            + ca
+            + ["-e", "CREATE DATABASE " + self.__destination_connect.db_name]
+        )
+        result = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        if result.returncode != 0:
+            raise Exception(
+                "Creating destination database failed. Details:\n{}".format(
+                    result.stderr
+                )
+            )
+
+        args = ["mysql", "-D", self.__destination_connect.db_name] + ca
+        result = subprocess.run(
+            args,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            input=commands_to_create_schema,
+        )
+        if result.returncode != 0:
+            raise Exception(
+                "Creating destination schema. Details:\n{}".format(result.stderr)
+            )
+
+        os.chdir(cur_path)
+
+    def teardown(self):
+        self.run_query_on_destination(
+            "DROP DATABASE IF EXISTS " + self.__destination_connect.db_name + ";"
+        )
+
+    def add_constraints(self):
+        # no-op for mysql
+        pass
+
+    def run_query_on_destination(self, command):
+        cur_path = os.getcwd()
+        mysql_bin_path = get_mysql_bin_path()
+        if mysql_bin_path != "":
+            os.chdir(mysql_bin_path)
+
+        ca = connection_args(self.__destination_connect)
+        args = ["mysql"] + ca + ["-e", command]
+        result = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        os.chdir(cur_path)
+        if result.returncode != 0:
+            raise Exception(
+                "Failed to run command '{}'. Details:\n{}".format(
+                    command, result.stderr
+                )
+            )
+
+
+def get_mysql_bin_path():
+    if "MYSQL_PATH" in os.environ:
+        mysql_bin_path = os.environ["MYSQL_PATH"]
+    else:
+        mysql_bin_path = ""
+    err = os.system(
+        '"'
+        + os.path.join(mysql_bin_path, "mysqldump")
+        + '"'
+        + " --help > "
+        + os.devnull
+    )
+    if err != 0:
+        raise Exception(
+            "Couldn't find MySQL utilities, consider specifying MYSQL_PATH environment variable if MySQL isn't "
+            + "in your PATH."
+        )
+    return mysql_bin_path
+
+
+def connection_args(connect):
+    host_arg = "--host={}".format(connect.host)
+    port_arg = "--port={}".format(connect.port)
+    user_arg = "--user={}".format(connect.user)
+    password_arg = "--password={}".format(connect.password)
+    return [host_arg, port_arg, user_arg, password_arg]
+
+
+# This is just for unit testing the creation and tear down processes
+if __name__ == "__main__":
+    config_reader.initialize()
+
+    config = config_reader.get_config()
+    src_connect = db_connect.DbConnect(
+        config_reader.DbType.MYSQL, config.source_db_connection_info
+    )
+    dest_connect = db_connect.DbConnect(
+        config_reader.DbType.MYSQL, config.destination_db_connection_info
+    )
+    msdbc = MySqlDatabaseCreator(src_connect, dest_connect)
+    msdbc.teardown()
+    msdbc.create()