use azure clients as context manager

microsoft · Jun 20, 2024 · b6aa998 · b6aa998
1 parent 20b7fa4
commit b6aa998
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/pcfuncs/ipban/__init__.py b/pcfuncs/ipban/__init__.py
@@ -2,7 +2,7 @@
 import logging
 
 import azure.functions as func
-from azure.data.tables import TableClient, TableServiceClient
+from azure.data.tables import TableServiceClient
 from azure.identity import DefaultAzureCredential
 from azure.monitor.query import LogsQueryClient
 
@@ -18,12 +18,10 @@ def main(mytimer: func.TimerRequest) -> None:
     )
     logger.info("Updating the ip ban list at %s", utc_timestamp)
     credential: DefaultAzureCredential = DefaultAzureCredential()
-    logs_query_client: LogsQueryClient = LogsQueryClient(credential)
-    table_service_client: TableServiceClient = TableServiceClient(
+    with LogsQueryClient(credential) as logs_query_client, TableServiceClient(
         endpoint=settings.storage_account_url, credential=credential
-    )
-    table_client: TableClient = table_service_client.create_table_if_not_exists(
+    ) as table_service_client, table_service_client.create_table_if_not_exists(
         settings.banned_ip_table
-    )
-    task: UpdateBannedIPTask = UpdateBannedIPTask(logs_query_client, table_client)
-    task.run()
+    ) as table_client:
+        task = UpdateBannedIPTask(logs_query_client, table_client)
+        task.run()