refactor(sqlalchemy): move _get_schema_using_query to base class

Author: cpcloud

ibis/backends/base/sql/alchemy/__init__.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+import abc
 import contextlib
 import getpass
 from operator import methodcaller
-from typing import TYPE_CHECKING, Any, Literal, Mapping
+from typing import TYPE_CHECKING, Any, Iterable, Literal, Mapping
 
 import sqlalchemy as sa
 
@@ -39,6 +40,8 @@
 if TYPE_CHECKING:
     import pandas as pd
 
+    import ibis.expr.datatypes as dt
+
 
 __all__ = (
     'BaseAlchemyBackend',
@@ -554,3 +557,11 @@ def _create_temp_view(self, view: sa.Table, definition: sa.sql.Selectable) -> No
             con.execute(compiled, **params)
         self._temp_views.add(raw_name)
         self._register_temp_view_cleanup(name, raw_name)
+
+    @abc.abstractmethod
+    def _metadata(self, query: str) -> Iterable[tuple[str, dt.DataType]]:
+        ...
+
+    def _get_schema_using_query(self, query: str) -> sch.Schema:
+        """Return an ibis Schema from a backend-specific SQL string."""
+        return sch.Schema.from_tuples(self._metadata(query))

ibis/backends/duckdb/__init__.py

@@ -471,10 +471,6 @@ def _metadata(self, query: str) -> Iterator[tuple[str, dt.DataType]]:
             ibis_type = parse(type)
             yield name, ibis_type(nullable=null.lower() == "yes")
 
-    def _get_schema_using_query(self, query: str) -> sch.Schema:
-        """Return an ibis Schema from a DuckDB SQL string."""
-        return sch.Schema.from_tuples(self._metadata(query))
-
     def _register_in_memory_table(self, table_op):
         df = table_op.data.to_frame()
         self.con.execute("register", (table_op.name, df))

ibis/backends/mssql/__init__.py

@@ -4,14 +4,16 @@
 
 import atexit
 import contextlib
-from typing import Literal
+from typing import TYPE_CHECKING, Iterable, Literal
 
 import sqlalchemy as sa
 
-import ibis.expr.schema as sch
 from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
 from ibis.backends.mssql.compiler import MsSqlCompiler
-from ibis.backends.mssql.datatypes import _FieldDescription, _type_from_result_set_info
+from ibis.backends.mssql.datatypes import _type_from_result_set_info
+
+if TYPE_CHECKING:
+    import ibis.expr.datatypes as dt
 
 
 class Backend(BaseAlchemyBackend):
@@ -52,17 +54,12 @@ def begin(self):
             finally:
                 bind.execute(f"SET DATEFIRST {previous_datefirst}")
 
-    def _get_schema_using_query(self, query):
+    def _metadata(self, query: str) -> Iterable[tuple[str, dt.DataType]]:
         with self.begin() as bind:
-            result = bind.execute(
+            for column in bind.execute(
                 f"EXEC sp_describe_first_result_set @tsql = N'{query}';"
-            )
-            result_set_info: list[_FieldDescription] = result.mappings().fetchall()
-        fields = [
-            (column['name'], _type_from_result_set_info(column))
-            for column in result_set_info
-        ]
-        return sch.Schema.from_tuples(fields)
+            ).mappings():
+                yield column["name"], _type_from_result_set_info(column)
 
     def _get_temp_view_definition(
         self,

ibis/backends/mysql/__init__.py

@@ -5,13 +5,12 @@
 import atexit
 import contextlib
 import warnings
-from typing import Literal
+from typing import Iterable, Literal
 
 import sqlalchemy as sa
 from sqlalchemy.dialects import mysql
 
 import ibis.expr.datatypes as dt
-import ibis.expr.schema as sch
 from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
 from ibis.backends.mysql.compiler import MySQLCompiler
 from ibis.backends.mysql.datatypes import _type_from_cursor_info
@@ -121,15 +120,14 @@ def begin(self):
                 except Exception as e:  # noqa: BLE001
                     warnings.warn(f"Couldn't reset MySQL timezone: {str(e)}")
 
-    def _get_schema_using_query(self, query: str) -> sch.Schema:
-        """Infer the schema of `query`."""
-        result = self.con.execute(f"SELECT * FROM ({query}) _ LIMIT 0")
-        cursor = result.cursor
-        fields = [
-            (field.name, _type_from_cursor_info(descr, field))
-            for descr, field in zip(cursor.description, cursor._result.fields)
-        ]
-        return sch.Schema.from_tuples(fields)
+    def _metadata(self, query: str) -> Iterable[tuple[str, dt.DataType]]:
+        with self.con.begin() as con:
+            result = con.execute(f"SELECT * FROM ({query}) _ LIMIT 0")
+            cursor = result.cursor
+            yield from (
+                (field.name, _type_from_cursor_info(descr, field))
+                for descr, field in zip(cursor.description, cursor._result.fields)
+            )
 
     def _get_temp_view_definition(
         self,

ibis/backends/postgres/__init__.py

@@ -3,17 +3,19 @@
 from __future__ import annotations
 
 import contextlib
-from typing import Literal
+from typing import TYPE_CHECKING, Iterable, Literal
 
 import sqlalchemy as sa
 
-import ibis.expr.schema as sch
 from ibis import util
 from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
 from ibis.backends.postgres.compiler import PostgreSQLCompiler
 from ibis.backends.postgres.datatypes import _get_type
 from ibis.backends.postgres.udf import udf as _udf
 
+if TYPE_CHECKING:
+    import ibis.expr.datatypes as dt
+
 
 class Backend(BaseAlchemyBackend):
     name = 'postgres'
@@ -172,7 +174,7 @@ def udf(
             language=language,
         )
 
-    def _get_schema_using_query(self, query: str) -> sch.Schema:
+    def _metadata(self, query: str) -> Iterable[tuple[str, dt.DataType]]:
         raw_name = util.guid()
         name = self.con.dialect.identifier_preparer.quote_identifier(raw_name)
         type_info_sql = f"""\
@@ -187,12 +189,9 @@ def _get_schema_using_query(self, query: str) -> sch.Schema:
 """
         with self.begin() as con:
             con.execute(f"CREATE TEMPORARY VIEW {name} AS {query}")
-            try:
-                type_info = con.execute(type_info_sql).fetchall()
-            finally:
-                con.execute(f"DROP VIEW {name}")
-        tuples = [(col, _get_type(typestr)) for col, typestr in type_info]
-        return sch.Schema.from_tuples(tuples)
+            type_info = con.execute(type_info_sql)
+            yield from ((col, _get_type(typestr)) for col, typestr in type_info)
+            con.execute(f"DROP VIEW IF EXISTS {name}")
 
     def _get_temp_view_definition(
         self,

ibis/backends/snowflake/__init__.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Iterable
 
 import snowflake.connector as sfc
 import sqlalchemy as sa
+import toolz
 from snowflake.sqlalchemy import URL
 
 import ibis.expr.operations as ops
@@ -20,6 +21,8 @@
 if TYPE_CHECKING:
     import pandas as pd
 
+    import ibis.expr.datatypes as dt
+
 
 _NATIVE_ARROW = True
 
@@ -117,16 +120,16 @@ def fetch_from_cursor(self, cursor, schema: sch.Schema) -> pd.DataFrame:
                 return schema.apply_to(df)
         return super().fetch_from_cursor(cursor, schema)
 
-    def _get_schema_using_query(self, query):
+    def _metadata(self, query: str) -> Iterable[tuple[str, dt.DataType]]:
         with self.begin() as bind:
             result = bind.execute(f"SELECT * FROM ({query}) t0 LIMIT 0")
             info_rows = bind.execute(f"DESCRIBE RESULT {result.cursor.sfqid!r}")
 
-        schema = {}
-        for name, raw_type, _, null, *_ in info_rows:
-            typ = parse(raw_type)
-            schema[name] = typ(nullable=null.upper() == "Y")
-        return sch.Schema.from_dict(schema)
+            for name, raw_type, null in toolz.pluck(
+                ["name", "type", "null?"], info_rows
+            ):
+                typ = parse(raw_type)
+                yield name, typ(nullable=null.upper() == "Y")
 
     def list_databases(self, like=None) -> list[str]:
         databases = [

ibis/backends/sqlite/__init__.py

@@ -18,22 +18,21 @@
 import sqlite3
 import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Iterable
 
 import sqlalchemy as sa
 from sqlalchemy.dialects.sqlite import DATETIME, TIMESTAMP
 
-if TYPE_CHECKING:
-    import ibis.expr.datatypes as dt
-    import ibis.expr.types as ir
-
-import ibis.expr.schema as sch
 from ibis.backends.base import Database
 from ibis.backends.base.sql.alchemy import BaseAlchemyBackend, to_sqla_type
 from ibis.backends.sqlite import udf
 from ibis.backends.sqlite.compiler import SQLiteCompiler
 from ibis.expr.schema import datatype
 
+if TYPE_CHECKING:
+    import ibis.expr.datatypes as dt
+    import ibis.expr.types as ir
+
 
 def to_datetime(value: str | None) -> datetime.datetime | None:
     """Convert a `str` to a `datetime` according to SQLite's rules.
@@ -217,7 +216,7 @@ def _table_from_schema(self, name, schema, database: str | None = None) -> sa.Ta
     def _current_schema(self) -> str | None:
         return self.current_database
 
-    def _get_schema_using_query(self, query: str) -> sch.Schema:
+    def _metadata(self, _: str) -> Iterable[tuple[str, dt.DataType]]:
         raise ValueError(
             "The SQLite backend cannot infer schemas from raw SQL - "
             "please specify the schema directly when calling `.sql` "

ibis/backends/trino/__init__.py

@@ -8,7 +8,6 @@
 import toolz
 
 import ibis.expr.datatypes as dt
-import ibis.expr.schema as sch
 from ibis import util
 from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
 from ibis.backends.trino.compiler import TrinoSQLCompiler
@@ -69,8 +68,3 @@ def _metadata(self, query: str) -> Iterator[tuple[str, dt.DataType]]:
             for name, type in toolz.pluck(["Column Name", "Type"], rows):
                 ibis_type = parse(type)
                 yield name, ibis_type(nullable=True)
-
-    def _get_schema_using_query(self, query: str) -> sch.Schema:
-        """Return an ibis Schema from a DuckDB SQL string."""
-        pairs = self._metadata(query)
-        return sch.Schema.from_tuples(pairs)