146 lines
5.6 KiB
Python
146 lines
5.6 KiB
Python
"""Загрузчик кварталов dataset_СДЕЛКИ через SSH-tunnel на localhost:15432.
|
||
|
||
Эквивалент 02_load_all_quarters.sh, но через psycopg вместо docker — удобнее
|
||
запускать с Windows / любой машины с уже открытым SSH-tunnel.
|
||
|
||
Idempotent: пропускает кварталы, уже залитые в rosreestr_deals.
|
||
|
||
Usage:
|
||
cd backend
|
||
uv run python ../data/sql/02b_load_via_tunnel.py 2024Q1 2024Q2
|
||
uv run python ../data/sql/02b_load_via_tunnel.py # все из JOBS
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
|
||
import psycopg
|
||
|
||
PG_DSN = os.environ.get(
|
||
"TUNNEL_DSN",
|
||
"postgresql://gendesign:2J2SBPMKuS998fiwhtQqDhMI@localhost:15432/gendesign",
|
||
)
|
||
|
||
RAW = Path(__file__).resolve().parent.parent / "raw"
|
||
|
||
# (source_quarter, period_start, csv_filename, delimiter)
|
||
JOBS: list[tuple[str, str, str, str]] = [
|
||
("2024Q1", "2024-01-01", "dataset_СДЕЛКИ_r-r_01-92_y_2024_q_1.csv", ";"),
|
||
("2024Q2", "2024-04-01", "dataset_СДЕЛКИ_r-r_01-92_y_2024_q_2.csv", ";"),
|
||
# Q3+ уже залиты, но оставляем как fallback / re-run target.
|
||
("2024Q3", "2024-07-01", "dataset_СДЕЛКИ_r_all_q_3.csv", ";"),
|
||
("2024Q4", "2024-10-01", "dataset_СДЕЛКИ_r-r_01-92_y_2024_q_4.csv", "~"),
|
||
("2025Q1", "2025-01-01", "dataset_СДЕЛКИ_r-r_01-92_y_2025_q_1.csv", "~"),
|
||
("2025Q2", "2025-04-01", "dataset_СДЕЛКИ_r-r_01-92_y_2025_q_2.csv", "~"),
|
||
("2025Q3", "2025-07-01", "dataset_СДЕЛКИ_r-r_01-92_y_2025_q_3.csv", "~"),
|
||
("2025Q4", "2025-10-01", "dataset_СДЕЛКИ_r-r_01-92_y_2025_q_4.csv", "~"),
|
||
("2026Q1", "2026-01-01", "dataset_СДЕЛКИ_r-r_01-92_y_2026_q_1.csv", "~"),
|
||
]
|
||
|
||
|
||
COPY_SQL = """
|
||
COPY rosreestr_deals_staging
|
||
FROM STDIN (FORMAT csv, HEADER true, DELIMITER %r, QUOTE '"', NULL '')
|
||
"""
|
||
|
||
INSERT_SQL = """
|
||
INSERT INTO rosreestr_deals (
|
||
source_quarter, period_start_date,
|
||
region_code, okato, district, city, quarter_cad_number, street,
|
||
realestate_type_code, wall_material_code, purpose_code,
|
||
year_build, floor, area,
|
||
doc_type, deal_price, currency, deal_count
|
||
)
|
||
SELECT
|
||
%s, period_start_date,
|
||
region_code, NULLIF(okato, '-'), NULLIF(district, ''), NULLIF(city, ''),
|
||
quarter_cad_number, NULLIF(street, ''),
|
||
NULLIF(realestate_type_code, ''),
|
||
NULLIF(wall_material_code, ''),
|
||
NULLIF(purpose_code, ''),
|
||
CASE WHEN year_build ~ '^[12][0-9]{3}$' THEN year_build::INT ELSE NULL END,
|
||
NULLIF(floor, ''),
|
||
CASE WHEN area > 0 AND area < 1e8 THEN area ELSE NULL END,
|
||
NULLIF(doc_type, ''),
|
||
CASE WHEN deal_price > 0 AND deal_price < 1e15 THEN deal_price ELSE NULL END,
|
||
COALESCE(NULLIF(currency, ''), 'рубль'),
|
||
COALESCE(deal_count, 1)
|
||
FROM rosreestr_deals_staging
|
||
WHERE region_code IS NOT NULL
|
||
AND quarter_cad_number IS NOT NULL
|
||
AND quarter_cad_number <> ''
|
||
AND doc_type IS NOT NULL
|
||
AND doc_type <> ''
|
||
AND period_start_date IS NOT NULL
|
||
"""
|
||
|
||
|
||
def main() -> int:
|
||
targets = sys.argv[1:] or [j[0] for j in JOBS]
|
||
selected = [j for j in JOBS if j[0] in targets]
|
||
if not selected:
|
||
print(f"Нет matching кварталов в JOBS. Targets: {targets}")
|
||
return 1
|
||
|
||
conn = psycopg.connect(PG_DSN, connect_timeout=15)
|
||
try:
|
||
for sq, period, fname, sep in selected:
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"SELECT COUNT(*) FROM rosreestr_deals WHERE source_quarter = %s",
|
||
(sq,),
|
||
)
|
||
already = cur.fetchone()[0]
|
||
if already:
|
||
print(f"=== {sq} SKIP — уже {already} строк (DELETE FROM rosreestr_deals WHERE source_quarter='{sq}' чтобы перезалить)")
|
||
continue
|
||
csv_path = RAW / fname
|
||
if not csv_path.exists():
|
||
print(f"=== {sq} SKIP — нет файла {csv_path}")
|
||
continue
|
||
print(f"=== {sq} ({fname}, sep='{sep}', {csv_path.stat().st_size/1e6:.1f} MB) ===")
|
||
t0 = time.time()
|
||
with conn.cursor() as cur:
|
||
cur.execute("TRUNCATE rosreestr_deals_staging")
|
||
copy_sql = (
|
||
f"COPY rosreestr_deals_staging FROM STDIN "
|
||
f"(FORMAT csv, HEADER true, DELIMITER '{sep}', QUOTE '\"', NULL '')"
|
||
)
|
||
with cur.copy(copy_sql) as copy:
|
||
with open(csv_path, "rb") as f:
|
||
while True:
|
||
chunk = f.read(8 * 1024 * 1024)
|
||
if not chunk:
|
||
break
|
||
copy.write(chunk)
|
||
t_copy = time.time() - t0
|
||
cur.execute(INSERT_SQL, (sq,))
|
||
inserted = cur.rowcount
|
||
conn.commit()
|
||
t_total = time.time() - t0
|
||
print(f" COPY {t_copy:.0f}s, INSERT {inserted} rows, total {t_total:.0f}s")
|
||
|
||
# Final summary
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"""
|
||
SELECT source_quarter, COUNT(*) AS rows, SUM(deal_count) AS deals
|
||
FROM rosreestr_deals
|
||
GROUP BY source_quarter
|
||
ORDER BY source_quarter
|
||
"""
|
||
)
|
||
print("\n=== Итоги по rosreestr_deals ===")
|
||
for sq, rows, deals in cur.fetchall():
|
||
print(f" {sq} rows={rows:>9} deals={deals:>10}")
|
||
finally:
|
||
conn.close()
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|