From 79e418bea3dbfa02e0faccd4c6d0b0f10cda6eab Mon Sep 17 00:00:00 2001
From: Eric Kidd <git@randomhacks.net>
Date: Wed, 18 Oct 2023 20:07:29 -0400
Subject: [PATCH] trino: Start transpiling arrays

---
 src/ast.rs               | 23 +++++++++++++++++++++--
 src/drivers/trino/mod.rs | 31 +++++++++++++++++++++++++++++--
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index 6ac9c0c..056d84b 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -1166,6 +1166,16 @@ impl Emit for ArrayExpression {
                 self.definition.emit(t, f)?;
                 self.delim2.with_token_str(")").emit(t, f)?;
             }
+            Target::Trino => {
+                if let Some(array_token) = &self.array_token {
+                    array_token.emit(t, f)?;
+                } else {
+                    write!(f, "ARRAY")?;
+                }
+                self.delim1.with_token_str("[").emit(t, f)?;
+                self.definition.emit(t, f)?;
+                self.delim2.with_token_str("]").emit(t, f)?;
+            }
             _ => self.emit_default(t, f)?,
         }
         Ok(())
@@ -1567,8 +1577,17 @@ impl Emit for DataType {
                 DataType::Timestamp(token) => {
                     token.with_token_str("TIMESTAMP WITH TIME ZONE").emit(t, f)
                 }
-                // TODO: Can we declare the element type?
-                DataType::Array { array_token, .. } => array_token.ensure_ws().emit(t, f),
+                DataType::Array {
+                    array_token,
+                    lt,
+                    data_type,
+                    gt,
+                } => {
+                    array_token.emit(t, f)?;
+                    lt.with_token_str("(").emit(t, f)?;
+                    data_type.emit(t, f)?;
+                    gt.with_token_str(")").emit(t, f)
+                }
                 // TODO: I think we can translate the column types?
                 DataType::Struct { struct_token, .. } => {
                     struct_token.with_token_str("ROW").emit(t, f)
diff --git a/src/drivers/trino/mod.rs b/src/drivers/trino/mod.rs
index 55bfca6..d6f609b 100644
--- a/src/drivers/trino/mod.rs
+++ b/src/drivers/trino/mod.rs
@@ -12,7 +12,7 @@ use crate::{
     ast::{self, Emit, Target},
     drivers::sqlite3::SQLite3String,
     errors::{format_err, Context, Error, Result},
-    transforms::{self, Transform},
+    transforms::{self, Transform, Udf},
 };
 
 use super::{sqlite3::SQLite3Ident, Column, Driver, DriverImpl, Locator};
@@ -20,6 +20,26 @@ use super::{sqlite3::SQLite3Ident, Column, Driver, DriverImpl, Locator};
 /// Our locator prefix.
 pub const TRINO_LOCATOR_PREFIX: &str = "trino:";
 
+// A `phf_map!` of BigQuery function names to native function names. Use
+// this for simple renaming.
+static FUNCTION_NAMES: phf::Map<&'static str, &'static str> = phf::phf_map! {
+    "ARRAY_LENGTH" => "CARDINALITY",
+};
+
+/// A `phf_map!` of BigQuery function names to UDFs.
+///
+/// TODO: I'm not even sure there's a way to define SQL UDFs in Trino.
+static UDFS: phf::Map<&'static str, &'static Udf> = phf::phf_map! {};
+
+/// Format a UDF.
+///
+/// TODO: I'm not even sure there's a way to define SQL UDFs in Trino.
+fn format_udf(udf: &Udf) -> String {
+    format!(
+        "CREATE OR REPLACE TEMP FUNCTION {} AS $$\n{}\n$$\n",
+        udf.decl, udf.sql
+    )
+}
 /// A locator for a Trino database. May or may not also work for Presto.
 #[derive(Debug)]
 pub struct TrinoLocator {
@@ -142,7 +162,14 @@ impl Driver for TrinoDriver {
     }
 
     fn transforms(&self) -> Vec<Box<dyn Transform>> {
-        vec![Box::new(transforms::OrReplaceToDropIfExists)]
+        vec![
+            Box::new(transforms::OrReplaceToDropIfExists),
+            Box::new(transforms::RenameFunctions::new(
+                &FUNCTION_NAMES,
+                &UDFS,
+                &format_udf,
+            )),
+        ]
     }
 
     #[tracing::instrument(skip(self))]