122 Zeilen
3.6 KiB
Python
122 Zeilen
3.6 KiB
Python
#!/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()) |