[feat] Extend organizations (multi-orgs, verified domains, sso providers)

api/ee/src/core/organizations/types.py

@@ -0,0 +1,129 @@
+from datetime import datetime
+from uuid import UUID
+from pydantic import BaseModel
+from typing import Optional, List, Dict, Any
+
+
+# ============================================================================
+# ORGANIZATION POLICIES
+# ============================================================================
+
+
+class OrganizationPolicy(BaseModel):
+    id: UUID
+    organization_id: UUID
+    allowed_methods: List[str]
+    invitation_only: bool
+    domains_only: bool
+    disable_root: bool
+    created_at: datetime
+    updated_at: Optional[datetime]
+
+    class Config:
+        from_attributes = True
+
+
+class OrganizationPolicyCreate(BaseModel):
+    organization_id: UUID
+    allowed_methods: List[str] = ["email:otp", "social:*"]
+    invitation_only: bool = True
+    domains_only: bool = False
+    disable_root: bool = False
+
+
+class OrganizationPolicyUpdate(BaseModel):
+    allowed_methods: Optional[List[str]] = None
+    invitation_only: Optional[bool] = None
+    domains_only: Optional[bool] = None
+    disable_root: Optional[bool] = None
+
+
+# ============================================================================
+# ORGANIZATION DOMAINS
+# ============================================================================
+
+
+class OrganizationDomain(BaseModel):
+    id: UUID
+    organization_id: UUID
+    domain: str
+    verified: bool
+    verification_token: Optional[str]
+    created_at: datetime
+    updated_at: Optional[datetime]
+
+    class Config:
+        from_attributes = True
+
+
+class OrganizationDomainCreate(BaseModel):
+    organization_id: UUID
+    domain: str
+    verification_token: Optional[str] = None
+
+
+# ============================================================================
+# ORGANIZATION PROVIDERS
+# ============================================================================
+
+
+class OrganizationProvider(BaseModel):
+    id: UUID
+    organization_id: UUID
+    slug: str
+    name: str
+    description: Optional[str]
+    enabled: bool
+    domain_id: Optional[UUID]
+    config: Dict[str, Any]
+    created_at: datetime
+    updated_at: Optional[datetime]
+
+    class Config:
+        from_attributes = True
+
+
+class OrganizationProviderCreate(BaseModel):
+    organization_id: UUID
+    slug: str
+    name: str
+    description: Optional[str] = None
+    enabled: bool = True
+    domain_id: Optional[UUID] = None
+    config: Dict[str, Any]
+
+
+class OrganizationProviderUpdate(BaseModel):
+    name: Optional[str] = None
+    description: Optional[str] = None
+    enabled: Optional[bool] = None
+    domain_id: Optional[UUID] = None
+    config: Optional[Dict[str, Any]] = None
+
+
+# ============================================================================
+# ORGANIZATION INVITATIONS
+# ============================================================================
+
+
+class OrganizationInvitation(BaseModel):
+    id: UUID
+    organization_id: UUID
+    email: str
+    role: str
+    token: str
+    status: str
+    expires_at: Optional[datetime]
+    created_at: datetime
+    updated_at: Optional[datetime]
+
+    class Config:
+        from_attributes = True
+
+
+class OrganizationInvitationCreate(BaseModel):
+    organization_id: UUID
+    email: str
+    role: str
+    token: str
+    expires_at: Optional[datetime] = None

api/ee/src/dbs/postgres/organizations/dbas.py

@@ -0,0 +1,128 @@
+import uuid_utils.compat as uuid
+from sqlalchemy import Column, String, UUID, Boolean, ARRAY
+from sqlalchemy.dialects.postgresql import JSONB
+
+from oss.src.dbs.postgres.shared.dbas import (
+    LegacyLifecycleDBA,
+    HeaderDBA,
+    OrganizationScopeDBA,
+)
+
+
+class OrganizationPolicyDBA(OrganizationScopeDBA, LegacyLifecycleDBA):
+    __abstract__ = True
+
+    id = Column(
+        UUID(as_uuid=True),
+        primary_key=True,
+        default=uuid.uuid7,
+        unique=True,
+        nullable=False,
+    )
+    allowed_methods = Column(
+        ARRAY(String),
+        nullable=False,
+        server_default="{}",
+    )
+    invitation_only = Column(
+        Boolean,
+        nullable=False,
+        server_default="true",
+    )
+    domains_only = Column(
+        Boolean,
+        nullable=False,
+        server_default="false",
+    )
+    disable_root = Column(
+        Boolean,
+        nullable=False,
+        server_default="false",
+    )
+
+
+class OrganizationDomainDBA(OrganizationScopeDBA, LegacyLifecycleDBA):
+    __abstract__ = True
+
+    id = Column(
+        UUID(as_uuid=True),
+        primary_key=True,
+        default=uuid.uuid7,
+        unique=True,
+        nullable=False,
+    )
+    domain = Column(
+        String,
+        nullable=False,
+    )
+    verified = Column(
+        Boolean,
+        nullable=False,
+        server_default="false",
+    )
+    verification_token = Column(
+        String,
+        nullable=True,
+    )
+
+
+class OrganizationProviderDBA(OrganizationScopeDBA, HeaderDBA, LegacyLifecycleDBA):
+    __abstract__ = True
+
+    id = Column(
+        UUID(as_uuid=True),
+        primary_key=True,
+        default=uuid.uuid7,
+        unique=True,
+        nullable=False,
+    )
+    slug = Column(
+        String,
+        nullable=False,
+    )
+    enabled = Column(
+        Boolean,
+        nullable=False,
+        server_default="true",
+    )
+    domain_id = Column(
+        UUID(as_uuid=True),
+        nullable=True,
+    )
+    config = Column(
+        JSONB(none_as_null=True),
+        nullable=False,
+    )
+
+
+class OrganizationInvitationDBA(OrganizationScopeDBA, LegacyLifecycleDBA):
+    __abstract__ = True
+
+    id = Column(
+        UUID(as_uuid=True),
+        primary_key=True,
+        default=uuid.uuid7,
+        unique=True,
+        nullable=False,
+    )
+    email = Column(
+        String,
+        nullable=False,
+    )
+    role = Column(
+        String,
+        nullable=False,
+    )
+    token = Column(
+        String,
+        nullable=False,
+    )
+    status = Column(
+        String,
+        nullable=False,
+        server_default="pending",
+    )
+    expires_at = Column(
+        UUID(as_uuid=True),
+        nullable=True,
+    )

api/ee/src/dbs/postgres/organizations/dbes.py

@@ -0,0 +1,99 @@
+from sqlalchemy import (
+    ForeignKeyConstraint,
+    UniqueConstraint,
+    Index,
+)
+
+from oss.src.dbs.postgres.shared.base import Base
+from ee.src.dbs.postgres.organizations.dbas import (
+    OrganizationPolicyDBA,
+    OrganizationDomainDBA,
+    OrganizationProviderDBA,
+    OrganizationInvitationDBA,
+)
+
+
+class OrganizationPolicyDBE(Base, OrganizationPolicyDBA):
+    __tablename__ = "organization_policies"
+
+    __table_args__ = (
+        ForeignKeyConstraint(
+            ["organization_id"],
+            ["organizations.id"],
+            ondelete="CASCADE",
+        ),
+        UniqueConstraint(
+            "organization_id",
+            name="uq_organization_policies_org",
+        ),
+    )
+
+
+class OrganizationDomainDBE(Base, OrganizationDomainDBA):
+    __tablename__ = "organization_domains"
+
+    __table_args__ = (
+        ForeignKeyConstraint(
+            ["organization_id"],
+            ["organizations.id"],
+            ondelete="CASCADE",
+        ),
+        UniqueConstraint(
+            "domain",
+            name="uq_organization_domains_domain",
+        ),
+        Index(
+            "ix_organization_domains_org_verified",
+            "organization_id",
+            "verified",
+        ),
+    )
+
+
+class OrganizationProviderDBE(Base, OrganizationProviderDBA):
+    __tablename__ = "organization_providers"
+
+    __table_args__ = (
+        ForeignKeyConstraint(
+            ["organization_id"],
+            ["organizations.id"],
+            ondelete="CASCADE",
+        ),
+        ForeignKeyConstraint(
+            ["domain_id"],
+            ["organization_domains.id"],
+            ondelete="SET NULL",
+        ),
+        UniqueConstraint(
+            "organization_id",
+            "slug",
+            name="uq_organization_providers_org_slug",
+        ),
+        Index(
+            "ix_organization_providers_org_enabled",
+            "organization_id",
+            "enabled",
+        ),
+    )
+
+
+class OrganizationInvitationDBE(Base, OrganizationInvitationDBA):
+    __tablename__ = "organization_invitations"
+
+    __table_args__ = (
+        ForeignKeyConstraint(
+            ["organization_id"],
+            ["organizations.id"],
+            ondelete="CASCADE",
+        ),
+        UniqueConstraint(
+            "token",
+            name="uq_organization_invitations_token",
+        ),
+        Index(
+            "ix_organization_invitations_org_email_status",
+            "organization_id",
+            "email",
+            "status",
+        ),
+    )

api/ee/src/main.py

@@ -11,6 +11,7 @@
 from ee.src.core.subscriptions.service import SubscriptionsService
 
 from ee.src.apis.fastapi.billing.router import SubscriptionsRouter
+from oss.src.apis.fastapi.auth import router as auth_router
 
 # DBS --------------------------------------------------------------------------
 
@@ -66,6 +67,13 @@ def extend_main(app: FastAPI):
         prefix="/workspaces",
     )
 
+    # Auth router at root level (no /api prefix) for OAuth callbacks
+    app.include_router(
+        auth_router.router,
+        prefix="/auth",
+        tags=["Auth"],
+    )
+
     # --------------------------------------------------------------------------
 
     return app