101 lines
2.9 KiB
Go
101 lines
2.9 KiB
Go
|
//go:build llm_script
|
||
|
|
|
||
|
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"database/sql"
|
||
|
|
"flag"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"net/http"
|
||
|
|
"os"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type zhipuCodingPlanImportConfig struct {
|
||
|
|
OverviewURL string
|
||
|
|
PromotionURL string
|
||
|
|
OverviewFixture string
|
||
|
|
PromotionFixture string
|
||
|
|
DryRun bool
|
||
|
|
Timeout time.Duration
|
||
|
|
}
|
||
|
|
|
||
|
|
func main() {
|
||
|
|
loadSubscriptionImportEnv()
|
||
|
|
|
||
|
|
var overviewURL string
|
||
|
|
var promotionURL string
|
||
|
|
var overviewFixture string
|
||
|
|
var promotionFixture string
|
||
|
|
var dryRun bool
|
||
|
|
var timeoutSeconds int
|
||
|
|
|
||
|
|
flag.StringVar(&overviewURL, "overview-url", defaultZhipuCodingPlanOverviewURL, "智谱 Coding Plan 概览 URL")
|
||
|
|
flag.StringVar(&promotionURL, "promotion-url", defaultZhipuCodingPlanPromotionURL, "智谱 Coding Plan 活动页 URL")
|
||
|
|
flag.StringVar(&overviewFixture, "overview-fixture", "", "智谱 Coding Plan 概览样例文件")
|
||
|
|
flag.StringVar(&promotionFixture, "promotion-fixture", "", "智谱 Coding Plan 活动页样例文件")
|
||
|
|
flag.BoolVar(&dryRun, "dry-run", false, "仅解析并打印摘要，不写入数据库")
|
||
|
|
flag.IntVar(&timeoutSeconds, "timeout", 20, "请求超时（秒）")
|
||
|
|
flag.Parse()
|
||
|
|
|
||
|
|
cfg := zhipuCodingPlanImportConfig{
|
||
|
|
OverviewURL: overviewURL,
|
||
|
|
PromotionURL: promotionURL,
|
||
|
|
OverviewFixture: overviewFixture,
|
||
|
|
PromotionFixture: promotionFixture,
|
||
|
|
DryRun: dryRun,
|
||
|
|
Timeout: time.Duration(timeoutSeconds) * time.Second,
|
||
|
|
}
|
||
|
|
|
||
|
|
var db *sql.DB
|
||
|
|
var err error
|
||
|
|
if !cfg.DryRun {
|
||
|
|
db, err = subscriptionImportDB()
|
||
|
|
if err != nil {
|
||
|
|
fmt.Fprintf(os.Stderr, "open db: %v\n", err)
|
||
|
|
os.Exit(1)
|
||
|
|
}
|
||
|
|
defer db.Close()
|
||
|
|
}
|
||
|
|
|
||
|
|
if err := runZhipuCodingPlanImport(cfg, db, os.Stdout); err != nil {
|
||
|
|
fmt.Fprintf(os.Stderr, "import_zhipu_coding_plan: %v\n", err)
|
||
|
|
os.Exit(1)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func runZhipuCodingPlanImport(cfg zhipuCodingPlanImportConfig, db *sql.DB, out io.Writer) error {
|
||
|
|
client := &http.Client{Timeout: cfg.Timeout}
|
||
|
|
overviewRaw, err := fetchSubscriptionPage(cfg.OverviewURL, cfg.OverviewFixture, client)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
promotionRaw, err := fetchSubscriptionPage(cfg.PromotionURL, cfg.PromotionFixture, client)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
records, err := parseZhipuCodingPlanCatalog(overviewRaw, promotionRaw)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if cfg.DryRun {
|
||
|
|
_, err = fmt.Fprintf(out, "source=zhipu-coding-plan-import plans=%d provider=%s operator=%s dry_run=true\n", len(records), records[0].ProviderName, records[0].OperatorName)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if db == nil {
|
||
|
|
return fmt.Errorf("db is required when dry-run=false")
|
||
|
|
}
|
||
|
|
if err := upsertSubscriptionImportRecords(db, records); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
var tableRows int
|
||
|
|
if err := db.QueryRow(`SELECT COUNT(*) FROM subscription_plan`).Scan(&tableRows); err != nil {
|
||
|
|
return fmt.Errorf("count subscription_plan: %w", err)
|
||
|
|
}
|
||
|
|
_, err = fmt.Fprintf(out, "source=zhipu-coding-plan-import plans=%d provider=%s operator=%s table_rows=%d dry_run=false\n", len(records), records[0].ProviderName, records[0].OperatorName, tableRows)
|
||
|
|
return err
|
||
|
|
}
|