Files
eryao/backend/alembic/versions/20260407_0002_add_invite_codes.py
T
2026-04-10 10:40:44 +08:00

330 lines
11 KiB
Python

"""add invite_codes table and profiles referred_by
Revision ID: 20260407_0002
Revises: 20260407_0001
Create Date: 2026-04-07 00:00:00
"""
from typing import Sequence, Union
from alembic import op
revision: str = "20260407_0002"
down_revision: Union[str, Sequence[str], None] = "20260407_0001"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
op.execute(
"""
CREATE TABLE invite_codes (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
code VARCHAR(6) NOT NULL UNIQUE CHECK (code ~ '^[ABCDEFGHJKMNPQRSTUVWXYZ23456789]{6}$'),
owner_id UUID REFERENCES profiles(id) ON DELETE SET NULL,
status VARCHAR(20) NOT NULL DEFAULT 'active' CHECK (status IN ('active', 'disabled', 'expired')),
used_count INTEGER NOT NULL DEFAULT 0 CHECK (used_count >= 0),
max_uses INTEGER CHECK (max_uses IS NULL OR max_uses >= 1),
expires_at TIMESTAMPTZ NULL,
reward_config JSONB NOT NULL DEFAULT '{}',
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
)
"""
)
op.execute("CREATE INDEX ix_invite_codes_owner_id ON invite_codes(owner_id)")
op.execute(
"CREATE INDEX ix_invite_codes_code ON invite_codes(code) WHERE status = 'active'"
)
op.execute("ALTER TABLE invite_codes ENABLE ROW LEVEL SECURITY")
op.execute("DROP POLICY IF EXISTS invite_codes_all_denied ON invite_codes")
op.execute(
"CREATE POLICY invite_codes_all_denied ON invite_codes FOR ALL USING (false)"
)
op.execute(
"""
ALTER TABLE profiles ADD COLUMN referred_by UUID REFERENCES profiles(id) ON DELETE SET NULL
"""
)
op.execute("CREATE INDEX ix_profiles_referred_by ON profiles(referred_by)")
op.execute(
"""
CREATE OR REPLACE FUNCTION public.generate_invite_code()
RETURNS TEXT
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = ''
AS $$
DECLARE
chars TEXT := 'ABCDEFGHJKMNPQRSTUVWXYZ23456789';
result TEXT := '';
i INT;
BEGIN
FOR i IN 1..6 LOOP
result := result || substr(chars, floor(random() * length(chars) + 1)::int, 1);
END LOOP;
RETURN result;
END;
$$;
"""
)
op.execute(
"""
CREATE OR REPLACE FUNCTION public.initialize_profile_and_invite_code_on_signup()
RETURNS trigger
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = public
AS $$
DECLARE
v_username text;
v_ledger_id uuid;
v_event_id text;
v_invite_code text;
v_referrer_id uuid;
v_attempts int := 0;
invite_code_value text;
BEGIN
v_username := 'user_' || substring(md5(new.id::text || clock_timestamp()::text || random()::text) from 1 for 6);
v_ledger_id := md5(new.id::text || 'ledger' || clock_timestamp()::text || random()::text)::uuid;
v_event_id := 'register:' || new.id::text;
INSERT INTO public.profiles (id, username, avatar_url, bio, settings)
VALUES (
new.id,
v_username,
null,
null,
jsonb_build_object(
'version', 1,
'preferences', jsonb_build_object(
'interface_language', 'zh-CN',
'ai_language', 'zh-CN',
'timezone', 'Asia/Shanghai',
'country', 'CN'
),
'privacy', jsonb_build_object('profile_visibility', 'public'),
'notification', jsonb_build_object(
'allow_notifications', true,
'allow_vibration', true
)
)
)
ON CONFLICT (id) DO NOTHING;
INSERT INTO public.user_points (
user_id,
balance,
frozen_balance,
lifetime_earned,
lifetime_spent,
version
)
VALUES (new.id, 60, 0, 60, 0, 0)
ON CONFLICT (user_id) DO NOTHING;
INSERT INTO public.points_ledger (
id,
user_id,
direction,
amount,
balance_after,
change_type,
biz_type,
biz_id,
event_id,
operator_id,
metadata
)
VALUES (
v_ledger_id,
new.id,
1,
60,
60,
'register',
null,
null,
v_event_id,
null,
jsonb_build_object(
'schema_version', 1,
'reason_code', 'REGISTER_WELCOME',
'operator_type', 'system',
'run_id', v_event_id,
'request_id', null,
'ext', jsonb_build_object('source', 'auth_signup')
)
)
ON CONFLICT (user_id, event_id) DO NOTHING;
LOOP
BEGIN
v_invite_code := public.generate_invite_code();
INSERT INTO public.invite_codes (code, owner_id, status, used_count, max_uses, expires_at, reward_config)
VALUES (
v_invite_code,
new.id,
'active',
0,
NULL,
NULL,
'{}'::jsonb
);
EXIT;
EXCEPTION WHEN unique_violation THEN
v_attempts := v_attempts + 1;
IF v_attempts >= 100 THEN
RAISE EXCEPTION 'Failed to generate unique invite code after 100 attempts';
END IF;
END;
END LOOP;
invite_code_value := new.raw_user_meta_data ->> 'invite_code';
IF invite_code_value IS NOT NULL AND length(invite_code_value) = 6 THEN
invite_code_value := upper(invite_code_value);
IF invite_code_value ~ '^[ABCDEFGHJKMNPQRSTUVWXYZ23456789]{6}$' THEN
UPDATE public.invite_codes
SET used_count = used_count + 1
WHERE code = invite_code_value
AND status = 'active'
AND (max_uses IS NULL OR used_count < max_uses)
AND (expires_at IS NULL OR expires_at > NOW())
RETURNING owner_id INTO v_referrer_id;
IF v_referrer_id IS NOT NULL THEN
UPDATE public.profiles
SET referred_by = v_referrer_id
WHERE id = new.id;
END IF;
END IF;
END IF;
RETURN NEW;
END;
$$;
"""
)
op.execute("DROP TRIGGER IF EXISTS on_auth_user_created ON auth.users")
op.execute(
"""
CREATE TRIGGER on_auth_user_created
AFTER INSERT ON auth.users
FOR EACH ROW
EXECUTE FUNCTION public.initialize_profile_and_invite_code_on_signup()
"""
)
def downgrade() -> None:
op.execute("DROP TRIGGER IF EXISTS on_auth_user_created ON auth.users")
op.execute(
"""
CREATE OR REPLACE FUNCTION public.initialize_profile_and_points_on_signup()
RETURNS trigger
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = public
AS $$
DECLARE
v_username text;
v_ledger_id uuid;
v_event_id text;
BEGIN
v_username := 'user_' || substring(md5(new.id::text || clock_timestamp()::text || random()::text) from 1 for 6);
v_ledger_id := md5(new.id::text || 'ledger' || clock_timestamp()::text || random()::text)::uuid;
v_event_id := 'register:' || new.id::text;
INSERT INTO public.profiles (id, username, avatar_url, bio, settings)
VALUES (
new.id,
v_username,
null,
null,
jsonb_build_object(
'version', 1,
'preferences', jsonb_build_object(
'interface_language', 'zh-CN',
'ai_language', 'zh-CN',
'timezone', 'Asia/Shanghai',
'country', 'CN'
),
'privacy', jsonb_build_object('profile_visibility', 'public'),
'notification', jsonb_build_object('push_enabled', true)
)
)
ON CONFLICT (id) DO NOTHING;
INSERT INTO public.user_points (
user_id,
balance,
frozen_balance,
lifetime_earned,
lifetime_spent,
version
)
VALUES (new.id, 60, 0, 60, 0, 0)
ON CONFLICT (user_id) DO NOTHING;
INSERT INTO public.points_ledger (
id,
user_id,
direction,
amount,
balance_after,
change_type,
biz_type,
biz_id,
event_id,
operator_id,
metadata
)
VALUES (
v_ledger_id,
new.id,
1,
60,
60,
'register',
null,
null,
v_event_id,
null,
jsonb_build_object(
'schema_version', 1,
'reason_code', 'REGISTER_WELCOME',
'operator_type', 'system',
'run_id', v_event_id,
'request_id', null,
'ext', jsonb_build_object('source', 'auth_signup')
)
)
ON CONFLICT (user_id, event_id) DO NOTHING;
RETURN NEW;
END;
$$;
"""
)
op.execute(
"CREATE TRIGGER on_auth_user_created AFTER INSERT ON auth.users FOR EACH ROW EXECUTE FUNCTION public.initialize_profile_and_points_on_signup()"
)
op.execute("DROP FUNCTION IF EXISTS public.generate_invite_code()")
op.execute("DROP INDEX IF EXISTS ix_profiles_referred_by")
op.execute("ALTER TABLE profiles DROP COLUMN IF EXISTS referred_by")
op.execute("DROP POLICY IF EXISTS invite_codes_all_denied ON invite_codes")
op.execute("ALTER TABLE invite_codes DISABLE ROW LEVEL SECURITY")
op.execute("DROP INDEX IF EXISTS ix_invite_codes_code")
op.execute("DROP INDEX IF EXISTS ix_invite_codes_owner_id")
op.execute("DROP TABLE IF EXISTS invite_codes")