Coverage for app/sync.py: 79%
94 statements
« prev ^ index » next coverage.py v7.6.10, created at 2026-06-04 22:29 +0000
1import sys
2import time
4from google.cloud import bigquery
5from sqlalchemy import Engine, text
7from app.core.config import settings
8from app.core.db import read_engine, write_engine
9from app.core.logger import logger
10from app.utils import parse_dates, start_of_week
13def sync_pypi_packages(
14 read_engine: Engine, write_engine: Engine, limit: int | None = None
15) -> None:
16 logger.info("Syncing PyPI packages to SQLite...")
18 with read_engine.connect() as conn:
19 result = conn.execute(
20 text("select max(package_uploaded_at) from pypi_packages")
21 )
22 last_uploaded_at = result.scalar()
24 client = bigquery.Client()
26 select_query = f"""
27 select
28 package_name,
29 latest_package_version,
30 package_summary,
31 package_home_page,
32 package_download_url,
33 format_timestamp("%Y-%m-%dT%H:%M:%SZ", package_uploaded_at) as package_uploaded_at,
34 format_timestamp("%Y-%m-%dT%H:%M:%SZ", current_datetime("UTC")) as synced_at
35 from
36 pypacktrends-prod.dbt.pypi_packages
37 where
38 package_uploaded_at > '{last_uploaded_at or "1900-01-01T00:00:00Z"}'
39 order by
40 package_name
41 """
43 if limit is not None:
44 select_query += f"limit {limit}"
46 logger.info(f"Query most recent distributions: {select_query}")
48 BATCH_SIZE = 5000
50 job_config = bigquery.QueryJobConfig(
51 labels={
52 "application": "pypacktrends",
53 "component": "sync",
54 "type": "packages",
55 "environment": settings.ENVIRONMENT,
56 }
57 )
58 rows = client.query(select_query, job_config=job_config).result(
59 page_size=BATCH_SIZE
60 )
61 total_rows = rows.total_rows
63 if total_rows < 1:
64 logger.info("No rows to insert")
65 return
67 logger.info(f"Total rows to insert: {total_rows}")
69 upsert_sql = text("""
70 insert or replace into pypi_packages
71 (package_name, latest_package_version, package_summary, package_home_page, package_download_url, package_uploaded_at, synced_at)
72 values
73 (:package_name, :latest_package_version, :package_summary, :package_home_page, :package_download_url, :package_uploaded_at, :synced_at)
74 """)
76 start_time = time.time()
78 with write_engine.begin() as conn:
79 row_count = 0
80 batch_count = 0
81 for page in rows.pages:
82 batch_count += 1
83 start_batch_time = time.time()
85 packages = [dict(row.items()) for row in page]
87 row_count += len(packages)
89 logger.info(f"Inserting {row_count} of {total_rows} rows...")
91 conn.execute(upsert_sql, packages)
93 batch_time = time.time() - start_batch_time
94 logger.info(f"Batch {batch_count} inserted in {batch_time:.2f} seconds.")
96 total_time = time.time() - start_time
97 logger.info(f"Sync completed in {total_time:.2f} seconds.")
100def sync_pypi_downloads(
101 write_engine: Engine, start_date: str, end_date: str, limit: int | None = None
102) -> None:
103 start_week, end_week = start_of_week(start_date), start_of_week(end_date)
105 logger.info(
106 f"Syncing PyPI downloads to SQLite for weeks {start_week} to {end_week}"
107 )
109 client = bigquery.Client()
111 select_query = f"""
112 select
113 package_name,
114 package_downloaded_week,
115 downloads,
116 cumulative_downloads,
117 first_distribution_week,
118 weeks_since_first_distribution,
119 format_timestamp('%Y-%m-%dT%H:%M:%SZ', current_datetime("UTC")) as synced_at
120 from
121 pypacktrends-prod.dbt.pypi_package_downloads_weekly_metrics
122 where
123 package_downloaded_week between '{start_week}' and '{end_week}'
124 """
126 if limit is not None:
127 select_query += f"limit {limit}"
129 logger.info(f"Query downloads: {select_query}")
131 BATCH_SIZE = 5000
132 job_config = bigquery.QueryJobConfig(
133 labels={
134 "application": "pypacktrends",
135 "component": "sync",
136 "type": "downloads",
137 "environment": settings.ENVIRONMENT,
138 }
139 )
140 rows = client.query(select_query, job_config=job_config).result(
141 page_size=BATCH_SIZE
142 )
143 total_rows = rows.total_rows
145 if total_rows < 1:
146 logger.info("No rows to insert")
147 return
149 logger.info(f"Total rows to insert: {total_rows}")
151 delete_sql = text("""
152 delete from pypi_package_downloads_weekly_metrics
153 where package_downloaded_week between :start_week and :end_week
154 """)
156 insert_sql = text("""
157 insert into pypi_package_downloads_weekly_metrics (
158 package_name,
159 package_downloaded_week,
160 downloads,
161 cumulative_downloads,
162 first_distribution_week,
163 weeks_since_first_distribution,
164 synced_at
165 )
166 values (
167 :package_name,
168 :package_downloaded_week,
169 :downloads,
170 :cumulative_downloads,
171 :first_distribution_week,
172 :weeks_since_first_distribution,
173 :synced_at
174 )
175 """)
177 start_time = time.time()
179 with write_engine.begin() as conn:
180 conn.execute(delete_sql, {"start_week": start_week, "end_week": end_week})
181 row_count = 0
182 batch_count = 0
183 for page in rows.pages:
184 batch_count += 1
185 start_batch_time = time.time()
187 packages = [dict(row.items()) for row in page]
188 row_count += len(packages)
190 logger.info(f"Inserting {row_count} of {total_rows} rows...")
192 conn.execute(insert_sql, packages)
194 batch_time = time.time() - start_batch_time
195 logger.info(f"Batch {batch_count} inserted in {batch_time:.2f} seconds.")
197 total_time = time.time() - start_time
198 logger.info(f"Sync completed in {total_time:.2f} seconds.")
201def main() -> None:
202 if len(sys.argv) != 4:
203 logger.error("Usage: python sync.py <entity> <start_date> <end_date>")
204 sys.exit(1)
206 _, entity, start_date, end_date = sys.argv
208 match entity:
209 case "packages":
210 sync_pypi_packages(read_engine, write_engine)
211 case "downloads":
212 sync_pypi_downloads(write_engine, *parse_dates((start_date, end_date)))
213 case "all":
214 sync_pypi_packages(read_engine, write_engine)
215 sync_pypi_downloads(write_engine, *parse_dates((start_date, end_date)))
216 case _:
217 logger.error("Unknown entity. Please use 'packages' or 'downloads'")
218 sys.exit(1)
221if __name__ == "__main__":
222 main()