pvw-cli 1.0.6__py3-none-any.whl → 1.0.9__py3-none-any.whl

purviewcli/client/endpoint.py

@@ -22,10 +22,15 @@ def get_data(http_dict):
         account_name = os.getenv(
             "PURVIEW_ACCOUNT_NAME", http_dict.get("account_name", "test-purview-account")
         )
+        
+        # Get account ID from environment (optional)
+        account_id = os.getenv("PURVIEW_ACCOUNT_ID")
 
         # Create config
         config = SyncPurviewConfig(
-            account_name=account_name, azure_region=os.getenv("AZURE_REGION", "public")
+            account_name=account_name, 
+            azure_region=os.getenv("AZURE_REGION", "public"),
+            account_id=account_id
         )
 
         # Create synchronous client
@@ -39,6 +44,13 @@ def get_data(http_dict):
             json=http_dict.get("payload"),
         )
 
+        # The synchronous client returns a wrapper dict like
+        # {"status": "success", "data": <json>, "status_code": 200}
+        # Normalize to return the raw JSON payload when available so
+        # calling code (which expects the API JSON) works consistently
+        if isinstance(result, dict) and result.get("status") == "success" and "data" in result:
+            return result.get("data")
+
         return result
 
     except Exception as e:

purviewcli/client/sync_client.py

@@ -13,9 +13,10 @@ from azure.core.exceptions import ClientAuthenticationError
 class SyncPurviewConfig:
     """Simple synchronous config"""
 
-    def __init__(self, account_name: str, azure_region: str = "public"):
+    def __init__(self, account_name: str, azure_region: str = "public", account_id: Optional[str] = None):
         self.account_name = account_name
         self.azure_region = azure_region
+        self.account_id = account_id  # Optional Purview account ID for UC endpoints
 
 
 class SyncPurviewClient:
@@ -24,7 +25,7 @@ class SyncPurviewClient:
     def __init__(self, config: SyncPurviewConfig):
         self.config = config
 
-        # Set up endpoints based on Azure region
+        # Set up regular Purview API endpoints based on Azure region, using account name in the URL
         if config.azure_region and config.azure_region.lower() == "china":
             self.base_url = f"https://{config.account_name}.purview.azure.cn"
             self.auth_scope = "https://purview.azure.cn/.default"
@@ -35,11 +36,49 @@ class SyncPurviewClient:
             self.base_url = f"https://{config.account_name}.purview.azure.com"
             self.auth_scope = "https://purview.azure.net/.default"
 
+        # Set up Unified Catalog endpoint using Purview account ID format
+        self.account_id = config.account_id or self._get_purview_account_id()
+        self.uc_base_url = f"https://{self.account_id}-api.purview-service.microsoft.com"
+        self.uc_auth_scope = "73c2949e-da2d-457a-9607-fcc665198967/.default"
+
         self._token = None
+        self._uc_token = None
         self._credential = None
 
-    def _get_authentication_token(self):
-        """Get Azure authentication token"""
+    def _get_purview_account_id(self):
+        """Get Purview account ID from Atlas endpoint URL"""
+        account_id = os.getenv("PURVIEW_ACCOUNT_ID")
+        if not account_id:
+            import subprocess
+            try:
+                # Get the Atlas catalog endpoint and extract account ID from it
+                result = subprocess.run([
+                    "az", "purview", "account", "show", 
+                    "--name", self.config.account_name,
+                    "--resource-group", os.getenv("PURVIEW_RESOURCE_GROUP", "fabric-artifacts"),
+                    "--query", "endpoints.catalog", 
+                    "-o", "tsv"
+                ], capture_output=True, text=True, check=True)
+                atlas_url = result.stdout.strip()
+                
+                if atlas_url and "-api.purview-service.microsoft.com" in atlas_url:
+                    account_id = atlas_url.split("://")[1].split("-api.purview-service.microsoft.com")[0]
+                else:
+                    raise Exception(f"Could not extract account ID from Atlas URL: {atlas_url}")
+            except Exception as e:
+                # For Unified Catalog, the account ID is typically the Azure Tenant ID
+                try:
+                    tenant_result = subprocess.run([
+                        "az", "account", "show", "--query", "tenantId", "-o", "tsv"
+                    ], capture_output=True, text=True, check=True)
+                    account_id = tenant_result.stdout.strip()
+                    print(f"Info: Using Tenant ID as Purview Account ID for Unified Catalog: {account_id}")
+                except Exception:
+                    raise Exception(f"Could not determine Purview account ID. For Unified Catalog, this is typically your Azure Tenant ID. Please set PURVIEW_ACCOUNT_ID environment variable. Error: {e}")
+        return account_id
+
+    def _get_authentication_token(self, for_unified_catalog=False):
+        """Get Azure authentication token for regular Purview or Unified Catalog APIs"""
         try:
             # Try different authentication methods in order of preference
 
@@ -56,10 +95,15 @@ class SyncPurviewClient:
                 # 2. Use default credential (managed identity, VS Code, CLI, etc.)
                 self._credential = DefaultAzureCredential()
 
-            # Get the token
-            token = self._credential.get_token(self.auth_scope)
-            self._token = token.token
-            return self._token
+            # Get the appropriate token based on the API type
+            if for_unified_catalog:
+                token = self._credential.get_token(self.uc_auth_scope)
+                self._uc_token = token.token
+                return self._uc_token
+            else:
+                token = self._credential.get_token(self.auth_scope)
+                self._token = token.token
+                return self._token
 
         except ClientAuthenticationError as e:
             raise Exception(f"Azure authentication failed: {str(e)}")
@@ -69,12 +113,25 @@ class SyncPurviewClient:
     def make_request(self, method: str, endpoint: str, **kwargs) -> Dict:
         """Make actual HTTP request to Microsoft Purview"""
         try:
-            # Get authentication token
-            if not self._token:
-                self._get_authentication_token()            # Prepare the request
-            url = f"{self.base_url}{endpoint}"
+            # Determine if this is a Unified Catalog request
+            is_unified_catalog = endpoint.startswith('/datagovernance/catalog')
+            
+            # Get the appropriate authentication token and base URL
+            if is_unified_catalog:
+                if not self._uc_token:
+                    self._get_authentication_token(for_unified_catalog=True)
+                token = self._uc_token
+                base_url = self.uc_base_url
+            else:
+                if not self._token:
+                    self._get_authentication_token(for_unified_catalog=False)
+                token = self._token
+                base_url = self.base_url
+            
+            # Prepare the request
+            url = f"{base_url}{endpoint}"
             headers = {
-                "Authorization": f"Bearer {self._token}",
+                "Authorization": f"Bearer {token}",
                 "Content-Type": "application/json",
                 "User-Agent": "purviewcli/2.0",
             }
@@ -89,7 +146,7 @@ class SyncPurviewClient:
                 timeout=30,
             )
             # Handle the response
-            if response.status_code == 200:
+            if response.status_code in [200, 201]:
                 try:
                     data = response.json()
                     return {"status": "success", "data": data, "status_code": response.status_code}
@@ -115,7 +172,7 @@ class SyncPurviewClient:
                     timeout=30,
                 )
 
-                if response.status_code == 200:
+                if response.status_code in [200, 201]:
                     try:
                         data = response.json()
                         return {