db-condenser 1.0.0__tar.gz

db_condenser-1.0.0/LICENSE

@@ -0,0 +1,9 @@
+Copyright 2019, Tonic AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+

db_condenser-1.0.0/PKG-INFO

@@ -0,0 +1,119 @@
+Metadata-Version: 2.4
+Name: db-condenser
+Version: 1.0.0
+Summary: Config-driven database subsetting tool that creates representative samples while preserving referential integrity. Modernized fork of Tonic's Condenser with performance and feature improvements.
+Keywords: database,subset,postgresql,sampling,data,migration,testing
+Author: Thomas Khuu
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Intended Audience :: Developers
+Classifier: Environment :: Console
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: faker>=38.2.0,<39
+Requires-Dist: mysql-connector-python>=9.5.0,<10
+Requires-Dist: psycopg[binary]>=3.2,<4
+Requires-Dist: toposort==1.10
+Requires-Python: >=3.10
+Project-URL: Repository, https://github.com/tkhuu01/db-condenser
+Description-Content-Type: text/markdown
+
+# Database Condenser
+
+A config-driven database subsetting tool for PostgreSQL and MySQL, forked from
+Tonic's Condenser.
+
+Some changes from the original Condenser:
+
+* Concurrent thread pool usage to speed up subsetting
+* Optional temp-table strategy for memory-efficient ID batching
+* Optional Postgres COPY protocol for faster transfer
+* Automatic sequence reset after subsetting
+* Built on psycopg3 and managed with astral's uv
+
+Subsetting data is the process of taking a representative sample of your data
+in a manner that preserves the integrity of your database, e.g., give me 5% of
+my users. If you do this naively, e.g., just grab 5% of all the tables in your
+database, most likely, your database will break foreign key constraints. At
+best, you’ll end up with a statistically non-representative data sample.
+
+One common use-case is to scale down a production database to a more reasonable
+size so that it can be used in staging, test, and development environments. This
+can be done to save costs and, when used in tandem with PII removal, can be
+quite powerful as a productivity enhancer. Another example is copying specific
+rows from one database and placing them into another while maintaining referential
+integrity.
+
+You can find more about the original Condenser details
+[here](https://www.tonic.ai/blog/condenser-a-database-subsetting-tool) and
+[here](https://www.tonic.ai/blog/condenser-v2/).
+
+## Installation
+
+Six steps to set up from source, assuming Python 3.10+:
+
+1. Install [astral-uv](https://docs.astral.sh/uv/getting-started/installation/)
+
+2. Install Postgres and/or MySQL database tools. For Postgres we need `pg_dump`
+and `psql` tools; they need to be on your `$PATH` or point to them with
+`$POSTGRES_PATH`. For MySQL we need `mysqldump` and `mysql`, they can be on your
+`$PATH` or point to them with `$MYSQL_PATH`.
+
+3. Clone this project locally.
+
+4. Install the project with `uv sync --frozen`
+
+5. Setup your configuration and save it in `config.json`. The provided
+`config.json.example` has the skeleton of what you need to provide: source and
+destination database connection details, as well as subsetting goals in
+`initial_targets`. Here's an example that will collect 10% of a table
+named `public.target_table`.
+
+    ```
+    "initial_targets": [
+        {
+            "table": "public.target_table",
+            "percent": 10
+        }
+    ]
+    ```
+
+    There may be more required configuration depending on your database, but
+    simple databases should be easy. See the CONFIG.md for more details,
+    and `config.json.example_all` for all of the options in a single config file.
+
+6. Run! `$ uv run subset`
+
+Or, from PyPI:
+
+```bash
+pip install db-condenser
+```
+
+```bash
+pipx install db-condenser
+```
+
+## Running
+
+Almost all the configuration is in the `config.json` file, so running it is as simple as
+
+```bash
+uv run subset
+```
+
+Three command-line arguments are supported:
+
+`-v`: Verbose output. Useful for performance debugging. Lists almost every
+query made, and its speed.
+
+`--no-constraints`: For Postgres this will not add constraints found in the source
+database to the destination database. This option has no effect for MySQL.
+
+`-y`: Skip confirmation on subsetting to destination host if it isn't localhost
+or 127.0.0.1

db_condenser-1.0.0/README.md

@@ -0,0 +1,94 @@
+# Database Condenser
+
+A config-driven database subsetting tool for PostgreSQL and MySQL, forked from
+Tonic's Condenser.
+
+Some changes from the original Condenser:
+
+* Concurrent thread pool usage to speed up subsetting
+* Optional temp-table strategy for memory-efficient ID batching
+* Optional Postgres COPY protocol for faster transfer
+* Automatic sequence reset after subsetting
+* Built on psycopg3 and managed with astral's uv
+
+Subsetting data is the process of taking a representative sample of your data
+in a manner that preserves the integrity of your database, e.g., give me 5% of
+my users. If you do this naively, e.g., just grab 5% of all the tables in your
+database, most likely, your database will break foreign key constraints. At
+best, you’ll end up with a statistically non-representative data sample.
+
+One common use-case is to scale down a production database to a more reasonable
+size so that it can be used in staging, test, and development environments. This
+can be done to save costs and, when used in tandem with PII removal, can be
+quite powerful as a productivity enhancer. Another example is copying specific
+rows from one database and placing them into another while maintaining referential
+integrity.
+
+You can find more about the original Condenser details
+[here](https://www.tonic.ai/blog/condenser-a-database-subsetting-tool) and
+[here](https://www.tonic.ai/blog/condenser-v2/).
+
+## Installation
+
+Six steps to set up from source, assuming Python 3.10+:
+
+1. Install [astral-uv](https://docs.astral.sh/uv/getting-started/installation/)
+
+2. Install Postgres and/or MySQL database tools. For Postgres we need `pg_dump`
+and `psql` tools; they need to be on your `$PATH` or point to them with
+`$POSTGRES_PATH`. For MySQL we need `mysqldump` and `mysql`, they can be on your
+`$PATH` or point to them with `$MYSQL_PATH`.
+
+3. Clone this project locally.
+
+4. Install the project with `uv sync --frozen`
+
+5. Setup your configuration and save it in `config.json`. The provided
+`config.json.example` has the skeleton of what you need to provide: source and
+destination database connection details, as well as subsetting goals in
+`initial_targets`. Here's an example that will collect 10% of a table
+named `public.target_table`.
+
+    ```
+    "initial_targets": [
+        {
+            "table": "public.target_table",
+            "percent": 10
+        }
+    ]
+    ```
+
+    There may be more required configuration depending on your database, but
+    simple databases should be easy. See the CONFIG.md for more details,
+    and `config.json.example_all` for all of the options in a single config file.
+
+6. Run! `$ uv run subset`
+
+Or, from PyPI:
+
+```bash
+pip install db-condenser
+```
+
+```bash
+pipx install db-condenser
+```
+
+## Running
+
+Almost all the configuration is in the `config.json` file, so running it is as simple as
+
+```bash
+uv run subset
+```
+
+Three command-line arguments are supported:
+
+`-v`: Verbose output. Useful for performance debugging. Lists almost every
+query made, and its speed.
+
+`--no-constraints`: For Postgres this will not add constraints found in the source
+database to the destination database. This option has no effect for MySQL.
+
+`-y`: Skip confirmation on subsetting to destination host if it isn't localhost
+or 127.0.0.1

db_condenser-1.0.0/pyproject.toml

@@ -0,0 +1,54 @@
+[project]
+name = "db-condenser"
+description = "Config-driven database subsetting tool that creates representative samples while preserving referential integrity. Modernized fork of Tonic's Condenser with performance and feature improvements."
+authors = [{ name = "Thomas Khuu" }]
+license = "MIT"
+license-files = ["LICENSE"]
+version = "1.0.0"
+readme = "README.md"
+classifiers = [
+    "Intended Audience :: Developers",
+    "Environment :: Console",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+requires-python = ">=3.10"
+keywords = [
+    "database",
+    "subset",
+    "postgresql",
+    "sampling",
+    "data",
+    "migration",
+    "testing",
+]
+dependencies = [
+    "faker>=38.2.0,<39",
+    "mysql-connector-python>=9.5.0,<10",
+    "psycopg[binary]>=3.2,<4",
+    "toposort==1.10",
+]
+
+[project.urls]
+Repository = "https://github.com/tkhuu01/db-condenser"
+
+[project.scripts]
+subset = "db_condenser.direct_subset:main"
+
+[dependency-groups]
+dev = [
+    "pytest>=8",
+    "ruff==0.15.11",
+]
+
+[tool.ruff.lint]
+extend-select = ["I"]
+
+[build-system]
+requires = ["uv_build>=0.11.0,<0.12.0"]
+build-backend = "uv_build"

db_condenser-1.0.0/src/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-1.0.0/src/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-1.0.0/src/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-1.0.0/src/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)