data-flow-control 0.1.2__tar.gz → 0.1.3__tar.gz

{data_flow_control-0.1.2 → data_flow_control-0.1.3}/Cargo.lock

@@ -1594,7 +1594,7 @@ dependencies = [
 
 [[package]]
 name = "passant-cli"
-version = "0.1.2"
+version = "0.1.3"
 dependencies = [
  "anyhow",
  "assert_cmd",
@@ -1606,7 +1606,7 @@ dependencies = [
 
 [[package]]
 name = "passant-core"
-version = "0.1.2"
+version = "0.1.3"
 dependencies = [
  "anyhow",
  "criterion",
@@ -1624,7 +1624,7 @@ dependencies = [
 
 [[package]]
 name = "passant-py"
-version = "0.1.2"
+version = "0.1.3"
 dependencies = [
  "passant-core",
  "pyo3",

{data_flow_control-0.1.2 → data_flow_control-0.1.3}/Cargo.toml

@@ -3,7 +3,7 @@ members = ["passant-core", "passant-py"]
 resolver = "2"
 
 [workspace.package]
-version = "0.1.2"
+version = "0.1.3"
 edition = "2024"
 license = "MIT"
 authors = ["Charlie Summers"]

{data_flow_control-0.1.2 → data_flow_control-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-flow-control
-Version: 0.1.2
+Version: 0.1.3
 Requires-Dist: boto3>=1.26.0
 Requires-Dist: duckdb>=1.3.0
 Requires-Dist: langchain>=1.3,<2

{data_flow_control-0.1.2 → data_flow_control-0.1.3}/passant-core/src/catalog.rs

@@ -16,6 +16,7 @@ use crate::policy::{PolicyIr, Resolution};
 use crate::policy_compile::{
     ParsedPolicyConstraint, parse_policy_constraint, parse_policy_constraint_with_registry,
 };
+use crate::sql::collect_query_projection_column_names;
 
 /// Catalog facts supplied at policy registration time (from any adapter snapshot).
 #[derive(Debug, Default, Clone)]
@@ -389,6 +390,24 @@ fn validate_pgn_policy_parsed(
         dimension_columns.insert(TableKey::new(alias), columns.clone());
         dimension_columns.insert(TableKey::new(base), columns);
     }
+    for (alias, query_sql) in dimension_queries {
+        let columns = collect_query_projection_column_names(query_sql)
+            .map_err(|err| {
+                RewriteError::catalog_with_context(
+                    ErrorKind::UnknownColumn,
+                    format!(
+                        "Dimension subquery '{alias}' has invalid projection for catalog validation: {err}"
+                    ),
+                    Some(alias.clone()),
+                    None,
+                    Some(constraint.to_string()),
+                    Some("catalog_validation"),
+                )
+            })?
+            .into_iter()
+            .collect::<HashSet<_>>();
+        dimension_columns.insert(TableKey::new(alias), columns);
+    }
 
     let source_names = sources
         .iter()
@@ -495,6 +514,7 @@ fn validate_pgn_policy_parsed(
             &table_key,
             dimension_tables,
             dimension_aliases,
+            dimension_queries,
             &dimension_columns,
         ) {
             if !columns.contains(&column_key) {
@@ -544,6 +564,7 @@ fn resolve_dimension_columns<'a>(
     table_key: &TableKey,
     dimension_tables: &[String],
     dimension_aliases: &HashMap<String, String>,
+    dimension_queries: &HashMap<String, String>,
     dimension_columns: &'a HashMap<TableKey, HashSet<String>>,
 ) -> Option<&'a HashSet<String>> {
     if dimension_tables
@@ -552,7 +573,13 @@ fn resolve_dimension_columns<'a>(
     {
         return dimension_columns.get(table_key);
     }
+    if dimension_queries.contains_key(table_key.as_str()) {
+        return dimension_columns.get(table_key);
+    }
     if let Some(base) = dimension_aliases.get(table_key.as_str()) {
+        if dimension_queries.contains_key(base) {
+            return dimension_columns.get(table_key);
+        }
         return dimension_columns.get(&TableKey::new(base));
     }
     dimension_columns.get(table_key)
@@ -1054,6 +1081,63 @@ mod tests {
             .expect("dimension alias column should validate");
     }
 
+    #[test]
+    fn validates_constraint_columns_through_dimension_subquery_alias() {
+        let catalog = sample_catalog();
+        let mut dimension_aliases = HashMap::new();
+        dimension_aliases.insert("u".to_string(), "u".to_string());
+        let mut dimension_queries = HashMap::new();
+        dimension_queries.insert(
+            "u".to_string(),
+            "(SELECT user_id FROM session_user)".to_string(),
+        );
+        let policy = PolicyIr::Pgn {
+            sources: vec!["foo".into()],
+            required_sources: Vec::new(),
+            dimension_tables: Vec::new(),
+            dimension_aliases,
+            dimension_queries,
+            sink: None,
+            sink_alias: None,
+            source_aliases: HashMap::new(),
+            constraint: "max(foo.id) > 1 AND u.user_id = 1".into(),
+            on_fail: Resolution::Remove,
+            description: None,
+        };
+        catalog
+            .validate_policy(&policy, &sample_registry())
+            .expect("dimension subquery alias column should validate");
+    }
+
+    #[test]
+    fn rejects_unknown_dimension_subquery_column() {
+        let catalog = sample_catalog();
+        let mut dimension_aliases = HashMap::new();
+        dimension_aliases.insert("u".to_string(), "u".to_string());
+        let mut dimension_queries = HashMap::new();
+        dimension_queries.insert(
+            "u".to_string(),
+            "(SELECT user_id FROM session_user)".to_string(),
+        );
+        let policy = PolicyIr::Pgn {
+            sources: vec!["foo".into()],
+            required_sources: Vec::new(),
+            dimension_tables: Vec::new(),
+            dimension_aliases,
+            dimension_queries,
+            sink: None,
+            sink_alias: None,
+            source_aliases: HashMap::new(),
+            constraint: "max(foo.id) > 1 AND u.missing = 1".into(),
+            on_fail: Resolution::Remove,
+            description: None,
+        };
+        let err = catalog
+            .validate_policy(&policy, &sample_registry())
+            .unwrap_err();
+        assert_eq!(err.kind(), ErrorKind::UnknownColumn);
+    }
+
     #[test]
     fn rejects_unknown_dimension_column() {
         let mut catalog = sample_catalog();

{data_flow_control-0.1.2 → data_flow_control-0.1.3}/passant-core/src/policy_store/compiled.rs

@@ -140,6 +140,7 @@ impl PolicyStore {
                 &compiled.ast,
                 policy.dimension_tables(),
                 policy.dimension_aliases(),
+                policy.dimension_queries(),
                 &source_key_set,
             )
         });
@@ -226,6 +227,7 @@ pub(crate) fn compile_branch_policy(
                 &constraint_ast,
                 policy.dimension_tables(),
                 policy.dimension_aliases(),
+                policy.dimension_queries(),
                 &source_key_set,
             )
         };

{data_flow_control-0.1.2 → data_flow_control-0.1.3}/passant-core/src/rewriter/dimensions.rs

@@ -56,6 +56,7 @@ pub(crate) fn inject_policy_dimensions(
             constraint,
             dimension_tables,
             dimension_aliases,
+            dimension_queries,
             &source_keys,
         )?
     };
@@ -175,11 +176,9 @@ pub(crate) fn inject_policy_dimensions(
             append_dimension_join(select, factor, join.on)?;
             joined_keys.insert(table_key);
         } else {
-            skipped.push(alias.clone());
-            warnings.push(format!(
-                "Dimension subquery '{alias}' was not joined: no equality link to the query or \
-                 another dimension in the policy constraint and catalog row count is not 1"
-            ));
+            warnings.push(expansion_warning(alias));
+            append_dimension_join(select, factor, None)?;
+            joined_keys.insert(table_key);
         }
     }
 
@@ -405,9 +404,10 @@ pub(crate) fn compile_dimension_join_plan(
     constraint_ast: &Expr,
     dimension_tables: &[String],
     dimension_aliases: &HashMap<String, String>,
+    dimension_queries: &HashMap<String, String>,
     source_keys: &HashSet<TableKey>,
 ) -> Option<DimensionJoinPlan> {
-    if dimension_tables.is_empty() {
+    if dimension_tables.is_empty() && dimension_queries.is_empty() {
         return None;
     }
     let mut dimension_keys = HashSet::new();
@@ -419,6 +419,9 @@ pub(crate) fn compile_dimension_join_plan(
             }
         }
     }
+    for alias in dimension_queries.keys() {
+        dimension_keys.insert(TableKey::new(alias));
+    }
     let mut source_conditions = HashMap::new();
     let mut dimension_edges = Vec::new();
     collect_dimension_join_graph(
@@ -454,17 +457,24 @@ fn dimension_join_graph(
     constraint: &str,
     dimension_tables: &[String],
     dimension_aliases: &HashMap<String, String>,
+    dimension_queries: &HashMap<String, String>,
     source_keys: &HashSet<TableKey>,
 ) -> Result<(HashMap<TableKey, Expr>, Vec<DimensionJoinEdge>), RewriteError> {
     let constraint = preprocess_policy_constraint(constraint);
     let expr = parse_projection_expr(&constraint).map_err(|err| {
         RewriteError::unsupported_statement(format!("dimension join analysis: {err}"))
     })?;
-    let plan = compile_dimension_join_plan(&expr, dimension_tables, dimension_aliases, source_keys)
-        .unwrap_or(DimensionJoinPlan {
-            source_conditions: HashMap::new(),
-            dimension_edges: Vec::new(),
-        });
+    let plan = compile_dimension_join_plan(
+        &expr,
+        dimension_tables,
+        dimension_aliases,
+        dimension_queries,
+        source_keys,
+    )
+    .unwrap_or(DimensionJoinPlan {
+        source_conditions: HashMap::new(),
+        dimension_edges: Vec::new(),
+    });
     Ok((plan.source_conditions, plan.dimension_edges))
 }
 

{data_flow_control-0.1.2 → data_flow_control-0.1.3}/passant-core/src/sql/columns.rs

@@ -1,10 +1,13 @@
 //! Qualified-column extraction from constraint ASTs.
 
 use sqlparser::ast::{
-    Array, Expr, Function, FunctionArg, FunctionArgExpr, FunctionArguments, Map, Subscript,
+    Array, Expr, Function, FunctionArg, FunctionArgExpr, FunctionArguments, Map, SelectItem,
+    SetExpr, Statement, Subscript,
 };
 
-use crate::identifiers::QualifiedColumn;
+use crate::diagnostics::RewriteError;
+use crate::identifiers::{ColumnName, QualifiedColumn};
+use crate::parser::parse_query;
 
 pub fn collect_qualified_columns_from_expr(expr: &Expr) -> Vec<QualifiedColumn> {
     let mut found = Vec::new();
@@ -12,6 +15,66 @@ pub fn collect_qualified_columns_from_expr(expr: &Expr) -> Vec<QualifiedColumn>
     found
 }
 
+/// Output column names from a dimension subquery SELECT list.
+pub fn collect_query_projection_column_names(query_sql: &str) -> Result<Vec<String>, RewriteError> {
+    let select = dimension_subquery_select(query_sql)?;
+    let mut names = Vec::new();
+    for item in &select.projection {
+        match item {
+            SelectItem::UnnamedExpr(expr) => {
+                let Some(name) = projection_column_name(expr) else {
+                    return Err(RewriteError::unsupported_statement(format!(
+                        "dimension subquery projection must name columns explicitly: {item}"
+                    )));
+                };
+                names.push(ColumnName::new(name).key());
+            }
+            SelectItem::ExprWithAlias { alias, .. } => {
+                names.push(ColumnName::new(alias.value.as_str()).key());
+            }
+            SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _) => {
+                return Err(RewriteError::unsupported_statement(
+                    "dimension subquery projections with * are not supported for catalog validation",
+                ));
+            }
+        }
+    }
+    if names.is_empty() {
+        return Err(RewriteError::unsupported_statement(
+            "dimension subquery must project at least one column",
+        ));
+    }
+    Ok(names)
+}
+
+fn dimension_subquery_select(query_sql: &str) -> Result<sqlparser::ast::Select, RewriteError> {
+    let statement = parse_query(query_sql)?;
+    let Statement::Query(query) = statement else {
+        return Err(RewriteError::unsupported_statement(
+            "dimension subquery must be a query",
+        ));
+    };
+    select_from_set_expr(*query.body)
+}
+
+fn select_from_set_expr(body: SetExpr) -> Result<sqlparser::ast::Select, RewriteError> {
+    match body {
+        SetExpr::Select(select) => Ok(*select),
+        SetExpr::Query(query) => select_from_set_expr(*query.body),
+        other => Err(RewriteError::unsupported_statement(format!(
+            "dimension subquery must be a SELECT, got {other}"
+        ))),
+    }
+}
+
+fn projection_column_name(expr: &Expr) -> Option<String> {
+    match expr {
+        Expr::Identifier(ident) => Some(ident.value.clone()),
+        Expr::CompoundIdentifier(parts) => parts.last().map(|ident| ident.value.clone()),
+        _ => None,
+    }
+}
+
 fn visit_expr(expr: &Expr, found: &mut Vec<QualifiedColumn>) {
     if let Some(column) = QualifiedColumn::from_expr(expr) {
         found.push(column);

{data_flow_control-0.1.2 → data_flow_control-0.1.3}/passant-core/src/sql/mod.rs

@@ -25,5 +25,5 @@ pub use builders::{
     query_from_select, sanitize_projection_alias, scalar_subquery, statement_from_query,
     string_concat, string_literal, table_alias, table_factor, with_ctes, wrap_table_with_filter,
 };
-pub use columns::collect_qualified_columns_from_expr;
+pub use columns::{collect_qualified_columns_from_expr, collect_query_projection_column_names};
 pub use expr::{rename_table_refs, replace_expr_subtrees, unqualify_table_refs};

{data_flow_control-0.1.2 → data_flow_control-0.1.3}/passant-core/tests/gaps/dimensions.rs

@@ -132,6 +132,43 @@ fn dimension_alias_qualifiers_rewrite_after_injection() {
     );
 }
 
+#[test]
+fn dimension_subquery_alias_in_constraint_rewrites() {
+    let mut catalog = TableCatalog::new();
+    catalog.register_table("foo", vec!["id".to_string()]);
+    catalog.register_table("session_user", vec!["user_id".to_string()]);
+    catalog.register_table_row_count("session_user", 1);
+    let mut dimension_aliases = HashMap::new();
+    dimension_aliases.insert("u".to_string(), "u".to_string());
+    let mut dimension_queries = HashMap::new();
+    dimension_queries.insert(
+        "u".to_string(),
+        "(SELECT user_id FROM session_user)".to_string(),
+    );
+    let policy = PolicyIr::Pgn {
+        sources: vec!["foo".to_string()],
+        required_sources: Vec::new(),
+        dimension_tables: Vec::new(),
+        dimension_aliases,
+        dimension_queries,
+        sink: None,
+        sink_alias: None,
+        source_aliases: HashMap::new(),
+        constraint: "max(foo.id) > 0 AND u.user_id = 1".to_string(),
+        on_fail: Resolution::Remove,
+        description: None,
+    };
+    let sql = rewrite_with_catalog("SELECT foo.id FROM foo", &[policy], catalog);
+    assert!(
+        sql.contains("u.user_id = 1"),
+        "expected dimension subquery alias in rewritten SQL: {sql}"
+    );
+    assert!(
+        sql.contains("SELECT user_id FROM session_user") || sql.contains("session_user"),
+        "expected dimension subquery injected: {sql}"
+    );
+}
+
 #[test]
 fn catalog_rejects_unknown_dimension_table() {
     let mut catalog = TableCatalog::new();

{data_flow_control-0.1.2 → data_flow_control-0.1.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "data-flow-control"
-version = "0.1.2"
+version = "0.1.3"
 description = "Rust-backed Data Flow Control rewrite engine"
 license = "MIT"
 license-files = ["LICENSE"]