more duckdb data types, use sets instead (#7326)

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR fixes any issues, list them here by number (e.g., Fixes
#123).
-->

This is based on
https://github.com/marimo-team/codemirror-sql/blob/caa7c664135988b634f55a3e57a1327a5ffeede2/src/dialects/duckdb/duckdb.ts

## 🔍 Description of Changes

<!--
Detail the specific changes made in this pull request. Explain the
problem addressed and how it was resolved. If applicable, provide before
and after comparisons, screenshots, or any relevant details to help
reviewers understand the changes easily.
-->

## 📋 Checklist

- [x] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [ ] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [x] I have added tests for the changes made.
- [x] I have run the code and verified that it works as expected.

Author: Light2Dark

marimo/_data/get_datasets.py

@@ -453,83 +453,112 @@ def _get_duckdb_database_names(
         return []
 
 
+_INTEGER_TYPES = {
+    "tinyint",
+    "smallint",
+    "integer",
+    "bigint",
+    "hugeint",
+    "integral",
+    "long",
+    "short",
+    "signed",
+    "oid",
+    "varint",
+    "int",
+    "int1",
+    "int2",
+    "int4",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "int128",
+    "ubigint",
+    "uhugeint",
+    "usmallint",
+    "utinyint",
+}
+_NUMERIC_TYPES = {"float", "real", "double", "decimal", "numeric", "dec"}
+_BOOLEAN_TYPES = {"boolean", "bool", "logical"}
+_STRING_TYPES = {
+    "varchar",
+    "char",
+    "bpchar",
+    "text",
+    "string",
+    "blob",
+    "uuid",
+    "guid",
+    "nvarchar",
+}
+_TIME_TYPES = {"time", "time with time zone", "timetz"}
+_DATETIME_TYPES = {"datetime", "interval"}
+_BINARY_TYPES = {"bit", "bitstring", "binary", "varbinary", "bytea"}
+_UNKNOWN_TYPES = {
+    "row",
+    "geometry",
+    # Null type (can occur when attaching databases or with unknown column types)
+    "null",
+    '"null"',
+}
+
+
 def _db_type_to_data_type(db_type: str) -> DataType:
     """Convert a DuckDB type to a Marimo data type.
     Reference: https://duckdb.org/docs/stable/sql/data_types/overview
+    Latest types: https://github.com/marimo-team/codemirror-sql/blob/caa7c664135988b634f55a3e57a1327a5ffeede2/src/dialects/duckdb/duckdb.ts
     """
     db_type = db_type.lower()
-    # Numeric types
-    if db_type in [
-        "tinyint",
-        "smallint",
-        "integer",
-        "bigint",
-        "hugeint",
-        "utinyint",
-        "usmallint",
-        "uinteger",
-        "ubigint",
-        "uhugeint",
-    ]:
+
+    # Check for exact matches first, then patterns
+
+    if db_type in _INTEGER_TYPES or db_type.startswith("uint"):
         return "integer"
+
     if (
-        db_type
-        in [
-            "float",
-            "real",
-            "double",
-            "decimal",
-            "numeric",
-        ]
+        db_type in _NUMERIC_TYPES
         or db_type.startswith("decimal")
         or db_type.startswith("float")
     ):
         return "number"
-    # Boolean type
-    if db_type == "boolean":
+
+    if db_type in _BOOLEAN_TYPES:
         return "boolean"
-    # String types
-    if db_type in [
-        "varchar",
-        "char",
-        "bpchar",
-        "text",
-        "string",
-        "blob",
-        "uuid",
-    ]:
+
+    if db_type in _STRING_TYPES:
         return "string"
-    # Date and Time types
+
     if db_type == "date":
         return "date"
-    if db_type in ["time", "time with time zone"]:
+    if db_type in _TIME_TYPES:
         return "time"
-    if db_type in ["datetime", "interval"] or db_type.startswith("timestamp"):
+    if db_type in _DATETIME_TYPES or db_type.startswith("timestamp"):
         return "datetime"
+
+    # Binary types (represented as string)
+    if db_type in _BINARY_TYPES:
+        return "string"
+
+    # Enum types (represented as string)
+    if db_type == "enum" or db_type.startswith("enum"):
+        return "string"
+
     # Nested types
-    if "[" in db_type and "]" in db_type:
-        return "unknown"
     if (
         db_type.startswith("union")
         or db_type.startswith("map")
         or db_type.startswith("struct")
         or db_type.startswith("list")
         or db_type.startswith("array")
         or db_type.startswith("json")
+        or ("[" in db_type and "]" in db_type)
     ):
         return "unknown"
-    # Special types
-    if db_type == "bit":
-        return "string"  # Representing bit as string
-    if db_type == "enum" or db_type.startswith("enum"):
-        return "string"  # Representing enum as string
-    # Geometry types
-    if db_type == "geometry":
-        return "unknown"
-    # Null type (can occur when attaching databases or with unknown column types)
-    if db_type == "null" or db_type == '"null"':
+
+    # Other special types
+    if db_type in _UNKNOWN_TYPES:
         return "unknown"
 
     LOGGER.warning("Unknown DuckDB type: %s", db_type)
-    # Unknown type
     return "unknown"

tests/_data/test_get_datasets.py

@@ -628,23 +628,54 @@ def test_db_type_to_data_type_various() -> None:
     # Integer types
     assert _db_type_to_data_type("integer") == "integer"
     assert _db_type_to_data_type("bigint") == "integer"
+    assert _db_type_to_data_type("int128") == "integer"
+    assert _db_type_to_data_type("integral") == "integer"
+    assert _db_type_to_data_type("long") == "integer"
+    assert _db_type_to_data_type("short") == "integer"
+    assert _db_type_to_data_type("signed") == "integer"
+    assert _db_type_to_data_type("oid") == "integer"
+    assert _db_type_to_data_type("varint") == "integer"
+
+    # Unsigned integers
+    assert _db_type_to_data_type("utinyint") == "integer"
+    assert _db_type_to_data_type("usmallint") == "integer"
+    assert _db_type_to_data_type("uinteger") == "integer"
+    assert _db_type_to_data_type("ubigint") == "integer"
+    assert _db_type_to_data_type("uhugeint") == "integer"
+    assert _db_type_to_data_type("uint128") == "integer"
 
     # Numeric types
     assert _db_type_to_data_type("float") == "number"
     assert _db_type_to_data_type("double") == "number"
+    assert _db_type_to_data_type("float4") == "number"
+    assert _db_type_to_data_type("dec") == "number"
+    assert _db_type_to_data_type("decimal") == "number"
 
     # String types
     assert _db_type_to_data_type("varchar") == "string"
     assert _db_type_to_data_type("text") == "string"
+    assert _db_type_to_data_type("blob") == "string"
+    assert _db_type_to_data_type("guid") == "string"
+    assert _db_type_to_data_type("nvarchar") == "string"
+
+    # Binary types (represented as string)
+    assert _db_type_to_data_type("binary") == "string"
+    assert _db_type_to_data_type("varbinary") == "string"
+    assert _db_type_to_data_type("bytea") == "string"
 
     # Boolean
     assert _db_type_to_data_type("boolean") == "boolean"
+    assert _db_type_to_data_type("bool") == "boolean"
+    assert _db_type_to_data_type("logical") == "boolean"
 
     # Date/Time
     assert _db_type_to_data_type("date") == "date"
     assert _db_type_to_data_type("timestamp") == "datetime"
+    assert _db_type_to_data_type("timestamptz") == "datetime"
+    assert _db_type_to_data_type("timetz") == "time"
 
     # Special types
     assert _db_type_to_data_type("geometry") == "unknown"
     assert _db_type_to_data_type("null") == "unknown"
     assert _db_type_to_data_type("json") == "unknown"
+    assert _db_type_to_data_type("row") == "unknown"