103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
"""
|
|
DB-related helpers for device discovery:
|
|
- Auto-detect the CAPS DB connection from a remote config file.
|
|
- Fetch rack host IPs from the CAPS PostgreSQL database.
|
|
"""
|
|
|
|
import logging
|
|
|
|
import psycopg
|
|
|
|
from .config_parser import parse_db_uri
|
|
from .ssh_client import ssh_run_multi, open_ssh_tunnel
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def autodetect_db_from_server(
|
|
server_ip: str,
|
|
ssh_port: int,
|
|
ssh_options: list,
|
|
server_config_path: str = "/etc/caps/config.ini",
|
|
) -> dict | None:
|
|
"""SSH into server, read its caps config, extract DB connection params.
|
|
|
|
Returns a dict with keys: host, port, user, password, dbname.
|
|
Returns None if no db_uri is found.
|
|
"""
|
|
from configparser import RawConfigParser
|
|
|
|
try:
|
|
config_text = await ssh_run_multi(
|
|
server_ip, ssh_port, ssh_options, f"cat {server_config_path}"
|
|
)
|
|
except Exception:
|
|
return None
|
|
|
|
parser = RawConfigParser(strict=False)
|
|
parser.read_string(config_text)
|
|
|
|
for section in list(parser.sections()) + ["DEFAULT"]:
|
|
try:
|
|
db_uri = parser.get(section, "db_uri", fallback=None)
|
|
if db_uri and db_uri.startswith("postgresql"):
|
|
return parse_db_uri(db_uri)
|
|
except Exception:
|
|
continue
|
|
|
|
db_uri = parser.defaults().get("db_uri")
|
|
if db_uri and db_uri.startswith("postgresql"):
|
|
return parse_db_uri(db_uri)
|
|
|
|
return None
|
|
|
|
|
|
async def get_rack_hosts(
|
|
host: str,
|
|
port: int,
|
|
user: str,
|
|
password: str,
|
|
dbname: str,
|
|
table: str,
|
|
via_tunnel: bool = False,
|
|
tunnel_host: str | None = None,
|
|
tunnel_ssh_port: int = 22,
|
|
ssh_options: list | None = None,
|
|
) -> list[tuple[str, str, str, str | None]]:
|
|
"""Return list of (external_id, host_ip, rack_name, rack_type) from the racks table.
|
|
|
|
Supports an optional SSH tunnel when the DB is not directly reachable.
|
|
"""
|
|
if via_tunnel and tunnel_host and ssh_options:
|
|
conn = await open_ssh_tunnel(tunnel_host, tunnel_ssh_port, ssh_options, host, port)
|
|
async with conn:
|
|
listener = await conn.forward_local_port("127.0.0.1", 0, host, port)
|
|
local_port = listener.get_port()
|
|
dsn = (
|
|
f"host=127.0.0.1 port={local_port} dbname={dbname} "
|
|
f"user={user} password={password} connect_timeout=10"
|
|
)
|
|
return await _fetch_rack_rows(dsn, table)
|
|
else:
|
|
dsn = f"host={host} port={port} dbname={dbname} user={user} password={password} connect_timeout=10"
|
|
return await _fetch_rack_rows(dsn, table)
|
|
|
|
|
|
async def _fetch_rack_rows(dsn: str, table: str) -> list[tuple[str, str, str, str | None]]:
|
|
async with await psycopg.AsyncConnection.connect(dsn) as conn:
|
|
cursor = await conn.execute(
|
|
f"SELECT id, data->>'host' AS host, name, data->>'rack_type' AS rack_type "
|
|
f"FROM \"{table}\" WHERE data->>'host' IS NOT NULL"
|
|
)
|
|
rows = await cursor.fetchall()
|
|
|
|
result = []
|
|
for r in rows:
|
|
if not r[1]:
|
|
continue
|
|
ext_id = str(r[0])
|
|
short_name = str(r[2]) if r[2] else None
|
|
# Combine CAPS id with short name: "lane-1-1 (01)"
|
|
rack_name = f"{ext_id} ({short_name})" if short_name and short_name != ext_id else ext_id
|
|
result.append((ext_id, str(r[1]), rack_name, r[3] or None))
|
|
return result
|