indexer: Index all recently uploaded projects

Author: zsol

py_wtf/cli/__init__.py

@@ -8,6 +8,8 @@
 import json
 import logging
 import shutil
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
 from functools import partial, wraps
 from pathlib import Path
 
@@ -22,7 +24,7 @@
 from py_wtf.indexer import index_dir, index_file, index_project
 from py_wtf.indexer.pypi import parse_deps, parse_upload_time
 from py_wtf.logging import setup_logging
-from py_wtf.repository import converter, ProjectRepository
+from py_wtf.repository import converter, METADATA_FILENAME, ProjectRepository
 from py_wtf.types import (
     Documentation,
     FQName,
@@ -135,6 +137,53 @@ async def index_dir_cmd(dir: str) -> None:
     rich.print(f"Found {cnt} modules in total.")
 
 
+@py_wtf.command(name="index-since")
+@click.option("--since", type=click.DateTime(), required=True)
+@click.argument("directory")
+@coroutine
+async def index_since(directory: str, since: datetime) -> None:
+    from google.cloud import bigquery
+
+    client = bigquery.Client(project="pypinfo-214114")
+    time_format = "%Y-%m-%d %H:%M:%S"
+    threadpool = ThreadPoolExecutor()
+    rows = await asyncio.get_running_loop().run_in_executor(
+        threadpool,
+        client.query_and_wait,
+        f"""
+        SELECT distinct name
+        FROM
+            `bigquery-public-data.pypi.distribution_metadata`
+        WHERE
+            TIMESTAMP(upload_time) >= TIMESTAMP("{since.strftime(time_format)}")
+        """,
+    )
+    rows = list(rows)
+    logger.info(f"Found {len(rows)} new projects to index")
+    out_dir = Path(directory)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    async with httpx.AsyncClient() as client:
+        resp = await client.get(f"https://py.wtf/_index/{METADATA_FILENAME}")
+        resp.raise_for_status()
+        (out_dir / METADATA_FILENAME).write_bytes(resp.content)
+
+    logger.info("Fetched prod index")
+    repo = ProjectRepository(out_dir)
+    rets = await asyncio.gather(
+        *[
+            repo.get(ProjectName(row.name), partial(index_project, repo=repo))
+            for row in rows
+        ],
+        return_exceptions=True,
+    )
+    for ret in rets:
+        if isinstance(ret, Exception):
+            logger.exception(ret)
+    logger.info("Done indexing")
+    repo.update_index()
+    logger.info("Wrote new index")
+
+
 @py_wtf.command()
 @click.argument("dir", required=False)
 @coroutine

py_wtf/repository.py

@@ -2,7 +2,7 @@
 import logging
 from asyncio import Future
 from collections import Counter, defaultdict
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from pathlib import Path
 from time import time
 from typing import AsyncIterable, Callable, Tuple
@@ -71,33 +71,36 @@ async def get(
             raise ValueError(f"{key} was never yielded by {factory}")
         return fut.result()
 
-    def generate_index(self, timestamp: int | None) -> Index:
+    def generate_index(
+        self, timestamp: int | None, skip_hydration: bool = False
+    ) -> Index:
         max_counts = 5
         latest_project_mtimes: list[Tuple[ProjectName, float]] = []
         all_project_names: list[ProjectName] = []
         dep_counts: Counter[str] = Counter()
 
-        for item in self.directory.iterdir():
-            if item.suffix != ".json" or not item.is_file():
-                continue
-            name = ProjectName(item.stem)
-            all_project_names.append(name)
-            if name not in self._cache:
-                self._load_from_disk(name)
-            project = self._cache[name].result()
-
-            for dep in project.metadata.dependencies:
-                dep_counts[dep] += 1
-
-            mtime = project.metadata.upload_time
-            if len(latest_project_mtimes) < max_counts:
-                latest_project_mtimes.append((name, mtime))
-                continue
-            for i in reversed(range(len(latest_project_mtimes))):
-                if mtime < latest_project_mtimes[i][1]:
-                    latest_project_mtimes.insert(i + 1, (name, mtime))
-            while len(latest_project_mtimes) > max_counts:
-                latest_project_mtimes.pop()
+        if not skip_hydration:
+            for item in self.directory.iterdir():
+                if item.suffix != ".json" or not item.is_file():
+                    continue
+                name = ProjectName(item.stem)
+                all_project_names.append(name)
+                if name not in self._cache:
+                    self._load_from_disk(name)
+                project = self._cache[name].result()
+
+                for dep in project.metadata.dependencies:
+                    dep_counts[dep] += 1
+
+                mtime = project.metadata.upload_time
+                if len(latest_project_mtimes) < max_counts:
+                    latest_project_mtimes.append((name, mtime))
+                    continue
+                for i in reversed(range(len(latest_project_mtimes))):
+                    if mtime < latest_project_mtimes[i][1]:
+                        latest_project_mtimes.insert(i + 1, (name, mtime))
+                while len(latest_project_mtimes) > max_counts:
+                    latest_project_mtimes.pop()
 
         latest_projects = [
             self._cache[name].result().metadata for name, _ in latest_project_mtimes
@@ -120,3 +123,20 @@ def generate_index(self, timestamp: int | None) -> Index:
     def write_index(self, timestamp: int | None = None) -> None:
         metadata = self.directory / METADATA_FILENAME
         metadata.write_text(converter.dumps(self.generate_index(timestamp)))
+
+    def update_index(self) -> None:
+        metadata = self.directory / METADATA_FILENAME
+        index = converter.loads(metadata.read_text(), Index)
+        new_index = self.generate_index(int(time()), skip_hydration=True)
+        index = replace(
+            index,
+            generated_at=new_index.generated_at,
+            latest_projects=sorted(
+                index.latest_projects + new_index.latest_projects,
+                key=lambda m: -m.upload_time,
+            ),
+            all_project_names=sorted(
+                {*index.all_project_names, *new_index.all_project_names}
+            ),
+        )
+        metadata.write_text(converter.dumps(index))

pyproject.toml

@@ -21,6 +21,7 @@ dependencies = [
   "appdirs==1.4.4",
   "cattrs==22.1.0",
   "click>=8.1.3",
+  "google-cloud-bigquery>=3.23.1",
   "httpx==0.23.0",
   "rich>=12.5.1",
   "libcst>=1.1.0",