wip: grib

internal/pkg/grib/downloader.go

@@ -0,0 +1,69 @@
+package grib
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"time"
+
+	"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 fmt.Errorf("bad status %s", 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()
+}