semantic-link-labs 0.8.11__py3-none-any.whl → 0.9.1__py3-none-any.whl

sempy_labs/tom/_model.py

@@ -20,6 +20,8 @@ import sempy_labs._icons as icons
 from sempy.fabric.exceptions import FabricHTTPException
 import ast
 from uuid import UUID
+import sempy_labs._authentication as auth
+
 
 if TYPE_CHECKING:
     import Microsoft.AnalysisServices.Tabular
@@ -44,19 +46,91 @@ class TOMWrapper:
 
     def __init__(self, dataset, workspace, readonly):
 
-        (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
-        (dataset_name, dataset_id) = resolve_dataset_name_and_id(dataset, workspace_id)
-        self._dataset_id = dataset_id
-        self._dataset_name = dataset_name
-        self._workspace_name = workspace_name
-        self._workspace_id = workspace_id
+        self._is_azure_as = False
+        prefix = "asazure"
+        prefix_full = f"{prefix}://"
+        read_write = ":rw"
+        self._token_provider = auth.token_provider.get()
+
+        # Azure AS workspace logic
+        if workspace is not None and workspace.startswith(prefix_full):
+            # Set read or read/write accordingly
+            if readonly is False and not workspace.endswith(read_write):
+                workspace += read_write
+            elif readonly is True and workspace.endswith(read_write):
+                workspace = workspace[: -len(read_write)]
+            self._workspace_name = workspace
+            self._workspace_id = workspace
+            self._dataset_id = dataset
+            self._dataset_name = dataset
+            self._is_azure_as = True
+            if self._token_provider is None:
+                raise ValueError(
+                    f"{icons.red_dot} A token provider must be provided when connecting to an Azure AS workspace."
+                )
+        else:
+            (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
+            (dataset_name, dataset_id) = resolve_dataset_name_and_id(
+                dataset, workspace_id
+            )
+            self._dataset_id = dataset_id
+            self._dataset_name = dataset_name
+            self._workspace_name = workspace_name
+            self._workspace_id = workspace_id
         self._readonly = readonly
         self._tables_added = []
 
-        self._tom_server = fabric.create_tom_server(
-            readonly=readonly, workspace=workspace_id
-        )
-        self.model = self._tom_server.Databases[dataset_id].Model
+        # No token provider (standard authentication)
+        if self._token_provider is None:
+            self._tom_server = fabric.create_tom_server(
+                readonly=readonly, workspace=workspace_id
+            )
+        # Service Principal Authentication for Azure AS via token provider
+        elif self._is_azure_as:
+            import Microsoft.AnalysisServices.Tabular as TOM
+
+            # Extract region from the workspace
+            match = re.search(rf"{prefix_full}(.*?).{prefix}", self._workspace_name)
+            if match:
+                region = match.group(1)
+            if self._token_provider is None:
+                raise ValueError(
+                    f"{icons.red_dot} A token provider must be provided when connecting to Azure Analysis Services."
+                )
+            token = self._token_provider(audience="asazure", region=region)
+            connection_str = f'Provider=MSOLAP;Data Source={self._workspace_name};Password="{token}";Persist Security Info=True;Impersonation Level=Impersonate'
+            self._tom_server = TOM.Server()
+            self._tom_server.Connect(connection_str)
+        # Service Principal Authentication for Power BI via token provider
+        else:
+            from sempy.fabric._client._utils import _build_adomd_connection_string
+            import Microsoft.AnalysisServices.Tabular as TOM
+            from Microsoft.AnalysisServices import AccessToken
+            from sempy.fabric._token_provider import (
+                create_on_access_token_expired_callback,
+                ConstantTokenProvider,
+            )
+            from System import Func
+
+            token = self._token_provider(audience="pbi")
+            self._tom_server = TOM.Server()
+            get_access_token = create_on_access_token_expired_callback(
+                ConstantTokenProvider(token)
+            )
+            self._tom_server.AccessToken = get_access_token(None)
+            self._tom_server.OnAccessTokenExpired = Func[AccessToken, AccessToken](
+                get_access_token
+            )
+            workspace_url = f"powerbi://api.powerbi.com/v1.0/myorg/{workspace}"
+            connection_str = _build_adomd_connection_string(
+                workspace_url, readonly=readonly
+            )
+            self._tom_server.Connect(connection_str)
+
+        if self._is_azure_as:
+            self.model = self._tom_server.Databases.GetByName(self._dataset_name).Model
+        else:
+            self.model = self._tom_server.Databases[dataset_id].Model
 
         self._table_map = {}
         self._column_map = {}
@@ -4404,8 +4478,6 @@ class TOMWrapper:
         if isinstance(measure_name, str):
             measure_name = [measure_name]
 
-        client = fabric.FabricRestClient()
-
         if len(measure_name) > max_batch_size:
             measure_lists = [
                 measure_name[i : i + max_batch_size]
@@ -4446,6 +4518,7 @@ class TOMWrapper:
                     "modelItems"
                 ].append(new_item)
 
+            client = fabric.FabricRestClient()
             response = client.post("/explore/v202304/nl2nl/completions", json=payload)
             if response.status_code != 200:
                 raise FabricHTTPException(response)
@@ -4667,11 +4740,15 @@ class TOMWrapper:
 @log
 @contextmanager
 def connect_semantic_model(
-    dataset: str | UUID, readonly: bool = True, workspace: Optional[str] = None
+    dataset: str | UUID,
+    readonly: bool = True,
+    workspace: Optional[str | UUID] = None,
 ) -> Iterator[TOMWrapper]:
     """
     Connects to the Tabular Object Model (TOM) within a semantic model.
 
+    Service Principal Authentication is supported (see `here <https://github.com/microsoft/semantic-link-labs/blob/main/notebooks/Service%20Principal.ipynb>`_ for examples).
+
     Parameters
     ----------
     dataset : str | uuid.UUID
@@ -4679,10 +4756,10 @@ def connect_semantic_model(
     readonly: bool, default=True
         Whether the connection is read-only or read/write. Setting this to False enables read/write which saves the changes made back to the server.
     workspace : str | uuid.UUID, default=None
-        The Fabric workspace name or ID.
+        The Fabric workspace name or ID. Also supports Azure Analysis Services (Service Principal Authentication required).
+        If connecting to Azure Analysis Services, enter the workspace parameter in the following format: 'asazure://<region>.asazure.windows.net/<server_name>'.
         Defaults to None which resolves to the workspace of the attached lakehouse
         or if no lakehouse attached, resolves to the workspace of the notebook.
-
     Returns
     -------
     typing.Iterator[TOMWrapper]
@@ -4692,7 +4769,11 @@ def connect_semantic_model(
     # initialize .NET to make sure System and Microsoft.AnalysisServices.Tabular is defined
     sempy.fabric._client._utils._init_analysis_services()
 
-    tw = TOMWrapper(dataset=dataset, workspace=workspace, readonly=readonly)
+    tw = TOMWrapper(
+        dataset=dataset,
+        workspace=workspace,
+        readonly=readonly,
+    )
     try:
         yield tw
     finally: