SNOW-630797 Support table_type in pandas_tools (#1216)

* SNOW-630797 Support table_type in pandas_tools

Description

table_type could support transient table as well, which is more
flexible comparing to create_temp_table

Testing

Added integ test

* Add negative test

* type hint + revert undesired test param change

* import Literal from typing_extensions

* refactor tests

* assert success + various table name

* use random_string

Author: sfc-gh-sfan

src/snowflake/connector/pandas_tools.py

@@ -8,11 +8,14 @@
 import os
 import random
 import string
+import warnings
 from functools import partial
 from logging import getLogger
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Callable, Iterable, Iterator, Sequence, TypeVar
 
+from typing_extensions import Literal
+
 from snowflake.connector import ProgrammingError
 from snowflake.connector.options import pandas
 
@@ -49,6 +52,7 @@ def write_pandas(
     auto_create_table: bool = False,
     create_temp_table: bool = False,
     overwrite: bool = False,
+    table_type: Literal["", "temp", "temporary", "transient"] = "",
 ) -> tuple[
     bool,
     int,
@@ -102,10 +106,12 @@ def write_pandas(
             I.e. identifiers will be coerced to uppercase by Snowflake.  (Default value = True)
         auto_create_table: When true, will automatically create a table with corresponding columns for each column in
             the passed in DataFrame. The table will not be created if it already exists
-        create_temp_table: Will make the auto-created table as a temporary table
+        create_temp_table: (Deprecated) Will make the auto-created table as a temporary table
         overwrite: When true, and if auto_create_table is true, then it drops the table. Otherwise, it
         truncates the table. In both cases it will replace the existing contents of the table with that of the passed in
             Pandas DataFrame.
+        table_type: The table type of to-be-created table. The supported table types include ``temp``/``temporary``
+            and ``transient``. Empty means permanent table as per SQL convention.
 
     Returns:
         Returns the COPY INTO command's results to verify ingestion in the form of a tuple of whether all chunks were
@@ -124,6 +130,22 @@ def write_pandas(
                 compression, compression_map.keys()
             )
         )
+
+    if create_temp_table:
+        warnings.warn(
+            "create_temp_table is deprecated, we still respect this parameter when it is True but "
+            'please consider using `table_type="temp"` instead',
+            DeprecationWarning,
+            # warnings.warn -> write_pandas
+            stacklevel=2,
+        )
+        table_type = "temp"
+
+    if table_type and table_type.lower() not in ["temp", "temporary", "transient"]:
+        raise ValueError(
+            "Unsupported table type. Expected table types: temp/temporary, transient"
+        )
+
     if quote_identifiers:
         location = (
             (('"' + database + '".') if database else "")
@@ -224,7 +246,7 @@ def write_pandas(
             [f"{quote}{c}{quote} {column_type_mapping[c]}" for c in df.columns]
         )
         create_table_sql = (
-            f"CREATE {'TEMP ' if create_temp_table else ''}TABLE IF NOT EXISTS {location} "
+            f"CREATE {table_type.upper()} TABLE IF NOT EXISTS {location} "
             f"({create_table_columns})"
             f" /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
         )

test/integ/pandas/test_pandas_tools.py

@@ -219,6 +219,79 @@ def test_write_pandas(
             cnx.execute_string(drop_sql)
 
 
+@pytest.mark.parametrize("table_type", ["", "temp", "temporary", "transient"])
+def test_write_pandas_table_type(
+    conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
+    table_type: str,
+):
+    with conn_cnx() as cnx:
+        table_name = random_string(5, "write_pandas_table_type_")
+        drop_sql = f"DROP TABLE IF EXISTS {table_name}"
+        try:
+            success, _, _, _ = write_pandas(
+                cnx,
+                sf_connector_version_df.get(),
+                table_name,
+                table_type=table_type,
+                auto_create_table=True,
+            )
+            table_info = (
+                cnx.cursor(DictCursor)
+                .execute(f"show tables like '{table_name}'")
+                .fetchall()
+            )
+            assert success
+            if not table_type:
+                expected_table_kind = "TABLE"
+            elif table_type == "temp":
+                expected_table_kind = "TEMPORARY"
+            else:
+                expected_table_kind = table_type.upper()
+            assert table_info[0]["kind"] == expected_table_kind
+        finally:
+            cnx.execute_string(drop_sql)
+
+
+def test_write_pandas_create_temp_table_deprecation_warning(
+    conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
+):
+    with conn_cnx() as cnx:
+        table_name = random_string(5, "driver_versions_")
+        drop_sql = f"DROP TABLE IF EXISTS {table_name}"
+        try:
+            with pytest.deprecated_call(match="create_temp_table is deprecated"):
+                success, _, _, _ = write_pandas(
+                    cnx,
+                    sf_connector_version_df.get(),
+                    table_name,
+                    create_temp_table=True,
+                    auto_create_table=True,
+                )
+
+            assert success
+            table_info = (
+                cnx.cursor(DictCursor)
+                .execute(f"show tables like '{table_name}'")
+                .fetchall()
+            )
+            assert table_info[0]["kind"] == "TEMPORARY"
+        finally:
+            cnx.execute_string(drop_sql)
+
+
+def test_invalid_table_type_write_pandas(
+    conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
+):
+    with conn_cnx() as cnx:
+        with pytest.raises(ValueError, match="Unsupported table type"):
+            write_pandas(
+                cnx,
+                sf_connector_version_df.get(),
+                "invalid_table_type",
+                table_type="invalid",
+            )
+
+
 @pytest.mark.parametrize("quote_identifiers", [True, False])
 def test_location_building_db_schema(conn_cnx, quote_identifiers: bool):
     """This tests that write_pandas constructs location correctly with database, schema and table name."""