add policy and decorators for keyvault.secrets

sdk/keyvault/azure-keyvault-secrets/azure/keyvault/secrets/_client.py

@@ -13,6 +13,7 @@
     from typing import Any, Dict, Generator, Mapping, Optional
 
 from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
+from azure.core.tracing.decorator import distributed_trace
 
 from ._shared import KeyVaultClientBase
 from ._models import Secret, DeletedSecret, SecretAttributes
@@ -32,6 +33,7 @@ class SecretClient(KeyVaultClientBase):
 
     # pylint:disable=protected-access
 
+    @distributed_trace
     def get_secret(self, name, version=None, **kwargs):
         # type: (str, str, Mapping[str, Any]) -> Secret
         """Get a specified secret from the vault.
@@ -59,6 +61,7 @@ def get_secret(self, name, version=None, **kwargs):
         )
         return Secret._from_secret_bundle(bundle)
 
+    @distributed_trace
     def set_secret(
         self, name, value, content_type=None, enabled=None, not_before=None, expires=None, tags=None, **kwargs
     ):
@@ -101,6 +104,7 @@ def set_secret(
         )
         return Secret._from_secret_bundle(bundle)
 
+    @distributed_trace
     def update_secret(
         self, name, version=None, content_type=None, enabled=None, not_before=None, expires=None, tags=None, **kwargs
     ):
@@ -151,6 +155,7 @@ def update_secret(
         )
         return SecretAttributes._from_secret_bundle(bundle)  # pylint: disable=protected-access
 
+    @distributed_trace
     def list_secrets(self, **kwargs):
         # type: (Mapping[str, Any]) -> Generator[SecretAttributes]
         """List secrets in the vault.
@@ -177,6 +182,7 @@ def list_secrets(self, **kwargs):
         pages = self._client.get_secrets(self._vault_url, maxresults=max_page_size, **kwargs)
         return (SecretAttributes._from_secret_item(item) for item in pages)
 
+    @distributed_trace
     def list_secret_versions(self, name, **kwargs):
         # type: (str, Mapping[str, Any]) -> Generator[SecretAttributes]
         """List all versions of the specified secret.
@@ -203,6 +209,7 @@ def list_secret_versions(self, name, **kwargs):
         pages = self._client.get_secret_versions(self._vault_url, name, maxresults=max_page_size, **kwargs)
         return (SecretAttributes._from_secret_item(item) for item in pages)
 
+    @distributed_trace
     def backup_secret(self, name, **kwargs):
         # type: (str, Mapping[str, Any]) -> bytes
         """Backs up the specified secret.
@@ -230,6 +237,7 @@ def backup_secret(self, name, **kwargs):
         )
         return backup_result.value
 
+    @distributed_trace
     def restore_secret(self, backup, **kwargs):
         # type: (bytes, Mapping[str, Any]) -> SecretAttributes
         """Restore a backed up secret to the vault.
@@ -254,6 +262,7 @@ def restore_secret(self, backup, **kwargs):
         bundle = self._client.restore_secret(self.vault_url, backup, error_map={409: ResourceExistsError}, **kwargs)
         return SecretAttributes._from_secret_bundle(bundle)
 
+    @distributed_trace
     def delete_secret(self, name, **kwargs):
         # type: (str, Mapping[str, Any]) -> DeletedSecret
         """Deletes a secret from the vault.
@@ -279,6 +288,7 @@ def delete_secret(self, name, **kwargs):
         bundle = self._client.delete_secret(self.vault_url, name, error_map={404: ResourceNotFoundError}, **kwargs)
         return DeletedSecret._from_deleted_secret_bundle(bundle)
 
+    @distributed_trace
     def get_deleted_secret(self, name, **kwargs):
         # type: (str, Mapping[str, Any]) -> DeletedSecret
         """Gets the specified deleted secret.
@@ -303,6 +313,7 @@ def get_deleted_secret(self, name, **kwargs):
         bundle = self._client.get_deleted_secret(self.vault_url, name, error_map={404: ResourceNotFoundError}, **kwargs)
         return DeletedSecret._from_deleted_secret_bundle(bundle)
 
+    @distributed_trace
     def list_deleted_secrets(self, **kwargs):
         # type: (Mapping[str, Any]) -> Generator[DeletedSecret]
         """Lists deleted secrets of the vault.
@@ -328,6 +339,7 @@ def list_deleted_secrets(self, **kwargs):
         pages = self._client.get_deleted_secrets(self._vault_url, maxresults=max_page_size, **kwargs)
         return (DeletedSecret._from_deleted_secret_item(item) for item in pages)
 
+    @distributed_trace
     def purge_deleted_secret(self, name, **kwargs):
         # type: (str, Mapping[str, Any]) -> None
         """Permanently deletes the specified secret.
@@ -350,6 +362,7 @@ def purge_deleted_secret(self, name, **kwargs):
         """
         self._client.purge_deleted_secret(self.vault_url, name, **kwargs)
 
+    @distributed_trace
     def recover_deleted_secret(self, name, **kwargs):
         # type: (str, Mapping[str, Any]) -> SecretAttributes
         """Recovers the deleted secret to the latest version.

sdk/keyvault/azure-keyvault-secrets/azure/keyvault/secrets/_shared/async_client_base.py

@@ -6,6 +6,7 @@
 from azure.core.async_paging import AsyncPagedMixin
 from azure.core.configuration import Configuration
 from azure.core.pipeline import AsyncPipeline
+from azure.core.pipeline.policies.distributed_tracing import DistributedTracingPolicy
 from azure.core.pipeline.transport import AsyncioRequestsTransport, HttpTransport
 from msrest.serialization import Model
 
@@ -101,6 +102,7 @@ def _build_pipeline(config: Configuration, transport: HttpTransport, **kwargs: A
             config.retry_policy,
             config.authentication_policy,
             config.logging_policy,
+            DistributedTracingPolicy(),
         ]
 
         if transport is None:

sdk/keyvault/azure-keyvault-secrets/azure/keyvault/secrets/_shared/client_base.py

@@ -15,6 +15,7 @@
     from azure.core.pipeline.transport import HttpTransport
 
 from .challenge_auth_policy import ChallengeAuthPolicy
+from azure.core.pipeline.policies.distributed_tracing import DistributedTracingPolicy
 
 
 KEY_VAULT_SCOPE = "https://vault.azure.net/.default"
@@ -73,6 +74,7 @@ def _build_pipeline(self, config, transport, **kwargs):
             config.retry_policy,
             config.authentication_policy,
             config.logging_policy,
+            DistributedTracingPolicy(),
         ]
 
         if transport is None:

sdk/keyvault/azure-keyvault-secrets/azure/keyvault/secrets/aio/_client.py

@@ -6,6 +6,8 @@
 from typing import Any, AsyncIterable, Mapping, Optional, Dict
 
 from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
+from azure.core.tracing.decorator import distributed_trace
+from azure.core.tracing.decorator_async import distributed_trace_async
 
 from azure.keyvault.secrets._models import Secret, DeletedSecret, SecretAttributes
 from .._shared import AsyncKeyVaultClientBase, AsyncPagingAdapter
@@ -25,6 +27,7 @@ class SecretClient(AsyncKeyVaultClientBase):
 
     # pylint:disable=protected-access
 
+    @distributed_trace_async
     async def get_secret(self, name: str, version: Optional[str] = None, **kwargs: Mapping[str, Any]) -> Secret:
         """Get a specified secret from the vault.
 
@@ -51,6 +54,7 @@ async def get_secret(self, name: str, version: Optional[str] = None, **kwargs: M
         )
         return Secret._from_secret_bundle(bundle)
 
+    @distributed_trace_async
     async def set_secret(
         self,
         name: str,
@@ -99,6 +103,7 @@ async def set_secret(
         )
         return Secret._from_secret_bundle(bundle)
 
+    @distributed_trace_async
     async def update_secret(
         self,
         name: str,
@@ -155,6 +160,7 @@ async def update_secret(
         )
         return SecretAttributes._from_secret_bundle(bundle)  # pylint: disable=protected-access
 
+    @distributed_trace
     def list_secrets(self, **kwargs: Mapping[str, Any]) -> AsyncIterable[SecretAttributes]:
         """List secrets in the vault.
 
@@ -180,6 +186,7 @@ def list_secrets(self, **kwargs: Mapping[str, Any]) -> AsyncIterable[SecretAttri
         iterable = AsyncPagingAdapter(pages, SecretAttributes._from_secret_item)
         return iterable
 
+    @distributed_trace
     def list_secret_versions(self, name: str, **kwargs: Mapping[str, Any]) -> AsyncIterable[SecretAttributes]:
         """List all versions of the specified secret.
 
@@ -205,6 +212,7 @@ def list_secret_versions(self, name: str, **kwargs: Mapping[str, Any]) -> AsyncI
         iterable = AsyncPagingAdapter(pages, SecretAttributes._from_secret_item)
         return iterable
 
+    @distributed_trace_async
     async def backup_secret(self, name: str, **kwargs: Mapping[str, Any]) -> bytes:
         """Backs up the specified secret.
 
@@ -230,6 +238,7 @@ async def backup_secret(self, name: str, **kwargs: Mapping[str, Any]) -> bytes:
         )
         return backup_result.value
 
+    @distributed_trace_async
     async def restore_secret(self, backup: bytes, **kwargs: Mapping[str, Any]) -> SecretAttributes:
         """Restores a backed up secret to a vault.
 
@@ -254,6 +263,7 @@ async def restore_secret(self, backup: bytes, **kwargs: Mapping[str, Any]) -> Se
         )
         return SecretAttributes._from_secret_bundle(bundle)
 
+    @distributed_trace_async
     async def delete_secret(self, name: str, **kwargs: Mapping[str, Any]) -> DeletedSecret:
         """Deletes a secret from the vault.
 
@@ -279,6 +289,7 @@ async def delete_secret(self, name: str, **kwargs: Mapping[str, Any]) -> Deleted
         )
         return DeletedSecret._from_deleted_secret_bundle(bundle)
 
+    @distributed_trace_async
     async def get_deleted_secret(self, name: str, **kwargs: Mapping[str, Any]) -> DeletedSecret:
         """Gets the specified deleted secret.
 
@@ -303,6 +314,7 @@ async def get_deleted_secret(self, name: str, **kwargs: Mapping[str, Any]) -> De
         )
         return DeletedSecret._from_deleted_secret_bundle(bundle)
 
+    @distributed_trace
     def list_deleted_secrets(self, **kwargs: Mapping[str, Any]) -> AsyncIterable[DeletedSecret]:
         """Lists deleted secrets of the vault.
 
@@ -327,6 +339,7 @@ def list_deleted_secrets(self, **kwargs: Mapping[str, Any]) -> AsyncIterable[Del
         iterable = AsyncPagingAdapter(pages, DeletedSecret._from_deleted_secret_item)
         return iterable
 
+    @distributed_trace_async
     async def purge_deleted_secret(self, name: str, **kwargs: Mapping[str, Any]) -> None:
         """Permanently deletes the specified secret.
 
@@ -347,6 +360,7 @@ async def purge_deleted_secret(self, name: str, **kwargs: Mapping[str, Any]) ->
         """
         await self._client.purge_deleted_secret(self.vault_url, name, **kwargs)
 
+    @distributed_trace_async
     async def recover_deleted_secret(self, name: str, **kwargs: Mapping[str, Any]) -> SecretAttributes:
         """Recovers the deleted secret to the latest version.