github.com/replit/upm@v0.0.0-20240423230255-9ce4fc3ea24c/internal/backends/python/gen_pypi_map/download_stats.go (about)

     1  package main
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"os"
     7  
     8  	"cloud.google.com/go/bigquery"
     9  	"golang.org/x/net/context"
    10  	"google.golang.org/api/iterator"
    11  )
    12  
    13  func FetchBQDownloads(gcp string, filePath string) error {
    14  	bqCache, err := GetPypiStats(gcp)
    15  	if err != nil {
    16  		return err
    17  	}
    18  
    19  	writer, err := os.Create(filePath)
    20  	if err != nil {
    21  		return err
    22  	}
    23  
    24  	cacheEncoder := json.NewEncoder(writer)
    25  	cacheEncoder.SetIndent("", "  ")
    26  	err = cacheEncoder.Encode(bqCache)
    27  
    28  	writer.Close()
    29  	return err
    30  }
    31  
    32  func LoadDownloadStats(path string) (map[string]int, error) {
    33  	file, err := os.Open(path)
    34  	if err != nil {
    35  		return nil, fmt.Errorf("could not open cache: %v", err)
    36  	}
    37  
    38  	packages := map[string]int{}
    39  
    40  	decoder := json.NewDecoder(file)
    41  	err = decoder.Decode(&packages)
    42  
    43  	return packages, err
    44  }
    45  
    46  func SaveDownloadStats(path string, cache PackageCache) {
    47  	writer, err := os.Create(path)
    48  	if err != nil {
    49  		fmt.Fprintf(os.Stderr, "Failed to open cache file for writing %v\n", err)
    50  	}
    51  	defer writer.Close()
    52  
    53  	cacheEncoder := json.NewEncoder(writer)
    54  	cacheEncoder.SetIndent("", "  ")
    55  	err = cacheEncoder.Encode(cache)
    56  	if err != nil {
    57  		panic(err)
    58  	}
    59  }
    60  
    61  func GetPypiStats(projectID string) (map[string]int, error) {
    62  	ctx := context.Background()
    63  	client, err := bigquery.NewClient(ctx, projectID)
    64  	if err != nil {
    65  		return map[string]int{}, fmt.Errorf("Failed to connect to bigquery [%s]: %v", projectID, err)
    66  	}
    67  
    68  	q := client.Query(
    69  		`select project as name, count(*) as downloads
    70  		from ` + "`bigquery-public-data.pypi.file_downloads`" + ` downloads
    71  		where
    72  		Date(timestamp) >= DATE_SUB(CURRENT_DATE(), INTERVAL 30 DAY)
    73  		group by project
    74  		order by downloads desc;`)
    75  
    76  	it, err := q.Read(ctx)
    77  	if err != nil {
    78  		return map[string]int{}, fmt.Errorf("Failed to query downloads: %v", err)
    79  	}
    80  
    81  	packages := map[string]int{}
    82  	for {
    83  		var info BqPackageInfo
    84  		err := it.Next(&info)
    85  		if err == iterator.Done {
    86  			break
    87  		}
    88  		if err != nil {
    89  			return packages, err
    90  		}
    91  		packages[info.Name] = info.Downloads
    92  	}
    93  
    94  	return packages, nil
    95  }