afanasyev.aa/predictor
1
Fork 0
forked from gsn/predictor
Code Pull requests Activity

feat: s3 download

Browse source
This commit is contained in:
Anatoly Antonov 2025-10-20 19:10:07 +09:00
parent a850615e1f
commit c4f355a32e
15 changed files with 590 additions and 109 deletions
Download patch file Download diff file Expand all files Collapse all files

71
internal/pkg/grib/grib.go
View file

@ -23,22 +23,13 @@ type Service interface {
GetStatus() (ready bool, lastUpdate time.Time, isFresh bool, errMsg string)
}
type ServiceConfig struct {
Dir string
TTL time.Duration
CacheTTL time.Duration
Parallel int
Client *http.Client
DatasetURL string
}
type service struct {
cfg ServiceConfig
cfg *Config
cache memCache
data atomic.Pointer[dataset]
}
func New(cfg ServiceConfig) (Service, error) {
func New(cfg *Config) (Service, error) {
if cfg.TTL == 0 {
cfg.TTL = 24 * time.Hour
}
@ -135,8 +126,7 @@ func (s *service) Update(ctx context.Context) error {
}
}
dl := Downloader{Dir: s.cfg.Dir, Parallel: s.cfg.Parallel, Client: s.cfg.Client, DatasetURL: s.cfg.DatasetURL}
run := nearestRun(time.Now().UTC().Add(-4 * time.Hour))
run := nearestRun(time.Now().UTC().Add(-24 * time.Hour))
// Check if we already have this run
cubePath := filepath.Join(s.cfg.Dir, run.Format("20060102_15")) + ".cube"
@ -156,9 +146,26 @@ func (s *service) Update(ctx context.Context) error {
}
}
// Download new data
if err := dl.Run(ctx, run); err != nil {
return err
// Download new data using S3 or HTTP
var downloadErr error
if s.cfg.UseS3 {
s3dl, err := NewS3Downloader(s.cfg.Dir, s.cfg.Parallel, s.cfg.S3Bucket, s.cfg.S3Region)
if err != nil {
return errcodes.Wrap(err, "failed to create S3 downloader")
}
downloadErr = s3dl.Run(ctx, run)
} else {
dl := Downloader{
Dir: s.cfg.Dir,
Parallel: s.cfg.Parallel,
Client: http.DefaultClient,
DatasetURL: s.cfg.DatasetURL,
}
downloadErr = dl.Run(ctx, run)
}
if downloadErr != nil {
return downloadErr
}
// Assemble cube if it doesn't exist
@ -179,8 +186,8 @@ func (s *service) Update(ctx context.Context) error {
}
func assembleCube(dir string, run time.Time, cubePath string) error {
const sizePerVar = 17 * 34 * 361 * 720 * 4
total := int64(sizePerVar * 2)
const sizePerVar = 97 * 47 * 361 * 720 * 4 // 97 time steps (0-96 hours), 47 pressure levels
total := int64(sizePerVar * 3) // 3 variables: gh, u, v
f, err := os.Create(cubePath)
if err != nil {
return err
@ -214,24 +221,30 @@ func assembleCube(dir string, run time.Time, cubePath string) error {
}
for _, m := range messages {
// Check if this is a wind component (u or v)
// Check if this is a wind component (u or v) or geopotential height
// ParameterCategory 2 = momentum, ParameterNumber 2 = u-wind, 3 = v-wind
// ParameterCategory 3 = mass, ParameterNumber 5 = geopotential height
if m.Section4.ProductDefinitionTemplateNumber != 0 {
continue
}
product := m.Section4.ProductDefinitionTemplate
if product.ParameterCategory != 2 {
continue
}
var varIdx int
switch product.ParameterNumber {
case 2: // u-wind
// Match tawhiri variable order: ['gh', 'u', 'v'] (indices 0, 1, 2)
if product.ParameterCategory == 2 {
switch product.ParameterNumber {
case 2: // u-wind
varIdx = 1
case 3: // v-wind
varIdx = 2
default:
continue
}
} else if product.ParameterCategory == 3 && product.ParameterNumber == 5 {
// geopotential height
varIdx = 0
case 3: // v-wind
varIdx = 1
default:
} else {
continue
}
@ -253,7 +266,7 @@ func assembleCube(dir string, run time.Time, cubePath string) error {
for i, v := range vals {
binary.LittleEndian.PutUint32(raw[i*4:], math.Float32bits(float32(v)))
}
base := int64(varIdx*sizePerVar + (ti*34+pIdx)*361*720*4)
base := int64(varIdx*sizePerVar + (ti*47+pIdx)*361*720*4)
copy(mm[base:base+int64(len(raw))], raw)
}
}
@ -266,7 +279,7 @@ func (s *service) Extract(ctx context.Context, lat, lon, alt float64, ts time.Ti
if d == nil {
return zero, errcodes.ErrNoDataset
}
if ts.Before(time.Unix(d.runUTC, 0)) || ts.After(time.Unix(d.runUTC, 0).Add(48*time.Hour)) {
if ts.Before(time.Unix(d.runUTC, 0)) || ts.After(time.Unix(d.runUTC, 0).Add(96*time.Hour)) {
return zero, errcodes.ErrOutOfBounds
}