Skip to content

Commit bdecb1f

Browse files
psteinroepsteinroe
committed
refactor: run async in parallel
1 parent afe3947 commit bdecb1f

File tree

7 files changed

+101
-79
lines changed

7 files changed

+101
-79
lines changed

‎crates/pg_completions/src/relevance.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ impl CompletionRelevance<'_> {
7070
Some(ct) => ct,
7171
};
7272

73-
let has_mentioned_tables = ctx.mentioned_relations.len() > 0;
73+
let has_mentioned_tables = !ctx.mentioned_relations.is_empty();
7474

7575
self.score += match self.data {
7676
CompletionRelevanceData::Table(_) => match clause_type {

‎crates/pg_completions/src/test_helper.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ pub(crate) async fn get_test_deps(
5151
.set_language(tree_sitter_sql::language())
5252
.expect("Error loading sql language");
5353

54-
let tree = parser.parse(&input.to_string(), None).unwrap();
54+
let tree = parser.parse(input.to_string(), None).unwrap();
5555

5656
(tree, schema_cache)
5757
}

‎crates/pg_treesitter_queries/src/lib.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ impl<'a> TreeSitterQueriesExecutor<'a> {
2121

2222
#[allow(private_bounds)]
2323
pub fn add_query_results<Q: Query<'a>>(&mut self) {
24-
let mut results = Q::execute(self.root_node, &self.stmt);
24+
let mut results = Q::execute(self.root_node, self.stmt);
2525
self.results.append(&mut results);
2626
}
2727

@@ -104,9 +104,9 @@ where
104104
let mut parser = tree_sitter::Parser::new();
105105
parser.set_language(tree_sitter_sql::language()).unwrap();
106106

107-
let tree = parser.parse(&sql, None).unwrap();
107+
let tree = parser.parse(sql, None).unwrap();
108108

109-
let mut executor = TreeSitterQueriesExecutor::new(tree.root_node(), &sql);
109+
let mut executor = TreeSitterQueriesExecutor::new(tree.root_node(), sql);
110110

111111
executor.add_query_results::<RelationMatch>();
112112

@@ -152,7 +152,7 @@ on sq1.id = pt.id;
152152
let mut parser = tree_sitter::Parser::new();
153153
parser.set_language(tree_sitter_sql::language()).unwrap();
154154

155-
let tree = parser.parse(&sql, None).unwrap();
155+
let tree = parser.parse(sql, None).unwrap();
156156

157157
// trust me bro
158158
let range = {
@@ -172,7 +172,7 @@ on sq1.id = pt.id;
172172
cursor.node().range()
173173
};
174174

175-
let mut executor = TreeSitterQueriesExecutor::new(tree.root_node(), &sql);
175+
let mut executor = TreeSitterQueriesExecutor::new(tree.root_node(), sql);
176176

177177
executor.add_query_results::<RelationMatch>();
178178

‎crates/pg_treesitter_queries/src/queries/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ pub enum QueryResult<'a> {
77
Relation(RelationMatch<'a>),
88
}
99

10-
impl<'a> QueryResult<'a> {
10+
impl QueryResult<'_> {
1111
pub fn within_range(&self, range: &tree_sitter::Range) -> bool {
1212
match self {
1313
Self::Relation(rm) => {

‎crates/pg_treesitter_queries/src/queries/relations.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use crate::{Query, QueryResult};
55
use super::QueryTryFrom;
66

77
static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
8-
static QUERY_STR: &'static str = r#"
8+
static QUERY_STR: &str = r#"
99
(relation
1010
(object_reference
1111
.
@@ -15,7 +15,7 @@ static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
1515
)+
1616
)
1717
"#;
18-
tree_sitter::Query::new(tree_sitter_sql::language(), &QUERY_STR).expect("Invalid TS Query")
18+
tree_sitter::Query::new(tree_sitter_sql::language(), QUERY_STR).expect("Invalid TS Query")
1919
});
2020

2121
#[derive(Debug)]
@@ -24,7 +24,7 @@ pub struct RelationMatch<'a> {
2424
pub(crate) table: tree_sitter::Node<'a>,
2525
}
2626

27-
impl<'a> RelationMatch<'a> {
27+
impl RelationMatch<'_> {
2828
pub fn get_schema(&self, sql: &str) -> Option<String> {
2929
let str = self
3030
.schema
@@ -48,7 +48,7 @@ impl<'a> TryFrom<&'a QueryResult<'a>> for &'a RelationMatch<'a> {
4848

4949
fn try_from(q: &'a QueryResult<'a>) -> Result<Self, Self::Error> {
5050
match q {
51-
QueryResult::Relation(r) => Ok(&r),
51+
QueryResult::Relation(r) => Ok(r),
5252

5353
#[allow(unreachable_patterns)]
5454
_ => Err("Invalid QueryResult type".into()),

‎crates/pg_workspace_new/src/workspace/server.rs‎

Lines changed: 87 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ use analyser::AnalyserVisitorBuilder;
44
use change::StatementChange;
55
use dashmap::{DashMap, DashSet};
66
use document::{Document, Statement};
7+
use futures::{stream, StreamExt};
78
use pg_analyse::{AnalyserOptions, AnalysisFilter};
89
use pg_analyser::{Analyser, AnalyserConfig, AnalyserContext};
910
use pg_diagnostics::{serde::Diagnostic as SDiagnostic, Diagnostic, DiagnosticExt, Severity};
@@ -13,6 +14,7 @@ use pg_schema_cache::SchemaCache;
1314
use pg_typecheck::TypecheckParams;
1415
use sqlx::PgPool;
1516
use std::sync::LazyLock;
17+
use text_size::TextRange;
1618
use tokio::runtime::Runtime;
1719
use tracing::info;
1820
use tree_sitter::TreeSitterStore;
@@ -338,72 +340,95 @@ impl Workspace for WorkspaceServer {
338340
filter,
339341
});
340342

341-
let con = self.connection.read().unwrap().get_pool();
342-
343-
let diagnostics: Vec<SDiagnostic> = doc
344-
.iter_statements_with_text_and_range()
345-
.flat_map(|(stmt, r, text)| {
346-
let mut stmt_diagnostics = self.pg_query.get_diagnostics(&stmt);
347-
348-
let ast = self.pg_query.get_ast(&stmt);
349-
let tree = self.tree_sitter.get_parse_tree(&stmt);
350-
351-
if let Some(ast) = ast {
352-
stmt_diagnostics.extend(
353-
analyser
354-
.run(AnalyserContext { root: &ast })
355-
.into_iter()
356-
.map(SDiagnostic::new)
357-
.collect::<Vec<_>>(),
358-
);
359-
360-
if let Some(con) = con.clone() {
361-
let text = text.to_string();
362-
let ast_clone = ast.clone();
363-
let tree_clone = tree.clone();
364-
if let Ok(typecheck_result) = run_async(async move {
365-
pg_typecheck::check_sql(TypecheckParams {
366-
conn: &con,
367-
sql: &text,
368-
ast: &ast_clone,
369-
tree: tree_clone.as_deref(),
370-
})
371-
.await
372-
}) {
373-
if let Some(typecheck_result) = typecheck_result.map(SDiagnostic::new) {
374-
stmt_diagnostics.push(typecheck_result);
343+
let mut diagnostics: Vec<SDiagnostic> = vec![];
344+
345+
// run diagnostics for each statement in parallel if its mostly i/o work
346+
if let Some(pool) = self.connection.read().unwrap().get_pool() {
347+
let typecheck_params: Vec<_> = doc
348+
.iter_statements_with_text_and_range()
349+
.map(|(stmt, range, text)| {
350+
let ast = self.pg_query.get_ast(&stmt);
351+
let tree = self.tree_sitter.get_parse_tree(&stmt);
352+
(text.to_string(), ast, tree, *range)
353+
})
354+
.collect();
355+
356+
let pool_clone = pool.clone();
357+
let path_clone = params.path.clone();
358+
let async_results = run_async(async move {
359+
stream::iter(typecheck_params)
360+
.map(|(text, ast, tree, range)| {
361+
let pool = pool_clone.clone();
362+
let path = path_clone.clone();
363+
async move {
364+
if let Some(ast) = ast {
365+
pg_typecheck::check_sql(TypecheckParams {
366+
conn: &pool,
367+
sql: &text,
368+
ast: &ast,
369+
tree: tree.as_deref(),
370+
})
371+
.await
372+
.map(|d| {
373+
let r = d.location().span.map(|span| span + range.start());
374+
375+
d.with_file_path(path.as_path().display().to_string())
376+
.with_file_span(r.unwrap_or(range))
377+
})
378+
} else {
379+
None
375380
}
376381
}
377-
}
378-
}
379-
380-
stmt_diagnostics
381-
.into_iter()
382-
.map(|d| {
383-
// We do now check if the severity of the diagnostics should be changed.
384-
// The configuration allows to change the severity of the diagnostics emitted by rules.
385-
let severity = d
386-
.category()
387-
.filter(|category| category.name().starts_with("lint/"))
388-
.map_or_else(
389-
|| d.severity(),
390-
|category| {
391-
settings
392-
.as_ref()
393-
.get_severity_from_rule_code(category)
394-
.unwrap_or(Severity::Warning)
395-
},
396-
);
397-
398-
SDiagnostic::new(
399-
d.with_file_path(params.path.as_path().display().to_string())
400-
.with_file_span(r)
401-
.with_severity(severity),
402-
)
403382
})
383+
.buffer_unordered(10)
404384
.collect::<Vec<_>>()
405-
})
406-
.collect();
385+
.await
386+
})?;
387+
388+
for result in async_results.into_iter().flatten() {
389+
diagnostics.push(SDiagnostic::new(result));
390+
}
391+
}
392+
393+
diagnostics.extend(doc.iter_statements_with_range().flat_map(|(stmt, r)| {
394+
let mut stmt_diagnostics = self.pg_query.get_diagnostics(&stmt);
395+
396+
let ast = self.pg_query.get_ast(&stmt);
397+
398+
if let Some(ast) = ast {
399+
stmt_diagnostics.extend(
400+
analyser
401+
.run(AnalyserContext { root: &ast })
402+
.into_iter()
403+
.map(SDiagnostic::new)
404+
.collect::<Vec<_>>(),
405+
);
406+
}
407+
408+
stmt_diagnostics
409+
.into_iter()
410+
.map(|d| {
411+
let severity = d
412+
.category()
413+
.filter(|category| category.name().starts_with("lint/"))
414+
.map_or_else(
415+
|| d.severity(),
416+
|category| {
417+
settings
418+
.as_ref()
419+
.get_severity_from_rule_code(category)
420+
.unwrap_or(Severity::Warning)
421+
},
422+
);
423+
424+
SDiagnostic::new(
425+
d.with_file_path(params.path.as_path().display().to_string())
426+
.with_file_span(r)
427+
.with_severity(severity),
428+
)
429+
})
430+
.collect::<Vec<_>>()
431+
}));
407432

408433
let errors = diagnostics
409434
.iter()

‎crates/pg_workspace_new/src/workspace/server/typecheck.rs‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,12 @@
1-
use std::sync::Arc;
2-
31
// pub conn: &'a PgPool,
42
// pub sql: &'a str,
53
// pub ast: &'a pg_query_ext::NodeEnum,
64
// pub tree: Option<&'a tree_sitter::Tree>,
75
use dashmap::DashMap;
86
use pg_diagnostics::serde::Diagnostic as SDiagnostic;
9-
use pg_query_ext::diagnostics::*;
10-
use pg_typecheck::{check_sql, TypecheckDiagnostic};
7+
use pg_typecheck::TypecheckDiagnostic;
118

12-
use super::{change::ModifiedStatement, document::Statement};
9+
use super::document::Statement;
1310

1411
pub struct TypecheckStore {
1512
diagnostics: DashMap<Statement, TypecheckDiagnostic>,

0 commit comments

Comments
 (0)