custom type mapping

.sqlxrc.sample.json

@@ -14,6 +14,21 @@
       "DB_NAME": "postgres",
       "PG_SEARCH_PATH": "public,myschema"
     },
+    "db_custom_types": {
+      "DB_TYPE": "postgres",
+      "DB_HOST": "127.0.0.1",
+      "DB_PORT": 54321,
+      "DB_USER": "postgres",
+      "DB_PASS": "postgres",
+      "DB_NAME": "postgres",
+      "PG_SEARCH_PATH": "public,myschema",
+      "type_mapping": {
+        "bigint": "string",
+        "interval": "string",
+        "timestamp": { "type": "DateTime", "import": "import type { DateTime } from 'luxon'" },
+        "timestamptz": { "type": "Dayjs", "import": "import type { Dayjs } from 'dayjs'" }
+      }
+    },
     "db_mysql": {
       "DB_TYPE": "mysql",
       "DB_HOST": "127.0.0.1",

Cargo.toml

@@ -29,7 +29,7 @@ colored = "3.1.1"
 mysql_async = { version = "0.36.2", features = ["minimal"] }
 tokio-postgres = "0.7.17"
 tokio = { version = "1.52.1", features = ["rt-multi-thread", "macros", "default"]}
-async-recursion = "1.1.1"
+async-recursion = "1.2.0"
 bb8 = "0.9.1"
 rusqlite = { version = "0.39", features = ["bundled"] }
 log = "0.4.29"

book/docs/connect/config-file.md

@@ -116,6 +116,52 @@ Supported fields of each connection include
 - `PG_SEARCH_PATH`: PostgreSQL schema search path (default is "$user,public") [https://www.postgresql.org/docs/current/ddl-schemas.html#DDL-SCHEMAS-PATH](https://www.postgresql.org/docs/current/ddl-schemas.html#DDL-SCHEMAS-PATH)
 - `POOL_SIZE`: Size of the connection pool to establish per connection type
 - `CONNECTION_TIMEOUT`: Timeout in second of Database connection attempt
+- `type_mapping`: Overrides the generated TypeScript type of database column types (see below)
+
+#### type_mapping
+
+By default, SQLX-TS translates each database column type into a built-in TypeScript type (e.g. `bigint` -> `number`).
+You can override this per connection by mapping a database column type to any TypeScript type. A mapping is either
+a type, or an object with `type` and an optional `import`; the import statement is added at the top of the generated
+types file.
+
+```json
+{
+  "connections": {
+    "default": {
+      "DB_TYPE": "postgres",
+      "DB_URL": "postgres://postgres:postgres@127.0.0.1:5432/mydb",
+      "type_mapping": {
+        "bigint": "string",
+        "numeric": "string | number",
+        "_int8": "string[]",
+        "timestamp": { "type": "DateTime", "import": "import type { DateTime } from 'luxon'" }
+      }
+    }
+  }
+}
+```
+
+**Matching**
+
+- Keys are matched case-insensitively against the column type reported by the database, e.g. `bigint`, `DATETIME`
+- An exact match wins, otherwise type modifiers are ignored, so `varchar` matches `VARCHAR(255)` and `bigint unsigned` matches `bigint(20) unsigned`
+- PostgreSQL: both the `data_type` (e.g. `timestamp without time zone`) and the `udt_name` (e.g. `timestamp`, `int8`, or the name of an enum or extension type such as `citext`) are matched. Domain types are matched by their underlying type
+- PostgreSQL arrays: array columns are reported as `ARRAY`, so the mapping of the element type does not apply. Map the array's `udt_name` instead, which is the element type prefixed with `_`, e.g. `"_int8": "string[]"`
+- MySQL: both the full `COLUMN_TYPE` (e.g. `tinyint(1)`, `bigint unsigned`) and the `DATA_TYPE` (e.g. `tinyint`, `bigint`) are matched, with `COLUMN_TYPE` taking priority. This lets you map `"tinyint(1)": "boolean"` separately from `"tinyint": "number"`
+
+**Generated types**
+
+- Nullable columns still produce `| null`, e.g. `string | null`
+- The mapping applies to both query results and parameters
+- `@result` and `@param` [annotations](../type-generation/annotations.md) take priority over the mapping
+- The mapping only applies to values typed from a table column, including expressions derived from one such as `COALESCE(col, 0)`. Types from the SQL itself, such as `CAST(col AS BIGINT)` or column definitions of a table-valued function (`jsonb_to_recordset($1) AS t(id BIGINT)`), are not mapped
+
+**Imports**
+
+- Each import is written once per generated file, including when all types are generated into a single file with `--generate-path`
+- Multi-line imports are written on a single line
+- Imports are de-duplicated by their text, so use the same import statement for every mapping that imports the same type. For example, `import type { DateTime } from 'luxon'` and `import type { DateTime } from "luxon"` are written as two imports, which TypeScript reports as a duplicate identifier
 
 ### generate_types
 

playpen/db/mysql_migration.sql

@@ -223,3 +223,11 @@ INSERT INTO json_test_data (name, data, metadata) VALUES
 ('analytics',
  '{"date": "2024-01-15", "metrics": {"visitors": 1500, "pageViews": 4500, "bounceRate": 0.35, "sources": {"organic": 850, "direct": 400, "referral": 250}}}',
  '{"source": "analytics", "version": "1.0"}');
+
+CREATE TABLE type_mapping_test (
+  id INT PRIMARY KEY,
+  is_active TINYINT(1) NOT NULL,
+  flags TINYINT NOT NULL,
+  big_id BIGINT UNSIGNED NOT NULL,
+  amount BIGINT NOT NULL
+);

playpen/db/mysql_migration_5_6.sql

@@ -209,3 +209,11 @@ INSERT INTO json_test_data (name, data, metadata) VALUES
 ('analytics',
  '{"date": "2024-01-15", "metrics": {"visitors": 1500, "pageViews": 4500, "bounceRate": 0.35, "sources": {"organic": 850, "direct": 400, "referral": 250}}}',
  '{"source": "analytics", "version": "1.0"}');
+
+CREATE TABLE type_mapping_test (
+  id INT PRIMARY KEY,
+  is_active TINYINT(1) NOT NULL,
+  flags TINYINT NOT NULL,
+  big_id BIGINT UNSIGNED NOT NULL,
+  amount BIGINT NOT NULL
+);

src/common/config.rs

@@ -7,6 +7,7 @@ use regex::Regex;
 use serde;
 use serde::{Deserialize, Serialize};
 use serde_json;
+use serde_json::Value as JsonValue;
 use std::collections::HashMap;
 use std::fs;
 use std::path::PathBuf;
@@ -20,6 +21,69 @@ pub struct SqlxConfig {
   pub connections: HashMap<String, DbConnectionConfig>,
 }
 
+#[derive(Clone, Debug)]
+pub enum CustomTypeMapping {
+  Simple(String),
+  WithImport { type_name: String, import: String },
+}
+
+impl<'de> Deserialize<'de> for CustomTypeMapping {
+  fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+  where
+    D: serde::Deserializer<'de>,
+  {
+    let value = JsonValue::deserialize(deserializer)?;
+    let non_empty = |type_name: &str| -> Result<String, D::Error> {
+      let type_name = type_name.trim();
+      if type_name.is_empty() {
+        return Err(serde::de::Error::custom("type in type_mapping cannot be empty"));
+      }
+      Ok(type_name.to_string())
+    };
+    match value {
+      JsonValue::String(s) => Ok(CustomTypeMapping::Simple(non_empty(&s)?)),
+      JsonValue::Object(map) => {
+        let type_name = map
+          .get("type")
+          .and_then(|v| v.as_str())
+          .ok_or_else(|| serde::de::Error::missing_field("type"))?;
+        let type_name = non_empty(type_name)?;
+        let import = match map.get("import") {
+          None | Some(JsonValue::Null) => None,
+          Some(JsonValue::String(import)) if import.trim().is_empty() => None,
+          Some(JsonValue::String(import)) => Some(import.to_string()),
+          Some(_) => return Err(serde::de::Error::custom("import in type_mapping must be a string")),
+        };
+        match import {
+          Some(import) => Ok(CustomTypeMapping::WithImport { type_name, import }),
+          None => Ok(CustomTypeMapping::Simple(type_name)),
+        }
+      }
+      _ => Err(serde::de::Error::custom(
+        "Expected a string or an object for CustomTypeMapping",
+      )),
+    }
+  }
+}
+
+impl Serialize for CustomTypeMapping {
+  fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+  where
+    S: serde::Serializer,
+  {
+    match self {
+      CustomTypeMapping::Simple(s) => serializer.serialize_str(s),
+      CustomTypeMapping::WithImport { type_name, import } => {
+        use serde::ser::SerializeMap;
+        let mut map = serializer.serialize_map(Some(2))?;
+        map.serialize_entry("type", type_name)?;
+        map.serialize_entry("import", import)?;
+        map.end()
+      }
+    }
+  }
+}
+
 pub const fn default_bool<const V: bool>() -> bool {
   V
 }
@@ -58,6 +122,47 @@ pub struct DbConnectionConfig {
   pub pool_size: u32,
   #[serde(rename = "CONNECTION_TIMEOUT", default = "default_connection_timeout")]
   pub connection_timeout: u64,
+  #[serde(default)]
+  pub type_mapping: Option<HashMap<String, CustomTypeMapping>>,
+}
+
+impl DbConnectionConfig {
+  pub fn find_type_mapping(&self, db_types: &[&str]) -> Option<&CustomTypeMapping> {
+    let type_mapping = self.type_mapping.as_ref()?;
+    let normalized: HashMap<String, &CustomTypeMapping> = type_mapping
+      .iter()
+      .map(|(k, v)| (k.trim().to_lowercase(), v))
+      .collect();
+
+    db_types.iter().find_map(|db_type| {
+      let db_type = db_type
+        .split_whitespace()
+        .collect::<Vec<_>>()
+        .join(" ")
+        .to_lowercase();
+      normalized
+        .get(&db_type)
+        .or_else(|| normalized.get(&strip_type_modifiers(&db_type)))
+        .copied()
+    })
+  }
+}
+
+/// e.g. `bigint(20) unsigned` -> `bigint unsigned`, `varchar(255)` -> `varchar`
+fn strip_type_modifiers(db_type: &str) -> String {
+  let mut depth = 0;
+  let stripped: String = db_type
+    .chars()
+    .filter(|c| {
+      match c {
+        '(' => depth += 1,
+        ')' => depth -= 1,
+        _ => return depth == 0,
+      }
+      false
+    })
+    .collect();
+  stripped.split_whitespace().collect::<Vec<_>>().join(" ")
 }
 
 fn default_pool_size() -> u32 {
@@ -317,6 +422,8 @@ impl Config {
       .or_else(|| Some(default_connection_timeout()))
       .unwrap();
 
+    let type_mapping = default_config.and_then(|x| x.type_mapping.clone());
+
     DbConnectionConfig {
       db_type: db_type.to_owned(),
       db_host,
@@ -328,6 +435,7 @@ impl Config {
       pg_search_path: pg_search_path.to_owned(),
       pool_size,
       connection_timeout,
+      type_mapping,
     }
   }
 

src/common/config.test.rs

@@ -0,0 +1,79 @@
+#[cfg(test)]
+mod tests {
+  use crate::common::config::{CustomTypeMapping, DbConnectionConfig};
+
+  fn connection_with_mapping(type_mapping: &str) -> DbConnectionConfig {
+    serde_json::from_str(&format!(r#"{{ "DB_TYPE": "mysql", "type_mapping": {type_mapping} }}"#)).unwrap()
+  }
+
+  fn mapped_type(conn: &DbConnectionConfig, db_types: &[&str]) -> Option<String> {
+    conn.find_type_mapping(db_types).map(|mapping| match mapping {
+      CustomTypeMapping::Simple(type_name) | CustomTypeMapping::WithImport { type_name, .. } => type_name.to_owned(),
+    })
+  }
+
+  #[test]
+  fn should_match_case_insensitively() {
+    let conn = connection_with_mapping(r#"{ "BigInt": "string" }"#);
+    assert_eq!(mapped_type(&conn, &["BIGINT"]), Some("string".to_string()));
+  }
+
+  #[test]
+  fn should_ignore_type_modifiers() {
+    let conn = connection_with_mapping(r#"{ "varchar": "Name", "bigint unsigned": "string" }"#);
+    assert_eq!(mapped_type(&conn, &["VARCHAR(255)"]), Some("Name".to_string()));
+    assert_eq!(mapped_type(&conn, &["bigint(20) unsigned"]), Some("string".to_string()));
+    assert_eq!(mapped_type(&conn, &["bigint  unsigned"]), Some("string".to_string()));
+  }
+
+  #[test]
+  fn should_prefer_exact_match_and_earlier_candidates() {
+    let conn = connection_with_mapping(r#"{ "tinyint(1)": "boolean", "tinyint": "number", "bigint": "bigint" }"#);
+    assert_eq!(
+      mapped_type(&conn, &["tinyint(1)", "tinyint"]),
+      Some("boolean".to_string())
+    );
+    assert_eq!(
+      mapped_type(&conn, &["tinyint(4)", "tinyint"]),
+      Some("number".to_string())
+    );
+    assert_eq!(
+      mapped_type(&conn, &["bigint(20) unsigned", "bigint"]),
+      Some("bigint".to_string())
+    );
+  }
+
+  #[test]
+  fn should_return_none_without_mapping() {
+    let conn = connection_with_mapping("null");
+    assert_eq!(mapped_type(&conn, &["bigint"]), None);
+    let conn = connection_with_mapping(r#"{ "bigint": "string" }"#);
+    assert_eq!(mapped_type(&conn, &["integer"]), None);
+  }
+
+  #[test]
+  fn should_parse_object_mapping_with_optional_import() {
+    let conn = connection_with_mapping(
+      r#"{ "a": { "type": "A" }, "b": { "type": "B", "import": "" }, "c": { "type": "C", "import": "import type { C } from 'c'" } }"#,
+    );
+    let mapping = conn.type_mapping.unwrap();
+    assert!(matches!(&mapping["a"], CustomTypeMapping::Simple(t) if t == "A"));
+    assert!(matches!(&mapping["b"], CustomTypeMapping::Simple(t) if t == "B"));
+    assert!(matches!(&mapping["c"], CustomTypeMapping::WithImport { type_name, .. } if type_name == "C"));
+  }
+
+  #[test]
+  fn should_reject_invalid_mappings() {
+    for type_mapping in [
+      r#"{ "bigint": "" }"#,
+      r#"{ "bigint": { "type": " " } }"#,
+      r#"{ "bigint": { "import": "import x from 'x'" } }"#,
+      r#"{ "bigint": { "type": "X", "import": 1 } }"#,
+      r#"{ "bigint": 1 }"#,
+    ] {
+      let result: Result<DbConnectionConfig, _> =
+        serde_json::from_str(&format!(r#"{{ "DB_TYPE": "mysql", "type_mapping": {type_mapping} }}"#));
+      assert!(result.is_err(), "Expected {type_mapping} to be rejected");
+    }
+  }
+}

src/common/lazy.rs

@@ -47,7 +47,7 @@ pub static DB_CONN_CACHE: LazyLock<HashMap<String, Arc<Mutex<DBConn>>>> = LazyLo
             .await
             .expect(&ERR_DB_CONNECTION_ISSUE);
 
-          DBConn::MySQLPooledConn(Mutex::new(pool))
+          DBConn::MySQLPooledConn(Mutex::new(pool), connection.to_string())
         })
       }),
       DatabaseType::Sqlite => task::block_in_place(|| {
@@ -61,7 +61,7 @@ pub static DB_CONN_CACHE: LazyLock<HashMap<String, Arc<Mutex<DBConn>>>> = LazyLo
             .await
             .expect(&ERR_DB_CONNECTION_ISSUE);
 
-          DBConn::SqliteConn(Mutex::new(pool))
+          DBConn::SqliteConn(Mutex::new(pool), connection.to_string())
         })
       }),
       DatabaseType::Postgres => task::block_in_place(|| {
@@ -75,10 +75,10 @@ pub static DB_CONN_CACHE: LazyLock<HashMap<String, Arc<Mutex<DBConn>>>> = LazyLo
             .await
             .expect(&ERR_DB_CONNECTION_ISSUE);
 
-          let db_conn = DBConn::PostgresConn(Mutex::new(pool));
+          let db_conn = DBConn::PostgresConn(Mutex::new(pool), connection.to_string());
 
           let conn = match &db_conn {
-            DBConn::PostgresConn(conn) => conn,
+            DBConn::PostgresConn(conn, _) => conn,
             _ => panic!("Invalid connection type"),
           };
 

src/common/mod.rs

@@ -5,6 +5,9 @@ use swc_common::MultiSpan;
 
 pub mod cli;
 pub mod config;
+#[cfg(test)]
+#[path = "./config.test.rs"]
+mod config_test;
 pub mod dotenv;
 pub mod lazy;
 pub mod types;