translate_query.rs

use async_recursion::async_recursion;
use sqlparser::ast::{FunctionArg, FunctionArgExpr, Query, Select, SelectItem, SetExpr, TableFactor, TableWithJoins};
use std::collections::HashMap;

use super::expressions::{
  translate_expr::translate_expr, translate_table_with_joins::translate_table_with_joins,
  translate_wildcard_expr::translate_wildcard_expr,
};
use crate::ts_generator::sql_parser::quoted_strings::{DisplayIndent, DisplayTableAlias};
use crate::{
  core::connection::DBConn,
  ts_generator::{
    errors::TsGeneratorError,
    sql_parser::expressions::translate_table_with_joins::get_default_table,
    types::ts_query::{TsFieldType, TsQuery},
  },
};

pub async fn translate_select(
  ts_query: &mut TsQuery,
  table_with_joins: &Option<Vec<TableWithJoins>>,
  select: &Select,
  db_conn: &DBConn,
  alias: Option<&str>,
  is_selection: bool,
) -> Result<(), TsGeneratorError> {
  let projection = select.clone().projection;

  // We create a new table with joins within the scope of this select (it could be within a subquery)
  // full_table_with_joins should contain all tables from the parent and keeping it within this query's scope
  let mut full_table_with_joins: Vec<TableWithJoins> = vec![];
  let child_table_with_joins = select.clone().from;

  // The most inner table should be extended first as it will be the default table within the subquery
  full_table_with_joins.extend(child_table_with_joins.clone());

  if table_with_joins.is_some() {
    let table_with_joins = table_with_joins.as_ref().unwrap();
    full_table_with_joins.extend(table_with_joins.clone());
  }

  let full_table_with_joins = &Some(full_table_with_joins.clone());

  // Process table functions in FROM clause to extract parameters and column definitions
  // Note: Table-valued functions (e.g., jsonb_to_recordset) can be parsed as either TableFactor::Table with args or TableFactor::Function.
  for twj in &child_table_with_joins {
    match &twj.relation {
      TableFactor::Table {
        args: Some(table_fn_args),
        alias,
        ..
      } => {
        // Extract parameters from function arguments
        for arg in &table_fn_args.args {
          if let FunctionArg::Unnamed(FunctionArgExpr::Expr(expr))
          | FunctionArg::Named {
            arg: FunctionArgExpr::Expr(expr),
            ..
          } = arg
          {
            // Process the expression to extract any parameters
            translate_expr(expr, &None, full_table_with_joins, None, ts_query, db_conn, false).await?;
          }
        }

        // Extract column definitions from alias if present
        // e.g., AS t(id INT, name TEXT)
        if let Some(alias) = alias {
          let table_name = DisplayTableAlias(alias).to_string();
          let mut columns = HashMap::new();

          for col_def in &alias.columns {
            let col_name = DisplayIndent(&col_def.name).to_string();
            if let Some(data_type) = &col_def.data_type {
              let ts_type = TsFieldType::from_sqlparser_datatype(data_type);
              columns.insert(col_name, ts_type);
            }
          }

          if !columns.is_empty() {
            ts_query.table_valued_function_columns.insert(table_name, columns);
          }
        }
      }
      TableFactor::Function { args, alias, .. } => {
        // Handle LATERAL functions
        for arg in args {
          if let FunctionArg::Unnamed(FunctionArgExpr::Expr(expr))
          | FunctionArg::Named {
            arg: FunctionArgExpr::Expr(expr),
            ..
          } = arg
          {
            translate_expr(expr, &None, full_table_with_joins, None, ts_query, db_conn, false).await?;
          }
        }

        // Extract column definitions from LATERAL function alias if present
        if let Some(alias) = alias {
          let table_name = DisplayTableAlias(alias).to_string();
          let mut columns = HashMap::new();

          for col_def in &alias.columns {
            let col_name = DisplayIndent(&col_def.name).to_string();
            if let Some(data_type) = &col_def.data_type {
              let ts_type = TsFieldType::from_sqlparser_datatype(data_type);
              columns.insert(col_name, ts_type);
            }
          }

          if !columns.is_empty() {
            ts_query.table_valued_function_columns.insert(table_name, columns);
          }
        }
      }
      _ => {}
    }
  }

  // Handle all select projects and figure out each field's type
  for select_item in projection {
    // Determine the default table name within the scope of this select item
    let table_name_owned: Option<String>;
    let mut table_name: Option<&str> = None;
    if full_table_with_joins.is_some() && !full_table_with_joins.as_ref().unwrap().is_empty() {
      table_name_owned = Some(
        translate_table_with_joins(full_table_with_joins, &select_item).map_err(|_| {
          TsGeneratorError::UnknownErrorWhileProcessingTableWithJoins(format!(
            "Default FROM table is not found from the query: {}. Ensure your query has a valid FROM clause.",
            select
          ))
        })?,
      );
      table_name = table_name_owned.as_deref();
    }

    match &select_item {
      SelectItem::UnnamedExpr(unnamed_expr) => {
        translate_expr(
          unnamed_expr,
          &table_name,
          full_table_with_joins,
          alias,
          ts_query,
          db_conn,
          is_selection,
        )
        .await?;
      }
      SelectItem::ExprWithAlias { expr, alias } => {
        let alias = DisplayIndent(alias).to_string();

        translate_expr(
          expr,
          &table_name,
          full_table_with_joins,
          Some(alias.as_str()),
          ts_query,
          db_conn,
          is_selection,
        )
        .await?;
      }
      SelectItem::QualifiedWildcard(_, _) => {
        // TODO: If there's are two tables and two qualifieid wildcards are provided
        // It will simply generate types for both tables' columns
        // Should we namespace each field based on the table alias? e.g. table1_field1, table2_field1
        if is_selection {
          translate_wildcard_expr(select, ts_query, db_conn).await?;
        }
      }
      SelectItem::Wildcard(_) => {
        if is_selection {
          translate_wildcard_expr(select, ts_query, db_conn).await?;
        }
      }
    }
  }

  // If there's any WHERE statements, process it
  if let Some(selection) = &select.selection {
    let current_scope_table_name = get_default_table(&child_table_with_joins);
    let current_scope_table_name = current_scope_table_name.as_str();
    translate_expr(
      selection,
      &Some(current_scope_table_name),
      full_table_with_joins,
      None,
      ts_query,
      db_conn,
      false,
    )
    .await?;
  }
  Ok(())
}

/// Translates a query and workout ts_query's results and params
#[async_recursion]
pub async fn translate_query(
  ts_query: &mut TsQuery,
  // this parameter is used to stack table_with_joins while recursing through subqueries
  // If there is only 1 entry of table_with_joins, it means it's processing the top level entity
  table_with_joins: &Option<Vec<TableWithJoins>>,
  query: &Query,
  db_conn: &DBConn,
  alias: Option<&str>,
  is_selection: bool,
) -> Result<(), TsGeneratorError> {
  // Process CTEs (WITH clause) before the main query body.
  // Each CTE is processed to extract its output columns, which are then registered
  // as virtual table columns so the main query can reference them.
  if let Some(with) = &query.with {
    for cte in &with.cte_tables {
      let cte_name = DisplayIndent(&cte.alias.name).to_string();
      let mut cte_ts_query = TsQuery::new(cte_name.clone());
      translate_query(&mut cte_ts_query, &None, &cte.query, db_conn, None, true).await?;

      // Merge parameters from the CTE body into the outer query's params.
      // Parameters inside CTE bodies ($1, $2, ?) belong to the overall query's parameter list.
      for (idx, types) in &cte_ts_query.params {
        ts_query.params.insert(*idx, types.clone());
      }

      // Extract the CTE's output columns and register them as virtual table columns
      // so the outer query can look them up just like table-valued function columns
      let cte_columns: HashMap<String, TsFieldType> = cte_ts_query
        .result
        .into_iter()
        .map(|(col_name, types)| {
          // Take the first non-null type; fall back to Any if only Null is present
          let ts_type = types
            .into_iter()
            .find(|t| *t != TsFieldType::Null)
            .unwrap_or(TsFieldType::Any);
          (col_name, ts_type)
        })
        .collect();

      ts_query.table_valued_function_columns.insert(cte_name, cte_columns);
    }
  }

  let body = *query.body.clone();
  match body {
    SetExpr::Select(select) => {
      translate_select(ts_query, table_with_joins, &select, db_conn, alias, is_selection).await
    }
    _ => Err(TsGeneratorError::Unknown(format!(
      "Unknown query type while processing query: {query}"
    ))),
  }
}