From 750434db00aa6ca99000a0928b18e987cebd2996 Mon Sep 17 00:00:00 2001
From: Bhumika Goel <bhumika.goel@snowflake.com>
Date: Thu, 11 Jul 2024 14:40:37 -0500
Subject: [PATCH] make contextmanager for use_*_warehouse

---
 .../cli/plugins/nativeapp/manager.py          | 47 +++++++++++--------
 1 file changed, 27 insertions(+), 20 deletions(-)

diff --git a/src/snowflake/cli/plugins/nativeapp/manager.py b/src/snowflake/cli/plugins/nativeapp/manager.py
index 9852f8632e..69cd65cb38 100644
--- a/src/snowflake/cli/plugins/nativeapp/manager.py
+++ b/src/snowflake/cli/plugins/nativeapp/manager.py
@@ -17,6 +17,7 @@
 import json
 import os
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from functools import cached_property
 from pathlib import Path
 from textwrap import dedent
@@ -218,33 +219,39 @@ def stage_schema(self) -> Optional[str]:
     def package_warehouse(self) -> Optional[str]:
         return self.na_project.package_warehouse
 
-    def use_package_warehouse(self) -> str:
+    @contextmanager
+    def use_package_warehouse(self):
         if self.package_warehouse:
-            return self.package_warehouse
-        raise ClickException(
-            dedent(
-                f"""\
-            Application package warehouse cannot be empty.
-            Please provide a value for it in your connection information or your project definition file.
-            """
+            with self.use_warehouse(self.package_warehouse):
+                yield
+        else:
+            raise ClickException(
+                dedent(
+                    f"""\
+                Application package warehouse cannot be empty.
+                Please provide a value for it in your connection information or your project definition file.
+                """
+                )
             )
-        )
 
     @property
     def application_warehouse(self) -> Optional[str]:
         return self.na_project.application_warehouse
 
-    def use_application_warehouse(self) -> str:
+    @contextmanager
+    def use_application_warehouse(self):
         if self.application_warehouse:
-            return self.application_warehouse
-        raise ClickException(
-            dedent(
-                f"""\
-            Application warehouse cannot be empty.
-            Please provide a value for it in your connection information or your project definition file.
-            """
+            with self.use_warehouse(self.application_warehouse):
+                yield
+        else:
+            raise ClickException(
+                dedent(
+                    f"""\
+                Application warehouse cannot be empty.
+                Please provide a value for it in your connection information or your project definition file.
+                """
+                )
             )
-        )
 
     @property
     def project_identifier(self) -> str:
@@ -508,7 +515,7 @@ def _application_object_to_str(self, obj: ApplicationOwnedObject) -> str:
     def get_snowsight_url(self) -> str:
         """Returns the URL that can be used to visit this app via Snowsight."""
         name = identifier_for_url(self.app_name)
-        with self.use_warehouse(self.use_application_warehouse()):
+        with self.use_application_warehouse():
             return make_snowsight_url(self._conn, f"/#/apps/application/{name}")
 
     def create_app_package(self) -> None:
@@ -582,7 +589,7 @@ def _apply_package_scripts(self) -> None:
                 raise InvalidPackageScriptError(relpath, e)
 
         # once we're sure all the templates expanded correctly, execute all of them
-        with self.use_warehouse(self.use_package_warehouse()):
+        with self.use_package_warehouse():
             try:
                 for i, queries in enumerate(queued_queries):
                     cc.step(f"Applying package script: {self.package_scripts[i]}")