From 9abf45921cfe9226d55b7ba26ff34350dfa73569 Mon Sep 17 00:00:00 2001 From: Julian Date: Wed, 28 May 2025 08:48:39 +0200 Subject: [PATCH 01/23] so far --- crates/pgt_schema_cache/src/columns.rs | 1 - crates/pgt_schema_cache/src/policies.rs | 1 - crates/pgt_test_utils/src/test_database.rs | 89 +++++++++++++++++++++- 3 files changed, 85 insertions(+), 6 deletions(-) diff --git a/crates/pgt_schema_cache/src/columns.rs b/crates/pgt_schema_cache/src/columns.rs index 60d422fd..943cf9ca 100644 --- a/crates/pgt_schema_cache/src/columns.rs +++ b/crates/pgt_schema_cache/src/columns.rs @@ -83,7 +83,6 @@ impl SchemaCacheItem for Column { #[cfg(test)] mod tests { use pgt_test_utils::test_database::get_new_test_db; - use sqlx::Executor; use crate::{SchemaCache, columns::ColumnClassKind}; diff --git a/crates/pgt_schema_cache/src/policies.rs b/crates/pgt_schema_cache/src/policies.rs index 85cd7821..770d1a7b 100644 --- a/crates/pgt_schema_cache/src/policies.rs +++ b/crates/pgt_schema_cache/src/policies.rs @@ -81,7 +81,6 @@ impl SchemaCacheItem for Policy { #[cfg(test)] mod tests { use pgt_test_utils::test_database::get_new_test_db; - use sqlx::Executor; use crate::{SchemaCache, policies::PolicyCommand}; diff --git a/crates/pgt_test_utils/src/test_database.rs b/crates/pgt_test_utils/src/test_database.rs index 67415c4a..a8ed9a4d 100644 --- a/crates/pgt_test_utils/src/test_database.rs +++ b/crates/pgt_test_utils/src/test_database.rs @@ -1,9 +1,85 @@ -use sqlx::{Executor, PgPool, postgres::PgConnectOptions}; +use std::ops::Deref; + +use sqlx::{ + Executor, PgPool, + postgres::{PgConnectOptions, PgQueryResult}, +}; use uuid::Uuid; +#[derive(Debug)] +pub struct TestDb { + pool: PgPool, + roles: Vec, +} + +#[derive(Debug)] +pub struct RoleWithArgs { + role: String, + args: Vec, +} + +impl TestDb { + pub async fn execute(&self, sql: &str) -> Result { + if sql.to_ascii_lowercase().contains("create role") { + panic!("Please setup roles via the `setup_roles` method.") + } + self.pool.execute(sql).await + } + + pub async fn setup_roles( + &mut self, + roles: Vec, + ) -> Result { + self.roles = roles.iter().map(|r| &r.role).cloned().collect(); + + let role_statements: Vec = roles + .into_iter() + .map(|r| { + format!( + r#" + if not exists ( + select from pg_catalog.pg_roles + where rolname = '{0}' + ) then + create role {0} {1}; + end if; + "#, + r.role, + r.args.join(" ") + ) + }) + .collect(); + + let query = format!( + r#" + do $$ + begin + {} + end $$; + "#, + role_statements.join("\n") + ); + + println!("{}", query); + + self.execute(&query).await + } + + pub fn get_roles(&self) -> &[String] { + &self.roles + } +} + +impl Deref for TestDb { + type Target = PgPool; + fn deref(&self) -> &Self::Target { + &self.pool + } +} + // TODO: Work with proper config objects instead of a connection_string. // With the current implementation, we can't parse the password from the connection string. -pub async fn get_new_test_db() -> PgPool { +pub async fn get_new_test_db() -> TestDb { dotenv::dotenv().expect("Unable to load .env file for tests"); let connection_string = std::env::var("DATABASE_URL").expect("DATABASE_URL not set"); @@ -36,7 +112,12 @@ pub async fn get_new_test_db() -> PgPool { .await .expect("Failed to create test database."); - sqlx::PgPool::connect_with(options_without_db_name.database(&database_name)) + let pool = sqlx::PgPool::connect_with(options_without_db_name.database(&database_name)) .await - .expect("Could not connect to test database") + .expect("Could not connect to test database"); + + TestDb { + pool, + roles: vec![], + } } From 8d4836e60389cc929f9a9266ea54e176534fd272 Mon Sep 17 00:00:00 2001 From: Julian Date: Wed, 28 May 2025 09:00:18 +0200 Subject: [PATCH 02/23] yayyy --- crates/pgt_schema_cache/src/policies.rs | 39 +++++++---------- crates/pgt_schema_cache/src/roles.rs | 51 ++++++++++------------ crates/pgt_test_utils/src/test_database.rs | 8 ++-- 3 files changed, 41 insertions(+), 57 deletions(-) diff --git a/crates/pgt_schema_cache/src/policies.rs b/crates/pgt_schema_cache/src/policies.rs index 770d1a7b..808a2b7e 100644 --- a/crates/pgt_schema_cache/src/policies.rs +++ b/crates/pgt_schema_cache/src/policies.rs @@ -80,26 +80,29 @@ impl SchemaCacheItem for Policy { #[cfg(test)] mod tests { - use pgt_test_utils::test_database::get_new_test_db; + use pgt_test_utils::test_database::{RoleWithArgs, get_new_test_db}; use crate::{SchemaCache, policies::PolicyCommand}; #[tokio::test] async fn loads_policies() { - let test_db = get_new_test_db().await; - - let setup = r#" - do $$ - begin - if not exists ( - select from pg_catalog.pg_roles - where rolname = 'admin' - ) then - create role admin; - end if; - end $$; + let mut test_db = get_new_test_db().await; + test_db + .setup_roles(vec![ + RoleWithArgs { + role: "admin".into(), + args: vec![], + }, + RoleWithArgs { + role: "owner".into(), + args: vec![], + }, + ]) + .await + .expect("Unable to setup admin roles"); + let setup = r#" create table public.users ( id serial primary key, name varchar(255) not null @@ -130,16 +133,6 @@ mod tests { to admin with check (true); - do $$ - begin - if not exists ( - select from pg_catalog.pg_roles - where rolname = 'owner' - ) then - create role owner; - end if; - end $$; - create schema real_estate; create table real_estate.properties ( diff --git a/crates/pgt_schema_cache/src/roles.rs b/crates/pgt_schema_cache/src/roles.rs index c212b791..eef02c96 100644 --- a/crates/pgt_schema_cache/src/roles.rs +++ b/crates/pgt_schema_cache/src/roles.rs @@ -22,41 +22,34 @@ impl SchemaCacheItem for Role { #[cfg(test)] mod tests { use crate::SchemaCache; - use pgt_test_utils::test_database::get_new_test_db; - use sqlx::Executor; + use pgt_test_utils::test_database::{RoleWithArgs, get_new_test_db}; #[tokio::test] async fn loads_roles() { - let test_db = get_new_test_db().await; - - let setup = r#" - do $$ - begin - if not exists ( - select from pg_catalog.pg_roles - where rolname = 'test_super' - ) then - create role test_super superuser createdb login bypassrls; - end if; - if not exists ( - select from pg_catalog.pg_roles - where rolname = 'test_nologin' - ) then - create role test_nologin; - end if; - if not exists ( - select from pg_catalog.pg_roles - where rolname = 'test_login' - ) then - create role test_login login; - end if; - end $$; - "#; + let mut test_db = get_new_test_db().await; test_db - .execute(setup) + .setup_roles(vec![ + RoleWithArgs { + role: "test_super".into(), + args: vec![ + "superuser".into(), + "createdb".into(), + "login".into(), + "bypassrls".into(), + ], + }, + RoleWithArgs { + role: "test_nologin".into(), + args: vec![], + }, + RoleWithArgs { + role: "test_login".into(), + args: vec!["login".into()], + }, + ]) .await - .expect("Failed to setup test database"); + .expect("Unable to set up roles."); let cache = SchemaCache::load(&test_db) .await diff --git a/crates/pgt_test_utils/src/test_database.rs b/crates/pgt_test_utils/src/test_database.rs index a8ed9a4d..b52b4167 100644 --- a/crates/pgt_test_utils/src/test_database.rs +++ b/crates/pgt_test_utils/src/test_database.rs @@ -14,8 +14,8 @@ pub struct TestDb { #[derive(Debug)] pub struct RoleWithArgs { - role: String, - args: Vec, + pub role: String, + pub args: Vec, } impl TestDb { @@ -60,9 +60,7 @@ impl TestDb { role_statements.join("\n") ); - println!("{}", query); - - self.execute(&query).await + self.pool.execute(query.as_str()).await } pub fn get_roles(&self) -> &[String] { From 363c2ea2831ae8a16d63fe61881cf22ecee0139a Mon Sep 17 00:00:00 2001 From: Julian Date: Wed, 28 May 2025 09:11:33 +0200 Subject: [PATCH 03/23] setup roles --- crates/pgt_test_utils/src/test_database.rs | 35 ++++++++++++++++------ 1 file changed, 26 insertions(+), 9 deletions(-) diff --git a/crates/pgt_test_utils/src/test_database.rs b/crates/pgt_test_utils/src/test_database.rs index b52b4167..3101d6ae 100644 --- a/crates/pgt_test_utils/src/test_database.rs +++ b/crates/pgt_test_utils/src/test_database.rs @@ -1,4 +1,8 @@ -use std::ops::Deref; +use std::{ + collections::HashSet, + ops::Deref, + sync::{LazyLock, Mutex}, +}; use sqlx::{ Executor, PgPool, @@ -6,10 +10,11 @@ use sqlx::{ }; use uuid::Uuid; +static DB_ROLES: LazyLock>> = LazyLock::new(|| Mutex::new(HashSet::new())); + #[derive(Debug)] pub struct TestDb { pool: PgPool, - roles: Vec, } #[derive(Debug)] @@ -30,7 +35,13 @@ impl TestDb { &mut self, roles: Vec, ) -> Result { - self.roles = roles.iter().map(|r| &r.role).cloned().collect(); + { + let roles: Vec = roles.iter().map(|r| &r.role).cloned().collect(); + let mut set = DB_ROLES.lock().unwrap(); + for role in roles { + set.insert(role); + } + } let role_statements: Vec = roles .into_iter() @@ -63,8 +74,17 @@ impl TestDb { self.pool.execute(query.as_str()).await } - pub fn get_roles(&self) -> &[String] { - &self.roles + pub fn get_roles(&self) -> Vec { + let mut roles = vec![]; + + { + let set = DB_ROLES.lock().unwrap(); + for role in set.iter() { + roles.push(role.clone()); + } + } + + roles } } @@ -114,8 +134,5 @@ pub async fn get_new_test_db() -> TestDb { .await .expect("Could not connect to test database"); - TestDb { - pool, - roles: vec![], - } + TestDb { pool } } From 0e1f0a806816ba512f16a9d751430292cc92c800 Mon Sep 17 00:00:00 2001 From: Julian Date: Wed, 28 May 2025 09:56:29 +0200 Subject: [PATCH 04/23] use distinct method --- crates/pgt_test_utils/src/test_database.rs | 39 +++++++++++++++++----- 1 file changed, 31 insertions(+), 8 deletions(-) diff --git a/crates/pgt_test_utils/src/test_database.rs b/crates/pgt_test_utils/src/test_database.rs index 3101d6ae..8600af83 100644 --- a/crates/pgt_test_utils/src/test_database.rs +++ b/crates/pgt_test_utils/src/test_database.rs @@ -12,6 +12,13 @@ use uuid::Uuid; static DB_ROLES: LazyLock>> = LazyLock::new(|| Mutex::new(HashSet::new())); +fn add_roles(roles: Vec) { + let mut set = DB_ROLES.lock().unwrap(); + for role in roles { + set.insert(role); + } +} + #[derive(Debug)] pub struct TestDb { pool: PgPool, @@ -35,13 +42,8 @@ impl TestDb { &mut self, roles: Vec, ) -> Result { - { - let roles: Vec = roles.iter().map(|r| &r.role).cloned().collect(); - let mut set = DB_ROLES.lock().unwrap(); - for role in roles { - set.insert(role); - } - } + let role_names: Vec = roles.iter().map(|r| &r.role).cloned().collect(); + add_roles(role_names); let role_statements: Vec = roles .into_iter() @@ -84,8 +86,25 @@ impl TestDb { } } + roles.sort(); + roles } + + async fn init_roles(&self) { + let results = sqlx::query!("select rolname from pg_catalog.pg_roles;") + .fetch_all(&self.pool) + .await + .unwrap(); + + let roles: Vec = results + .iter() + .filter_map(|r| r.rolname.as_ref()) + .cloned() + .collect(); + + add_roles(roles); + } } impl Deref for TestDb { @@ -134,5 +153,9 @@ pub async fn get_new_test_db() -> TestDb { .await .expect("Could not connect to test database"); - TestDb { pool } + let db = TestDb { pool }; + + db.init_roles().await; + + db } From 8e4d17cef3be76bc91ca60a15d0950a10d9b134b Mon Sep 17 00:00:00 2001 From: Julian Date: Wed, 28 May 2025 09:59:53 +0200 Subject: [PATCH 05/23] =?UTF-8?q?better=E2=80=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crates/pgt_test_utils/src/test_database.rs | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/crates/pgt_test_utils/src/test_database.rs b/crates/pgt_test_utils/src/test_database.rs index 8600af83..7d7791d2 100644 --- a/crates/pgt_test_utils/src/test_database.rs +++ b/crates/pgt_test_utils/src/test_database.rs @@ -1,7 +1,7 @@ use std::{ collections::HashSet, ops::Deref, - sync::{LazyLock, Mutex}, + sync::{LazyLock, Mutex, MutexGuard}, }; use sqlx::{ @@ -76,19 +76,8 @@ impl TestDb { self.pool.execute(query.as_str()).await } - pub fn get_roles(&self) -> Vec { - let mut roles = vec![]; - - { - let set = DB_ROLES.lock().unwrap(); - for role in set.iter() { - roles.push(role.clone()); - } - } - - roles.sort(); - - roles + pub fn get_roles(&self) -> MutexGuard<'_, HashSet> { + DB_ROLES.lock().unwrap() } async fn init_roles(&self) { From d1d8453432559acb0a20e05a8bd649f8fe33b5c9 Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 29 May 2025 08:04:35 +0200 Subject: [PATCH 06/23] sqlx prepare --- ...053db65ea6a7529e2cb97b2d3432a18aff6ba.json | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 .sqlx/query-b0504a4340264403ad43d05c60d053db65ea6a7529e2cb97b2d3432a18aff6ba.json diff --git a/.sqlx/query-b0504a4340264403ad43d05c60d053db65ea6a7529e2cb97b2d3432a18aff6ba.json b/.sqlx/query-b0504a4340264403ad43d05c60d053db65ea6a7529e2cb97b2d3432a18aff6ba.json new file mode 100644 index 00000000..dfc842b7 --- /dev/null +++ b/.sqlx/query-b0504a4340264403ad43d05c60d053db65ea6a7529e2cb97b2d3432a18aff6ba.json @@ -0,0 +1,20 @@ +{ + "db_name": "PostgreSQL", + "query": "select rolname from pg_catalog.pg_roles;", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "rolname", + "type_info": "Name" + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + true + ] + }, + "hash": "b0504a4340264403ad43d05c60d053db65ea6a7529e2cb97b2d3432a18aff6ba" +} From 9021bc02b457758d09511bafd4a70710433ba59a Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 29 May 2025 08:10:48 +0200 Subject: [PATCH 07/23] ok --- crates/pgt_completions/src/test_helper.rs | 1 - crates/pgt_lsp/tests/server.rs | 1 - crates/pgt_schema_cache/src/tables.rs | 1 - crates/pgt_schema_cache/src/triggers.rs | 1 - crates/pgt_typecheck/src/typed_identifier.rs | 1 - crates/pgt_typecheck/tests/diagnostics.rs | 1 - 6 files changed, 6 deletions(-) diff --git a/crates/pgt_completions/src/test_helper.rs b/crates/pgt_completions/src/test_helper.rs index 937c11af..19c8e966 100644 --- a/crates/pgt_completions/src/test_helper.rs +++ b/crates/pgt_completions/src/test_helper.rs @@ -2,7 +2,6 @@ use std::fmt::Display; use pgt_schema_cache::SchemaCache; use pgt_test_utils::test_database::get_new_test_db; -use sqlx::Executor; use crate::{CompletionItem, CompletionItemKind, CompletionParams, complete}; diff --git a/crates/pgt_lsp/tests/server.rs b/crates/pgt_lsp/tests/server.rs index 581ea1fe..746e829f 100644 --- a/crates/pgt_lsp/tests/server.rs +++ b/crates/pgt_lsp/tests/server.rs @@ -19,7 +19,6 @@ use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::Value; use serde_json::{from_value, to_value}; -use sqlx::Executor; use std::any::type_name; use std::fmt::Display; use std::time::Duration; diff --git a/crates/pgt_schema_cache/src/tables.rs b/crates/pgt_schema_cache/src/tables.rs index a0a40d6a..98b0be3e 100644 --- a/crates/pgt_schema_cache/src/tables.rs +++ b/crates/pgt_schema_cache/src/tables.rs @@ -81,7 +81,6 @@ impl SchemaCacheItem for Table { mod tests { use crate::{SchemaCache, tables::TableKind}; use pgt_test_utils::test_database::get_new_test_db; - use sqlx::Executor; #[tokio::test] async fn includes_views_in_query() { diff --git a/crates/pgt_schema_cache/src/triggers.rs b/crates/pgt_schema_cache/src/triggers.rs index 0a5241d6..80660008 100644 --- a/crates/pgt_schema_cache/src/triggers.rs +++ b/crates/pgt_schema_cache/src/triggers.rs @@ -127,7 +127,6 @@ impl SchemaCacheItem for Trigger { #[cfg(test)] mod tests { use pgt_test_utils::test_database::get_new_test_db; - use sqlx::Executor; use crate::{ SchemaCache, diff --git a/crates/pgt_typecheck/src/typed_identifier.rs b/crates/pgt_typecheck/src/typed_identifier.rs index 5efe0421..49152c2e 100644 --- a/crates/pgt_typecheck/src/typed_identifier.rs +++ b/crates/pgt_typecheck/src/typed_identifier.rs @@ -232,7 +232,6 @@ fn resolve_type<'a>( #[cfg(test)] mod tests { use pgt_test_utils::test_database::get_new_test_db; - use sqlx::Executor; #[tokio::test] async fn test_apply_identifiers() { diff --git a/crates/pgt_typecheck/tests/diagnostics.rs b/crates/pgt_typecheck/tests/diagnostics.rs index 9628962d..740669e7 100644 --- a/crates/pgt_typecheck/tests/diagnostics.rs +++ b/crates/pgt_typecheck/tests/diagnostics.rs @@ -5,7 +5,6 @@ use pgt_console::{ use pgt_diagnostics::PrintDiagnostic; use pgt_test_utils::test_database::get_new_test_db; use pgt_typecheck::{TypecheckParams, check_sql}; -use sqlx::Executor; async fn test(name: &str, query: &str, setup: Option<&str>) { let test_db = get_new_test_db().await; From 0753b50f0a7318a4b502447467f71aefbb802606 Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 29 May 2025 08:35:35 +0200 Subject: [PATCH 08/23] better --- crates/pgt_completions/src/test_helper.rs | 12 +++---- crates/pgt_lsp/tests/server.rs | 33 ++++++++----------- crates/pgt_test_utils/src/lib.rs | 2 ++ .../testdb_migrations/0001-setup-roles.sql | 23 +++++++++++++ 4 files changed, 45 insertions(+), 25 deletions(-) create mode 100644 crates/pgt_test_utils/testdb_migrations/0001-setup-roles.sql diff --git a/crates/pgt_completions/src/test_helper.rs b/crates/pgt_completions/src/test_helper.rs index 19c8e966..c8516108 100644 --- a/crates/pgt_completions/src/test_helper.rs +++ b/crates/pgt_completions/src/test_helper.rs @@ -1,7 +1,7 @@ use std::fmt::Display; use pgt_schema_cache::SchemaCache; -use pgt_test_utils::test_database::get_new_test_db; +use sqlx::{Executor, PgPool}; use crate::{CompletionItem, CompletionItemKind, CompletionParams, complete}; @@ -35,9 +35,8 @@ impl Display for InputQuery { pub(crate) async fn get_test_deps( setup: &str, input: InputQuery, + test_db: &PgPool, ) -> (tree_sitter::Tree, pgt_schema_cache::SchemaCache) { - let test_db = get_new_test_db().await; - test_db .execute(setup) .await @@ -206,8 +205,9 @@ pub(crate) async fn assert_complete_results( query: &str, assertions: Vec, setup: &str, + pool: &PgPool, ) { - let (tree, cache) = get_test_deps(setup, query.into()).await; + let (tree, cache) = get_test_deps(setup, query.into(), pool).await; let params = get_test_params(&tree, &cache, query.into()); let items = complete(params); @@ -240,8 +240,8 @@ pub(crate) async fn assert_complete_results( }); } -pub(crate) async fn assert_no_complete_results(query: &str, setup: &str) { - let (tree, cache) = get_test_deps(setup, query.into()).await; +pub(crate) async fn assert_no_complete_results(query: &str, setup: &str, pool: &PgPool) { + let (tree, cache) = get_test_deps(setup, query.into(), pool).await; let params = get_test_params(&tree, &cache, query.into()); let items = complete(params); diff --git a/crates/pgt_lsp/tests/server.rs b/crates/pgt_lsp/tests/server.rs index 746e829f..19b65b06 100644 --- a/crates/pgt_lsp/tests/server.rs +++ b/crates/pgt_lsp/tests/server.rs @@ -13,12 +13,13 @@ use pgt_configuration::database::PartialDatabaseConfiguration; use pgt_fs::MemoryFileSystem; use pgt_lsp::LSPServer; use pgt_lsp::ServerFactory; -use pgt_test_utils::test_database::get_new_test_db; use pgt_workspace::DynRef; use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::Value; use serde_json::{from_value, to_value}; +use sqlx::Executor; +use sqlx::PgPool; use std::any::type_name; use std::fmt::Display; use std::time::Duration; @@ -344,11 +345,10 @@ async fn basic_lifecycle() -> Result<()> { Ok(()) } -#[tokio::test] -async fn test_database_connection() -> Result<()> { +#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] +async fn test_database_connection(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); let mut fs = MemoryFileSystem::default(); - let test_db = get_new_test_db().await; let setup = r#" create table public.users ( @@ -456,11 +456,10 @@ async fn server_shutdown() -> Result<()> { Ok(()) } -#[tokio::test] -async fn test_completions() -> Result<()> { +#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] +async fn test_completions(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); let mut fs = MemoryFileSystem::default(); - let test_db = get_new_test_db().await; let setup = r#" create table public.users ( @@ -557,11 +556,10 @@ async fn test_completions() -> Result<()> { Ok(()) } -#[tokio::test] -async fn test_issue_271() -> Result<()> { +#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] +async fn test_issue_271(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); let mut fs = MemoryFileSystem::default(); - let test_db = get_new_test_db().await; let setup = r#" create table public.users ( @@ -759,11 +757,10 @@ async fn test_issue_271() -> Result<()> { Ok(()) } -#[tokio::test] -async fn test_execute_statement() -> Result<()> { +#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] +async fn test_execute_statement(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); let mut fs = MemoryFileSystem::default(); - let test_db = get_new_test_db().await; let database = test_db .connect_options() @@ -898,11 +895,10 @@ async fn test_execute_statement() -> Result<()> { Ok(()) } -#[tokio::test] -async fn test_issue_281() -> Result<()> { +#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] +async fn test_issue_281(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); let mut fs = MemoryFileSystem::default(); - let test_db = get_new_test_db().await; let setup = r#" create table public.users ( @@ -982,11 +978,10 @@ async fn test_issue_281() -> Result<()> { Ok(()) } -#[tokio::test] -async fn test_issue_303() -> Result<()> { +#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] +async fn test_issue_303(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); let mut fs = MemoryFileSystem::default(); - let test_db = get_new_test_db().await; let setup = r#" create table public.users ( diff --git a/crates/pgt_test_utils/src/lib.rs b/crates/pgt_test_utils/src/lib.rs index 4d6d3070..935975ca 100644 --- a/crates/pgt_test_utils/src/lib.rs +++ b/crates/pgt_test_utils/src/lib.rs @@ -1 +1,3 @@ pub mod test_database; + +pub static MIGRATIONS: sqlx::migrate::Migrator = sqlx::migrate!("./testdb_migrations"); diff --git a/crates/pgt_test_utils/testdb_migrations/0001-setup-roles.sql b/crates/pgt_test_utils/testdb_migrations/0001-setup-roles.sql new file mode 100644 index 00000000..67e6dfec --- /dev/null +++ b/crates/pgt_test_utils/testdb_migrations/0001-setup-roles.sql @@ -0,0 +1,23 @@ +do $$ +begin +if not exists ( + select from pg_catalog.pg_roles + where rolname = 'admin' +) then + create role admin superuser createdb login bypassrls; +end if; + +if not exists ( + select from pg_catalog.pg_roles + where rolname = 'test_login' +) then + create role test_login login; +end if; + +if not exists ( + select from pg_catalog.pg_roles + where rolname = 'test_nologin' +) then + create role test_nologin; +end if; +end; \ No newline at end of file From 1be61b4efdbbfbbb58578cb5ebbee096a4ac3b20 Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 29 May 2025 10:07:03 +0200 Subject: [PATCH 09/23] ok --- .../pgt_completions/src/providers/columns.rs | 137 ++++++++++------ .../src/providers/functions.rs | 26 +-- .../pgt_completions/src/providers/policies.rs | 14 +- .../pgt_completions/src/providers/schemas.rs | 16 +- .../pgt_completions/src/providers/tables.rs | 113 ++++++++----- .../src/relevance/filtering.rs | 26 +-- .../pgt_completions/src/relevance/scoring.rs | 27 +++- crates/pgt_completions/src/test_helper.rs | 16 +- crates/pgt_schema_cache/src/columns.rs | 8 +- crates/pgt_test_utils/src/lib.rs | 2 - crates/pgt_test_utils/src/test_database.rs | 150 ------------------ 11 files changed, 246 insertions(+), 289 deletions(-) delete mode 100644 crates/pgt_test_utils/src/test_database.rs diff --git a/crates/pgt_completions/src/providers/columns.rs b/crates/pgt_completions/src/providers/columns.rs index da6d23bc..b1dcbdf7 100644 --- a/crates/pgt_completions/src/providers/columns.rs +++ b/crates/pgt_completions/src/providers/columns.rs @@ -44,6 +44,8 @@ pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut Completio mod tests { use std::vec; + use sqlx::{Executor, PgPool}; + use crate::{ CompletionItem, CompletionItemKind, complete, test_helper::{ @@ -66,8 +68,8 @@ mod tests { } } - #[tokio::test] - async fn completes_columns() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn completes_columns(pool: PgPool) { let setup = r#" create schema private; @@ -87,6 +89,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + let queries: Vec = vec![ TestCase { message: "correctly prefers the columns of present tables", @@ -121,7 +125,7 @@ mod tests { ]; for q in queries { - let (tree, cache) = get_test_deps(setup, q.get_input_query()).await; + let (tree, cache) = get_test_deps(None, q.get_input_query(), &pool).await; let params = get_test_params(&tree, &cache, q.get_input_query()); let results = complete(params); @@ -137,8 +141,8 @@ mod tests { } } - #[tokio::test] - async fn shows_multiple_columns_if_no_relation_specified() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn shows_multiple_columns_if_no_relation_specified(pool: PgPool) { let setup = r#" create schema private; @@ -158,6 +162,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + let case = TestCase { query: format!(r#"select n{};"#, CURSOR_POS), description: "", @@ -165,11 +171,11 @@ mod tests { message: "", }; - let (tree, cache) = get_test_deps(setup, case.get_input_query()).await; + let (tree, cache) = get_test_deps(None, case.get_input_query(), &pool).await; let params = get_test_params(&tree, &cache, case.get_input_query()); let mut items = complete(params); - let _ = items.split_off(6); + let _ = items.split_off(4); #[derive(Eq, PartialEq, Debug)] struct LabelAndDesc { @@ -190,8 +196,6 @@ mod tests { ("narrator", "public.audio_books"), ("narrator_id", "private.audio_books"), ("id", "public.audio_books"), - ("name", "Schema: pg_catalog"), - ("nameconcatoid", "Schema: pg_catalog"), ] .into_iter() .map(|(label, schema)| LabelAndDesc { @@ -203,8 +207,8 @@ mod tests { assert_eq!(labels, expected); } - #[tokio::test] - async fn suggests_relevant_columns_without_letters() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn suggests_relevant_columns_without_letters(pool: PgPool) { let setup = r#" create table users ( id serial primary key, @@ -221,7 +225,7 @@ mod tests { description: "", }; - let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await; + let (tree, cache) = get_test_deps(Some(setup), test_case.get_input_query(), &pool).await; let params = get_test_params(&tree, &cache, test_case.get_input_query()); let results = complete(params); @@ -251,8 +255,8 @@ mod tests { ); } - #[tokio::test] - async fn ignores_cols_in_from_clause() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn ignores_cols_in_from_clause(pool: PgPool) { let setup = r#" create schema private; @@ -271,7 +275,7 @@ mod tests { description: "", }; - let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await; + let (tree, cache) = get_test_deps(Some(setup), test_case.get_input_query(), &pool).await; let params = get_test_params(&tree, &cache, test_case.get_input_query()); let results = complete(params); @@ -282,8 +286,8 @@ mod tests { ); } - #[tokio::test] - async fn prefers_columns_of_mentioned_tables() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn prefers_columns_of_mentioned_tables(pool: PgPool) { let setup = r#" create schema private; @@ -304,6 +308,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + assert_complete_results( format!(r#"select {} from users"#, CURSOR_POS).as_str(), vec![ @@ -312,7 +318,8 @@ mod tests { CompletionAssertion::Label("id2".into()), CompletionAssertion::Label("name2".into()), ], - setup, + None, + &pool, ) .await; @@ -324,7 +331,8 @@ mod tests { CompletionAssertion::Label("id1".into()), CompletionAssertion::Label("name1".into()), ], - setup, + None, + &pool, ) .await; @@ -332,13 +340,14 @@ mod tests { assert_complete_results( format!(r#"select sett{} from private.users"#, CURSOR_POS).as_str(), vec![CompletionAssertion::Label("user_settings".into())], - setup, + None, + &pool, ) .await; } - #[tokio::test] - async fn filters_out_by_aliases() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn filters_out_by_aliases(pool: PgPool) { let setup = r#" create schema auth; @@ -357,6 +366,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + // test in SELECT clause assert_complete_results( format!( @@ -374,7 +385,8 @@ mod tests { CompletionAssertion::Label("title".to_string()), CompletionAssertion::Label("user_id".to_string()), ], - setup, + None, + &pool, ) .await; @@ -396,13 +408,14 @@ mod tests { CompletionAssertion::Label("title".to_string()), CompletionAssertion::Label("user_id".to_string()), ], - setup, + None, + &pool, ) .await; } - #[tokio::test] - async fn does_not_complete_cols_in_join_clauses() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn does_not_complete_cols_in_join_clauses(pool: PgPool) { let setup = r#" create schema auth; @@ -435,13 +448,14 @@ mod tests { CompletionAssertion::LabelAndKind("posts".to_string(), CompletionItemKind::Table), CompletionAssertion::LabelAndKind("users".to_string(), CompletionItemKind::Table), ], - setup, + Some(setup), + &pool, ) .await; } - #[tokio::test] - async fn completes_in_join_on_clause() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn completes_in_join_on_clause(pool: PgPool) { let setup = r#" create schema auth; @@ -460,6 +474,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + assert_complete_results( format!( "select u.id, auth.posts.content from auth.users u join auth.posts on u.{}", @@ -472,7 +488,8 @@ mod tests { CompletionAssertion::LabelAndKind("email".to_string(), CompletionItemKind::Column), CompletionAssertion::LabelAndKind("name".to_string(), CompletionItemKind::Column), ], - setup, + None, + &pool, ) .await; @@ -488,13 +505,14 @@ mod tests { CompletionAssertion::LabelAndKind("email".to_string(), CompletionItemKind::Column), CompletionAssertion::LabelAndKind("name".to_string(), CompletionItemKind::Column), ], - setup, + None, + &pool, ) .await; } - #[tokio::test] - async fn prefers_not_mentioned_columns() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn prefers_not_mentioned_columns(pool: PgPool) { let setup = r#" create schema auth; @@ -513,6 +531,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + assert_complete_results( format!( "select {} from public.one o join public.two on o.id = t.id;", @@ -526,7 +546,8 @@ mod tests { CompletionAssertion::Label("d".to_string()), CompletionAssertion::Label("e".to_string()), ], - setup, + None, + &pool, ) .await; @@ -546,7 +567,8 @@ mod tests { CompletionAssertion::Label("z".to_string()), CompletionAssertion::Label("a".to_string()), ], - setup, + None, + &pool, ) .await; @@ -562,7 +584,8 @@ mod tests { CompletionAssertion::LabelAndDesc("id".to_string(), "public.two".to_string()), CompletionAssertion::Label("z".to_string()), ], - setup, + None, + &pool, ) .await; @@ -574,13 +597,14 @@ mod tests { ) .as_str(), vec![CompletionAssertion::Label("z".to_string())], - setup, + None, + &pool, ) .await; } - #[tokio::test] - async fn suggests_columns_in_insert_clause() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn suggests_columns_in_insert_clause(pool: PgPool) { let setup = r#" create table instruments ( id bigint primary key generated always as identity, @@ -595,6 +619,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + // We should prefer the instrument columns, even though they // are lower in the alphabet @@ -605,7 +631,8 @@ mod tests { CompletionAssertion::Label("name".to_string()), CompletionAssertion::Label("z".to_string()), ], - setup, + None, + &pool, ) .await; @@ -615,14 +642,16 @@ mod tests { CompletionAssertion::Label("name".to_string()), CompletionAssertion::Label("z".to_string()), ], - setup, + None, + &pool, ) .await; assert_complete_results( format!("insert into instruments (id, {}, name)", CURSOR_POS).as_str(), vec![CompletionAssertion::Label("z".to_string())], - setup, + None, + &pool, ) .await; @@ -637,20 +666,22 @@ mod tests { CompletionAssertion::Label("id".to_string()), CompletionAssertion::Label("z".to_string()), ], - setup, + None, + &pool, ) .await; // no completions in the values list! assert_no_complete_results( format!("insert into instruments (id, name) values ({})", CURSOR_POS).as_str(), - setup, + None, + &pool, ) .await; } - #[tokio::test] - async fn suggests_columns_in_where_clause() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn suggests_columns_in_where_clause(pool: PgPool) { let setup = r#" create table instruments ( id bigint primary key generated always as identity, @@ -666,6 +697,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + assert_complete_results( format!("select name from instruments where {} ", CURSOR_POS).as_str(), vec![ @@ -674,7 +707,8 @@ mod tests { CompletionAssertion::Label("name".into()), CompletionAssertion::Label("z".into()), ], - setup, + None, + &pool, ) .await; @@ -689,7 +723,8 @@ mod tests { CompletionAssertion::KindNotExists(CompletionItemKind::Column), CompletionAssertion::KindNotExists(CompletionItemKind::Schema), ], - setup, + None, + &pool, ) .await; @@ -705,7 +740,8 @@ mod tests { CompletionAssertion::Label("name".into()), CompletionAssertion::Label("z".into()), ], - setup, + None, + &pool, ) .await; @@ -721,7 +757,8 @@ mod tests { CompletionAssertion::Label("id".into()), CompletionAssertion::Label("name".into()), ], - setup, + None, + &pool, ) .await; } diff --git a/crates/pgt_completions/src/providers/functions.rs b/crates/pgt_completions/src/providers/functions.rs index f1b57e8c..2bc4f331 100644 --- a/crates/pgt_completions/src/providers/functions.rs +++ b/crates/pgt_completions/src/providers/functions.rs @@ -65,13 +65,15 @@ fn get_completion_text(ctx: &CompletionContext, func: &Function) -> CompletionTe #[cfg(test)] mod tests { + use sqlx::PgPool; + use crate::{ CompletionItem, CompletionItemKind, complete, test_helper::{CURSOR_POS, get_test_deps, get_test_params}, }; - #[tokio::test] - async fn completes_fn() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn completes_fn(pool: PgPool) { let setup = r#" create or replace function cool() returns trigger @@ -86,7 +88,7 @@ mod tests { let query = format!("select coo{}", CURSOR_POS); - let (tree, cache) = get_test_deps(setup, query.as_str().into()).await; + let (tree, cache) = get_test_deps(Some(setup), query.as_str().into(), &pool).await; let params = get_test_params(&tree, &cache, query.as_str().into()); let results = complete(params); @@ -98,8 +100,8 @@ mod tests { assert_eq!(label, "cool"); } - #[tokio::test] - async fn prefers_fn_if_invocation() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn prefers_fn_if_invocation(pool: PgPool) { let setup = r#" create table coos ( id serial primary key, @@ -119,7 +121,7 @@ mod tests { let query = format!(r#"select * from coo{}()"#, CURSOR_POS); - let (tree, cache) = get_test_deps(setup, query.as_str().into()).await; + let (tree, cache) = get_test_deps(Some(setup), query.as_str().into(), &pool).await; let params = get_test_params(&tree, &cache, query.as_str().into()); let results = complete(params); @@ -132,8 +134,8 @@ mod tests { assert_eq!(kind, CompletionItemKind::Function); } - #[tokio::test] - async fn prefers_fn_in_select_clause() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn prefers_fn_in_select_clause(pool: PgPool) { let setup = r#" create table coos ( id serial primary key, @@ -153,7 +155,7 @@ mod tests { let query = format!(r#"select coo{}"#, CURSOR_POS); - let (tree, cache) = get_test_deps(setup, query.as_str().into()).await; + let (tree, cache) = get_test_deps(Some(setup), query.as_str().into(), &pool).await; let params = get_test_params(&tree, &cache, query.as_str().into()); let results = complete(params); @@ -166,8 +168,8 @@ mod tests { assert_eq!(kind, CompletionItemKind::Function); } - #[tokio::test] - async fn prefers_function_in_from_clause_if_invocation() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn prefers_function_in_from_clause_if_invocation(pool: PgPool) { let setup = r#" create table coos ( id serial primary key, @@ -187,7 +189,7 @@ mod tests { let query = format!(r#"select * from coo{}()"#, CURSOR_POS); - let (tree, cache) = get_test_deps(setup, query.as_str().into()).await; + let (tree, cache) = get_test_deps(Some(setup), query.as_str().into(), &pool).await; let params = get_test_params(&tree, &cache, query.as_str().into()); let results = complete(params); diff --git a/crates/pgt_completions/src/providers/policies.rs b/crates/pgt_completions/src/providers/policies.rs index a4d3a9bb..216fcefa 100644 --- a/crates/pgt_completions/src/providers/policies.rs +++ b/crates/pgt_completions/src/providers/policies.rs @@ -59,10 +59,12 @@ pub fn complete_policies<'a>(ctx: &CompletionContext<'a>, builder: &mut Completi #[cfg(test)] mod tests { + use sqlx::{Executor, PgPool}; + use crate::test_helper::{CURSOR_POS, CompletionAssertion, assert_complete_results}; - #[tokio::test] - async fn completes_within_quotation_marks() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn completes_within_quotation_marks(pool: PgPool) { let setup = r#" create schema private; @@ -84,13 +86,16 @@ mod tests { with check (true); "#; + pool.execute(setup).await.unwrap(); + assert_complete_results( format!("alter policy \"{}\" on private.users;", CURSOR_POS).as_str(), vec![ CompletionAssertion::Label("read for public users disallowed".into()), CompletionAssertion::Label("write for public users allowed".into()), ], - setup, + None, + &pool, ) .await; @@ -99,7 +104,8 @@ mod tests { vec![CompletionAssertion::Label( "write for public users allowed".into(), )], - setup, + None, + &pool, ) .await; } diff --git a/crates/pgt_completions/src/providers/schemas.rs b/crates/pgt_completions/src/providers/schemas.rs index 02d2fd0c..561da0f8 100644 --- a/crates/pgt_completions/src/providers/schemas.rs +++ b/crates/pgt_completions/src/providers/schemas.rs @@ -27,13 +27,15 @@ pub fn complete_schemas<'a>(ctx: &'a CompletionContext, builder: &mut Completion #[cfg(test)] mod tests { + use sqlx::PgPool; + use crate::{ CompletionItemKind, test_helper::{CURSOR_POS, CompletionAssertion, assert_complete_results}, }; - #[tokio::test] - async fn autocompletes_schemas() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn autocompletes_schemas(pool: PgPool) { let setup = r#" create schema private; create schema auth; @@ -75,13 +77,14 @@ mod tests { CompletionItemKind::Schema, ), ], - setup, + Some(setup), + &pool, ) .await; } - #[tokio::test] - async fn suggests_tables_and_schemas_with_matching_keys() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn suggests_tables_and_schemas_with_matching_keys(pool: PgPool) { let setup = r#" create schema ultimate; @@ -99,7 +102,8 @@ mod tests { CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), CompletionAssertion::LabelAndKind("ultimate".into(), CompletionItemKind::Schema), ], - setup, + Some(setup), + &pool, ) .await; } diff --git a/crates/pgt_completions/src/providers/tables.rs b/crates/pgt_completions/src/providers/tables.rs index 2102d41c..3fbee8f1 100644 --- a/crates/pgt_completions/src/providers/tables.rs +++ b/crates/pgt_completions/src/providers/tables.rs @@ -42,6 +42,8 @@ pub fn complete_tables<'a>(ctx: &'a CompletionContext, builder: &mut CompletionB #[cfg(test)] mod tests { + use sqlx::{Executor, PgPool}; + use crate::{ CompletionItem, CompletionItemKind, complete, test_helper::{ @@ -50,8 +52,8 @@ mod tests { }, }; - #[tokio::test] - async fn autocompletes_simple_table() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn autocompletes_simple_table(pool: PgPool) { let setup = r#" create table users ( id serial primary key, @@ -62,7 +64,7 @@ mod tests { let query = format!("select * from u{}", CURSOR_POS); - let (tree, cache) = get_test_deps(setup, query.as_str().into()).await; + let (tree, cache) = get_test_deps(Some(setup), query.as_str().into(), &pool).await; let params = get_test_params(&tree, &cache, query.as_str().into()); let items = complete(params); @@ -77,8 +79,8 @@ mod tests { ) } - #[tokio::test] - async fn autocompletes_table_alphanumerically() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn autocompletes_table_alphanumerically(pool: PgPool) { let setup = r#" create table addresses ( id serial primary key @@ -93,6 +95,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + let test_cases = vec![ (format!("select * from u{}", CURSOR_POS), "users"), (format!("select * from e{}", CURSOR_POS), "emails"), @@ -100,7 +104,7 @@ mod tests { ]; for (query, expected_label) in test_cases { - let (tree, cache) = get_test_deps(setup, query.as_str().into()).await; + let (tree, cache) = get_test_deps(None, query.as_str().into(), &pool).await; let params = get_test_params(&tree, &cache, query.as_str().into()); let items = complete(params); @@ -116,8 +120,8 @@ mod tests { } } - #[tokio::test] - async fn autocompletes_table_with_schema() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn autocompletes_table_with_schema(pool: PgPool) { let setup = r#" create schema customer_support; create schema private; @@ -135,6 +139,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + let test_cases = vec![ (format!("select * from u{}", CURSOR_POS), "user_y"), // user_y is preferred alphanumerically (format!("select * from private.u{}", CURSOR_POS), "user_z"), @@ -145,7 +151,7 @@ mod tests { ]; for (query, expected_label) in test_cases { - let (tree, cache) = get_test_deps(setup, query.as_str().into()).await; + let (tree, cache) = get_test_deps(None, query.as_str().into(), &pool).await; let params = get_test_params(&tree, &cache, query.as_str().into()); let items = complete(params); @@ -161,8 +167,8 @@ mod tests { } } - #[tokio::test] - async fn prefers_table_in_from_clause() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn prefers_table_in_from_clause(pool: PgPool) { let setup = r#" create table coos ( id serial primary key, @@ -182,7 +188,7 @@ mod tests { let query = format!(r#"select * from coo{}"#, CURSOR_POS); - let (tree, cache) = get_test_deps(setup, query.as_str().into()).await; + let (tree, cache) = get_test_deps(Some(setup), query.as_str().into(), &pool).await; let params = get_test_params(&tree, &cache, query.as_str().into()); let items = complete(params); @@ -195,8 +201,8 @@ mod tests { assert_eq!(kind, CompletionItemKind::Table); } - #[tokio::test] - async fn suggests_tables_in_update() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn suggests_tables_in_update(pool: PgPool) { let setup = r#" create table coos ( id serial primary key, @@ -204,13 +210,16 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + assert_complete_results( format!("update {}", CURSOR_POS).as_str(), vec![CompletionAssertion::LabelAndKind( "public".into(), CompletionItemKind::Schema, )], - setup, + None, + &pool, ) .await; @@ -220,12 +229,17 @@ mod tests { "coos".into(), CompletionItemKind::Table, )], - setup, + None, + &pool, ) .await; - assert_no_complete_results(format!("update public.coos {}", CURSOR_POS).as_str(), setup) - .await; + assert_no_complete_results( + format!("update public.coos {}", CURSOR_POS).as_str(), + None, + &pool, + ) + .await; assert_complete_results( format!("update coos set {}", CURSOR_POS).as_str(), @@ -233,7 +247,8 @@ mod tests { CompletionAssertion::Label("id".into()), CompletionAssertion::Label("name".into()), ], - setup, + None, + &pool, ) .await; @@ -243,13 +258,14 @@ mod tests { CompletionAssertion::Label("id".into()), CompletionAssertion::Label("name".into()), ], - setup, + None, + &pool, ) .await; } - #[tokio::test] - async fn suggests_tables_in_delete() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn suggests_tables_in_delete(pool: PgPool) { let setup = r#" create table coos ( id serial primary key, @@ -257,7 +273,9 @@ mod tests { ); "#; - assert_no_complete_results(format!("delete {}", CURSOR_POS).as_str(), setup).await; + pool.execute(setup).await.unwrap(); + + assert_no_complete_results(format!("delete {}", CURSOR_POS).as_str(), None, &pool).await; assert_complete_results( format!("delete from {}", CURSOR_POS).as_str(), @@ -265,14 +283,16 @@ mod tests { CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema), CompletionAssertion::LabelAndKind("coos".into(), CompletionItemKind::Table), ], - setup, + None, + &pool, ) .await; assert_complete_results( format!("delete from public.{}", CURSOR_POS).as_str(), vec![CompletionAssertion::Label("coos".into())], - setup, + None, + &pool, ) .await; @@ -282,13 +302,14 @@ mod tests { CompletionAssertion::Label("id".into()), CompletionAssertion::Label("name".into()), ], - setup, + None, + &pool, ) .await; } - #[tokio::test] - async fn suggests_tables_in_join() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn suggests_tables_in_join(pool: PgPool) { let setup = r#" create schema auth; @@ -315,13 +336,14 @@ mod tests { CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), // self-join CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), ], - setup, + Some(setup), + &pool, ) .await; } - #[tokio::test] - async fn suggests_tables_in_alter_and_drop_statements() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn suggests_tables_in_alter_and_drop_statements(pool: PgPool) { let setup = r#" create schema auth; @@ -340,6 +362,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + assert_complete_results( format!("alter table {}", CURSOR_POS).as_str(), vec![ @@ -348,7 +372,8 @@ mod tests { CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), ], - setup, + None, + &pool, ) .await; @@ -360,7 +385,8 @@ mod tests { CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), ], - setup, + None, + &pool, ) .await; @@ -372,7 +398,8 @@ mod tests { CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), ], - setup, + None, + &pool, ) .await; @@ -384,13 +411,14 @@ mod tests { CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), // self-join CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), ], - setup, + None, + &pool, ) .await; } - #[tokio::test] - async fn suggests_tables_in_insert_into() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn suggests_tables_in_insert_into(pool: PgPool) { let setup = r#" create schema auth; @@ -401,6 +429,8 @@ mod tests { ); "#; + pool.execute(setup).await.unwrap(); + assert_complete_results( format!("insert into {}", CURSOR_POS).as_str(), vec![ @@ -408,7 +438,8 @@ mod tests { CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema), CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), ], - setup, + None, + &pool, ) .await; @@ -418,7 +449,8 @@ mod tests { "users".into(), CompletionItemKind::Table, )], - setup, + None, + &pool, ) .await; @@ -434,7 +466,8 @@ mod tests { CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema), CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), ], - setup, + None, + &pool, ) .await; } diff --git a/crates/pgt_completions/src/relevance/filtering.rs b/crates/pgt_completions/src/relevance/filtering.rs index 5323e2bc..0be9e48a 100644 --- a/crates/pgt_completions/src/relevance/filtering.rs +++ b/crates/pgt_completions/src/relevance/filtering.rs @@ -218,12 +218,14 @@ impl CompletionFilter<'_> { #[cfg(test)] mod tests { + use sqlx::{Executor, PgPool}; + use crate::test_helper::{ CURSOR_POS, CompletionAssertion, assert_complete_results, assert_no_complete_results, }; - #[tokio::test] - async fn completion_after_asterisk() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn completion_after_asterisk(pool: PgPool) { let setup = r#" create table users ( id serial primary key, @@ -232,7 +234,9 @@ mod tests { ); "#; - assert_no_complete_results(format!("select * {}", CURSOR_POS).as_str(), setup).await; + pool.execute(setup).await.unwrap(); + + assert_no_complete_results(format!("select * {}", CURSOR_POS).as_str(), None, &pool).await; // if there s a COMMA after the asterisk, we're good assert_complete_results( @@ -242,19 +246,21 @@ mod tests { CompletionAssertion::Label("email".into()), CompletionAssertion::Label("id".into()), ], - setup, + None, + &pool, ) .await; } - #[tokio::test] - async fn completion_after_create_table() { - assert_no_complete_results(format!("create table {}", CURSOR_POS).as_str(), "").await; + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn completion_after_create_table(pool: PgPool) { + assert_no_complete_results(format!("create table {}", CURSOR_POS).as_str(), None, &pool) + .await; } - #[tokio::test] - async fn completion_in_column_definitions() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn completion_in_column_definitions(pool: PgPool) { let query = format!(r#"create table instruments ( {} )"#, CURSOR_POS); - assert_no_complete_results(query.as_str(), "").await; + assert_no_complete_results(query.as_str(), None, &pool).await; } } diff --git a/crates/pgt_completions/src/relevance/scoring.rs b/crates/pgt_completions/src/relevance/scoring.rs index 2fe12511..a8c89f50 100644 --- a/crates/pgt_completions/src/relevance/scoring.rs +++ b/crates/pgt_completions/src/relevance/scoring.rs @@ -187,6 +187,16 @@ impl CompletionScore<'_> { } } + fn get_item_name(&self) -> &str { + match self.data { + CompletionRelevanceData::Table(t) => t.name.as_str(), + CompletionRelevanceData::Function(f) => f.name.as_str(), + CompletionRelevanceData::Column(c) => c.name.as_str(), + CompletionRelevanceData::Schema(s) => s.name.as_str(), + CompletionRelevanceData::Policy(p) => p.name.as_str(), + } + } + fn get_schema_name(&self) -> &str { match self.data { CompletionRelevanceData::Function(f) => f.schema.as_str(), @@ -234,19 +244,30 @@ impl CompletionScore<'_> { } fn check_is_user_defined(&mut self) { - let schema = self.get_schema_name().to_string(); + let schema_name = self.get_schema_name().to_string(); let system_schemas = ["pg_catalog", "information_schema", "pg_toast"]; - if system_schemas.contains(&schema.as_str()) { + if system_schemas.contains(&schema_name.as_str()) { self.score -= 20; } // "public" is the default postgres schema where users // create objects. Prefer it by a slight bit. - if schema.as_str() == "public" { + if schema_name.as_str() == "public" { self.score += 2; } + + let item_name = self.get_item_name().to_string(); + let table_name = self.get_table_name(); + + // migrations shouldn't pop up on top + if item_name.contains("migrations") + || table_name.is_some_and(|t| t.contains("migrations")) + || schema_name.contains("migrations") + { + self.score -= 15; + } } fn check_columns_in_stmt(&mut self, ctx: &CompletionContext) { diff --git a/crates/pgt_completions/src/test_helper.rs b/crates/pgt_completions/src/test_helper.rs index c8516108..61decce0 100644 --- a/crates/pgt_completions/src/test_helper.rs +++ b/crates/pgt_completions/src/test_helper.rs @@ -33,14 +33,16 @@ impl Display for InputQuery { } pub(crate) async fn get_test_deps( - setup: &str, + setup: Option<&str>, input: InputQuery, test_db: &PgPool, ) -> (tree_sitter::Tree, pgt_schema_cache::SchemaCache) { - test_db - .execute(setup) - .await - .expect("Failed to execute setup query"); + if let Some(setup) = setup { + test_db + .execute(setup) + .await + .expect("Failed to execute setup query"); + } let schema_cache = SchemaCache::load(&test_db) .await @@ -204,7 +206,7 @@ impl CompletionAssertion { pub(crate) async fn assert_complete_results( query: &str, assertions: Vec, - setup: &str, + setup: Option<&str>, pool: &PgPool, ) { let (tree, cache) = get_test_deps(setup, query.into(), pool).await; @@ -240,7 +242,7 @@ pub(crate) async fn assert_complete_results( }); } -pub(crate) async fn assert_no_complete_results(query: &str, setup: &str, pool: &PgPool) { +pub(crate) async fn assert_no_complete_results(query: &str, setup: Option<&str>, pool: &PgPool) { let (tree, cache) = get_test_deps(setup, query.into(), pool).await; let params = get_test_params(&tree, &cache, query.into()); let items = complete(params); diff --git a/crates/pgt_schema_cache/src/columns.rs b/crates/pgt_schema_cache/src/columns.rs index 943cf9ca..786ca5d4 100644 --- a/crates/pgt_schema_cache/src/columns.rs +++ b/crates/pgt_schema_cache/src/columns.rs @@ -82,14 +82,12 @@ impl SchemaCacheItem for Column { #[cfg(test)] mod tests { - use pgt_test_utils::test_database::get_new_test_db; + use sqlx::{Executor, PgPool}; use crate::{SchemaCache, columns::ColumnClassKind}; - #[tokio::test] - async fn loads_columns() { - let test_db = get_new_test_db().await; - + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn loads_columns(test_db: PgPool) { let setup = r#" create table public.users ( id serial primary key, diff --git a/crates/pgt_test_utils/src/lib.rs b/crates/pgt_test_utils/src/lib.rs index 935975ca..e21c6ce4 100644 --- a/crates/pgt_test_utils/src/lib.rs +++ b/crates/pgt_test_utils/src/lib.rs @@ -1,3 +1 @@ -pub mod test_database; - pub static MIGRATIONS: sqlx::migrate::Migrator = sqlx::migrate!("./testdb_migrations"); diff --git a/crates/pgt_test_utils/src/test_database.rs b/crates/pgt_test_utils/src/test_database.rs deleted file mode 100644 index 7d7791d2..00000000 --- a/crates/pgt_test_utils/src/test_database.rs +++ /dev/null @@ -1,150 +0,0 @@ -use std::{ - collections::HashSet, - ops::Deref, - sync::{LazyLock, Mutex, MutexGuard}, -}; - -use sqlx::{ - Executor, PgPool, - postgres::{PgConnectOptions, PgQueryResult}, -}; -use uuid::Uuid; - -static DB_ROLES: LazyLock>> = LazyLock::new(|| Mutex::new(HashSet::new())); - -fn add_roles(roles: Vec) { - let mut set = DB_ROLES.lock().unwrap(); - for role in roles { - set.insert(role); - } -} - -#[derive(Debug)] -pub struct TestDb { - pool: PgPool, -} - -#[derive(Debug)] -pub struct RoleWithArgs { - pub role: String, - pub args: Vec, -} - -impl TestDb { - pub async fn execute(&self, sql: &str) -> Result { - if sql.to_ascii_lowercase().contains("create role") { - panic!("Please setup roles via the `setup_roles` method.") - } - self.pool.execute(sql).await - } - - pub async fn setup_roles( - &mut self, - roles: Vec, - ) -> Result { - let role_names: Vec = roles.iter().map(|r| &r.role).cloned().collect(); - add_roles(role_names); - - let role_statements: Vec = roles - .into_iter() - .map(|r| { - format!( - r#" - if not exists ( - select from pg_catalog.pg_roles - where rolname = '{0}' - ) then - create role {0} {1}; - end if; - "#, - r.role, - r.args.join(" ") - ) - }) - .collect(); - - let query = format!( - r#" - do $$ - begin - {} - end $$; - "#, - role_statements.join("\n") - ); - - self.pool.execute(query.as_str()).await - } - - pub fn get_roles(&self) -> MutexGuard<'_, HashSet> { - DB_ROLES.lock().unwrap() - } - - async fn init_roles(&self) { - let results = sqlx::query!("select rolname from pg_catalog.pg_roles;") - .fetch_all(&self.pool) - .await - .unwrap(); - - let roles: Vec = results - .iter() - .filter_map(|r| r.rolname.as_ref()) - .cloned() - .collect(); - - add_roles(roles); - } -} - -impl Deref for TestDb { - type Target = PgPool; - fn deref(&self) -> &Self::Target { - &self.pool - } -} - -// TODO: Work with proper config objects instead of a connection_string. -// With the current implementation, we can't parse the password from the connection string. -pub async fn get_new_test_db() -> TestDb { - dotenv::dotenv().expect("Unable to load .env file for tests"); - - let connection_string = std::env::var("DATABASE_URL").expect("DATABASE_URL not set"); - let password = std::env::var("DB_PASSWORD").unwrap_or("postgres".into()); - - let options_from_conn_str: PgConnectOptions = connection_string - .parse() - .expect("Invalid Connection String"); - - let host = options_from_conn_str.get_host(); - assert!( - host == "localhost" || host == "127.0.0.1", - "Running tests against non-local database!" - ); - - let options_without_db_name = PgConnectOptions::new() - .host(host) - .port(options_from_conn_str.get_port()) - .username(options_from_conn_str.get_username()) - .password(&password); - - let postgres = sqlx::PgPool::connect_with(options_without_db_name.clone()) - .await - .expect("Unable to connect to test postgres instance"); - - let database_name = Uuid::new_v4().to_string(); - - postgres - .execute(format!(r#"create database "{}";"#, database_name).as_str()) - .await - .expect("Failed to create test database."); - - let pool = sqlx::PgPool::connect_with(options_without_db_name.database(&database_name)) - .await - .expect("Could not connect to test database"); - - let db = TestDb { pool }; - - db.init_roles().await; - - db -} From e55271f804a91bb1f88f06e4ca4ca718fa90474d Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 29 May 2025 10:17:09 +0200 Subject: [PATCH 10/23] ok --- crates/pgt_schema_cache/src/columns.rs | 2 +- crates/pgt_schema_cache/src/policies.rs | 23 +++----------- crates/pgt_schema_cache/src/roles.rs | 32 +++----------------- crates/pgt_schema_cache/src/schema_cache.rs | 8 ++--- crates/pgt_schema_cache/src/tables.rs | 15 ++++----- crates/pgt_schema_cache/src/triggers.rs | 15 ++++----- crates/pgt_typecheck/src/typed_identifier.rs | 8 ++--- crates/pgt_typecheck/tests/diagnostics.rs | 11 +++---- 8 files changed, 32 insertions(+), 82 deletions(-) diff --git a/crates/pgt_schema_cache/src/columns.rs b/crates/pgt_schema_cache/src/columns.rs index 786ca5d4..01f9b41c 100644 --- a/crates/pgt_schema_cache/src/columns.rs +++ b/crates/pgt_schema_cache/src/columns.rs @@ -126,7 +126,7 @@ mod tests { let public_schema_columns = cache .columns .iter() - .filter(|c| c.schema_name.as_str() == "public") + .filter(|c| c.schema_name.as_str() == "public" && !c.table_name.contains("migrations")) .count(); assert_eq!(public_schema_columns, 4); diff --git a/crates/pgt_schema_cache/src/policies.rs b/crates/pgt_schema_cache/src/policies.rs index 808a2b7e..b754725b 100644 --- a/crates/pgt_schema_cache/src/policies.rs +++ b/crates/pgt_schema_cache/src/policies.rs @@ -80,28 +80,13 @@ impl SchemaCacheItem for Policy { #[cfg(test)] mod tests { - use pgt_test_utils::test_database::{RoleWithArgs, get_new_test_db}; - use crate::{SchemaCache, policies::PolicyCommand}; - - #[tokio::test] - async fn loads_policies() { - let mut test_db = get_new_test_db().await; + use sqlx::{Executor, PgPool}; - test_db - .setup_roles(vec![ - RoleWithArgs { - role: "admin".into(), - args: vec![], - }, - RoleWithArgs { - role: "owner".into(), - args: vec![], - }, - ]) - .await - .expect("Unable to setup admin roles"); + use crate::{SchemaCache, policies::PolicyCommand}; + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn loads_policies(test_db: PgPool) { let setup = r#" create table public.users ( id serial primary key, diff --git a/crates/pgt_schema_cache/src/roles.rs b/crates/pgt_schema_cache/src/roles.rs index eef02c96..636af34a 100644 --- a/crates/pgt_schema_cache/src/roles.rs +++ b/crates/pgt_schema_cache/src/roles.rs @@ -21,36 +21,12 @@ impl SchemaCacheItem for Role { #[cfg(test)] mod tests { - use crate::SchemaCache; - use pgt_test_utils::test_database::{RoleWithArgs, get_new_test_db}; - - #[tokio::test] - async fn loads_roles() { - let mut test_db = get_new_test_db().await; + use sqlx::PgPool; - test_db - .setup_roles(vec![ - RoleWithArgs { - role: "test_super".into(), - args: vec![ - "superuser".into(), - "createdb".into(), - "login".into(), - "bypassrls".into(), - ], - }, - RoleWithArgs { - role: "test_nologin".into(), - args: vec![], - }, - RoleWithArgs { - role: "test_login".into(), - args: vec!["login".into()], - }, - ]) - .await - .expect("Unable to set up roles."); + use crate::SchemaCache; + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn loads_roles(test_db: PgPool) { let cache = SchemaCache::load(&test_db) .await .expect("Failed to load Schema Cache"); diff --git a/crates/pgt_schema_cache/src/schema_cache.rs b/crates/pgt_schema_cache/src/schema_cache.rs index 516b37e6..8fb9683b 100644 --- a/crates/pgt_schema_cache/src/schema_cache.rs +++ b/crates/pgt_schema_cache/src/schema_cache.rs @@ -93,14 +93,12 @@ pub trait SchemaCacheItem { #[cfg(test)] mod tests { - use pgt_test_utils::test_database::get_new_test_db; + use sqlx::PgPool; use crate::SchemaCache; - #[tokio::test] - async fn it_loads() { - let test_db = get_new_test_db().await; - + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn it_loads(test_db: PgPool) { SchemaCache::load(&test_db) .await .expect("Couldnt' load Schema Cache"); diff --git a/crates/pgt_schema_cache/src/tables.rs b/crates/pgt_schema_cache/src/tables.rs index 98b0be3e..16b86c54 100644 --- a/crates/pgt_schema_cache/src/tables.rs +++ b/crates/pgt_schema_cache/src/tables.rs @@ -79,13 +79,12 @@ impl SchemaCacheItem for Table { #[cfg(test)] mod tests { - use crate::{SchemaCache, tables::TableKind}; - use pgt_test_utils::test_database::get_new_test_db; + use sqlx::{Executor, PgPool}; - #[tokio::test] - async fn includes_views_in_query() { - let test_db = get_new_test_db().await; + use crate::{SchemaCache, tables::TableKind}; + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn includes_views_in_query(test_db: PgPool) { let setup = r#" create table public.base_table ( id serial primary key, @@ -115,10 +114,8 @@ mod tests { assert_eq!(view.schema, "public"); } - #[tokio::test] - async fn includes_materialized_views_in_query() { - let test_db = get_new_test_db().await; - + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn includes_materialized_views_in_query(test_db: PgPool) { let setup = r#" create table public.base_table ( id serial primary key, diff --git a/crates/pgt_schema_cache/src/triggers.rs b/crates/pgt_schema_cache/src/triggers.rs index 80660008..2b2a3aff 100644 --- a/crates/pgt_schema_cache/src/triggers.rs +++ b/crates/pgt_schema_cache/src/triggers.rs @@ -126,17 +126,16 @@ impl SchemaCacheItem for Trigger { #[cfg(test)] mod tests { - use pgt_test_utils::test_database::get_new_test_db; + + use sqlx::{Executor, PgPool}; use crate::{ SchemaCache, triggers::{TriggerAffected, TriggerEvent, TriggerTiming}, }; - #[tokio::test] - async fn loads_triggers() { - let test_db = get_new_test_db().await; - + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn loads_triggers(test_db: PgPool) { let setup = r#" create table public.users ( id serial primary key, @@ -218,10 +217,8 @@ mod tests { assert_eq!(delete_trigger.proc_name, "log_user_insert"); } - #[tokio::test] - async fn loads_instead_and_truncate_triggers() { - let test_db = get_new_test_db().await; - + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn loads_instead_and_truncate_triggers(test_db: PgPool) { let setup = r#" create table public.docs ( id serial primary key, diff --git a/crates/pgt_typecheck/src/typed_identifier.rs b/crates/pgt_typecheck/src/typed_identifier.rs index 49152c2e..710b2fe9 100644 --- a/crates/pgt_typecheck/src/typed_identifier.rs +++ b/crates/pgt_typecheck/src/typed_identifier.rs @@ -231,10 +231,10 @@ fn resolve_type<'a>( #[cfg(test)] mod tests { - use pgt_test_utils::test_database::get_new_test_db; + use sqlx::{Executor, PgPool}; - #[tokio::test] - async fn test_apply_identifiers() { + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn test_apply_identifiers(test_db: PgPool) { let input = "select v_test + fn_name.custom_type.v_test2 + $3 + custom_type.v_test3 + fn_name.v_test2 + enum_type"; let identifiers = vec![ @@ -294,8 +294,6 @@ mod tests { }, ]; - let test_db = get_new_test_db().await; - let setup = r#" CREATE TYPE "public"."custom_type" AS ( v_test2 integer, diff --git a/crates/pgt_typecheck/tests/diagnostics.rs b/crates/pgt_typecheck/tests/diagnostics.rs index 740669e7..9bf5d786 100644 --- a/crates/pgt_typecheck/tests/diagnostics.rs +++ b/crates/pgt_typecheck/tests/diagnostics.rs @@ -3,12 +3,10 @@ use pgt_console::{ markup, }; use pgt_diagnostics::PrintDiagnostic; -use pgt_test_utils::test_database::get_new_test_db; use pgt_typecheck::{TypecheckParams, check_sql}; +use sqlx::{Executor, PgPool}; -async fn test(name: &str, query: &str, setup: Option<&str>) { - let test_db = get_new_test_db().await; - +async fn test(name: &str, query: &str, setup: Option<&str>, test_db: &PgPool) { if let Some(setup) = setup { test_db .execute(setup) @@ -57,8 +55,8 @@ async fn test(name: &str, query: &str, setup: Option<&str>) { }); } -#[tokio::test] -async fn invalid_column() { +#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] +async fn invalid_column(pool: PgPool) { test( "invalid_column", "select id, unknown from contacts;", @@ -72,6 +70,7 @@ async fn invalid_column() { ); "#, ), + &pool, ) .await; } From 9cd04ddb11571770ce7b4e0f925e7035b7bd1c56 Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 29 May 2025 10:19:02 +0200 Subject: [PATCH 11/23] ok --- crates/pgt_completions/src/test_helper.rs | 2 +- crates/pgt_typecheck/tests/diagnostics.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/pgt_completions/src/test_helper.rs b/crates/pgt_completions/src/test_helper.rs index 61decce0..1bd5229c 100644 --- a/crates/pgt_completions/src/test_helper.rs +++ b/crates/pgt_completions/src/test_helper.rs @@ -44,7 +44,7 @@ pub(crate) async fn get_test_deps( .expect("Failed to execute setup query"); } - let schema_cache = SchemaCache::load(&test_db) + let schema_cache = SchemaCache::load(test_db) .await .expect("Failed to load Schema Cache"); diff --git a/crates/pgt_typecheck/tests/diagnostics.rs b/crates/pgt_typecheck/tests/diagnostics.rs index 9bf5d786..a7448503 100644 --- a/crates/pgt_typecheck/tests/diagnostics.rs +++ b/crates/pgt_typecheck/tests/diagnostics.rs @@ -19,7 +19,7 @@ async fn test(name: &str, query: &str, setup: Option<&str>, test_db: &PgPool) { .set_language(tree_sitter_sql::language()) .expect("Error loading sql language"); - let schema_cache = pgt_schema_cache::SchemaCache::load(&test_db) + let schema_cache = pgt_schema_cache::SchemaCache::load(test_db) .await .expect("Failed to load Schema Cache"); From 3fab2bf82869818cb50dd53af7e7247d0673f858 Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 29 May 2025 10:57:00 +0200 Subject: [PATCH 12/23] rename --- .../{0001-setup-roles.sql => 0001_setup-roles.sql} | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) rename crates/pgt_test_utils/testdb_migrations/{0001-setup-roles.sql => 0001_setup-roles.sql} (97%) diff --git a/crates/pgt_test_utils/testdb_migrations/0001-setup-roles.sql b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql similarity index 97% rename from crates/pgt_test_utils/testdb_migrations/0001-setup-roles.sql rename to crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql index 67e6dfec..109466ca 100644 --- a/crates/pgt_test_utils/testdb_migrations/0001-setup-roles.sql +++ b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql @@ -1,5 +1,6 @@ do $$ begin + if not exists ( select from pg_catalog.pg_roles where rolname = 'admin' @@ -20,4 +21,6 @@ if not exists ( ) then create role test_nologin; end if; -end; \ No newline at end of file + +end +$$; \ No newline at end of file From c270f6b7e58ec86948f63e96aa971e56cd96a922 Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 29 May 2025 11:16:04 +0200 Subject: [PATCH 13/23] fix test --- crates/pgt_schema_cache/src/policies.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/pgt_schema_cache/src/policies.rs b/crates/pgt_schema_cache/src/policies.rs index b754725b..af5e72d4 100644 --- a/crates/pgt_schema_cache/src/policies.rs +++ b/crates/pgt_schema_cache/src/policies.rs @@ -125,10 +125,10 @@ mod tests { owner_id int not null ); - create policy owner_policy + create policy test_nologin_policy on real_estate.properties for update - to owner + to test_nologin using (owner_id = current_user::int); "#; @@ -186,13 +186,13 @@ mod tests { let owner_policy = cache .policies .iter() - .find(|p| p.name == "owner_policy") + .find(|p| p.name == "test_nologin_policy") .unwrap(); assert_eq!(owner_policy.table_name, "properties"); assert_eq!(owner_policy.schema_name, "real_estate"); assert!(owner_policy.is_permissive); assert_eq!(owner_policy.command, PolicyCommand::Update); - assert_eq!(owner_policy.role_names, vec!["owner"]); + assert_eq!(owner_policy.role_names, vec!["test_nologin"]); assert_eq!( owner_policy.security_qualification, Some("(owner_id = (CURRENT_USER)::integer)".into()) From 5500285718d918b4bbb3edaff5af0fa3c69a306c Mon Sep 17 00:00:00 2001 From: Julian Date: Sat, 31 May 2025 08:56:54 +0200 Subject: [PATCH 14/23] does locking fix it? --- crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql index 109466ca..3d92c6b5 100644 --- a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql +++ b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql @@ -1,6 +1,8 @@ do $$ begin +select pg_advisory_lock(12345); + if not exists ( select from pg_catalog.pg_roles where rolname = 'admin' @@ -22,5 +24,7 @@ if not exists ( create role test_nologin; end if; +select pg_advisory_unlock(12345); + end $$; \ No newline at end of file From 9a4f099b03751100483d697d9f3b959bd64b61bb Mon Sep 17 00:00:00 2001 From: Julian Date: Sat, 31 May 2025 09:40:40 +0200 Subject: [PATCH 15/23] sp33d --- crates/pgt_schema_cache/src/roles.rs | 2 +- .../testdb_migrations/0001_setup-roles.sql | 43 +++++++++++-------- 2 files changed, 25 insertions(+), 20 deletions(-) diff --git a/crates/pgt_schema_cache/src/roles.rs b/crates/pgt_schema_cache/src/roles.rs index 636af34a..33187078 100644 --- a/crates/pgt_schema_cache/src/roles.rs +++ b/crates/pgt_schema_cache/src/roles.rs @@ -33,7 +33,7 @@ mod tests { let roles = &cache.roles; - let super_role = roles.iter().find(|r| r.name == "test_super").unwrap(); + let super_role = roles.iter().find(|r| r.name == "admin").unwrap(); assert!(super_role.is_super_user); assert!(super_role.can_create_db); assert!(super_role.can_login); diff --git a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql index 3d92c6b5..022b2111 100644 --- a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql +++ b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql @@ -1,30 +1,35 @@ do $$ begin -select pg_advisory_lock(12345); +if (select count(*) from pg_catalog.pg_roles where rolname in ('admin', 'test_login','test_nologin')) != 3 then -if not exists ( - select from pg_catalog.pg_roles - where rolname = 'admin' -) then - create role admin superuser createdb login bypassrls; -end if; + perform pg_advisory_lock(12345); -if not exists ( - select from pg_catalog.pg_roles - where rolname = 'test_login' -) then - create role test_login login; -end if; + if not exists ( + select from pg_catalog.pg_roles + where rolname = 'admin' + ) then + create role admin superuser createdb login bypassrls; + end if; + + if not exists ( + select from pg_catalog.pg_roles + where rolname = 'test_login' + ) then + create role test_login login; + end if; + + if not exists ( + select from pg_catalog.pg_roles + where rolname = 'test_nologin' + ) then + create role test_nologin; + end if; + + perform pg_advisory_unlock(12345); -if not exists ( - select from pg_catalog.pg_roles - where rolname = 'test_nologin' -) then - create role test_nologin; end if; -select pg_advisory_unlock(12345); end $$; \ No newline at end of file From 8e90a4992d7d6957229a1ca56d20d974198eb5f4 Mon Sep 17 00:00:00 2001 From: Julian Date: Sat, 31 May 2025 10:01:41 +0200 Subject: [PATCH 16/23] is it the name maybe --- crates/pgt_schema_cache/src/policies.rs | 22 +++++++++---------- crates/pgt_schema_cache/src/roles.rs | 2 +- .../testdb_migrations/0001_setup-roles.sql | 6 ++--- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/crates/pgt_schema_cache/src/policies.rs b/crates/pgt_schema_cache/src/policies.rs index af5e72d4..8e2ee4d7 100644 --- a/crates/pgt_schema_cache/src/policies.rs +++ b/crates/pgt_schema_cache/src/policies.rs @@ -112,10 +112,10 @@ mod tests { to public with check (true); - create policy admin_policy + create policy owner_policy on public.users for all - to admin + to owner with check (true); create schema real_estate; @@ -170,18 +170,18 @@ mod tests { assert_eq!(public_policy.security_qualification, Some("true".into())); assert_eq!(public_policy.with_check, None); - let admin_policy = cache + let owner_policy = cache .policies .iter() - .find(|p| p.name == "admin_policy") + .find(|p| p.name == "owner_policy") .unwrap(); - assert_eq!(admin_policy.table_name, "users"); - assert_eq!(admin_policy.schema_name, "public"); - assert!(admin_policy.is_permissive); - assert_eq!(admin_policy.command, PolicyCommand::All); - assert_eq!(admin_policy.role_names, vec!["admin"]); - assert_eq!(admin_policy.security_qualification, None); - assert_eq!(admin_policy.with_check, Some("true".into())); + assert_eq!(owner_policy.table_name, "users"); + assert_eq!(owner_policy.schema_name, "public"); + assert!(owner_policy.is_permissive); + assert_eq!(owner_policy.command, PolicyCommand::All); + assert_eq!(owner_policy.role_names, vec!["owner"]); + assert_eq!(owner_policy.security_qualification, None); + assert_eq!(owner_policy.with_check, Some("true".into())); let owner_policy = cache .policies diff --git a/crates/pgt_schema_cache/src/roles.rs b/crates/pgt_schema_cache/src/roles.rs index 33187078..7ced66f9 100644 --- a/crates/pgt_schema_cache/src/roles.rs +++ b/crates/pgt_schema_cache/src/roles.rs @@ -33,7 +33,7 @@ mod tests { let roles = &cache.roles; - let super_role = roles.iter().find(|r| r.name == "admin").unwrap(); + let super_role = roles.iter().find(|r| r.name == "owner").unwrap(); assert!(super_role.is_super_user); assert!(super_role.can_create_db); assert!(super_role.can_login); diff --git a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql index 022b2111..85e22cb5 100644 --- a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql +++ b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql @@ -1,15 +1,15 @@ do $$ begin -if (select count(*) from pg_catalog.pg_roles where rolname in ('admin', 'test_login','test_nologin')) != 3 then +if (select count(*) from pg_catalog.pg_roles where rolname in ('owner', 'test_login','test_nologin')) != 3 then perform pg_advisory_lock(12345); if not exists ( select from pg_catalog.pg_roles - where rolname = 'admin' + where rolname = 'owner' ) then - create role admin superuser createdb login bypassrls; + create role owner superuser createdb login bypassrls; end if; if not exists ( From 1de5f4ad5a7639958f3ad15ccfaf0ce9b06da873 Mon Sep 17 00:00:00 2001 From: Julian Date: Sat, 31 May 2025 10:59:15 +0200 Subject: [PATCH 17/23] try cmd --- .github/workflows/pull_request.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index f79392b7..8aa24265 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -157,7 +157,10 @@ jobs: # running containers via `services` only works on linux # https://github.com/actions/runner/issues/1866 - name: Setup postgres + id: postgres uses: ikalnytskyi/action-setup-postgres@v7 + - name: Print Roles + run: psql ${{ steps.postgres.outputs.connection-uri }} -c "select rolname from pg_roles;" - name: Run tests run: cargo test --workspace From ed3751e1102d3e9e5031099bb3dadb7e3789c979 Mon Sep 17 00:00:00 2001 From: Julian Date: Sat, 31 May 2025 11:58:53 +0200 Subject: [PATCH 18/23] this probably fixes it --- .../testdb_migrations/0001_setup-roles.sql | 42 +++++++------------ 1 file changed, 15 insertions(+), 27 deletions(-) diff --git a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql index 85e22cb5..840a1829 100644 --- a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql +++ b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql @@ -1,35 +1,23 @@ do $$ begin -if (select count(*) from pg_catalog.pg_roles where rolname in ('owner', 'test_login','test_nologin')) != 3 then - - perform pg_advisory_lock(12345); - - if not exists ( - select from pg_catalog.pg_roles - where rolname = 'owner' - ) then - create role owner superuser createdb login bypassrls; - end if; - - if not exists ( - select from pg_catalog.pg_roles - where rolname = 'test_login' - ) then - create role test_login login; - end if; - - if not exists ( - select from pg_catalog.pg_roles - where rolname = 'test_nologin' - ) then - create role test_nologin; - end if; - - perform pg_advisory_unlock(12345); +begin + create role owner superuser createdb login bypassrls; +exception when duplicate_object then + null; +end; -end if; +begin + create role test_login login; +exception when duplicate_object then + null; +end; +begin + create role test_nologin; +exception when duplicate_object then + null; +end; end $$; \ No newline at end of file From 4d854c5075003deb9c834dd5431a0dc51030d7b9 Mon Sep 17 00:00:00 2001 From: Julian Date: Sat, 31 May 2025 13:36:31 +0200 Subject: [PATCH 19/23] maybe its the cache --- .github/workflows/pull_request.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index 8aa24265..e0f633ae 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -149,8 +149,8 @@ jobs: uses: ./.github/actions/free-disk-space - name: Install toolchain uses: moonrepo/setup-rust@v1 - with: - cache-base: main + # with: + # cache-base: main env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} From e02a8292c90f81f350d1fcf8394cb5fafc463f01 Mon Sep 17 00:00:00 2001 From: Julian Date: Sat, 31 May 2025 16:12:00 +0200 Subject: [PATCH 20/23] verify that changing squeel ha an effect --- .../testdb_migrations/0001_setup-roles.sql | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql index 840a1829..514643d8 100644 --- a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql +++ b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql @@ -1,23 +1,23 @@ do $$ begin -begin - create role owner superuser createdb login bypassrls; -exception when duplicate_object then - null; -end; +-- begin +-- create role owner superuser createdb login bypassrls; +-- exception when duplicate_object then +-- null; +-- end; -begin - create role test_login login; -exception when duplicate_object then - null; -end; +-- begin +-- create role test_login login; +-- exception when duplicate_object then +-- null; +-- end; -begin - create role test_nologin; -exception when duplicate_object then - null; -end; +-- begin +-- create role test_nologin; +-- exception when duplicate_object then +-- null; +-- end; end $$; \ No newline at end of file From d23a67827a96c13ace065f6677f63a39effbe865 Mon Sep 17 00:00:00 2001 From: Julian Date: Sat, 31 May 2025 16:18:20 +0200 Subject: [PATCH 21/23] ok --- .../testdb_migrations/0001_setup-roles.sql | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql index 514643d8..840a1829 100644 --- a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql +++ b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql @@ -1,23 +1,23 @@ do $$ begin --- begin --- create role owner superuser createdb login bypassrls; --- exception when duplicate_object then --- null; --- end; +begin + create role owner superuser createdb login bypassrls; +exception when duplicate_object then + null; +end; --- begin --- create role test_login login; --- exception when duplicate_object then --- null; --- end; +begin + create role test_login login; +exception when duplicate_object then + null; +end; --- begin --- create role test_nologin; --- exception when duplicate_object then --- null; --- end; +begin + create role test_nologin; +exception when duplicate_object then + null; +end; end $$; \ No newline at end of file From 7e29755866cfa047fd97564029ddd8ec97c04891 Mon Sep 17 00:00:00 2001 From: Julian Date: Sat, 31 May 2025 17:02:14 +0200 Subject: [PATCH 22/23] ok then --- .github/workflows/pull_request.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index e0f633ae..8aa24265 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -149,8 +149,8 @@ jobs: uses: ./.github/actions/free-disk-space - name: Install toolchain uses: moonrepo/setup-rust@v1 - # with: - # cache-base: main + with: + cache-base: main env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} From fb1537a62d3166bb1716a0bf30c180c6eafb8b9d Mon Sep 17 00:00:00 2001 From: Julian Date: Sat, 31 May 2025 17:13:57 +0200 Subject: [PATCH 23/23] omg --- .../testdb_migrations/0001_setup-roles.sql | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql index 840a1829..1f1d50b3 100644 --- a/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql +++ b/crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql @@ -3,20 +3,29 @@ begin begin create role owner superuser createdb login bypassrls; -exception when duplicate_object then - null; +exception + when duplicate_object then + null; + when unique_violation then + null; end; begin create role test_login login; -exception when duplicate_object then - null; +exception + when duplicate_object then + null; + when unique_violation then + null; end; begin create role test_nologin; -exception when duplicate_object then - null; +exception + when duplicate_object then + null; + when unique_violation then + null; end; end