416 lines
15 KiB
Python
416 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
# /// script
|
|
# requires-python = ">=3.11"
|
|
# dependencies = []
|
|
# ///
|
|
"""flight-search skill — CLI for managing flight searches and results in SQLite."""
|
|
|
|
import argparse
|
|
import json
|
|
import sqlite3
|
|
import sys
|
|
from datetime import date, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
DB_PATH = (
|
|
Path(__file__).resolve().parent.parent.parent.parent / "db" / "flight_search.sqlite"
|
|
)
|
|
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS searches (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
origin TEXT NOT NULL,
|
|
destination TEXT NOT NULL,
|
|
dep_date TEXT NOT NULL,
|
|
ret_date TEXT NOT NULL,
|
|
adults INTEGER NOT NULL DEFAULT 1,
|
|
max_layovers INTEGER NOT NULL DEFAULT 1,
|
|
min_stay_days INTEGER NOT NULL DEFAULT 1,
|
|
flex_days INTEGER NOT NULL DEFAULT 0,
|
|
currency TEXT NOT NULL DEFAULT 'CZK',
|
|
status TEXT NOT NULL DEFAULT 'active',
|
|
created_at TEXT NOT NULL,
|
|
completed_at TEXT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS results (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
search_id INTEGER NOT NULL REFERENCES searches(id) ON DELETE CASCADE,
|
|
airline TEXT NOT NULL,
|
|
route TEXT NOT NULL,
|
|
dep_date TEXT NOT NULL,
|
|
ret_date TEXT NOT NULL,
|
|
dep_time TEXT,
|
|
arr_time TEXT,
|
|
layovers INTEGER NOT NULL DEFAULT 0,
|
|
layover_info TEXT,
|
|
duration TEXT,
|
|
price REAL,
|
|
price_currency TEXT NOT NULL DEFAULT 'CZK',
|
|
price_czk REAL NOT NULL,
|
|
booking_url TEXT,
|
|
found_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_results_search ON results(search_id);
|
|
CREATE INDEX IF NOT EXISTS idx_results_route ON results(route);
|
|
"""
|
|
|
|
|
|
def _connect() -> sqlite3.Connection:
|
|
DB_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
conn = sqlite3.connect(DB_PATH)
|
|
conn.row_factory = sqlite3.Row
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
conn.execute("PRAGMA foreign_keys=ON")
|
|
conn.executescript(SCHEMA)
|
|
return conn
|
|
|
|
|
|
def _now() -> str:
|
|
from datetime import datetime, timezone
|
|
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
# ── create-search ──────────────────────────────────────────────────────────
|
|
|
|
|
|
def cmd_create_search(args: argparse.Namespace) -> None:
|
|
dep = date.fromisoformat(args.dep_date)
|
|
ret = date.fromisoformat(args.ret_date)
|
|
flex = args.flex_days
|
|
min_stay = args.min_stay_days
|
|
adults = args.adults
|
|
origin = args.origin.upper()
|
|
destination = args.destination.upper()
|
|
currency = args.currency.upper()
|
|
|
|
# Generate date combinations respecting min_stay
|
|
dep_range = [dep + timedelta(days=d) for d in range(-flex, flex + 1)]
|
|
ret_range = [ret + timedelta(days=d) for d in range(-flex, flex + 1)]
|
|
|
|
combos = []
|
|
for d in dep_range:
|
|
for r in ret_range:
|
|
if (r - d).days >= min_stay:
|
|
combos.append((d, r))
|
|
|
|
# Sort by proximity to original dates, then by total stay duration
|
|
combos.sort(key=lambda dr: (abs((dr[0] - dep).days) + abs((dr[1] - ret).days), (dr[1] - dr[0]).days))
|
|
|
|
# Cap at 15 URLs
|
|
combos = combos[:15]
|
|
|
|
urls = []
|
|
for d, r in combos:
|
|
url = (
|
|
f"https://www.cz.kayak.com/flights/{origin}-{destination}"
|
|
f"/{d.strftime('%Y-%m-%d')}/{r.strftime('%Y-%m-%d')}"
|
|
f"/{adults}adults?sort=price_a"
|
|
)
|
|
if args.max_layovers is not None:
|
|
url += f"&fs=stops=-{args.max_layovers + 1}"
|
|
urls.append(url)
|
|
|
|
conn = _connect()
|
|
try:
|
|
cur = conn.execute(
|
|
"""INSERT INTO searches
|
|
(origin, destination, dep_date, ret_date, adults, max_layovers,
|
|
min_stay_days, flex_days, currency, status, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'active', ?)""",
|
|
(
|
|
origin,
|
|
destination,
|
|
args.dep_date,
|
|
args.ret_date,
|
|
adults,
|
|
args.max_layovers if args.max_layovers is not None else 1,
|
|
min_stay,
|
|
flex,
|
|
currency,
|
|
_now(),
|
|
),
|
|
)
|
|
search_id = cur.lastrowid
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
output = {
|
|
"search_id": search_id,
|
|
"origin": origin,
|
|
"destination": destination,
|
|
"date_combinations": len(combos),
|
|
"urls": urls,
|
|
}
|
|
print(json.dumps(output, indent=2, ensure_ascii=False))
|
|
|
|
|
|
# ── add-result ──────────────────────────────────────────────────────────────
|
|
|
|
|
|
def cmd_add_result(args: argparse.Namespace) -> None:
|
|
conn = _connect()
|
|
try:
|
|
conn.execute(
|
|
"""INSERT INTO results
|
|
(search_id, airline, route, dep_date, ret_date, dep_time, arr_time,
|
|
layovers, layover_info, duration, price, price_currency, price_czk,
|
|
booking_url, found_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
|
(
|
|
args.search_id,
|
|
args.airline,
|
|
args.route,
|
|
args.dep_date,
|
|
args.ret_date,
|
|
args.dep_time,
|
|
args.arr_time,
|
|
args.layovers if args.layovers is not None else 0,
|
|
args.layover_info,
|
|
args.duration,
|
|
args.price,
|
|
args.price_currency if args.price_currency else "CZK",
|
|
args.price_czk if args.price_czk is not None else args.price,
|
|
args.booking_url,
|
|
_now(),
|
|
),
|
|
)
|
|
conn.commit()
|
|
print(json.dumps({"ok": True, "search_id": args.search_id}))
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ── add-results (batch from JSON) ──────────────────────────────────────────
|
|
|
|
|
|
def cmd_add_results(args: argparse.Namespace) -> None:
|
|
if args.json_file:
|
|
with open(args.json_file) as f:
|
|
data = json.load(f)
|
|
else:
|
|
data = json.load(sys.stdin)
|
|
|
|
if not isinstance(data, list):
|
|
data = [data]
|
|
|
|
conn = _connect()
|
|
try:
|
|
count = 0
|
|
for row in data:
|
|
conn.execute(
|
|
"""INSERT INTO results
|
|
(search_id, airline, route, dep_date, ret_date, dep_time, arr_time,
|
|
layovers, layover_info, duration, price, price_currency, price_czk,
|
|
booking_url, found_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
|
(
|
|
row.get("search_id"),
|
|
row.get("airline"),
|
|
row.get("route"),
|
|
row.get("dep_date"),
|
|
row.get("ret_date"),
|
|
row.get("dep_time"),
|
|
row.get("arr_time"),
|
|
row.get("layovers", 0),
|
|
row.get("layover_info"),
|
|
row.get("duration"),
|
|
row.get("price"),
|
|
row.get("price_currency", "CZK"),
|
|
row.get("price_czk", row.get("price")),
|
|
row.get("booking_url"),
|
|
_now(),
|
|
),
|
|
)
|
|
count += 1
|
|
conn.commit()
|
|
print(json.dumps({"ok": True, "added": count}))
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ── results ─────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def cmd_results(args: argparse.Namespace) -> None:
|
|
conn = _connect()
|
|
try:
|
|
rows = conn.execute(
|
|
"""SELECT r.*, s.origin, s.destination, s.adults
|
|
FROM results r
|
|
JOIN searches s ON s.id = r.search_id
|
|
WHERE r.search_id = ?
|
|
ORDER BY r.price_czk ASC
|
|
LIMIT ?""",
|
|
(args.search_id, args.top or 5),
|
|
).fetchall()
|
|
|
|
if not rows:
|
|
print("(no results found)")
|
|
return
|
|
|
|
for i, row in enumerate(rows, 1):
|
|
layover = ""
|
|
if row["layover_info"]:
|
|
layover = f" ({row['layover_info']})"
|
|
stops = f"{row['layovers']} stop{'' if row['layovers'] == 1 else 's'}"
|
|
price_str = f"{row['price_czk']:,.0f} {row['price_currency']}"
|
|
if row["price_currency"] != "CZK" and row["price"] != row["price_czk"]:
|
|
price_str = f"{row['price']:,.0f} {row['price_currency']} → {row['price_czk']:,.0f} CZK"
|
|
|
|
print(
|
|
f"#{i} {row['airline']} {row['route']} "
|
|
f"{row['dep_date']} → {row['ret_date']} "
|
|
f"{row['duration'] or '?'} {stops}{layover} {price_str}"
|
|
)
|
|
if row["booking_url"]:
|
|
print(f" {row['booking_url']}")
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ── list-searches ───────────────────────────────────────────────────────────
|
|
|
|
|
|
def cmd_list_searches(args: argparse.Namespace) -> None:
|
|
conn = _connect()
|
|
try:
|
|
status_filter = args.status or "active"
|
|
if status_filter == "all":
|
|
rows = conn.execute(
|
|
"SELECT * FROM searches ORDER BY created_at DESC"
|
|
).fetchall()
|
|
else:
|
|
rows = conn.execute(
|
|
"SELECT * FROM searches WHERE status = ? ORDER BY created_at DESC",
|
|
(status_filter,),
|
|
).fetchall()
|
|
|
|
if not rows:
|
|
print("(no searches)")
|
|
return
|
|
|
|
for row in rows:
|
|
result_count = conn.execute(
|
|
"SELECT COUNT(*) FROM results WHERE search_id = ?", (row["id"],)
|
|
).fetchone()[0]
|
|
print(
|
|
f"#{row['id']} {row['origin']}→{row['destination']} "
|
|
f"{row['dep_date']} → {row['ret_date']} "
|
|
f"{row['adults']} adults flex ±{row['flex_days']}d "
|
|
f"min stay {row['min_stay_days']}d "
|
|
f"max {row['max_layovers']} stop(s) "
|
|
f"[{row['status']}] {result_count} results"
|
|
)
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ── delete-search ───────────────────────────────────────────────────────────
|
|
|
|
|
|
def cmd_delete_search(args: argparse.Namespace) -> None:
|
|
conn = _connect()
|
|
try:
|
|
# Check exists
|
|
row = conn.execute(
|
|
"SELECT id, origin, destination FROM searches WHERE id = ?",
|
|
(args.search_id,),
|
|
).fetchone()
|
|
if not row:
|
|
print(json.dumps({"error": f"search {args.search_id} not found"}))
|
|
sys.exit(1)
|
|
|
|
result_count = conn.execute(
|
|
"SELECT COUNT(*) FROM results WHERE search_id = ?", (args.search_id,)
|
|
).fetchone()[0]
|
|
|
|
conn.execute("DELETE FROM results WHERE search_id = ?", (args.search_id,))
|
|
conn.execute("DELETE FROM searches WHERE id = ?", (args.search_id,))
|
|
conn.commit()
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"deleted": True,
|
|
"search_id": args.search_id,
|
|
"route": f"{row['origin']}→{row['destination']}",
|
|
"results_removed": result_count,
|
|
}
|
|
)
|
|
)
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ── main ────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Flight search CLI")
|
|
sub = parser.add_subparsers(dest="command", required=True)
|
|
|
|
# create-search
|
|
p_create = sub.add_parser("create-search", help="Create a new flight search")
|
|
p_create.add_argument("--origin", required=True, help="Origin IATA code (e.g. PRG)")
|
|
p_create.add_argument("--destination", required=True, help="Destination IATA code (e.g. PTY)")
|
|
p_create.add_argument("--dep-date", required=True, help="Departure date YYYY-MM-DD")
|
|
p_create.add_argument("--ret-date", required=True, help="Return date YYYY-MM-DD")
|
|
p_create.add_argument("--adults", type=int, default=1, help="Number of adults (default 1)")
|
|
p_create.add_argument("--max-layovers", type=int, default=None, help="Max additional stops (default 1)")
|
|
p_create.add_argument("--min-stay-days", type=int, default=1, help="Minimum stay in days (default 1)")
|
|
p_create.add_argument("--flex-days", type=int, default=0, help="±days for flexible dates (default 0)")
|
|
p_create.add_argument("--currency", default="CZK", help="Currency code (default CZK)")
|
|
|
|
# add-result
|
|
p_add = sub.add_parser("add-result", help="Add a single flight result")
|
|
p_add.add_argument("--search-id", type=int, required=True)
|
|
p_add.add_argument("--airline", required=True)
|
|
p_add.add_argument("--route", required=True, help="e.g. PRG→AMS→PTY / PTY→AMS→PRG")
|
|
p_add.add_argument("--dep-date", required=True)
|
|
p_add.add_argument("--ret-date", required=True)
|
|
p_add.add_argument("--dep-time", default=None)
|
|
p_add.add_argument("--arr-time", default=None)
|
|
p_add.add_argument("--layovers", type=int, default=None)
|
|
p_add.add_argument("--layover-info", default=None, help="e.g. AMS 2h15m")
|
|
p_add.add_argument("--duration", default=None, help="e.g. 12h30m")
|
|
p_add.add_argument("--price", type=float, required=True)
|
|
p_add.add_argument("--price-currency", default=None)
|
|
p_add.add_argument("--price-czk", type=float, default=None)
|
|
p_add.add_argument("--booking-url", default=None)
|
|
|
|
# add-results (batch)
|
|
p_batch = sub.add_parser("add-results", help="Add multiple results from JSON")
|
|
p_batch.add_argument("--json-file", default=None, help="Path to JSON file (default: stdin)")
|
|
|
|
# results
|
|
p_res = sub.add_parser("results", help="Show results for a search")
|
|
p_res.add_argument("search_id", type=int)
|
|
p_res.add_argument("--top", type=int, default=5, help="Show top N results (default 5)")
|
|
|
|
# list-searches
|
|
p_list = sub.add_parser("list-searches", help="List all searches")
|
|
p_list.add_argument("--status", default="active", help="Filter by status (active/completed/all)")
|
|
|
|
# delete-search
|
|
p_del = sub.add_parser("delete-search", help="Delete a search and its results")
|
|
p_del.add_argument("search_id", type=int)
|
|
|
|
args = parser.parse_args()
|
|
|
|
commands = {
|
|
"create-search": cmd_create_search,
|
|
"add-result": cmd_add_result,
|
|
"add-results": cmd_add_results,
|
|
"results": cmd_results,
|
|
"list-searches": cmd_list_searches,
|
|
"delete-search": cmd_delete_search,
|
|
}
|
|
commands[args.command](args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |