From 9c8e9d523ac42d0dc9604b4a291831a6260a3923 Mon Sep 17 00:00:00 2001 From: anyongjin Date: Wed, 7 Oct 2026 21:23:21 +0800 Subject: [PATCH] update --- .gitignore | 5 +- README.md | 2 +- cbpred/latest_predict.py | 109 +++++++++++++++++++++++++----- cbpred/pipeline.py | 96 +++++++++++++++++++++------ cmd/cbpred-export/main.go | 62 ++++++++++++++--- cmd/cbpred-export/main_test.go | 95 ++++++++++++++++++++++++++ docs/usage.md | 19 ++++-- go.mod | 2 +- scripts/update_and_predict.sh | 3 + tests/test_latest_predict.py | 117 +++++++++++++++++++++++++++++++++ tests/test_pipeline.py | 100 +++++++++++++++++++++++++++- 11 files changed, 557 insertions(+), 53 deletions(-) mode change 100755 => 100644 scripts/update_and_predict.sh diff --git a/.gitignore b/.gitignore index 99a0935..eed4de3 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,10 @@ __pycache__/ *.py[cod] .pytest_cache/ +config/banbot.local.yml +.cbpred-latest-predict.lock +go.work +go.work.sum data/raw/ data/manifest.json artifacts/ @@ -9,4 +13,3 @@ artifacts_*/ reports/ predictions/ data_*/ - diff --git a/README.md b/README.md index 44b2774..42caeb0 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ ## 快速开始 -运行环境需要 Python 3、Go、可访问的 banbot 源码和 banbot 配置中的 QuestDB。默认路径是： +运行环境需要 Python 3、Go、可访问的 banbot 源码和数据库。脚本优先使用已有 TimescaleDB 配置，支持 `--banbot-config`、`--dsn` 和 `--db-type timescale`。Linux 的 banbot 路径通常是： ```text /data/ban/banbot diff --git a/cbpred/latest_predict.py b/cbpred/latest_predict.py index 94679db..cd21fe7 100644 --- a/cbpred/latest_predict.py +++ b/cbpred/latest_predict.py @@ -2,19 +2,34 @@ from __future__ import annotations import argparse import csv -import fcntl +import errno import json +import os import socket import subprocess import tempfile +import time from datetime import datetime, timezone from pathlib import Path -from urllib.parse import urlparse + +if os.name == "nt": + import msvcrt +else: + import fcntl import joblib import pandas as pd +import yaml +from psycopg.conninfo import conninfo_to_dict -from .pipeline import load_database_url, parse_time_ms, predict_from_csv, prepare_data, timeframe_step_ms +from .pipeline import ( + effective_banbot_config, + load_database_url, + parse_time_ms, + predict_from_csv, + prepare_data, + timeframe_step_ms, +) REPO_ROOT = Path(__file__).resolve().parents[1] @@ -32,6 +47,22 @@ OUTPUT_COLUMNS = ( ) +def _acquire_update_lock(lock) -> None: + if os.name != "nt": + fcntl.flock(lock, fcntl.LOCK_EX) + return + # Windows locks a byte range; every caller locks the first byte. + while True: + lock.seek(0) + try: + msvcrt.locking(lock.fileno(), msvcrt.LK_NBLCK, 1) + return + except OSError as exc: + if exc.errno not in {errno.EACCES, errno.EAGAIN, errno.EDEADLK}: + raise + time.sleep(0.1) + + def _utc_label(timestamp_ms: int) -> str: return datetime.fromtimestamp(timestamp_ms / 1000, timezone.utc).strftime( "%Y-%m-%d %H:%M:%S UTC" @@ -39,8 +70,35 @@ def _utc_label(timestamp_ms: int) -> str: def _database_endpoint(dsn: str) -> tuple[str, int]: - parsed = urlparse(dsn) - return parsed.hostname or "127.0.0.1", parsed.port or 8812 + parameters = conninfo_to_dict(dsn) + host = parameters.get("host") or parameters.get("hostaddr") or "127.0.0.1" + port = parameters.get("port") or "5432" + return host.split(",")[0], int(port.split(",")[0] or 5432) + + +def default_banbot_dir() -> Path: + configured = os.environ.get("CBPRED_BANBOT_DIR") + if configured: + return Path(configured) + candidates = ( + REPO_ROOT.parent.parent / "ban/banbot", + REPO_ROOT.parent / "banbot", + Path("/data/ban/banbot"), + ) + return next((path for path in candidates if path.is_dir()), candidates[-1]) + + +def default_banbot_config(banbot_dir: Path) -> Path: + configured = os.environ.get("CBPRED_BANBOT_CONFIG") + if configured: + return Path(configured) + candidates = ( + REPO_ROOT / "config/banbot.local.yml", + banbot_dir.parent / "data/config.timescale.yml", + banbot_dir / "biz/config.timescale.yml", + banbot_dir / "biz/config.yml", + ) + return next((path for path in candidates if path.is_file()), candidates[-1]) def output_columns_for_timeframe(timeframe: str) -> tuple[str, ...]: @@ -98,28 +156,40 @@ def _validate_freshness( def ensure_database( banbot_config: Path, dsn: str | None, + db_type: str = "auto", ) -> None: database_url = load_database_url(banbot_config, dsn) + if db_type == "auto" and dsn is None: + config = yaml.safe_load(banbot_config.read_text(encoding="utf-8")) or {} + db_type = (config.get("database") or {}).get("db_type") or "auto" + # PostgreSQL connections are validated by prepare_data via libpq, which also + # supports Unix sockets, multiple hosts and connection-service definitions. + if db_type not in {"auto", "questdb"}: + return + parameters = conninfo_to_dict(database_url) + if any("," in parameters.get(key, "") for key in ("host", "hostaddr", "port")): + return host, port = _database_endpoint(database_url) + if db_type == "auto" and port != 8812: + return try: with socket.create_connection((host, port), timeout=1): return except OSError: pass - if port != 8812: - return - print("[latest-predict] QuestDB is unavailable; letting banbot start it", flush=True) start_ms = int(datetime.now(timezone.utc).timestamp() * 1000) - 8 * 60 * 60 * 1000 end_ms = int(datetime.now(timezone.utc).timestamp() * 1000) - with tempfile.TemporaryDirectory(prefix="cbpred-db-bootstrap-") as output_dir: + with tempfile.TemporaryDirectory(prefix="cbpred-db-bootstrap-") as output_dir, effective_banbot_config( + banbot_config, database_url, "questdb" + ) as effective_config: subprocess.run( [ "go", "run", "./cmd/cbpred-export", "--config", - str(banbot_config), + str(effective_config), "--data-dir", str(banbot_config.parent), "--pairs", @@ -222,11 +292,12 @@ def build_predictions(data_dir: Path, artifact: Path, output_dir: Path) -> dict[ def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Update banbot K-lines and predict with the current model") - parser.add_argument("--banbot-dir", type=Path, default=Path("/data/ban/banbot")) + parser.add_argument("--banbot-dir", type=Path, default=default_banbot_dir()) parser.add_argument( "--banbot-config", type=Path, - default=Path("/data/ban/banbot/biz/config.yml"), + default=None, + help="banbot config (default: local config or existing TimescaleDB config)", ) parser.add_argument("--data-dir", type=Path, default=None) parser.add_argument("--artifact", type=Path, default=None) @@ -241,8 +312,11 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--min-age-days", type=int, default=730) parser.add_argument("--chunk-size", type=int, default=100) parser.add_argument("--end", default=None, help="UTC date/ISO timestamp or milliseconds") - parser.add_argument("--dsn", default=None) - parser.add_argument("--db-type", choices=["auto", "questdb", "timescale"], default="auto") + parser.add_argument("--dsn", default=os.environ.get("CBPRED_DSN")) + parser.add_argument( + "--db-type", choices=["auto", "questdb", "timescale"], + default=os.environ.get("CBPRED_DB_TYPE", "auto"), + ) parser.add_argument("--exchange", default="binance") parser.add_argument("--timeframe", choices=["8h", "1d"], default="8h") return parser @@ -250,6 +324,9 @@ def build_parser() -> argparse.ArgumentParser: def main() -> None: args = build_parser().parse_args() + args.banbot_config = ( + args.banbot_config or default_banbot_config(args.banbot_dir) + ).resolve() data_dir = (args.data_dir or default_data_dir(args.timeframe)).resolve() artifact = (args.artifact or default_artifact(args.timeframe)).resolve() output_dir = ( @@ -258,8 +335,8 @@ def main() -> None: _validate_artifact_timeframe(artifact, args.timeframe) lock_path = args.banbot_config.parent / ".cbpred-latest-predict.lock" with lock_path.open("a+") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - ensure_database(args.banbot_config, args.dsn) + _acquire_update_lock(lock) + ensure_database(args.banbot_config, args.dsn, args.db_type) manifest = prepare_data( out_dir=data_dir, banbot_dir=args.banbot_dir, diff --git a/cbpred/pipeline.py b/cbpred/pipeline.py index 82b84f2..35f3d83 100644 --- a/cbpred/pipeline.py +++ b/cbpred/pipeline.py @@ -8,10 +8,12 @@ import re import shlex import subprocess import sys +import tempfile +from contextlib import contextmanager from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path -from typing import Any, Iterable +from typing import Any, Iterable, Iterator import joblib import numpy as np @@ -141,13 +143,61 @@ def detect_db_type(conn, requested: str = "auto") -> str: return requested with conn.cursor() as cur: try: - cur.execute("SELECT 1 FROM exsymbol_q LIMIT 1") - return "questdb" - except Exception: + cur.execute("SELECT to_regclass('exsymbol'), to_regclass('exsymbol_q')") + timescale_table, questdb_table = cur.fetchone() + except Exception as exc: conn.rollback() + # QuestDB speaks the PostgreSQL wire protocol but lacks to_regclass. + message = str(exc).lower() + if getattr(exc, "sqlstate", None) != "42883" and not ( + "to_regclass" in message + and any(text in message for text in ("unknown function", "does not exist", "unsupported", "no signature")) + ): + raise + else: + if timescale_table is not None: + return "timescale" + if questdb_table is not None: + return "questdb" + return "timescale" + for table, backend in (("exsymbol", "timescale"), ("exsymbol_q", "questdb")): + try: + cur.execute(f"SELECT 1 FROM {table} LIMIT 1") + return backend + except Exception as exc: + conn.rollback() + message = str(exc).lower() + if getattr(exc, "sqlstate", None) != "42P01" and not ( + table in message and any(text in message for text in ("does not exist", "not found")) + ): + raise return "timescale" +@contextmanager +def effective_banbot_config(config_path: Path, dsn: str, db_type: str) -> Iterator[Path]: + """Give Python and Go the same database without changing the original config.""" + config_path = config_path.resolve() + config = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {} + database = dict(config.get("database") or {}) + database.update(url=dsn, db_type=db_type) + config["database"] = database + config["market_type"] = "linear" + temporary_path = None + try: + # Keep relative assets and --datadir anchored to the original directory. + with tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", suffix=".yml", prefix=".cbpred-db-", + dir=config_path.parent, delete=False, + ) as handle: + temporary_path = Path(handle.name) + yaml.safe_dump(config, handle, allow_unicode=True, sort_keys=False) + yield temporary_path + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + + def fetch_symbols( conn, db_type: str, @@ -250,7 +300,10 @@ def run_banbot_download( str(end_ms), ] try: - subprocess.run(command, cwd=banbot_dir, check=True, stderr=subprocess.PIPE, text=True) + subprocess.run( + command, cwd=banbot_dir, check=True, + stderr=subprocess.PIPE, text=True, encoding="utf-8", + ) except subprocess.CalledProcessError as exc: invalid_symbols = set(INVALID_SYMBOL_RE.findall(exc.stderr or "")) rejected = [item for item in chunk if item.symbol in invalid_symbols] @@ -391,27 +444,28 @@ def prepare_data( ) finally: conn.close() - if not skip_download: - symbols, unavailable = run_banbot_download( + with effective_banbot_config(banbot_config, database_url, actual_db_type) as effective_config: + if not skip_download: + symbols, unavailable = run_banbot_download( + symbols, + banbot_dir, + effective_config, + history_start_ms, + end_ms, + chunk_size, + banbot_cmd, + ) + excluded.extend({**symbol.__dict__, "reason": "unavailable_on_exchange"} for symbol in unavailable) + run_banbot_export( symbols, banbot_dir, - banbot_config, + effective_config, + raw_dir, history_start_ms, end_ms, - chunk_size, - banbot_cmd, + banbot_export_cmd, + timeframe, ) - excluded.extend({**symbol.__dict__, "reason": "unavailable_on_exchange"} for symbol in unavailable) - run_banbot_export( - symbols, - banbot_dir, - banbot_config, - raw_dir, - history_start_ms, - end_ms, - banbot_export_cmd, - timeframe, - ) manifest_symbols: list[dict[str, Any]] = [] for symbol in symbols: path = raw_dir / symbol_filename(symbol) diff --git a/cmd/cbpred-export/main.go b/cmd/cbpred-export/main.go index f5ce3d1..a7c1efd 100644 --- a/cmd/cbpred-export/main.go +++ b/cmd/cbpred-export/main.go @@ -6,6 +6,7 @@ import ( "fmt" "os" "path/filepath" + "sort" "strconv" "strings" @@ -64,12 +65,21 @@ func main() { } fail(err) } - _, rows, err := orm.GetSeries(exs, *timeframe, *startMS, *endMS, 0, false) + _, rows, err := orm.GetSeries(exs, "1h", *startMS, *endMS, 0, false) if err != nil { fail(err) } + bars := make([]*orm.SeriesOHLCV, 0, len(rows)) + for _, row := range rows { + bar, err := row.OHLCV(exs) + if err != nil { + fail(err) + } + bars = append(bars, bar) + } + bars = aggregateHourly(bars, *timeframe) name := fmt.Sprintf("%d_%s.csv", exs.ID, strings.NewReplacer("/", "_", ":", "_").Replace(pair)) - if err := writeCSV(filepath.Join(*outDir, name), rows, exs); err != nil { + if err := writeCSV(filepath.Join(*outDir, name), bars); err != nil { fail(err) } } @@ -79,7 +89,47 @@ func isUnavailablePair(err *errs.Error) bool { return err != nil && err.Code == errs.CodeNoMarketForPair } -func writeCSV(path string, rows []*orm.DataSeries, exs *orm.ExSymbol) error { +// aggregateHourly emits only complete UTC buckets with exactly one bar per hour. +func aggregateHourly(rows []*orm.SeriesOHLCV, timeframe string) []*orm.SeriesOHLCV { + const hourMS int64 = 60 * 60 * 1000 + hours := 8 + if timeframe == "1d" { + hours = 24 + } + stepMS := int64(hours) * hourMS + ordered := append([]*orm.SeriesOHLCV(nil), rows...) + sort.Slice(ordered, func(i, j int) bool { return ordered[i].Time < ordered[j].Time }) + var output []*orm.SeriesOHLCV + for first := 0; first < len(ordered); { + bucket := ordered[first].Time / stepMS * stepMS + end := first + 1 + for end < len(ordered) && ordered[end].Time/stepMS*stepMS == bucket { + end++ + } + complete := end-first == hours + for i := first; complete && i < end; i++ { + complete = ordered[i].Time == bucket+int64(i-first)*hourMS + } + if complete { + bar := *ordered[first] + bar.Time, bar.EndMS, bar.TimeFrame, bar.Closed = bucket, bucket+stepMS, timeframe, true + for _, row := range ordered[first+1 : end] { + bar.High = max(bar.High, row.High) + bar.Low = min(bar.Low, row.Low) + bar.Close = row.Close + bar.Volume += row.Volume + bar.Quote += row.Quote + bar.BuyVolume += row.BuyVolume + bar.TradeNum += row.TradeNum + } + output = append(output, &bar) + } + first = end + } + return output +} + +func writeCSV(path string, rows []*orm.SeriesOHLCV) error { file, err := os.Create(path) if err != nil { return err @@ -89,11 +139,7 @@ func writeCSV(path string, rows []*orm.DataSeries, exs *orm.ExSymbol) error { if err := writer.Write(columns); err != nil { return err } - for _, row := range rows { - bar, err := row.OHLCV(exs) - if err != nil { - return err - } + for _, bar := range rows { values := []string{ strconv.FormatInt(bar.Time, 10), strconv.FormatFloat(bar.Open, 'f', -1, 64), diff --git a/cmd/cbpred-export/main_test.go b/cmd/cbpred-export/main_test.go index 6b77ff0..a10b88e 100644 --- a/cmd/cbpred-export/main_test.go +++ b/cmd/cbpred-export/main_test.go @@ -1,8 +1,13 @@ package main import ( + "encoding/csv" + "os" + "path/filepath" + "reflect" "testing" + "github.com/banbox/banbot/orm" "github.com/banbox/banexg/errs" ) @@ -14,3 +19,93 @@ func TestIsUnavailablePair(t *testing.T) { t.Fatal("non-market errors must not be skipped") } } + +const testHourMS int64 = 3600000 + +func hourlyBars(startHour, count int) []*orm.SeriesOHLCV { + rows := make([]*orm.SeriesOHLCV, 0, count) + for i := startHour; i < startHour+count; i++ { + rows = append(rows, &orm.SeriesOHLCV{ + Time: int64(i) * testHourMS, EndMS: int64(i+1) * testHourMS, + Open: float64(i + 10), High: float64(i + 20), Low: float64(i), Close: float64(i + 12), + Volume: 2, Quote: 3, BuyVolume: 1, TradeNum: 4, + }) + } + return rows +} + +func TestAggregateHourlyDailyAndCSV(t *testing.T) { + // Reversed inputs still produce UTC-aligned daily history in time order. + rows := hourlyBars(0, 72) + for i, j := 0, len(rows)-1; i < j; i, j = i+1, j-1 { + rows[i], rows[j] = rows[j], rows[i] + } + bars := aggregateHourly(rows, "1d") + if len(bars) != 3 { + t.Fatalf("daily bars = %d, want 3", len(bars)) + } + for i, bar := range bars { + if bar.Time != int64(i*24)*testHourMS || bar.EndMS != int64((i+1)*24)*testHourMS || bar.TimeFrame != "1d" || !bar.Closed { + t.Fatalf("unexpected daily metadata: %+v", bar) + } + } + path := filepath.Join(t.TempDir(), "daily.csv") + if err := writeCSV(path, bars); err != nil { + t.Fatal(err) + } + file, err := os.Open(path) + if err != nil { + t.Fatal(err) + } + defer file.Close() + contents, err := csv.NewReader(file).ReadAll() + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(contents[0], columns) { + t.Fatalf("unexpected CSV columns: %v", contents[0]) + } + want := []string{"0", "10", "43", "0", "35", "48", "72", "24", "96"} + if !reflect.DeepEqual(contents[1], want) { + t.Fatalf("CSV first bar = %v, want %v", contents[1], want) + } +} + +func TestAggregateHourlyEightHour(t *testing.T) { + bars := aggregateHourly(hourlyBars(0, 24), "8h") + if len(bars) != 3 { + t.Fatalf("eight-hour bars = %d, want 3", len(bars)) + } + for i, bar := range bars { + if bar.Time != int64(i*8)*testHourMS || bar.Open != float64(i*8+10) || bar.Close != float64(i*8+19) || bar.Volume != 16 || bar.Quote != 24 || bar.BuyVolume != 8 || bar.TradeNum != 32 { + t.Fatalf("unexpected eight-hour bar: %+v", bar) + } + } +} + +func TestAggregateHourlyRejectsIncompleteBuckets(t *testing.T) { + missing := hourlyBars(0, 48) + missing = append(missing[:12:12], missing[13:]...) + gapWithDuplicate := hourlyBars(0, 48) + gapWithDuplicate[12] = gapWithDuplicate[11] + duplicate := hourlyBars(0, 48) + duplicate = append(duplicate, duplicate[12]) + for name, rows := range map[string][]*orm.SeriesOHLCV{ + "missing hour": missing, + "gap with duplicate": gapWithDuplicate, + "extra duplicate": duplicate, + "partial start and end": hourlyBars(1, 70), + "misaligned hour": append(hourlyBars(0, 23), &orm.SeriesOHLCV{Time: 23*testHourMS + 1}), + } { + t.Run(name, func(t *testing.T) { + bars := aggregateHourly(rows, "1d") + want := 1 + if name == "misaligned hour" { + want = 0 + } + if len(bars) != want || (want == 1 && bars[0].Time != 24*testHourMS) { + t.Fatalf("unexpected complete buckets: %+v", bars) + } + }) + } +} diff --git a/docs/usage.md b/docs/usage.md index 44b7bca..fc94586 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -9,9 +9,11 @@ - Python 3 和本项目的虚拟环境； - `pandas`、`numpy`、`scikit-learn`、`joblib` 等项目依赖； - Go 和 `/data/ban/banbot` 源码； -- banbot 配置中的 QuestDB 数据库，以及对 Binance 数据的访问。 +- banbot 配置中的 TimescaleDB（优先）或 QuestDB 数据库，以及对 Binance 数据的访问。 -默认配置文件是 `/data/ban/banbot/biz/config.yml`。如果路径不同，可通过 `--banbot-config` 和 `--banbot-dir` 覆盖。也可以直接传入 `--dsn`。 +一键脚本优先使用 `config/banbot.local.yml`，其次使用 banbot 同级 `data/config.timescale.yml` 或 `biz/config.timescale.yml` 中的 TimescaleDB 配置，最后使用 `biz/config.yml`。本机 banbot 目录会自动查找，也可通过 `--banbot-dir`、`--banbot-config`，或 `CBPRED_BANBOT_DIR`、`CBPRED_BANBOT_CONFIG` 环境变量指定。 + +`--dsn`（或 `CBPRED_DSN`）覆盖数据库连接；`--db-type timescale`（或 `CBPRED_DB_TYPE=timescale`）明确使用 TimescaleDB。自动检测优先选择 `exsymbol` 表对应的 TimescaleDB。Python 筛选、Go 下载和导出使用同一连接和数据库类型；连接失败会报错。只有 QuestDB 配置会尝试自动启动 QuestDB。 安装： @@ -28,14 +30,23 @@ python3 -m venv .venv 脚本调用 `python -m cbpred.latest_predict`，按以下顺序执行： -1. 检查配置中的 QuestDB 是否可连接；QuestDB 默认不可用时，通过 banbot 的导出命令触发其启动流程。 +1. 检查所选数据库是否可连接；TimescaleDB 不可用时直接报错，QuestDB 配置可通过 banbot 的导出命令触发其启动流程。 2. 从数据库筛选 Binance `linear` 市场、结算币为 USDT、上市至少 730 天且未退市的品种。 3. 调用 banbot 下载最近 36 个月的 1H K 线。 -4. 调用 banbot 导出请求周期的完整 K 线；默认是 8H，也可以通过 `--timeframe 1d` 导出日线。 +4. 调用 banbot 读取 1H 历史，并按 UTC 边界聚合请求周期的完整 K 线；默认是 8H，也可以通过 `--timeframe 1d` 导出日线。缺少小时数据的周期会被排除。 5. 写入 `data/manifest.json` 和 `data/raw/*.csv`。 6. 加载 `artifacts/v1/model.joblib`，对清单中的每个品种使用最新完整 8H K 线预测。 7. 按最新 8H USDT 成交额降序写入 `predictions/latest_v1_8h/predictions.csv`，并写入 `report.json`。 +分别生成 8H 和日线预测： + +```bash +./scripts/update_and_predict.sh --timeframe 8h --db-type timescale +./scripts/update_and_predict.sh --timeframe 1d --db-type timescale +``` + +两个 CSV 分别位于 `predictions/latest_v1_8h/predictions.csv` 和 `predictions/latest_v1_1d/predictions.csv`。其中 `price_level_reached_80pct` 是预测的 80% 命中位置，`price_level_cap_80pct` 是回调上限参考位置。 + 脚本的默认参数可以覆盖： ```bash diff --git a/go.mod b/go.mod index a6c94e3..234e920 100644 --- a/go.mod +++ b/go.mod @@ -111,4 +111,4 @@ require ( modernc.org/sqlite v1.39.1 // indirect ) -replace github.com/banbox/banbot => /data/ban/banbot +replace github.com/banbox/banbot => ../../ban/banbot diff --git a/scripts/update_and_predict.sh b/scripts/update_and_predict.sh old mode 100755 new mode 100644 index a152786..4a66805 --- a/scripts/update_and_predict.sh +++ b/scripts/update_and_predict.sh @@ -3,6 +3,9 @@ set -euo pipefail REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" PYTHON_BIN="$REPO_ROOT/.venv/bin/python" +if [[ ! -x "$PYTHON_BIN" && -x "$REPO_ROOT/.venv/Scripts/python.exe" ]]; then + PYTHON_BIN="$REPO_ROOT/.venv/Scripts/python.exe" +fi if [[ ! -x "$PYTHON_BIN" ]]; then printf 'missing virtualenv Python: %s\n' "$PYTHON_BIN" >&2 diff --git a/tests/test_latest_predict.py b/tests/test_latest_predict.py index be6e016..7a54fdf 100644 --- a/tests/test_latest_predict.py +++ b/tests/test_latest_predict.py @@ -1,4 +1,7 @@ import json +import os +import subprocess +import sys import tempfile import unittest from pathlib import Path @@ -9,15 +12,129 @@ import pandas as pd from cbpred.latest_predict import ( OUTPUT_COLUMNS, REPO_ROOT, + _acquire_update_lock, + _database_endpoint, + build_parser, build_predictions, default_artifact, + default_banbot_config, + default_banbot_dir, default_data_dir, default_output_dir, + ensure_database, output_columns_for_timeframe, ) class LatestPredictTests(unittest.TestCase): + def test_database_url_without_port_uses_postgresql_default(self): + self.assertEqual(_database_endpoint("postgresql://localhost/qdb"), ("localhost", 5432)) + + def test_database_endpoint_supports_libpq_connection_formats(self): + for dsn in ( + "host=timescale.example port=5433 dbname=qdb", + "postgresql:///qdb?host=timescale.example&port=5433", + "host=timescale.example,backup.example port=5433,5434 dbname=qdb", + ): + self.assertEqual(_database_endpoint(dsn), ("timescale.example", 5433)) + + def test_banbot_defaults_support_checkout_at_drive_root(self): + with patch.dict(os.environ, {}, clear=True), patch( + "cbpred.latest_predict.REPO_ROOT", Path(Path.cwd().anchor) / "cbpred" + ): + self.assertIsInstance(default_banbot_dir(), Path) + + def test_default_config_prefers_timescale_and_respects_explicit_environment(self): + with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ, {}, clear=True): + root = Path(directory) + banbot_dir = root / "banbot" + (banbot_dir / "biz").mkdir(parents=True) + (root / "data").mkdir() + legacy = banbot_dir / "biz/config.yml" + legacy.write_text("database: {}\n") + timescale = root / "data/config.timescale.yml" + timescale.write_text("database: {db_type: timescale}\n") + with patch("cbpred.latest_predict.REPO_ROOT", root / "cbpred"): + self.assertEqual(default_banbot_config(banbot_dir), timescale) + with patch.dict(os.environ, {"CBPRED_BANBOT_CONFIG": str(legacy)}): + self.assertEqual(default_banbot_config(banbot_dir), legacy) + + def test_parser_reads_database_environment(self): + with patch.dict(os.environ, {"CBPRED_DSN": "postgresql://localhost/qdb", "CBPRED_DB_TYPE": "timescale"}): + args = build_parser().parse_args([]) + self.assertEqual(args.dsn, "postgresql://localhost/qdb") + self.assertEqual(args.db_type, "timescale") + + def test_unavailable_timescale_never_bootstraps_questdb(self): + with patch("cbpred.latest_predict.socket.create_connection", side_effect=OSError), patch( + "cbpred.latest_predict.subprocess.run" + ) as bootstrap: + ensure_database(Path("unused.yml"), "postgresql://localhost:8812/qdb", "timescale") + bootstrap.assert_not_called() + + def test_configured_timescale_never_bootstraps_questdb_in_auto_mode(self): + with tempfile.TemporaryDirectory() as directory: + config = Path(directory) / "config.yml" + config.write_text("database:\n db_type: timescale\n url: postgresql://localhost:8812/qdb\n") + with patch("cbpred.latest_predict.socket.create_connection", side_effect=OSError), patch( + "cbpred.latest_predict.subprocess.run" + ) as bootstrap: + ensure_database(config, None) + bootstrap.assert_not_called() + + def test_postgresql_backend_aliases_never_bootstrap_questdb(self): + with patch("cbpred.latest_predict.socket.create_connection", side_effect=OSError), patch( + "cbpred.latest_predict.subprocess.run" + ) as bootstrap: + for backend in ("timescaledb", "postgres", "postgresql", "postgresql+timescale"): + ensure_database(Path("unused.yml"), "postgresql://localhost:8812/qdb", backend) + bootstrap.assert_not_called() + + def test_multihost_dsn_is_validated_by_libpq_without_bootstrap(self): + with patch("cbpred.latest_predict.socket.create_connection") as probe, patch( + "cbpred.latest_predict.subprocess.run" + ) as bootstrap: + for backend in ("auto", "timescale"): + ensure_database( + Path("unused.yml"), + "host=primary,secondary port=8812,5433 dbname=qdb", + backend, + ) + probe.assert_not_called() + bootstrap.assert_not_called() + + def test_update_lock_serializes_processes_and_releases_on_close(self): + with tempfile.TemporaryDirectory() as directory: + lock_path = Path(directory) / "update.lock" + script = ( + "import sys\n" + "from cbpred.latest_predict import _acquire_update_lock\n" + "with open(sys.argv[1], 'a+') as lock:\n" + " print('ready', flush=True)\n" + " _acquire_update_lock(lock)\n" + " print('acquired', flush=True)\n" + ) + with lock_path.open("a+") as lock: + _acquire_update_lock(lock) + process = subprocess.Popen( + [sys.executable, "-c", script, str(lock_path)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + self.assertEqual(process.stdout.readline().strip(), "ready") + with self.assertRaises(subprocess.TimeoutExpired): + process.communicate(timeout=0.3) + lock.close() + stdout, stderr = process.communicate(timeout=10) + self.assertEqual(process.returncode, 0, stderr) + self.assertEqual(stdout.strip(), "acquired") + finally: + if process.poll() is None: + process.kill() + process.communicate() + def test_output_columns_are_stable(self): self.assertEqual( OUTPUT_COLUMNS, diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 5fe4bb5..e4b2173 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -2,13 +2,15 @@ from __future__ import annotations import json import subprocess +import sys import tempfile import unittest from pathlib import Path -from unittest.mock import patch +from unittest.mock import MagicMock, patch import numpy as np import pandas as pd +import yaml from cbpred.pipeline import ( CSV_COLUMNS, @@ -25,6 +27,9 @@ from cbpred.pipeline import ( run_training, Symbol, run_banbot_download, + detect_db_type, + effective_banbot_config, + prepare_data, ) @@ -44,6 +49,99 @@ def make_hour_rows(start_ms: int, hours: int, phase: float = 0.0): class PipelineTests(unittest.TestCase): + def test_download_reads_utf8_stderr_on_windows(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + downloader = root / "downloader.py" + downloader.write_text( + "import sys\n" + "sys.stderr.buffer.write('\\u4e00'.encode('utf-8'))\n", + encoding="utf-8", + ) + symbol = Symbol(1, "binance", "linear", "BTC/USDT:USDT", 1, 0) + downloaded, unavailable = run_banbot_download( + [symbol], root, root / "config.yml", 1, 2, + banbot_cmd=f'"{sys.executable}" "{downloader}"', + ) + self.assertEqual(downloaded, [symbol]) + self.assertEqual(unavailable, []) + + def test_auto_database_prefers_timescale_when_both_tables_exist(self): + conn = MagicMock() + cur = conn.cursor.return_value.__enter__.return_value + cur.fetchone.return_value = ("exsymbol", "exsymbol_q") + self.assertEqual(detect_db_type(conn), "timescale") + cur.execute.assert_called_once_with("SELECT to_regclass('exsymbol'), to_regclass('exsymbol_q')") + + def test_auto_database_detects_questdb_postgres_tables(self): + conn = MagicMock() + conn.cursor.return_value.__enter__.return_value.fetchone.return_value = (None, "exsymbol_q") + self.assertEqual(detect_db_type(conn), "questdb") + + def test_auto_database_falls_back_for_questdb_without_to_regclass(self): + conn = MagicMock() + cur = conn.cursor.return_value.__enter__.return_value + cur.execute.side_effect = [ + RuntimeError("unknown function: to_regclass"), + RuntimeError("table does not exist [table=exsymbol]"), + None, + ] + self.assertEqual(detect_db_type(conn), "questdb") + self.assertEqual(conn.rollback.call_count, 2) + + def test_auto_database_does_not_hide_permission_errors(self): + conn = MagicMock() + conn.cursor.return_value.__enter__.return_value.execute.side_effect = RuntimeError("permission denied") + with self.assertRaisesRegex(RuntimeError, "permission denied"): + detect_db_type(conn) + + def test_effective_config_is_removed_after_failure_and_original_unchanged(self): + with tempfile.TemporaryDirectory() as directory: + config_path = Path(directory) / "config.yml" + original = "database:\n url: original\n db_type: questdb\nassets: ./assets\n" + config_path.write_text(original) + with self.assertRaisesRegex(RuntimeError, "failed"): + with effective_banbot_config(config_path, "selected", "timescale") as effective: + self.assertEqual(effective.parent, config_path.parent) + self.assertTrue(effective.exists()) + raise RuntimeError("failed") + self.assertFalse(effective.exists()) + self.assertEqual(config_path.read_text(), original) + + def test_prepare_propagates_selected_database_to_download_and_export(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + config_path = root / "config.yml" + original = "database:\n url: original\n db_type: questdb\n max_conn: 4\nmarket_type: spot\nassets: ./assets\n" + config_path.write_text(original) + symbol = Symbol(1, "binance", "linear", "BTC/USDT:USDT", 1, 0) + commands = [] + config_paths = [] + + def run(command, **kwargs): + commands.append(command) + path = Path(command[command.index("--config") + 1]) + config_paths.append(path) + effective = yaml.safe_load(path.read_text(encoding="utf-8")) + self.assertEqual(effective["database"], {"url": "selected-secret", "db_type": "timescale", "max_conn": 4}) + self.assertEqual(effective["assets"], "./assets") + self.assertEqual(effective["market_type"], "linear") + self.assertEqual(path.parent, config_path.parent) + self.assertNotIn("selected-secret", command) + + conn = MagicMock() + with patch("cbpred.pipeline.connect_database", return_value=conn) as connect, \ + patch("cbpred.pipeline.detect_db_type", return_value="timescale"), \ + patch("cbpred.pipeline.fetch_symbols", return_value=([symbol], [])), \ + patch("cbpred.pipeline.subprocess.run", side_effect=run): + manifest = prepare_data(root / "data", root, config_path, dsn="selected-secret", end_ms=EIGHT_H_MS * 1000) + connect.assert_called_once_with("selected-secret") + self.assertEqual(len(commands), 2) + self.assertEqual(config_paths[0], config_paths[1]) + self.assertFalse(config_paths[0].exists()) + self.assertEqual(config_path.read_text(), original) + self.assertEqual(manifest["db_type"], "timescale") + def test_download_skips_invalid_current_symbol_and_retries_chunk(self): valid = Symbol(1, "binance", "linear", "BTC/USDT:USDT", 1, 0) stale = Symbol(2, "binance", "linear", "SXP/USDT:USDT", 1, 0)