Skip to main content

shape_ast/parser/
data_sources.rs

1//! Parser for data source and query declarations
2
3use crate::ast::data_sources::{DataSourceDecl, QueryDecl};
4use crate::error::{Result, ShapeError};
5use pest::iterators::Pair;
6
7use super::expressions;
8use super::types::parse_type_annotation;
9use super::{Rule, pair_location, pair_span};
10
11/// Parse a datasource declaration: `datasource Name: Type = expr`
12pub fn parse_datasource_def(pair: Pair<Rule>) -> Result<DataSourceDecl> {
13    let pair_loc = pair_location(&pair);
14    let mut inner = pair.into_inner();
15
16    // Parse name
17    let name_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
18        message: "expected datasource name".to_string(),
19        location: Some(pair_loc.clone()),
20    })?;
21    let name_span = pair_span(&name_pair);
22    let name = name_pair.as_str().to_string();
23
24    // Parse type annotation (DataSource<T>)
25    let type_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
26        message: "expected type annotation after ':'".to_string(),
27        location: Some(pair_loc.clone()),
28    })?;
29    let schema = parse_type_annotation(type_pair)?;
30
31    // Parse provider expression
32    let expr_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
33        message: "expected provider expression after '='".to_string(),
34        location: Some(pair_loc),
35    })?;
36    let provider_expr = expressions::parse_expression(expr_pair)?;
37
38    Ok(DataSourceDecl {
39        name,
40        name_span,
41        schema,
42        provider_expr,
43    })
44}
45
46/// Parse a query declaration: `query Name: Type = expr`
47pub fn parse_query_decl(pair: Pair<Rule>) -> Result<QueryDecl> {
48    let pair_loc = pair_location(&pair);
49    let mut inner = pair.into_inner();
50
51    // Parse name
52    let name_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
53        message: "expected query name".to_string(),
54        location: Some(pair_loc.clone()),
55    })?;
56    let name_span = pair_span(&name_pair);
57    let name = name_pair.as_str().to_string();
58
59    // Parse type annotation (Query<T, Params>)
60    let type_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
61        message: "expected type annotation after ':'".to_string(),
62        location: Some(pair_loc.clone()),
63    })?;
64    let type_ann = parse_type_annotation(type_pair)?;
65
66    // Extract output_schema and params_schema from generic type
67    // Expected: Query<OutputType, ParamsType>
68    let (output_schema, params_schema) = match &type_ann {
69        crate::ast::TypeAnnotation::Generic { args, .. } if args.len() == 2 => {
70            (args[0].clone(), args[1].clone())
71        }
72        _ => {
73            // Single type: treat as output schema with empty params
74            (type_ann, crate::ast::TypeAnnotation::Object(vec![]))
75        }
76    };
77
78    // Parse initializer expression (e.g., sql(DB, "SELECT ..."))
79    let expr_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
80        message: "expected expression after '='".to_string(),
81        location: Some(pair_loc.clone()),
82    })?;
83    let init_expr = expressions::parse_expression(expr_pair)?;
84
85    // Extract source_name and sql from the expression
86    // For now, store them as defaults - the compiler will resolve these
87    let (source_name, sql, sql_span) = extract_query_parts(&init_expr, &pair_loc)?;
88
89    Ok(QueryDecl {
90        name,
91        name_span,
92        output_schema,
93        params_schema,
94        source_name,
95        sql,
96        sql_span,
97    })
98}
99
100/// Try to extract source name and SQL from a sql(Source, "...") expression.
101/// Falls back to defaults if the expression doesn't match this pattern.
102fn extract_query_parts(
103    expr: &crate::ast::Expr,
104    loc: &crate::error::SourceLocation,
105) -> Result<(String, String, crate::ast::Span)> {
106    // Try to match sql(Source, "SELECT ...")
107    if let crate::ast::Expr::FunctionCall {
108        name, args, span, ..
109    } = expr
110    {
111        if name == "sql" && args.len() == 2 {
112            let source = if let crate::ast::Expr::Identifier(src, _) = &args[0] {
113                src.clone()
114            } else {
115                return Err(ShapeError::ParseError {
116                    message: "expected data source identifier as first argument to sql()"
117                        .to_string(),
118                    location: Some(loc.clone()),
119                });
120            };
121            let sql_str = if let crate::ast::Expr::Literal(crate::ast::Literal::String(s), _) =
122                &args[1]
123            {
124                s.clone()
125            } else {
126                return Err(ShapeError::ParseError {
127                    message: "expected SQL string literal as second argument to sql()".to_string(),
128                    location: Some(loc.clone()),
129                });
130            };
131            return Ok((source, sql_str, *span));
132        }
133    }
134
135    // For non-sql expressions (e.g., provider functions), store empty defaults
136    Ok((String::new(), String::new(), crate::ast::Span::new(0, 0)))
137}