Skip to content

Add database field support to migrations #48

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nlenepveu wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add database field support to migrations #48

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 12 additions & 9 deletions src/houseplant/clickhouse_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
50 changes: 33 additions & 17 deletions src/houseplant/houseplant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
{
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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: |
Expand Down Expand 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
Expand All @@ -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:
Expand Down