Test pr 418 (#419)

* feat: Add temporary option to the create stage command. * update changelog * Update tests/test_create.py Co-authored-by: Adam Ling <[email protected]> --------- Co-authored-by: DanCardin <[email protected]> Co-authored-by: Adam Ling <[email protected]>
snowflakedb · Jun 6, 2023 · c200f97 · c200f97
1 parent c563913
commit c200f97
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 3 deletions.
diff --git a/DESCRIPTION.md b/DESCRIPTION.md
@@ -9,6 +9,10 @@ Source code is also available at:
 
 # Release Notes
 
+- v1.4.8(TBD)
+
+  - Added opiton to create a temporary stage command.
+
 - v1.4.7(Mar 22, 2023)
 
   - Re-applied the application name of driver connection `SnowflakeConnection` to `SnowflakeSQLAlchemy`.

diff --git a/src/snowflake/sqlalchemy/base.py b/src/snowflake/sqlalchemy/base.py
@@ -534,11 +534,12 @@ def visit_create_stage(self, create_stage, **kw):
         """
         This visitor will create the SQL representation for a CREATE STAGE command.
         """
-        return "CREATE {}STAGE {}{} URL={}".format(
-            "OR REPLACE " if create_stage.replace_if_exists else "",
+        return "CREATE {or_replace}{temporary}STAGE {}{} URL={}".format(
             create_stage.stage.namespace,
             create_stage.stage.name,
             repr(create_stage.container),
+            or_replace="OR REPLACE " if create_stage.replace_if_exists else "",
+            temporary="TEMPORARY " if create_stage.temporary else "",
         )
 
     def visit_create_file_format(self, file_format, **kw):

diff --git a/src/snowflake/sqlalchemy/custom_commands.py b/src/snowflake/sqlalchemy/custom_commands.py
@@ -482,9 +482,10 @@ class CreateStage(DDLElement):
 
     __visit_name__ = "create_stage"
 
-    def __init__(self, container, stage, replace_if_exists=False):
+    def __init__(self, container, stage, replace_if_exists=False, *, temporary=False):
         super().__init__()
         self.container = container
+        self.temporary = temporary
         self.stage = stage
         self.replace_if_exists = replace_if_exists
 

diff --git a/tests/test_create.py b/tests/test_create.py
@@ -54,6 +54,16 @@ def test_create_stage(sql_compiler):
     )
     assert actual == expected
 
+    create_stage = CreateStage(stage=stage, container=container, temporary=True)
+    # validate that the resulting SQL is as expected
+    actual = sql_compiler(create_stage)
+    expected = (
+        "CREATE TEMPORARY STAGE MY_DB.MY_SCHEMA.AZURE_STAGE "
+        "URL='azure://myaccount.blob.core.windows.net/my-container' "
+        "CREDENTIALS=(AZURE_SAS_TOKEN='saas_token')"
+    )
+    assert actual == expected
+
 
 def test_create_csv_format(sql_compiler):
     """