55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
import os
|
|
import sys
|
|
import yaml
|
|
import psycopg
|
|
from psycopg import sql
|
|
|
|
def main():
|
|
mapping_file = os.getenv("SCHEMA_MAPPING_FILE", "/rig/datadir/datasource/schema_mapping.yaml")
|
|
|
|
if len(sys.argv) >= 2:
|
|
mapping_file = sys.argv[1]
|
|
|
|
db_uri = os.getenv("TARGET_DB_URI", "postgresql://postgres:postgres@postgres:5432/postgres")
|
|
|
|
print(f"Loading schema mapping from {mapping_file}...")
|
|
with open(mapping_file, 'r') as f:
|
|
config = yaml.safe_load(f)
|
|
|
|
with psycopg.connect(db_uri) as conn:
|
|
with conn.cursor() as cur:
|
|
for table in config.get('tables', []):
|
|
source = table['source']
|
|
target = table['target']
|
|
columns = table['columns']
|
|
|
|
print(f"Migrating {source} to {target}...")
|
|
|
|
if table.get('truncate', False):
|
|
try:
|
|
cur.execute(f"TRUNCATE TABLE {target} RESTART IDENTITY CASCADE;")
|
|
except psycopg.errors.UndefinedTable:
|
|
print(f" Target table {target} does not exist yet; skipping truncate.")
|
|
conn.rollback()
|
|
|
|
target_cols = ", ".join(columns.keys())
|
|
source_exprs = ", ".join(columns.values())
|
|
|
|
query = f"""
|
|
INSERT INTO {target} ({target_cols})
|
|
SELECT {source_exprs}
|
|
FROM {source}
|
|
"""
|
|
try:
|
|
cur.execute(query)
|
|
print(f" Successfully migrated {source}.")
|
|
except Exception as e:
|
|
print(f" Failed to migrate {source}: {e}")
|
|
conn.rollback()
|
|
raise
|
|
conn.commit()
|
|
print("Schema migration complete.")
|
|
|
|
if __name__ == '__main__':
|
|
main()
|