feat(backends): add from_connection to top level

deepyaman · Jul 22, 2024 · 5f846a5 · 5f846a5
1 parent 1f4fd37
commit 5f846a5
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/ibis/backends/snowflake/__init__.py b/ibis/backends/snowflake/__init__.py
@@ -85,7 +85,7 @@ class Backend(SQLBackend, CanCreateCatalog, CanCreateDatabase, CanCreateSchema):
     supports_python_udfs = True
 
     _latest_udf_python_version = (3, 10)
-    _top_level_methods = ("from_snowpark",)
+    _top_level_methods = ("from_connection", "from_snowpark")
 
     def __init__(self, *args, _from_snowpark: bool = False, **kwargs) -> None:
         super().__init__(*args, **kwargs)

diff --git a/ibis/backends/sql/__init__.py b/ibis/backends/sql/__init__.py
@@ -70,6 +70,8 @@ class SQLBackend(BaseBackend, _DatabaseSchemaHandler):
     compiler: ClassVar[SQLGlotCompiler]
     name: ClassVar[str]
 
+    _top_level_methods = ("from_connection",)
+
     @property
     def dialect(self) -> sg.Dialect:
         return self.compiler.dialect

diff --git a/ibis/backends/tests/test_client.py b/ibis/backends/tests/test_client.py
@@ -1620,7 +1620,9 @@ def test_insert_using_col_name_not_position(con, first_row, second_row, monkeypa
 DEFAULT_CON_ATTR = "con"
 
 
+@pytest.mark.parametrize("top_level", [True, False])
 @pytest.mark.never(["dask", "pandas", "polars"], reason="don't have connection concept")
-def test_from_connection(con):
-    new_con = type(con).from_connection(getattr(con, CON_ATTR.get(con.name, "con")))
+def test_from_connection(con, top_level):
+    backend = getattr(ibis, con.name) if top_level else type(con)
+    new_con = backend.from_connection(getattr(con, CON_ATTR.get(con.name, "con")))
     assert {"astronauts", "batting", "diamonds"} <= set(new_con.list_tables())