|
1 | 1 | """ Main module for the FastAPI application. """ |
2 | 2 |
|
3 | | -import asyncio |
4 | 3 | import os |
5 | | -import shutil |
6 | | -import time |
7 | | -from contextlib import asynccontextmanager |
8 | | -from pathlib import Path |
9 | 4 |
|
10 | 5 | from api_analytics.fastapi import Analytics |
11 | 6 | from dotenv import load_dotenv |
12 | 7 | from fastapi import FastAPI, Request |
13 | | -from fastapi.responses import FileResponse, HTMLResponse, Response |
| 8 | +from fastapi.responses import FileResponse, HTMLResponse |
14 | 9 | from fastapi.staticfiles import StaticFiles |
15 | | -from fastapi.templating import Jinja2Templates |
16 | | -from slowapi import _rate_limit_exceeded_handler |
17 | 10 | from slowapi.errors import RateLimitExceeded |
18 | 11 | from starlette.middleware.trustedhost import TrustedHostMiddleware |
19 | 12 |
|
|
24 | 17 | # Load environment variables from .env file |
25 | 18 | load_dotenv() |
26 | 19 |
|
27 | | - |
28 | | -async def remove_old_repositories(): |
29 | | - """ |
30 | | - Background task that runs periodically to clean up old repository directories. |
31 | | -
|
32 | | - This task: |
33 | | - - Scans the TMP_BASE_PATH directory every 60 seconds |
34 | | - - Removes directories older than DELETE_REPO_AFTER seconds |
35 | | - - Before deletion, logs repository URLs to history.txt if a matching .txt file exists |
36 | | - - Handles errors gracefully if deletion fails |
37 | | -
|
38 | | - The repository URL is extracted from the first .txt file in each directory, |
39 | | - assuming the filename format: "owner-repository.txt" |
40 | | - """ |
41 | | - while True: |
42 | | - try: |
43 | | - if not TMP_BASE_PATH.exists(): |
44 | | - await asyncio.sleep(60) |
45 | | - continue |
46 | | - |
47 | | - current_time = time.time() |
48 | | - |
49 | | - for folder in TMP_BASE_PATH.iterdir(): |
50 | | - if not folder.is_dir(): |
51 | | - continue |
52 | | - |
53 | | - # Skip if folder is not old enough |
54 | | - if current_time - folder.stat().st_ctime <= DELETE_REPO_AFTER: |
55 | | - continue |
56 | | - |
57 | | - await process_folder(folder) |
58 | | - |
59 | | - except Exception as e: |
60 | | - print(f"Error in remove_old_repositories: {e}") |
61 | | - |
62 | | - await asyncio.sleep(60) |
63 | | - |
64 | | - |
65 | | -async def process_folder(folder: Path) -> None: |
66 | | - """ |
67 | | - Process a single folder for deletion and logging. |
68 | | -
|
69 | | - Parameters |
70 | | - ---------- |
71 | | - folder : Path |
72 | | - The path to the folder to be processed. |
73 | | - """ |
74 | | - # Try to log repository URL before deletion |
75 | | - try: |
76 | | - txt_files = [f for f in folder.iterdir() if f.suffix == ".txt"] |
77 | | - |
78 | | - # Extract owner and repository name from the filename |
79 | | - if txt_files and "-" in (filename := txt_files[0].stem): |
80 | | - owner, repo = filename.split("-", 1) |
81 | | - repo_url = f"{owner}/{repo}" |
82 | | - with open("history.txt", mode="a", encoding="utf-8") as history: |
83 | | - history.write(f"{repo_url}\n") |
84 | | - |
85 | | - except Exception as e: |
86 | | - print(f"Error logging repository URL for {folder}: {e}") |
87 | | - |
88 | | - # Delete the folder |
89 | | - try: |
90 | | - shutil.rmtree(folder) |
91 | | - except Exception as e: |
92 | | - print(f"Error deleting {folder}: {e}") |
93 | | - |
94 | | - |
95 | | -@asynccontextmanager |
96 | | -async def lifespan(_: FastAPI): |
97 | | - """ |
98 | | - Lifecycle manager for the FastAPI application. |
99 | | - Handles startup and shutdown events. |
100 | | -
|
101 | | - Parameters |
102 | | - ---------- |
103 | | - _ : FastAPI |
104 | | - The FastAPI application instance (unused). |
105 | | -
|
106 | | - Yields |
107 | | - ------- |
108 | | - None |
109 | | - Yields control back to the FastAPI application while the background task runs. |
110 | | - """ |
111 | | - task = asyncio.create_task(remove_old_repositories()) |
112 | | - |
113 | | - yield |
114 | | - # Cancel the background task on shutdown |
115 | | - task.cancel() |
116 | | - try: |
117 | | - await task |
118 | | - except asyncio.CancelledError: |
119 | | - pass |
120 | | - |
121 | | - |
122 | 20 | # Initialize the FastAPI application with lifespan |
123 | 21 | app = FastAPI(lifespan=lifespan) |
124 | 22 | app.state.limiter = limiter |
125 | 23 |
|
126 | | - |
127 | | -async def rate_limit_exception_handler(request: Request, exc: Exception) -> Response: |
128 | | - """ |
129 | | - Custom exception handler for rate-limiting errors. |
130 | | -
|
131 | | - Parameters |
132 | | - ---------- |
133 | | - request : Request |
134 | | - The incoming HTTP request. |
135 | | - exc : Exception |
136 | | - The exception raised, expected to be RateLimitExceeded. |
137 | | -
|
138 | | - Returns |
139 | | - ------- |
140 | | - Response |
141 | | - A response indicating that the rate limit has been exceeded. |
142 | | -
|
143 | | - Raises |
144 | | - ------ |
145 | | - exc |
146 | | - If the exception is not a RateLimitExceeded error, it is re-raised. |
147 | | - """ |
148 | | - if isinstance(exc, RateLimitExceeded): |
149 | | - # Delegate to the default rate limit handler |
150 | | - return _rate_limit_exceeded_handler(request, exc) |
151 | | - # Re-raise other exceptions |
152 | | - raise exc |
153 | | - |
154 | | - |
155 | 24 | # Register the custom exception handler for rate limits |
156 | 25 | app.add_exception_handler(RateLimitExceeded, rate_limit_exception_handler) |
157 | 26 |
|
|
0 commit comments