data_migration.py

import glob, os
import pathlib
from django.apps import apps
from django.db.models.fields.related import OneToOneField, ForeignKey
from collections import OrderedDict
from django.template import Template, Context
from django.db import connection


def execute_raw_sql(query):
    with connection.cursor() as cursor:
        cursor.execute(query)


def get_parent_tables(model):
    fields = model._meta.get_fields()
    parent_tables = set()
    related_fields = [
        OneToOneField,
        ForeignKey
    ]
    for field in fields:
        if field.__class__ in related_fields:
            related_model = field.related_model
            tbl = related_model._meta.db_table
            parent_tables.add(tbl)
            parent_tables.update(get_parent_tables(related_model))
    return parent_tables


def get_sql_templates():
    pwd = str(pathlib.Path('__file__').parent.absolute())
    os.chdir(pwd)
    return glob.glob("*.sql")


def get_model_table_map():
    models = apps.get_models()
    table_map = {}
    for model in models:
        table = model._meta.db_table
        table_map[table] = model
        # m2m tables
        m2m_fields = model._meta.many_to_many
        for field in m2m_fields:
            m2m_cls = getattr(getattr(model,field.name), 'through')
            table = m2m_cls._meta.db_table
            table_map[table] = m2m_cls
    return table_map


def get_tables_with_order():
    templates = get_sql_templates()
    model_table_map = get_model_table_map()
    tbl_data = {}
    for template in templates:
        table_name = template.split('.')[0]
        if table_name == 'django_migrations':
            continue
        model = model_table_map[table_name]
        parent_tables = get_parent_tables(model)
        tbl_data[table_name] = {'parent_tables': parent_tables, 'count': len(parent_tables)}
    return sorted(tbl_data.items(), key=lambda item: item[1]['count'])


def migrate_data_for_retailer(retailer_name, schema_name):
    tables_with_order = get_tables_with_order()
    for tbl, dep in tables_with_order:
        print(tbl, dep)
        pwd = str(pathlib.Path('__file__').parent.absolute())
        tmpl = os.path.join(pwd, f'{tbl}.sql')
        print(f"file path: {tmpl}")
        with open(tmpl) as f:
            tmpl_str = f.read()
            template = Template(tmpl_str)
            context = Context({
                'retailer_name': retailer_name,
                'schema_name': schema_name
            })
            sql = template.render(context)
            print(f"executing sql:\n {sql}\n")
            execute_raw_sql(sql)
            print("success...")