import os
import psycopg2
from dotenv import load_dotenv

load_dotenv()

conn = psycopg2.connect(
    host=os.getenv('DB_HOST'),
    port=os.getenv('DB_PORT'),
    dbname=os.getenv('DB_NAME'),
    user=os.getenv('DB_USER'),
    password=os.getenv('DB_PASSWORD')
)

cur = conn.cursor()

# 年度フィールドを追加
try:
    cur.execute("ALTER TABLE income_tax_table ADD COLUMN tax_year INTEGER")
    print("✓ tax_year カラムを追加")
except Exception as e:
    print(f"カラム追加スキップ（既に存在する可能性）: {e}")
    conn.rollback()

# 既存データに年度を設定
try:
    cur.execute("UPDATE income_tax_table SET tax_year = EXTRACT(YEAR FROM valid_from) WHERE tax_year IS NULL")
    print("✓ 既存データに年度を設定")
except Exception as e:
    print(f"エラー: {e}")

# インデックスを追加
try:
    cur.execute("CREATE INDEX IF NOT EXISTS idx_income_tax_year ON income_tax_table(tax_year)")
    print("✓ インデックスを追加")
except Exception as e:
    print(f"インデックス追加エラー: {e}")

conn.commit()
cur.close()
conn.close()

print("\n完了！")