Add PyPI Package Manager Support

core/config.py

@@ -12,6 +12,7 @@
 class PackageManager(Enum):
     CRATES = "crates"
     HOMEBREW = "homebrew"
+    PYPI = "pypi"
 
 
 TEST = env_vars("TEST", "false")
@@ -20,6 +21,7 @@ class PackageManager(Enum):
 SOURCES = {
     PackageManager.CRATES: "https://static.crates.io/db-dump.tar.gz",
     PackageManager.HOMEBREW: "https://github.com/Homebrew/homebrew-core/tree/master/Formula",  # noqa
+    PackageManager.PYPI: "https://pypi.org/simple/",  # Base URL for PyPI's Simple API
 }
 
 # The three configuration values URLTypes, DependencyTypes, and UserTypes will query the

core/db.py

@@ -128,15 +128,15 @@ def update_caches(
     ):
         if update_packages:
             self._update_cache(
-                self.package_cache, Package, "import_id", "id", items, "crate_id"
+                self.package_cache, Package, "import_id", "id", items, "import_id"
             )
         if update_users:
             self._update_cache(
                 self.user_cache, User, "import_id", "id", items, "owner_id"
             )
         if update_versions:
             self._update_cache(
-                self.version_cache, Version, "import_id", "id", items, "version_id"
+                self.version_cache, Version, "import_id", "id", items, "import_id"
             )
         if update_licenses:
             self._update_cache(
@@ -159,10 +159,17 @@ def insert_versions(self, version_generator: Iterable[dict[str, str]]):
             self._insert_batch(Version, versions)
 
     def _process_version(self, item: Dict[str, str]):
-        package_id = self.package_cache.get(item["crate_id"])
-        if not package_id:
-            self.logger.warn(f"package {item['crate_id']} not found")
-            return None
+        # FIXME: this is a hack, the import_id of a version shouldn't be the same as the package's import_id
+        # but the original logic here tries to use the package's import_id as the version's import_id
+        # this is a temporary fix (for the implementation of pypi)
+        # to fix this completely, we need to update the logic of the crates transformer too
+        if item["package_id"]:
+            package_id = item["package_id"]
+        else:
+            package_id = self.package_cache.get(item["import_id"])
+            if not package_id:
+                self.logger.warn(f"package {item['import_id']} not found")
+                return None
 
         license_id = self.license_cache.get(item["license"])
         if not license_id:
@@ -201,11 +208,30 @@ def insert_dependencies(self, dependency_generator: Iterable[dict[str, str]]):
             self._insert_batch(DependsOn, dependencies)
 
     def _process_depends_on(self, item: Dict[str, str]):
-        return DependsOn(
-            version_id=self.version_cache[item["version_id"]],
-            dependency_id=self.package_cache[item["crate_id"]],
-            semver_range=item["semver_range"],
-        ).to_dict()
+        version_id = self.version_cache.get(item["version_id"])
+
+        # in case the version cannot be found from the cache
+        if not version_id:
+            # we need to fetch from the database
+            version = self.select_version_by_import_id(item["version_id"])
+            if not version:
+                self.logger.warn(f"version {item['version_id']} not found")
+                return None
+            version_id = version.id
+            self.version_cache[item["version_id"]] = version_id
+
+        # Create base dependency object
+        depends_on = {
+            "version_id": self.version_cache[item["version_id"]],
+            "dependency_id": self.package_cache[item["import_id"]],
+            "semver_range": item["semver_range"]
+        }
+
+        # Add dependency_type_id if provided
+        if "dependency_type_id" in item:
+            depends_on["dependency_type_id"] = item["dependency_type_id"]
+
+        return DependsOn(**depends_on).to_dict()
 
     def insert_users(self, user_generator: Iterable[dict[str, str]], source_id: UUID):
         def process_user(item: Dict[str, str]):
@@ -245,13 +271,13 @@ def _process_user_package(self, item: Dict[str, str]):
             self.logger.warn(f"user {item['owner_id']} not found")
             return None
 
-        if item["crate_id"] not in self.package_cache:
-            self.logger.warn(f"package {item['crate_id']} not found")
+        if item["import_id"] not in self.package_cache:
+            self.logger.warn(f"package {item['import_id']} not found")
             return None
 
         return UserPackage(
             user_id=self.user_cache[item["owner_id"]],
-            package_id=self.package_cache[item["crate_id"]],
+            package_id=self.package_cache[item["import_id"]],
         ).to_dict()
 
     def insert_user_versions(
@@ -476,6 +502,18 @@ def select_version_by_import_id(self, import_id: str) -> Version | None:
             result = session.query(Version).filter_by(import_id=import_id).first()
             if result:
                 return result
+
+    def select_latest_version_by_import_id(self, import_id: str) -> Version | None:
+        with self.session() as session:
+            # First get the package
+            package = session.query(Package).filter_by(import_id=import_id).first()
+            if not package:
+                return None
+
+            # Then get the latest version for this package
+            result = session.query(Version).filter_by(package_id=package.id).order_by(Version.version.desc()).first()
+            if result:
+                return result
 
     def select_package_manager_name_by_id(self, id: UUID) -> str | None:
         with self.session() as session:

core/models/__init__.py

@@ -189,13 +189,18 @@ class DependsOn(Base):
     dependency_type: Mapped["DependsOnType"] = relationship()
 
     def to_dict(self):
-        return {
+        result = {
             "version_id": self.version_id,
             "dependency_id": self.dependency_id,
-            # "dependency_type_id": self.dependency_type_id,
             "semver_range": self.semver_range,
         }
 
+        # if dependency_type_id is provided, include it
+        if hasattr(self, 'dependency_type_id') and self.dependency_type_id is not None:
+            result["dependency_type_id"] = self.dependency_type_id
+
+        return result
+
 
 class DependsOnType(Base):
     __tablename__ = "depends_on_types"

docker-compose.yml

@@ -69,6 +69,26 @@ services:
       alembic:
         condition: service_completed_successfully
 
+  pypi:
+    build:
+      context: .
+      dockerfile: ./package_managers/pypi/Dockerfile
+    environment:
+      - CHAI_DATABASE_URL=postgresql://postgres:s3cr3t@db:5432/chai
+      - NO_CACHE=${NO_CACHE:-false}
+      - PYTHONPATH=/
+      - DEBUG=${DEBUG:-false}
+      - TEST=${TEST:-false}
+      - FETCH=${FETCH:-true}
+      - FREQUENCY=${FREQUENCY:-24}
+    volumes:
+      - ./data/pypi:/data/pypi
+    depends_on:
+      db:
+        condition: service_healthy
+      alembic:
+        condition: service_completed_successfully
+
   api:
     build:
       context: ./api

package_managers/crates/transformer.py

@@ -92,7 +92,7 @@ def dependencies(self) -> Generator[Dict[str, str], None, None]:
 
             yield {
                 "version_id": start_id,
-                "crate_id": end_id,
+                "import_id": end_id,
                 "semver_range": req,
                 "dependency_type": dependency_type,
             }
@@ -130,7 +130,7 @@ def user_packages(self) -> Generator[Dict[str, str], None, None]:
             owner_id = row["owner_id"]
 
             yield {
-                "crate_id": crate_id,
+                "import_id": crate_id,
                 "owner_id": owner_id,
             }
 

package_managers/pypi/Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY package_managers/pypi/requirements.txt .
+RUN pip install -r requirements.txt
+
+COPY . .
+
+CMD ["python", "-m", "package_managers.pypi.main"]

package_managers/pypi/README.md

@@ -0,0 +1,31 @@
+# PyPI
+
+The PyPI service processes package data from the Python Package Index (PyPI) and transforms it into CHAI's normalized format. It uses PyPI's JSON data dumps for efficient bulk processing.
+
+## Getting Started
+
+To run the PyPI service, use the following commands:
+
+```bash
+docker compose build pypi
+docker compose run pypi
+```
+
+## Execution Steps
+
+The PyPI loader follows these steps:
+
+1. Initialization: Sets up configuration and database connection
+2. Fetching: Downloads the latest PyPI JSON data dump if `FETCH` is true
+3. Transformation: Converts PyPI's JSON format into CHAI's schema
+4. Loading: Inserts transformed data into the database:
+   - Packages
+   - Users
+   - User Packages
+   - URLs
+   - Package URLs
+   - Versions
+   - Dependencies
+5. Cleanup: Removes temporary files if `NO_CACHE` is true
+
+The main execution logic is in the `run_pipeline` function in `main.py`.