warehouse audit APIs

Author: m-kovalsky

src/sempy_labs/__init__.py

@@ -101,6 +101,9 @@
     get_warehouse_columns,
     get_warehouse_tables,
     get_warehouse_connection_string,
+    get_warehouse_sql_audit_settings,
+    update_warehouse_sql_audit_settings,
+    set_warehouse_audit_actions_and_group,
 )
 from ._data_pipelines import (
     list_data_pipelines,
@@ -605,4 +608,7 @@
     "set_workspace_network_communication_policy",
     "get_warehouse_connection_string",
     "list_data_access_roles",
+    "get_warehouse_sql_audit_settings",
+    "update_warehouse_sql_audit_settings",
+    "set_warehouse_audit_actions_and_group",
 ]

src/sempy_labs/_warehouses.py

@@ -1,14 +1,15 @@
-from ._helper_functions import (
+from sempy_labs._helper_functions import (
     resolve_item_id,
     resolve_workspace_name_and_id,
     _base_api,
     _create_dataframe,
     _update_dataframe_datatypes,
     delete_item,
     resolve_workspace_id,
+    resolve_item_name_and_id,
 )
 import pandas as pd
-from typing import Optional
+from typing import Optional, List
 import sempy_labs._icons as icons
 from uuid import UUID
 from sempy._utils._log import log
@@ -287,3 +288,172 @@ def get_warehouse_connection_string(
     response = _base_api(request=url, client="fabric_sp")
 
     return response.json().get("connectionString")
+
+
+@log
+def get_warehouse_sql_audit_settings(
+    warehouse: str | UUID, workspace: Optional[str | UUID] = None
+) -> pd.DataFrame:
+    """
+    Shows the SQL audit settings of a Fabric warehouse.
+
+    This is a wrapper function for the following API: `SQL Audit Settings - Get SQL Audit Settings <https://learn.microsoft.com/rest/api/fabric/warehouse/sql-audit-settings/get-sql-audit-settings>`_.
+
+    Service Principal Authentication is supported (see `here <https://learn.microsoft.com/rest/api/fabric/warehouse/sql-audit-settings/get-sql-audit-settings#service-principal-authentication>`_).
+
+    Parameters
+    ----------
+    warehouse : str | uuid.UUID
+        Name or ID of the Fabric warehouse.
+    workspace : str | uuid.UUID, default=None
+        The Fabric workspace name or ID.
+        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
+    -------
+    pandas.DataFrame
+        A pandas dataframe containing the SQL audit settings of the specified warehouse.
+    """
+    workspace_id = resolve_workspace_id(workspace)
+    warehouse_id = resolve_item_id(
+        item=warehouse, type="Warehouse", workspace=workspace
+    )
+
+    columns = {
+        "State": "string",
+        "Retention Days": "int",
+        "Audit Actions And Group": "list",
+    }
+
+    df = _create_dataframe(columns=columns)
+
+    response = _base_api(
+        request=f"/v1/workspaces/{workspace_id}/warehouses/{warehouse_id}/settings/sqlAudit",
+        client="fabric_sp",
+    ).json()
+
+    rows = []
+    rows.append(
+        {
+            "State": response.get("state"),
+            "Retention Days": response.get("retentionDays"),
+            "Audit Actions And Group": response.get("auditActionsAndGroups"),
+        }
+    )
+
+    if rows:
+        df = pd.DataFrame(rows, columns=list(columns.keys()))
+        _update_dataframe_datatypes(dataframe=df, column_map=columns)
+
+    return df
+
+
+@log
+def update_warehouse_sql_audit_settings(
+    warehouse: str | UUID,
+    workspace: Optional[str | UUID] = None,
+    retention_days: Optional[int] = None,
+    state: Optional[str] = None,
+):
+    """
+    Update settings associated with the warehouse.
+
+    This is a wrapper function for the following API: SQL Audit Settings - Update SQL Audit Settings <https://learn.microsoft.com/rest/api/fabric/warehouse/sql-audit-settings/update-sql-audit-settings>`_.
+
+    Service Principal Authentication is supported (see `here <https://learn.microsoft.com/rest/api/fabric/warehouse/sql-audit-settings/get-sql-audit-settings#service-principal-authentication>`_).
+
+    Parameters
+    ----------
+    warehouse : str | uuid.UUID
+        Name or ID of the Fabric warehouse.
+    workspace : str | uuid.UUID, default=None
+        The Fabric workspace name or ID.
+        Defaults to None which resolves to the workspace of the attached lakehouse
+        or if no lakehouse attached, resolves to the workspace of the notebook.
+    """
+
+    (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
+    (warehouse_name, warehouse_id) = resolve_item_name_and_id(
+        item=warehouse, type="Warehouse", workspace=workspace
+    )
+
+    payload = {}
+    if retention_days is not None:
+        if not isinstance(retention_days, int) or retention_days < 0:
+            raise ValueError(
+                f"{icons.red_dot} retention_days must be a non-negative integer."
+            )
+        payload["retentionDays"] = retention_days
+    if state is not None:
+        state = state.capitalize()
+        if state not in ["Enabled", "Disabled"]:
+            raise ValueError(
+                f"{icons.red_dot} state must be either 'Enabled' or 'Disabled'."
+            )
+        payload["state"] = state
+
+    if not payload:
+        print(
+            f"{icons.info} No updates were made as neither retention_days nor state were provided."
+        )
+        return
+
+    _base_api(
+        request=f"/v1/workspaces/{workspace_id}/warehouses/{warehouse_id}/settings/sqlAudit",
+        client="fabric_sp",
+        method="patch",
+        payload=payload,
+    )
+
+    print(
+        f"{icons.green_dot} The SQL audit settings for the '{warehouse_name}' warehouse within the '{workspace_name}' workspace have been updated accordingly."
+    )
+
+
+@log
+def set_warehouse_audit_actions_and_group(
+    warehouse: str | UUID,
+    sql_audit_groups: List[str],
+    workspace: Optional[str | UUID] = None,
+):
+    """
+    Update the audit actions and groups for this warehouse.
+
+    This is a wrapper function for the following API: SQL Audit Settings - Set Audit Actions And Groups <https://learn.microsoft.com/rest/api/fabric/warehouse/sql-audit-settings/set-audit-actions-and-groups>`_.
+
+    Service Principal Authentication is supported (see `here <https://learn.microsoft.com/rest/api/fabric/warehouse/sql-audit-settings/get-sql-audit-settings#service-principal-authentication>`_).
+
+    Parameters
+    ----------
+    warehouse : str | uuid.UUID
+        Name or ID of the Fabric warehouse.
+    workspace : str | uuid.UUID, default=None
+        The Fabric workspace name or ID.
+        Defaults to None which resolves to the workspace of the attached lakehouse
+        or if no lakehouse attached, resolves to the workspace of the notebook.
+    """
+
+    (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
+    (warehouse_name, warehouse_id) = resolve_item_name_and_id(
+        item=warehouse, type="Warehouse", workspace=workspace
+    )
+    if (
+        not sql_audit_groups
+        or not isinstance(sql_audit_groups, list)
+        or not all(isinstance(item, str) for item in sql_audit_groups)
+    ):
+        raise ValueError(
+            f"{icons.red_dot} sql_audit_groups must be a non-empty list of strings."
+        )
+
+    _base_api(
+        request=f"/v1/workspaces/{workspace_id}/warehouses/{warehouse_id}/settings/sqlAudit/setAuditActionsAndGroups",
+        client="fabric_sp",
+        method="post",
+        payload=sql_audit_groups,
+    )
+
+    print(
+        f"{icons.green_dot} The SQL audit actions and groups for the '{warehouse_name}' warehouse within the '{workspace_name}' workspace have been updated accordingly."
+    )