diff --git a/spark-analytics/base/webview.yaml b/spark-analytics/base/webview.yaml index 7da0c59..49b8292 100755 --- a/spark-analytics/base/webview.yaml +++ b/spark-analytics/base/webview.yaml @@ -76,14 +76,22 @@ data: def get_conn(): return psycopg2.connect(host=DB_HOST, port=DB_PORT, dbname=DB_NAME, user=DB_USER, password=DB_PASS) - def table_row_count(table: str) -> int: + def all_row_counts() -> dict: + counts = {} try: with get_conn() as conn: with conn.cursor() as cur: - cur.execute(f"SELECT COUNT(*) FROM {table}") - return cur.fetchone()[0] + for _, _, table in DATASETS: + try: + cur.execute(f"SELECT COUNT(*) FROM {table}") + counts[table] = cur.fetchone()[0] + except Exception: + conn.rollback() + counts[table] = -1 except Exception: - return -1 + for _, _, table in DATASETS: + counts.setdefault(table, -1) + return counts def fetch_rows(table: str, limit: int = 10_000): with get_conn() as conn: @@ -110,9 +118,10 @@ data: ) def index_html() -> str: + counts = all_row_counts() cards = "" for slug, label, table in DATASETS: - count = table_row_count(table) + count = counts.get(table, -1) if count > 0: badge = f'{count} rows' elif count == 0: @@ -179,7 +188,10 @@ data: self.send_header("Content-Type", "text/html; charset=utf-8") self.send_header("Content-Length", str(len(data))) self.end_headers() - self.wfile.write(data) + try: + self.wfile.write(data) + except BrokenPipeError: + pass if __name__ == "__main__": server = HTTPServer(("0.0.0.0", PORT), Handler)