From 34fe5a015a9590c0de8111c820d2de0f9a079ddb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jan=20B=C3=ADlek?= Date: Fri, 5 Jun 2026 10:29:16 +0200 Subject: [PATCH 1/5] =?UTF-8?q?SNOW-3194269:=20snowCD=20migration=20?= =?UTF-8?q?=E2=80=94=20per-endpoint=20connectivity=20diagnostics?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `snow connection test --enable-diag` now runs SnowCD-style live connectivity checks instead of just printing the path to the connector's text report. SnowCD itself is being marked unsupported (snowflake-prod-docs#22135) and Snowflake CLI is the documented replacement. What it does: - Fetches `SYSTEM$ALLOWLIST()` and merges `SYSTEM$ALLOWLIST_PRIVATELINK()` on top, deduped by (type, host, port). Falls back to (conn.host, 443) if the public allowlist call fails (low-priv role). - Probes every resolvable endpoint: DNS + TCP, plus a TLS handshake on port 443 with cert issuer + notAfter extracted. Latency measured with time.perf_counter. Wildcards and empty hosts marked Skipped. - Honours REQUESTS_CA_BUNDLE and SSL_CERT_FILE env vars for custom CA roots, matching what snowflake-connector-python reads. - Streams `Checking : ✅ (12.3 ms)` per endpoint, then prints a per-endpoint table (url | type | status | latency_ms | issuer | cert_expires) and a `Results: X Healthy, Y Unhealthy out of Z. N skipped (non-resolvable patterns)` summary line. - Network-policy snapshot: CURRENT_IP_ADDRESS() for the live IP, SHOW PARAMETERS LIKE 'NETWORK_POLICY' at account + user scope (user-level wins), DESC NETWORK POLICY for allow/block lists, DESC NETWORK RULE per referenced rule (mode/type/values). Privilege errors surface as a `Note:` row instead of silently showing empty lists. - --format json carries the same data structurally under Diagnostic.{checks, healthy, unhealthy, skipped, tested, network_policy}. Backward compatible: the connector's enable_diag=True still runs and the legacy SnowflakeConnectionTestReport.txt is still produced; the existing `Diag Report Location` line in the connection summary is preserved. Existing test_connection_test_diag_report passes unchanged. Tests: 29 new in tests/test_connection_diagnostic.py covering wildcard skip, DNS / TCP / TLS failures, port-80 TCP-only, latency capture, allowlist file vs query branches, malformed-entry filtering, fallback on permission errors, PrivateLink merge + dedupe + silent failure, CA bundle resolution priority, and the network-policy snapshot (user > account precedence, DESC failure mode, query-error survival). All 96 existing tests/test_connection.py tests continue to pass. Live-tested against a real account (TABLE format, JSON format, --diag-allowlist-path, REQUESTS_CA_BUNDLE pickup). --- RELEASE-NOTES.md | 1 + .../cli/_plugins/connection/commands.py | 112 ++++ .../cli/_plugins/connection/diagnostic.py | 517 ++++++++++++++++++ tests/test_connection_diagnostic.py | 464 ++++++++++++++++ 4 files changed, 1094 insertions(+) create mode 100644 src/snowflake/cli/_plugins/connection/diagnostic.py create mode 100644 tests/test_connection_diagnostic.py diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 94b57548f5..449c5fb538 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -19,6 +19,7 @@ ## Deprecations ## New additions +* `snow connection test --enable-diag` now runs SnowCD-style per-endpoint connectivity checks on the live connection. It merges `SYSTEM$ALLOWLIST()` with `SYSTEM$ALLOWLIST_PRIVATELINK()` (deduped by `(type, host, port)`; the privatelink call is best-effort and skipped on deployments where it is unavailable), probes each resolvable endpoint (DNS, TCP, plus a TLS handshake on port 443) with measured latency, and prints a streaming `Checking : ✅` line per endpoint, followed by a per-endpoint table (`url | type | status | latency_ms | issuer | cert_expires`), an effective-network-policy summary (account/user policy, current IP, allowed/blocked IP and rule lists), and a `Results: X Healthy, Y Unhealthy out of Z. N skipped (non-resolvable patterns)` line. Wildcard hostnames are counted as `Skipped`. The TLS probe respects the `REQUESTS_CA_BUNDLE` and `SSL_CERT_FILE` env vars (matching `snowflake-connector-python`) for custom CA bundles. JSON output (`--format json`) carries the same data under a `Diagnostic` key. The legacy `SnowflakeConnectionTestReport.txt` written by `snowflake-connector-python` is still produced. Replaces the deprecated SnowCD tool. * Added a `--protocol` option (plus matching `SNOWFLAKE_PROTOCOL` env var and `protocol` config key) to `snow connection add` and the global connection overrides. This allows selecting `http` or `https` as the connection protocol without editing `config.toml`, which is primarily useful for local development against `http` deployments. * Added `snow helpers generate-project-schema` command to emit a JSON Schema for the `snowflake.yml` project definition file. The output follows [JSON Schema Draft 2020-12](https://json-schema.org/draft/2020-12/schema) and can be plugged into supported editors (e.g. VS Code via the YAML extension) or CI pipelines to get completion and to catch structural mistakes (unknown keys, wrong types, missing required fields) before a deploy. The schema is generated from the CLI's own models, so some cross-field and semantic checks are still only applied at load/deploy time. Use `--definition-version` to select the project definition version (`1`, `1.1`, or `2`; defaults to `2`) and `--output-file`/`-o` to write the schema to a file. diff --git a/src/snowflake/cli/_plugins/connection/commands.py b/src/snowflake/cli/_plugins/connection/commands.py index 0218ee9af3..5487785004 100644 --- a/src/snowflake/cli/_plugins/connection/commands.py +++ b/src/snowflake/cli/_plugins/connection/commands.py @@ -28,6 +28,11 @@ ) from click.core import ParameterSource # type: ignore from snowflake import connector +from snowflake.cli._plugins.connection.diagnostic import ( + collect_network_policy, + run_diagnostic, + status_line, +) from snowflake.cli._plugins.connection.util import ( strip_if_value_present, ) @@ -417,10 +422,117 @@ def test( result["Diag Report Location"] = os.path.join( conn_ctx.diag_log_path, "SnowflakeConnectionTestReport.txt" ) + return _connection_test_with_diag(conn, conn_ctx, result) return ObjectResult(result) +def _connection_test_with_diag( + conn, + conn_ctx, + connection_summary: dict, +) -> CommandResult: + """Run the SnowCD-style per-endpoint diagnostic and assemble the full result. + + Streams `Checking : ` lines for human (TABLE) output and + appends a per-endpoint table plus a `Results: ...` summary line. For JSON + output the streaming is suppressed and the structured payload carries the + diagnostic data instead. + """ + from snowflake.cli.api.output.formats import OutputFormat + from snowflake.cli.api.output.types import MultipleResults + + is_table_output = get_cli_context().output_format == OutputFormat.TABLE + + if is_table_output: + cli_console.message("Fetching allowlist from Snowflake...") + + def _stream(check): + if is_table_output: + cli_console.message(status_line(check)) + + report = run_diagnostic( + conn=conn, + allowlist_path=conn_ctx.diag_allowlist_path, + on_check=_stream, + ) + + if is_table_output: + cli_console.message("Inspecting network policies...") + policy = collect_network_policy(conn, user=conn.user) + + diagnostic_payload = { + "checks": [c.to_dict() for c in report.checks], + "healthy": report.healthy, + "unhealthy": report.unhealthy, + "skipped": report.skipped, + "tested": report.tested, + "network_policy": policy.to_dict(), + } + + if not is_table_output: + # JSON / CSV output gets the structured payload nested in the + # connection summary so consumers can read everything off one object. + connection_summary["Diagnostic"] = diagnostic_payload + return ObjectResult(connection_summary) + + # TABLE output: keep the connection summary clean. The per-endpoint + # table, network policy block, and summary line below carry the data. + results = MultipleResults() + results.add(ObjectResult(connection_summary)) + tested_rows = [ + { + "url": c.host, + "type": c.type, + "status": c.status, + "latency_ms": c.latency_ms if c.latency_ms is not None else "", + "issuer": c.cert_issuer or "", + "cert_expires": c.cert_expires or "", + } + for c in report.checks + if c.status != "Skipped" + ] + if tested_rows: + results.add(CollectionResult(tested_rows)) + if policy.has_policy(): + policy_summary = { + "Effective network policy": policy.effective_policy, + "Source": "user" if policy.user_policy else "account", + "Account-level": policy.account_policy or "(none)", + "User-level": policy.user_policy or "(none)", + "Current IP": policy.current_ip or "(unknown)", + "Allowed IPs": ", ".join(policy.allowed_ip_list) or "(none)", + "Blocked IPs": ", ".join(policy.blocked_ip_list) or "(none)", + "Allowed network rules": ", ".join(policy.allowed_rule_list) or "(none)", + "Blocked network rules": ", ".join(policy.blocked_rule_list) or "(none)", + } + if policy.error: + policy_summary["Note"] = policy.error + results.add(ObjectResult(policy_summary)) + if policy.rules: + results.add( + CollectionResult( + [ + { + "rule": r.name, + "mode": r.mode, + "type": r.type, + "values": ", ".join(r.values), + } + for r in policy.rules + ] + ) + ) + elif policy.current_ip: + results.add( + MessageResult( + f"No network policy in effect (current IP: {policy.current_ip})." + ) + ) + results.add(MessageResult(report.summary_line())) + return results + + @app.command(requires_connection=False) def set_default( name: str = typer.Argument( diff --git a/src/snowflake/cli/_plugins/connection/diagnostic.py b/src/snowflake/cli/_plugins/connection/diagnostic.py new file mode 100644 index 0000000000..5416a9880c --- /dev/null +++ b/src/snowflake/cli/_plugins/connection/diagnostic.py @@ -0,0 +1,517 @@ +# Copyright (c) 2024 Snowflake Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Connectivity diagnostics for `snow connection test --enable-diag`. + +Replicates the per-endpoint checks that SnowCD used to provide, sourcing +endpoints from `SYSTEM$ALLOWLIST()` (or a JSON file passed via +`--diag-allowlist-path`). Each resolvable endpoint is probed with a TCP +connect, and TLS port 443 endpoints additionally have their certificate +issuer and expiry recorded. +""" +from __future__ import annotations + +import json +import logging +import os +import socket +import ssl +import time +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, Callable, Iterable, Literal, Optional + +from click.exceptions import ClickException +from snowflake.cli._plugins.connection.util import ALLOWLIST_QUERY +from snowflake.connector import SnowflakeConnection +from snowflake.connector.cursor import DictCursor + +log = logging.getLogger(__name__) + +DEFAULT_TIMEOUT_SECONDS: float = 5.0 +TLS_PORT: int = 443 + +Status = Literal["Healthy", "Unhealthy", "Skipped"] + + +@dataclass +class EndpointCheck: + host: str + port: int + type: str # noqa: A003 — intentional public field name, matches `SYSTEM$ALLOWLIST()` JSON + status: Status + error: Optional[str] = None + cert_issuer: Optional[str] = None + cert_expires: Optional[str] = None + latency_ms: Optional[float] = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class DiagnosticReport: + checks: list[EndpointCheck] = field(default_factory=list) + + @property + def healthy(self) -> int: + return sum(1 for c in self.checks if c.status == "Healthy") + + @property + def unhealthy(self) -> int: + return sum(1 for c in self.checks if c.status == "Unhealthy") + + @property + def skipped(self) -> int: + return sum(1 for c in self.checks if c.status == "Skipped") + + @property + def tested(self) -> int: + return self.healthy + self.unhealthy + + def summary_line(self) -> str: + return ( + f"Results: {self.healthy} Healthy, {self.unhealthy} Unhealthy " + f"out of {self.tested} endpoints. " + f"{self.skipped} skipped (non-resolvable patterns)." + ) + + +ALLOWLIST_PRIVATELINK_QUERY = "SELECT SYSTEM$ALLOWLIST_PRIVATELINK()" + + +def _query_allowlist(conn: SnowflakeConnection, sql: str, key: str) -> list[dict]: + """Run an allowlist-returning system function and parse its single JSON cell. + + Returns `[]` on any failure (permission denied, function not present in + the deployment, malformed JSON). The caller must not let this raise. + """ + try: + *_, cursor = conn.execute_string(sql, cursor_class=DictCursor) + row = cursor.fetchone() + if not row: + return [] + raw = row.get(key) + if not raw: + return [] + parsed = json.loads(raw) + return parsed if isinstance(parsed, list) else [] + except Exception: + log.debug("Allowlist query failed: %s", sql, exc_info=True) + return [] + + +def _dedupe_entries(entries: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Drop duplicate `(type, host, port)` triples, preserving first-seen order.""" + seen: set[tuple[str, str, Any]] = set() + out: list[dict[str, Any]] = [] + for e in entries: + key = (str(e.get("type", "")), str(e.get("host", "")), e.get("port")) + if key in seen: + continue + seen.add(key) + out.append(e) + return out + + +def load_allowlist( + conn: SnowflakeConnection, allowlist_path: Optional[Path] +) -> list[dict[str, Any]]: + """Return the raw allowlist as a list of `{type, host, port}` dicts. + + If `allowlist_path` is given, parse it as JSON and use it directly. + Otherwise merge `SYSTEM$ALLOWLIST()` and `SYSTEM$ALLOWLIST_PRIVATELINK()` + on the open connection. PrivateLink entries are appended after the public + set; duplicate `(type, host, port)` triples are dropped. Either query may + fail (low-priv role, function unavailable in deployment) — failures are + silent and we return whatever we managed to collect, with the public + allowlist's failure handled by `run_diagnostic` (fall back to the + connection host) and the privatelink failure simply skipping that source. + """ + if allowlist_path is not None: + try: + payload = json.loads(Path(allowlist_path).read_text()) + except (OSError, json.JSONDecodeError) as exc: + raise ClickException( + f"Could not read allowlist file {allowlist_path}: {exc}" + ) + if not isinstance(payload, list): + raise ClickException( + f"Allowlist file {allowlist_path} must contain a JSON array." + ) + return payload + + # Public allowlist: let exceptions propagate so run_diagnostic's + # fallback-to-conn.host branch fires when the role lacks privileges. + *_, cursor = conn.execute_string(ALLOWLIST_QUERY, cursor_class=DictCursor) + public = json.loads(cursor.fetchone()["SYSTEM$ALLOWLIST()"]) + privatelink = _query_allowlist( + conn, ALLOWLIST_PRIVATELINK_QUERY, "SYSTEM$ALLOWLIST_PRIVATELINK()" + ) + if not isinstance(public, list): + public = [] + return _dedupe_entries([*public, *privatelink]) + + +def is_resolvable(host: str) -> bool: + """Return False for hostnames the OS resolver can never satisfy. + + Allowlist entries can include wildcard patterns (e.g. + `*.region.snowflakecomputing.com`) and bare placeholder values; those are + counted as `Skipped` rather than `Unhealthy`. + """ + if not host: + return False + if "*" in host: + return False + return True + + +def _resolve_cafile() -> Optional[str]: + """Honour the same CA bundle env vars `snowflake-connector-python` reads. + + `REQUESTS_CA_BUNDLE` wins (matches `requests` and the connector); falls + back to `SSL_CERT_FILE` (the OpenSSL convention). Returns `None` if + neither points at a readable file, in which case `ssl.create_default_context()` + will use the system trust store. + """ + for var in ("REQUESTS_CA_BUNDLE", "SSL_CERT_FILE"): + path = os.environ.get(var) + if path and os.path.isfile(path): + return path + return None + + +def _probe_tls(sock: socket.socket, host: str) -> tuple[Optional[str], Optional[str]]: + """Wrap `sock` in TLS and return `(issuer, not_after)` from the peer cert. + + Caller is responsible for closing the underlying socket. + """ + context = ssl.create_default_context(cafile=_resolve_cafile()) + context.minimum_version = ssl.TLSVersion.TLSv1_2 + with context.wrap_socket(sock, server_hostname=host) as ssock: + cert: Any = ssock.getpeercert() or {} + issuer_raw: Any = cert.get("issuer") or () + issuer_cn: Optional[str] = None + for rdn in issuer_raw: + for entry in rdn: + if ( + isinstance(entry, tuple) + and len(entry) == 2 + and entry[0] == "organizationName" + ): + issuer_cn = entry[1] + break + if issuer_cn: + break + not_after = cert.get("notAfter") + return issuer_cn, not_after if isinstance(not_after, str) else None + + +def check_endpoint( + host: str, + port: int, + type_: str, + timeout: float = DEFAULT_TIMEOUT_SECONDS, +) -> EndpointCheck: + """Probe a single endpoint. Never raises; failures map to Unhealthy.""" + if not is_resolvable(host): + return EndpointCheck( + host=host, + port=port, + type=type_, + status="Skipped", + error="non-resolvable pattern", + ) + + try: + # getaddrinfo + connect together cover DNS and TCP reachability. + start = time.perf_counter() + sock = socket.create_connection((host, port), timeout=timeout) + latency_ms = round((time.perf_counter() - start) * 1000, 1) + except OSError as exc: + return EndpointCheck( + host=host, + port=port, + type=type_, + status="Unhealthy", + error=str(exc), + ) + + issuer = expires = None + try: + if port == TLS_PORT: + try: + issuer, expires = _probe_tls(sock, host) + except (ssl.SSLError, OSError) as exc: + return EndpointCheck( + host=host, + port=port, + type=type_, + status="Unhealthy", + error=f"TLS handshake failed: {exc}", + ) + finally: + try: + sock.close() + except OSError: + pass + + return EndpointCheck( + host=host, + port=port, + type=type_, + status="Healthy", + cert_issuer=issuer, + cert_expires=expires, + latency_ms=latency_ms, + ) + + +def _normalise_entries( + entries: Iterable[dict[str, Any]], +) -> list[tuple[str, int, str]]: + """Pull (host, port, type) triples out of allowlist entries, dropping malformed rows.""" + out: list[tuple[str, int, str]] = [] + for entry in entries: + host = entry.get("host") + port = entry.get("port", TLS_PORT) + type_ = entry.get("type", "UNKNOWN") + if not isinstance(host, str): + continue + try: + port_int = int(port) + except (TypeError, ValueError): + continue + out.append((host, port_int, str(type_))) + return out + + +def run_diagnostic( + conn: SnowflakeConnection, + allowlist_path: Optional[Path], + on_check: Callable[[EndpointCheck], None] = lambda _check: None, + timeout: float = DEFAULT_TIMEOUT_SECONDS, +) -> DiagnosticReport: + """Load the allowlist, probe each entry, return a full report. + + Calls `on_check(...)` once per endpoint in input order; this is how the CLI + streams `Checking : ✅` lines as the run progresses. + + If `SYSTEM$ALLOWLIST()` cannot be queried (no file path and the call + fails — typically a permission error for low-privilege roles), the + diagnostic falls back to checking just `(conn.host, 443)`. + """ + try: + allowlist = load_allowlist(conn, allowlist_path) + except ClickException: + raise + except Exception as exc: + log.warning( + "Could not call SYSTEM$ALLOWLIST(); falling back to the connection host.", + exc_info=True, + ) + allowlist = [ + {"type": "SNOWFLAKE_DEPLOYMENT", "host": conn.host, "port": TLS_PORT} + ] + # Surface the fallback so users notice it in --enable-diag output. + log.info("Allowlist fetch failed: %s", exc) + + report = DiagnosticReport() + for host, port, type_ in _normalise_entries(allowlist): + check = check_endpoint(host, port, type_, timeout=timeout) + report.checks.append(check) + on_check(check) + return report + + +def status_line(check: EndpointCheck) -> str: + """Return the streaming-line representation: `Checking : `.""" + icon = {"Healthy": "✅", "Unhealthy": "❌", "Skipped": "⏭"}[check.status] + suffix = f" ({check.error})" if check.status == "Unhealthy" and check.error else "" + if check.status == "Healthy" and check.latency_ms is not None: + suffix = f" ({check.latency_ms} ms)" + return f"Checking {check.type}: {check.host} {icon}{suffix}" + + +# --------------------------------------------------------------------------- # +# Network policy / IP rule diagnostic +# --------------------------------------------------------------------------- # + + +@dataclass +class NetworkRule: + name: str + mode: str # INGRESS | EGRESS | INTERNAL_STAGE + type: str # noqa: A003 — IPV4 | HOST_PORT | AWSVPCEID | AZURELINKID | PRIVATE_HOST_PORT + values: list[str] = field(default_factory=list) + + +@dataclass +class NetworkPolicySnapshot: + """Snapshot of network-policy state at the time of `snow connection test`. + + `account_policy` and `user_policy` are the names returned by + `SHOW PARAMETERS LIKE 'NETWORK_POLICY'` at each scope; the effective + policy is `user_policy or account_policy` per Snowflake precedence rules. + """ + + current_ip: Optional[str] = None + account_policy: Optional[str] = None + user_policy: Optional[str] = None + effective_policy: Optional[str] = None + allowed_ip_list: list[str] = field(default_factory=list) + blocked_ip_list: list[str] = field(default_factory=list) + allowed_rule_list: list[str] = field(default_factory=list) + blocked_rule_list: list[str] = field(default_factory=list) + rules: list[NetworkRule] = field(default_factory=list) + error: Optional[str] = None + + def has_policy(self) -> bool: + return self.effective_policy is not None + + def to_dict(self) -> dict[str, Any]: + return { + "current_ip": self.current_ip, + "account_policy": self.account_policy, + "user_policy": self.user_policy, + "effective_policy": self.effective_policy, + "allowed_ip_list": self.allowed_ip_list, + "blocked_ip_list": self.blocked_ip_list, + "allowed_network_rule_list": self.allowed_rule_list, + "blocked_network_rule_list": self.blocked_rule_list, + "rules": [asdict(r) for r in self.rules], + "error": self.error, + } + + +def _safe_query(conn: SnowflakeConnection, sql: str) -> Optional[list[dict]]: + """Run a query through `execute_string`; return rows as dicts, or None on error. + + Network-policy diagnostics are best-effort: a low-priv role may not be + able to call `SHOW PARAMETERS` at account scope or `DESCRIBE NETWORK + POLICY`. Fail closed and let the caller surface the gap. + """ + try: + *_, cursor = conn.execute_string(sql, cursor_class=DictCursor) + return list(cursor) + except Exception as exc: + log.debug("Network-policy query failed: %s\n%s", sql, exc) + return None + + +def _scalar_value(rows: Optional[list[dict]], column: str) -> Optional[str]: + if not rows: + return None + val = rows[0].get(column) + return str(val) if val not in (None, "") else None + + +def _split_csv(value: Optional[str]) -> list[str]: + if not value: + return [] + return [v.strip() for v in value.split(",") if v.strip()] + + +def _describe_network_policy( + conn: SnowflakeConnection, name: str, snapshot: NetworkPolicySnapshot +) -> bool: + """Populate snapshot allow/block lists from `DESC NETWORK POLICY `. + + Returns True if the DESC succeeded (any row came back), False otherwise. + """ + rows = _safe_query(conn, f"DESC NETWORK POLICY IDENTIFIER('{name}')") + if not rows: + return False + by_name: dict[str, str] = {} + for row in rows: + prop = row.get("name") or row.get("NAME") + val = row.get("value") or row.get("VALUE") + if prop and val is not None: + by_name[str(prop).upper()] = str(val) + snapshot.allowed_ip_list = _split_csv(by_name.get("ALLOWED_IP_LIST")) + snapshot.blocked_ip_list = _split_csv(by_name.get("BLOCKED_IP_LIST")) + snapshot.allowed_rule_list = _split_csv(by_name.get("ALLOWED_NETWORK_RULE_LIST")) + snapshot.blocked_rule_list = _split_csv(by_name.get("BLOCKED_NETWORK_RULE_LIST")) + return True + + +def _describe_network_rule( + conn: SnowflakeConnection, qualified_name: str +) -> Optional[NetworkRule]: + rows = _safe_query(conn, f"DESC NETWORK RULE IDENTIFIER('{qualified_name}')") + if not rows: + return None + head = rows[0] + keys = {k.lower(): k for k in head.keys()} + mode = head.get(keys.get("mode", ""), "") or "" + type_ = head.get(keys.get("type", ""), "") or "" + raw_values = head.get(keys.get("value_list", ""), "") or "" + return NetworkRule( + name=qualified_name, + mode=str(mode), + type=str(type_), + values=_split_csv(str(raw_values)), + ) + + +def collect_network_policy( + conn: SnowflakeConnection, user: Optional[str] = None +) -> NetworkPolicySnapshot: + """Best-effort snapshot of the active network policy and its referenced rules. + + Sources, in order of precedence: + 1. `SHOW PARAMETERS LIKE 'NETWORK_POLICY' FOR USER ` (user-level wins) + 2. `SHOW PARAMETERS LIKE 'NETWORK_POLICY' IN ACCOUNT` + 3. `DESC NETWORK POLICY ` for the inline allow/block lists + 4. `DESC NETWORK RULE ` for referenced rule values + + Always returns a snapshot; failures populate `snapshot.error` instead. + """ + snapshot = NetworkPolicySnapshot() + + # CURRENT_IP_ADDRESS() is the documented function. SYSTEM$GET_CLIENT_IP + # exists in some internal deployments but is not generally available. + ip_rows = _safe_query(conn, "SELECT CURRENT_IP_ADDRESS() AS IP") + if ip_rows: + snapshot.current_ip = _scalar_value(ip_rows, "IP") + + account_rows = _safe_query(conn, "SHOW PARAMETERS LIKE 'NETWORK_POLICY' IN ACCOUNT") + snapshot.account_policy = _scalar_value(account_rows, "value") + + if user: + user_rows = _safe_query( + conn, + f"SHOW PARAMETERS LIKE 'NETWORK_POLICY' FOR USER IDENTIFIER('{user}')", + ) + snapshot.user_policy = _scalar_value(user_rows, "value") + + snapshot.effective_policy = snapshot.user_policy or snapshot.account_policy + + if snapshot.effective_policy: + described = _describe_network_policy(conn, snapshot.effective_policy, snapshot) + if not described: + snapshot.error = ( + f"Could not DESC NETWORK POLICY {snapshot.effective_policy} " + "(role lacks privilege?). Allowed/blocked lists not shown." + ) + for rule_name in ( + *snapshot.allowed_rule_list, + *snapshot.blocked_rule_list, + ): + rule = _describe_network_rule(conn, rule_name) + if rule is not None: + snapshot.rules.append(rule) + + return snapshot diff --git a/tests/test_connection_diagnostic.py b/tests/test_connection_diagnostic.py new file mode 100644 index 0000000000..f40222e9a9 --- /dev/null +++ b/tests/test_connection_diagnostic.py @@ -0,0 +1,464 @@ +# Copyright (c) 2024 Snowflake Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for `snow connection test --enable-diag` SnowCD-style diagnostics.""" +from __future__ import annotations + +import json +import socket +import ssl +from pathlib import Path +from tempfile import NamedTemporaryFile +from unittest import mock + +import pytest +from click.exceptions import ClickException +from snowflake.cli._plugins.connection.diagnostic import ( + DiagnosticReport, + EndpointCheck, + NetworkPolicySnapshot, + check_endpoint, + collect_network_policy, + is_resolvable, + load_allowlist, + run_diagnostic, + status_line, +) + +ALLOWLIST_FIXTURE = [ + { + "type": "SNOWFLAKE_DEPLOYMENT", + "host": "acct.snowflakecomputing.com", + "port": 443, + }, + {"type": "OCSP_RESPONDER", "host": "ocsp.x.com", "port": 80}, + {"type": "STAGE", "host": "*.region.snowflakecomputing.com", "port": 443}, +] + + +@pytest.mark.parametrize( + "host, expected", + [ + ("acct.snowflakecomputing.com", True), + ("sfc-stage.s3.amazonaws.com", True), + ("*.region.snowflakecomputing.com", False), + ("", False), + ], +) +def test_is_resolvable(host, expected): + assert is_resolvable(host) is expected + + +def test_status_line_includes_latency_for_healthy(): + line = status_line( + EndpointCheck("x.com", 443, "SNOWFLAKE_DEPLOYMENT", "Healthy", latency_ms=12.3) + ) + assert "✅" in line + assert "x.com" in line + assert "12.3 ms" in line + + +def test_status_line_includes_error_for_unhealthy(): + line = status_line( + EndpointCheck( + "x.com", 443, "SNOWFLAKE_DEPLOYMENT", "Unhealthy", error="DNS fail" + ) + ) + assert "❌" in line + assert "DNS fail" in line + + +def test_load_allowlist_from_file(): + with NamedTemporaryFile("w+", suffix=".json", delete=False) as f: + f.write(json.dumps(ALLOWLIST_FIXTURE)) + path = Path(f.name) + conn = mock.MagicMock() + out = load_allowlist(conn, path) + assert out == ALLOWLIST_FIXTURE + conn.execute_string.assert_not_called() + + +def test_load_allowlist_file_invalid_json_raises(): + with NamedTemporaryFile("w+", suffix=".json", delete=False) as f: + f.write("not json") + path = Path(f.name) + with pytest.raises(ClickException): + load_allowlist(mock.MagicMock(), path) + + +def test_load_allowlist_from_query_when_no_path(): + cursor = mock.MagicMock() + cursor.fetchone.return_value = {"SYSTEM$ALLOWLIST()": json.dumps(ALLOWLIST_FIXTURE)} + conn = mock.MagicMock() + conn.execute_string.return_value = (cursor,) + assert load_allowlist(conn, None) == ALLOWLIST_FIXTURE + + +def test_load_allowlist_merges_privatelink_entries(): + """SYSTEM$ALLOWLIST_PRIVATELINK() entries are appended after the public set.""" + pl_extra = [ + { + "type": "SNOWFLAKE_DEPLOYMENT", + "host": "acct.privatelink.snowflakecomputing.com", + "port": 443, + } + ] + + def execute_string(sql, cursor_class=None): + cur = mock.MagicMock() + if "PRIVATELINK" in sql: + cur.fetchone.return_value = { + "SYSTEM$ALLOWLIST_PRIVATELINK()": json.dumps(pl_extra) + } + else: + cur.fetchone.return_value = { + "SYSTEM$ALLOWLIST()": json.dumps(ALLOWLIST_FIXTURE) + } + return (cur,) + + conn = mock.MagicMock() + conn.execute_string.side_effect = execute_string + out = load_allowlist(conn, None) + assert out == ALLOWLIST_FIXTURE + pl_extra + + +def test_load_allowlist_dedupes_overlap_with_privatelink(): + """A (type, host, port) appearing in both lists collapses to one entry.""" + overlap = ALLOWLIST_FIXTURE[0] # SNOWFLAKE_DEPLOYMENT acct... + + def execute_string(sql, cursor_class=None): + cur = mock.MagicMock() + if "PRIVATELINK" in sql: + cur.fetchone.return_value = { + "SYSTEM$ALLOWLIST_PRIVATELINK()": json.dumps([overlap]) + } + else: + cur.fetchone.return_value = { + "SYSTEM$ALLOWLIST()": json.dumps(ALLOWLIST_FIXTURE) + } + return (cur,) + + conn = mock.MagicMock() + conn.execute_string.side_effect = execute_string + out = load_allowlist(conn, None) + # Same length — the duplicate from PL was dropped. + assert len(out) == len(ALLOWLIST_FIXTURE) + + +def test_load_allowlist_silently_skips_failing_privatelink_query(): + """If SYSTEM$ALLOWLIST_PRIVATELINK is unavailable the public list is still returned.""" + + def execute_string(sql, cursor_class=None): + if "PRIVATELINK" in sql: + raise RuntimeError("Function not enabled in this deployment") + cur = mock.MagicMock() + cur.fetchone.return_value = { + "SYSTEM$ALLOWLIST()": json.dumps(ALLOWLIST_FIXTURE) + } + return (cur,) + + conn = mock.MagicMock() + conn.execute_string.side_effect = execute_string + assert load_allowlist(conn, None) == ALLOWLIST_FIXTURE + + +def test_check_endpoint_skips_wildcards(): + result = check_endpoint("*.x.com", 443, "STAGE") + assert result.status == "Skipped" + assert result.error == "non-resolvable pattern" + + +def test_check_endpoint_marks_unhealthy_on_dns_failure(): + with mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + side_effect=socket.gaierror("Name or service not known"), + ): + result = check_endpoint("nonexistent.invalid", 443, "SNOWFLAKE_DEPLOYMENT") + assert result.status == "Unhealthy" + assert "Name or service" in (result.error or "") + assert result.latency_ms is None + + +def test_check_endpoint_marks_unhealthy_on_tls_failure(): + fake_sock = mock.MagicMock() + with mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + return_value=fake_sock, + ), mock.patch( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + side_effect=ssl.SSLError("cert verify failed"), + ): + result = check_endpoint("x.com", 443, "SNOWFLAKE_DEPLOYMENT") + assert result.status == "Unhealthy" + assert "TLS handshake failed" in (result.error or "") + + +def test_check_endpoint_records_latency_and_cert(): + fake_sock = mock.MagicMock() + with mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + return_value=fake_sock, + ), mock.patch( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + return_value=("DigiCert Inc", "Jan 19 23:59:59 2027 GMT"), + ): + result = check_endpoint("x.com", 443, "SNOWFLAKE_DEPLOYMENT") + assert result.status == "Healthy" + assert result.cert_issuer == "DigiCert Inc" + assert result.cert_expires == "Jan 19 23:59:59 2027 GMT" + assert result.latency_ms is not None + assert result.latency_ms >= 0 + + +def test_check_endpoint_skips_tls_for_port_80(): + fake_sock = mock.MagicMock() + with mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + return_value=fake_sock, + ), mock.patch( + "snowflake.cli._plugins.connection.diagnostic._probe_tls" + ) as probe_tls: + result = check_endpoint("crl.example.com", 80, "CRL_DISTRIBUTION_POINT") + probe_tls.assert_not_called() + assert result.status == "Healthy" + assert result.cert_issuer is None + assert result.cert_expires is None + + +def test_diagnostic_report_counts(): + report = DiagnosticReport( + checks=[ + EndpointCheck("a", 443, "X", "Healthy"), + EndpointCheck("b", 443, "X", "Healthy"), + EndpointCheck("c", 443, "X", "Unhealthy"), + EndpointCheck("d", 443, "X", "Skipped"), + ] + ) + assert report.healthy == 2 + assert report.unhealthy == 1 + assert report.skipped == 1 + assert report.tested == 3 + assert ( + report.summary_line() + == "Results: 2 Healthy, 1 Unhealthy out of 3 endpoints. 1 skipped (non-resolvable patterns)." + ) + + +def test_run_diagnostic_streams_per_endpoint(monkeypatch): + cursor = mock.MagicMock() + cursor.fetchone.return_value = {"SYSTEM$ALLOWLIST()": json.dumps(ALLOWLIST_FIXTURE)} + conn = mock.MagicMock() + conn.execute_string.return_value = (cursor,) + + fake_sock = mock.MagicMock() + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + lambda *a, **kw: fake_sock, + ) + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + lambda sock, host: ("DigiCert Inc", "Jan 19 23:59:59 2027 GMT"), + ) + + streamed = [] + report = run_diagnostic(conn, allowlist_path=None, on_check=streamed.append) + + assert len(streamed) == 3 + assert report.healthy == 2 + assert report.skipped == 1 + assert [c.type for c in streamed] == [ + "SNOWFLAKE_DEPLOYMENT", + "OCSP_RESPONDER", + "STAGE", + ] + + +def test_run_diagnostic_falls_back_when_allowlist_query_fails(monkeypatch): + conn = mock.MagicMock() + conn.host = "fallback.snowflakecomputing.com" + conn.execute_string.side_effect = RuntimeError( + "permission denied on SYSTEM$ALLOWLIST" + ) + + fake_sock = mock.MagicMock() + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + lambda *a, **kw: fake_sock, + ) + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + lambda sock, host: (None, None), + ) + + report = run_diagnostic(conn, allowlist_path=None) + assert report.tested == 1 + assert report.checks[0].host == "fallback.snowflakecomputing.com" + assert report.checks[0].type == "SNOWFLAKE_DEPLOYMENT" + + +def test_run_diagnostic_normalises_malformed_entries(monkeypatch): + cursor = mock.MagicMock() + cursor.fetchone.return_value = { + "SYSTEM$ALLOWLIST()": json.dumps( + [ + {"type": "X", "host": "good.com", "port": 443}, + {"type": "X", "host": None, "port": 443}, + {"type": "X", "host": "bad.com", "port": "not-a-number"}, + ] + ) + } + conn = mock.MagicMock() + conn.execute_string.return_value = (cursor,) + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + lambda *a, **kw: mock.MagicMock(), + ) + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + lambda *a, **kw: (None, None), + ) + report = run_diagnostic(conn, allowlist_path=None) + assert [c.host for c in report.checks] == ["good.com"] + + +def _scripted_cursor(rows): + cur = mock.MagicMock() + cur.__iter__ = lambda self: iter(rows) + return cur + + +def test_collect_network_policy_user_overrides_account(): + conn = mock.MagicMock() + + def execute_string(sql, cursor_class=None): + if "CURRENT_IP_ADDRESS" in sql: + return (_scripted_cursor([{"IP": "1.2.3.4"}]),) + if "NETWORK_POLICY' IN ACCOUNT" in sql: + return (_scripted_cursor([{"value": "ACCT_POLICY"}]),) + if "NETWORK_POLICY' FOR USER" in sql: + return (_scripted_cursor([{"value": "USER_POLICY"}]),) + if "DESC NETWORK POLICY" in sql: + return ( + _scripted_cursor( + [ + {"name": "ALLOWED_IP_LIST", "value": "0.0.0.0/0"}, + {"name": "BLOCKED_IP_LIST", "value": ""}, + {"name": "ALLOWED_NETWORK_RULE_LIST", "value": "DB.SCHEMA.NR1"}, + {"name": "BLOCKED_NETWORK_RULE_LIST", "value": ""}, + ] + ), + ) + if "DESC NETWORK RULE" in sql: + return ( + _scripted_cursor( + [ + { + "mode": "INGRESS", + "type": "IPV4", + "value_list": "1.0.0.0/8,2.0.0.0/8", + } + ] + ), + ) + return (_scripted_cursor([]),) + + conn.execute_string.side_effect = execute_string + snap = collect_network_policy(conn, user="alice") + assert snap.current_ip == "1.2.3.4" + assert snap.account_policy == "ACCT_POLICY" + assert snap.user_policy == "USER_POLICY" + assert snap.effective_policy == "USER_POLICY" + assert snap.allowed_ip_list == ["0.0.0.0/0"] + assert snap.allowed_rule_list == ["DB.SCHEMA.NR1"] + assert len(snap.rules) == 1 + assert snap.rules[0].values == ["1.0.0.0/8", "2.0.0.0/8"] + assert snap.rules[0].mode == "INGRESS" + assert snap.rules[0].type == "IPV4" + + +def test_collect_network_policy_handles_no_policy(): + conn = mock.MagicMock() + conn.execute_string.return_value = (_scripted_cursor([]),) + snap = collect_network_policy(conn, user="alice") + assert snap.has_policy() is False + assert snap.effective_policy is None + + +def test_collect_network_policy_survives_query_errors(): + conn = mock.MagicMock() + conn.execute_string.side_effect = RuntimeError("insufficient privileges") + snap = collect_network_policy(conn, user="alice") + assert isinstance(snap, NetworkPolicySnapshot) + assert snap.has_policy() is False + + +def test_resolve_cafile_prefers_requests_ca_bundle(tmp_path, monkeypatch): + from snowflake.cli._plugins.connection.diagnostic import _resolve_cafile + + primary = tmp_path / "primary.pem" + primary.write_text("primary") + secondary = tmp_path / "secondary.pem" + secondary.write_text("secondary") + monkeypatch.setenv("REQUESTS_CA_BUNDLE", str(primary)) + monkeypatch.setenv("SSL_CERT_FILE", str(secondary)) + assert _resolve_cafile() == str(primary) + + +def test_resolve_cafile_falls_back_to_ssl_cert_file(tmp_path, monkeypatch): + from snowflake.cli._plugins.connection.diagnostic import _resolve_cafile + + bundle = tmp_path / "bundle.pem" + bundle.write_text("bundle") + monkeypatch.delenv("REQUESTS_CA_BUNDLE", raising=False) + monkeypatch.setenv("SSL_CERT_FILE", str(bundle)) + assert _resolve_cafile() == str(bundle) + + +def test_resolve_cafile_returns_none_when_env_points_at_missing_file(monkeypatch): + from snowflake.cli._plugins.connection.diagnostic import _resolve_cafile + + monkeypatch.setenv("REQUESTS_CA_BUNDLE", "/nonexistent/ca.pem") + monkeypatch.delenv("SSL_CERT_FILE", raising=False) + assert _resolve_cafile() is None + + +def test_resolve_cafile_returns_none_when_unset(monkeypatch): + from snowflake.cli._plugins.connection.diagnostic import _resolve_cafile + + monkeypatch.delenv("REQUESTS_CA_BUNDLE", raising=False) + monkeypatch.delenv("SSL_CERT_FILE", raising=False) + assert _resolve_cafile() is None + + +def test_probe_tls_passes_cafile_to_ssl_context(tmp_path, monkeypatch): + """_probe_tls plumbs the resolved CA bundle into ssl.create_default_context.""" + bundle = tmp_path / "bundle.pem" + bundle.write_text("bundle") + monkeypatch.setenv("REQUESTS_CA_BUNDLE", str(bundle)) + + fake_context = mock.MagicMock() + fake_ssock = mock.MagicMock() + fake_ssock.__enter__ = lambda self: fake_ssock + fake_ssock.__exit__ = lambda self, *a: False + fake_ssock.getpeercert.return_value = {} + fake_context.wrap_socket.return_value = fake_ssock + + with mock.patch( + "snowflake.cli._plugins.connection.diagnostic.ssl.create_default_context", + return_value=fake_context, + ) as create_default_context: + from snowflake.cli._plugins.connection.diagnostic import _probe_tls + + _probe_tls(mock.MagicMock(), "x.com") + create_default_context.assert_called_once_with(cafile=str(bundle)) From 6aa2a814194ed09ae185b3a109bd3df68834fe45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jan=20B=C3=ADlek?= Date: Fri, 5 Jun 2026 11:11:48 +0200 Subject: [PATCH 2/5] SNOW-3194269: simplify diagnostic module and CLI wiring Apply A Philosophy of Software Design simplifications to the SnowCD migration code: diagnostic.py - Drop pass-through to_dict() on EndpointCheck and NetworkPolicySnapshot; use dataclasses.asdict() at call sites. - Rename NetworkPolicySnapshot.{allowed,blocked}_rule_list to {allowed,blocked}_network_rule_list to match the JSON contract directly, eliminating the rename layer. - Unify _query_allowlist (returned []) and _safe_query (returned None) into _query_json_list and _query_rows with one consistent return shape. - Replace exception-based fallback in run_diagnostic with empty-list fallback: load_allowlist always returns a list; run_diagnostic falls back to (conn.host, 443) when empty + no allowlist file. - Drop the side-effect-and-return pattern in _describe_network_policy; return parsed dict, apply at call site. - Flatten triple-nested issuer extraction into _issuer_org_name helper. - Normalise DESC row keys once at the top of _query_rows instead of case-insensitive lookups inline at every site. - Tighten is_resolvable to one expression. commands.py - Split _connection_test_with_diag into orchestrator + pure _diagnostic_table_results presentation function. - Replace _stream closure with a one-line conditional lambda. - Use asdict() for the JSON payload; drop the diagnostic_payload temp. All 29 new + 96 existing connection tests still pass. --- .../cli/_plugins/connection/commands.py | 71 ++-- .../cli/_plugins/connection/diagnostic.py | 384 ++++++++---------- tests/test_connection_diagnostic.py | 49 ++- 3 files changed, 242 insertions(+), 262 deletions(-) diff --git a/src/snowflake/cli/_plugins/connection/commands.py b/src/snowflake/cli/_plugins/connection/commands.py index 5487785004..1037751ae2 100644 --- a/src/snowflake/cli/_plugins/connection/commands.py +++ b/src/snowflake/cli/_plugins/connection/commands.py @@ -16,6 +16,7 @@ import logging import os.path +from dataclasses import asdict from pathlib import Path from typing import Optional @@ -27,8 +28,13 @@ UsageError, ) from click.core import ParameterSource # type: ignore +from snowflake.connector import ProgrammingError +from snowflake.connector.constants import CONNECTIONS_FILE + from snowflake import connector from snowflake.cli._plugins.connection.diagnostic import ( + DiagnosticReport, + NetworkPolicySnapshot, collect_network_policy, run_diagnostic, status_line, @@ -79,8 +85,6 @@ ObjectResult, ) from snowflake.cli.api.secure_path import SecurePath -from snowflake.connector import ProgrammingError -from snowflake.connector.constants import CONNECTIONS_FILE app = SnowTyperFactory( name="connection", @@ -432,54 +436,58 @@ def _connection_test_with_diag( conn_ctx, connection_summary: dict, ) -> CommandResult: - """Run the SnowCD-style per-endpoint diagnostic and assemble the full result. + """Run the SnowCD-style per-endpoint diagnostic and assemble the result. + + For TABLE output: streams `Checking : ` lines as the + run progresses, then appends a per-endpoint table, a network-policy block, + and a `Results: ...` summary line. - Streams `Checking : ` lines for human (TABLE) output and - appends a per-endpoint table plus a `Results: ...` summary line. For JSON - output the streaming is suppressed and the structured payload carries the - diagnostic data instead. + For JSON / CSV output: returns the structured diagnostic nested under a + `Diagnostic` key so consumers can read everything off one object. """ from snowflake.cli.api.output.formats import OutputFormat - from snowflake.cli.api.output.types import MultipleResults is_table_output = get_cli_context().output_format == OutputFormat.TABLE if is_table_output: cli_console.message("Fetching allowlist from Snowflake...") - - def _stream(check): - if is_table_output: - cli_console.message(status_line(check)) - report = run_diagnostic( conn=conn, allowlist_path=conn_ctx.diag_allowlist_path, - on_check=_stream, + on_check=lambda c: ( + cli_console.message(status_line(c)) if is_table_output else None + ), ) if is_table_output: cli_console.message("Inspecting network policies...") policy = collect_network_policy(conn, user=conn.user) - diagnostic_payload = { - "checks": [c.to_dict() for c in report.checks], - "healthy": report.healthy, - "unhealthy": report.unhealthy, - "skipped": report.skipped, - "tested": report.tested, - "network_policy": policy.to_dict(), - } - if not is_table_output: - # JSON / CSV output gets the structured payload nested in the - # connection summary so consumers can read everything off one object. - connection_summary["Diagnostic"] = diagnostic_payload + connection_summary["Diagnostic"] = { + "checks": [asdict(c) for c in report.checks], + "healthy": report.healthy, + "unhealthy": report.unhealthy, + "skipped": report.skipped, + "tested": report.tested, + "network_policy": asdict(policy), + } return ObjectResult(connection_summary) - # TABLE output: keep the connection summary clean. The per-endpoint - # table, network policy block, and summary line below carry the data. + return _diagnostic_table_results(connection_summary, report, policy) + + +def _diagnostic_table_results( + connection_summary: dict, + report: DiagnosticReport, + policy: NetworkPolicySnapshot, +) -> CommandResult: + """Assemble the TABLE-format output: summary, per-endpoint rows, policy block, totals.""" + from snowflake.cli.api.output.types import MultipleResults + results = MultipleResults() results.add(ObjectResult(connection_summary)) + tested_rows = [ { "url": c.host, @@ -494,6 +502,7 @@ def _stream(check): ] if tested_rows: results.add(CollectionResult(tested_rows)) + if policy.has_policy(): policy_summary = { "Effective network policy": policy.effective_policy, @@ -503,8 +512,10 @@ def _stream(check): "Current IP": policy.current_ip or "(unknown)", "Allowed IPs": ", ".join(policy.allowed_ip_list) or "(none)", "Blocked IPs": ", ".join(policy.blocked_ip_list) or "(none)", - "Allowed network rules": ", ".join(policy.allowed_rule_list) or "(none)", - "Blocked network rules": ", ".join(policy.blocked_rule_list) or "(none)", + "Allowed network rules": ", ".join(policy.allowed_network_rule_list) + or "(none)", + "Blocked network rules": ", ".join(policy.blocked_network_rule_list) + or "(none)", } if policy.error: policy_summary["Note"] = policy.error diff --git a/src/snowflake/cli/_plugins/connection/diagnostic.py b/src/snowflake/cli/_plugins/connection/diagnostic.py index 5416a9880c..de5053c05c 100644 --- a/src/snowflake/cli/_plugins/connection/diagnostic.py +++ b/src/snowflake/cli/_plugins/connection/diagnostic.py @@ -14,12 +14,12 @@ """Connectivity diagnostics for `snow connection test --enable-diag`. -Replicates the per-endpoint checks that SnowCD used to provide, sourcing -endpoints from `SYSTEM$ALLOWLIST()` (or a JSON file passed via -`--diag-allowlist-path`). Each resolvable endpoint is probed with a TCP -connect, and TLS port 443 endpoints additionally have their certificate -issuer and expiry recorded. +Replicates the per-endpoint checks SnowCD used to provide. Endpoints come from +`SYSTEM$ALLOWLIST()` (or a JSON file via `--diag-allowlist-path`); each one is +probed with TCP connect, plus a TLS handshake on port 443 to capture the cert +issuer and expiry. """ + from __future__ import annotations import json @@ -28,19 +28,21 @@ import socket import ssl import time -from dataclasses import asdict, dataclass, field +from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Iterable, Literal, Optional from click.exceptions import ClickException -from snowflake.cli._plugins.connection.util import ALLOWLIST_QUERY from snowflake.connector import SnowflakeConnection from snowflake.connector.cursor import DictCursor +from snowflake.cli._plugins.connection.util import ALLOWLIST_QUERY + log = logging.getLogger(__name__) DEFAULT_TIMEOUT_SECONDS: float = 5.0 TLS_PORT: int = 443 +ALLOWLIST_PRIVATELINK_QUERY = "SELECT SYSTEM$ALLOWLIST_PRIVATELINK()" Status = Literal["Healthy", "Unhealthy", "Skipped"] @@ -49,32 +51,32 @@ class EndpointCheck: host: str port: int - type: str # noqa: A003 — intentional public field name, matches `SYSTEM$ALLOWLIST()` JSON + type: str # noqa: A003 — public field name; matches `SYSTEM$ALLOWLIST()` JSON status: Status error: Optional[str] = None cert_issuer: Optional[str] = None cert_expires: Optional[str] = None latency_ms: Optional[float] = None - def to_dict(self) -> dict[str, Any]: - return asdict(self) - @dataclass class DiagnosticReport: checks: list[EndpointCheck] = field(default_factory=list) + def _count(self, status: Status) -> int: + return sum(1 for c in self.checks if c.status == status) + @property def healthy(self) -> int: - return sum(1 for c in self.checks if c.status == "Healthy") + return self._count("Healthy") @property def unhealthy(self) -> int: - return sum(1 for c in self.checks if c.status == "Unhealthy") + return self._count("Unhealthy") @property def skipped(self) -> int: - return sum(1 for c in self.checks if c.status == "Skipped") + return self._count("Skipped") @property def tested(self) -> int: @@ -88,28 +90,28 @@ def summary_line(self) -> str: ) -ALLOWLIST_PRIVATELINK_QUERY = "SELECT SYSTEM$ALLOWLIST_PRIVATELINK()" +# --------------------------------------------------------------------------- # +# Allowlist loading +# --------------------------------------------------------------------------- # -def _query_allowlist(conn: SnowflakeConnection, sql: str, key: str) -> list[dict]: - """Run an allowlist-returning system function and parse its single JSON cell. +def _query_json_list(conn: SnowflakeConnection, sql: str, key: str) -> list[dict]: + """Run an allowlist-returning system function; return its JSON list, or [] on any failure. - Returns `[]` on any failure (permission denied, function not present in - the deployment, malformed JSON). The caller must not let this raise. + Failures are silent because either query may be unavailable: low-priv roles + can't call `SYSTEM$ALLOWLIST()`, and `SYSTEM$ALLOWLIST_PRIVATELINK()` isn't + enabled in every deployment. The caller decides what to do with an empty + result (typically: fall back to the connection host). """ try: *_, cursor = conn.execute_string(sql, cursor_class=DictCursor) row = cursor.fetchone() - if not row: - return [] - raw = row.get(key) - if not raw: - return [] - parsed = json.loads(raw) - return parsed if isinstance(parsed, list) else [] + raw = row.get(key) if row else None + parsed = json.loads(raw) if raw else [] except Exception: log.debug("Allowlist query failed: %s", sql, exc_info=True) return [] + return parsed if isinstance(parsed, list) else [] def _dedupe_entries(entries: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -118,10 +120,9 @@ def _dedupe_entries(entries: list[dict[str, Any]]) -> list[dict[str, Any]]: out: list[dict[str, Any]] = [] for e in entries: key = (str(e.get("type", "")), str(e.get("host", "")), e.get("port")) - if key in seen: - continue - seen.add(key) - out.append(e) + if key not in seen: + seen.add(key) + out.append(e) return out @@ -130,14 +131,11 @@ def load_allowlist( ) -> list[dict[str, Any]]: """Return the raw allowlist as a list of `{type, host, port}` dicts. - If `allowlist_path` is given, parse it as JSON and use it directly. - Otherwise merge `SYSTEM$ALLOWLIST()` and `SYSTEM$ALLOWLIST_PRIVATELINK()` - on the open connection. PrivateLink entries are appended after the public - set; duplicate `(type, host, port)` triples are dropped. Either query may - fail (low-priv role, function unavailable in deployment) — failures are - silent and we return whatever we managed to collect, with the public - allowlist's failure handled by `run_diagnostic` (fall back to the - connection host) and the privatelink failure simply skipping that source. + If `allowlist_path` is given, parse it as JSON. Otherwise merge + `SYSTEM$ALLOWLIST()` and `SYSTEM$ALLOWLIST_PRIVATELINK()` from the open + connection, deduping `(type, host, port)`. Either query can fail silently + (low-priv role, function unavailable); the caller treats an empty result + as a signal to fall back to the connection host. """ if allowlist_path is not None: try: @@ -152,39 +150,28 @@ def load_allowlist( ) return payload - # Public allowlist: let exceptions propagate so run_diagnostic's - # fallback-to-conn.host branch fires when the role lacks privileges. - *_, cursor = conn.execute_string(ALLOWLIST_QUERY, cursor_class=DictCursor) - public = json.loads(cursor.fetchone()["SYSTEM$ALLOWLIST()"]) - privatelink = _query_allowlist( + public = _query_json_list(conn, ALLOWLIST_QUERY, "SYSTEM$ALLOWLIST()") + privatelink = _query_json_list( conn, ALLOWLIST_PRIVATELINK_QUERY, "SYSTEM$ALLOWLIST_PRIVATELINK()" ) - if not isinstance(public, list): - public = [] return _dedupe_entries([*public, *privatelink]) -def is_resolvable(host: str) -> bool: - """Return False for hostnames the OS resolver can never satisfy. +# --------------------------------------------------------------------------- # +# Endpoint probing +# --------------------------------------------------------------------------- # - Allowlist entries can include wildcard patterns (e.g. - `*.region.snowflakecomputing.com`) and bare placeholder values; those are - counted as `Skipped` rather than `Unhealthy`. - """ - if not host: - return False - if "*" in host: - return False - return True + +def is_resolvable(host: str) -> bool: + """False for hostnames the OS resolver can never satisfy (wildcards, empty).""" + return bool(host) and "*" not in host def _resolve_cafile() -> Optional[str]: - """Honour the same CA bundle env vars `snowflake-connector-python` reads. + """Honour `REQUESTS_CA_BUNDLE` then `SSL_CERT_FILE`, matching the connector. - `REQUESTS_CA_BUNDLE` wins (matches `requests` and the connector); falls - back to `SSL_CERT_FILE` (the OpenSSL convention). Returns `None` if - neither points at a readable file, in which case `ssl.create_default_context()` - will use the system trust store. + Returns `None` if neither points at a readable file, in which case + `ssl.create_default_context()` uses the system trust store. """ for var in ("REQUESTS_CA_BUNDLE", "SSL_CERT_FILE"): path = os.environ.get(var) @@ -193,30 +180,27 @@ def _resolve_cafile() -> Optional[str]: return None -def _probe_tls(sock: socket.socket, host: str) -> tuple[Optional[str], Optional[str]]: - """Wrap `sock` in TLS and return `(issuer, not_after)` from the peer cert. - - Caller is responsible for closing the underlying socket. - """ - context = ssl.create_default_context(cafile=_resolve_cafile()) - context.minimum_version = ssl.TLSVersion.TLSv1_2 - with context.wrap_socket(sock, server_hostname=host) as ssock: - cert: Any = ssock.getpeercert() or {} - issuer_raw: Any = cert.get("issuer") or () - issuer_cn: Optional[str] = None - for rdn in issuer_raw: +def _issuer_org_name(cert: dict) -> Optional[str]: + """Extract organizationName from a peer cert's issuer RDN sequence.""" + for rdn in cert.get("issuer") or (): for entry in rdn: if ( isinstance(entry, tuple) and len(entry) == 2 and entry[0] == "organizationName" ): - issuer_cn = entry[1] - break - if issuer_cn: - break + return entry[1] + return None + + +def _probe_tls(sock: socket.socket, host: str) -> tuple[Optional[str], Optional[str]]: + """Wrap `sock` in TLS; return `(issuer_org, not_after)` from the peer cert.""" + context = ssl.create_default_context(cafile=_resolve_cafile()) + context.minimum_version = ssl.TLSVersion.TLSv1_2 + with context.wrap_socket(sock, server_hostname=host) as ssock: + cert: Any = ssock.getpeercert() or {} not_after = cert.get("notAfter") - return issuer_cn, not_after if isinstance(not_after, str) else None + return _issuer_org_name(cert), not_after if isinstance(not_after, str) else None def check_endpoint( @@ -228,40 +212,30 @@ def check_endpoint( """Probe a single endpoint. Never raises; failures map to Unhealthy.""" if not is_resolvable(host): return EndpointCheck( - host=host, - port=port, - type=type_, - status="Skipped", - error="non-resolvable pattern", + host, port, type_, "Skipped", error="non-resolvable pattern" ) try: - # getaddrinfo + connect together cover DNS and TCP reachability. start = time.perf_counter() sock = socket.create_connection((host, port), timeout=timeout) latency_ms = round((time.perf_counter() - start) * 1000, 1) except OSError as exc: - return EndpointCheck( - host=host, - port=port, - type=type_, - status="Unhealthy", - error=str(exc), - ) + return EndpointCheck(host, port, type_, "Unhealthy", error=str(exc)) - issuer = expires = None try: if port == TLS_PORT: try: issuer, expires = _probe_tls(sock, host) except (ssl.SSLError, OSError) as exc: return EndpointCheck( - host=host, - port=port, - type=type_, - status="Unhealthy", + host, + port, + type_, + "Unhealthy", error=f"TLS handshake failed: {exc}", ) + else: + issuer = expires = None finally: try: sock.close() @@ -269,10 +243,10 @@ def check_endpoint( pass return EndpointCheck( - host=host, - port=port, - type=type_, - status="Healthy", + host, + port, + type_, + "Healthy", cert_issuer=issuer, cert_expires=expires, latency_ms=latency_ms, @@ -282,19 +256,17 @@ def check_endpoint( def _normalise_entries( entries: Iterable[dict[str, Any]], ) -> list[tuple[str, int, str]]: - """Pull (host, port, type) triples out of allowlist entries, dropping malformed rows.""" + """Pull `(host, port, type)` triples from allowlist entries, dropping malformed rows.""" out: list[tuple[str, int, str]] = [] for entry in entries: host = entry.get("host") - port = entry.get("port", TLS_PORT) - type_ = entry.get("type", "UNKNOWN") if not isinstance(host, str): continue try: - port_int = int(port) + port = int(entry.get("port", TLS_PORT)) except (TypeError, ValueError): continue - out.append((host, port_int, str(type_))) + out.append((host, port, str(entry.get("type", "UNKNOWN")))) return out @@ -306,27 +278,22 @@ def run_diagnostic( ) -> DiagnosticReport: """Load the allowlist, probe each entry, return a full report. - Calls `on_check(...)` once per endpoint in input order; this is how the CLI - streams `Checking : ✅` lines as the run progresses. + Calls `on_check(...)` once per endpoint in input order so the CLI can + stream `Checking : ✅` lines as the run progresses. - If `SYSTEM$ALLOWLIST()` cannot be queried (no file path and the call - fails — typically a permission error for low-privilege roles), the - diagnostic falls back to checking just `(conn.host, 443)`. + If the queryable allowlist comes back empty (typically permission denied + on `SYSTEM$ALLOWLIST()` for low-privilege roles) and no `allowlist_path` + was supplied, fall back to checking just `(conn.host, 443)`. """ - try: - allowlist = load_allowlist(conn, allowlist_path) - except ClickException: - raise - except Exception as exc: + allowlist = load_allowlist(conn, allowlist_path) + if not allowlist and allowlist_path is None: log.warning( - "Could not call SYSTEM$ALLOWLIST(); falling back to the connection host.", - exc_info=True, + "Allowlist empty (likely permission denied on SYSTEM$ALLOWLIST()); " + "falling back to the connection host." ) allowlist = [ {"type": "SNOWFLAKE_DEPLOYMENT", "host": conn.host, "port": TLS_PORT} ] - # Surface the fallback so users notice it in --enable-diag output. - log.info("Allowlist fetch failed: %s", exc) report = DiagnosticReport() for host, port, type_ in _normalise_entries(allowlist): @@ -337,16 +304,19 @@ def run_diagnostic( def status_line(check: EndpointCheck) -> str: - """Return the streaming-line representation: `Checking : `.""" + """Format a streaming line: `Checking : [ (detail)]`.""" icon = {"Healthy": "✅", "Unhealthy": "❌", "Skipped": "⏭"}[check.status] - suffix = f" ({check.error})" if check.status == "Unhealthy" and check.error else "" if check.status == "Healthy" and check.latency_ms is not None: suffix = f" ({check.latency_ms} ms)" + elif check.status == "Unhealthy" and check.error: + suffix = f" ({check.error})" + else: + suffix = "" return f"Checking {check.type}: {check.host} {icon}{suffix}" # --------------------------------------------------------------------------- # -# Network policy / IP rule diagnostic +# Network policy snapshot # --------------------------------------------------------------------------- # @@ -362,9 +332,9 @@ class NetworkRule: class NetworkPolicySnapshot: """Snapshot of network-policy state at the time of `snow connection test`. - `account_policy` and `user_policy` are the names returned by - `SHOW PARAMETERS LIKE 'NETWORK_POLICY'` at each scope; the effective - policy is `user_policy or account_policy` per Snowflake precedence rules. + `account_policy` and `user_policy` are the raw values from + `SHOW PARAMETERS LIKE 'NETWORK_POLICY'` at each scope; `effective_policy` + is `user_policy or account_policy` per Snowflake precedence. """ current_ip: Optional[str] = None @@ -373,96 +343,77 @@ class NetworkPolicySnapshot: effective_policy: Optional[str] = None allowed_ip_list: list[str] = field(default_factory=list) blocked_ip_list: list[str] = field(default_factory=list) - allowed_rule_list: list[str] = field(default_factory=list) - blocked_rule_list: list[str] = field(default_factory=list) + allowed_network_rule_list: list[str] = field(default_factory=list) + blocked_network_rule_list: list[str] = field(default_factory=list) rules: list[NetworkRule] = field(default_factory=list) error: Optional[str] = None def has_policy(self) -> bool: return self.effective_policy is not None - def to_dict(self) -> dict[str, Any]: - return { - "current_ip": self.current_ip, - "account_policy": self.account_policy, - "user_policy": self.user_policy, - "effective_policy": self.effective_policy, - "allowed_ip_list": self.allowed_ip_list, - "blocked_ip_list": self.blocked_ip_list, - "allowed_network_rule_list": self.allowed_rule_list, - "blocked_network_rule_list": self.blocked_rule_list, - "rules": [asdict(r) for r in self.rules], - "error": self.error, - } - - -def _safe_query(conn: SnowflakeConnection, sql: str) -> Optional[list[dict]]: - """Run a query through `execute_string`; return rows as dicts, or None on error. - - Network-policy diagnostics are best-effort: a low-priv role may not be - able to call `SHOW PARAMETERS` at account scope or `DESCRIBE NETWORK - POLICY`. Fail closed and let the caller surface the gap. + +def _query_rows(conn: SnowflakeConnection, sql: str) -> list[dict]: + """Run a query and return rows as case-folded dicts; `[]` on any failure. + + Network-policy diagnostics are best-effort — a low-priv role may not be + able to `SHOW PARAMETERS` at account scope or `DESC NETWORK POLICY`. Fail + closed and let the caller surface the gap. """ try: *_, cursor = conn.execute_string(sql, cursor_class=DictCursor) - return list(cursor) + return [{str(k).lower(): v for k, v in row.items()} for row in cursor] except Exception as exc: log.debug("Network-policy query failed: %s\n%s", sql, exc) - return None + return [] -def _scalar_value(rows: Optional[list[dict]], column: str) -> Optional[str]: +def _scalar_value(rows: list[dict], column: str) -> Optional[str]: if not rows: return None - val = rows[0].get(column) + val = rows[0].get(column.lower()) return str(val) if val not in (None, "") else None def _split_csv(value: Optional[str]) -> list[str]: - if not value: - return [] - return [v.strip() for v in value.split(",") if v.strip()] + return [v.strip() for v in (value or "").split(",") if v.strip()] def _describe_network_policy( - conn: SnowflakeConnection, name: str, snapshot: NetworkPolicySnapshot -) -> bool: - """Populate snapshot allow/block lists from `DESC NETWORK POLICY `. - - Returns True if the DESC succeeded (any row came back), False otherwise. - """ - rows = _safe_query(conn, f"DESC NETWORK POLICY IDENTIFIER('{name}')") + conn: SnowflakeConnection, name: str +) -> Optional[dict[str, list[str]]]: + """Return parsed allow/block lists from `DESC NETWORK POLICY`, or None on failure.""" + rows = _query_rows(conn, f"DESC NETWORK POLICY IDENTIFIER('{name}')") if not rows: - return False - by_name: dict[str, str] = {} - for row in rows: - prop = row.get("name") or row.get("NAME") - val = row.get("value") or row.get("VALUE") - if prop and val is not None: - by_name[str(prop).upper()] = str(val) - snapshot.allowed_ip_list = _split_csv(by_name.get("ALLOWED_IP_LIST")) - snapshot.blocked_ip_list = _split_csv(by_name.get("BLOCKED_IP_LIST")) - snapshot.allowed_rule_list = _split_csv(by_name.get("ALLOWED_NETWORK_RULE_LIST")) - snapshot.blocked_rule_list = _split_csv(by_name.get("BLOCKED_NETWORK_RULE_LIST")) - return True + return None + by_name = { + str(r["name"]).upper(): str(r["value"]) + for r in rows + if r.get("name") and r.get("value") is not None + } + return { + "allowed_ip_list": _split_csv(by_name.get("ALLOWED_IP_LIST")), + "blocked_ip_list": _split_csv(by_name.get("BLOCKED_IP_LIST")), + "allowed_network_rule_list": _split_csv( + by_name.get("ALLOWED_NETWORK_RULE_LIST") + ), + "blocked_network_rule_list": _split_csv( + by_name.get("BLOCKED_NETWORK_RULE_LIST") + ), + } def _describe_network_rule( conn: SnowflakeConnection, qualified_name: str ) -> Optional[NetworkRule]: - rows = _safe_query(conn, f"DESC NETWORK RULE IDENTIFIER('{qualified_name}')") + rows = _query_rows(conn, f"DESC NETWORK RULE IDENTIFIER('{qualified_name}')") if not rows: return None head = rows[0] - keys = {k.lower(): k for k in head.keys()} - mode = head.get(keys.get("mode", ""), "") or "" - type_ = head.get(keys.get("type", ""), "") or "" - raw_values = head.get(keys.get("value_list", ""), "") or "" return NetworkRule( name=qualified_name, - mode=str(mode), - type=str(type_), - values=_split_csv(str(raw_values)), + mode=str(head.get("mode") or ""), + type=str(head.get("type") or ""), + values=_split_csv(str(head.get("value_list") or "")), ) @@ -471,47 +422,54 @@ def collect_network_policy( ) -> NetworkPolicySnapshot: """Best-effort snapshot of the active network policy and its referenced rules. - Sources, in order of precedence: - 1. `SHOW PARAMETERS LIKE 'NETWORK_POLICY' FOR USER ` (user-level wins) + Order of precedence: + 1. `SHOW PARAMETERS LIKE 'NETWORK_POLICY' FOR USER ` (user wins) 2. `SHOW PARAMETERS LIKE 'NETWORK_POLICY' IN ACCOUNT` - 3. `DESC NETWORK POLICY ` for the inline allow/block lists + 3. `DESC NETWORK POLICY ` for inline allow/block lists 4. `DESC NETWORK RULE ` for referenced rule values Always returns a snapshot; failures populate `snapshot.error` instead. """ snapshot = NetworkPolicySnapshot() - # CURRENT_IP_ADDRESS() is the documented function. SYSTEM$GET_CLIENT_IP - # exists in some internal deployments but is not generally available. - ip_rows = _safe_query(conn, "SELECT CURRENT_IP_ADDRESS() AS IP") - if ip_rows: - snapshot.current_ip = _scalar_value(ip_rows, "IP") - - account_rows = _safe_query(conn, "SHOW PARAMETERS LIKE 'NETWORK_POLICY' IN ACCOUNT") - snapshot.account_policy = _scalar_value(account_rows, "value") - + snapshot.current_ip = _scalar_value( + _query_rows(conn, "SELECT CURRENT_IP_ADDRESS() AS IP"), "IP" + ) + snapshot.account_policy = _scalar_value( + _query_rows(conn, "SHOW PARAMETERS LIKE 'NETWORK_POLICY' IN ACCOUNT"), "value" + ) if user: - user_rows = _safe_query( - conn, - f"SHOW PARAMETERS LIKE 'NETWORK_POLICY' FOR USER IDENTIFIER('{user}')", + snapshot.user_policy = _scalar_value( + _query_rows( + conn, + f"SHOW PARAMETERS LIKE 'NETWORK_POLICY' FOR USER IDENTIFIER('{user}')", + ), + "value", ) - snapshot.user_policy = _scalar_value(user_rows, "value") - snapshot.effective_policy = snapshot.user_policy or snapshot.account_policy - if snapshot.effective_policy: - described = _describe_network_policy(conn, snapshot.effective_policy, snapshot) - if not described: - snapshot.error = ( - f"Could not DESC NETWORK POLICY {snapshot.effective_policy} " - "(role lacks privilege?). Allowed/blocked lists not shown." - ) - for rule_name in ( - *snapshot.allowed_rule_list, - *snapshot.blocked_rule_list, - ): - rule = _describe_network_rule(conn, rule_name) - if rule is not None: - snapshot.rules.append(rule) + if not snapshot.effective_policy: + return snapshot + + described = _describe_network_policy(conn, snapshot.effective_policy) + if described is None: + snapshot.error = ( + f"Could not DESC NETWORK POLICY {snapshot.effective_policy} " + "(role lacks privilege?). Allowed/blocked lists not shown." + ) + return snapshot + + snapshot.allowed_ip_list = described["allowed_ip_list"] + snapshot.blocked_ip_list = described["blocked_ip_list"] + snapshot.allowed_network_rule_list = described["allowed_network_rule_list"] + snapshot.blocked_network_rule_list = described["blocked_network_rule_list"] + + for rule_name in ( + *snapshot.allowed_network_rule_list, + *snapshot.blocked_network_rule_list, + ): + rule = _describe_network_rule(conn, rule_name) + if rule is not None: + snapshot.rules.append(rule) return snapshot diff --git a/tests/test_connection_diagnostic.py b/tests/test_connection_diagnostic.py index f40222e9a9..aaa0b48880 100644 --- a/tests/test_connection_diagnostic.py +++ b/tests/test_connection_diagnostic.py @@ -13,6 +13,7 @@ # limitations under the License. """Tests for `snow connection test --enable-diag` SnowCD-style diagnostics.""" + from __future__ import annotations import json @@ -24,6 +25,7 @@ import pytest from click.exceptions import ClickException + from snowflake.cli._plugins.connection.diagnostic import ( DiagnosticReport, EndpointCheck, @@ -192,12 +194,15 @@ def test_check_endpoint_marks_unhealthy_on_dns_failure(): def test_check_endpoint_marks_unhealthy_on_tls_failure(): fake_sock = mock.MagicMock() - with mock.patch( - "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", - return_value=fake_sock, - ), mock.patch( - "snowflake.cli._plugins.connection.diagnostic._probe_tls", - side_effect=ssl.SSLError("cert verify failed"), + with ( + mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + return_value=fake_sock, + ), + mock.patch( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + side_effect=ssl.SSLError("cert verify failed"), + ), ): result = check_endpoint("x.com", 443, "SNOWFLAKE_DEPLOYMENT") assert result.status == "Unhealthy" @@ -206,12 +211,15 @@ def test_check_endpoint_marks_unhealthy_on_tls_failure(): def test_check_endpoint_records_latency_and_cert(): fake_sock = mock.MagicMock() - with mock.patch( - "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", - return_value=fake_sock, - ), mock.patch( - "snowflake.cli._plugins.connection.diagnostic._probe_tls", - return_value=("DigiCert Inc", "Jan 19 23:59:59 2027 GMT"), + with ( + mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + return_value=fake_sock, + ), + mock.patch( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + return_value=("DigiCert Inc", "Jan 19 23:59:59 2027 GMT"), + ), ): result = check_endpoint("x.com", 443, "SNOWFLAKE_DEPLOYMENT") assert result.status == "Healthy" @@ -223,12 +231,15 @@ def test_check_endpoint_records_latency_and_cert(): def test_check_endpoint_skips_tls_for_port_80(): fake_sock = mock.MagicMock() - with mock.patch( - "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", - return_value=fake_sock, - ), mock.patch( - "snowflake.cli._plugins.connection.diagnostic._probe_tls" - ) as probe_tls: + with ( + mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + return_value=fake_sock, + ), + mock.patch( + "snowflake.cli._plugins.connection.diagnostic._probe_tls" + ) as probe_tls, + ): result = check_endpoint("crl.example.com", 80, "CRL_DISTRIBUTION_POINT") probe_tls.assert_not_called() assert result.status == "Healthy" @@ -380,7 +391,7 @@ def execute_string(sql, cursor_class=None): assert snap.user_policy == "USER_POLICY" assert snap.effective_policy == "USER_POLICY" assert snap.allowed_ip_list == ["0.0.0.0/0"] - assert snap.allowed_rule_list == ["DB.SCHEMA.NR1"] + assert snap.allowed_network_rule_list == ["DB.SCHEMA.NR1"] assert len(snap.rules) == 1 assert snap.rules[0].values == ["1.0.0.0/8", "2.0.0.0/8"] assert snap.rules[0].mode == "INGRESS" From 3940bd583328846fe08220db840a5befbf38e25f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jan=20B=C3=ADlek?= Date: Fri, 5 Jun 2026 12:05:11 +0200 Subject: [PATCH 3/5] SNOW-3194269: align diagnostic with project conventions - Route DESC NETWORK POLICY/RULE and SHOW PARAMETERS FOR USER through FQN.sql_identifier and identifier_to_show_like_pattern; add _safe_fqn() so a malformed server-supplied identifier no longer crashes the run. - Swap ClickException for CliError per the migration doc; drop the now-unused click import. - Read --diag-allowlist-path through SecurePath with an explicit file_size_limit_mb so the read is logged and bounded. - Sanitize every server-derived string (host, IP, policy/rule names, allow/block lists, cert metadata) before display via sanitize_for_terminal; add _safe() helper in commands.py. - Lift the inline OutputFormat / MultipleResults imports to module top. - Tests: expect CliError on bad JSON, assert IDENTIFIER('...') wrapping on the user/policy/rule SQL, and add an ANSI-stripping check on status_line for a server-controlled host. --- .../cli/_plugins/connection/commands.py | 74 +++++++++++------ .../cli/_plugins/connection/diagnostic.py | 82 ++++++++++++++----- tests/test_connection_diagnostic.py | 28 +++++-- 3 files changed, 132 insertions(+), 52 deletions(-) diff --git a/src/snowflake/cli/_plugins/connection/commands.py b/src/snowflake/cli/_plugins/connection/commands.py index 1037751ae2..c2c15f4ade 100644 --- a/src/snowflake/cli/_plugins/connection/commands.py +++ b/src/snowflake/cli/_plugins/connection/commands.py @@ -28,9 +28,6 @@ UsageError, ) from click.core import ParameterSource # type: ignore -from snowflake.connector import ProgrammingError -from snowflake.connector.constants import CONNECTIONS_FILE - from snowflake import connector from snowflake.cli._plugins.connection.diagnostic import ( DiagnosticReport, @@ -78,13 +75,18 @@ from snowflake.cli.api.config_ng.masking import mask_sensitive_value from snowflake.cli.api.console import cli_console from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, ObjectType +from snowflake.cli.api.output.formats import OutputFormat from snowflake.cli.api.output.types import ( CollectionResult, CommandResult, MessageResult, + MultipleResults, ObjectResult, ) +from snowflake.cli.api.sanitizers import sanitize_for_terminal from snowflake.cli.api.secure_path import SecurePath +from snowflake.connector import ProgrammingError +from snowflake.connector.constants import CONNECTIONS_FILE app = SnowTyperFactory( name="connection", @@ -445,8 +447,6 @@ def _connection_test_with_diag( For JSON / CSV output: returns the structured diagnostic nested under a `Diagnostic` key so consumers can read everything off one object. """ - from snowflake.cli.api.output.formats import OutputFormat - is_table_output = get_cli_context().output_format == OutputFormat.TABLE if is_table_output: @@ -477,25 +477,36 @@ def _connection_test_with_diag( return _diagnostic_table_results(connection_summary, report, policy) +def _safe(value: Optional[str], fallback: str = "") -> str: + """Return `value` with ANSI escapes stripped, or `fallback` when missing.""" + if value is None or value == "": + return fallback + return sanitize_for_terminal(value) or fallback + + def _diagnostic_table_results( connection_summary: dict, report: DiagnosticReport, policy: NetworkPolicySnapshot, ) -> CommandResult: - """Assemble the TABLE-format output: summary, per-endpoint rows, policy block, totals.""" - from snowflake.cli.api.output.types import MultipleResults + """Assemble the TABLE-format output: summary, per-endpoint rows, policy block, totals. + Every server-derived string (host, IP, policy/rule names, allow/block lists, + cert metadata) is run through `sanitize_for_terminal` before display so a + crafted response from `SYSTEM$ALLOWLIST()` or `DESC NETWORK POLICY` cannot + inject ANSI escape sequences into the user's terminal. + """ results = MultipleResults() results.add(ObjectResult(connection_summary)) tested_rows = [ { - "url": c.host, - "type": c.type, + "url": _safe(c.host), + "type": _safe(c.type), "status": c.status, "latency_ms": c.latency_ms if c.latency_ms is not None else "", - "issuer": c.cert_issuer or "", - "cert_expires": c.cert_expires or "", + "issuer": _safe(c.cert_issuer), + "cert_expires": _safe(c.cert_expires), } for c in report.checks if c.status != "Skipped" @@ -505,30 +516,43 @@ def _diagnostic_table_results( if policy.has_policy(): policy_summary = { - "Effective network policy": policy.effective_policy, + "Effective network policy": _safe(policy.effective_policy), "Source": "user" if policy.user_policy else "account", - "Account-level": policy.account_policy or "(none)", - "User-level": policy.user_policy or "(none)", - "Current IP": policy.current_ip or "(unknown)", - "Allowed IPs": ", ".join(policy.allowed_ip_list) or "(none)", - "Blocked IPs": ", ".join(policy.blocked_ip_list) or "(none)", - "Allowed network rules": ", ".join(policy.allowed_network_rule_list) + "Account-level": _safe(policy.account_policy, "(none)"), + "User-level": _safe(policy.user_policy, "(none)"), + "Current IP": _safe(policy.current_ip, "(unknown)"), + "Allowed IPs": ", ".join( + sanitize_for_terminal(v) or "" for v in policy.allowed_ip_list + ) + or "(none)", + "Blocked IPs": ", ".join( + sanitize_for_terminal(v) or "" for v in policy.blocked_ip_list + ) or "(none)", - "Blocked network rules": ", ".join(policy.blocked_network_rule_list) + "Allowed network rules": ", ".join( + sanitize_for_terminal(v) or "" for v in policy.allowed_network_rule_list + ) + or "(none)", + "Blocked network rules": ", ".join( + sanitize_for_terminal(v) or "" for v in policy.blocked_network_rule_list + ) or "(none)", } if policy.error: - policy_summary["Note"] = policy.error + # `error` interpolates the server-supplied policy name. + policy_summary["Note"] = _safe(policy.error) results.add(ObjectResult(policy_summary)) if policy.rules: results.add( CollectionResult( [ { - "rule": r.name, - "mode": r.mode, - "type": r.type, - "values": ", ".join(r.values), + "rule": _safe(r.name), + "mode": _safe(r.mode), + "type": _safe(r.type), + "values": ", ".join( + sanitize_for_terminal(v) or "" for v in r.values + ), } for r in policy.rules ] @@ -537,7 +561,7 @@ def _diagnostic_table_results( elif policy.current_ip: results.add( MessageResult( - f"No network policy in effect (current IP: {policy.current_ip})." + f"No network policy in effect (current IP: {_safe(policy.current_ip)})." ) ) results.add(MessageResult(report.summary_line())) diff --git a/src/snowflake/cli/_plugins/connection/diagnostic.py b/src/snowflake/cli/_plugins/connection/diagnostic.py index de5053c05c..2d965d2790 100644 --- a/src/snowflake/cli/_plugins/connection/diagnostic.py +++ b/src/snowflake/cli/_plugins/connection/diagnostic.py @@ -32,17 +32,22 @@ from pathlib import Path from typing import Any, Callable, Iterable, Literal, Optional -from click.exceptions import ClickException +from snowflake.cli._plugins.connection.util import ALLOWLIST_QUERY +from snowflake.cli.api.exceptions import CliError +from snowflake.cli.api.identifiers import FQN +from snowflake.cli.api.project.util import identifier_to_show_like_pattern +from snowflake.cli.api.sanitizers import sanitize_for_terminal +from snowflake.cli.api.secure_path import SecurePath from snowflake.connector import SnowflakeConnection from snowflake.connector.cursor import DictCursor -from snowflake.cli._plugins.connection.util import ALLOWLIST_QUERY - log = logging.getLogger(__name__) DEFAULT_TIMEOUT_SECONDS: float = 5.0 TLS_PORT: int = 443 ALLOWLIST_PRIVATELINK_QUERY = "SELECT SYSTEM$ALLOWLIST_PRIVATELINK()" +ALLOWLIST_FILE_SIZE_LIMIT_MB: float = 1.0 +NETWORK_POLICY_PARAM_NAME = "NETWORK_POLICY" Status = Literal["Healthy", "Unhealthy", "Skipped"] @@ -139,14 +144,20 @@ def load_allowlist( """ if allowlist_path is not None: try: - payload = json.loads(Path(allowlist_path).read_text()) + payload = json.loads( + SecurePath(allowlist_path).read_text( + file_size_limit_mb=ALLOWLIST_FILE_SIZE_LIMIT_MB + ) + ) except (OSError, json.JSONDecodeError) as exc: - raise ClickException( - f"Could not read allowlist file {allowlist_path}: {exc}" + raise CliError( + f"Could not read allowlist file {allowlist_path}: {exc}. " + "Check the path and that the file is valid JSON." ) if not isinstance(payload, list): - raise ClickException( - f"Allowlist file {allowlist_path} must contain a JSON array." + raise CliError( + f"Allowlist file {allowlist_path} must contain a JSON array " + "of {type, host, port} entries." ) return payload @@ -304,15 +315,22 @@ def run_diagnostic( def status_line(check: EndpointCheck) -> str: - """Format a streaming line: `Checking : [ (detail)]`.""" + """Format a streaming line: `Checking : [ (detail)]`. + + Server-derived fields (`host`, `type`) are sanitised before display so + crafted ANSI sequences in an allowlist response cannot rewrite the + terminal. + """ icon = {"Healthy": "✅", "Unhealthy": "❌", "Skipped": "⏭"}[check.status] + host = sanitize_for_terminal(check.host) or "" + type_ = sanitize_for_terminal(check.type) or "" if check.status == "Healthy" and check.latency_ms is not None: suffix = f" ({check.latency_ms} ms)" elif check.status == "Unhealthy" and check.error: - suffix = f" ({check.error})" + suffix = f" ({sanitize_for_terminal(check.error)})" else: suffix = "" - return f"Checking {check.type}: {check.host} {icon}{suffix}" + return f"Checking {type_}: {host} {icon}{suffix}" # --------------------------------------------------------------------------- # @@ -378,11 +396,23 @@ def _split_csv(value: Optional[str]) -> list[str]: return [v.strip() for v in (value or "").split(",") if v.strip()] +def _safe_fqn(name: str) -> Optional[FQN]: + """Parse an identifier, returning `None` if it fails the FQN regex.""" + try: + return FQN.from_string(name) + except Exception: + log.debug("Could not parse identifier as FQN: %r", name) + return None + + def _describe_network_policy( conn: SnowflakeConnection, name: str ) -> Optional[dict[str, list[str]]]: """Return parsed allow/block lists from `DESC NETWORK POLICY`, or None on failure.""" - rows = _query_rows(conn, f"DESC NETWORK POLICY IDENTIFIER('{name}')") + fqn = _safe_fqn(name) + if fqn is None: + return None + rows = _query_rows(conn, f"DESC NETWORK POLICY {fqn.sql_identifier}") if not rows: return None by_name = { @@ -405,7 +435,10 @@ def _describe_network_policy( def _describe_network_rule( conn: SnowflakeConnection, qualified_name: str ) -> Optional[NetworkRule]: - rows = _query_rows(conn, f"DESC NETWORK RULE IDENTIFIER('{qualified_name}')") + fqn = _safe_fqn(qualified_name) + if fqn is None: + return None + rows = _query_rows(conn, f"DESC NETWORK RULE {fqn.sql_identifier}") if not rows: return None head = rows[0] @@ -435,17 +468,22 @@ def collect_network_policy( snapshot.current_ip = _scalar_value( _query_rows(conn, "SELECT CURRENT_IP_ADDRESS() AS IP"), "IP" ) + policy_pattern = identifier_to_show_like_pattern(NETWORK_POLICY_PARAM_NAME) snapshot.account_policy = _scalar_value( - _query_rows(conn, "SHOW PARAMETERS LIKE 'NETWORK_POLICY' IN ACCOUNT"), "value" + _query_rows(conn, f"SHOW PARAMETERS LIKE {policy_pattern} IN ACCOUNT"), + "value", ) - if user: - snapshot.user_policy = _scalar_value( - _query_rows( - conn, - f"SHOW PARAMETERS LIKE 'NETWORK_POLICY' FOR USER IDENTIFIER('{user}')", - ), - "value", - ) + if isinstance(user, str) and user.strip(): + user_fqn = _safe_fqn(user) + if user_fqn is not None: + snapshot.user_policy = _scalar_value( + _query_rows( + conn, + f"SHOW PARAMETERS LIKE {policy_pattern} " + f"FOR USER {user_fqn.sql_identifier}", + ), + "value", + ) snapshot.effective_policy = snapshot.user_policy or snapshot.account_policy if not snapshot.effective_policy: diff --git a/tests/test_connection_diagnostic.py b/tests/test_connection_diagnostic.py index aaa0b48880..3ee7c9acf1 100644 --- a/tests/test_connection_diagnostic.py +++ b/tests/test_connection_diagnostic.py @@ -24,8 +24,6 @@ from unittest import mock import pytest -from click.exceptions import ClickException - from snowflake.cli._plugins.connection.diagnostic import ( DiagnosticReport, EndpointCheck, @@ -37,6 +35,7 @@ run_diagnostic, status_line, ) +from snowflake.cli.api.exceptions import CliError ALLOWLIST_FIXTURE = [ { @@ -81,6 +80,16 @@ def test_status_line_includes_error_for_unhealthy(): assert "DNS fail" in line +def test_status_line_strips_ansi_from_server_host(): + """Server-supplied host must not be able to inject ANSI escapes into the terminal.""" + malicious = "\x1b[31mevil.com\x1b[0m" + line = status_line( + EndpointCheck(malicious, 443, "SNOWFLAKE_DEPLOYMENT", "Healthy", latency_ms=1.0) + ) + assert "\x1b" not in line + assert "evil.com" in line + + def test_load_allowlist_from_file(): with NamedTemporaryFile("w+", suffix=".json", delete=False) as f: f.write(json.dumps(ALLOWLIST_FIXTURE)) @@ -95,7 +104,7 @@ def test_load_allowlist_file_invalid_json_raises(): with NamedTemporaryFile("w+", suffix=".json", delete=False) as f: f.write("not json") path = Path(f.name) - with pytest.raises(ClickException): + with pytest.raises(CliError): load_allowlist(mock.MagicMock(), path) @@ -351,13 +360,15 @@ def _scripted_cursor(rows): def test_collect_network_policy_user_overrides_account(): conn = mock.MagicMock() + seen_sql: list[str] = [] def execute_string(sql, cursor_class=None): + seen_sql.append(sql) if "CURRENT_IP_ADDRESS" in sql: return (_scripted_cursor([{"IP": "1.2.3.4"}]),) - if "NETWORK_POLICY' IN ACCOUNT" in sql: + if "SHOW PARAMETERS" in sql and "IN ACCOUNT" in sql: return (_scripted_cursor([{"value": "ACCT_POLICY"}]),) - if "NETWORK_POLICY' FOR USER" in sql: + if "SHOW PARAMETERS" in sql and "FOR USER" in sql: return (_scripted_cursor([{"value": "USER_POLICY"}]),) if "DESC NETWORK POLICY" in sql: return ( @@ -396,6 +407,13 @@ def execute_string(sql, cursor_class=None): assert snap.rules[0].values == ["1.0.0.0/8", "2.0.0.0/8"] assert snap.rules[0].mode == "INGRESS" assert snap.rules[0].type == "IPV4" + # Identifiers must be wrapped in IDENTIFIER('...') — never bare-interpolated. + user_sql = next(s for s in seen_sql if "FOR USER" in s) + assert "IDENTIFIER('alice')" in user_sql + desc_sql = next(s for s in seen_sql if "DESC NETWORK POLICY" in s) + assert "IDENTIFIER('USER_POLICY')" in desc_sql + rule_sql = next(s for s in seen_sql if "DESC NETWORK RULE" in s) + assert "IDENTIFIER('DB.SCHEMA.NR1')" in rule_sql def test_collect_network_policy_handles_no_policy(): From 9063d37397e90d43ad05b719ba8db682b7d2a654 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jan=20B=C3=ADlek?= Date: Fri, 5 Jun 2026 12:23:08 +0200 Subject: [PATCH 4/5] SNOW-3194269: surface connector probe before silent wait MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `--enable-diag` triggers ConnectionDiagnostic inside snowflake- connector-python during connect() — a silent multi-second probe of every endpoint that ran before our own streaming output started. Print "Running connector connectivity probe (this may take a moment)..." ahead of the connection so the wait is explained. TABLE output only; JSON/CSV stay clean. --- src/snowflake/cli/_plugins/connection/commands.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/snowflake/cli/_plugins/connection/commands.py b/src/snowflake/cli/_plugins/connection/commands.py index c2c15f4ade..0001b54bb4 100644 --- a/src/snowflake/cli/_plugins/connection/commands.py +++ b/src/snowflake/cli/_plugins/connection/commands.py @@ -391,6 +391,17 @@ def test( # Test connection cli_context = get_cli_context() + conn_ctx = cli_context.connection_context + + # The connector's `enable_diag=True` runs ConnectionDiagnostic during + # connect() — a silent multi-second probe of every endpoint. Surface a + # heads-up so the wait isn't unexplained; our own per-endpoint stream + # comes later via `_connection_test_with_diag`. + if conn_ctx.enable_diag and cli_context.output_format == OutputFormat.TABLE: + cli_console.message( + "Running connector connectivity probe (this may take a moment)..." + ) + conn = cli_context.connection # Test session attributes @@ -412,7 +423,6 @@ def test( except ProgrammingError as err: raise ClickException(str(err)) - conn_ctx = cli_context.connection_context result = { "Connection name": conn_ctx.connection_name, "Status": "OK", From 819f183aa7b3b0e622deb0151cac7d3ff9827663 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jan=20B=C3=ADlek?= Date: Fri, 5 Jun 2026 12:27:22 +0200 Subject: [PATCH 5/5] SNOW-3194269: tighten connector-probe heads-up message --- src/snowflake/cli/_plugins/connection/commands.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/snowflake/cli/_plugins/connection/commands.py b/src/snowflake/cli/_plugins/connection/commands.py index 0001b54bb4..71b645efa2 100644 --- a/src/snowflake/cli/_plugins/connection/commands.py +++ b/src/snowflake/cli/_plugins/connection/commands.py @@ -398,9 +398,7 @@ def test( # heads-up so the wait isn't unexplained; our own per-endpoint stream # comes later via `_connection_test_with_diag`. if conn_ctx.enable_diag and cli_context.output_format == OutputFormat.TABLE: - cli_console.message( - "Running connector connectivity probe (this may take a moment)..." - ) + cli_console.message("Running connector connectivity probe...") conn = cli_context.connection