Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add copy_grants to write_pandas #2050

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat: add copy_grants to write_pandas #2050

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
013c53d
feat: add copy_grants to write_pandas
surpatean Sep 10, 2024
095f7c7
chore: change EOL of test_pandas_tools.py back to LF
surpatean Sep 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 23 additions & 6 deletions src/snowflake/connector/pandas_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,7 @@ def write_pandas(
overwrite: bool = False,
table_type: Literal["", "temp", "temporary", "transient"] = "",
use_logical_type: bool | None = None,
copy_grants: bool | None = None,
**kwargs: Any,
) -> tuple[
bool,
Expand Down Expand Up @@ -237,14 +238,17 @@ def write_pandas(
the passed in DataFrame. The table will not be created if it already exists
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.
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.
use_logical_type: Boolean that specifies whether to use Parquet logical types. With this file format option,
Snowflake can interpret Parquet logical types during data loading. To enable Parquet logical types,
set use_logical_type as True. Set to None to use Snowflakes default. For more information, see:
https://docs.snowflake.com/en/sql-reference/sql/create-file-format
copy_grants: When true and when both overwrite is true and auto_create_table is true, the grants of the table
being overwritten will be copied to the new table. Otherwise, the new table will be created with only the
default/future grants defined on the schema/database.
Comment on lines +249 to +251
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please add the fact that create table privileges are required with this option.



Returns:
Expand Down Expand Up @@ -315,6 +319,14 @@ def write_pandas(
else:
sql_use_logical_type = " USE_LOGICAL_TYPE = FALSE"

if copy_grants and (not overwrite or not auto_create_table):
warnings.warn(
"copy_grants is only applied when used in combination with "
"overwrite and auto_create_table",
UserWarning,
stacklevel=2,
)

cursor = conn.cursor()
stage_location = _create_temp_stage(
cursor,
Expand Down Expand Up @@ -444,10 +456,15 @@ def drop_object(name: str, object_type: str) -> None:
name=table_name,
quote_identifiers=quote_identifiers,
)
drop_object(original_table_location, "table")
rename_table_sql = f"ALTER TABLE {target_table_location} RENAME TO {original_table_location} /* Python:snowflake.connector.pandas_tools.write_pandas() */"
logger.debug(f"rename table with '{rename_table_sql}'")
cursor.execute(rename_table_sql, _is_internal=True)
clone_table_sql = (
f"CREATE OR REPLACE TABLE {original_table_location} "
f"CLONE {target_table_location} "
f"{'COPY GRANTS' if copy_grants else ''}"
f" /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
)
logger.debug(f"clone table with '{clone_table_sql}'")
cursor.execute(clone_table_sql, _is_internal=True)
drop_object(target_table_location, "table")
Comment on lines -447 to +467
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this could be a breaking change since not all users my have create table privileges. Can you refactor the code like so

if clone_table_sql:
    clone_table_sql = "create or replace ..."
    logger.debug()
    cursor.execute()
    drop_object()
else:
    rename_table_sql = "alter table ..."
    logger.debug()
    cursor.execute()

except ProgrammingError:
if overwrite and auto_create_table:
# drop table only if we created a new one with a random name
Expand Down
71 changes: 71 additions & 0 deletions test/integ/pandas/test_pandas_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -407,6 +407,77 @@ def test_write_pandas_table_type(
cnx.execute_string(drop_sql)


@pytest.mark.parametrize("copy_grants", [True, False, None])
def test_write_pandas_copy_grants(
conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
copy_grants: bool | None,
):
with conn_cnx() as cnx:
table_name = random_string(5, "test_write_pandas_copy_grants_")
drop_sql = f"DROP TABLE IF EXISTS {table_name}"
try:
success, _, _, _ = write_pandas(
cnx,
sf_connector_version_df.get(),
table_name,
auto_create_table=True,
overwrite=True,
quote_identifiers=False,
)
assert success

# Get the default number of grants on new tables in the schema
count_default_grants = (
cnx.cursor(DictCursor)
.execute(f"show grants on table {table_name}")
.rowcount
)

# Add dummy grants on the table
(
cnx.cursor(DictCursor)
.execute(f"grant references on table {table_name} to role {cnx.role}")
)
(
cnx.cursor(DictCursor)
.execute(f"grant references on table {table_name} to role SYSADMIN")
)
count_current_grants = (
cnx.cursor(DictCursor)
.execute(f"show grants on table {table_name}")
.rowcount
)

# Sanity check, grants count should have increased
assert count_current_grants > count_default_grants

# Re-create the table
success, _, _, _ = write_pandas(
cnx,
sf_connector_version_df.get(),
table_name,
auto_create_table=True,
overwrite=True,
quote_identifiers=False,
copy_grants=copy_grants,
)
assert success

count_current_grants = (
cnx.cursor(DictCursor)
.execute(f"show grants on table {table_name}")
.rowcount
)
if copy_grants:
# Grants count should be more than default, because the added ones were copied
assert count_current_grants > count_default_grants
else:
# Grants should match the default for tables in the schema
assert count_current_grants == count_default_grants
finally:
cnx.execute_string(drop_sql)


def test_write_pandas_create_temp_table_deprecation_warning(
conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
):
Expand Down