Coverage for app/publish.py: 0%

1import os 

2import sys 

3import time 

4 

5import duckdb 

6from duckdb import DuckDBPyConnection 

7from google.cloud import bigquery 

8 

9from app.core.config import settings 

10from app.core.logger import logger 

11from app.utils import parse_dates 

12 

13 

14class DuckDBClient: 

15 def __init__( 

16 self, 

17 gcp_project: str, 

18 r2_access_key_id: str, 

19 r2_secret_access_key: str, 

20 r2_account_id: str, 

21 r2_bucket_name: str, 

22 r2_custom_domain: str, 

23 ): 

24 self.gcp_project = gcp_project 

25 self.r2_access_key_id = r2_access_key_id 

26 self.r2_secret_access_key = r2_secret_access_key 

27 self.r2_account_id = r2_account_id 

28 self.r2_bucket_path = f"r2://{r2_bucket_name}" 

29 self.r2_custom_domain = f"https://{r2_custom_domain}" 

30 self.conn = self._init_duckdb_client() 

31 

32 def _init_duckdb_client(self) -> DuckDBPyConnection: 

33 conn = duckdb.connect("/tmp/duckdb_github_actions.db") 

34 sql = f""" 

35CREATE SECRET ( 

36 TYPE r2, 

37 KEY_ID '{self.r2_access_key_id}', 

38 SECRET '{self.r2_secret_access_key}', 

39 ACCOUNT_ID '{self.r2_account_id}' 

40); 

41 

42CREATE OR REPLACE TABLE pypi_downloads ( 

43 package_name VARCHAR, 

44 package_downloaded_week DATE, 

45 downloads BIGINT, 

46 cumulative_downloads BIGINT, 

47 first_distribution_week DATE, 

48 weeks_since_first_distribution INTEGER, 

49 published_at TIMESTAMP, 

50 week DATE 

51); 

52""" 

53 conn.execute(sql) 

54 logger.info("DuckDB client initialized") 

55 return conn 

56 

57 

58def publish_pypi_downloads_manifest(duckdb_client: DuckDBClient) -> None: 

59 export_path = f"{duckdb_client.r2_bucket_path}/pypi-weekly-downloads" 

60 manifest_path = f"{export_path}/manifest.parquet" 

61 

62 sql = f""" 

63CREATE OR REPLACE TABLE pypi_downloads_manifest AS 

64with urls as ( 

65 select distinct 

66 regexp_extract(filename, 'week=([0-9\\-]+)', 1) AS week, 

67 replace( 

68 replace(filename, '{duckdb_client.r2_bucket_path}', '{duckdb_client.r2_custom_domain}'), 

69 '=', 

70 '%3D' 

71 ) as url 

72 from 

73 read_parquet('{export_path}/*/*.parquet') 

74) 

75 

76select * from urls order by 1; 

77 

78COPY pypi_downloads_manifest TO '{manifest_path}' 

79(FORMAT parquet, OVERWRITE_OR_IGNORE); 

80""" 

81 duckdb_client.conn.execute(sql) 

82 logger.info(f"Manifest updated: {manifest_path}") 

83 

84 

85def publish_pypi_downloads( 

86 duckdb_client: DuckDBClient, start_date: str, end_date: str 

87) -> None: 

88 export_path = f"{duckdb_client.r2_bucket_path}/pypi-weekly-downloads/" 

89 

90 client = bigquery.Client() 

91 

92 select_query = f""" 

93SELECT 

94 package_name, 

95 package_downloaded_week, 

96 downloads, 

97 cumulative_downloads, 

98 first_distribution_week, 

99 weeks_since_first_distribution, 

100 current_datetime('UTC') as published_at, 

101 package_downloaded_week as week 

102FROM 

103 `{duckdb_client.gcp_project}.dbt.pypi_package_downloads_weekly_metrics` 

104WHERE 

105 package_downloaded_week BETWEEN '{start_date}' AND '{end_date}' 

106ORDER BY 

107 package_downloaded_week, 

108 package_name 

109""" 

110 

111 logger.info(f"Query most recent distributions: {select_query}") 

112 

113 job_config = bigquery.QueryJobConfig( 

114 labels={ 

115 "application": "pypacktrends", 

116 "component": "publish", 

117 "type": "downloads", 

118 "environment": settings.ENVIRONMENT, 

119 } 

120 ) 

121 rows = client.query(select_query, job_config=job_config).result() 

122 total_rows = rows.total_rows 

123 

124 if total_rows < 1: 

125 logger.info("No rows to insert") 

126 return 

127 

128 logger.info(f"Total rows to insert: {total_rows}") 

129 

130 upsert_sql = """ 

131 INSERT INTO pypi_downloads 

132 (package_name, package_downloaded_week, downloads, cumulative_downloads, first_distribution_week, weeks_since_first_distribution, published_at, week) 

133 VALUES (?, ?, ?, ?, ?, ?, ?, ?) 

134 """ 

135 

136 start_time = time.time() 

137 row_count = 0 

138 batch_count = 0 

139 for page in rows.pages: 

140 batch_count += 1 

141 start_batch_time = time.time() 

142 

143 downloads = [ 

144 ( 

145 row["package_name"], 

146 row["package_downloaded_week"], 

147 row["downloads"], 

148 row["cumulative_downloads"], 

149 row["first_distribution_week"], 

150 row["weeks_since_first_distribution"], 

151 row["published_at"], 

152 row["week"], 

153 ) 

154 for row in page 

155 ] 

156 

157 row_count += len(downloads) 

158 

159 logger.info(f"Inserting {row_count} of {total_rows} rows...") 

160 

161 duckdb_client.conn.executemany(upsert_sql, downloads) 

162 

163 batch_time = time.time() - start_batch_time 

164 logger.info(f"Batch {batch_count} inserted in {batch_time:.2f} seconds.") 

165 

166 export_sql = f""" 

167COPY pypi_downloads TO '{export_path}' 

168(FORMAT parquet, PARTITION_BY (week), OVERWRITE_OR_IGNORE); 

169""" 

170 

171 duckdb_client.conn.execute(export_sql) 

172 total_time = time.time() - start_time 

173 logger.info( 

174 f"Successfully exported PyPI weekly downloads from {start_date} to {end_date} to R2 bucket path: {export_path} in {total_time:.2f} seconds." 

175 ) 

176 

177 

178def main() -> None: 

179 if len(sys.argv) != 3: 

180 logger.error("Usage: python publish.py <start_date> <end_date>") 

181 sys.exit(1) 

182 

183 _, start_date, end_date = sys.argv 

184 

185 gcp_project = os.getenv("GCP_PROJECT") 

186 

187 if gcp_project is None: 

188 raise ValueError("GCP_PROJECT environment variable is not set") 

189 

190 cloudflare_account_id = os.getenv("CLOUDFLARE_ACCOUNT_ID") 

191 if cloudflare_account_id is None: 

192 raise ValueError("CLOUDFLARE_ACCOUNT_ID environment variable is not set") 

193 

194 r2_access_key_id = os.getenv("CLOUDFLARE_R2_ACCESS_KEY_ID") 

195 if r2_access_key_id is None: 

196 raise ValueError("CLOUDFLARE_R2_ACCESS_KEY_ID environment variable is not set") 

197 

198 r2_secret_access_key = os.getenv("CLOUDFLARE_R2_SECRET_ACCESS_KEY") 

199 if r2_secret_access_key is None: 

200 raise ValueError( 

201 "CLOUDFLARE_R2_SECRET_ACCESS_KEY environment variable is not set" 

202 ) 

203 

204 r2_bucket_name = os.getenv("CLOUDFLARE_R2_PUBLIC_BUCKET_NAME") 

205 if r2_bucket_name is None: 

206 raise ValueError( 

207 "CLOUDFLARE_R2_PUBLIC_BUCKET_NAME environment variable is not set" 

208 ) 

209 

210 r2_custom_domain = os.getenv("CLOUDFLARE_R2_CUSTOM_DOMAIN") 

211 if r2_custom_domain is None: 

212 raise ValueError("CLOUDFLARE_R2_CUSTOM_DOMAIN environment variable is not set") 

213 

214 duckdb_client = DuckDBClient( 

215 gcp_project=gcp_project, 

216 r2_access_key_id=r2_access_key_id, 

217 r2_secret_access_key=r2_secret_access_key, 

218 r2_account_id=cloudflare_account_id, 

219 r2_bucket_name=r2_bucket_name, 

220 r2_custom_domain=r2_custom_domain, 

221 ) 

222 

223 publish_pypi_downloads(duckdb_client, *parse_dates((start_date, end_date))) 

224 publish_pypi_downloads_manifest(duckdb_client) 

225 

226 

227if __name__ == "__main__": 

228 main()