add u64 codec in mssql and sqlite

sqlpage · Nov 7, 2024 · 1dc8628 · 1dc8628
1 parent 659cf57
commit 1dc8628
Show file tree

Hide file tree

Showing 2 changed files with 120 additions and 0 deletions.
diff --git a/sqlx-core/src/mssql/types/uint.rs b/sqlx-core/src/mssql/types/uint.rs
@@ -28,3 +28,92 @@ impl Decode<'_, Mssql> for u8 {
         Ok(value.as_bytes()?[0])
     }
 }
+
+impl Type<Mssql> for u16 {
+    fn type_info() -> MssqlTypeInfo {
+        MssqlTypeInfo(TypeInfo::new(DataType::IntN, 4))
+    }
+
+    fn compatible(ty: &MssqlTypeInfo) -> bool {
+        matches!(ty.0.ty, DataType::Int | DataType::IntN) && ty.0.size == 4
+    }
+}
+
+impl Encode<'_, Mssql> for u16 {
+    fn encode_by_ref(&self, buf: &mut Vec<u8>) -> IsNull {
+        let value = i32::from(*self);
+        buf.extend(&value.to_le_bytes());
+        IsNull::No
+    }
+}
+
+impl Decode<'_, Mssql> for u16 {
+    fn decode(value: MssqlValueRef<'_>) -> Result<Self, BoxDynError> {
+        let bytes = value.as_bytes()?;
+        let val = i32::from_le_bytes(bytes.try_into()?);
+        u16::try_from(val).map_err(Into::into)
+    }
+}
+
+impl Type<Mssql> for u32 {
+    fn type_info() -> MssqlTypeInfo {
+        MssqlTypeInfo(TypeInfo::new(DataType::IntN, 8))
+    }
+
+    fn compatible(ty: &MssqlTypeInfo) -> bool {
+        matches!(ty.0.ty, DataType::BigInt | DataType::IntN) && ty.0.size == 8
+    }
+}
+
+impl Encode<'_, Mssql> for u32 {
+    fn encode_by_ref(&self, buf: &mut Vec<u8>) -> IsNull {
+        let value = i64::from(*self);
+        buf.extend(&value.to_le_bytes());
+        IsNull::No
+    }
+}
+
+impl Decode<'_, Mssql> for u32 {
+    fn decode(value: MssqlValueRef<'_>) -> Result<Self, BoxDynError> {
+        let bytes = value.as_bytes()?;
+        let val = i64::from_le_bytes(bytes.try_into()?);
+        u32::try_from(val).map_err(Into::into)
+    }
+}
+
+impl Type<Mssql> for u64 {
+    fn type_info() -> MssqlTypeInfo {
+        MssqlTypeInfo(TypeInfo::new(DataType::Numeric, 0))
+    }
+
+    fn compatible(ty: &MssqlTypeInfo) -> bool {
+        matches!(ty.0.ty, DataType::Numeric | DataType::Decimal)
+    }
+}
+
+impl Encode<'_, Mssql> for u64 {
+    fn encode_by_ref(&self, buf: &mut Vec<u8>) -> IsNull {
+        let bytes = self.to_le_bytes();
+        let scale = 0i8;
+        let len = 17u8;
+
+        buf.push(len);
+        buf.push(scale.to_le_bytes()[0]);
+        buf.extend(&bytes);
+        buf.extend(&[0u8; 8]);
+
+        IsNull::No
+    }
+}
+
+impl Decode<'_, Mssql> for u64 {
+    fn decode(value: MssqlValueRef<'_>) -> Result<Self, BoxDynError> {
+        let bytes = value.as_bytes()?;
+        if bytes.len() < 17 {
+            return Err("Invalid numeric value length".into());
+        }
+
+        let value_bytes = &bytes[2..10];
+        Ok(u64::from_le_bytes(value_bytes.try_into()?))
+    }
+}
diff --git a/sqlx-core/src/sqlite/types/uint.rs b/sqlx-core/src/sqlite/types/uint.rs
@@ -76,3 +76,34 @@ impl<'r> Decode<'r, Sqlite> for u32 {
         Ok(value.int64().try_into()?)
     }
 }
+
+impl Type<Sqlite> for u64 {
+    fn type_info() -> SqliteTypeInfo {
+        SqliteTypeInfo(DataType::Int64)
+    }
+
+    fn compatible(ty: &SqliteTypeInfo) -> bool {
+        matches!(ty.0, DataType::Int | DataType::Int64)
+    }
+}
+
+impl<'q> Encode<'q, Sqlite> for u64 {
+    fn encode_by_ref(&self, args: &mut Vec<SqliteArgumentValue<'q>>) -> IsNull {
+        match i64::try_from(*self) {
+            Ok(value) => {
+                args.push(SqliteArgumentValue::Int64(value));
+                IsNull::No
+            }
+            Err(_) => {
+                log::warn!("u64 value {} too large for sqlite, encoding as NULL", self);
+                IsNull::Yes
+            }
+        }
+    }
+}
+
+impl<'r> Decode<'r, Sqlite> for u64 {
+    fn decode(value: SqliteValueRef<'r>) -> Result<Self, BoxDynError> {
+        Ok(value.int64().try_into()?)
+    }
+}