Implement store and load for wrapped types (for Hive & Iceberg)

ARCHITECTURE.md

@@ -7,20 +7,21 @@ Here's a high level overview of how `joinery` works.
 Compilation procedes in several phases:
 
 1. [Tokenize](./src/tokenizer.rs).
-    - Split the source into identifiers, punctuation, literals, etc. All tokens contain the original source code, location information, and surrounding whitespace.
+   - Split the source into identifiers, punctuation, literals, etc. All tokens contain the original source code, location information, and surrounding whitespace.
 2. [Parse into AST](./src/ast.rs).
-    - We use the [`peg` crate](https://docs.rs/peg/). This is a [Parsing Expression Grammar](https://en.wikipedia.org/wiki/Parsing_expression_grammar) (PEG) parser. This is a bit _ad hoc_ as grammars go, but `peg` is a very nice library.
-    - We make heavy use of `#[derive]` macros to implement the AST types.
+   - We use the [`peg` crate](https://docs.rs/peg/). This is a [Parsing Expression Grammar](https://en.wikipedia.org/wiki/Parsing_expression_grammar) (PEG) parser. This is a bit _ad hoc_ as grammars go, but `peg` is a very nice library.
+   - We make heavy use of `#[derive]` macros to implement the AST types.
 3. [Check types](./src/infer/mod.rs).
-    - The internal type system is defined in [`src/types.rs`](./src/types.rs). This is distinct from the simplisitic "source level" type system parsed by [`src/ast.rs`](./src/ast.rs), and better suited to doing inference.
-    - Name lookup is handled in [`src/scopes.rs`](./src/scopes.rs). Note that SQL requires several different kinds of scopes.
+   - The internal type system is defined in [`src/types.rs`](./src/types.rs). This is distinct from the simplisitic "source level" type system parsed by [`src/ast.rs`](./src/ast.rs), and better suited to doing inference.
+   - Name lookup is handled in [`src/scopes.rs`](./src/scopes.rs). Note that SQL requires several different kinds of scopes.
+   - Type checking also needs to know about "memory" types (like Trino's UUID) versus "storage" types (like Trino's VARCHAR when using Hive, which doesn't allow storing UUID). And it needs make sure that all appropriate `LoadExpression` and `StoreExpression` values get inserted.
 4. [Apply transforms](./src/transforms/mod.rs).
-    - A list of transforms is supplied by each database driver.
-    - Transforms use Rust pattern-matching to match parts of the AST, and build new AST nodes using `sql_quote!`. Note that `sql_quote!` outputs _tokens_, so we need to call back into the parser. This is closely patterned after Rust programmatic macros using [`syn`](https://docs.rs/syn/) and [`quote`](https://docs.rs/quote/).
-    - After applying a transform, we _may_ need to check types again to support later transforms. This works a bit like an LLVM analysis pass, where specific transforms may indicate that the require types, and the harness ensures that valid types are available.
-    - The output of a transform must be structurally valid BigQuery SQL, though after a certain point it may no longer type check.
+   - A list of transforms is supplied by each database driver.
+   - Transforms use Rust pattern-matching to match parts of the AST, and build new AST nodes using `sql_quote!`. Note that `sql_quote!` outputs _tokens_, so we need to call back into the parser. This is closely patterned after Rust programmatic macros using [`syn`](https://docs.rs/syn/) and [`quote`](https://docs.rs/quote/).
+   - After applying a transform, we _may_ need to check types again to support later transforms. This works a bit like an LLVM analysis pass, where specific transforms may indicate that the require types, and the harness ensures that valid types are available.
+   - The output of a transform must be structurally valid BigQuery SQL, though after a certain point it may no longer type check.
 5. [Emit SQL](./src/ast.rs).
-    - This consumes AST nodes and emits them as database-specific strings. We prefer to do as much work as possible using AST transforms, but sometimes we can't represent database-specific features in the AST.
+   - This consumes AST nodes and emits them as database-specific strings. We prefer to do as much work as possible using AST transforms, but sometimes we can't represent database-specific features in the AST.
 6. [Run](./src/drivers/mod.rs).
    - This is a slightly dodgy layer that knows how to run SQL. Mostly it's intended for running our test suites, not for production use. Some of the Rust database drivers have problems reading complex data types back into Rust.
 

LOAD.md

@@ -0,0 +1,19 @@
+All SQL must be run through joinery.
+Every time that we create a table, we need to insert its BigQuery `CREATE TABLE name (col...)` into pg.
+That table looks like this:
+
+```sql
+PRIMARY KEY (bq_project, bq_dataset, bq_table_name)
+
+bq_project -- The BigQuery project, which gets mapped to a Trino catalog somehow
+bq_dataset -- The BigQuery dataset, which is a trino schema with the same name
+bq_table_name -- The BigQuery table, which is a trino table with the same name
+create_table_sql -- this is always a raw typed CREATE TABLE statement in BigQuery SQL. CREATE TABLE (my_col data_type, ...);
+```
+
+When we run another SQL query tomorrow, we need to make sure that we have access to the BigQuery table names and their `CREATE TABLE` SQL
+We can load those table definitions into our scope
+Then run type inference normally
+
+So when we try to access prod_gke.my_dataset.my_table,
+...we find a CREATE TABLE for it, parse it, and inject it in the scope

src/ast.rs

@@ -34,12 +34,14 @@ use crate::{
         trino::{TrinoString, KEYWORDS as TRINO_KEYWORDS},
     },
     errors::{format_err, Error, Result},
+    infer::{InferTypes, InsertStoreExpressions as _},
     known_files::{FileId, KnownFiles},
+    scope::{Scope, ScopeHandle},
     tokenizer::{
         tokenize_sql, EmptyFile, Ident, Keyword, Literal, LiteralValue, PseudoKeyword, Punct,
         RawToken, Span, Spanned, ToTokens, Token, TokenStream, TokenWriter,
     },
-    types::{StructType, TableType, ValueType},
+    types::{SimpleType, StructType, TableType, ValueType},
     util::{is_c_ident, AnsiIdent},
 };
 
@@ -621,6 +623,16 @@ pub struct SqlProgram {
     pub statements: NodeVec<Statement>,
 }
 
+impl SqlProgram {
+    /// Call `infer_types` for the first time, using the root scope, and doing
+    /// the one-time task of inserting [`StoreExpression`] where needed.
+    pub fn infer_types_for_first_time(&mut self) -> Result<(Option<TableType>, ScopeHandle)> {
+        self.insert_store_expressions()?;
+        let scope = Scope::root();
+        self.infer_types(&scope)
+    }
+}
+
 /// A statement in our abstract syntax tree.
 #[derive(Clone, Debug, Drive, DriveMut, Emit, EmitDefault, Spanned, ToTokens)]
 pub enum Statement {
@@ -825,7 +837,8 @@ pub enum Expression {
     Literal(Literal),
     BoolValue(Keyword),
     Null(Keyword),
-    Name(Name),
+    Name(NameExpression),
+    Store(StoreExpression),
     Cast(Cast),
     Is(IsExpression),
     In(InExpression),
@@ -855,8 +868,6 @@ pub enum Expression {
     FunctionCall(FunctionCall),
     Index(IndexExpression),
     FieldAccess(FieldAccessExpression),
-    Load(LoadExpression),
-    Store(StoreExpression),
 }
 
 impl Expression {
@@ -901,6 +912,111 @@ impl DatePart {
     }
 }
 
+/// A "load" expression, which transforms an SQL value from a "storage" type (eg
+/// "VARCHAR") to a "memory" type (eg "UUID"). Used for databases like Trino,
+/// where the storage types for a given connector may be more limited than the
+/// standard Trino memory types.
+///
+/// These are not found in the original parsed AST, but are added while
+/// transforming the AST.
+#[derive(Clone, Debug, Drive, DriveMut, EmitDefault, Spanned, ToTokens)]
+pub struct NameExpression {
+    /// **If** we need to do a load conversion, this will be the inferred memory
+    /// type.
+    #[emit(skip)]
+    #[to_tokens(skip)]
+    #[drive(skip)]
+    pub load_to_memory_type: Option<ValueType>,
+
+    /// Our underlying expression.
+    pub name: Name,
+}
+
+impl Emit for NameExpression {
+    fn emit(&self, t: Target, f: &mut TokenWriter<'_>) -> ::std::io::Result<()> {
+        match t {
+            // Target::BigQuery => {
+            //     f.write_token_start("%LOAD(")?;
+            //     self.name.emit(t, f)?;
+            //     f.write_token_start(")")
+            // }
+            Target::Trino(connector_type) if self.load_to_memory_type.is_some() => {
+                let bq_memory_type = self
+                    .load_to_memory_type
+                    .as_ref()
+                    .expect("memory_type should have been filled in by type inference");
+                let trino_memory_type =
+                    TrinoDataType::try_from(bq_memory_type).map_err(io::Error::other)?;
+                let transform = connector_type.storage_transform_for(&trino_memory_type);
+                let (prefix, suffix) = transform.load_prefix_and_suffix();
+
+                // Wrapping the expression in our prefix and suffix.
+                // If the expression was col_name containing '[1,2]' in Trino,
+                // BQ memory type -> JSON, Trino memory type -> JSON, Trino storage type -> VARCHAR
+                // The Trino storage type is dependent on what the connector can support.
+                // In this case, the wrapped version would be JSON_PARSE(col_name)
+                f.write_token_start(&prefix)?;
+                self.name.emit(t, f)?;
+                f.write_token_start(&suffix)
+            }
+            _ => self.name.emit(t, f),
+        }
+    }
+}
+
+/// A "store" expression, which transforms an SQL value from a "memory" type
+/// (eg "UUID") to a "storage" type (eg "VARCHAR"). Used for databases like
+/// Trino, where the storage types for a given connector may be more limited
+/// than the standard Trino memory types.
+///
+/// These are not found in the original parsed AST, but are added while
+/// transforming the AST.
+#[derive(Clone, Debug, Drive, DriveMut, EmitDefault, Spanned, ToTokens)]
+pub struct StoreExpression {
+    /// Inferred memory type.
+    #[emit(skip)]
+    #[to_tokens(skip)]
+    #[drive(skip)]
+    pub memory_type: Option<ValueType>,
+
+    /// Our underlying expression.
+    pub expression: Box<Expression>,
+}
+
+impl Emit for StoreExpression {
+    fn emit(&self, t: Target, f: &mut TokenWriter<'_>) -> ::std::io::Result<()> {
+        match t {
+            Target::BigQuery => {
+                f.write_token_start("%STORE(")?;
+                self.expression.emit(t, f)?;
+                f.write_token_start(")")
+            }
+            Target::Trino(connector_type) => {
+                let bq_memory_type = self
+                    .memory_type
+                    .as_ref()
+                    .expect("memory_type should have been filled in by type inference");
+
+                // If our bq_memory_type is NULL, we don't need to do any transforms because
+                // NULL is NULL in both storage and memory types and dbcrossbar_trino doesn't
+                // support NULL as a memory type.
+                if let ValueType::Simple(SimpleType::Null) = bq_memory_type {
+                    self.expression.emit(t, f)
+                } else {
+                    let trino_memory_type =
+                        TrinoDataType::try_from(bq_memory_type).map_err(io::Error::other)?;
+                    let transform = connector_type.storage_transform_for(&trino_memory_type);
+                    let (prefix, suffix) = transform.store_prefix_and_suffix();
+
+                    f.write_token_start(&prefix)?;
+                    self.expression.emit(t, f)?;
+                    f.write_token_start(&suffix)
+                }
+            }
+        }
+    }
+}
+
 /// A cast expression.
 #[derive(Clone, Debug, Drive, DriveMut, Emit, EmitDefault, Spanned, ToTokens)]
 pub struct Cast {
@@ -1633,93 +1749,6 @@ pub struct FieldAccessExpression {
     pub field_name: Ident,
 }
 
-/// A "load" expression, which transforms an SQL value from a "storage" type (eg
-/// "VARCHAR") to a "memory" type (eg "UUID"). Used for databases like Trino,
-/// where the storage types for a given connector may be more limited than the
-/// standard Trino memory types.
-///
-/// These are not found in the original parsed AST, but are added while
-/// transforming the AST.
-#[derive(Clone, Debug, Drive, DriveMut, EmitDefault, Spanned, ToTokens)]
-pub struct LoadExpression {
-    /// Inferred memory type.
-    #[emit(skip)]
-    #[to_tokens(skip)]
-    #[drive(skip)]
-    memory_type: Option<ValueType>,
-
-    /// Our underlying expression.
-    pub expression: Box<Expression>,
-}
-
-impl Emit for LoadExpression {
-    fn emit(&self, t: Target, f: &mut TokenWriter<'_>) -> ::std::io::Result<()> {
-        match t {
-            Target::Trino(connector_type) => {
-                let bq_memory_type = self
-                    .memory_type
-                    .as_ref()
-                    .expect("memory_type should have been filled in by type inference");
-                let trino_memory_type =
-                    TrinoDataType::try_from(bq_memory_type).map_err(io::Error::other)?;
-                let transform = connector_type.storage_transform_for(&trino_memory_type);
-                let (prefix, suffix) = transform.load_prefix_and_suffix();
-
-                // Wrapping the expression in our prefix and suffix.
-                // If the expression was col_name containing '[1,2]' in Trino,
-                // BQ memory type -> JSON, Trino memory type -> JSON, Trino storage type -> VARCHAR
-                // The Trino storage type is dependent on what the connector can support.
-                // In this case, the wrapped version would be JSON_PARSE(col_name)
-                f.write_token_start(&prefix)?;
-                self.expression.emit(t, f)?;
-                f.write_token_start(&suffix)
-            }
-            _ => self.emit_default(t, f),
-        }
-    }
-}
-
-/// A "store" expression, which transforms an SQL value from a "memory" type
-/// (eg "UUID") to a "storage" type (eg "VARCHAR"). Used for databases like
-/// Trino, where the storage types for a given connector may be more limited
-/// than the standard Trino memory types.
-///
-/// These are not found in the original parsed AST, but are added while
-/// transforming the AST.
-#[derive(Clone, Debug, Drive, DriveMut, EmitDefault, Spanned, ToTokens)]
-pub struct StoreExpression {
-    /// Inferred memory type.
-    #[emit(skip)]
-    #[to_tokens(skip)]
-    #[drive(skip)]
-    memory_type: Option<ValueType>,
-
-    /// Our underlying expression.
-    pub expression: Box<Expression>,
-}
-
-impl Emit for StoreExpression {
-    fn emit(&self, t: Target, f: &mut TokenWriter<'_>) -> ::std::io::Result<()> {
-        match t {
-            Target::Trino(connector_type) => {
-                let bq_memory_type = self
-                    .memory_type
-                    .as_ref()
-                    .expect("memory_type should have been filled in by type inference");
-                let trino_memory_type =
-                    TrinoDataType::try_from(bq_memory_type).map_err(io::Error::other)?;
-                let transform = connector_type.storage_transform_for(&trino_memory_type);
-                let (prefix, suffix) = transform.store_prefix_and_suffix();
-
-                f.write_token_start(&prefix)?;
-                self.expression.emit(t, f)?;
-                f.write_token_start(&suffix)
-            }
-            _ => self.emit_default(t, f),
-        }
-    }
-}
-
 /// An `AS` alias.
 #[derive(Clone, Debug, Drive, DriveMut, Emit, EmitDefault, Spanned, ToTokens)]
 pub struct Alias {
@@ -2418,7 +2447,7 @@ peg::parser! {
             // Things from here down might start with arbitrary identifiers, so
             // we need to be careful about the order.
             function_call:function_call() { Expression::FunctionCall(function_call) }
-            column_name:name() { Expression::Name(column_name) }
+            column_name:name() { Expression::Name(NameExpression { load_to_memory_type: None, name: column_name }) }
         }
 
         rule interval_expression() -> IntervalExpression

src/cmd/run.rs

@@ -3,10 +3,7 @@ use std::path::PathBuf;
 use clap::Parser;
 use tracing::instrument;
 
-use crate::{
-    ast::parse_sql, drivers, errors::Result, infer::InferTypes, known_files::KnownFiles,
-    scope::Scope,
-};
+use crate::{ast::parse_sql, drivers, errors::Result, known_files::KnownFiles};
 
 /// Run an SQL file using the specified database.
 #[derive(Debug, Parser)]
@@ -34,8 +31,7 @@ pub async fn cmd_run(files: &mut KnownFiles, opt: &RunOpt) -> Result<()> {
     let mut ast = parse_sql(files, file_id)?;
 
     // Run the type checker, but do not fail on errors.
-    let scope = Scope::root();
-    if let Err(err) = ast.infer_types(&scope) {
+    if let Err(err) = ast.infer_types_for_first_time() {
         err.emit(files);
         eprintln!("\nType checking failed. Manual fixes will probably be required!");
     }

src/cmd/sql_test.rs

@@ -17,9 +17,7 @@ use crate::{
     ast::{self, parse_sql, CreateTableStatement, CreateViewStatement, Target},
     drivers::{self, Driver},
     errors::{format_err, Context, Error, Result},
-    infer::InferTypes,
     known_files::{FileId, KnownFiles},
-    scope::Scope,
 };
 
 /// Run SQL tests from a directory.
@@ -148,8 +146,7 @@ async fn run_test(
     let mut ast = parse_sql(files, file_id)?;
 
     // Type check the AST.
-    let scope = Scope::root();
-    ast.infer_types(&scope)?;
+    ast.infer_types_for_first_time()?;
 
     //eprintln!("SQLite3: {}", ast.emit_to_string(Target::SQLite3));
     let output_tables = find_output_tables(&ast)?;

src/cmd/transpile.rs

@@ -9,9 +9,7 @@ use crate::{
     ast::{parse_sql, Emit},
     drivers,
     errors::Result,
-    infer::InferTypes,
     known_files::KnownFiles,
-    scope::Scope,
 };
 
 /// Run SQL tests from a directory.
@@ -38,8 +36,7 @@ pub async fn cmd_transpile(files: &mut KnownFiles, opt: &TranspileOpt) -> Result
     let mut ast = parse_sql(files, file_id)?;
 
     // Run the type checker, but do not fail on errors.
-    let scope = Scope::root();
-    if let Err(err) = ast.infer_types(&scope) {
+    if let Err(err) = ast.infer_types_for_first_time() {
         err.emit(files);
         eprintln!("\nType checking failed. Manual fixes will probably be required!");
     }

src/drivers/mod.rs

@@ -9,7 +9,7 @@ use tracing::{debug, trace};
 use crate::{
     ast::{self, Emit, Target},
     errors::{format_err, Error, Result},
-    infer::InferTypes,
+    infer::InferTypes as _,
     scope::Scope,
     transforms::{Transform, TransformExtra},
 };