From ad92b45017170a7531bf6f81ebe3af962dd7f5f1 Mon Sep 17 00:00:00 2001 From: Nicolas Date: Sat, 15 Nov 2025 00:47:37 +0100 Subject: [PATCH] Add database field support to migrations Allow migrations to specify which database tables should be created in by adding an optional `database` field to migration YAML files. Changes: - Extract `database` field from migration files in migrate_up() and migrate_down() - Add `database` to format_args for SQL template substitution - Update database query methods to accept optional database parameter - Modify update_schema() to query correct databases and use qualified table names - Include database field in generated migration template (commented out) The database field is optional and defaults to the current database connection (CLICKHOUSE_DB) when not specified, maintaining backward compatibility. This enables creating tables in different databases within the same ClickHouse instance, addressing the limitation where all tables were created in the default database. --- src/houseplant/clickhouse_client.py | 21 ++++++------ src/houseplant/houseplant.py | 50 +++++++++++++++++++---------- 2 files changed, 45 insertions(+), 26 deletions(-) diff --git a/src/houseplant/clickhouse_client.py b/src/houseplant/clickhouse_client.py index 13479c6..f5c8457 100644 --- a/src/houseplant/clickhouse_client.py +++ b/src/houseplant/clickhouse_client.py @@ -192,38 +192,41 @@ def get_latest_migration(self): """) return result[0][0] if result else None - def get_database_tables(self): + def get_database_tables(self, database=None): """Get the database tables with their engines, indexes and partitioning.""" - return self.client.execute(""" + database_expr = f"'{database}'" if database else "currentDatabase()" + return self.client.execute(f""" SELECT name FROM system.tables - WHERE database = currentDatabase() + WHERE database = {database_expr} AND position('MergeTree' IN engine) > 0 AND engine NOT IN ('MaterializedView', 'Dictionary') AND name != 'schema_migrations' ORDER BY name """) - def get_database_materialized_views(self): + def get_database_materialized_views(self, database=None): """Get the database materialized views.""" - return self.client.execute(""" + database_expr = f"'{database}'" if database else "currentDatabase()" + return self.client.execute(f""" SELECT name FROM system.tables - WHERE database = currentDatabase() + WHERE database = {database_expr} AND engine = 'MaterializedView' AND name != 'schema_migrations' ORDER BY name """) - def get_database_dictionaries(self): + def get_database_dictionaries(self, database=None): """Get the database dictionaries.""" - return self.client.execute(""" + database_expr = f"'{database}'" if database else "currentDatabase()" + return self.client.execute(f""" SELECT name FROM system.tables - WHERE database = currentDatabase() + WHERE database = {database_expr} AND engine = 'Dictionary' AND name != 'schema_migrations' ORDER BY name diff --git a/src/houseplant/houseplant.py b/src/houseplant/houseplant.py index b41a3bf..9c2a949 100644 --- a/src/houseplant/houseplant.py +++ b/src/houseplant/houseplant.py @@ -116,10 +116,12 @@ def migrate_up(self, version: str | None = None): ) return + database = migration.get("database", "").strip() or self.db.database + table_definition = migration.get("table_definition", "").strip() table_settings = migration.get("table_settings", "").strip() - format_args = {"table": table} + format_args = {"table": table, "database": database} if table_definition and table_settings: format_args.update( { @@ -211,10 +213,12 @@ def migrate_down(self, version: str | None = None): ) return + database = migration.get("database", "").strip() or self.db.database + # Get migration SQL based on environment migration_env = migration.get(self.env, {}) migration_sql = ( - migration_env.get("down", {}).format(table=table).strip() + migration_env.get("down", {}).format(table=table, database=database).strip() ) if migration_sql: @@ -249,6 +253,7 @@ def generate(self, name: str): f.write(f"""version: "{timestamp}" name: {migration_name} table: +# database: # Optional: specify database name (defaults to CLICKHOUSE_DB) development: &development up: | @@ -291,12 +296,8 @@ def update_schema(self): migration_files = get_migration_files() latest_version = applied_migrations[-1][0] if applied_migrations else "0" - # Get all database objects - tables = self.db.get_database_tables() - materialized_views = self.db.get_database_materialized_views() - dictionaries = self.db.get_database_dictionaries() - # Track processed tables to ensure first migration takes precedence + # Use qualified names (database.table) to avoid conflicts across databases processed_tables = set() # Group statements by type @@ -316,43 +317,58 @@ def update_schema(self): with open(migration_file) as f: migration_data = yaml.safe_load(f) - # Extract table name from migration + # Extract table name and database from migration table_name = migration_data.get("table") if not table_name: continue + database = migration_data.get("database", "").strip() or self.db.database + + # Build qualified table identifier for tracking + qualified_name = f"{database}.{table_name}" if database else table_name + # Skip if we've already processed this table - if table_name in processed_tables: + if qualified_name in processed_tables: continue + # Get database objects from the specified database + tables = self.db.get_database_tables(database) + materialized_views = self.db.get_database_materialized_views(database) + dictionaries = self.db.get_database_dictionaries(database) + # Check tables first for table in tables: if table[0] == table_name: + # Use qualified name in SHOW CREATE query when database is specified + table_ref = f"{database}.{table_name}" if database else table_name create_stmt = self.db.client.execute( - f"SHOW CREATE TABLE {table_name}" + f"SHOW CREATE TABLE {table_ref}" )[0][0] table_statements.append(create_stmt) - processed_tables.add(table_name) + processed_tables.add(qualified_name) + break # Then materialized views for mv in materialized_views: if mv[0] == table_name: - mv_name = mv[0] - create_stmt = self.db.client.execute(f"SHOW CREATE VIEW {mv_name}")[ + mv_ref = f"{database}.{table_name}" if database else table_name + create_stmt = self.db.client.execute(f"SHOW CREATE VIEW {mv_ref}")[ 0 ][0] mv_statements.append(create_stmt) - processed_tables.add(table_name) + processed_tables.add(qualified_name) + break # Finally dictionaries for ch_dict in dictionaries: if ch_dict[0] == table_name: - dict_name = ch_dict[0] + dict_ref = f"{database}.{table_name}" if database else table_name create_stmt = self.db.client.execute( - f"SHOW CREATE DICTIONARY {dict_name}" + f"SHOW CREATE DICTIONARY {dict_ref}" )[0][0] dict_statements.append(create_stmt) - processed_tables.add(table_name) + processed_tables.add(qualified_name) + break # Write schema file with open("ch/schema.sql", "w") as f: