Initial commit

v2_adminpanel/apply_partition_migration.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+"""
+Apply partition migration for license_heartbeats table.
+This script creates missing partitions for the current and future months.
+"""
+
+import psycopg2
+import os
+import sys
+from datetime import datetime, timedelta
+from dateutil.relativedelta import relativedelta
+
+def get_db_connection():
+    """Get database connection"""
+    return psycopg2.connect(
+        host=os.environ.get('POSTGRES_HOST', 'postgres'),
+        database=os.environ.get('POSTGRES_DB', 'v2_adminpanel'),
+        user=os.environ.get('POSTGRES_USER', 'postgres'),
+        password=os.environ.get('POSTGRES_PASSWORD', 'postgres')
+    )
+
+def create_partition(cursor, year, month):
+    """Create a partition for the given year and month"""
+    partition_name = f"license_heartbeats_{year}_{month:02d}"
+    start_date = f"{year}-{month:02d}-01"
+    
+    # Calculate end date (first day of next month)
+    if month == 12:
+        end_date = f"{year + 1}-01-01"
+    else:
+        end_date = f"{year}-{month + 1:02d}-01"
+    
+    # Check if partition already exists
+    cursor.execute("""
+        SELECT EXISTS (
+            SELECT 1 
+            FROM pg_tables 
+            WHERE tablename = %s
+        )
+    """, (partition_name,))
+    
+    exists = cursor.fetchone()[0]
+    
+    if not exists:
+        try:
+            cursor.execute(f"""
+                CREATE TABLE {partition_name} PARTITION OF license_heartbeats
+                FOR VALUES FROM ('{start_date}') TO ('{end_date}')
+            """)
+            print(f"✓ Created partition {partition_name}")
+            return True
+        except Exception as e:
+            print(f"✗ Error creating partition {partition_name}: {e}")
+            return False
+    else:
+        print(f"- Partition {partition_name} already exists")
+        return False
+
+def main():
+    """Main function"""
+    print("Applying license_heartbeats partition migration...")
+    print("-" * 50)
+    
+    try:
+        # Connect to database
+        conn = get_db_connection()
+        cursor = conn.cursor()
+        
+        # Check if license_heartbeats table exists
+        cursor.execute("""
+            SELECT EXISTS (
+                SELECT 1 
+                FROM information_schema.tables 
+                WHERE table_name = 'license_heartbeats'
+            )
+        """)
+        
+        if not cursor.fetchone()[0]:
+            print("✗ Error: license_heartbeats table does not exist!")
+            print("  Please run the init.sql script first.")
+            return 1
+        
+        # Get current date
+        current_date = datetime.now()
+        partitions_created = 0
+        
+        # Create partitions for the next 6 months (including current month)
+        for i in range(7):
+            target_date = current_date + relativedelta(months=i)
+            if create_partition(cursor, target_date.year, target_date.month):
+                partitions_created += 1
+        
+        # Commit changes
+        conn.commit()
+        
+        print("-" * 50)
+        print(f"✓ Migration complete. Created {partitions_created} new partitions.")
+        
+        # List all partitions
+        cursor.execute("""
+            SELECT tablename 
+            FROM pg_tables 
+            WHERE tablename LIKE 'license_heartbeats_%'
+            ORDER BY tablename
+        """)
+        
+        partitions = cursor.fetchall()
+        print(f"\nTotal partitions: {len(partitions)}")
+        for partition in partitions:
+            print(f"  - {partition[0]}")
+        
+        cursor.close()
+        conn.close()
+        
+        return 0
+        
+    except Exception as e:
+        print(f"✗ Error: {e}")
+        return 1
+
+if __name__ == "__main__":
+    sys.exit(main())