Skip to content

Commit c54be12

Browse files
kvlongekvlonge
committed
refactor(commands): call native helpers directly in freshness, janitor, check and diff
The freshness, janitor, check, diff and deferred-target code calls the native discovery, graph, selection, connection, target, progress, state and output helpers it uses as Rust functions instead of looking them up on the native module. SQL test chain resolution is called directly; SQL test planning stays a module lookup because tests observe it there. Native classes the commands construct (progress reporters, the event writer, the cost context and freshness errors) are still created through the module.
1 parent a417aca commit c54be12

10 files changed

Lines changed: 142 additions & 162 deletions

File tree

‎crates/sqlbuild-python/src/bindings/_helpers/check/check_command.rs‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,19 @@
11
//! `sqb check`: run the selected Python checks after their Python dependencies.
22
3+
use crate::bindings::_helpers::cli::progress::current_native_progress_projector;
4+
use crate::bindings::_helpers::cli::progress_text::write_execution_header;
35
use crate::bindings::_helpers::compile_driver::pipeline_entries::build_static_pipeline_result;
46
use crate::bindings::_helpers::compile_driver::python_calls::keywords;
57
use crate::bindings::_helpers::compile_driver::relation_targets::build_project_relation_targets;
68
use crate::bindings::_helpers::compile_driver::relation_targets::build_python_relation_targets;
79
use crate::bindings::_helpers::compile_driver::selection_phases::compile_static_command_context;
810
use crate::bindings::_helpers::dbt_interop::reference_resolver::resolve_external_sql_reference_resolver;
911
use crate::bindings::_helpers::graph::algorithms::transitive_closure_many;
12+
use crate::bindings::_helpers::graph::execution_gates::sql_ref_key;
1013
use crate::bindings::_helpers::graph::python_lifecycle::build_python_sql_run_lifecycle;
1114
use crate::bindings::_helpers::planner_entries::command_plans::build_relation_command_plan;
1215
use crate::bindings::_helpers::planner_phases::objects::object;
16+
use crate::bindings::_helpers::planner_selection::selectors::parse_project_selector;
1317
use crate::bindings::_helpers::planner_selection::unit_tests::reject_unit_test_selectors;
1418
use crate::bindings::_helpers::result_output::documents::write_execution_json_output;
1519
use pyo3::prelude::{
@@ -82,12 +86,12 @@ fn empty_tuple(py: Python<'_>) -> Bound<'_, PyAny> {
8286

8387
/// Refuse `+`, `~` and path selectors, splitting selectors as Python's `split` does.
8488
fn refuse_graph_selectors(py: Python<'_>, selectors: &[Bound<'_, PyAny>]) -> PyResult<()> {
85-
let parse: Bound<'_, PyAny> = registry::native(py, NativeAttr::ParseProjectSelector)?;
8689
let path_selector: Bound<'_, PyAny> = registry::class(py, RecordId::PathSelector)?;
8790
for selector in selectors {
8891
for token in selector.call_method0("split")?.try_iter()? {
8992
for part in token?.call_method1("split", (",",))?.try_iter()? {
90-
let parsed: Bound<'_, PyAny> = parse.call1((part?,))?;
93+
let parsed: Bound<'_, PyAny> =
94+
parse_project_selector(part?.downcast::<PyString>()?)?;
9195
if parsed.is_instance(&path_selector)?
9296
|| parsed.getattr("upstream")?.is_truthy()?
9397
|| parsed.getattr("downstream")?.is_truthy()?
@@ -197,10 +201,9 @@ fn compile_pipeline<'py>(
197201
let key_class: Bound<'py, PyAny> = registry::class(py, RecordId::CompiledObjectKey)?;
198202
let source_type: Bound<'py, PyAny> =
199203
imported(py, COMPILE_TYPES, "CompiledResourceType")?.getattr("SOURCE")?;
200-
let ref_key: Bound<'py, PyAny> = registry::native(py, NativeAttr::SqlRefKey)?;
201204
let mut keys: Vec<Bound<'py, PyAny>> = Vec::new();
202205
for reference in required_refs.try_iter()? {
203-
keys.push(ref_key.call1((reference?,))?);
206+
keys.push(sql_ref_key(py, &reference?)?);
204207
}
205208
for source_file in discovered.getattr("source_files")?.try_iter()? {
206209
for source in source_file?.getattr("source_entries")?.try_iter()? {
@@ -567,8 +570,7 @@ fn execute<'py>(
567570
let event_writer: Bound<'py, PyAny> =
568571
registry::native(py, NativeAttr::ExecutionEventWriter)?.call0()?;
569572
let executed: PyResult<Vec<Bound<'py, PyAny>>> = (|| {
570-
let projector: Bound<'py, PyAny> =
571-
registry::native(py, NativeAttr::CurrentNativeProgressProjector)?.call0()?;
573+
let projector: Bound<'py, PyAny> = current_native_progress_projector(py)?;
572574
if !projector.is_none() {
573575
let ordinals: Bound<'py, PyDict> = PyDict::new(py);
574576
let mut total: usize = 0;
@@ -657,7 +659,7 @@ fn run_check_command(py: Python<'_>, request: &Bound<'_, PyAny>) -> PyResult<i32
657659
header.set_item("target", py.None())?;
658660
header.set_item("concurrency", PyInt::new(py, 1))?;
659661
header.set_item("use_color", invocation.getattr("use_color")?)?;
660-
registry::native(py, NativeAttr::WriteExecutionHeader)?.call((), Some(&header))?;
662+
write_execution_header(py, Some(header))?;
661663
let prepared: Prepared<'_> = prepare(request, &invocation)?;
662664
let session: Bound<'_, PyAny> = call_with(
663665
&host::entry(py, UserCodeHost::BuildProviderSession)?,

‎crates/sqlbuild-python/src/bindings/_helpers/check/check_results.rs‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,10 @@ use sqlbuild_core::check_command::models::{CheckGrouping, CheckResultText, Check
1414
use sqlbuild_core::cli_text::models::CliStyle;
1515
use sqlbuild_core::json::models::JsonValue;
1616

17+
use crate::bindings::_helpers::cli::progress::current_native_progress_projector;
1718
use crate::bindings::_helpers::compile_driver::python_calls::KeywordArguments;
18-
use crate::bindings::_helpers::records::ids::NativeAttr;
19-
use crate::bindings::_helpers::records::registry;
2019
use crate::bindings::_helpers::result_output::integration_values::json_value;
20+
use crate::bindings::_helpers::result_output::terminal_index::current_terminal_event_index;
2121

2222
const WARN_SEVERITY: &str = "warn";
2323

@@ -124,8 +124,7 @@ pub(crate) fn write_results(
124124
let check: Bound<'_, PyAny> = check?;
125125
checks.insert(check.getattr("name")?.extract()?, check);
126126
}
127-
let projector: Bound<'_, PyAny> =
128-
registry::native(py, NativeAttr::CurrentNativeProgressProjector)?.call0()?;
127+
let projector: Bound<'_, PyAny> = current_native_progress_projector(py)?;
129128
let mut groups: Vec<(String, Vec<CheckResultText>)> = Vec::new();
130129
for result in results {
131130
let text: CheckResultText = check_result_text(result)?;
@@ -169,8 +168,7 @@ pub(crate) fn results_json(results: &[Bound<'_, PyAny>]) -> PyResult<String> {
169168
return check_json(&[]).map_err(|error| PyValueError::new_err(format!("{error:?}")));
170169
};
171170
let py: Python<'_> = first.py();
172-
let index: Bound<'_, PyAny> =
173-
registry::native(py, NativeAttr::CurrentTerminalEventIndex)?.call0()?;
171+
let index: Bound<'_, PyAny> = current_terminal_event_index(py)?;
174172
let mut texts: Vec<CheckResultText> = Vec::new();
175173
for result in results {
176174
let name: Bound<'_, PyAny> = result.getattr("node_name")?;

‎crates/sqlbuild-python/src/bindings/_helpers/compile_driver/deferral.rs‎

Lines changed: 6 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,9 @@ use pyo3::{pyfunction, wrap_pyfunction};
99
use sqlbuild_model_config::templates::main::deferred_target_segments::deferred_target_segments;
1010
use sqlbuild_model_config::templates::models::DeferredSegment;
1111

12+
use crate::bindings::_helpers::authored_values::project_vars::render_project_var_text;
1213
use crate::bindings::_helpers::compile_driver::python_calls::{KeywordArguments, call_with};
13-
use crate::bindings::_helpers::records::ids::{ErrorId, NativeAttr, RecordId};
14+
use crate::bindings::_helpers::records::ids::{ErrorId, RecordId};
1415
use crate::bindings::_helpers::records::registry;
1516

1617
const PRESERVE_TARGET_VALUE: &str = "preserve";
@@ -150,24 +151,15 @@ fn target_field<'py>(
150151
.iter()
151152
.map(|key| (*key).to_owned())
152153
.collect();
153-
let render: Bound<'py, PyAny> = registry::native(py, NativeAttr::RenderProjectVarText)?;
154154
let mut resolved: String = String::new();
155155
for segment in deferred_target_segments(&template, &logical, &context_keys, &variables) {
156156
match segment {
157157
DeferredSegment::Text(text) => resolved.push_str(&text),
158158
DeferredSegment::Variable(name) => {
159-
let text: String = call_with(
160-
&render,
161-
&[
162-
("value", effective_vars.get_item(name.as_str())?),
163-
(
164-
"label",
165-
PyString::new(py, &format!("deferred target variable '${{{name}}}'"))
166-
.into_any(),
167-
),
168-
],
169-
)?
170-
.extract()?;
159+
let text: String = render_project_var_text(
160+
&effective_vars.get_item(name.as_str())?,
161+
&format!("deferred target variable '${{{name}}}'"),
162+
)?;
171163
resolved.push_str(&text);
172164
}
173165
}

‎crates/sqlbuild-python/src/bindings/_helpers/compile_driver/diff.rs‎

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,12 @@ use sqlbuild_analysis::diff_selection::main::resolve_diff_selection::resolve_dif
1313
use sqlbuild_analysis::diff_selection::models::DiffGraph;
1414

1515
use crate::bindings::_helpers::analysis::model_reads;
16+
use crate::bindings::_helpers::compile_driver::project::build_compiled_project;
1617
use crate::bindings::_helpers::compile_driver::python_calls::{
17-
KeywordArguments, call_with, imported, keywords,
18+
KeywordArguments, imported, keywords,
1819
};
1920
use crate::bindings::_helpers::graph::views::project_lineage_views;
20-
use crate::bindings::_helpers::records::ids::{ErrorId, NativeAttr};
21+
use crate::bindings::_helpers::records::ids::ErrorId;
2122
use crate::bindings::_helpers::records::registry;
2223

2324
const COMPILE_TYPES_MODULE: &str = "sqlbuild.compiler.compile.types";
@@ -203,25 +204,27 @@ fn run_diff_pipeline<'py>(
203204
&["cli_vars", "external_sql_reference_resolver"],
204205
),
205206
)?;
206-
let build: Bound<'py, PyAny> = registry::native(py, NativeAttr::BuildCompiledProject)?;
207207
let compile = |target: &str| -> PyResult<Bound<'py, PyAny>> {
208-
call_with(
209-
&build,
210-
&[
211-
("discovered_inputs", arguments.value("discovered_inputs")?),
212-
("adapter", arguments.value("adapter")?),
213-
("selected_target", arguments.value(target)?),
214-
(
215-
"resolved_connection",
216-
arguments.value("resolved_connection")?,
217-
),
218-
("no_sql_validation", arguments.value("no_sql_validation")?),
219-
("cli_vars", arguments.value("cli_vars")?),
220-
(
221-
"external_sql_reference_resolver",
222-
arguments.value("external_sql_reference_resolver")?,
223-
),
224-
],
208+
build_compiled_project(
209+
py,
210+
Some(keywords(
211+
py,
212+
&[
213+
("discovered_inputs", arguments.value("discovered_inputs")?),
214+
("adapter", arguments.value("adapter")?),
215+
("selected_target", arguments.value(target)?),
216+
(
217+
"resolved_connection",
218+
arguments.value("resolved_connection")?,
219+
),
220+
("no_sql_validation", arguments.value("no_sql_validation")?),
221+
("cli_vars", arguments.value("cli_vars")?),
222+
(
223+
"external_sql_reference_resolver",
224+
arguments.value("external_sql_reference_resolver")?,
225+
),
226+
],
227+
)?),
225228
)
226229
};
227230
let left: Bound<'py, PyAny> = compile("from_target")?;

‎crates/sqlbuild-python/src/bindings/_helpers/freshness/freshness_command.rs‎

Lines changed: 40 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -15,15 +15,24 @@ use sqlbuild_core::freshness_text::main::render_freshness_text::render_freshness
1515
use sqlbuild_core::freshness_text::models::FreshnessSourceResult;
1616

1717
use crate::bindings::_helpers::command_runtime::adapters::resolve_cli_adapter;
18+
use crate::bindings::_helpers::command_runtime::connection::resolve_project_connection_config;
19+
use crate::bindings::_helpers::compile_driver::discovery::discover_project_inputs;
20+
use crate::bindings::_helpers::compile_driver::project::build_project_graph;
21+
use crate::bindings::_helpers::contracts::spec_resolution::resolve_effective_adapter_name;
1822
use crate::bindings::_helpers::freshness::freshness_observe::{PreviousState, observe_freshness};
1923
use crate::bindings::_helpers::freshness::observation::Observer;
2024
use crate::bindings::_helpers::freshness::state_io::{StateRead, read_latest};
25+
use crate::bindings::_helpers::planner_selection::selectors::{
26+
expand_project_upstream_keys, select_project_graph,
27+
};
28+
use crate::bindings::_helpers::planner_selection::unit_tests::reject_unit_test_selectors;
2129
use crate::bindings::_helpers::presentation_text::terminal::stdout_supports_color;
30+
use crate::bindings::_helpers::result_output::documents::write_execution_json_output;
2231

2332
const COMPILE_TYPES_MODULE: &str = "sqlbuild.compiler.compile.types";
2433
use crate::bindings::_helpers::native_host::adapter::AdapterHost;
2534
use crate::bindings::_helpers::native_host::host;
26-
use crate::bindings::_helpers::records::ids::{ErrorId, NativeAttr};
35+
use crate::bindings::_helpers::records::ids::ErrorId;
2736
use crate::bindings::_helpers::records::registry;
2837
use sqlbuild_core::constants::SOURCE_FRESHNESS_TABLE_NAME as STATE_TABLE_NAME;
2938
const STALE_REQUIRES_STATE_MESSAGE: &str = "freshness --fail-on-stale requires --state";
@@ -72,17 +81,11 @@ fn source_names_for_keys(
7281
let py: Python<'_> = graph.py();
7382
let source_type: Bound<'_, PyAny> =
7483
imported(py, COMPILE_TYPES_MODULE, "CompiledResourceType")?.getattr("SOURCE")?;
75-
let expand: Bound<'_, PyAny> = registry::native(py, NativeAttr::ExpandProjectUpstreamKeys)?;
7684
let upstream_deps: Bound<'_, PyAny> = graph.getattr("upstream_deps")?;
7785
for key in keys.try_iter()? {
7886
let key: Bound<'_, PyAny> = key?;
79-
let upstream: Bound<'_, PyAny> = expand.call(
80-
(),
81-
Some(&keywords(
82-
py,
83-
&[("key", &key), ("upstream_deps", &upstream_deps)],
84-
)?),
85-
)?;
87+
let upstream: Bound<'_, PyAny> =
88+
expand_project_upstream_keys(key.clone(), &upstream_deps)?.into_any();
8689
for candidate in std::iter::once(Ok(key.clone())).chain(upstream.try_iter()?) {
8790
let candidate: Bound<'_, PyAny> = candidate?;
8891
if candidate.getattr("resource_type")?.eq(&source_type)? {
@@ -99,13 +102,7 @@ fn graph_keys<'py>(
99102
) -> PyResult<Bound<'py, PyAny>> {
100103
let py: Python<'py> = graph.py();
101104
let empty: Bound<'py, PyAny> = PyTuple::empty(py).into_any();
102-
registry::native(py, NativeAttr::SelectProjectGraph)?.call(
103-
(),
104-
Some(&keywords(
105-
py,
106-
&[("graph", graph), ("select", selectors), ("exclude", &empty)],
107-
)?),
108-
)
105+
Ok(select_project_graph(graph, selectors, &empty)?.into_any())
109106
}
110107

111108
/// Selected source names, sorted: selected or all keys and their upstream sources, minus excluded.
@@ -206,19 +203,18 @@ fn observe(
206203
names: &[String],
207204
) -> PyResult<Vec<FreshnessSourceResult>> {
208205
let py: Python<'_> = adapter.py();
209-
let connection_config: Bound<'_, PyAny> =
210-
registry::native(py, NativeAttr::ResolveProjectConnectionConfig)?.call(
211-
(),
212-
Some(&keywords(
213-
py,
214-
&[
215-
("discovered_inputs", discovered),
216-
("project_dir", project_dir),
217-
("selected_target", &request.getattr("selected_target")?),
218-
("cli_vars", &request.getattr("cli_vars")?),
219-
],
220-
)?),
221-
)?;
206+
let connection_config: Bound<'_, PyAny> = resolve_project_connection_config(
207+
py,
208+
Some(keywords(
209+
py,
210+
&[
211+
("discovered_inputs", discovered),
212+
("project_dir", project_dir),
213+
("selected_target", &request.getattr("selected_target")?),
214+
("cli_vars", &request.getattr("cli_vars")?),
215+
],
216+
)?),
217+
)?;
222218
let connection: Bound<'_, PyAny> = adapter.call_method1("connect", (connection_config,))?;
223219
let observed = (|| {
224220
let project: Bound<'_, PyAny> = graph.getattr("project")?;
@@ -258,11 +254,12 @@ fn write_json_output(
258254
payload: &str,
259255
(json_output, json_output_path): (bool, &Bound<'_, PyAny>),
260256
) -> PyResult<()> {
261-
let options: Bound<'_, PyDict> = keywords(py, &[("json_output_path", json_output_path)])?;
262-
options.set_item("payload", format!("{payload}\n"))?;
263-
options.set_item("json_output", json_output)?;
264-
registry::native(py, NativeAttr::WriteExecutionJsonOutput)?.call((), Some(&options))?;
265-
Ok(())
257+
write_execution_json_output(
258+
py,
259+
PyString::new(py, &format!("{payload}\n")).into_any(),
260+
json_output,
261+
json_output_path.clone(),
262+
)
266263
}
267264

268265
/// Run one `sqb freshness` request.
@@ -281,36 +278,19 @@ fn run_freshness_command(py: Python<'_>, request: &Bound<'_, PyAny>) -> PyResult
281278
directory if !directory.is_none() => directory,
282279
_ => py.import("pathlib")?.getattr("Path")?.call_method0("cwd")?,
283280
};
284-
let discovered: Bound<'_, PyAny> = registry::native(py, NativeAttr::DiscoverProjectInputs)?
285-
.call((), Some(&keywords(py, &[("project_dir", &project_dir)])?))?;
281+
let discovered: Bound<'_, PyAny> =
282+
discover_project_inputs(py, Some(keywords(py, &[("project_dir", &project_dir)])?))?;
286283
let select: Bound<'_, PyAny> = request.getattr("select")?;
287284
let exclude: Bound<'_, PyAny> = request.getattr("exclude")?;
288-
registry::native(py, NativeAttr::RejectUnitTestSelectors)?.call(
289-
(),
290-
Some(&keywords(
291-
py,
292-
&[
293-
("select", &select),
294-
("exclude", &exclude),
295-
("discovered_inputs", &discovered),
296-
],
297-
)?),
285+
reject_unit_test_selectors(&select, &exclude, &discovered)?;
286+
let adapter_name: Bound<'_, PyAny> = resolve_effective_adapter_name(
287+
&discovered.getattr("project_config")?,
288+
&discovered.getattr("local_config")?,
298289
)?;
299-
let adapter_name: Bound<'_, PyAny> =
300-
registry::native(py, NativeAttr::ResolveEffectiveAdapterName)?.call(
301-
(),
302-
Some(&keywords(
303-
py,
304-
&[
305-
("project_config", &discovered.getattr("project_config")?),
306-
("local_config", &discovered.getattr("local_config")?),
307-
],
308-
)?),
309-
)?;
310290
let adapter: Bound<'_, PyAny> = resolve_cli_adapter(&adapter_name, &project_dir)?;
311-
let graph: Bound<'_, PyAny> = registry::native(py, NativeAttr::BuildProjectGraph)?.call(
312-
(),
313-
Some(&keywords(
291+
let graph: Bound<'_, PyAny> = build_project_graph(
292+
py,
293+
Some(keywords(
314294
py,
315295
&[
316296
("discovered_inputs", &discovered),

0 commit comments

Comments
 (0)