Merge branch 'main' into Greylist-check#61

pyproject.toml

@@ -4,18 +4,18 @@ description = "Distributes tasks to dragonfly-clients and handles package status
 authors = [
     { name = "jz9", email = "[email protected]" },
 ]
-requires-python = ">=3.11"
+requires-python = ">=3.11.4"
 license = { text = "MIT" }
 
 dependencies = [
     "fastapi>=0.95.1",
-    "pydantic==1.10.11",
+    "pydantic-settings>=2.0.0",
     "uvicorn[standard]>=0.22.0",
     "SQLAlchemy>=2.0.12",
     "python-dotenv>=1.0.0",
     "alembic>=1.10.4",
     "asyncpg>=0.27.0",
-    "letsbuilda-pypi>=4.0.0",
+    "letsbuilda-pypi[async]==5.0.0b2",
     "aiohttp>=3.8.4",
     "msgraph-core>=0.2.2",
     "azure-identity>=1.13.0",

src/mainframe/constants.py

@@ -1,15 +1,23 @@
 from os import getenv
 
-from pydantic import BaseSettings
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
 # Git SHA for Sentry
 GIT_SHA = getenv("GIT_SHA", "development")
 
 
-class Mainframe(BaseSettings):
-    class Config(BaseSettings.Config):
-        env_file = ".env"
+class EnvConfig(BaseSettings):
+    """Our default configuration for models that should load from .env files."""
 
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        env_nested_delimiter="__",
+        extra="ignore",
+    )
+
+
+class Mainframe(EnvConfig):
     client_origin_url: str = ""
     auth0_domain: str = ""
     auth0_audience: str = ""
@@ -27,10 +35,8 @@ class Config(BaseSettings.Config):
 mainframe_settings = Mainframe()  # pyright: ignore
 
 
-class _Sentry(BaseSettings):
-    class Config(BaseSettings.Config):
-        env_prefix = "sentry_"
-        env_file = ".env"
+class _Sentry(EnvConfig):
+    EnvConfig.model_config["env_prefix"] = "sentry_"
 
     dsn: str = ""
     environment: str = ""
@@ -40,10 +46,8 @@ class Config(BaseSettings.Config):
 Sentry = _Sentry()  # pyright: ignore
 
 
-class Microsoft(BaseSettings):
-    class Config(BaseSettings.Config):
-        env_prefix = "microsoft_"
-        env_file = ".env"
+class Microsoft(EnvConfig):
+    EnvConfig.model_config["env_prefix"] = "microsoft_"
 
     tenant_id: str
     client_id: str

src/mainframe/endpoints/package.py

@@ -3,7 +3,8 @@
 
 import structlog
 from fastapi import APIRouter, Depends, HTTPException, Request
-from letsbuilda.pypi import PyPIServices  # type: ignore
+from letsbuilda.pypi.async_client import PyPIServices  # type: ignore
+from letsbuilda.pypi.exceptions import PackageNotFoundError
 from sqlalchemy import select
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -187,15 +188,17 @@ async def batch_queue_package(
 
         try:
             package_metadata = await pypi_client.get_package_metadata(name, version)
-        except KeyError:
+        except PackageNotFoundError:
             continue
 
         scan = Scan(
-            name=package_metadata.info.name,
-            version=package_metadata.info.version,
+            name=package_metadata.title,
+            version=package_metadata.releases[0].version,
             status=Status.QUEUED,
             queued_by=auth.subject,
-            download_urls=[DownloadURL(url=url.url) for url in package_metadata.urls],
+            download_urls=[
+                DownloadURL(url=distribution.url) for distribution in package_metadata.releases[0].distributions
+            ],
         )
 
         rows.append(scan)
@@ -241,22 +244,24 @@ async def queue_package(
     pypi_client: PyPIServices = request.app.state.pypi_client
     try:
         package_metadata = await pypi_client.get_package_metadata(name, version)
-    except KeyError:
+    except PackageNotFoundError:
         error = HTTPException(404, detail=f"Package {name}@{version} was not found on PyPI")
         await log.aerror(
             f"Package {name}@{version} was not found on PyPI", error_message=error.detail, tag="package_not_found_pypi"
         )
         raise error
 
-    version = package_metadata.info.version  # Use latest version if not provided
+    version = package_metadata.releases[0].version  # Use latest version if not provided
     log = logger.bind(package={"name": name, "version": version})
 
     new_package = Scan(
         name=name,
         version=version,
         status=Status.QUEUED,
         queued_by=auth.subject,
-        download_urls=[DownloadURL(url=url.url) for url in package_metadata.urls],
+        download_urls=[
+            DownloadURL(url=distribution.url) for distribution in package_metadata.releases[0].distributions
+        ],
     )
 
     session.add(new_package)

src/mainframe/endpoints/report.py

@@ -4,7 +4,8 @@
 
 import structlog
 from fastapi import APIRouter, Depends, HTTPException, Request
-from letsbuilda.pypi import PyPIServices
+from letsbuilda.pypi.async_client import PyPIServices
+from letsbuilda.pypi.exceptions import PackageNotFoundError
 from msgraph.core import GraphClient
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -123,14 +124,14 @@ async def report_package(
     pypi_client: PyPIServices = request.app.state.pypi_client
     try:
         package_metadata = await pypi_client.get_package_metadata(name, version)
-    except KeyError:
+    except PackageNotFoundError:
         error = HTTPException(404, detail=f"Package `{name}@{version}` was not found on PyPI")
         await log.aerror(
             f"Package {name}@{version} was not found on PyPI", error_message=error.detail, tag="package_not_found_pypi"
         )
         raise error
 
-    version = package_metadata.info.version
+    version = package_metadata.releases[0].version
     log = logger.bind(package={"name": name, "version": version})
 
     query = select(Scan).where(Scan.name == name).options(selectinload(Scan.rules))

src/mainframe/server.py

@@ -11,7 +11,7 @@
 from asgi_correlation_id import CorrelationIdMiddleware
 from asgi_correlation_id.context import correlation_id
 from fastapi import FastAPI, Request, Response
-from letsbuilda.pypi import PyPIServices
+from letsbuilda.pypi.async_client import PyPIServices
 
 from mainframe.constants import GIT_SHA, Sentry
 from mainframe.database import async_session