From 8f557fa3b304c7d0d72d2fab51097797a107ce80 Mon Sep 17 00:00:00 2001
From: Shixuan Fan <shixuan.fan@snowflake.com>
Date: Thu, 20 Jul 2023 12:39:02 -0700
Subject: [PATCH] SNOW-853049 Add write_pandas test to a different schema
 (#958)

Description

This enables testing writing pandas DF to a different schema for both
client and stored proc.

Testing

integ
---
 tests/integ/test_pandas_to_df.py | 36 ++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/tests/integ/test_pandas_to_df.py b/tests/integ/test_pandas_to_df.py
index ee557aae59..595e2db211 100644
--- a/tests/integ/test_pandas_to_df.py
+++ b/tests/integ/test_pandas_to_df.py
@@ -11,8 +11,10 @@
 from pandas.testing import assert_frame_equal
 
 from snowflake.connector.errors import ProgrammingError
+from snowflake.snowpark import Row
 from snowflake.snowpark._internal.utils import (
     TempObjectType,
+    is_in_stored_procedure,
     random_name_for_temp_object,
     warning_dict,
 )
@@ -403,3 +405,37 @@ def test_special_name_quoting(
             ) in df_data
     finally:
         session.sql(drop_sql).collect()
+
+
+def test_write_to_different_schema(session):
+    pd_df = PandasDF(
+        [
+            (1, 4.5, "Nike"),
+            (2, 7.5, "Adidas"),
+            (3, 10.5, "Puma"),
+        ],
+        columns=["id".upper(), "foot_size".upper(), "shoe_make".upper()],
+    )
+    original_schema_name = session.get_current_schema()
+    test_schema_name = Utils.random_temp_schema()
+
+    try:
+        Utils.create_schema(session, test_schema_name)
+        # For owner's rights stored proc test, current schema does not change after creating a new schema
+        if not is_in_stored_procedure():
+            session.sql(f"use schema {original_schema_name}").collect()
+        assert session.get_current_schema() == original_schema_name
+        table_name = random_name_for_temp_object(TempObjectType.TABLE)
+        session.write_pandas(
+            pd_df,
+            table_name,
+            quote_identifiers=False,
+            schema=test_schema_name,
+            auto_create_table=True,
+        )
+        Utils.check_answer(
+            session.table(f"{test_schema_name}.{table_name}").sort("id"),
+            [Row(1, 4.5, "Nike"), Row(2, 7.5, "Adidas"), Row(3, 10.5, "Puma")],
+        )
+    finally:
+        Utils.drop_schema(session, test_schema_name)