from __future__ import annotations from sqlalchemy import inspect, text from sqlalchemy.engine import Engine TENANT_TABLES = { "client_users": None, "client_feature_access": None, "raw_materials": None, "raw_material_price_versions": None, "mixes": None, "mix_ingredients": None, "products": None, "scenarios": None, "costing_results": None, "process_cost_rules": None, "packaging_cost_rules": None, "freight_cost_rules": None, } def _has_column(engine: Engine, table_name: str, column_name: str) -> bool: inspector = inspect(engine) try: columns = inspector.get_columns(table_name) except Exception: return False return any(column["name"] == column_name for column in columns) def _add_tenant_column(engine: Engine, table_name: str) -> None: if _has_column(engine, table_name, "tenant_id"): return with engine.begin() as connection: connection.execute(text(f"ALTER TABLE {table_name} ADD COLUMN tenant_id VARCHAR(64)")) def _table_exists(engine: Engine, table_name: str) -> bool: return inspect(engine).has_table(table_name) def ensure_tenant_columns(engine: Engine) -> None: for table_name in TENANT_TABLES: if _table_exists(engine, table_name): _add_tenant_column(engine, table_name) def sync_tenant_ids(engine: Engine) -> None: if not _table_exists(engine, "client_accounts"): return with engine.begin() as connection: default_tenant = connection.execute( text("SELECT tenant_id FROM client_accounts ORDER BY id LIMIT 1") ).scalar_one_or_none() if not default_tenant: return statements = [ text( """ UPDATE client_users SET tenant_id = ( SELECT client_accounts.tenant_id FROM client_accounts WHERE client_accounts.id = client_users.client_account_id ) WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE client_feature_access SET tenant_id = ( SELECT client_accounts.tenant_id FROM client_accounts WHERE client_accounts.id = client_feature_access.client_account_id ) WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE raw_materials SET tenant_id = :default_tenant WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE raw_material_price_versions SET tenant_id = ( SELECT raw_materials.tenant_id FROM raw_materials WHERE raw_materials.id = raw_material_price_versions.raw_material_id ) WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE mixes SET tenant_id = COALESCE( ( SELECT client_accounts.tenant_id FROM client_accounts WHERE client_accounts.name = mixes.client_name ), :default_tenant ) WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE mix_ingredients SET tenant_id = ( SELECT mixes.tenant_id FROM mixes WHERE mixes.id = mix_ingredients.mix_id ) WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE products SET tenant_id = COALESCE( ( SELECT client_accounts.tenant_id FROM client_accounts WHERE client_accounts.name = products.client_name ), ( SELECT mixes.tenant_id FROM mixes WHERE mixes.id = products.mix_id ), :default_tenant ) WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE scenarios SET tenant_id = :default_tenant WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE costing_results SET tenant_id = COALESCE( ( SELECT products.tenant_id FROM products WHERE products.id = costing_results.product_id ), ( SELECT scenarios.tenant_id FROM scenarios WHERE scenarios.id = costing_results.scenario_id ), :default_tenant ) WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE process_cost_rules SET tenant_id = :default_tenant WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE packaging_cost_rules SET tenant_id = :default_tenant WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), text( """ UPDATE freight_cost_rules SET tenant_id = :default_tenant WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default' """ ), ] for statement in statements: connection.execute(statement, {"default_tenant": default_tenant})