DevOpsLab/app/alembic/versions/003_add_presets_table.py

"""add presets table

Revision ID: 003
Revises: 002
Create Date: 2024-01-01 12:00:00.000000

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = '003'
down_revision = '002_add_playbooks'
branch_labels = None
depends_on = None


def upgrade():
    # Создание таблицы presets
    op.create_table(
        'presets',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.String(), nullable=False),
        sa.Column('category', sa.String(), nullable=True, server_default='main'),
        sa.Column('description', sa.Text(), nullable=True),
        sa.Column('content', sa.Text(), nullable=False),
        sa.Column('docker_network', sa.String(), nullable=True),
        sa.Column('hosts', postgresql.JSON(astext_type=sa.Text()), nullable=True),
        sa.Column('images', postgresql.JSON(astext_type=sa.Text()), nullable=True),
        sa.Column('systemd_defaults', postgresql.JSON(astext_type=sa.Text()), nullable=True),
        sa.Column('kind_clusters', postgresql.JSON(astext_type=sa.Text()), nullable=True),
        sa.Column('status', sa.String(), nullable=True, server_default='active'),
        sa.Column('created_at', sa.DateTime(), nullable=False),
        sa.Column('updated_at', sa.DateTime(), nullable=False),
        sa.Column('created_by', sa.String(), nullable=True),
        sa.Column('updated_by', sa.String(), nullable=True),
        sa.Column('extra_data', postgresql.JSON(astext_type=sa.Text()), nullable=True),
        sa.PrimaryKeyConstraint('id')
    )
    op.create_index(op.f('ix_presets_id'), 'presets', ['id'], unique=False)
    op.create_index(op.f('ix_presets_name'), 'presets', ['name'], unique=True)
    op.create_index(op.f('ix_presets_category'), 'presets', ['category'], unique=False)
    
    # Основные preset'ы
    if presets_dir.exists():
        for preset_file in presets_dir.glob("*.yml"):
            if preset_file.name == "deploy.yml":
                continue
            
            try:
                with open(preset_file) as f:
                    content = f.read()
                    preset_data = yaml.safe_load(content) or {}
                
                # Извлечение описания из комментария
                description = None
                for line in content.split('\n'):
                    if line.strip().startswith('#description:'):
                        description = line.split('#description:')[1].strip()
                        break
                
                # Подсчет хостов
                hosts_count = len(preset_data.get('hosts', []))
                
                connection.execute(
                    sa.text("""
                        INSERT INTO presets (name, category, description, content, docker_network, hosts, images, systemd_defaults, kind_clusters, created_at, updated_at)
                        VALUES (:name, :category, :description, :content, :docker_network, :hosts, :images, :systemd_defaults, :kind_clusters, :created_at, :updated_at)
                        ON CONFLICT (name) DO NOTHING
                    """),
                    {
                        'name': preset_file.stem,
                        'category': 'main',
                        'description': description,
                        'content': content,
                        'docker_network': preset_data.get('docker_network'),
                        'hosts': preset_data.get('hosts', []),
                        'images': preset_data.get('images', {}),
                        'systemd_defaults': preset_data.get('systemd_defaults', {}),
                        'kind_clusters': preset_data.get('kind_clusters', []),
                        'created_at': datetime.utcnow(),
                        'updated_at': datetime.utcnow()
                    }
                )
            except Exception as e:
                print(f"Ошибка при импорте preset {preset_file.name}: {e}")
    
    # K8s preset'ы
    if k8s_presets_dir.exists():
        for preset_file in k8s_presets_dir.glob("*.yml"):
            try:
                with open(preset_file) as f:
                    content = f.read()
                    preset_data = yaml.safe_load(content) or {}
                
                # Извлечение описания из комментария
                description = None
                for line in content.split('\n'):
                    if line.strip().startswith('#description:'):
                        description = line.split('#description:')[1].strip()
                        break
                
                connection.execute(
                    sa.text("""
                        INSERT INTO presets (name, category, description, content, docker_network, hosts, images, systemd_defaults, kind_clusters, created_at, updated_at)
                        VALUES (:name, :category, :description, :content, :docker_network, :hosts, :images, :systemd_defaults, :kind_clusters, :created_at, :updated_at)
                        ON CONFLICT (name) DO NOTHING
                    """),
                    {
                        'name': preset_file.stem,
                        'category': 'k8s',
                        'description': description,
                        'content': content,
                        'docker_network': preset_data.get('docker_network'),
                        'hosts': preset_data.get('hosts', []),
                        'images': preset_data.get('images', {}),
                        'systemd_defaults': preset_data.get('systemd_defaults', {}),
                        'kind_clusters': preset_data.get('kind_clusters', []),
                        'created_at': datetime.utcnow(),
                        'updated_at': datetime.utcnow()
                    }
                )
            except Exception as e:
                print(f"Ошибка при импорте k8s preset {preset_file.name}: {e}")


def downgrade():
    op.drop_index(op.f('ix_presets_category'), table_name='presets')
    op.drop_index(op.f('ix_presets_name'), table_name='presets')
    op.drop_index(op.f('ix_presets_id'), table_name='presets')
    op.drop_table('presets')