70 lines
1.6 KiB
Go
70 lines
1.6 KiB
Go
package grib
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"git.intra.yksa.space/gsn/predictor/internal/pkg/errcodes"
|
|
"golang.org/x/sync/errgroup"
|
|
)
|
|
|
|
// NOMADS only.
|
|
const nomadsRoot = "https://nomads.ncep.noaa.gov/pub/data/nccf/com/gfs/prod"
|
|
|
|
type Downloader struct {
|
|
Dir string
|
|
Parallel int
|
|
Client *http.Client
|
|
}
|
|
|
|
func (d *Downloader) fileURL(run string, hour int, step int) string {
|
|
return fmt.Sprintf("%s/gfs.%s/%02d/atmos/gfs.t%02dz.pgrb2.0p50.f%03d", nomadsRoot, run, hour, hour, step)
|
|
}
|
|
|
|
func (d *Downloader) fetch(ctx context.Context, url, dst string) error {
|
|
if _, err := os.Stat(dst); err == nil {
|
|
return nil
|
|
}
|
|
tmp := dst + ".part"
|
|
f, err := os.Create(tmp)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer f.Close()
|
|
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
|
resp, err := d.Client.Do(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK {
|
|
return errcodes.Wrap(errcodes.ErrDownload, "bad status: "+resp.Status)
|
|
}
|
|
if _, err := io.Copy(f, resp.Body); err != nil {
|
|
return err
|
|
}
|
|
return os.Rename(tmp, dst)
|
|
}
|
|
|
|
func (d *Downloader) Run(ctx context.Context, run time.Time) error {
|
|
runStr := run.Format("20060102")
|
|
hour := run.Hour()
|
|
g, ctx := errgroup.WithContext(ctx)
|
|
sem := make(chan struct{}, d.Parallel)
|
|
for _, step := range steps {
|
|
step := step
|
|
sem <- struct{}{}
|
|
g.Go(func() error {
|
|
defer func() { <-sem }()
|
|
url := d.fileURL(runStr, hour, step)
|
|
dst := filepath.Join(d.Dir, fileName(run, step))
|
|
return d.fetch(ctx, url, dst)
|
|
})
|
|
}
|
|
return g.Wait()
|
|
}
|