From de3ddb6ad2e9782386bb2264edc1de66969ce89a Mon Sep 17 00:00:00 2001 From: Saurabh Kumar <70131915+Saurabhkmr98@users.noreply.github.com> Date: Thu, 30 Jan 2025 20:25:10 +0530 Subject: [PATCH] SILO - add silo data migration script (#2331) * add silo data migration script * add silo data migration django management command and bgtask * register silo data migration task * fix entity connection migration * accept single db uri for migration command --- .../plane/bgtasks/silo_data_migration_task.py | 313 ++++++++++++++++++ .../commands/silo_data_migration.py | 15 + apiserver/plane/settings/common.py | 1 + 3 files changed, 329 insertions(+) create mode 100644 apiserver/plane/bgtasks/silo_data_migration_task.py create mode 100644 apiserver/plane/db/management/commands/silo_data_migration.py diff --git a/apiserver/plane/bgtasks/silo_data_migration_task.py b/apiserver/plane/bgtasks/silo_data_migration_task.py new file mode 100644 index 0000000000..8019fb16d6 --- /dev/null +++ b/apiserver/plane/bgtasks/silo_data_migration_task.py @@ -0,0 +1,313 @@ +# Third party imports +from celery import shared_task + +# Module imports +from plane.utils.exception_logger import log_exception + +import psycopg2 +from psycopg2.extras import RealDictCursor +import uuid +from datetime import datetime +from django.db import transaction + +# Import your Django models - adjust import paths as needed +from plane.ee.models import ( + WorkspaceCredential, + WorkspaceConnection, + WorkspaceEntityConnection, + ImportJob, + ImportReport +) + +class DatabaseMigration: + def __init__(self, db_uri, batch_size = 10): + self.conn = psycopg2.connect(dsn=db_uri) + self.cursor = self.conn.cursor(cursor_factory=RealDictCursor) + self.batch_size = batch_size + + def _execute_with_error_handling(self, query, params=None): + try: + self.cursor.execute(query, params) + except psycopg2.Error as e: + log_exception(e) + raise + + def migrate_credentials(self): + print("Migrating credentials...") + + self._execute_with_error_handling("SELECT COUNT(*) as count FROM silo.credentials") + total_count = self.cursor.fetchone()['count'] + + total_batches = (total_count + self.batch_size - 1) // self.batch_size + records_processed = 0 + + for batch_num in range(total_batches): + offset = batch_num * self.batch_size + self._execute_with_error_handling(""" + SELECT * FROM silo.credentials + ORDER BY id + LIMIT %s OFFSET %s + """, (self.batch_size, offset)) + current_batch = self.cursor.fetchall() + + try: + with transaction.atomic(): + credentials = [ + WorkspaceCredential( + id=cred['id'], + source=cred['source'], + workspace_id=cred['workspace_id'], + user_id=cred['user_id'], + source_access_token=cred['source_access_token'], + source_refresh_token=cred['source_refresh_token'], + source_hostname=cred['source_hostname'], + target_access_token=cred['target_access_token'], + is_pat=cred['is_pat'], + is_active=cred['is_active'], + created_at=datetime.now(), + updated_at=datetime.now() + ) + for cred in current_batch + ] + WorkspaceCredential.objects.bulk_create(credentials, ignore_conflicts=True) + + records_processed += len(current_batch) + print(f"Processed credentials batch {batch_num + 1}/{total_batches} " + f"({records_processed}/{total_count} records)") + + except Exception as e: + print(f"Error processing credentials batch {batch_num + 1}: {str(e)}") + raise + + def migrate_workspace_connections(self): + print("Migrating workspace connections...") + + self._execute_with_error_handling("SELECT COUNT(*) as count FROM silo.workspace_connections") + total_count = self.cursor.fetchone()['count'] + + total_batches = (total_count + self.batch_size - 1) // self.batch_size + records_processed = 0 + + for batch_num in range(total_batches): + offset = batch_num * self.batch_size + self._execute_with_error_handling(""" + SELECT * FROM silo.workspace_connections + ORDER BY id + LIMIT %s OFFSET %s + """, (self.batch_size, offset)) + current_batch = self.cursor.fetchall() + + try: + with transaction.atomic(): + connections = [ + WorkspaceConnection( + id=conn['id'], + workspace_id=conn['workspace_id'], + credential_id=conn['credentials_id'], + target_hostname=conn['target_hostname'], + source_hostname=conn['source_hostname'], + connection_type=conn['connection_type'], + connection_id=conn['connection_id'], + connection_data=conn['connection_data'], + connection_slug=conn['connection_slug'], + config=conn['config'], + scopes=[], + created_at=conn['created_at'], + updated_at=conn['updated_at'] + ) + for conn in current_batch + ] + WorkspaceConnection.objects.bulk_create(connections, ignore_conflicts=True) + + records_processed += len(current_batch) + print(f"Processed connections batch {batch_num + 1}/{total_batches} " + f"({records_processed}/{total_count} records)") + + except Exception as e: + print(f"Error processing connections batch {batch_num + 1}: {str(e)}") + raise + + def migrate_entity_connections(self): + print("Migrating entity connections...") + + self._execute_with_error_handling("SELECT COUNT(*) as count FROM silo.entity_connections") + total_count = self.cursor.fetchone()['count'] + + total_batches = (total_count + self.batch_size - 1) // self.batch_size + records_processed = 0 + + for batch_num in range(total_batches): + offset = batch_num * self.batch_size + self._execute_with_error_handling(""" + SELECT * FROM silo.entity_connections + ORDER BY id + LIMIT %s OFFSET %s + """, (self.batch_size, offset)) + current_batch = self.cursor.fetchall() + + try: + with transaction.atomic(): + entities = [ + WorkspaceEntityConnection( + id=entity['id'], + workspace_id=entity['workspace_id'], + project_id=entity['project_id'], + workspace_connection_id=entity['workspace_connection_id'], + type=entity['connection_type'], + entity_type=entity['entity_data']['type'] if 'type' in entity['entity_data'] else None, + entity_id=entity['entity_id'], + entity_slug=entity['entity_slug'], + entity_data=entity['entity_data'], + config=entity['config'], + created_at=entity['created_at'], + updated_at=entity['updated_at'] + ) + for entity in current_batch + ] + WorkspaceEntityConnection.objects.bulk_create(entities, ignore_conflicts=True) + + records_processed += len(current_batch) + print(f"Processed entities batch {batch_num + 1}/{total_batches} " + f"({records_processed}/{total_count} records)") + + except Exception as e: + print(f"Error processing entities batch {batch_num + 1}: {str(e)}") + raise + + def migrate_jobs(self): + print("Migrating jobs...") + + # First create import reports + self._execute_with_error_handling("SELECT COUNT(*) as count FROM silo.job_configs") + total_count = self.cursor.fetchone()['count'] + + report_ids = [] + total_batches = (total_count + self.batch_size - 1) // self.batch_size + records_processed = 0 + + # Create reports in batches + for batch_num in range(total_batches): + offset = batch_num * self.batch_size + self._execute_with_error_handling(""" + SELECT * FROM silo.job_configs + ORDER BY id + LIMIT %s OFFSET %s + """, (self.batch_size, offset)) + current_batch = self.cursor.fetchall() + + try: + with transaction.atomic(): + reports = [] + for _ in current_batch: + report_id = str(uuid.uuid4()) + report_ids.append(report_id) + reports.append( + ImportReport( + id=report_id, + created_at=datetime.now(), + updated_at=datetime.now(), + batch_size=0, + total_batch_count=0, + imported_batch_count=0, + errored_batch_count=0, + total_issue_count=0, + imported_issue_count=0, + errored_issue_count=0, + total_page_count=0, + imported_page_count=0, + errored_page_count=0 + ) + ) + ImportReport.objects.bulk_create(reports, ignore_conflicts=True) + + records_processed += len(current_batch) + print(f"Processed reports batch {batch_num + 1}/{total_batches} " + f"({records_processed}/{total_count} records)") + + except Exception as e: + print(f"Error processing reports batch {batch_num + 1}: {str(e)}") + raise + + records_processed = 0 + for batch_num in range(total_batches): + offset = batch_num * self.batch_size + self._execute_with_error_handling(""" + SELECT j.*, jc.meta as config_meta + FROM silo.jobs j + LEFT JOIN silo.job_configs jc ON j.config_id = jc.id + ORDER BY j.id + LIMIT %s OFFSET %s + """, (self.batch_size, offset)) + current_batch = self.cursor.fetchall() + + try: + with transaction.atomic(): + jobs = [ + ImportJob( + id=job['id'], + source=job['migration_type'], + config=job['config_meta'], + project_id=job['project_id'], + workspace_id=job['workspace_id'], + initiator_id=job['initiator_id'], + report_id=report_ids[i] if i < len(report_ids) else str(uuid.uuid4()), + status=job['status'], + created_at=job['created_at'], + updated_at=job['updated_at'], + with_issue_types=False, + success_metadata={}, + error_metadata={} + ) + for i, job in enumerate(current_batch) + ] + ImportJob.objects.bulk_create(jobs, ignore_conflicts=True) + + records_processed += len(current_batch) + print(f"Processed jobs batch {batch_num + 1}/{total_batches} " + f"({records_processed}/{total_count} records)") + + except Exception as e: + print(f"Error processing jobs batch {batch_num + 1}: {str(e)}") + raise + + def migrate_all(self): + try: + # Start transaction + print("Starting migration process...") + self.conn.autocommit = False + + # Execute migrations + self.migrate_credentials() + self.migrate_workspace_connections() + self.migrate_entity_connections() + self.migrate_jobs() + + # Commit the transaction + self.conn.commit() + print("Migration completed successfully!") + + except Exception as e: + self.conn.rollback() + print(f"Migration failed, rolling back changes: {str(e)}") + raise + finally: + self.cursor.close() + self.conn.close() + print("Database connection closed.") + + def __del__(self): + try: + if hasattr(self, 'cursor') and self.cursor is not None: + self.cursor.close() + if hasattr(self, 'conn') and self.conn is not None: + self.conn.close() + except Exception: + pass + + +@shared_task +def schedule_silo_data_migration_task(db_uri, batch_size): + migration = DatabaseMigration(db_uri, batch_size) + migration.migrate_all() + + return "Migration completed successfully!" \ No newline at end of file diff --git a/apiserver/plane/db/management/commands/silo_data_migration.py b/apiserver/plane/db/management/commands/silo_data_migration.py new file mode 100644 index 0000000000..cbb6ce84cd --- /dev/null +++ b/apiserver/plane/db/management/commands/silo_data_migration.py @@ -0,0 +1,15 @@ +from django.core.management.base import BaseCommand +from plane.bgtasks.silo_data_migration_task import schedule_silo_data_migration_task + +class Command(BaseCommand): + help = 'Migrate data from old schema to new schema using Django models' + + def handle(self, *args, **options): + db_uri = input("Enter database uri: ") + batch_size = int(input("Enter the batch size: ")) + + try: + schedule_silo_data_migration_task.delay(db_uri, batch_size) + except Exception as e: + print(f'Migration failed: {str(e)}') + raise \ No newline at end of file diff --git a/apiserver/plane/settings/common.py b/apiserver/plane/settings/common.py index 6e44bfc4b3..e447855cf9 100644 --- a/apiserver/plane/settings/common.py +++ b/apiserver/plane/settings/common.py @@ -280,6 +280,7 @@ CELERY_IMPORTS = ( # issue version tasks "plane.bgtasks.issue_version_sync", "plane.bgtasks.issue_description_version_sync", + "plane.bgtasks.silo_data_migration_task" ) # Sentry Settings