feat: add invite code feature (create, validate, referrer tracking)
This commit is contained in:
qzl
@@ -0,0 +1,197 @@
|
||||
"""invite_codes table and profile referral fields
|
||||
|
||||
Revision ID: 202602270006
|
||||
Revises: 202602260005
|
||||
Create Date: 2026-02-27 10:00:00
|
||||
"""
|
||||
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "202602270006"
|
||||
down_revision: Union[str, Sequence[str], None] = "202602260005"
|
||||
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(8) NOT NULL UNIQUE,
|
||||
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 = public
|
||||
AS $$
|
||||
DECLARE
|
||||
chars TEXT := 'ABCDEFGHJKMNPQRSTUVWXYZ23456789';
|
||||
result TEXT := '';
|
||||
i INT;
|
||||
BEGIN
|
||||
FOR i IN 1..8 LOOP
|
||||
result := result || substr(chars, floor(random() * length(chars) + 1)::int, 1);
|
||||
END LOOP;
|
||||
RETURN result;
|
||||
END;
|
||||
$$;
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
|
||||
RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
SECURITY DEFINER
|
||||
SET search_path = public
|
||||
AS $$
|
||||
DECLARE
|
||||
invite_code_value TEXT;
|
||||
referrer_id UUID;
|
||||
new_code TEXT;
|
||||
attempts INT := 0;
|
||||
BEGIN
|
||||
INSERT INTO public.profiles (id, username, avatar_url, bio, settings, referred_by, created_at, updated_at)
|
||||
VALUES (
|
||||
NEW.id,
|
||||
COALESCE(
|
||||
NEW.raw_user_meta_data ->> 'username',
|
||||
split_part(NEW.email, '@', 1),
|
||||
'user_' || substring(NEW.id::text, 1, 8)
|
||||
),
|
||||
NULL,
|
||||
NULL,
|
||||
'{}'::jsonb,
|
||||
NULL,
|
||||
now(),
|
||||
now()
|
||||
)
|
||||
ON CONFLICT (id) DO NOTHING;
|
||||
|
||||
LOOP
|
||||
BEGIN
|
||||
new_code := public.generate_invite_code();
|
||||
INSERT INTO public.invite_codes (code, owner_id, status, used_count, max_uses, expires_at, reward_config)
|
||||
VALUES (
|
||||
new_code,
|
||||
NEW.id,
|
||||
'active',
|
||||
0,
|
||||
NULL,
|
||||
NULL,
|
||||
'{}'::jsonb
|
||||
);
|
||||
EXIT;
|
||||
EXCEPTION WHEN unique_violation THEN
|
||||
attempts := attempts + 1;
|
||||
IF 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) = 8 THEN
|
||||
invite_code_value := upper(invite_code_value);
|
||||
IF invite_code_value ~ '^[ABCDEFGHJKMNPQRSTUVWXYZ23456789]{8}$' 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 referrer_id;
|
||||
|
||||
IF referrer_id IS NOT NULL THEN
|
||||
UPDATE public.profiles
|
||||
SET referred_by = referrer_id
|
||||
WHERE id = NEW.id;
|
||||
END IF;
|
||||
END IF;
|
||||
END IF;
|
||||
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP FUNCTION IF EXISTS public.create_profile_for_new_user()")
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
|
||||
RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
SECURITY DEFINER
|
||||
SET search_path = public
|
||||
AS $$
|
||||
BEGIN
|
||||
INSERT INTO public.profiles (id, username, avatar_url, bio, settings, created_at, updated_at)
|
||||
VALUES (
|
||||
NEW.id,
|
||||
COALESCE(
|
||||
NEW.raw_user_meta_data ->> 'username',
|
||||
split_part(NEW.email, '@', 1),
|
||||
'user_' || substring(NEW.id::text, 1, 8)
|
||||
),
|
||||
NULL,
|
||||
NULL,
|
||||
'{}'::jsonb,
|
||||
now(),
|
||||
now()
|
||||
)
|
||||
ON CONFLICT (id) DO NOTHING;
|
||||
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$;
|
||||
"""
|
||||
)
|
||||
|
||||
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")
|
||||
Reference in New Issue
Block a user