Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import subprocess
- import re
- classes = ['SignupInfo']
- #
- # Get the relevant files and Django app names for the classes we need to augment
- #
- class_to_file_map = {}
- app_to_classes_map = {}
- for clazz in classes:
- grep_str = subprocess.check_output('grep -r "class ' + clazz + '" **/models.py', shell=True)
- matches = re.search('((.*)/models\.py).*', grep_str)
- app_path = matches.group(2)
- django_app = app_path.replace('/', '.')
- class_to_file_map[clazz] = matches.group(1)
- if django_app not in app_to_classes_map:
- app_to_classes_map[django_app] = []
- app_to_classes_map[django_app].append(clazz)
- #
- # Edit the class definitions to no longer have "db_index=False"
- #
- for clazz, file in class_to_file_map.items():
- f = open(file, 'r')
- # Read the lines, then edit the correct one
- lines = f.readlines()
- for i in range(len(lines)):
- if lines[i].find('class {}'.format(clazz)) != -1:
- while True:
- i += 1
- if lines[i].find('organization =') != -1:
- # Funky regex here to capture that db_index may be the first or last arg, which changes
- # the position of the comma that we need to delete.
- db_index_str = re.search('(db_index=False[,]?[\s]?|[,]?[\s]?db_index=False)', lines[i]).group(1)
- lines[i] = lines[i].replace(db_index_str, '')
- break
- break
- f.close()
- # Write the editted lines back to the file
- f = open(file, 'w+')
- f.writelines(lines)
- f.close()
- #
- # Have Django make the migrations for adding the indexes,
- # then have Django spit out the SQL for those migrations,
- # then yank the relevant index-creation SQL out of that and augment it to make the index concurrently.
- #
- app_to_index_creation_sql = {}
- django_migration_files_created = []
- django_migs_output = subprocess.check_output('python manage.py makemigrations', shell=True)
- for line in django_migs_output.split('\n'):
- mig_number_match = re.search('([0-9]{4})_', line)
- if not mig_number_match:
- continue
- mig_number = mig_number_match.group(1)
- mig_app_name = re.search('\s([\S]+)/migrations', line).group(1)
- django_migration_file_created = re.search('\s([\S]+\.py)', line).group(1)
- django_migration_files_created.append(django_migration_file_created)
- mig_sql = subprocess.check_output('python manage.py sqlmigrate {} {}'.format(mig_app_name, mig_number), shell=True)
- index_creation_sql_lines = re.findall('(CREATE INDEX [^;]+;)', mig_sql)
- for i in range(len(index_creation_sql_lines)):
- sql_line = index_creation_sql_lines[i]
- sql_line = sql_line.replace('INDEX', 'INDEX CONCURRENTLY')
- index_creation_sql_lines[i] = sql_line
- app_to_index_creation_sql[mig_app_name] = index_creation_sql_lines
- #
- # These migration files are no longer created, we're gonna make new ones from scratch. NUKE TIME
- #
- for mig_file in django_migration_files_created:
- subprocess.check_output('rm {}'.format(mig_file), shell=True)
- #
- # Old migration files: Who are you?
- # New migration files: I'm you but better
- #
- # Time to make new migration files with the SQL for making indexes concurrently
- #
- for django_app in app_to_classes_map.keys():
- makemig_output = subprocess.check_output('python manage.py makemigrations --empty ' + django_app, shell=True)
- new_migration_file = re.search('\s([\S]+/migrations.*\.py)', makemig_output).group(1)
- f = open(new_migration_file, 'r')
- lines = f.readlines()
- for i in range(len(lines)):
- line = lines[i]
- if line.startswith('class'):
- break
- lines.insert(i, '\n')
- for j, sql_line in enumerate(app_to_index_creation_sql[django_app]):
- lines.insert(i, 'INDEX_CREATION_SQL_{} = \'{}\'\n'.format(j, sql_line))
- lines.insert(i, 'from django.conf import settings\n')
- lines.append(' if not settings.INITIAL_DATABASE_SETUP:\n')
- for j, __ in enumerate(app_to_index_creation_sql[django_app]):
- lines.append(' operations.append(migrations.RunSQL(INDEX_CREATION_SQL_{}, migrations.RunSQL.noop))\n'.format(j))
- lines.append(' atomic = False')
- f.close()
- f = open(new_migration_file, 'w+')
- f.writelines(lines)
- f.close()
- #
- # Take a quick trip back in time to erase Django's memory of ever having db_index=False for the associated classes
- #
- for django_app, classes in app_to_classes_map.items():
- for clazz in classes:
- # We'll do this in a try/catch, since there's a chance that there have been
- # no AddField/AlterField operations on Organization for this class
- try:
- grep_output = subprocess.check_output('grep -i "\(addfield\|alterfield\|createmodel\)" {}/migrations/* -A 10 | grep -i organization -B 10 | grep -i {}'
- .format(django_app.replace('.', '/'),
- clazz),
- shell=True)
- except:
- continue
- files_to_alter = []
- for line in grep_output.split('\n'):
- if line:
- files_to_alter.append(re.search('([\S]+\.py)', line).group(1))
- for file in set(files_to_alter):
- f = open(file, 'r')
- lines = f.readlines()
- upcoming_edit_line = False
- for i in range(len(lines)):
- line = lines[i]
- found_index = max(line.find(clazz.lower()), line.find(clazz))
- if found_index != -1 and line[found_index - 1] == "'":
- upcoming_edit_line = True
- if upcoming_edit_line and line.find('db_index') != -1:
- # Same funky regex here as before to catch the fact that the comma we need
- # to delete may be before or after the arg but NOT BOTH
- db_index_str = re.search('(db_index=False[,]?[\s]?|[,]?[\s]?db_index=False)', line).group(1)
- lines[i] = line.replace(db_index_str, '')
- upcoming_edit_line = False
- f.close()
- f = open(file, 'w+')
- f.writelines(lines)
- f.close()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement