408 lines
13 KiB
Python
408 lines
13 KiB
Python
"""PostgreSQL persistence for the inventory domain."""
|
|
|
|
from collections.abc import Callable, Iterator, Mapping, Sequence
|
|
from contextlib import contextmanager
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
import psycopg
|
|
from psycopg.rows import dict_row
|
|
|
|
from app.config import settings
|
|
|
|
|
|
ConnectionFactory = Callable[[], psycopg.Connection]
|
|
|
|
TAG_TRIMMED_NAME_SQL = (
|
|
"regexp_replace(name, '^[[:space:]]+|[[:space:]]+$', '', 'g')"
|
|
)
|
|
|
|
MINER_SELECT = f"""
|
|
SELECT
|
|
miners.id,
|
|
miners.site_id,
|
|
miners.room_id,
|
|
miners.asset_tag,
|
|
miners.current_mac,
|
|
host(miners.current_ip) AS current_ip,
|
|
miners.hostname,
|
|
miners.miner_sn,
|
|
miners.control_board_sn,
|
|
miners.power_sn,
|
|
miners.model,
|
|
miners.firmware_version,
|
|
miners.lifecycle_status,
|
|
miners.status,
|
|
miners.notes,
|
|
miners.first_seen_at,
|
|
miners.last_seen_at,
|
|
miners.created_at,
|
|
miners.updated_at,
|
|
tag_data.tags
|
|
FROM miners
|
|
LEFT JOIN LATERAL (
|
|
SELECT COALESCE(
|
|
jsonb_agg(
|
|
jsonb_build_object(
|
|
'id', tags.id,
|
|
'name', tags.name,
|
|
'created_at', tags.created_at,
|
|
'updated_at', tags.updated_at
|
|
)
|
|
ORDER BY lower({TAG_TRIMMED_NAME_SQL}), tags.id
|
|
),
|
|
'[]'::jsonb
|
|
) AS tags
|
|
FROM miner_tags
|
|
JOIN tags ON tags.id = miner_tags.tag_id
|
|
WHERE miner_tags.miner_id = miners.id
|
|
) tag_data ON true
|
|
"""
|
|
|
|
MINER_WRITE_COLUMNS = {
|
|
"site_id",
|
|
"room_id",
|
|
"asset_tag",
|
|
"current_mac",
|
|
"current_ip",
|
|
"hostname",
|
|
"miner_sn",
|
|
"control_board_sn",
|
|
"power_sn",
|
|
"model",
|
|
"firmware_version",
|
|
"lifecycle_status",
|
|
"status",
|
|
"notes",
|
|
"first_seen_at",
|
|
"last_seen_at",
|
|
}
|
|
|
|
|
|
def default_connection_factory() -> psycopg.Connection:
|
|
return psycopg.connect(
|
|
host=settings.postgres_host,
|
|
port=settings.postgres_port,
|
|
dbname=settings.postgres_db,
|
|
user=settings.postgres_user,
|
|
password=settings.postgres_password,
|
|
connect_timeout=5,
|
|
)
|
|
|
|
|
|
class InventoryRepository:
|
|
"""Database-only inventory persistence with optional transaction binding."""
|
|
|
|
def __init__(
|
|
self,
|
|
connection_factory: ConnectionFactory | None = None,
|
|
*,
|
|
connection: psycopg.Connection | None = None,
|
|
):
|
|
self._connection_factory = connection_factory or default_connection_factory
|
|
self._connection = connection
|
|
|
|
@contextmanager
|
|
def transaction(self) -> Iterator["InventoryRepository"]:
|
|
if self._connection is not None:
|
|
yield self
|
|
return
|
|
|
|
connection = self._connection_factory()
|
|
transaction_repository = InventoryRepository(
|
|
self._connection_factory,
|
|
connection=connection,
|
|
)
|
|
try:
|
|
yield transaction_repository
|
|
connection.commit()
|
|
except Exception:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|
|
|
|
@contextmanager
|
|
def _connection_scope(
|
|
self,
|
|
*,
|
|
write: bool = False,
|
|
) -> Iterator[psycopg.Connection]:
|
|
if self._connection is not None:
|
|
yield self._connection
|
|
return
|
|
|
|
connection = self._connection_factory()
|
|
try:
|
|
yield connection
|
|
if write:
|
|
connection.commit()
|
|
except Exception:
|
|
if write:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|
|
|
|
@staticmethod
|
|
def _fetch_miner(
|
|
cur,
|
|
where_sql: str,
|
|
params: Sequence[Any],
|
|
*,
|
|
for_update: bool = False,
|
|
) -> dict | None:
|
|
lock_sql = " FOR UPDATE OF miners" if for_update else ""
|
|
cur.execute(
|
|
f"{MINER_SELECT} WHERE {where_sql}{lock_sql}",
|
|
params,
|
|
)
|
|
row = cur.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def get_miner_by_id(self, miner_id: UUID) -> dict | None:
|
|
with self._connection_scope() as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
return self._fetch_miner(cur, "miners.id = %s", (miner_id,))
|
|
|
|
def get_miner_by_id_for_update(self, miner_id: UUID) -> dict | None:
|
|
with self._connection_scope() as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
return self._fetch_miner(
|
|
cur,
|
|
"miners.id = %s",
|
|
(miner_id,),
|
|
for_update=True,
|
|
)
|
|
|
|
def get_miner_by_asset_id(self, asset_id: str) -> dict | None:
|
|
with self._connection_scope() as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
return self._fetch_miner(
|
|
cur,
|
|
"lower(trim(miners.asset_tag)) = lower(trim(%s))",
|
|
(asset_id,),
|
|
)
|
|
|
|
def list_miners(self, *, limit: int = 100, offset: int = 0) -> list[dict]:
|
|
return self.search_miners(limit=limit, offset=offset)
|
|
|
|
def search_miners(
|
|
self,
|
|
*,
|
|
query: str | None = None,
|
|
asset_id: str | None = None,
|
|
lifecycle_status: str | None = None,
|
|
tag_names: Sequence[str] = (),
|
|
limit: int = 100,
|
|
offset: int = 0,
|
|
) -> list[dict]:
|
|
conditions: list[str] = []
|
|
params: list[Any] = []
|
|
|
|
if query:
|
|
pattern = f"%{query}%"
|
|
conditions.append(
|
|
"""
|
|
(
|
|
miners.asset_tag ILIKE %s
|
|
OR miners.current_mac ILIKE %s
|
|
OR host(miners.current_ip) ILIKE %s
|
|
OR miners.model ILIKE %s
|
|
OR miners.miner_sn ILIKE %s
|
|
OR miners.control_board_sn ILIKE %s
|
|
OR miners.power_sn ILIKE %s
|
|
OR miners.notes ILIKE %s
|
|
)
|
|
"""
|
|
)
|
|
params.extend([pattern] * 8)
|
|
|
|
if asset_id:
|
|
conditions.append("lower(trim(miners.asset_tag)) = lower(trim(%s))")
|
|
params.append(asset_id)
|
|
|
|
if lifecycle_status:
|
|
conditions.append("miners.lifecycle_status = %s")
|
|
params.append(lifecycle_status)
|
|
|
|
for tag_name in tag_names:
|
|
conditions.append(
|
|
f"""
|
|
EXISTS (
|
|
SELECT 1
|
|
FROM miner_tags search_miner_tags
|
|
JOIN tags search_tags ON search_tags.id = search_miner_tags.tag_id
|
|
WHERE search_miner_tags.miner_id = miners.id
|
|
AND lower({TAG_TRIMMED_NAME_SQL.replace("name", "search_tags.name")})
|
|
= %s
|
|
)
|
|
"""
|
|
)
|
|
params.append(tag_name)
|
|
|
|
where_sql = " AND ".join(conditions) if conditions else "true"
|
|
params.extend([limit, offset])
|
|
|
|
with self._connection_scope() as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
f"""
|
|
{MINER_SELECT}
|
|
WHERE {where_sql}
|
|
ORDER BY
|
|
lower(trim(miners.asset_tag)) NULLS LAST,
|
|
miners.created_at,
|
|
miners.id
|
|
LIMIT %s OFFSET %s
|
|
""",
|
|
params,
|
|
)
|
|
return [dict(row) for row in cur.fetchall()]
|
|
|
|
def create_miner(self, values: Mapping[str, Any]) -> dict:
|
|
clean_values = {
|
|
key: value for key, value in values.items() if key in MINER_WRITE_COLUMNS
|
|
}
|
|
columns = list(clean_values)
|
|
params = [clean_values[column] for column in columns]
|
|
|
|
with self._connection_scope(write=True) as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
if columns:
|
|
placeholders = ", ".join(["%s"] * len(columns))
|
|
cur.execute(
|
|
f"""
|
|
INSERT INTO miners ({", ".join(columns)})
|
|
VALUES ({placeholders})
|
|
RETURNING id
|
|
""",
|
|
params,
|
|
)
|
|
else:
|
|
cur.execute("INSERT INTO miners DEFAULT VALUES RETURNING id")
|
|
miner_id = cur.fetchone()["id"]
|
|
return self._fetch_miner(cur, "miners.id = %s", (miner_id,))
|
|
|
|
def update_miner(self, miner_id: UUID, values: Mapping[str, Any]) -> dict | None:
|
|
clean_values = {
|
|
key: value for key, value in values.items() if key in MINER_WRITE_COLUMNS
|
|
}
|
|
if not clean_values:
|
|
return self.get_miner_by_id(miner_id)
|
|
|
|
assignments = [f"{column} = %s" for column in clean_values]
|
|
params = list(clean_values.values())
|
|
params.append(miner_id)
|
|
|
|
with self._connection_scope(write=True) as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
f"""
|
|
UPDATE miners
|
|
SET {", ".join(assignments)}, updated_at = now()
|
|
WHERE id = %s
|
|
RETURNING id
|
|
""",
|
|
params,
|
|
)
|
|
row = cur.fetchone()
|
|
if not row:
|
|
return None
|
|
return self._fetch_miner(cur, "miners.id = %s", (miner_id,))
|
|
|
|
def add_asset_id_history(
|
|
self,
|
|
miner_id: UUID,
|
|
old_asset_id: str | None,
|
|
new_asset_id: str,
|
|
changed_by: UUID,
|
|
) -> dict:
|
|
with self._connection_scope(write=True) as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO miner_asset_id_history (
|
|
miner_id,
|
|
old_asset_id,
|
|
new_asset_id,
|
|
changed_by
|
|
)
|
|
VALUES (%s, %s, %s, %s)
|
|
RETURNING
|
|
id,
|
|
miner_id,
|
|
old_asset_id,
|
|
new_asset_id,
|
|
changed_at,
|
|
changed_by
|
|
""",
|
|
(miner_id, old_asset_id, new_asset_id, changed_by),
|
|
)
|
|
return dict(cur.fetchone())
|
|
|
|
def list_asset_id_history(self, miner_id: UUID) -> list[dict]:
|
|
with self._connection_scope() as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT
|
|
id,
|
|
miner_id,
|
|
old_asset_id,
|
|
new_asset_id,
|
|
changed_at,
|
|
changed_by
|
|
FROM miner_asset_id_history
|
|
WHERE miner_id = %s
|
|
ORDER BY changed_at, id
|
|
""",
|
|
(miner_id,),
|
|
)
|
|
return [dict(row) for row in cur.fetchall()]
|
|
|
|
def list_tags(self) -> list[dict]:
|
|
with self._connection_scope() as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
f"""
|
|
SELECT id, name, created_at, updated_at
|
|
FROM tags
|
|
ORDER BY lower({TAG_TRIMMED_NAME_SQL}), id
|
|
"""
|
|
)
|
|
return [dict(row) for row in cur.fetchall()]
|
|
|
|
def add_tag(self, miner_id: UUID, tag_name: str) -> dict:
|
|
with self._connection_scope(write=True) as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
f"""
|
|
INSERT INTO tags (name)
|
|
VALUES (%s)
|
|
ON CONFLICT (lower({TAG_TRIMMED_NAME_SQL}))
|
|
DO UPDATE SET updated_at = tags.updated_at
|
|
RETURNING id, name, created_at, updated_at
|
|
""",
|
|
(tag_name,),
|
|
)
|
|
tag = dict(cur.fetchone())
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO miner_tags (miner_id, tag_id)
|
|
VALUES (%s, %s)
|
|
ON CONFLICT (miner_id, tag_id) DO NOTHING
|
|
""",
|
|
(miner_id, tag["id"]),
|
|
)
|
|
return tag
|
|
|
|
def remove_tag(self, miner_id: UUID, tag_id: UUID) -> bool:
|
|
with self._connection_scope(write=True) as connection:
|
|
with connection.cursor() as cur:
|
|
cur.execute(
|
|
"DELETE FROM miner_tags WHERE miner_id = %s AND tag_id = %s",
|
|
(miner_id, tag_id),
|
|
)
|
|
return cur.rowcount > 0
|