fix(snowflake): make sqlalchemy 2.0 compatible

Author: cpcloud

ci/schema/snowflake.sql

@@ -1,10 +1,16 @@
-CREATE OR REPLACE FILE FORMAT ibis_csv_fmt
+USE WAREHOUSE ibis_testing;
+DROP DATABASE IF EXISTS ibis_testing;
+CREATE DATABASE IF NOT EXISTS ibis_testing;
+CREATE SCHEMA IF NOT EXISTS ibis_testing.ibis_testing;
+USE SCHEMA ibis_testing.ibis_testing;
+
+CREATE OR REPLACE FILE FORMAT ibis_testing
     type = 'CSV'
     field_delimiter = ','
     skip_header = 1
     field_optionally_enclosed_by = '"';
 
-CREATE OR REPLACE STAGE ibis_testing_stage file_format = ibis_csv_fmt;
+CREATE OR REPLACE STAGE ibis_testing file_format = ibis_testing;
 
 CREATE OR REPLACE TABLE diamonds (
     "carat" FLOAT,

ibis/backends/snowflake/__init__.py

@@ -55,29 +55,18 @@ def _convert_kwargs(self, kwargs):
     @property
     def version(self) -> str:
         with self.begin() as con:
-            return con.execute("SELECT CURRENT_VERSION()").scalar()
+            return con.exec_driver_sql("SELECT CURRENT_VERSION()").scalar()
 
     def do_connect(
-        self,
-        user: str,
-        password: str,
-        account: str,
-        database: str,
-        **kwargs,
+        self, user: str, password: str, account: str, database: str, **kwargs: Any
     ):
         dbparams = dict(zip(("database", "schema"), database.split("/", 1)))
         if dbparams.get("schema") is None:
             raise ValueError(
                 "Schema must be non-None. Pass the schema as part of the "
                 f"database e.g., {dbparams['database']}/my_schema"
             )
-        url = URL(
-            account=account,
-            user=user,
-            password=password,
-            **dbparams,
-            **kwargs,
-        )
+        url = URL(account=account, user=user, password=password, **dbparams, **kwargs)
         self.database_name = dbparams["database"]
         return super().do_connect(sa.create_engine(url))
 
@@ -88,11 +77,11 @@ def _get_sqla_table(
         cols = []
         identifier = name if not schema else schema + "." + name
         with self.begin() as con:
-            cur = con.execute(sa.text(f"DESCRIBE TABLE {identifier}"))
-            for (colname, *_), colinfo in zip(cur, inspected):
-                del colinfo["name"]
+            cur = con.execute(sa.text(f"DESCRIBE TABLE {identifier}")).mappings()
+            for colname, colinfo in zip(toolz.pluck("name", cur), inspected):
+                colinfo["name"] = colname
                 colinfo["type_"] = colinfo.pop("type")
-                cols.append(sa.Column(colname, **colinfo, quote=True))
+                cols.append(sa.Column(**colinfo, quote=True))
         return sa.Table(
             name,
             self.meta,
@@ -139,6 +128,6 @@ def list_databases(self, like=None) -> list[str]:
                 row.database_name
                 for row in con.execute(
                     sa.text('SELECT database_name FROM information_schema.databases')
-                )
+                ).mappings()
             ]
         return self._filter_with_like(databases, like)

ibis/backends/snowflake/tests/conftest.py

@@ -11,22 +11,21 @@
 import sqlalchemy as sa
 
 import ibis
-from ibis.backends.conftest import TEST_TABLES
+from ibis.backends.conftest import TEST_TABLES, init_database
 from ibis.backends.tests.base import BackendTest, RoundAwayFromZero
 
 if TYPE_CHECKING:
     from ibis.backends.base import BaseBackend
 
 
 def copy_into(con, data_dir: Path, table: str) -> None:
-    stage = "ibis_testing_stage"
+    stage = "ibis_testing"
     csv = f"{table}.csv"
-    src = data_dir / csv
-    con.execute(sa.text(f"PUT file://{src.absolute()} @{stage}/{csv}"))
-    con.execute(
-        sa.text(
-            f"COPY INTO {table} FROM @{stage}/{csv} FILE_FORMAT = (FORMAT_NAME = ibis_csv_fmt)"
-        )
+    con.exec_driver_sql(
+        f"PUT file://{data_dir.joinpath(csv).absolute()} @{stage}/{csv}"
+    )
+    con.exec_driver_sql(
+        f"COPY INTO {table} FROM @{stage}/{csv} FILE_FORMAT = (FORMAT_NAME = ibis_testing)"
     )
 
 
@@ -36,10 +35,7 @@ def __init__(self, data_directory: Path) -> None:
 
     @staticmethod
     def _load_data(
-        data_dir,
-        script_dir,
-        database: str = "ibis_testing",
-        **_: Any,
+        data_dir, script_dir, database: str = "ibis_testing", **_: Any
     ) -> None:
         """Load test data into a Snowflake backend instance.
 
@@ -53,26 +49,24 @@ def _load_data(
 
         pytest.importorskip("snowflake.connector")
         pytest.importorskip("snowflake.sqlalchemy")
-        schema = (script_dir / 'schema' / 'snowflake.sql').read_text()
-
-        con = TestConf.connect(data_dir)
 
-        with con.begin() as con:
-            con.execute(sa.text("USE WAREHOUSE ibis_testing"))
-            con.execute(sa.text(f"DROP DATABASE IF EXISTS {database}"))
-            con.execute(sa.text(f"CREATE DATABASE IF NOT EXISTS {database}"))
-            con.execute(sa.text(f"CREATE SCHEMA IF NOT EXISTS {database}.ibis_testing"))
-            con.execute(sa.text(f"USE SCHEMA {database}.ibis_testing"))
+        if (snowflake_url := os.environ.get("SNOWFLAKE_URL")) is None:
+            pytest.skip("SNOWFLAKE_URL environment variable is not defined")
 
-            for stmt in filter(None, map(str.strip, schema.split(';'))):
-                con.execute(sa.text(stmt))
+        with script_dir.joinpath('schema', 'snowflake.sql').open() as schema:
+            con = init_database(
+                url=sa.engine.make_url(snowflake_url).set(database=""),
+                database=database,
+                schema=schema,
+            )
 
+        with con.begin() as c:
             # not much we can do to make this faster, but running these in
             # multiple threads seems to save about 2x
             with concurrent.futures.ThreadPoolExecutor() as exe:
                 for result in concurrent.futures.as_completed(
                     map(
-                        partial(exe.submit, partial(copy_into, con, data_dir)),
+                        partial(exe.submit, partial(copy_into, c, data_dir)),
                         TEST_TABLES,
                     )
                 ):