228 lines
7.1 KiB
Go
228 lines
7.1 KiB
Go
package main
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"flag"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"time"
|
||
|
||
dbpkg "weatherstation/internal/database"
|
||
"weatherstation/internal/rain"
|
||
)
|
||
|
||
// 简单的小时雨量（CMPA）按时间范围下载器：
|
||
// - 输入时间为北京时间（Asia/Shanghai）
|
||
// - 构造下载路径使用 UTC（本地整点 -8h）
|
||
// - 入库前通过 rain.StoreTileBytes 使用 URL 解析将 UTC 还原为北京时间并写库
|
||
// 用法示例：
|
||
//
|
||
// go run ./cmd/rainfetch --from "2025-10-07 09:00:00" --to "2025-10-07 11:00:00" \
|
||
// --tiles "7/40/102,7/40/104" --outdir rain_data
|
||
func main() {
|
||
var (
|
||
fromStr = flag.String("from", "", "起始时间（北京时间，YYYY-MM-DD HH:MM:SS 或 YYYY-MM-DD）")
|
||
toStr = flag.String("to", "", "结束时间（北京时间，YYYY-MM-DD HH:MM:SS 或 YYYY-MM-DD）")
|
||
tiles = flag.String("tiles", "7/40/102,7/40/104", "瓦片列表，逗号分隔，每项为 z/y/x，如 7/40/102")
|
||
outDir = flag.String("outdir", "rain_data", "保存目录（同时也会写入数据库）")
|
||
baseURL = flag.String("base", "https://image.data.cma.cn/tiles/China/CMPA_RT_China_0P01_HOR-PRE_GISJPG_Tiles/%Y%m%d/%H/%M/{z}/{y}/{x}.bin", "下载基础URL模板（UTC路径时间）")
|
||
dryRun = flag.Bool("dry", false, "仅打印将要下载的URL与目标，不实际下载写库")
|
||
)
|
||
flag.Parse()
|
||
|
||
if strings.TrimSpace(*fromStr) == "" || strings.TrimSpace(*toStr) == "" {
|
||
log.Fatalln("必须提供 --from 与 --to（北京时间）")
|
||
}
|
||
|
||
loc, _ := time.LoadLocation("Asia/Shanghai")
|
||
if loc == nil {
|
||
loc = time.FixedZone("CST", 8*3600)
|
||
}
|
||
|
||
parseCST := func(s string) (time.Time, error) {
|
||
s = strings.TrimSpace(s)
|
||
var t time.Time
|
||
var err error
|
||
if len(s) == len("2006-01-02") {
|
||
// 日期：按 00:00:00 处理
|
||
if tm, e := time.ParseInLocation("2006-01-02", s, loc); e == nil {
|
||
t = tm
|
||
} else {
|
||
err = e
|
||
}
|
||
} else {
|
||
t, err = time.ParseInLocation("2006-01-02 15:04:05", s, loc)
|
||
}
|
||
return t, err
|
||
}
|
||
|
||
start, err1 := parseCST(*fromStr)
|
||
end, err2 := parseCST(*toStr)
|
||
if err1 != nil || err2 != nil {
|
||
log.Fatalf("解析时间失败: from=%v to=%v", err1, err2)
|
||
}
|
||
if end.Before(start) {
|
||
log.Fatalln("结束时间需不小于起始时间")
|
||
}
|
||
|
||
// 小时步进（包含端点）：先对齐到小时
|
||
cur := start.Truncate(time.Hour)
|
||
end = end.Truncate(time.Hour)
|
||
|
||
// 解析 tiles 参数
|
||
type tcoord struct{ z, y, x int }
|
||
var tlist []tcoord
|
||
for _, part := range strings.Split(*tiles, ",") {
|
||
p := strings.TrimSpace(part)
|
||
if p == "" {
|
||
continue
|
||
}
|
||
var z, y, x int
|
||
if _, err := fmt.Sscanf(p, "%d/%d/%d", &z, &y, &x); err != nil {
|
||
log.Fatalf("无效的 tiles 项: %s", p)
|
||
}
|
||
tlist = append(tlist, tcoord{z, y, x})
|
||
}
|
||
if len(tlist) == 0 {
|
||
log.Fatalln("tiles 解析后为空")
|
||
}
|
||
|
||
if err := os.MkdirAll(*outDir, 0o755); err != nil {
|
||
log.Fatalf("创建输出目录失败: %v", err)
|
||
}
|
||
|
||
ctx := context.Background()
|
||
total := 0
|
||
success := 0
|
||
for !cur.After(end) {
|
||
// 本地整点（CST）→ UTC 路径时间
|
||
slotLocal := cur
|
||
slotUTC := slotLocal.Add(-8 * time.Hour).In(time.UTC)
|
||
dateStr := slotUTC.Format("20060102")
|
||
hh := slotUTC.Format("15")
|
||
mm := "00"
|
||
log.Printf("[rainfetch] 时次 local=%s, utc=%s", slotLocal.Format("2006-01-02 15:04"), slotUTC.Format("2006-01-02 15:04"))
|
||
|
||
for _, tc := range tlist {
|
||
total++
|
||
// 构造 URL
|
||
url := *baseURL
|
||
url = strings.ReplaceAll(url, "%Y%m%d", dateStr)
|
||
url = strings.ReplaceAll(url, "%H", hh)
|
||
url = strings.ReplaceAll(url, "%M", mm)
|
||
url = strings.ReplaceAll(url, "{z}", fmt.Sprintf("%d", tc.z))
|
||
url = strings.ReplaceAll(url, "{y}", fmt.Sprintf("%d", tc.y))
|
||
url = strings.ReplaceAll(url, "{x}", fmt.Sprintf("%d", tc.x))
|
||
|
||
fname := fmt.Sprintf("rain_z%d_y%d_x%d_%s.bin", tc.z, tc.y, tc.x, slotLocal.Format("20060102_1504"))
|
||
dest := filepath.Join(*outDir, fname)
|
||
|
||
if *dryRun {
|
||
log.Printf("[rainfetch] DRY url=%s -> %s", url, dest)
|
||
continue
|
||
}
|
||
|
||
// 若 DB 已有，则跳过
|
||
if ref, e := rain.ParseCMPATileURL(url); e == nil {
|
||
exists, e2 := databaseHas(ctx, ref.Product, ref.DT, tc.z, tc.y, tc.x)
|
||
if e2 == nil && exists {
|
||
log.Printf("[rainfetch] skip exists in DB z=%d y=%d x=%d dt=%s", tc.z, tc.y, tc.x, ref.DT.Format("2006-01-02 15:04"))
|
||
continue
|
||
}
|
||
}
|
||
|
||
if err := httpDownloadTo(ctx, url, dest); err != nil {
|
||
log.Printf("[rainfetch] 下载失败 z=%d y=%d x=%d: %v", tc.z, tc.y, tc.x, err)
|
||
continue
|
||
}
|
||
log.Printf("[rainfetch] 保存 %s", dest)
|
||
|
||
// 写库（使用 URL 解析 UTC → CST 后 upsert）
|
||
b, rerr := os.ReadFile(dest)
|
||
if rerr != nil {
|
||
log.Printf("[rainfetch] 读文件失败: %v", rerr)
|
||
continue
|
||
}
|
||
if err := rain.StoreTileBytes(ctx, url, b); err != nil {
|
||
log.Printf("[rainfetch] 入库失败: %v", err)
|
||
continue
|
||
}
|
||
success++
|
||
}
|
||
cur = cur.Add(1 * time.Hour)
|
||
}
|
||
log.Printf("[rainfetch] 完成：尝试 %d，成功 %d", total, success)
|
||
}
|
||
|
||
// 轻量 DB 存在性检查（避免引入内部 database 包到该命令）：
|
||
// 为了避免循环依赖，这里复制一份最小 SQL 调用；实际工程也可抽取共享函数。
|
||
// 但当前 repo 中 database.GetDB 在 internal/database 包内，雨量 API 直接使用它。
|
||
|
||
// 注意：为保持最小侵入，这里通过 rain.StoreTileBytes 完成入库；
|
||
// 仅在下载前进行“是否已存在”查询，避免重复下载。为此需要访问 internal/database。
|
||
|
||
func databaseHas(ctx context.Context, product string, dt time.Time, z, y, x int) (bool, error) {
|
||
const q = `SELECT 1 FROM rain_tiles WHERE product=$1 AND dt=$2 AND z=$3 AND y=$4 AND x=$5 LIMIT 1`
|
||
var one int
|
||
err := dbpkg.GetDB().QueryRowContext(ctx, q, product, dt, z, y, x).Scan(&one)
|
||
if err == sql.ErrNoRows {
|
||
return false, nil
|
||
}
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
return true, nil
|
||
}
|
||
|
||
func httpDownloadTo(ctx context.Context, url, dest string) error {
|
||
client := &http.Client{Timeout: 20 * time.Second}
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||
if err != nil {
|
||
return fmt.Errorf("build request: %w", err)
|
||
}
|
||
req.Header.Set("Referer", "https://data.cma.cn/")
|
||
req.Header.Set("Origin", "https://data.cma.cn")
|
||
req.Header.Set("User-Agent", "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Safari/537.36")
|
||
resp, err := client.Do(req)
|
||
if err != nil {
|
||
return fmt.Errorf("http get: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != http.StatusOK {
|
||
return fmt.Errorf("unexpected status: %d", resp.StatusCode)
|
||
}
|
||
tmp := dest + ".part"
|
||
f, err := os.Create(tmp)
|
||
if err != nil {
|
||
return fmt.Errorf("create temp: %w", err)
|
||
}
|
||
_, copyErr := io.Copy(f, resp.Body)
|
||
closeErr := f.Close()
|
||
if copyErr != nil {
|
||
_ = os.Remove(tmp)
|
||
return fmt.Errorf("write body: %w", copyErr)
|
||
}
|
||
if closeErr != nil {
|
||
_ = os.Remove(tmp)
|
||
return fmt.Errorf("close temp: %w", closeErr)
|
||
}
|
||
if err := os.Rename(tmp, dest); err != nil {
|
||
// Cross-device fallback
|
||
data, rerr := os.ReadFile(tmp)
|
||
if rerr != nil {
|
||
return fmt.Errorf("read temp: %w", rerr)
|
||
}
|
||
if werr := os.WriteFile(dest, data, 0o644); werr != nil {
|
||
return fmt.Errorf("write final: %w", werr)
|
||
}
|
||
_ = os.Remove(tmp)
|
||
}
|
||
return nil
|
||
}
|