# docker-compose.yml
version: '3.8'

services:
  django_test_db:
    build: .
    container_name: pg_partman
    environment:
      POSTGRES_USER: partman
      POSTGRES_PASSWORD: securepass
      POSTGRES_DB: parts_db
    volumes:
      - pgdata:/var/lib/postgresql/data
      - ./postgres-init.sql:/docker-entrypoint-initdb.d/init.sql
    # 映射到我本地端口是5444，方便本地调试
    ports:
      - "5444:5432"

volumes:
  pgdata: {}

CREATE EXTENSION pg_partman;

# Create your models here.
# parts/models.py
from django.db import models
from django.utils import timezone

class Part(models.Model):
    id = models.AutoField(primary_key=True)
    sn = models.CharField(max_length=32)
    create_time = models.DateTimeField(default=timezone.now)
    status = models.SmallIntegerField(choices=[
        (0, 'Pending'),
        (1, 'In Progress'),
        (2, 'Completed')
    ])

    class Meta:
        managed = False  # 禁用Django的自动建表
        db_table = 'parts'
        unique_together = (('id', 'create_time'),)

    @classmethod
    def create_partition(cls, year: int, month: int):
        """创建分区表的原生SQL操作"""
        from django.db import connection
        table_name = f'parts_{year}_{month:02d}'

        with connection.cursor() as cursor:
            cursor.execute(f"""
                CREATE TABLE {table_name} PARTITION OF parts
                FOR VALUES FROM ('{year}-{month:02d}-01') 
                TO ('{year}-{month+1:02d}-01');
            """)

from django.core.management.base import BaseCommand
from django.db import connection
from parts.models import Part
from datetime import datetime, timedelta


class Command(BaseCommand):
    help = 'Initialize the parts database with partitions'

    def handle(self, *args, **kwargs):
        try:
            with connection.cursor() as cursor:
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS parts (
                        id SERIAL,
                        sn VARCHAR(32) NOT NULL,
                        create_time TIMESTAMPTZ NOT NULL,
                        status SMALLINT,
                        PRIMARY KEY (id, create_time)
                    ) PARTITION BY RANGE (create_time);
                """)

            # 创建未来3个月的分区
            now = datetime.now()
            for month_offset in range(3):
                dt = now + timedelta(days=30*month_offset)
                Part.create_partition(dt.year, dt.month)
                self.stdout.write(self.style.SUCCESS(f'Created partition for {dt.year}-{dt.month:02d}'))

            self.stdout.write(self.style.SUCCESS('Database initialized successfully'))

        except Exception as e:
            self.stdout.write(self.style.ERROR(f'Error: {str(e)}'))

# 修改后的初始化命令（parts/management/commands/init_db.py）
class Command(BaseCommand):
    help = 'Initialize partitions with pg_partman'

    def handle(self, *args, **kwargs):
        try:
            with connection.cursor() as cursor:
                # 创建父表（如果尚未存在）
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS parts (
                        id SERIAL,
                        sn VARCHAR(32) NOT NULL,
                        create_time TIMESTAMPTZ NOT NULL,
                        status SMALLINT,
                        PRIMARY KEY (id, create_time)
                    ) PARTITION BY RANGE (create_time);
                """)

                # 配置 pg_partman 自动分区策略
                cursor.execute("""
                    SELECT partman.create_parent(
                        p_parent_table => 'public.parts',
                        p_control => 'create_time',  -- 分区键
                        p_type => 'native',          -- 使用原生分区
                        p_interval => '1 month',     -- 每月一个分区
                        p_premake => 3               -- 预创建未来3个月的分区
                    );
                """)

                # 设置保留策略（例如保留最近12个月）
                cursor.execute("""
                    UPDATE partman.part_config 
                    SET retention = '12 months', 
                        retention_keep_table = false  -- 自动删除旧分区
                    WHERE parent_table = 'public.parts';
                """)

            self.stdout.write(self.style.SUCCESS('Partitioning configured'))

        except Exception as e:
            self.stdout.write(self.style.ERROR(f'Error: {str(e)}'))

# 修改后的 Dockerfile
FROM postgres:14

# 安装插件
RUN apt-get update && \
    apt-get install -y \
    postgresql-14-partman \
    postgresql-14-cron \  # 添加 pg_cron
    && rm -rf /var/lib/apt/lists/*

###  允许 pg_cron 在启动时加载
RUN echo "shared_preload_libraries = 'pg_cron'" >> /usr/share/postgresql/postgresql.conf.sample
RUN echo "cron.database_name = 'postgres'" >> /usr/share/postgresql/postgresql.conf.sample

CREATE EXTENSION pg_cron;

-- 每天凌晨执行分区维护
SELECT cron.schedule(
    'pg_partman_maintenance',
    '0 3 * * *',  -- Cron 表达式（每天3点）
    $$CALL partman.run_maintenance_proc()$$
);

CALL partman.run_maintenance_proc();

-- 查看所有分区
SELECT * FROM partman.show_partitions('public.parts');

-- 查看配置
SELECT * FROM partman.part_config;