content: """add_sso_auth_provider_tables
  
  Revision ID: 20251202164905_07bea433f543
  Revises: 20250928_simple_global_cache
  Create Date: 2025-12-02 16:49:05.959012
  
  """
  from typing import Sequence, Union
  
  from alembic import op
  import sqlalchemy as sa
  from sqlalchemy.dialects import postgresql
  
  # revision identifiers, used by Alembic.
  revision: str = '20251202164905_07bea433f543'
  down_revision: Union[str, None] = '20250928_simple_global_cache'
  branch_labels: Union[str, Sequence[str], None] = None
  depends_on: Union[str, Sequence[str], None] = None
  
  
  def upgrade() -> None:
      # ### commands auto generated by Alembic - please adjust! ###
      op.create_table('auth_audit_log',
      sa.Column('id', sa.UUID(), nullable=False),
      sa.Column('user_id', sa.String(length=255), nullable=True),
      sa.Column('event_type', sa.String(length=50), nullable=False),
      sa.Column('provider_type', sa.String(length=50), nullable=True),
      sa.Column('status', sa.String(length=20), nullable=False),
      sa.Column('ip_address', sa.String(length=45), nullable=True),
      sa.Column('user_agent', sa.Text(), nullable=True),
      sa.Column('error_message', sa.Text(), nullable=True),
      sa.Column('created_at', sa.TIMESTAMP(timezone=True), nullable=False),
      sa.Column('extra_data', postgresql.JSONB(astext_type=sa.Text()), nullable=True),
      sa.ForeignKeyConstraint(['user_id'], ['users.uid'], ondelete='SET NULL'),
      sa.PrimaryKeyConstraint('id')
      )
      op.create_table('organization_sso_config',
      sa.Column('id', sa.UUID(), nullable=False),
      sa.Column('domain', sa.String(length=255), nullable=False),
      sa.Column('organization_name', sa.String(length=255), nullable=True),
      sa.Column('sso_provider', sa.String(length=50), nullable=False),
      sa.Column('sso_config', postgresql.JSONB(astext_type=sa.Text()), nullable=False),
      sa.Column('enforce_sso', sa.Boolean(), nullable=True),
      sa.Column('allow_other_providers', sa.Boolean(), nullable=True),
      sa.Column('configured_by', sa.String(length=255), nullable=True),
      sa.Column('configured_at', sa.TIMESTAMP(timezone=True), nullable=False),
      sa.Column('is_active', sa.Boolean(), nullable=True),
      sa.ForeignKeyConstraint(['configured_by'], ['users.uid'], ),
      sa.PrimaryKeyConstraint('id'),
      sa.UniqueConstraint('domain')
      )
      op.create_table('pending_provider_links',
      sa.Column('id', sa.UUID(), nullable=False),
      sa.Column('user_id', sa.String(length=255), nullable=False),
      sa.Column('provider_type', sa.String(length=50), nullable=False),
      sa.Column('provider_uid', sa.String(length=255), nullable=False),
      sa.Column('provider_data', postgresql.JSONB(astext_type=sa.Text()), nullable=False),
      sa.Column('token', sa.String(length=255), nullable=False),
      sa.Column('expires_at', sa.TIMESTAMP(timezone=True), nullable=False),
      sa.Column('created_at', sa.TIMESTAMP(timezone=True), nullable=False),
      sa.Column('ip_address', sa.String(length=45), nullable=True),
      sa.Column('user_agent', sa.Text(), nullable=True),
      sa.ForeignKeyConstraint(['user_id'], ['users.uid'], ondelete='CASCADE'),
      sa.PrimaryKeyConstraint('id'),
      sa.UniqueConstraint('token')
      )
      op.create_table('user_auth_providers',
      sa.Column('id', sa.UUID(), nullable=False),
      sa.Column('user_id', sa.String(length=255), nullable=False),
      sa.Column('provider_type', sa.String(length=50), nullable=False),
      sa.Column('provider_uid', sa.String(length=255), nullable=False),
      sa.Column('provider_data', postgresql.JSONB(astext_type=sa.Text()), nullable=True),
      sa.Column('access_token', sa.Text(), nullable=True),
      sa.Column('refresh_token', sa.Text(), nullable=True),
      sa.Column('token_expires_at', sa.TIMESTAMP(timezone=True), nullable=True),
      sa.Column('is_primary', sa.Boolean(), nullable=True),
      sa.Column('linked_at', sa.TIMESTAMP(timezone=True), nullable=False),
      sa.Column('last_used_at', sa.TIMESTAMP(timezone=True), nullable=True),
      sa.Column('linked_by_ip', sa.String(length=45), nullable=True),
      sa.Column('linked_by_user_agent', sa.Text(), nullable=True),
      sa.ForeignKeyConstraint(['user_id'], ['users.uid'], ondelete='CASCADE'),
      sa.PrimaryKeyConstraint('id'),
      sa.UniqueConstraint('provider_type', 'provider_uid', name='unique_provider_uid'),
      sa.UniqueConstraint('user_id', 'provider_type', name='unique_user_provider')
      )
      # Note: Skipping unrelated table changes detected by Alembic
      # These should be in separate migrations:
      # - inference_cache drop (separate feature)
      # - integrations index changes (separate feature)
      # - messages enum changes (separate feature)
      op.add_column('users', sa.Column('organization', sa.String(length=255), nullable=True))
      op.add_column('users', sa.Column('organization_name', sa.String(length=255), nullable=True))
  
      # Create indexes for better query performance
      op.create_index('ix_user_auth_providers_user_id', 'user_auth_providers', ['user_id'])
      op.create_index('ix_user_auth_providers_provider_type', 'user_auth_providers', ['provider_type'])
      op.create_index('ix_user_auth_providers_primary', 'user_auth_providers', ['user_id', 'is_primary'])
      op.create_index('ix_pending_links_token', 'pending_provider_links', ['token'])
      op.create_index('ix_pending_links_expires', 'pending_provider_links', ['expires_at'])
      op.create_index('ix_org_sso_domain', 'organization_sso_config', ['domain'])
      op.create_index('ix_audit_user', 'auth_audit_log', ['user_id'])
      op.create_index('ix_audit_created', 'auth_audit_log', ['created_at'])
      op.create_index('ix_audit_event', 'auth_audit_log', ['event_type'])
      op.create_index('ix_users_organization', 'users', ['organization'])
  
      # Migrate existing users' provider_info to user_auth_providers
      # This ensures existing GitHub OAuth users work with new system
      op.execute("""
          INSERT INTO user_auth_providers (
              id,
              user_id,
              provider_type,
              provider_uid,
              provider_data,
              access_token,
              is_primary,
              linked_at,
              last_used_at
          )
          SELECT
              gen_random_uuid(),
              uid,
              'firebase_github',
              uid,
              provider_info,
              CASE
                  WHEN provider_info IS NOT NULL AND provider_info ? 'access_token'
                  THEN provider_info->>'access_token'
                  ELSE NULL
              END,
              true,  -- Existing provider is primary
              created_at,
              last_login_at
          FROM users
          WHERE provider_info IS NOT NULL
          AND uid IS NOT NULL
          ON CONFLICT DO NOTHING
      """)
  
      # Extract organization from email domain
      # Filters out common personal email providers
      op.execute("""
          UPDATE users
          SET organization = CASE
              WHEN SPLIT_PART(email, '@', 2) IN (
                  'gmail.com', 'yahoo.com', 'hotmail.com', 'outlook.com',
                  'icloud.com', 'protonmail.com', 'mail.com', 'aol.com'
              ) THEN NULL
              ELSE SPLIT_PART(email, '@', 2)
          END
          WHERE organization IS NULL
      """)
  
      # ### end Alembic commands ###
  
  
  def downgrade() -> None:
      # ### commands auto generated by Alembic - please adjust! ###
      # Drop indexes first
      op.drop_index('ix_users_organization', table_name='users')
      op.drop_index('ix_audit_event', table_name='auth_audit_log')
      op.drop_index('ix_audit_created', table_name='auth_audit_log')
      op.drop_index('ix_audit_user', table_name='auth_audit_log')
      op.drop_index('ix_org_sso_domain', table_name='organization_sso_config')
      op.drop_index('ix_pending_links_expires', table_name='pending_provider_links')
      op.drop_index('ix_pending_links_token', table_name='pending_provider_links')
      op.drop_index('ix_user_auth_providers_primary', table_name='user_auth_providers')
      op.drop_index('ix_user_auth_providers_provider_type', table_name='user_auth_providers')
      op.drop_index('ix_user_auth_providers_user_id', table_name='user_auth_providers')
  
      # Drop columns
      op.drop_column('users', 'organization_name')
      op.drop_column('users', 'organization')
  
      # Note: Skipping unrelated rollback operations
      # These were in the autogenerated migration but should be separate
      op.drop_table('user_auth_providers')
      op.drop_table('pending_provider_links')
      op.drop_table('organization_sso_config')
      op.drop_table('auth_audit_log')
      # ### end Alembic commands ###
  