qe: Capabilities-based size optimizations for WASM engine (#4701)

* qe: Capabilities-based size optimizations for WASM engine

Core feature of this PR: `can_have_capability` function in `psl-core`
crate. For multi-connector builds it always returns `true` and has no
effect. If we are inside of a single connector build, however, it will
eveluate value based on the actual connector capabilities. In many
cases, this will allow optimizer to completely eliminate following code
if connector does not support specific feature.

Notable exeception to that were relation joins: for some reason,
optimizer can not eliminate those functions from SQLite bundle. So, this
is the only place in sql-query-connector where we have to introduce
conditionally compiled feature.

In order to take maximum advantage of this functionality, we have to
disable default features in `quaint` and `psl` crates and enable them
only in native engines.

Close prisma/team-orm#928

* Fix tests & lints

* Fix migrate tests

* Fix codspeed

* Now for real

* FFS

* Address some of the feedback

* Introduce reachble_only_with_capability macro

* Compile `returning` uncoditionally

* Restore benchmarks

* Fix benchmarks, but preserve all size gains

* Sigh

.github/workflows/codspeed.yml

@@ -28,10 +28,10 @@ jobs:
         run: cargo install cargo-codspeed
 
       - name: "Build the benchmark targets: schema"
-        run: cargo codspeed build -p schema
+        run: cargo codspeed build -p schema --features all_connectors
 
       - name: "Build the benchmark targets: request-handlers"
-        run: cargo codspeed build -p request-handlers
+        run: cargo codspeed build -p request-handlers --features native
 
       - name: Run the benchmarks
         uses: CodSpeedHQ/action@v2

Cargo.toml

@@ -76,16 +76,6 @@ url = { version = "2.5.0" }
 
 [workspace.dependencies.quaint]
 path = "quaint"
-features = [
-  "expose-drivers",
-  "fmt-sql",
-  "mssql",
-  "mysql",
-  "pooled",
-  "postgresql",
-  "sqlite",
-  "native",
-]
 
 [profile.dev.package.backtrace]
 opt-level = 3

nix/shell.nix

@@ -7,7 +7,8 @@ in
 {
   devShells.default = pkgs.mkShell {
     packages = with pkgs; [
-      devToolchain
+      # devToolchain
+      rustup
       llvmPackages_latest.bintools
 
       nodejs_20

psl/psl-core/Cargo.toml

@@ -3,6 +3,14 @@ edition = "2021"
 name = "psl-core"
 version = "0.1.0"
 
+[features]
+postgresql = []
+sqlite = []
+mysql = []
+cockroachdb = []
+mssql = []
+mongodb = []
+
 [dependencies]
 diagnostics = { path = "../diagnostics" }
 parser-database = { path = "../parser-database" }
@@ -25,3 +33,4 @@ hex = "0.4"
 # For the connector API.
 lsp-types = "0.91.1"
 url.workspace = true
+cfg-if = "1.0.0"

psl/psl-core/src/builtin_connectors/capabilities_support.rs

@@ -0,0 +1,93 @@
+use crate::datamodel_connector::{Connector, ConnectorCapabilities, ConnectorCapability};
+use cfg_if::cfg_if;
+
+cfg_if! {
+    // if built only for mysql
+    if #[cfg(all(feature="mysql", not(any(feature = "postgresql", feature="sqlite", feature = "cockroachdb", feature="mssql", feature="mongodb"))))] {
+        #[inline(always)]
+        const fn can_have_capability_impl(capability: ConnectorCapability) -> bool {
+            check_comptime_capability(super::mysql_datamodel_connector::CAPABILITIES, capability)
+        }
+
+        pub fn has_capability(_: &dyn Connector, capability: ConnectorCapability) -> bool {
+            can_have_capability_impl(capability)
+        }
+    // if built only for sqlite
+    } else if #[cfg(all(feature="sqlite", not(any(feature = "postgresql", feature="mysql", feature = "cockroachdb", feature="mssql", feature="mongodb"))))] {
+        #[inline(always)]
+        const fn can_have_capability_impl(capability: ConnectorCapability) -> bool {
+            check_comptime_capability(super::sqlite_datamodel_connector::CAPABILITIES, capability)
+        }
+
+        #[inline(always)]
+        pub fn has_capability(_: &dyn Connector, capability: ConnectorCapability) -> bool {
+            can_have_capability_impl(capability)
+        }
+    // if built only for postgresql
+    } else if #[cfg(all(feature="postgresql", not(any(feature = "sqlite", feature="mysql", feature = "cockroachdb", feature="mssql", feature="mongodb"))))] {
+        #[inline(always)]
+        const fn can_have_capability_impl(capability: ConnectorCapability) -> bool {
+            check_comptime_capability(super::postgres_datamodel_connector::CAPABILITIES, capability)
+        }
+
+        #[inline(always)]
+        pub fn has_capability(_: &dyn Connector, capability: ConnectorCapability) -> bool {
+            can_have_capability_impl(capability)
+        }
+    // any other build configuration
+    } else {
+        #[inline(always)]
+        const fn can_have_capability_impl(_: ConnectorCapability) -> bool {
+            true
+        }
+
+        #[inline(always)]
+        pub fn has_capability(connector: &dyn Connector, capability: ConnectorCapability) -> bool {
+            connector.capabilities().contains(capability)
+        }
+    }
+}
+
+/// Helper function for determining if engine, compiled with the current settings,
+/// can potentially have provided capability on. Useful for single-connector builds and can
+/// be used to exclude certain code that we know for sure can't be executed for current connector.
+/// Has no effect on multi-connector builds
+/// # Example
+/// ```ignore
+/// if !can_have_capability(ConnectorCapability::FullTextSearch) {
+///    unreachable!()
+/// }
+/// ... // if compiled for a single connector, optimizer will exclude the following code if connector does not support full text search
+/// ```
+#[inline(always)]
+pub const fn can_have_capability(cap: ConnectorCapability) -> bool {
+    can_have_capability_impl(cap)
+}
+
+/// Marks the code as reachable only by the connectors,
+/// having the specific capability.
+/// Optimizer usually can optimize the code away if none of the connectors
+/// current build supports the capability.
+///
+/// If we are within a single connector build that has no such capability,
+/// and the code marked with this macro is reached, it will panic.
+#[macro_export]
+macro_rules! reachable_only_with_capability {
+    ($cap: expr) => {
+        if !$crate::builtin_connectors::can_have_capability($cap) {
+            core::unreachable!()
+        }
+    };
+}
+
+#[inline(always)]
+#[allow(dead_code)] // not used if more than one connector is built
+const fn check_comptime_capability(capabilities: ConnectorCapabilities, cap: ConnectorCapability) -> bool {
+    (capabilities.bits_c() & (cap as u64)) > 0
+}
+
+#[inline(always)]
+pub const fn can_support_relation_load_strategy() -> bool {
+    can_have_capability(ConnectorCapability::LateralJoin)
+        || can_have_capability(ConnectorCapability::CorrelatedSubqueries)
+}

psl/psl-core/src/builtin_connectors/completions.rs

@@ -1,17 +1,14 @@
-use crate::datamodel_connector::format_completion_docs;
-use lsp_types::{
-    CompletionItem, CompletionItemKind, CompletionList, Documentation, InsertTextFormat, MarkupContent, MarkupKind,
-};
-
-pub(crate) fn extensions_completion(completion_list: &mut CompletionList) {
+#[cfg(feature = "postgresql")]
+pub(crate) fn extensions_completion(completion_list: &mut lsp_types::CompletionList) {
+    use lsp_types::*;
     completion_list.items.push(CompletionItem {
         label: "extensions".to_owned(),
         insert_text: Some("extensions = [$0]".to_owned()),
         insert_text_format: Some(InsertTextFormat::SNIPPET),
         kind: Some(CompletionItemKind::FIELD),
         documentation: Some(Documentation::MarkupContent(MarkupContent {
             kind: MarkupKind::Markdown,
-            value: format_completion_docs(
+            value: crate::datamodel_connector::format_completion_docs(
                 r#"extensions = [pg_trgm, postgis(version: "2.1")]"#,
                 r#"Enable PostgreSQL extensions. [Learn more](https://pris.ly/d/postgresql-extensions)"#,
                 None,
@@ -21,15 +18,22 @@ pub(crate) fn extensions_completion(completion_list: &mut CompletionList) {
     })
 }
 
-pub(crate) fn schemas_completion(completion_list: &mut CompletionList) {
+#[cfg(any(
+    feature = "postgresql",
+    feature = "cockroachdb",
+    feature = "mssql",
+    feature = "mysql"
+))]
+pub(crate) fn schemas_completion(completion_list: &mut lsp_types::CompletionList) {
+    use lsp_types::*;
     completion_list.items.push(CompletionItem {
         label: "schemas".to_owned(),
         insert_text: Some(r#"schemas = [$0]"#.to_owned()),
         insert_text_format: Some(InsertTextFormat::SNIPPET),
         kind: Some(CompletionItemKind::FIELD),
         documentation: Some(Documentation::MarkupContent(MarkupContent {
             kind: MarkupKind::Markdown,
-            value: format_completion_docs(
+            value: crate::datamodel_connector::format_completion_docs(
                 r#"schemas = ["foo", "bar", "baz"]"#,
                 "The list of database schemas. [Learn More](https://pris.ly/d/multi-schema-configuration)",
                 None,

psl/psl-core/src/builtin_connectors/mod.rs

@@ -1,27 +1,64 @@
+#[cfg(feature = "cockroachdb")]
 pub mod cockroach_datamodel_connector;
 pub mod completions;
 
+#[cfg(feature = "cockroachdb")]
 pub use cockroach_datamodel_connector::CockroachType;
+#[cfg(feature = "mongodb")]
 pub use mongodb::MongoDbType;
+#[cfg(feature = "mssql")]
 pub use mssql_datamodel_connector::{MsSqlType, MsSqlTypeParameter};
+#[cfg(feature = "mysql")]
 pub use mysql_datamodel_connector::MySqlType;
+#[cfg(feature = "postgresql")]
 pub use postgres_datamodel_connector::{PostgresDatasourceProperties, PostgresType};
 
+mod capabilities_support;
+#[cfg(feature = "mongodb")]
 mod mongodb;
+#[cfg(feature = "mssql")]
 mod mssql_datamodel_connector;
+#[cfg(feature = "mysql")]
 mod mysql_datamodel_connector;
 mod native_type_definition;
+#[cfg(feature = "postgresql")]
 mod postgres_datamodel_connector;
+#[cfg(feature = "sqlite")]
 mod sqlite_datamodel_connector;
 mod utils;
+pub use capabilities_support::{can_have_capability, can_support_relation_load_strategy, has_capability};
 
-use crate::{datamodel_connector::Connector, ConnectorRegistry};
+use crate::ConnectorRegistry;
 
-pub const POSTGRES: &'static dyn Connector = &postgres_datamodel_connector::PostgresDatamodelConnector;
-pub const COCKROACH: &'static dyn Connector = &cockroach_datamodel_connector::CockroachDatamodelConnector;
-pub const MYSQL: &'static dyn Connector = &mysql_datamodel_connector::MySqlDatamodelConnector;
-pub const SQLITE: &'static dyn Connector = &sqlite_datamodel_connector::SqliteDatamodelConnector;
-pub const MSSQL: &'static dyn Connector = &mssql_datamodel_connector::MsSqlDatamodelConnector;
-pub const MONGODB: &'static dyn Connector = &mongodb::MongoDbDatamodelConnector;
+#[cfg(feature = "postgresql")]
+pub const POSTGRES: &'static dyn crate::datamodel_connector::Connector =
+    &postgres_datamodel_connector::PostgresDatamodelConnector;
+#[cfg(feature = "cockroachdb")]
+pub const COCKROACH: &'static dyn crate::datamodel_connector::Connector =
+    &cockroach_datamodel_connector::CockroachDatamodelConnector;
+#[cfg(feature = "mysql")]
+pub const MYSQL: &'static dyn crate::datamodel_connector::Connector =
+    &mysql_datamodel_connector::MySqlDatamodelConnector;
+#[cfg(feature = "sqlite")]
+pub const SQLITE: &'static dyn crate::datamodel_connector::Connector =
+    &sqlite_datamodel_connector::SqliteDatamodelConnector;
+#[cfg(feature = "mssql")]
+pub const MSSQL: &'static dyn crate::datamodel_connector::Connector =
+    &mssql_datamodel_connector::MsSqlDatamodelConnector;
+#[cfg(feature = "mongodb")]
+pub const MONGODB: &'static dyn crate::datamodel_connector::Connector = &mongodb::MongoDbDatamodelConnector;
 
-pub static BUILTIN_CONNECTORS: ConnectorRegistry<'static> = &[POSTGRES, MYSQL, SQLITE, MSSQL, COCKROACH, MONGODB];
+pub static BUILTIN_CONNECTORS: ConnectorRegistry<'static> = &[
+    #[cfg(feature = "postgresql")]
+    POSTGRES,
+    #[cfg(feature = "mysql")]
+    MYSQL,
+    #[cfg(feature = "sqlite")]
+    SQLITE,
+    #[cfg(feature = "mssql")]
+    MSSQL,
+    #[cfg(feature = "cockroachdb")]
+    COCKROACH,
+    #[cfg(feature = "mongodb")]
+    MONGODB,
+];

psl/psl-core/src/builtin_connectors/mysql_datamodel_connector.rs

@@ -28,7 +28,7 @@ const TEXT_TYPE_NAME: &str = "Text";
 const MEDIUM_TEXT_TYPE_NAME: &str = "MediumText";
 const LONG_TEXT_TYPE_NAME: &str = "LongText";
 
-const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(ConnectorCapability::{
+pub const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(ConnectorCapability::{
     Enums |
     EnumArrayPush |
     Json |
@@ -55,6 +55,7 @@ const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(Connector
     MultiSchema |
 
     FullTextIndex |
+    FullTextSearch |
     FullTextSearchWithIndex |
     MultipleFullTextAttributesPerModel |
     ImplicitManyToManyRelation |

psl/psl-core/src/builtin_connectors/postgres_datamodel_connector.rs

@@ -26,7 +26,7 @@ const CONSTRAINT_SCOPES: &[ConstraintScope] = &[
     ConstraintScope::ModelPrimaryKeyKeyIndexForeignKey,
 ];
 
-const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(ConnectorCapability::{
+pub const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(ConnectorCapability::{
     AdvancedJsonNullability |
     AnyId |
     AutoIncrement |
@@ -39,6 +39,7 @@ const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(Connector
     CreateSkipDuplicates |
     Enums |
     EnumArrayPush |
+    FullTextSearch |
     FullTextSearchWithoutIndex |
     InsensitiveFilters |
     Json |

psl/psl-core/src/builtin_connectors/sqlite_datamodel_connector.rs

@@ -10,7 +10,7 @@ use enumflags2::BitFlags;
 
 const NATIVE_TYPE_CONSTRUCTORS: &[NativeTypeConstructor] = &[];
 const CONSTRAINT_SCOPES: &[ConstraintScope] = &[ConstraintScope::GlobalKeyIndex];
-const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(ConnectorCapability::{
+pub const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(ConnectorCapability::{
     AnyId |
     AutoIncrement |
     CompoundIds |

psl/psl-core/src/builtin_connectors/utils.rs

@@ -1,3 +1,4 @@
+#[cfg(any(feature = "postgresql", feature = "mysql"))]
 pub(crate) mod common {
     use chrono::*;
 
@@ -25,6 +26,7 @@ pub(crate) mod common {
     }
 }
 
+#[cfg(feature = "postgresql")]
 pub(crate) mod postgres {
     use chrono::*;
 
@@ -50,6 +52,7 @@ pub(crate) mod postgres {
     }
 }
 
+#[cfg(feature = "mysql")]
 pub(crate) mod mysql {
     use chrono::*;