[DOP-25470] Enable authentication using personal tokens

Author: dolfinus

data_rentgen/consumer/__init__.py

@@ -74,7 +74,7 @@ async def security_lifespan(context: ContextRepo):
                 await settings.kafka.security.destroy()
                 tg.cancel_scope.cancel()
         except ExceptionGroup as e:
-            for exception in e.exceptions:
+            for exception in e.exceptions:  # type: ignore[attr-defined]
                 raise exception from None
 
     return FastStream(

data_rentgen/consumer/settings/__init__.py

@@ -39,7 +39,10 @@ class ConsumerApplicationSettings(BaseSettings):
         DATA_RENTGEN__LOGGING__PRESET=json
     """  # noqa: E501
 
-    database: DatabaseSettings = Field(description=":ref:`Database settings <configuration-database>`")
+    database: DatabaseSettings = Field(
+        default_factory=DatabaseSettings,  # type: ignore[arg-type]
+        description=":ref:`Database settings <configuration-database>`",
+    )
     logging: LoggingSettings = Field(
         default_factory=LoggingSettings,
         description=":ref:`Logging settings <configuration-consumer-logging>`",

data_rentgen/http2kafka/__init__.py

@@ -9,13 +9,16 @@
 from faststream._compat import ExceptionGroup
 from faststream.kafka import KafkaBroker
 from faststream.kafka.publisher.asyncapi import AsyncAPIDefaultPublisher
+from sqlalchemy.ext.asyncio import AsyncSession
 
 import data_rentgen
+from data_rentgen.db.factory import session_generator
 from data_rentgen.http2kafka.router import router as openlineage_router
 from data_rentgen.http2kafka.settings import Http2KafkaApplicationSettings
 from data_rentgen.logging.setup_logging import setup_logging
 from data_rentgen.server.api.handlers import apply_exception_handlers
 from data_rentgen.server.middlewares import apply_middlewares
+from data_rentgen.server.providers.auth.personal_token_provider import PersonalTokenAuthProvider
 
 logger = logging.getLogger(__name__)
 
@@ -34,7 +37,7 @@ async def lifespan(app: FastAPI):
             await settings.kafka.security.destroy()
             tg.cancel_scope.cancel()
         except ExceptionGroup as e:
-            for exception in e.exceptions:
+            for exception in e.exceptions:  # type: ignore[attr-defined]
                 raise exception from None
 
 
@@ -67,6 +70,8 @@ def application_factory(settings: Http2KafkaApplicationSettings) -> FastAPI:
     application.state.settings = settings
     application.include_router(openlineage_router)
 
+    PersonalTokenAuthProvider.setup(application)
+
     # Reusing Server source code
     apply_exception_handlers(application)
     apply_middlewares(application, settings.server)
@@ -86,6 +91,7 @@ async def get_publisher():
         {
             Http2KafkaApplicationSettings: get_settings,
             AsyncAPIDefaultPublisher: get_publisher,
+            AsyncSession: session_generator(settings.database),  # type: ignore[dict-item]
         },
     )
     return application

data_rentgen/http2kafka/router/__init__.py

@@ -8,19 +8,22 @@
 from fastapi import APIRouter, Body, Depends, Request, Response
 from faststream.kafka.publisher.asyncapi import AsyncAPIDefaultPublisher
 
+from data_rentgen.db.models.user import User
 from data_rentgen.dependencies.stub import Stub
 from data_rentgen.http2kafka.router.gzip_route import SupportsGzipRoute
 from data_rentgen.openlineage.run_event import OpenLineageRunEvent
 from data_rentgen.openlineage.run_facets import OpenLineageParentRunFacet
 from data_rentgen.server.errors import get_error_responses
 from data_rentgen.server.errors.schemas import InvalidRequestSchema
+from data_rentgen.server.errors.schemas.not_authorized import NotAuthorizedSchema
+from data_rentgen.server.services.get_user import PersonalTokenPolicy, get_user
 
 logger = logging.getLogger(__name__)
 
 router = APIRouter(
     prefix="/v1/openlineage",
     tags=["OpenLineage"],
-    responses=get_error_responses(include={InvalidRequestSchema}),
+    responses=get_error_responses(include={InvalidRequestSchema, NotAuthorizedSchema}),
     route_class=SupportsGzipRoute,
 )
 
@@ -35,6 +38,7 @@ async def send_events_to_kafka(
     event: Annotated[OpenLineageRunEvent, Body()],
     request: Request,
     kafka_publisher: Annotated[AsyncAPIDefaultPublisher, Depends(Stub(AsyncAPIDefaultPublisher))],
+    current_user: Annotated[User, Depends(get_user(personal_token_policy=PersonalTokenPolicy.REQUIRE))],
 ):
     body_json_bytes = await request.body()
     logger.debug("Got 1 message (%dKiB)", len(body_json_bytes) / 1024)
@@ -45,6 +49,7 @@ async def send_events_to_kafka(
         correlation_id=correlation_id.get(),
         headers={
             "content-type": "application/json",
+            "reported-by-user-name": current_user.name,
         },
         # add event to message accumulator. messages are send in a background task.
         # do not wait until the message is actually send

data_rentgen/http2kafka/settings.py

@@ -6,7 +6,9 @@
 
 from data_rentgen.consumer.settings.kafka import KafkaSettings
 from data_rentgen.consumer.settings.producer import ProducerSettings
+from data_rentgen.db.settings import DatabaseSettings
 from data_rentgen.logging.settings import LoggingSettings
+from data_rentgen.server.settings.auth import AuthSettings
 from data_rentgen.server.settings.server import ServerSettings
 
 
@@ -41,6 +43,14 @@ class Http2KafkaApplicationSettings(BaseSettings):
         DATA_RENTGEN__PRODUCER__MAIN_TOPIC="input.runs"
     """  # noqa: E501
 
+    auth: AuthSettings = Field(
+        default_factory=AuthSettings,
+        description=":ref:`Authentication settings <configuration-server-authentication>`",
+    )
+    database: DatabaseSettings = Field(
+        default_factory=DatabaseSettings,  # type: ignore[arg-type]
+        description=":ref:`Database settings <configuration-database>`",
+    )
     logging: LoggingSettings = Field(
         default_factory=LoggingSettings,
         description=":ref:`Logging settings <configuration-server-logging>`",

data_rentgen/server/__init__.py

@@ -12,6 +12,7 @@
 from data_rentgen.server.api.handlers import apply_exception_handlers
 from data_rentgen.server.api.router import api_router
 from data_rentgen.server.middlewares import apply_middlewares
+from data_rentgen.server.providers.auth.personal_token_provider import PersonalTokenAuthProvider
 from data_rentgen.server.settings import ServerApplicationSettings
 
 if TYPE_CHECKING:
@@ -36,11 +37,16 @@ def application_factory(settings: ServerApplicationSettings) -> FastAPI:
     apply_exception_handlers(application)
     auth_class: type[AuthProvider] = settings.auth.provider  # type: ignore[assignment]
     auth_class.setup(application)
+
+    PersonalTokenAuthProvider.setup(application)
     apply_middlewares(application, settings.server)
 
+    async def get_settings():
+        return settings
+
     application.dependency_overrides.update(
         {
-            ServerApplicationSettings: lambda: settings,
+            ServerApplicationSettings: get_settings,
             AsyncSession: session_generator(settings.database),  # type: ignore[dict-item]
         },
     )

data_rentgen/server/api/handlers.py

@@ -95,15 +95,14 @@ def application_exception_handler(request: Request, exc: ApplicationError) -> Re
 
 
 def not_authorized_redirect_exception_handler(request: Request, exc: RedirectError) -> Response:
-    logger.debug("Redirect user to keycloak")
     response = get_response_for_exception(RedirectError)
     if not response:
         return unknown_exception_handler(request, exc)
+
     content = response.schema(  # type: ignore[call-arg]
         message=exc.message,
         details=exc.details,
     )
-
     return exception_json_response(
         status=response.status,
         content=content,

data_rentgen/server/api/v1/router/auth.py

@@ -22,7 +22,7 @@
     KeycloakAuthProvider,
 )
 from data_rentgen.server.schemas.v1.auth import AuthTokenSchema
-from data_rentgen.server.services import get_user
+from data_rentgen.server.services import PersonalTokenPolicy, get_user
 
 router = APIRouter(
     prefix="/auth",
@@ -35,8 +35,10 @@
     summary="Get auth token",
     responses=get_error_responses(
         include={
-            InvalidRequestSchema,
+            NotAuthorizedSchema,
             NotAuthorizedRedirectSchema,
+            InvalidRequestSchema,
+            NotImplementedErrorSchema,
         },
     ),
 )
@@ -56,6 +58,7 @@ async def token(
     summary="Handle redirect callback from OAuth2 provider",
     responses=get_error_responses(
         include={
+            NotAuthorizedSchema,
             InvalidRequestSchema,
             NotImplementedErrorSchema,
         },
@@ -90,7 +93,7 @@ async def auth_callback(
 )
 async def logout(
     request: Request,
-    current_user: Annotated[User, Depends(get_user())],
+    current_user: Annotated[User, Depends(get_user(personal_token_policy=PersonalTokenPolicy.DENY))],
     auth_provider: Annotated[KeycloakAuthProvider, Depends(Stub(AuthProvider))],
 ):
     refresh_token = request.session.get("refresh_token", None)

data_rentgen/server/api/v1/router/personal_token.py

@@ -13,6 +13,7 @@
 
 from data_rentgen.db.models import User  # noqa: TC001
 from data_rentgen.server.errors import get_error_responses
+from data_rentgen.server.providers.auth.personal_token_provider import PersonalTokenAuthProvider  # noqa: TC001
 from data_rentgen.server.schemas.v1 import (
     PageResponseV1,
     PersonalTokenCreatedDetailedResponseV1,
@@ -23,7 +24,7 @@
     PersonalTokenResponseV1,
     PersonalTokenScopeV1,
 )
-from data_rentgen.server.services import PersonalTokenService, get_user
+from data_rentgen.server.services import PersonalTokenPolicy, PersonalTokenService, get_user
 
 router = APIRouter(
     prefix="/personal-tokens",
@@ -50,8 +51,9 @@ async def get_personal_tokens(
 @router.post("")
 async def create_personal_token(
     token_params: PersonalTokenCreateRequestV1,
-    current_user: Annotated[User, Depends(get_user())],
+    current_user: Annotated[User, Depends(get_user(personal_token_policy=PersonalTokenPolicy.DENY))],
     user_token_service: Annotated[PersonalTokenService, Depends()],
+    personal_token_auth_provider: Annotated[PersonalTokenAuthProvider, Depends()],
 ) -> PersonalTokenCreatedDetailedResponseV1:
     async with user_token_service:
         token = await user_token_service.create(
@@ -72,16 +74,17 @@ async def create_personal_token(
             since=token.since,
             until=token.until,
         ),
-        content="TODO",
+        content=personal_token_auth_provider.generate_jwt(user=current_user, token=token),
     )
 
 
 @router.patch("/{token_id}")
 async def reset_personal_token(
     token_id: UUID,
     new_token_params: PersonalTokenResetRequestV1,
-    current_user: Annotated[User, Depends(get_user())],
+    current_user: Annotated[User, Depends(get_user(personal_token_policy=PersonalTokenPolicy.DENY))],
     user_token_service: Annotated[PersonalTokenService, Depends()],
+    personal_token_auth_provider: Annotated[PersonalTokenAuthProvider, Depends()],
 ) -> PersonalTokenCreatedDetailedResponseV1:
     async with user_token_service:
         old_token = await user_token_service.revoke(current_user, token_id)
@@ -102,14 +105,14 @@ async def reset_personal_token(
             since=new_token.since,
             until=new_token.until,
         ),
-        content="TODO",
+        content=personal_token_auth_provider.generate_jwt(user=current_user, token=new_token),
     )
 
 
 @router.delete("/{token_id}", status_code=HTTPStatus.NO_CONTENT)
 async def revoke_personal_token(
     token_id: UUID,
-    current_user: Annotated[User, Depends(get_user())],
+    current_user: Annotated[User, Depends(get_user(personal_token_policy=PersonalTokenPolicy.DENY))],
     user_token_service: Annotated[PersonalTokenService, Depends()],
 ):
     async with user_token_service:

data_rentgen/server/errors/registration.py

@@ -20,11 +20,18 @@ class APIErrorResponse(NamedTuple):
 _responses_by_status_code: dict[int, APIErrorResponse] = {}
 
 
-def register_error_response(exception: type[Exception], status: http.HTTPStatus):
+def register_error_response(
+    exception: type[Exception],
+    status: http.HTTPStatus,
+):
     """Register mapping between exception, status code and JSON body schema."""
 
     def wrapper(cls):
-        response = APIErrorResponse(status.value, status.phrase, cls)
+        response = APIErrorResponse(
+            status=status.value,
+            description=status.phrase,
+            schema=cls,
+        )
         _responses_by_exception[exception] = response
         _responses_by_status_code[status.value] = response
         return cls