snowglobe-cli 0.1.0__py3-none-any.whl

snowglobe/collectors/query_history.py

@@ -0,0 +1,46 @@
+
+from typing import List
+from snowglobe.queries.query_history import get_warehouse_queries
+from snowglobe.models.query import QueryStats
+
+
+class QueryCollector:
+    """
+    Collects Snowflake query metadata.
+    """
+
+    def __init__(self, connection, days):
+        """
+        connection: SnowflakeReadOnly instance
+        """
+        self.connection = connection
+        self.days = days
+
+    def warehouse_queries(self) -> List[QueryStats]:
+        """
+        Collect queries against warehouses over X days
+        """
+        query_history: List[QueryStats] = []
+
+        with self.connection:
+            qh = get_warehouse_queries(days=self.days)
+            rows = self.connection.query(qh)
+            for q in rows:
+                result = QueryStats(
+                    query_id=q["QUERY_ID"],
+                    user_name=q["USER_NAME"],
+                    warehouse_name=q["WAREHOUSE_NAME"],
+                    warehouse_size=q["WAREHOUSE_SIZE"],
+                    query_text=q["QUERY_TEXT"],
+                    query_tag=q["QUERY_TAG"],
+                    query_type=q["QUERY_TYPE"],
+                    bytes_scanned=q["BYTES_SCANNED"],
+                    execution_time_sec=q["EXECUTION_TIME_SEC"],
+                    start_time=q["START_TIME"],
+                    warehouse_multiplier=q["WAREHOUSE_MULTIPLIER"],
+                    estimated_credits=q["ESTIMATED_CREDITS"],
+                )
+                query_history.append(result)
+
+        return query_history
+

snowglobe/collectors/query_profile.py

@@ -0,0 +1,101 @@
+import json
+import re
+from typing import List
+from snowglobe.models.query import QueryProfile
+
+# Snowflake query IDs are UUIDs in the format: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
+QUERY_ID_PATTERN = re.compile(r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$', re.IGNORECASE)
+
+
+def _validate_query_id(query_id: str) -> str:
+    """Validate query_id matches Snowflake UUID format to prevent injection."""
+    if not QUERY_ID_PATTERN.match(query_id):
+        raise ValueError(f"Invalid query ID format: '{query_id}'. Expected UUID format.")
+    return query_id
+
+
+class QueryProfileCollector:
+
+    def __init__(self, connection):
+        self.connection = connection
+
+    def fetch_sql_text(self, query_id: str):
+        query_id = _validate_query_id(query_id)
+
+        with self.connection:
+            sql = f"SELECT query_text FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY WHERE query_id = '{query_id}'"
+            result = self.connection.query(sql)
+
+        if not result:
+            raise ValueError(f"No query found with id {query_id}")
+
+        return result[0]["QUERY_TEXT"]
+
+    def fetch_query_profile(self, query_id: str) -> List[QueryProfile]:
+        query_id = _validate_query_id(query_id)
+
+        with self.connection:
+            sql = f"SELECT * FROM TABLE(GET_QUERY_OPERATOR_STATS('{query_id}'))"
+            try:
+                rows = self.connection.query(sql)
+            except Exception as e:
+                err = str(e).lower()
+                if "profile expired" in err:
+                    raise ValueError(
+                        f"Query profile for '{query_id}' has expired. "
+                        "Snowflake retains profiles for ~14 days. Try a more recent query."
+                    ) from None
+                raise
+
+            query_profile: List[QueryProfile] = []
+
+            for q in rows:
+                operator_statistics = json.loads(q["OPERATOR_STATISTICS"]) if q["OPERATOR_STATISTICS"] else {}
+                execution_time_breakdown = json.loads(q["EXECUTION_TIME_BREAKDOWN"]) if q["EXECUTION_TIME_BREAKDOWN"] else {}
+                operator_attributes = json.loads(q["OPERATOR_ATTRIBUTES"]) if q["OPERATOR_ATTRIBUTES"] else {}
+                parent_operators = json.loads(q["PARENT_OPERATORS"]) if q["PARENT_OPERATORS"] else []
+
+                result = QueryProfile(
+                    query_id=q["QUERY_ID"],
+                    step_id=q["STEP_ID"],
+                    operator_id=q["OPERATOR_ID"],
+                    parent_operators=parent_operators,
+                    operator_type=q["OPERATOR_TYPE"],
+                    operator_statistics=operator_statistics,
+                    execution_time_breakdown=execution_time_breakdown,
+                    operator_attributes=operator_attributes,
+                )
+
+                query_profile.append(result)
+
+        return query_profile
+
+    def fetch_query_insights(self, query_id: str) -> List[dict]:
+        """
+        Fetch Snowflake-native query insights from ACCOUNT_USAGE.QUERY_INSIGHTS.
+        Returns a list of insight dicts. Empty list if no insights or view unavailable.
+        Note: latency can be up to 90 minutes after query execution.
+        """
+        query_id = _validate_query_id(query_id)
+
+        try:
+            with self.connection:
+                sql = f"""
+                SELECT INSIGHT_TYPE_ID, MESSAGE, SUGGESTIONS, IS_OPPORTUNITY, INSIGHT_TOPIC
+                FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_INSIGHTS
+                WHERE QUERY_ID = '{query_id}'
+                """
+                rows = self.connection.query(sql)
+
+            insights = []
+            for row in rows:
+                insights.append({
+                    "type_id": row["INSIGHT_TYPE_ID"],
+                    "message": row["MESSAGE"],
+                    "suggestions": row["SUGGESTIONS"],
+                    "is_opportunity": row["IS_OPPORTUNITY"],
+                    "topic": row["INSIGHT_TOPIC"],
+                })
+            return insights
+        except Exception:
+            return []

snowglobe/config/loader.py

@@ -0,0 +1,42 @@
+import os
+import yaml
+from pathlib import Path
+from typing import Dict, Any
+
+class ProfileNotFound(Exception):
+    pass
+
+class SnowglobeConfig:
+    CONFIG_PATH = Path.home() / ".snowglobe" / "config.yaml"
+
+    def __init__(self):
+        self.profiles: Dict[str, Dict[str, Any]] = {}
+        self.load()
+
+    
+    def _expand_env(self, obj: Any) -> Any:
+        """Recursively expand environment variables in strings, preserving types."""
+        if isinstance(obj, str):
+            return os.path.expandvars(obj)
+        if isinstance(obj, list):
+            return [self._expand_env(x) for x in obj]
+        if isinstance(obj, dict):
+            return {k: self._expand_env(v) for k, v in obj.items()}
+        return obj
+
+    def load(self):
+        if not self.CONFIG_PATH.exists():
+            raise FileNotFoundError(f"Config file not found at {self.CONFIG_PATH}")
+        with open(self.CONFIG_PATH, "r") as f:
+            raw_profiles = yaml.safe_load(f) or {}
+
+        for profile_name, values in raw_profiles.items():
+            self.profiles[profile_name] = self._expand_env(values)
+
+    def get_profile(self, profile_name: str) -> Dict[str, Any]:
+        if profile_name not in self.profiles:
+            raise ProfileNotFound(f"Profile '{profile_name}' not found in config")
+        return self.profiles[profile_name]
+
+    def list_profiles(self):
+        return list(self.profiles.keys())