| 1 | from typing import Any |
| 2 | from typing import Dict |
| 3 | from typing import List |
| 4 | |
| 5 | from fastapi import HTTPException |
| 6 | from influxdb_client.client.influxdb_client_async import InfluxDBClientAsync |
| 7 | from loguru import logger |
| 8 | |
| 9 | from app.connectors.influxdb.schema.alerts import InfluxDBAlert |
| 10 | from app.connectors.influxdb.schema.alerts import InfluxDBAlertsResponse |
| 11 | from app.connectors.utils import get_connector_info_from_db |
| 12 | from app.db.db_session import get_db_session |
| 13 | |
| 14 | |
| 15 | async def verify_influxdb_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]: |
| 16 | """ |
| 17 | Verifies the connection to InfluxDB service. |
| 18 | |
| 19 | Returns: |
| 20 | dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful. |
| 21 | """ |
| 22 | logger.info(f"Verifying the InfluxDB connection to {attributes['connector_url']}") |
| 23 | influxdb_client = InfluxDBClientAsync( |
| 24 | url=attributes["connector_url"], |
| 25 | token=attributes["connector_api_key"], |
| 26 | org=await get_influxdb_organization(), |
| 27 | verify_ssl=False, |
| 28 | ) |
| 29 | try: |
| 30 | ping = await influxdb_client.ping() |
| 31 | logger.info(f"Response from InfluxDB: {ping}") |
| 32 | if ping: |
| 33 | logger.info(f"Connection to {attributes['connector_url']} successful") |
| 34 | # Now try to fetch alerts |
| 35 | await get_alerts() |
| 36 | # logger.info(f"Alerts from InfluxDB: {alerts}") |
| 37 | return { |
| 38 | "connectionSuccessful": True, |
| 39 | "message": "InfluxDB connection successful", |
| 40 | } |
| 41 | else: |
| 42 | logger.error(f"Connection to {attributes['connector_url']} failed") |
| 43 | return { |
| 44 | "connectionSuccessful": False, |
| 45 | "message": f"Connection to {attributes['connector_url']} failed", |
| 46 | } |
| 47 | except Exception as e: |
| 48 | logger.error( |
| 49 | f"Connection to {attributes['connector_url']} failed with error: {e}", |
| 50 | ) |
| 51 | return { |
| 52 | "connectionSuccessful": False, |
| 53 | "message": f"Connection to {attributes['connector_url']} failed with error: {e}", |
| 54 | } |
| 55 | finally: |
| 56 | # Make sure to close the client session |
| 57 | await influxdb_client.close() |
| 58 | |
| 59 | |
| 60 | async def verify_influxdb_connection(connector_name: str) -> str: |
| 61 | """ |
| 62 | Returns the authentication token for the InfluxDB service. |
| 63 | |
| 64 | Returns: |
| 65 | str: Authentication token for the InfluxDB service. |
| 66 | """ |
| 67 | async with get_db_session() as session: # This will correctly enter the context manager |
| 68 | attributes = await get_connector_info_from_db(connector_name, session) |
| 69 | logger.info(f"Verifying the InfluxDB connection to {attributes['connector_url']}") |
| 70 | if attributes is None: |
| 71 | logger.error("No InfluxDB connector found in the database") |
| 72 | return None |
| 73 | return await verify_influxdb_credentials(attributes) |
| 74 | |
| 75 | |
| 76 | async def create_influxdb_client(connector_name: str) -> InfluxDBClientAsync: |
| 77 | """ |
| 78 | Returns an InfluxDBClientAsync client for the InfluxDB service. |
| 79 | |
| 80 | Returns: |
| 81 | InfluxDBClientAsync: InfluxDBClientAsync client for the InfluxDB service. |
| 82 | """ |
| 83 | # attributes = get_connector_info_from_db(connector_name) |
| 84 | async with get_db_session() as session: # This will correctly enter the context manager |
| 85 | attributes = await get_connector_info_from_db(connector_name, session) |
| 86 | if attributes is None: |
| 87 | raise HTTPException( |
| 88 | status_code=500, |
| 89 | detail=f"No {connector_name} connector found in the database", |
| 90 | ) |
| 91 | try: |
| 92 | return InfluxDBClientAsync( |
| 93 | url=attributes["connector_url"], |
| 94 | token=attributes["connector_api_key"], |
| 95 | org=await get_influxdb_organization(), |
| 96 | verify_ssl=False, |
| 97 | ) |
| 98 | except Exception as e: |
| 99 | raise HTTPException( |
| 100 | status_code=500, |
| 101 | detail=f"Failed to create InfluxDB client: {e}", |
| 102 | ) |
| 103 | |
| 104 | |
| 105 | async def get_influxdb_organization() -> str: |
| 106 | """ |
| 107 | Read the `connector_extra_data` from the database and return the organization name. |
| 108 | which is the first item. For example: `SOCFORTRESS,telegraf`. |
| 109 | """ |
| 110 | async with get_db_session() as session: # This will correctly enter the context manager |
| 111 | attributes = await get_connector_info_from_db("InfluxDB", session) |
| 112 | if attributes is None: |
| 113 | raise HTTPException( |
| 114 | status_code=500, |
| 115 | detail="No InfluxDB connector found in the database", |
| 116 | ) |
| 117 | return attributes["connector_extra_data"].split(",")[0] |
| 118 | |
| 119 | |
| 120 | # ! RUN A TEST QUERY TO FETCH ALERTS AND VERIFY THE CONNECTION |
| 121 | # Constants |
| 122 | BUCKET_NAME = "_monitoring" |
| 123 | |
| 124 | |
| 125 | def construct_query() -> str: |
| 126 | """Constructs the InfluxDB query. |
| 127 | |
| 128 | Returns: |
| 129 | str: The constructed InfluxDB query. |
| 130 | """ |
| 131 | return """ |
| 132 | from(bucket: "{bucket_name}") |
| 133 | |> range(start: -1h, stop: now()) |
| 134 | |> filter(fn: (r) => r._measurement == "statuses" and r._field == "_message") |
| 135 | |> filter(fn: (r) => exists r._check_id and exists r._value and exists r._check_name and exists r._level) |
| 136 | |> keep(columns: ["_time", "_value", "_check_id", "_check_name", "_level"]) |
| 137 | |> rename(columns: {{ "_time": "time", "_value": "message", "_check_id": "checkID", "_check_name": "checkName", "_level": "level" }}) |
| 138 | |> group() |
| 139 | |> sort(columns: ["time"], desc: true) |
| 140 | |> limit(n: 100, offset: 29) |
| 141 | """.format( |
| 142 | bucket_name=BUCKET_NAME, |
| 143 | ) |
| 144 | |
| 145 | |
| 146 | async def process_alert_records(result) -> List[InfluxDBAlert]: |
| 147 | """Processes alert records from InfluxDB query result. |
| 148 | |
| 149 | Args: |
| 150 | result: The query result from InfluxDB. |
| 151 | |
| 152 | Returns: |
| 153 | A list of InfluxDBAlert objects representing the processed alert records. |
| 154 | """ |
| 155 | alerts = [] |
| 156 | for table in result: |
| 157 | for record in table.records: |
| 158 | alert = InfluxDBAlert( |
| 159 | time=record.values.get("time").isoformat() if record.values.get("time") else None, |
| 160 | message=record.values.get("message"), |
| 161 | checkID=record.values.get("checkID"), |
| 162 | checkName=record.values.get("checkName"), |
| 163 | level=record.values.get("level"), |
| 164 | ) |
| 165 | alerts.append(alert) |
| 166 | return alerts |
| 167 | |
| 168 | |
| 169 | async def get_alerts() -> InfluxDBAlertsResponse: |
| 170 | """Fetches alerts from InfluxDB and returns them. |
| 171 | |
| 172 | Returns: |
| 173 | InfluxDBAlertsResponse: The response object containing the fetched alerts. |
| 174 | |
| 175 | Raises: |
| 176 | HTTPException: If there is an error fetching the alerts. |
| 177 | """ |
| 178 | client = await create_influxdb_client("InfluxDB") |
| 179 | try: |
| 180 | query = construct_query() |
| 181 | logger.info(f"Fetching alerts from InfluxDB: {query}") |
| 182 | result = await client.query_api().query(org=await get_influxdb_organization(), query=query) |
| 183 | logger.info(f"Alerts from InfluxDB: {result}") |
| 184 | alerts = await process_alert_records(result) |
| 185 | return InfluxDBAlertsResponse(alerts=alerts, success=True, message="Alerts fetched successfully") |
| 186 | except Exception as e: |
| 187 | raise HTTPException( |
| 188 | status_code=500, |
| 189 | detail=f"Failed to fetch alerts from InfluxDB: {e}", |
| 190 | ) |
| 191 | finally: |
| 192 | # Make sure to close the client session |
| 193 | await client.close() |