Feature/windows notifications (#7)
Reviewed-on: #7 Co-authored-by: Florian <flo@meergans.org>
This commit was merged in pull request #7.
This commit is contained in:
1 parent
cd26070a3b
commit
2b69f3146e
6 files changed
+241
-74
No files matched your search
+92
-34
@@ -1,29 +1,39 @@
|
||||
from fastapi import FastAPI, Depends, HTTPException, Request, Response
|
||||
import asyncio
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Dict
|
||||
|
||||
import httpx
|
||||
import uvicorn
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.security.api_key import APIKeyHeader
|
||||
from pydantic import BaseModel, Field
|
||||
from starlette.exceptions import HTTPException as StarletteHTTPException
|
||||
from typing import Dict
|
||||
from pydantic import BaseModel
|
||||
from validator import verify_api_key
|
||||
from db_module import get_db, create_connection_pool, close_connection_pool, start_healthcheck_thread
|
||||
from simple_logger_handler import setup_logger, LOG_LEVEL
|
||||
from rabbitmq_handler import RabbitMQProducer
|
||||
import uvicorn
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from db_module import close_connection_pool, create_connection_pool, get_db, start_healthcheck_thread
|
||||
from metrics_server import REQUEST_COUNTER
|
||||
import asyncio
|
||||
from uvicorn_logger_config import LOGGING_CONFIG
|
||||
from rabbitmq_handler import RabbitMQProducer
|
||||
from secret_manager import cleanup_secret_files
|
||||
from simple_logger_handler import LOG_LEVEL, setup_logger
|
||||
from uvicorn_logger_config import LOGGING_CONFIG
|
||||
from validator import verify_api_key
|
||||
|
||||
logger = setup_logger(__name__)
|
||||
producer = RabbitMQProducer()
|
||||
api_key_header_internal = APIKeyHeader(name="X-API-Key-Internal")
|
||||
MAX_PAGES = 5
|
||||
|
||||
class Notification(BaseModel):
|
||||
receipent_user_id : int
|
||||
receipent_user_id : int = Field(gt=0) #ToDo fix among all services
|
||||
message : Dict
|
||||
|
||||
class PullResponse(BaseModel):
|
||||
messages: list[dict]
|
||||
count: int
|
||||
has_more: bool
|
||||
|
||||
class PullRequest(BaseModel):
|
||||
user_id: int = Field(gt=0)
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
@@ -34,25 +44,42 @@ async def lifespan(app: FastAPI):
|
||||
start_healthcheck_thread()
|
||||
logger.info("MySQL healthcheck thread started.")
|
||||
|
||||
logger.info("Starting RabbitMQ producer...")
|
||||
await producer.connect()
|
||||
app.state.rmq_producer = producer
|
||||
logger.info("[FastAPI] RabbitMQ producer initialized.")
|
||||
producer = RabbitMQProducer()
|
||||
PULL_SERVICE_URL = os.getenv(
|
||||
"BACKEND_PULL_NOTIFICATIONS_URL",
|
||||
"http://backend-pull-notifications.app-notifications.svc.cluster.local:8101",
|
||||
)
|
||||
|
||||
SECRET_PATHS = frozenset({
|
||||
"/etc/secrets/db_username",
|
||||
"/etc/secrets/db_password",
|
||||
"/etc/secrets/rmq_username",
|
||||
"/etc/secrets/rmq_password"
|
||||
})
|
||||
cleanup_secret_files(SECRET_PATHS)
|
||||
try:
|
||||
logger.info("Starting RabbitMQ producer...")
|
||||
await producer.connect()
|
||||
app.state.rmq_producer = producer
|
||||
logger.info("[FastAPI] RabbitMQ producer initialized.")
|
||||
|
||||
yield
|
||||
logger.info("Closing RabbitMQ producer...")
|
||||
await producer.close()
|
||||
app.state.http_client = httpx.AsyncClient (
|
||||
base_url=PULL_SERVICE_URL,
|
||||
timeout=httpx.Timeout(10.0,connect=3.0)
|
||||
)
|
||||
logger.info("[FastAPI] HTTP client initialized.")
|
||||
|
||||
logger.info("Closing MySQL connection pool...")
|
||||
close_connection_pool()
|
||||
SECRET_PATHS = frozenset({
|
||||
"/etc/secrets/db_username",
|
||||
"/etc/secrets/db_password",
|
||||
"/etc/secrets/rmq_username",
|
||||
"/etc/secrets/rmq_password"
|
||||
})
|
||||
cleanup_secret_files(SECRET_PATHS)
|
||||
|
||||
yield
|
||||
finally:
|
||||
logger.info("Closing RabbitMQ producer...")
|
||||
await producer.close()
|
||||
|
||||
logger.info("Closing MySQL connection pool...")
|
||||
close_connection_pool()
|
||||
|
||||
await app.state.http_client.aclose()
|
||||
logger.info("[FastAPI] Closing Http client")
|
||||
|
||||
api = FastAPI(
|
||||
title="Internal Notifier API",
|
||||
@@ -78,10 +105,10 @@ async def prometheus_middleware(request: Request, call_next):
|
||||
return response
|
||||
|
||||
def verify_api_key_dependency_internal(db=Depends(get_db), api_key: str = Depends(api_key_header_internal)) -> str:
|
||||
logger.debug(f"Verifying internal API key: {api_key}")
|
||||
cursor = db.cursor()
|
||||
cursor.execute("SELECT program_name, api_key FROM internal_api_keys WHERE status = 'active'")
|
||||
for program_name, hashed_key in cursor.fetchall():
|
||||
with db.cursor() as cursor:
|
||||
cursor.execute("SELECT program_name, api_key FROM internal_api_keys WHERE status = 'active'")
|
||||
rows = cursor.fetchall()
|
||||
for program_name, hashed_key in rows:
|
||||
if verify_api_key(api_key=api_key, hashed=hashed_key):
|
||||
logger.debug(f"API key verified for program: {program_name}")
|
||||
return program_name
|
||||
@@ -101,7 +128,7 @@ async def custom_http_exception_handler(request,exc):
|
||||
content={"detail": exc.detail}
|
||||
)
|
||||
|
||||
@api.post("/internal/receive-notifications")
|
||||
@api.post("/internal/receive-notifications") #ToDo fix among all services
|
||||
async def receive_notifications(
|
||||
notification_data: Notification,
|
||||
request: Request,
|
||||
@@ -118,6 +145,37 @@ async def receive_notifications(
|
||||
logger.info("Successfully queued for delivery to RMQ")
|
||||
return {"status": "queued"}
|
||||
|
||||
@api.post("/internal/pull-notifications", response_model=PullResponse)
|
||||
async def pull_notifications(
|
||||
pull_request: PullRequest,
|
||||
request: Request,
|
||||
program_name: str = Depends(verify_api_key_dependency_internal),
|
||||
):
|
||||
user_id = pull_request.user_id
|
||||
logger.debug(f"[FastAPI] User id {user_id} is requesting messages")
|
||||
client: httpx.AsyncClient = request.app.state.http_client
|
||||
|
||||
messages: list[dict] = []
|
||||
has_more = False
|
||||
|
||||
try:
|
||||
for _ in range(MAX_PAGES):
|
||||
resp = await client.post("/internal/pull-notifications", json={"user_id": user_id})
|
||||
resp.raise_for_status()
|
||||
data = PullResponse.model_validate(resp.json())
|
||||
|
||||
messages.extend(data.messages)
|
||||
has_more = data.has_more
|
||||
if not has_more:
|
||||
break
|
||||
except (httpx.HTTPError, ValueError) as e:
|
||||
logger.error(f"[FastAPI] Pull service failed for user_id={user_id}: {e!r}")
|
||||
status = 504 if isinstance(e, httpx.TimeoutException) else 502
|
||||
raise HTTPException(status_code=status, detail="Notification service unavailable") from e
|
||||
|
||||
logger.info(f"[FastAPI] Returning {len(messages)} message(s) for user_id={user_id}, has_more={has_more}")
|
||||
return PullResponse(messages=messages, count=len(messages), has_more=has_more)
|
||||
|
||||
async def start_servers():
|
||||
logger.debug("Starting FastAPI and metrics servers")
|
||||
config_main = uvicorn.Config("main:api", host="0.0.0.0", port=8101, log_level=LOG_LEVEL.lower(), log_config=LOGGING_CONFIG)
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from fastapi import FastAPI, Response
|
||||
from prometheus_client import generate_latest, CONTENT_TYPE_LATEST, Counter
|
||||
from prometheus_client import CONTENT_TYPE_LATEST, Counter, generate_latest
|
||||
|
||||
metrics_api = FastAPI(title="Metrics Server", description="Prometheus metrics endpoint")
|
||||
|
||||
|
||||
+116
-39
@@ -1,72 +1,151 @@
|
||||
import asyncio
|
||||
import aio_pika
|
||||
from aio_pika.exceptions import AMQPException
|
||||
from secret_manager import return_credentials
|
||||
import os
|
||||
from simple_logger_handler import setup_logger
|
||||
import json
|
||||
import os
|
||||
from enum import IntEnum
|
||||
from typing import NamedTuple
|
||||
from urllib.parse import quote
|
||||
|
||||
import aio_pika
|
||||
from aio_pika.exceptions import AMQPException, DeliveryError
|
||||
|
||||
from secret_manager import return_credentials
|
||||
from simple_logger_handler import setup_logger
|
||||
|
||||
logger = setup_logger(__name__)
|
||||
MAX_UNROUTABLE_ATTEMPTS = 3
|
||||
DRAIN_TIMEOUT_SECONDS = 10 # how long close() waits for buffered messages
|
||||
|
||||
rmq_username = return_credentials("/etc/secrets/rmq_username")
|
||||
rmq_password = return_credentials("/etc/secrets/rmq_password")
|
||||
rmq_host = os.getenv("BACKEND_API_INTERNAL_RMQ_HOST", "localhost")
|
||||
rmq_vhost = os.getenv("BACKEND_API_INTERNAL_RMQ_VHOST", "app_notifications")
|
||||
rmq_exchange = os.getenv("BACKEND_API_INTERNAL_RMQ_EXCHANGE", "app_notifications")
|
||||
|
||||
RABBITMQ_URL = f"amqp://{rmq_username}:{rmq_password}@{rmq_host}/{rmq_vhost}"
|
||||
class Kind(IntEnum):
|
||||
PUSH = 0
|
||||
PULL = 1
|
||||
|
||||
class OutgoingMessage(NamedTuple):
|
||||
user_id: int
|
||||
routing_key: str
|
||||
message: dict
|
||||
attempts: int
|
||||
kind: Kind
|
||||
|
||||
class RabbitMQProducer:
|
||||
def __init__(self, url=RABBITMQ_URL, exchange_name=rmq_exchange,rmq_host=rmq_host):
|
||||
self.url = url
|
||||
self.rmq_host = rmq_host
|
||||
self.exchange_name = exchange_name
|
||||
"""
|
||||
Buffered RabbitMQ producer. publish() only enqueues in memory; a background
|
||||
task sends each message to the exchange, once with the push key and once with
|
||||
the pull key per user.
|
||||
"""
|
||||
def __init__(self):
|
||||
self.rmq_host = os.getenv("RMQ_HOST", "localhost")
|
||||
self.rmq_vhost = os.getenv("RMQ_VHOST", "app_notifications")
|
||||
self.rmq_exchange = os.getenv("RMQ_EXCHANGE", "app_notifications")
|
||||
self.pull_routing_prefix = os.getenv("PULL_NOTIFICATIONS_RMQ_ROUTING_PREFIX", "notify.pull")
|
||||
self.push_routing_prefix = os.getenv("PUSH_NOTIFICATIONS_RMQ_ROUTING_PREFIX", "notify.user")
|
||||
self.pull_queue_prefix = os.getenv("PULL_NOTIFICATIONS_RMQ_QUEUE_PREFIX", "notifications-pull")
|
||||
self.push_queue_name = os.getenv("PUSH_NOTIFICATIONS_RMQ_QUEUE", "notifications")
|
||||
|
||||
self.connection: aio_pika.RobustConnection | None = None
|
||||
self.channel: aio_pika.RobustChannel | None = None
|
||||
self.exchange: aio_pika.Exchange | None = None
|
||||
self._queue: asyncio.Queue[tuple[int, dict]] = asyncio.Queue()
|
||||
self._queue: asyncio.Queue[OutgoingMessage] = asyncio.Queue()
|
||||
self._flush_task: asyncio.Task | None = None
|
||||
self._closing = False
|
||||
self._ready = asyncio.Event()
|
||||
|
||||
async def connect(self):
|
||||
username = return_credentials("/etc/secrets/rmq_username")
|
||||
password = return_credentials("/etc/secrets/rmq_password")
|
||||
url = (
|
||||
f"amqp://{quote(username, safe='')}:{quote(password, safe='')}"
|
||||
f"@{self.rmq_host}/{quote(self.rmq_vhost, safe='')}"
|
||||
)
|
||||
|
||||
logger.info(f"[RabbitMQ] Connecting to {self.rmq_host}...")
|
||||
self.connection = await aio_pika.connect_robust(self.url)
|
||||
self.connection = await aio_pika.connect_robust(url)
|
||||
self.channel = await self.connection.channel(publisher_confirms=True)
|
||||
self.exchange = await self.channel.declare_exchange(
|
||||
self.exchange_name, aio_pika.ExchangeType.TOPIC, durable=True
|
||||
self.rmq_exchange, aio_pika.ExchangeType.TOPIC, durable=True
|
||||
)
|
||||
logger.info(f"[RabbitMQ] Connected and exchange '{self.exchange_name}' ready.")
|
||||
logger.info(f"[RabbitMQ] Connected and exchange '{self.rmq_exchange}' ready.")
|
||||
push_queue = await self.channel.declare_queue(self.push_queue_name, durable=True)
|
||||
await push_queue.bind(self.exchange, routing_key=f"{self.push_routing_prefix}.*")
|
||||
self._ready.set()
|
||||
self._flush_task = asyncio.create_task(self._flush_queue_loop())
|
||||
logger.debug("[RabbitMQ] Flush queue loop started")
|
||||
|
||||
async def publish(self, routing_key: int, message: dict):
|
||||
await self._queue.put((routing_key, message))
|
||||
logger.debug(f"[RabbitMQ] Queued message for {routing_key}")
|
||||
async def ensure_user_queue(self, user_id: int) -> aio_pika.Queue:
|
||||
"""
|
||||
Declare the user's durable queue (with expiry and message TTL) and bind it to
|
||||
the exchange with the user's routing key. Safe to call repeatedly, creates it if it expired, and resets its
|
||||
idle-expiry timer.
|
||||
"""
|
||||
logger.debug("[RabbitMQ] Declaring user pull queue for user:{user_id}")
|
||||
queue = await self.channel.declare_queue(
|
||||
f"{self.pull_queue_prefix}.{user_id}",
|
||||
durable=True,
|
||||
arguments={
|
||||
"x-expires": 30 * 24 * 3600 * 1000,
|
||||
"x-message-ttl": 7 * 24 * 3600 * 1000,
|
||||
},
|
||||
)
|
||||
await queue.bind(self.exchange, routing_key=f"{self.pull_routing_prefix}.{user_id}")
|
||||
return queue
|
||||
|
||||
async def publish(self, user_id: str | int, message: dict):
|
||||
"""
|
||||
Buffer a message for a user under two routing keys (push and pull). Works while RabbitMQ is down.
|
||||
Raises:
|
||||
- RuntimeError if the producer is closing
|
||||
"""
|
||||
if self._closing:
|
||||
raise RuntimeError("[RabbitMQ] Producer is closing")
|
||||
user_id = int(user_id)
|
||||
await self._queue.put(OutgoingMessage(user_id, f"{self.push_routing_prefix}.{user_id}", message, 0, Kind.PUSH))
|
||||
await self._queue.put(OutgoingMessage(user_id, f"{self.pull_routing_prefix}.{user_id}", message, 0, Kind.PULL))
|
||||
logger.debug(f"[RabbitMQ] Queued message for {user_id}")
|
||||
|
||||
async def _flush_queue_loop(self):
|
||||
"""Background task: wait until connected, then take buffered messages one at a
|
||||
time and publish them as persistent JSON.
|
||||
|
||||
- Declares the user's pull queue first when the key is a pull key.
|
||||
- Returned/unroutable messages (DeliveryError) are retried up to
|
||||
MAX_UNROUTABLE_ATTEMPTS times, then dropped with an error log.
|
||||
- Connection errors are retried without limit after a 2 second pause.
|
||||
- Any other error (e.g. non-serializable payload) drops the message so the
|
||||
loop keeps running.
|
||||
"""
|
||||
logger.debug("[RabbitMQ] Waiting for ready signal before flushing queue")
|
||||
await self._ready.wait()
|
||||
while True:
|
||||
routing_key, message = await self._queue.get()
|
||||
logger.debug(f"[RabbitMQ] Publishing message to routing key {routing_key}: {message}")
|
||||
outgoing_message = await self._queue.get()
|
||||
logger.debug(f"[RabbitMQ] Publishing message to routing key {outgoing_message.routing_key}: {outgoing_message.message}")
|
||||
try:
|
||||
if outgoing_message.kind == Kind.PULL:
|
||||
await self.ensure_user_queue(outgoing_message.user_id)
|
||||
await self.exchange.publish(
|
||||
aio_pika.Message(
|
||||
body=json.dumps(message).encode(),
|
||||
body=json.dumps(outgoing_message.message).encode(),
|
||||
content_type="application/json",
|
||||
delivery_mode=aio_pika.DeliveryMode.PERSISTENT
|
||||
),
|
||||
routing_key=f"notify.user.{routing_key}",
|
||||
routing_key=outgoing_message.routing_key,
|
||||
mandatory=True
|
||||
)
|
||||
logger.debug(f"[RabbitMQ] Successfully published to notify.user.{routing_key}")
|
||||
logger.debug(f"[RabbitMQ] Successfully published to {outgoing_message.routing_key}")
|
||||
|
||||
except DeliveryError as e:
|
||||
attempts = outgoing_message.attempts + 1
|
||||
if attempts >= MAX_UNROUTABLE_ATTEMPTS:
|
||||
logger.error(f"[RabbitMQ] Dropping message for {outgoing_message.routing_key} after {attempts} attempts: {e}")
|
||||
else:
|
||||
await asyncio.sleep(2)
|
||||
await self._queue.put(outgoing_message._replace(attempts=attempts))
|
||||
|
||||
except (AMQPException, OSError) as e:
|
||||
logger.warning(f"[RabbitMQ] Publish failed, requeuing message: {e}", exc_info=True)
|
||||
logger.warning(f"[RabbitMQ] Failed to publish to {outgoing_message.routing_key}: {e}. Requeuing message.", exc_info=True)
|
||||
await asyncio.sleep(2)
|
||||
await self._queue.put((routing_key, message))
|
||||
await self._queue.put(outgoing_message)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[RabbitMQ] Dropping unpublishable message for {outgoing_message.routing_key}: {e}", exc_info=True)
|
||||
finally:
|
||||
self._queue.task_done()
|
||||
|
||||
@@ -74,20 +153,18 @@ class RabbitMQProducer:
|
||||
logger.info("[RabbitMQ] Closing producer...")
|
||||
self._closing = True
|
||||
if self._flush_task and not self._flush_task.done():
|
||||
try:
|
||||
await asyncio.wait_for(self._queue.join(), timeout=DRAIN_TIMEOUT_SECONDS)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(f"[RabbitMQ] {self._queue.qsize()} message(s) still buffered, dropping on shutdown")
|
||||
|
||||
self._flush_task.cancel()
|
||||
try:
|
||||
await self._flush_task
|
||||
logger.debug("[RabbitMQ] Flush task cancelled successfully")
|
||||
except asyncio.CancelledError:
|
||||
logger.debug("[RabbitMQ] Flush task cancellation caught")
|
||||
logger.debug("[RabbitMQ] Flush task cancelled")
|
||||
|
||||
if self.connection:
|
||||
await self.connection.close()
|
||||
logger.info("[RabbitMQ] Connection closed")
|
||||
|
||||
|
||||
async def main():
|
||||
producer = RabbitMQProducer()
|
||||
await producer.connect()
|
||||
await producer.publish("1", {"type": "notification", "content": "Test message"})
|
||||
await producer.close()
|
||||
|
||||
Reference in new issue
Block a user