ref: make Pipeline generic on the Session type too

src/sentry/auth/helper.py

@@ -35,6 +35,7 @@
 from sentry.auth.partnership_configs import ChannelName
 from sentry.auth.provider import MigratingIdentityId, Provider
 from sentry.auth.providers.fly.provider import FlyOAuth2Provider
+from sentry.auth.store import FLOW_LOGIN, FLOW_SETUP_PROVIDER, AuthHelperSessionStore
 from sentry.auth.superuser import is_active_superuser
 from sentry.hybridcloud.models.outbox import outbox_context
 from sentry.locks import locks
@@ -48,7 +49,7 @@
     RpcOrganizationMemberFlags,
     organization_service,
 )
-from sentry.pipeline import Pipeline, PipelineSessionStore
+from sentry.pipeline import Pipeline
 from sentry.pipeline.provider import PipelineProvider
 from sentry.signals import sso_enabled, user_signup
 from sentry.tasks.auth.auth import email_missing_links_control
@@ -58,7 +59,6 @@
 from sentry.utils.hashlib import md5_text
 from sentry.utils.http import absolute_uri
 from sentry.utils.retries import TimedRetryPolicy
-from sentry.utils.session_store import redis_property
 from sentry.utils.urls import add_params_to_url
 from sentry.web.forms.accounts import AuthenticationForm
 from sentry.web.helpers import render_to_response
@@ -83,27 +83,6 @@
 ERR_INVALID_IDENTITY = _("The provider did not return a valid user identity.")
 
 
-class AuthHelperSessionStore(PipelineSessionStore):
-    redis_namespace = "auth"
-
-    @property
-    def session_key(self) -> str:
-        return "auth_key"
-
-    flow = redis_property("flow")
-    referrer = redis_property("referrer")
-
-    def mark_session(self) -> None:
-        super().mark_session()
-        self.request.session.modified = True
-
-    def is_valid(self) -> bool:
-        return super().is_valid() and self.flow in (
-            AuthHelper.FLOW_LOGIN,
-            AuthHelper.FLOW_SETUP_PROVIDER,
-        )
-
-
 @dataclass
 class AuthIdentityHandler:
     # SSO auth handler
@@ -650,7 +629,7 @@ def handle_new_user(self) -> AuthIdentity:
         return auth_identity
 
 
-class AuthHelper(Pipeline[AuthProvider]):
+class AuthHelper(Pipeline[AuthProvider, AuthHelperSessionStore]):
     """
     Helper class which is passed into AuthView's.
 
@@ -671,11 +650,6 @@ class AuthHelper(Pipeline[AuthProvider]):
        it with their account (thus creating a new account).
     """
 
-    # logging in or registering
-    FLOW_LOGIN = 1
-    # configuring the provider
-    FLOW_SETUP_PROVIDER = 2
-
     pipeline_name = "pipeline"
     provider_manager = manager
     provider_model_cls = AuthProvider
@@ -731,7 +705,7 @@ def __init__(
 
     def get_provider(
         self, provider_key: str | None, *, organization: RpcOrganization | None
-    ) -> PipelineProvider[AuthProvider]:
+    ) -> PipelineProvider[AuthProvider, AuthHelperSessionStore]:
         if self.provider_model:
             return self.provider_model.get_provider()
         elif provider_key:
@@ -741,15 +715,15 @@ def get_provider(
 
     def get_pipeline_views(self) -> Sequence[View]:
         assert isinstance(self.provider, Provider)
-        if self.flow == self.FLOW_LOGIN:
+        if self.flow == FLOW_LOGIN:
             return self.provider.get_auth_pipeline()
-        elif self.flow == self.FLOW_SETUP_PROVIDER:
+        elif self.flow == FLOW_SETUP_PROVIDER:
             return self.provider.get_setup_pipeline()
         else:
             raise NotImplementedError
 
     def is_valid(self) -> bool:
-        return super().is_valid() and self.state.flow in (self.FLOW_LOGIN, self.FLOW_SETUP_PROVIDER)
+        return super().is_valid() and self.state.flow in (FLOW_LOGIN, FLOW_SETUP_PROVIDER)
 
     def get_initial_state(self) -> Mapping[str, Any]:
         state = dict(super().get_initial_state())
@@ -772,10 +746,10 @@ def finish_pipeline(self) -> HttpResponseBase:
         except IdentityNotValid as error:
             return self.error(str(error) or ERR_INVALID_IDENTITY)
 
-        if self.state.flow == self.FLOW_LOGIN:
+        if self.state.flow == FLOW_LOGIN:
             # create identity and authenticate the user
             response = self._finish_login_pipeline(identity)
-        elif self.state.flow == self.FLOW_SETUP_PROVIDER:
+        elif self.state.flow == FLOW_SETUP_PROVIDER:
             # Configuring the SSO Auth provider
             response = self._finish_setup_pipeline(identity)
         else:
@@ -921,11 +895,11 @@ def _finish_setup_pipeline(self, identity: Mapping[str, Any]) -> HttpResponseRed
     def error(self, message: str | _StrPromise) -> HttpResponseRedirect:
         redirect_uri = "/"
 
-        if self.state.flow == self.FLOW_LOGIN:
+        if self.state.flow == FLOW_LOGIN:
             # create identity and authenticate the user
             redirect_uri = reverse("sentry-auth-organization", args=[self.organization.slug])
 
-        elif self.state.flow == self.FLOW_SETUP_PROVIDER:
+        elif self.state.flow == FLOW_SETUP_PROVIDER:
             redirect_uri = reverse(
                 "sentry-organization-auth-settings", args=[self.organization.slug]
             )

src/sentry/auth/provider.py

@@ -10,6 +10,7 @@
 from django.utils.encoding import force_str
 
 from sentry.auth.services.auth.model import RpcAuthProvider
+from sentry.auth.store import AuthHelperSessionStore
 from sentry.auth.view import AuthView
 from sentry.models.authidentity import AuthIdentity
 from sentry.models.authprovider import AuthProvider
@@ -39,7 +40,7 @@ def __str__(self) -> str:
         return force_str(self.id)
 
 
-class Provider(PipelineProvider[AuthProvider], abc.ABC):
+class Provider(PipelineProvider[AuthProvider, AuthHelperSessionStore], abc.ABC):
     """
     A provider indicates how authenticate should happen for a given service,
     including its configuration and basic identity management.
@@ -84,7 +85,7 @@ def get_pipeline_views(self) -> Sequence[AuthView]:
 
     # TODO: state should be Mapping[str, Any]?
     # Must be reconciled with sentry.pipeline.base.Pipeline.fetch_state
-    def build_config(self, state: Any) -> Mapping[str, Any]:
+    def build_config(self, state: Any) -> dict[str, Any]:
         """
         Return a mapping containing provider configuration.
 

src/sentry/auth/providers/saml2/provider.py

@@ -21,6 +21,7 @@
 from sentry import features, options
 from sentry.auth.exceptions import IdentityNotValid
 from sentry.auth.provider import Provider
+from sentry.auth.store import FLOW_LOGIN
 from sentry.auth.view import AuthView
 from sentry.models.authidentity import AuthIdentity
 from sentry.models.authprovider import AuthProvider
@@ -118,7 +119,7 @@ def dispatch(self, request: HttpRequest, organization_slug: str) -> HttpResponse
             request=request,
             organization=(org_context.organization),
             auth_provider=auth_provider,
-            flow=AuthHelper.FLOW_LOGIN,
+            flow=FLOW_LOGIN,
         )
 
         helper.initialize()

src/sentry/auth/store.py

@@ -0,0 +1,25 @@
+from sentry.pipeline.store import PipelineSessionStore
+from sentry.utils.session_store import redis_property
+
+# logging in or registering
+FLOW_LOGIN = 1
+# configuring the provider
+FLOW_SETUP_PROVIDER = 2
+
+
+class AuthHelperSessionStore(PipelineSessionStore):
+    redis_namespace = "auth"
+
+    @property
+    def session_key(self) -> str:
+        return "auth_key"
+
+    flow = redis_property("flow")
+    referrer = redis_property("referrer")
+
+    def mark_session(self) -> None:
+        super().mark_session()
+        self.request.session.modified = True
+
+    def is_valid(self) -> bool:
+        return super().is_valid() and self.flow in (FLOW_LOGIN, FLOW_SETUP_PROVIDER)

src/sentry/identity/base.py

@@ -2,12 +2,12 @@
 import logging
 from typing import Any
 
+from sentry.identity.pipeline_types import IdentityPipelineProviderT
 from sentry.identity.services.identity.model import RpcIdentity
-from sentry.pipeline import PipelineProvider
-from sentry.users.models.identity import Identity, IdentityProvider
+from sentry.users.models.identity import Identity
 
 
-class Provider(PipelineProvider[IdentityProvider], abc.ABC):
+class Provider(IdentityPipelineProviderT, abc.ABC):
     """
     A provider indicates how identity authenticate should happen for a given service.
     """

src/sentry/identity/bitbucket/provider.py

@@ -2,23 +2,20 @@
 from django.http.response import HttpResponseBase, HttpResponseRedirect
 
 from sentry.identity.base import Provider
-from sentry.pipeline import Pipeline, PipelineView
-from sentry.users.models.identity import IdentityProvider
+from sentry.identity.pipeline_types import IdentityPipelineT, IdentityPipelineViewT
 from sentry.utils.http import absolute_uri
 
 
 class BitbucketIdentityProvider(Provider):
     key = "bitbucket"
     name = "Bitbucket"
 
-    def get_pipeline_views(self) -> list[PipelineView[IdentityProvider]]:
+    def get_pipeline_views(self) -> list[IdentityPipelineViewT]:
         return [BitbucketLoginView()]
 
 
-class BitbucketLoginView(PipelineView[IdentityProvider]):
-    def dispatch(
-        self, request: HttpRequest, pipeline: Pipeline[IdentityProvider]
-    ) -> HttpResponseBase:
+class BitbucketLoginView(IdentityPipelineViewT):
+    def dispatch(self, request: HttpRequest, pipeline: IdentityPipelineT) -> HttpResponseBase:
         from sentry.integrations.base import IntegrationDomain
         from sentry.integrations.utils.metrics import (
             IntegrationPipelineViewEvent,

src/sentry/identity/oauth2.py

@@ -18,6 +18,7 @@
 from sentry.auth.exceptions import IdentityNotValid
 from sentry.exceptions import NotRegistered
 from sentry.http import safe_urlopen, safe_urlread
+from sentry.identity.pipeline_types import IdentityPipelineT, IdentityPipelineViewT
 from sentry.identity.services.identity.model import RpcIdentity
 from sentry.integrations.base import IntegrationDomain
 from sentry.integrations.utils.metrics import (
@@ -26,9 +27,8 @@
     IntegrationPipelineViewEvent,
     IntegrationPipelineViewType,
 )
-from sentry.pipeline import Pipeline, PipelineView
 from sentry.shared_integrations.exceptions import ApiError
-from sentry.users.models.identity import Identity, IdentityProvider
+from sentry.users.models.identity import Identity
 from sentry.utils.http import absolute_uri
 
 from .base import Provider
@@ -40,7 +40,7 @@
 ERR_TOKEN_RETRIEVAL = "Failed to retrieve token from the upstream service."
 
 
-def _redirect_url(pipeline: Pipeline[IdentityProvider]) -> str:
+def _redirect_url(pipeline: IdentityPipelineT) -> str:
     associate_url = reverse(
         "sentry-extension-setup",
         kwargs={
@@ -119,7 +119,7 @@ def get_oauth_scopes(self):
     def get_refresh_token_headers(self):
         return None
 
-    def get_pipeline_views(self) -> list[PipelineView[IdentityProvider]]:
+    def get_pipeline_views(self) -> list[IdentityPipelineViewT]:
         return [
             OAuth2LoginView(
                 authorize_url=self.get_oauth_authorize_url(),
@@ -243,7 +243,7 @@ def record_event(event: IntegrationPipelineViewType, provider: str):
     )
 
 
-class OAuth2LoginView(PipelineView[IdentityProvider]):
+class OAuth2LoginView(IdentityPipelineViewT):
     authorize_url: str | None = None
     client_id: str | None = None
     scope = ""
@@ -273,9 +273,7 @@ def get_authorize_params(self, state, redirect_uri):
         }
 
     @method_decorator(csrf_exempt)
-    def dispatch(
-        self, request: HttpRequest, pipeline: Pipeline[IdentityProvider]
-    ) -> HttpResponseBase:
+    def dispatch(self, request: HttpRequest, pipeline: IdentityPipelineT) -> HttpResponseBase:
         with record_event(IntegrationPipelineViewType.OAUTH_LOGIN, pipeline.provider.key).capture():
             for param in ("code", "error", "state"):
                 if param in request.GET:
@@ -295,7 +293,7 @@ def dispatch(
             return HttpResponseRedirect(redirect_uri)
 
 
-class OAuth2CallbackView(PipelineView[IdentityProvider]):
+class OAuth2CallbackView(IdentityPipelineViewT):
     access_token_url: str | None = None
     client_id: str | None = None
     client_secret: str | None = None
@@ -319,7 +317,7 @@ def get_token_params(self, code, redirect_uri):
         }
 
     def exchange_token(
-        self, request: HttpRequest, pipeline: Pipeline[IdentityProvider], code: str
+        self, request: HttpRequest, pipeline: IdentityPipelineT, code: str
     ) -> dict[str, str]:
         with record_event(
             IntegrationPipelineViewType.TOKEN_EXCHANGE, pipeline.provider.key
@@ -363,9 +361,7 @@ def exchange_token(
                     "error_description": "We were not able to parse a JSON response, please try again.",
                 }
 
-    def dispatch(
-        self, request: HttpRequest, pipeline: Pipeline[IdentityProvider]
-    ) -> HttpResponseBase:
+    def dispatch(self, request: HttpRequest, pipeline: IdentityPipelineT) -> HttpResponseBase:
         with record_event(
             IntegrationPipelineViewType.OAUTH_CALLBACK, pipeline.provider.key
         ).capture() as lifecycle:

src/sentry/identity/pipeline.py

@@ -1,16 +1,18 @@
+from __future__ import annotations
+
 from django.contrib import messages
 from django.http import HttpResponseRedirect
 from django.urls import reverse
 from django.utils.translation import gettext_lazy as _
 
 from sentry import features, options
+from sentry.identity.pipeline_types import IdentityPipelineProviderT, IdentityPipelineT
 from sentry.integrations.base import IntegrationDomain
 from sentry.integrations.utils.metrics import (
     IntegrationPipelineViewEvent,
     IntegrationPipelineViewType,
 )
 from sentry.organizations.services.organization.model import RpcOrganization
-from sentry.pipeline import Pipeline, PipelineProvider
 from sentry.users.models.identity import Identity, IdentityProvider
 from sentry.utils import metrics
 
@@ -19,15 +21,15 @@
 IDENTITY_LINKED = _("Your {identity_provider} account has been associated with your Sentry account")
 
 
-class IdentityProviderPipeline(Pipeline[IdentityProvider]):
+class IdentityProviderPipeline(IdentityPipelineT):
     pipeline_name = "identity_provider"
     provider_manager = default_manager
     provider_model_cls = IdentityProvider
 
     # TODO(iamrajjoshi): Delete this after Azure DevOps migration is complete
     def get_provider(
         self, provider_key: str, *, organization: RpcOrganization | None
-    ) -> PipelineProvider[IdentityProvider]:
+    ) -> IdentityPipelineProviderT:
         if provider_key == "vsts" and features.has(
             "organizations:migrate-azure-devops-integration", organization
         ):

src/sentry/identity/pipeline_types.py

@@ -0,0 +1,11 @@
+from typing import TypeAlias
+
+from sentry.pipeline.base import Pipeline
+from sentry.pipeline.provider import PipelineProvider
+from sentry.pipeline.store import PipelineSessionStore
+from sentry.pipeline.views.base import PipelineView
+from sentry.users.models.identity import IdentityProvider
+
+IdentityPipelineT: TypeAlias = Pipeline[IdentityProvider, PipelineSessionStore]
+IdentityPipelineProviderT: TypeAlias = PipelineProvider[IdentityProvider, PipelineSessionStore]
+IdentityPipelineViewT: TypeAlias = PipelineView[IdentityProvider, PipelineSessionStore]