Added req for sdz and customer name in schema and ts files

Author: Akash Verma

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

@@ -11,8 +11,6 @@
 """
 Client to interact with BurstIQ LifeGraph APIs
 """
-import base64
-import json
 import traceback
 from datetime import datetime, timedelta
 from typing import Any, Dict, List, Optional
@@ -49,34 +47,49 @@ def __init__(self, config: BurstIQConnection):
         Args:
             config: BurstIQConnection configuration
         """
-        self.realm_name = config.realmName
-        self.username = config.username
-        self.password = config.password.get_secret_value()
-        self.client_id = getattr(config, "clientId", "burst")
-
-        self.auth_server_url = getattr(config, "authServerUrl", AUTH_SERVER_BASE)
+        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
 
-        # Extract customer and SDZ names from token after authentication
-        self.customer_name: Optional[str] = None
-        self.sdz_name: Optional[str] = None
+    def test_authenticate(self):
+        """
+        Explicitly test authentication with BurstIQ.
+        This is used during test_connection to validate credentials.
 
-        # Authenticate on initialization
+        Raises:
+            Exception: If authentication fails
+        """
         self._authenticate()
 
     def _authenticate(self):
         """Authenticate with BurstIQ and get access token"""
-        token_url = f"{self.auth_server_url}/realms/{self.realm_name}/protocol/openid-connect/token"
+        # Get configuration values
+        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:
+            raise ValueError("username is required for authentication")
+        if not password:
+            raise ValueError("password is required for authentication")
+
+        auth_server_url = getattr(self.config, "authServerUrl", AUTH_SERVER_BASE)
+        client_id = getattr(self.config, "clientId", "burst")
+        token_url = (
+            f"{auth_server_url}/realms/{realm_name}/protocol/openid-connect/token"
+        )
 
         payload = {
-            "client_id": self.client_id,
+            "client_id": client_id,
             "grant_type": "password",
-            "username": self.username,
-            "password": self.password,
+            "username": username,
+            "password": password.get_secret_value(),
         }
 
         headers = {"Content-Type": "application/x-www-form-urlencoded"}
@@ -98,69 +111,34 @@ def _authenticate(self):
                 seconds=expires_in - 60
             )  # 60s buffer
 
-            # Decode token to extract customer and SDZ names
-            self._decode_and_extract_claims()
+            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"
             )
-            logger.info(f"Customer: {self.customer_name}, SDZ: {self.sdz_name}")
+            if customer_name and sdz_name:
+                logger.info(f"Customer: {customer_name}, SDZ: {sdz_name}")
 
         except Exception as exc:
             logger.error(f"Authentication failed: {exc}")
             logger.debug(traceback.format_exc())
             raise Exception("Failed to authenticate with BurstIQ") from exc
 
-    def _decode_and_extract_claims(self):
-        """Decode JWT token and extract customer/SDZ names"""
-        if not self.access_token:
-            return
-
-        try:
-            # Decode JWT payload (second part of token)
-            parts = self.access_token.split(".")
-            if len(parts) != 3:
-                logger.warning("Invalid JWT token format")
-                return
-
-            payload = parts[1]
-            # Add padding if needed
-            padding = 4 - (len(payload) % 4)
-            if padding != 4:
-                payload += "=" * padding
-
-            decoded_bytes = base64.urlsafe_b64decode(payload)
-            decoded_token = json.loads(decoded_bytes)
-
-            # Extract customer name from issuer (realm)
-            iss = decoded_token.get("iss", "")
-            if "/realms/" in iss:
-                self.customer_name = iss.split("/realms/")[-1].strip("/")
-
-            # Extract SDZ name from audience or resource_access
-            aud = decoded_token.get("aud")
-            if aud:
-                self.sdz_name = aud if isinstance(aud, str) else aud[0]
-
-            # If not in aud, check resource_access keys
-            if not self.sdz_name:
-                resource_access = decoded_token.get("resource_access", {})
-                if resource_access:
-                    self.sdz_name = list(resource_access.keys())[0]
-
-        except Exception as exc:
-            logger.warning(f"Failed to decode token: {exc}")
-            logger.debug(traceback.format_exc())
-
     def _get_auth_header(self) -> Dict[str, str]:
         """
-        Get authentication headers with current access token
+        Get authentication headers with current access token.
+        Authenticates on first call if not already authenticated.
 
         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
-        if self.token_expires_at and datetime.now() >= self.token_expires_at:
+        elif self.token_expires_at and datetime.now() >= self.token_expires_at:
             logger.info("Access token expired, re-authenticating...")
             self._authenticate()
 
@@ -170,11 +148,14 @@ def _get_auth_header(self) -> Dict[str, str]:
             "Accept": "application/json",
         }
 
-        # Add BurstIQ-specific headers
-        if self.customer_name:
-            headers["biq_customer_name"] = self.customer_name
-        if self.sdz_name:
-            headers["biq_sdz_name"] = self.sdz_name
+        # Add BurstIQ-specific headers from config
+        customer_name = getattr(self.config, "biqCustomerName", None)
+        sdz_name = getattr(self.config, "biqSdzName", None)
+
+        if customer_name:
+            headers["biq_customer_name"] = customer_name
+        if sdz_name:
+            headers["biq_sdz_name"] = sdz_name
 
         return headers
 

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

@@ -66,6 +66,10 @@ def test_connection(
         TestConnectionResult
     """
 
+    def test_authenticate():
+        """Test authentication with BurstIQ credentials"""
+        client.test_authenticate()
+
     def test_get_dictionaries():
         """Test fetching dictionaries from BurstIQ"""
         dictionaries = client.get_dictionaries(limit=1)
@@ -75,11 +79,12 @@ def test_get_dictionaries():
     def test_get_edges():
         """Test fetching edges used for lineage"""
         edges = client.get_edges(limit=1)
-        if not edges:
-            raise ConnectionError("Failed to fetch edges from BurstIQ")
+        # Edges might not exist, so don't fail if empty
+        logger.info(f"Found {len(edges)} edges in BurstIQ")
 
     test_fn = {
-        "CheckAccess": test_get_dictionaries,
+        "CheckAccess": test_authenticate,
+        "GetDictionaries": test_get_dictionaries,
         "GetEdges": test_get_edges,
     }
 

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

@@ -430,17 +430,27 @@ def get_table_constraints(
         # Process attributes for foreign keys (reference relationships)
         for attribute in dictionary.attributes:
             if attribute.referenceDictionaryName:
-                # This is a foreign key relationship
-                table_constraints.append(
-                    TableConstraint(
-                        constraintType=ConstraintType.FOREIGN_KEY,
-                        columns=[attribute.name],
-                        referredColumns=[
-                            attribute.name
-                        ],  # Assume same column name in referenced table
-                    )
+                # Build FQN for the referred column in the referenced table
+                referred_col_fqn = fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Table,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                    schema_name=self.context.get().database_schema,
+                    table_name=attribute.referenceDictionaryName,
+                    column_name=attribute.name,  # Assuming same column name in referenced table
                 )
 
+                # Only add foreign key constraint if FQN was successfully built
+                if referred_col_fqn:
+                    table_constraints.append(
+                        TableConstraint(
+                            constraintType=ConstraintType.FOREIGN_KEY,
+                            columns=[attribute.name],
+                            referredColumns=[referred_col_fqn],
+                        )
+                    )
+
         return table_constraints if table_constraints else None
 
     def yield_table(

openmetadata-service/src/main/resources/json/data/testConnections/database/burstiq.json

@@ -6,15 +6,22 @@
   "steps": [
     {
       "name": "CheckAccess",
-      "description": "Validate that we can properly authenticate with BurstIQ using the provided credentials and access the LifeGraph API.",
-      "errorMessage": "Failed to connect to BurstIQ LifeGraph. Please validate the username, password, and realm name.",
+      "description": "Validate that we can properly authenticate with BurstIQ using the provided credentials.",
+      "errorMessage": "Failed to authenticate with BurstIQ. Please validate the username, password, and realm name are correct.",
+      "shortCircuit": true,
+      "mandatory": true
+    },
+    {
+      "name": "GetDictionaries",
+      "description": "Validate that we can fetch dictionaries (tables) from BurstIQ LifeGraph API.",
+      "errorMessage": "Failed to fetch dictionaries from BurstIQ LifeGraph. Please validate that the user has sufficient privileges to read metadata.",
       "shortCircuit": true,
       "mandatory": true
     },
     {
       "name": "GetEdges",
-      "description": "List all the Edges available to the user in the BurstIQ LifeGraph platform.",
-      "errorMessage": "Failed to fetch Edges. Please validate that the user has sufficient privileges to read metadata from BurstIQ LifeGraph.",
+      "description": "List all the Edges available to the user in the BurstIQ LifeGraph platform for lineage extraction.",
+      "errorMessage": "Failed to fetch Edges. Lineage extraction may not be available but basic metadata ingestion will still work.",
       "mandatory": false
     }
   ]

openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/burstIQConnection.json

@@ -36,6 +36,16 @@
       "description": "BurstIQ Keycloak realm name (e.g., 'ems' from https://auth.burstiq.com/realms/ems).",
       "type": "string"
     },
+    "biqSdzName": {
+      "title": "BurstIQ SDZ Name",
+      "description": "BurstIQ Secure Data Zone (SDZ) name for API requests.",
+      "type": "string"
+    },
+    "biqCustomerName": {
+      "title": "BurstIQ Customer Name",
+      "description": "BurstIQ customer name for API requests.",
+      "type": "string"
+    },
     "tableFilterPattern": {
       "title": "Table Filter Pattern",
       "description": "Regex to only include/exclude dictionaries (tables) that matches the pattern.",
@@ -51,5 +61,5 @@
     }
   },
   "additionalProperties": false,
-  "required": ["username", "password", "realmName"]
+  "required": ["username", "password", "realmName", "biqSdzName", "biqCustomerName"]
 }

openmetadata-ui/src/main/resources/ui/public/locales/en-US/Database/BurstIQ.md

@@ -25,6 +25,12 @@ BurstIQ Keycloak realm name (e.g., 'ems' from https://auth.burstiq.com/realms/em
 $$
 
 $$section
-### Table Filter Pattern $(id="tableFilterPattern")
-Regex to only include/exclude dictionaries (tables) that matches the pattern.
+### BurstIQ SDZ Name $(id="biqSdzName")
+BurstIQ Secure Data Zone (SDZ) name for API requests.
 $$
+
+$$section
+### BurstIQ Customer Name $(id="biqCustomerName")
+BurstIQ customer name for API requests.
+$$
+

openmetadata-ui/src/main/resources/ui/src/generated/api/automations/createWorkflow.ts

@@ -1500,6 +1500,14 @@ export interface ConfigObject {
      * admin tables will be fetched.
      */
     includeSystemTables?: boolean;
+    /**
+     * BurstIQ customer name for API requests.
+     */
+    biqCustomerName?: string;
+    /**
+     * BurstIQ Secure Data Zone (SDZ) name for API requests.
+     */
+    biqSdzName?: string;
     /**
      * BurstIQ Keycloak realm name (e.g., 'ems' from https://auth.burstiq.com/realms/ems).
      */

openmetadata-ui/src/main/resources/ui/src/generated/api/services/createDatabaseService.ts

@@ -993,6 +993,14 @@ export interface ConfigObject {
      * admin tables will be fetched.
      */
     includeSystemTables?: boolean;
+    /**
+     * BurstIQ customer name for API requests.
+     */
+    biqCustomerName?: string;
+    /**
+     * BurstIQ Secure Data Zone (SDZ) name for API requests.
+     */
+    biqSdzName?: string;
     /**
      * BurstIQ Keycloak realm name (e.g., 'ems' from https://auth.burstiq.com/realms/ems).
      */

openmetadata-ui/src/main/resources/ui/src/generated/api/services/ingestionPipelines/createIngestionPipeline.ts

@@ -4178,6 +4178,14 @@ export interface ConfigObject {
      * admin tables will be fetched.
      */
     includeSystemTables?: boolean;
+    /**
+     * BurstIQ customer name for API requests.
+     */
+    biqCustomerName?: string;
+    /**
+     * BurstIQ Secure Data Zone (SDZ) name for API requests.
+     */
+    biqSdzName?: string;
     /**
      * BurstIQ Keycloak realm name (e.g., 'ems' from https://auth.burstiq.com/realms/ems).
      */

openmetadata-ui/src/main/resources/ui/src/generated/entity/automations/testServiceConnection.ts

@@ -1382,6 +1382,14 @@ export interface ConfigObject {
      * admin tables will be fetched.
      */
     includeSystemTables?: boolean;
+    /**
+     * BurstIQ customer name for API requests.
+     */
+    biqCustomerName?: string;
+    /**
+     * BurstIQ Secure Data Zone (SDZ) name for API requests.
+     */
+    biqSdzName?: string;
     /**
      * BurstIQ Keycloak realm name (e.g., 'ems' from https://auth.burstiq.com/realms/ems).
      */