Skip to content

Commit

Permalink
fix(snowflake-snowpark): disable the reconnect method (#8969)
Browse files Browse the repository at this point in the history
It's not yet feasible to call `reconnect` with an Ibis Snowflake backend initialized from a Snowpark connection, so this PR disables that code path.
  • Loading branch information
cpcloud authored Apr 15, 2024
1 parent 4b476ba commit e31eded
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 0 deletions.
9 changes: 9 additions & 0 deletions ibis/backends/snowflake/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,7 @@ def do_connect(self, create_object_udfs: bool = True, **kwargs: Any):
session_parameters=session_parameters,
create_object_udfs=create_object_udfs,
)
self._from_snowpark = False

def _setup_session(self, *, con, session_parameters, create_object_udfs: bool):
self.con = con
Expand Down Expand Up @@ -320,8 +321,16 @@ def from_snowpark(cls, session, *, create_object_udfs: bool = True) -> Backend:
session_parameters={},
create_object_udfs=create_object_udfs,
)
backend._from_snowpark = True
return backend

def reconnect(self) -> None:
if self._from_snowpark:
raise com.IbisError(
"Reconnecting is not supported when using a Snowpark session"
)
super().reconnect()

def _get_udf_source(self, udf_node: ops.ScalarUDF):
name = type(udf_node).__name__
signature = ", ".join(
Expand Down
6 changes: 6 additions & 0 deletions ibis/backends/tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import importlib
import inspect
import json
import os
import re
import string
import subprocess
Expand Down Expand Up @@ -313,6 +314,11 @@ def test_create_table_from_schema(con, new_schema, temp_table):
reason="`tbl_properties` is required when creating table with schema",
)
def test_create_temporary_table_from_schema(con_no_data, new_schema):
if con_no_data.name == "snowflake" and os.environ.get("SNOWFLAKE_SNOWPARK"):
with pytest.raises(com.IbisError, match="Reconnecting is not supported"):
con_no_data.reconnect()
return

temp_table = gen_name(f"test_{con_no_data.name}_tmp")
table = con_no_data.create_table(temp_table, schema=new_schema, temp=True)

Expand Down

0 comments on commit e31eded

Please sign in to comment.