Feature: burstiq profiler (#26983)

* Feature: burstiq profiler

* Fix: Added burstiq profiler interface

* Update generated TypeScript types

* Fix: Resolved comments

* Fix: pycheckstyle

* Fixed failing test, and addressed comment

* Added sampler test

* Fixed sampler changes

* Fix: resolved comments

* Fix: Failing tests

* Fix: Added caching to reduce profiler runtime, fixed md file

* Fix: minor improvement

---------

Co-authored-by: Akash Verma <akashverma@Akashs-MacBook-Pro-2.local>
Co-authored-by: Teddy <teddy.crepineau@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: 4 people

ingestion/src/metadata/ingestion/source/database/burstiq/client.py

@@ -20,15 +20,19 @@
 from metadata.generated.schema.entity.services.connections.database.burstIQConnection import (
     BurstIQConnection,
 )
+from metadata.ingestion.source.database.burstiq.models import (
+    BurstIQDictionary,
+    BurstIQEdge,
+    SdzMetricsResponse,
+    TokenResponse,
+    TQLRecord,
+)
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
-AUTH_TIMEOUT = (10, 30)  # 10s connect, 30s read for authentication
-API_TIMEOUT = (
-    10,
-    120,
-)  # 10s connect, 120s read for API calls (handles 600+ dictionaries)
+AUTH_TIMEOUT = (10, 30)
+API_TIMEOUT = (10, 120)
 
 AUTH_SERVER_BASE = "https://auth.burstiq.com"
 API_BASE_URL = "https://api.burstiq.com"
@@ -41,18 +45,12 @@ class BurstIQClient:
     """
 
     def __init__(self, config: BurstIQConnection):
-        """
-        Initialize BurstIQ Client
-
-        Args:
-            config: BurstIQConnection configuration
-        """
         self.config = config
         self.api_base_url = getattr(config, "apiUrl", API_BASE_URL).rstrip("/")
 
-        # Token management
         self.access_token: Optional[str] = None
         self.token_expires_at: Optional[datetime] = None
+        self._chain_metrics: Optional[Dict[str, int]] = None
 
     def test_authenticate(self):
         """
@@ -65,13 +63,11 @@ def test_authenticate(self):
         self._authenticate()
 
     def _authenticate(self):
-        """Authenticate with BurstIQ and get access token"""
-        # Get configuration values
+        """Authenticate with BurstIQ and store the access token."""
         realm_name = getattr(self.config, "realmName", None)
         username = getattr(self.config, "username", None)
         password = getattr(self.config, "password", None)
 
-        # Validate required fields
         if not realm_name:
             raise ValueError("realmName is required for authentication")
         if not username:
@@ -92,30 +88,27 @@ def _authenticate(self):
             "password": password.get_secret_value(),
         }
 
-        headers = {"Content-Type": "application/x-www-form-urlencoded"}
-
         try:
             logger.info(f"Authenticating with BurstIQ at: {token_url}")
             response = requests.post(
-                token_url, data=payload, headers=headers, timeout=AUTH_TIMEOUT
+                token_url,
+                data=payload,
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+                timeout=AUTH_TIMEOUT,
             )
             response.raise_for_status()
 
-            token_data = response.json()
-
-            self.access_token = token_data.get("access_token")
-
-            # Calculate token expiration
-            expires_in = token_data.get("expires_in", 3600)
+            token = TokenResponse.model_validate(response.json())
+            self.access_token = token.access_token
             self.token_expires_at = datetime.now() + timedelta(
-                seconds=expires_in - 60
-            )  # 60s buffer
+                seconds=token.expires_in - 60
+            )
 
             customer_name = getattr(self.config, "biqCustomerName", None)
             sdz_name = getattr(self.config, "biqSdzName", None)
 
             logger.info(
-                f"Authentication successful. Token expires in {expires_in} seconds"
+                f"Authentication successful. Token expires in {token.expires_in} seconds"
             )
             if customer_name and sdz_name:
                 logger.info(f"Customer: {customer_name}, SDZ: {sdz_name}")
@@ -127,17 +120,14 @@ def _authenticate(self):
 
     def _get_auth_header(self) -> Dict[str, str]:
         """
-        Get authentication headers with current access token.
-        Authenticates on first call if not already authenticated.
+        Get authentication headers, refreshing the token if necessary.
 
         Returns:
             Dictionary of headers
         """
-        # Authenticate if not already done (lazy authentication)
         if not self.access_token:
             logger.info("No access token found, authenticating...")
             self._authenticate()
-        # Check if token needs refresh
         elif self.token_expires_at and datetime.now() >= self.token_expires_at:
             logger.info("Access token expired, re-authenticating...")
             self._authenticate()
@@ -148,20 +138,20 @@ def _get_auth_header(self) -> Dict[str, str]:
             "Accept": "application/json",
         }
 
-        # Add BurstIQ-specific headers from config
         customer_name = getattr(self.config, "biqCustomerName", None)
         sdz_name = getattr(self.config, "biqSdzName", None)
+        system_wallet_id = getattr(self.config, "biqSystemWalletId", None)
 
         if customer_name:
             headers["biq_customer_name"] = customer_name
         if sdz_name:
             headers["biq_sdz_name"] = sdz_name
+        if system_wallet_id:
+            headers["biq_system_wallet_id"] = system_wallet_id
 
         return headers
 
-    def _make_request(
-        self, method: str, endpoint: str, **kwargs
-    ) -> Optional[Dict[str, Any]]:
+    def _make_request(self, method: str, endpoint: str, **kwargs) -> Optional[Any]:
         """
         Make HTTP request to BurstIQ API
 
@@ -178,11 +168,9 @@ def _make_request(
         url = f"{self.api_base_url}/{endpoint.lstrip('/')}"
         headers = self._get_auth_header()
 
-        # Merge with any additional headers provided
         if "headers" in kwargs:
             headers.update(kwargs.pop("headers"))
 
-        # Log request params for debugging
         params = kwargs.get("params", {})
         logger.debug(f"Making {method} request to {url} with params: {params}")
 
@@ -199,14 +187,12 @@ def _make_request(
 
             response.raise_for_status()
 
-            # Parse JSON response
             json_data = response.json()
 
-            # Log response size
             if isinstance(json_data, list):
                 logger.debug(f"Received {len(json_data)} items in response")
             else:
-                logger.debug(f"Received single item response")
+                logger.debug("Received single item response")
 
             return json_data
 
@@ -229,15 +215,15 @@ def _make_request(
             logger.debug(traceback.format_exc())
             raise
 
-    def get_dictionaries(self, limit: Optional[int] = None) -> List[Dict[str, Any]]:
+    def get_dictionaries(self, limit: Optional[int] = None) -> List[BurstIQDictionary]:
         """
         Fetch all data dictionaries from BurstIQ
 
         Args:
             limit: Optional limit on number of dictionaries to fetch
 
         Returns:
-            List of dictionary objects
+            List of BurstIQDictionary model instances
         """
         params = {}
         if limit:
@@ -249,22 +235,26 @@ def get_dictionaries(self, limit: Optional[int] = None) -> List[Dict[str, Any]]:
         if data is None:
             return []
 
-        dictionaries = data if isinstance(data, list) else [data]
+        raw_items = data if isinstance(data, list) else [data]
+        dictionaries = [BurstIQDictionary.model_validate(item) for item in raw_items]
         logger.info(f"Found {len(dictionaries)} dictionaries")
         return dictionaries
 
-    def get_dictionary_by_name(self, name: str) -> Optional[Dict[str, Any]]:
+    def get_dictionary_by_name(self, name: str) -> Optional[BurstIQDictionary]:
         """
         Get a specific dictionary by name
 
         Args:
             name: Dictionary name
 
         Returns:
-            Dictionary object or None
+            BurstIQDictionary instance or None
         """
         logger.debug(f"Fetching dictionary: {name}")
-        return self._make_request("GET", f"/api/metadata/dictionary/{name}")
+        data = self._make_request("GET", f"/api/metadata/dictionary/{name}")
+        if data is None:
+            return None
+        return BurstIQDictionary.model_validate(data)
 
     def get_edges(
         self,
@@ -273,7 +263,7 @@ def get_edges(
         to_dictionary: Optional[str] = None,
         limit: Optional[int] = None,
         skip: Optional[int] = None,
-    ) -> List[Dict[str, Any]]:
+    ) -> List[BurstIQEdge]:
         """
         Query edge definitions (lineage relationships) from BurstIQ
 
@@ -285,7 +275,7 @@ def get_edges(
             skip: Optional number of edges to skip (pagination)
 
         Returns:
-            List of edge definition objects
+            List of BurstIQEdge model instances
         """
         params = {}
         if name:
@@ -307,10 +297,64 @@ def get_edges(
         if data is None:
             return []
 
-        edges = data if isinstance(data, list) else [data]
+        raw_items = data if isinstance(data, list) else [data]
+        edges = [BurstIQEdge.model_validate(item) for item in raw_items]
         logger.info(f"Found {len(edges)} edge definitions")
         return edges
 
+    def get_chain_metrics(self) -> Dict[str, int]:
+        """
+        Fetch asset counts per chain from BurstIQ metrics endpoint.
+
+        Returns:
+            Dict mapping chain name to asset (row) count
+        """
+        if self._chain_metrics is not None:
+            return self._chain_metrics
+        logger.info("Fetching chain metrics from BurstIQ...")
+        data = self._make_request("GET", "/api/metrics/sdz")
+        if data is None:
+            return {}
+        metrics = SdzMetricsResponse.model_validate(data)
+        self._chain_metrics = {
+            name: chain.assets for name, chain in metrics.chainMetrics.items()
+        }
+        return self._chain_metrics
+
+    def get_records_by_tql(
+        self, chain: str, limit: int, skip: int = 0
+    ) -> List[Dict[str, Any]]:
+        """
+        Fetch data records from a chain using TQL (Temporal Query Language).
+
+        Args:
+            chain: Chain (dictionary) name to query
+            limit: Maximum number of records to fetch
+            skip: Number of records to skip (for pagination)
+
+        Returns:
+            List of flat record dicts (data envelope unwrapped)
+        """
+        tql = f"FROM {chain} SKIP {skip} LIMIT {limit} SELECT data.*"
+        logger.info(f"Fetching records for chain '{chain}' via TQL (limit={limit})")
+        try:
+            raw = self._make_request(
+                "POST", "/api/graphchain/query", json={"query": tql}
+            )
+        except Exception as exc:
+            logger.warning(f"TQL query failed for chain '{chain}': {exc}")
+            return []
+
+        if not isinstance(raw, list):
+            return []
+
+        records = [
+            TQLRecord.model_validate(item).to_record()
+            for item in raw
+            if isinstance(item, dict)
+        ]
+        logger.info(f"Fetched {len(records)} records for chain '{chain}'")
+        return records
+
     def close(self):
-        """Cleanup method - no session to close when using plain requests"""
         pass

ingestion/src/metadata/ingestion/source/database/burstiq/connection.py

@@ -11,7 +11,8 @@
 """
 Source connection handler for BurstIQ
 """
-from typing import Optional
+import hashlib
+from typing import Dict, Optional
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -30,18 +31,24 @@
 
 logger = ingestion_logger()
 
+_CLIENT_CACHE: Dict[str, BurstIQClient] = {}
+
 
 def get_connection(connection: BurstIQConnection) -> BurstIQClient:
     """
-    Create BurstIQ client connection
-
-    Args:
-        connection: BurstIQConnection configuration
+    Create or return a cached BurstIQ client connection.
 
-    Returns:
-        BurstIQClient instance
+    Caching avoids re-authentication on every table during profiler ingestion,
+    where SamplerInterface.__init__ calls get_ssl_connection once per table.
+    Using id(connection) was unreliable because each table deserialization
+    produces a new object with a different id. A SHA-256 digest of the
+    serialised config is used as the key: collision-resistant but never
+    stores plaintext credentials in the cache keys.
     """
-    return BurstIQClient(config=connection)
+    key = hashlib.sha256(connection.model_dump_json().encode()).hexdigest()
+    if key not in _CLIENT_CACHE:
+        _CLIENT_CACHE[key] = BurstIQClient(config=connection)
+    return _CLIENT_CACHE[key]
 
 
 def test_connection(