qe: use one db per test on cockroachdb

.github/workflows/query-engine.yml

@@ -10,6 +10,7 @@ on:
       - 'LICENSE'
       - 'CODEOWNERS'
       - 'renovate.json'
+      - 'introspection-engine/**'
 
 jobs:
   rust-vitess-tests:
@@ -35,7 +36,10 @@ jobs:
             single_threaded: false
             connector: "sqlserver"
             version: "2022"
-
+          - name: "cockroach_22_1_0"
+            single_threaded: false
+            connector: "cockroachdb"
+            version: "22.1.0"
     env:
       LOG_LEVEL: "info"
       LOG_QUERIES: "y"
@@ -49,7 +53,7 @@ jobs:
       TEST_CONNECTOR: ${{ matrix.database.connector }}
       TEST_CONNECTOR_VERSION: ${{ matrix.database.version }}
 
-    runs-on: buildjet-16vcpu-ubuntu-2004
+    runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
 

docker-compose.yml

@@ -3,11 +3,13 @@ services:
   cockroach_22_1_0:
     image: prismagraphql/cockroachdb-custom:22.1.0
     restart: always
-    command: start-single-node --insecure
+    command: start-single-node --insecure --store=type=mem,size=3GiB
     ports:
       - "26257:26257"
     networks:
       - databases
+    tmpfs: /cockroach/cockroach-data
+
 
   cockroach_21_2_0_patched:
     image: prismagraphql/cockroachdb-custom:21.2.0-patched

migration-engine/qe-setup/Cargo.toml

@@ -16,3 +16,5 @@ mongodb = "2.3.0"
 tempfile = "3.3.0"
 url = "2"
 quaint.workspace = true
+once_cell = "1.17.0"
+crossbeam-channel = "0.5.6"

migration-engine/qe-setup/src/cockroachdb.rs

@@ -0,0 +1,75 @@
+use migration_core::migration_connector::{ConnectorError, ConnectorResult};
+use quaint::{prelude::*, single::Quaint};
+use url::Url;
+
+pub(crate) async fn cockroach_setup(url: String, prisma_schema: &str) -> ConnectorResult<()> {
+    let mut url = Url::parse(&url).map_err(ConnectorError::url_parse_error)?;
+    let quaint_url = quaint::connector::PostgresUrl::new(url.clone()).unwrap();
+    let db_name = quaint_url.dbname();
+    let conn = create_admin_conn(&mut url).await?;
+
+    let query = format!(
+        r#"
+        DROP DATABASE IF EXISTS "{db_name}";
+        CREATE DATABASE "{db_name}";
+        "#
+    );
+
+    conn.raw_cmd(&query).await.unwrap();
+    crate::diff_and_apply(prisma_schema).await;
+
+    drop_db_when_thread_exits(url, db_name);
+
+    Ok(())
+}
+
+async fn create_admin_conn(url: &mut Url) -> ConnectorResult<Quaint> {
+    url.set_path("/postgres");
+    Ok(Quaint::new(url.as_ref()).await.unwrap())
+}
+
+fn drop_db_when_thread_exits(admin_url: Url, db_name: &str) {
+    use crossbeam_channel::*;
+    use once_cell::sync::OnceCell;
+    use std::{cell::RefCell, thread};
+    use test_setup::runtime::run_with_thread_local_runtime as tok;
+
+    // === Dramatis Personæ ===
+
+    // DB_DROP_THREAD: A thread that drops databases.
+    static DB_DROP_THREAD: OnceCell<Sender<String>> = OnceCell::new();
+
+    let sender = DB_DROP_THREAD.get_or_init(|| {
+        let (sender, receiver) = unbounded::<String>();
+
+        thread::spawn(move || {
+            let mut admin_url = admin_url;
+            let conn = tok(create_admin_conn(&mut admin_url)).unwrap();
+            loop {
+                // Receive new databases to drop.
+                for msg in receiver.iter() {
+                    tok(conn.raw_cmd(&msg)).unwrap();
+                }
+            }
+        });
+
+        sender
+    });
+
+    // NOTIFIER: a thread local that notifies DB_DROP_THREAD when dropped.
+    struct Notifier(String, Sender<String>);
+
+    impl Drop for Notifier {
+        fn drop(&mut self) {
+            self.1.send(std::mem::take(&mut self.0)).unwrap()
+        }
+    }
+
+    thread_local! {
+        static NOTIFIER: RefCell<Option<Notifier>> = RefCell::new(None);
+    }
+
+    NOTIFIER.with(move |cell| {
+        *cell.borrow_mut() = Some(Notifier(format!("DROP DATABASE \"{db_name}\""), sender.clone()));
+    });
+}

migration-engine/qe-setup/src/lib.rs

@@ -1,13 +1,16 @@
 //! Query Engine test setup.
+
 #![allow(clippy::await_holding_lock)]
+
+mod cockroachdb;
 mod mongodb;
 mod mssql;
 mod mysql;
 mod postgres;
 
 pub use migration_core::migration_connector::ConnectorError;
 
-use self::{mongodb::*, mssql::*, mysql::*, postgres::*};
+use self::{cockroachdb::*, mongodb::*, mssql::*, mysql::*, postgres::*};
 use enumflags2::BitFlags;
 use migration_core::{
     json_rpc::types::*,
@@ -40,9 +43,10 @@ pub async fn setup(prisma_schema: &str, db_schemas: &[&str]) -> ConnectorResult<
     let (source, url, _preview_features) = parse_configuration(prisma_schema)?;
 
     match &source.active_provider {
-        provider if [POSTGRES.provider_name(), COCKROACH.provider_name()].contains(provider) => {
+        provider if [POSTGRES.provider_name()].contains(provider) => {
             postgres_setup(url, prisma_schema, db_schemas).await?
         }
+        provider if COCKROACH.is_provider(provider) => cockroach_setup(url, prisma_schema).await?,
         provider if MSSQL.is_provider(provider) => mssql_setup(url, prisma_schema, db_schemas).await?,
         provider if MYSQL.is_provider(provider) => {
             mysql_reset(&url).await?;
@@ -66,7 +70,7 @@ pub async fn teardown(prisma_schema: &str, db_schemas: &[&str]) -> ConnectorResu
     let (source, url, _) = parse_configuration(prisma_schema)?;
 
     match &source.active_provider {
-        provider if [POSTGRES.provider_name(), COCKROACH.provider_name()].contains(provider) => {
+        provider if [POSTGRES.provider_name()].contains(provider) => {
             postgres_teardown(&url, db_schemas).await?;
         }
 
@@ -76,6 +80,7 @@ pub async fn teardown(prisma_schema: &str, db_schemas: &[&str]) -> ConnectorResu
                 MSSQL.provider_name(),
                 MYSQL.provider_name(),
                 MONGODB.provider_name(),
+                COCKROACH.provider_name(),
             ]
             .contains(provider) => {}
 

query-engine/connector-test-kit-rs/query-engine-tests/tests/new/metrics.rs

@@ -30,6 +30,7 @@ mod metrics {
             Sqlite => assert_eq!(total_queries, 9),
             SqlServer(_) => assert_eq!(total_queries, 15),
             MongoDb(_) => assert_eq!(total_queries, 5),
+            ConnectorVersion::CockroachDb => assert_eq!(total_queries, 10),
             _ => assert_eq!(total_queries, 11),
         }
         assert_eq!(total_operations, 2);