[{'Text': '"""Snowflake and ClickHouse database query tools."""\n\nimport json\nfrom typing import List, Optional\nfrom fastmcp import Context\nfrom point_topic_mcp.core.context_assembly import get_ontology_summary, list_datasets, assemble_context\nfrom point_topic_mcp.core.utils import dynamic_docstring, check_env_vars\nfrom dotenv import load_dotenv\n\nload_dotenv()\n\n# Snowflake database tools (require credentials)\n# if False: # Disable all snowflake tools as we migrate to ontology-first approach, but keep code for reference and potential future use.\nif check_env_vars(\'snowflake_tools\', [\'SNOWFLAKE_USER\', \'SNOWFLAKE_PASSWORD\']):\n    @dynamic_docstring([("{DATASETS}", list_datasets)])\n    def assemble_dataset_context(\n        dataset_names: List[str],\n        ctx: Optional[Context] = None\n    ) -> str:\n        """\n        Assemble full context (instructions, schema, examples) for one or more datasets.\n\n        This is essential before executing a query, for the agent to understand how to query the datasets.\n\n        Args:\n            dataset_names: List of dataset names to include (e.g., [\'upc\', \'upc_take_up\'])\n\n        {DATASETS}\n\n        Returns the complete context needed for querying these datasets.\n        """\n        return assemble_context(dataset_names)\n\n    async def execute_query(\n        sql_query: str,\n        ctx: Optional[Context] = None,\n    ) -> str:\n        """\n        Execute safe SQL queries against the Snowflake database.\n        Only read-only queries allowed (SELECT, WITH, SHOW, DESCRIBE, EXPLAIN).\n\n        Multiple queries can be executed in one call by separating them with semicolons (;).\n        Each query will be validated and executed separately, with clearly labeled results.\n\n        You must first assemble the context for the datasets you are querying using the assemble_dataset_context tool.\n\n        Args:\n            sql_query: The SQL query or queries to execute (separated by semicolons for multiple queries)\n\n        Returns:\n            Results starting with `query_id:<id>` on the first line, followed by CSV data.\n            The query_id can be passed to cancel_query() to cancel a running query,\n            or to get_query_status() to check its current state.\n            For multiple queries, each result section includes its own query_id.\n\n        Examples:\n            Single query: "SELECT COUNT(*) FROM table1"\n            Multiple queries: "SELECT COUNT(*) FROM table1; SELECT AVG(price) FROM table2"\n        """\n        from point_topic_mcp.connectors.snowflake import SnowflakeDB, non_empty_sql_statements\n\n        query_list = non_empty_sql_statements(sql_query)\n\n        # multi-query: use existing path (cancellation not supported per-query)\n        if len(query_list) > 1:\n            sf = SnowflakeDB()\n            sf.connect()\n            result = sf.execute_safe_queries(sql_query)\n            sf.close_connection()\n            return result\n\n        # single query: submit, emit query_id immediately via ctx, then fetch\n        sf = SnowflakeDB()\n        sf.connect()\n        try:\n            cursor, query_id = sf.submit_async_query(sql_query)\n        except ValueError as e:\n            sf.close_connection()\n            return f"ERROR: {e}"\n        except Exception as e:\n            sf.close_connection()\n            return f"DATABASE ERROR: {str(e)}"\n\n        # Emit query_id mid-execution via MCP notifications/progress (not logging).\n        # UPC Query Agent and other clients can subscribe to progress on tools/call.\n        if ctx:\n            report = getattr(ctx, "report_progress", None)\n            if report is None:\n                raise RuntimeError(\n                    "execute_query requires Context.report_progress (FastMCP progress API). "\n                    "Upgrade FastMCP or fix Context injection."\n                )\n            payload = json.dumps({"event": "query_executing", "query_id": query_id})\n            await report(progress=0, total=None, message=payload)\n\n        # run blocking poll loop in a thread so the event loop stays free\n        # to flush progress notifications to the client\n        import asyncio\n        result = await asyncio.to_thread(sf.poll_and_fetch, cursor, query_id)\n        sf.close_connection()\n        return result\n\n    def get_query_status(\n        query_id: str,\n        ctx: Optional[Context] = None\n    ) -> str:\n        """\n        Check the current status of a Snowflake query.\n\n        Useful for verifying a query is still running before cancelling it,\n        or confirming a cancellation completed.\n\n        Args:\n            query_id: The query ID returned in the first line of execute_query results.\n\n        Returns:\n            JSON with query_id and status: RUNNING, SUCCESS, CANCELLED, or ERROR:<message>\n        """\n        import json\n        from point_topic_mcp.connectors.snowflake import SnowflakeDB\n\n        sf = SnowflakeDB()\n        sf.connect()\n        status = sf.check_query_status(query_id)\n        sf.close_connection()\n        return json.dumps({"query_id": query_id, "status": status})\n\n    def cancel_query(\n        query_id: str,\n        ctx: Optional[Context] = None\n    ) -> str:\n        """\n        Cancel a running Snowflake query.\n\n        The query_id is returned in the first line of execute_query results.\n        Intended for human-initiated cancellation (e.g. a cancel button in the UI)\n        when a query is taking too long or was submitted in error.\n\n        Args:\n            query_id: The query ID to cancel.\n\n        Returns:\n            JSON confirming cancellation.\n        """\n        import json\n        from point_topic_mcp.connectors.snowflake import SnowflakeDB\n\n        sf = SnowflakeDB()\n        sf.connect()\n        sf.cancel_query(query_id)\n        sf.close_connection()\n        return json.dumps({"query_id": query_id, "status": "cancelled"})\n\n    def describe_table(table_name: str, ctx: Optional[Context] = None) -> str:\n        """\n        Describe a table in the Snowflake database.\n\n        Note the schema is already included in the assemble context function\n\n        Args:\n            table_name: The name of the table to describe.\n            Use the full database and schema name.\n            e.g. "upc_core.reports.upc_output"\n\n        Returns:\n            The schema as CSV string\n        """\n        from point_topic_mcp.connectors.snowflake import SnowflakeDB\n        sf = SnowflakeDB()\n        sf.connect()\n        result = sf.describe_table(table_name)\n        sf.close_connection()\n        return result\n\n    def get_la_code(la_name: str, ctx: Optional[Context] = None) -> str:\n        """\n        Get the LA code for a given LA name.\n\n        Args:\n            la_name: The name of the LA to get the code for\n\n        Returns:\n            The LA code for the given LA name\n\n        Example:\n            get_la_code("Westminster") -> "E09000033"\n        """\n        from point_topic_mcp.connectors.snowflake import SnowflakeDB\n\n        sql_query = f"select distinct la_code from upc_core.reports.upc_output where lower(la_name) like lower(\'{la_name}\')"\n        sf = SnowflakeDB()\n        sf.connect()\n        result = sf.execute_safe_query(sql_query)\n        sf.close_connection()\n        return result\n\n    def get_la_list_full(la_name: str, ctx: Optional[Context] = None) -> str:\n        """\n        Get the full list of LA codes and names.\n\n        Can be used if the get_la_code tool doesn\'t match the LA name.\n\n        Returns the full list of LA codes and names in CSV format.\n        """\n        from point_topic_mcp.connectors.snowflake import SnowflakeDB\n\n        sql_query = "select distinct la_code, la_name from upc_core.reports.upc_output"\n        sf = SnowflakeDB()\n        sf.connect()\n        result = sf.execute_safe_query(sql_query)\n        sf.close_connection()\n        return result\n\n\n    def get_dataset_status(ctx: Optional[Context] = None) -> str:\n        """\n        Get the current status of all UPC datasets including latest available dates.\n\n        This tells you:\n        - The latest UPC data snapshot date (upc_reported_at)\n        - The latest take-up data snapshot date (takeup_reported_at)\n        - The latest forecast data snapshot date (forecast_reported_at)\n        - Which datasets are currently available\n\n        Use this to determine:\n        - What date to filter for "current" or "latest" queries\n        - Whether time-series tables are needed for historical analysis\n        - Data freshness and availability\n\n        Returns:\n            CSV with dataset version and latest dates for each dataset\n        """\n        from point_topic_mcp.connectors.snowflake import SnowflakeDB\n\n        sql_query = "select * from upc_client._status.upc_status"\n        sf = SnowflakeDB()\n        sf.connect()\n        result = sf.execute_safe_query(sql_query)\n        sf.close_connection()\n        return result\n\n\n# ── ClickHouse ontology tools ────────────────────────────────────────────────\n# Requires CLICKHOUSE_HOST + CLICKHOUSE_PASSWORD env vars,\n# OR will auto-fetch from AWS Secrets Manager (ptserver-secret) if running on\n# the server (109.169.53.90).\n\n# def _clickhouse_available() -> bool:\n#     """True if ClickHouse creds are resolvable (env vars or Secrets Manager)."""\n#     if os.environ.get("CLICKHOUSE_HOST") and os.environ.get("CLICKHOUSE_PASSWORD"):\n#         return True\n#     # Try Secrets Manager silently\n#     try:\n#         import boto3, json as _json\n#         secret = _json.loads(\n#             boto3.client("secretsmanager", region_name="eu-west-1")\n#                 .get_secret_value(SecretId="ptserver-secret")["SecretString"]\n#         )\n#         return bool(secret.get("CLICKHOUSE_HOST") and secret.get("CLICKHOUSE_PASSWORD"))\n#     except Exception:\n#         return False\n\nif False:\n# if check_env_vars(\'clickhouse_tools\', [\'CLICKHOUSE_HOST\', \'CLICKHOUSE_PASSWORD\']) or _clickhouse_available():\n    @dynamic_docstring([("SUMMARY", get_ontology_summary)])\n    def assemble_ontology_context(\n        ctx: Optional[Context] = None,\n    ) -> str:\n        """\n        Assemble the full context (instructions, schema, SQL examples) for the\n        Point Topic Ontology dataset stored in ClickHouse.\n\n        Call this before executing any ontology query so the agent understands\n        the schema, available tables, and correct SQL dialect.\n\n        Returns the complete context needed for querying the ontology.\n\n        Summary: {SUMMARY}\n        """\n        return assemble_context(["ontology"])\n\n    def execute_ontology_query(\n        sql_query: str,\n        ctx: Optional[Context] = None,\n    ) -> str:\n        """\n        Execute a read-only SQL query against the Point Topic Ontology ClickHouse database.\n\n        Only SELECT/WITH/SHOW/DESCRIBE queries are allowed.\n        Call assemble_ontology_context first to understand the schema.\n\n        The ontology database contains UK telecom market data:\n        organisations, networks, ISP-network relationships, broadband services,\n        coverage measurements, and more.\n\n        Args:\n            sql_query: SQL query to run. Table names are unqualified lowercase\n                       (e.g. foaf_organization, cto_uses_network).\n                       Use toDate() for date comparisons and today() for current date.\n\n        Returns:\n            CSV results string, or an error message prefixed with ERROR:/DATABASE ERROR:\n        """\n        from point_topic_mcp.connectors.clickhouse import ClickHouseDB\n\n        db = ClickHouseDB()\n        try:\n            db.connect()\n        except RuntimeError as e:\n            return f"DATABASE ERROR: {e}"\n        except Exception as e:\n            return f"DATABASE ERROR: Failed to connect to ClickHouse: {e}"\n        \n        try:\n            return db.execute_safe_queries(sql_query)\n        finally:\n            db.close_connection()\n\n    def describe_ontology_table(\n        table_name: str,\n        ctx: Optional[Context] = None,\n    ) -> str:\n        """\n        Show the columns and types for a table in the ontology ClickHouse database.\n\n        Useful for exploring the schema or verifying column names before querying.\n\n        Args:\n            table_name: Unqualified table name, e.g. \'foaf_organization\', \'cto_measurement_record\'\n\n        Returns:\n            CSV with columns: name, type, comment\n        """\n        from point_topic_mcp.connectors.clickhouse import ClickHouseDB\n\n        db = ClickHouseDB()\n        try:\n            db.connect()\n        except RuntimeError as e:\n            return f"DATABASE ERROR: {e}"\n        except Exception as e:\n            return f"DATABASE ERROR: Failed to connect to ClickHouse: {e}"\n        \n        try:\n            return db.describe_table(table_name)\n        finally:\n            db.close_connection()\n'}]