From 8c18f912a20b35b9e24c772bf18b4eec50cb983a Mon Sep 17 00:00:00 2001
From: Nethra C Sashikar <sashikarnethra@gmail.com>
Date: Thu, 19 Dec 2024 09:42:25 -0800
Subject: [PATCH] add experimental decorator for firewallsku (#38657)

* add experimental decorator

* mark firewallsku as experimental field

* remove duplicate import

* isort

* format file

---------

Co-authored-by: Nethra Sashikar <nsashikar@microsoft.com>
---
 .../ai/ml/_schema/workspace/networking.py     | 27 ++++++++++++-------
 .../azure/ai/ml/constants/_workspace.py       |  3 ++-
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py
index 9c805b783e80..f228ee3e8ee3 100644
--- a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py
+++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py
@@ -7,6 +7,7 @@
 from marshmallow import EXCLUDE, fields
 from marshmallow.decorators import post_load, pre_dump
 
+from azure.ai.ml._schema import ExperimentalField
 from azure.ai.ml._schema.core.fields import NestedField, StringTransformedEnum, UnionField
 from azure.ai.ml._schema.core.schema_meta import PatchedSchemaMeta
 from azure.ai.ml._utils.utils import _snake_to_camel, camel_to_snake
@@ -47,7 +48,12 @@ def createdestobject(self, data, **kwargs):
         category = data.get("category", OutboundRuleCategory.USER_DEFINED)
         name = data.get("name")
         status = data.get("status", None)
-        return FqdnDestination(name=name, destination=dest, category=_snake_to_camel(category), status=status)
+        return FqdnDestination(
+            name=name,
+            destination=dest,
+            category=_snake_to_camel(category),
+            status=status,
+        )
 
 
 class ServiceTagDestinationSchema(metaclass=PatchedSchemaMeta):
@@ -185,13 +191,15 @@ class ManagedNetworkSchema(metaclass=PatchedSchemaMeta):
         ),
         allow_none=True,
     )
-    firewall_sku = StringTransformedEnum(
-        allowed_values=[
-            FirewallSku.STANDARD,
-            FirewallSku.BASIC,
-        ],
-        casing_transform=camel_to_snake,
-        metadata={"description": "Firewall sku for FQDN rules in AllowOnlyApprovedOutbound mode"},
+    firewall_sku = ExperimentalField(
+        StringTransformedEnum(
+            allowed_values=[
+                FirewallSku.STANDARD,
+                FirewallSku.BASIC,
+            ],
+            casing_transform=camel_to_snake,
+            metadata={"description": "Firewall sku for FQDN rules in AllowOnlyApprovedOutbound mode"},
+        )
     )
     network_id = fields.Str(required=False, dump_only=True)
     status = NestedField(ManagedNetworkStatusSchema, allow_none=False, unknown=EXCLUDE)
@@ -211,5 +219,6 @@ def make(self, data, **kwargs):
             )
         else:
             return ManagedNetwork(
-                isolation_mode=_snake_to_camel(data["isolation_mode"]), firewall_sku=firewall_sku_value
+                isolation_mode=_snake_to_camel(data["isolation_mode"]),
+                firewall_sku=firewall_sku_value,
             )
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_workspace.py b/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_workspace.py
index bf0479c01823..359e211dd7c0 100644
--- a/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_workspace.py
+++ b/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_workspace.py
@@ -3,8 +3,8 @@
 # ---------------------------------------------------------
 from enum import Enum
 
-from azure.core import CaseInsensitiveEnumMeta
 from azure.ai.ml._utils._experimental import experimental
+from azure.core import CaseInsensitiveEnumMeta
 
 
 class ManagedServiceIdentityType(str, Enum, metaclass=CaseInsensitiveEnumMeta):
@@ -24,6 +24,7 @@ class IsolationMode:
     ALLOW_ONLY_APPROVED_OUTBOUND = "AllowOnlyApprovedOutbound"
 
 
+@experimental
 class FirewallSku:
     """Firewall Sku for FQDN rules in AllowOnlyApprovedOutbound."""