330 lines
11 KiB
Python
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, 100, 0, 100, 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,
|
|
100,
|
|
100,
|
|
'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, 100, 0, 100, 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,
|
|
100,
|
|
100,
|
|
'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")
|