refactor(snowflake): add flags to supplemental JavaScript UDFs

ibis-project · Apr 16, 2023 · 054add4 · 054add4
1 parent 0092304
commit 054add4
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 16 deletions.
diff --git a/ibis/backends/snowflake/__init__.py b/ibis/backends/snowflake/__init__.py
@@ -136,19 +136,6 @@ def _VARIANT_to_python(self, _):
 }
 
 
-def _make_udf(name, defn, *, quote) -> str:
-    signature = ", ".join(
-        f'{quote(argname)} {sa.types.to_instance(typ)}'
-        for argname, typ in defn["inputs"].items()
-    )
-    return f"""\
-CREATE FUNCTION IF NOT EXISTS {name}({signature})
-RETURNS {sa.types.to_instance(defn["returns"])}
-LANGUAGE JAVASCRIPT
-AS
-$$ {defn["source"]} $$"""
-
-
 class Backend(BaseAlchemyBackend):
     name = "snowflake"
     compiler = SnowflakeCompiler
@@ -165,7 +152,27 @@ def _convert_kwargs(self, kwargs):
     @property
     def version(self) -> str:
         with self.begin() as con:
-            return con.exec_driver_sql("SELECT CURRENT_VERSION()").scalar()
+            return con.execute(sa.select(sa.func.current_version())).scalar()
+
+    def _compile_sqla_type(self, typ) -> str:
+        return sa.types.to_instance(typ).compile(dialect=self.con.dialect)
+
+    def _make_udf(self, name: str, defn) -> str:
+        dialect = self.con.dialect
+        quote = dialect.preparer(dialect).quote_identifier
+        signature = ", ".join(
+            f"{quote(argname)} {self._compile_sqla_type(typ)}"
+            for argname, typ in defn["inputs"].items()
+        )
+        return_type = self._compile_sqla_type(defn["returns"])
+        return f"""\
+CREATE FUNCTION IF NOT EXISTS {name}({signature})
+RETURNS {return_type}
+LANGUAGE JAVASCRIPT
+RETURNS NULL ON NULL INPUT
+IMMUTABLE
+AS
+$$ {defn["source"]} $$"""
 
     def do_connect(
         self,
@@ -216,7 +223,7 @@ def connect(dbapi_connection, connection_record):
                 try:
                     cur.execute("CREATE DATABASE IF NOT EXISTS ibis_udfs")
                     for name, defn in _SNOWFLAKE_MAP_UDFS.items():
-                        cur.execute(_make_udf(name, defn, quote=quote))
+                        cur.execute(self._make_udf(name, defn))
                     cur.execute(f"USE SCHEMA {quote(database)}.{quote(schema)}")
                 except Exception as e:  # noqa: BLE001
                     warnings.warn(

diff --git a/ibis/expr/types/maps.py b/ibis/expr/types/maps.py
@@ -449,5 +449,5 @@ def map(
     └──────────────────────┘
     """
     if values is None:
-        keys, values = list(keys.keys()), list(keys.values())
+        keys, values = tuple(keys.keys()), tuple(keys.values())
     return ops.Map(keys, values).to_expr()