fix: bson should handle nulls (#2485)

GlareDB · Feb 1, 2024 · 5dd4864 · 5dd4864
1 parent c264733
commit 5dd4864
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 0 deletions.
diff --git a/crates/datasources/src/bson/builder.rs b/crates/datasources/src/bson/builder.rs
@@ -241,11 +241,16 @@ macro_rules! append_scalar {
 fn append_value(val: RawBsonRef, typ: &DataType, col: &mut dyn ArrayBuilder) -> Result<()> {
     // So robust
     match (val, typ) {
+        // null
+        (RawBsonRef::Null, _) => append_null(typ, col)?,
+        (RawBsonRef::Undefined, _) => append_null(typ, col)?,
+
         // Boolean
         (RawBsonRef::Boolean(v), DataType::Boolean) => append_scalar!(BooleanBuilder, col, v),
         (RawBsonRef::Boolean(v), DataType::Utf8) => {
             append_scalar!(StringBuilder, col, v.to_string())
         }
+
         // Double
         (RawBsonRef::Double(v), DataType::Int32) => append_scalar!(Int32Builder, col, v as i32),
         (RawBsonRef::Double(v), DataType::Int64) => append_scalar!(Int64Builder, col, v as i64),

diff --git a/scripts/create-test-mongo-db.sh b/scripts/create-test-mongo-db.sh
@@ -40,6 +40,11 @@ docker exec $CONTAINER_ID mongoimport \
        "mongodb://localhost:27017/${DB_NAME}" \
        /tmp/bikeshare_stations.csv 1>&2
 
+# insert fixture data for a null handling regression test.
+docker exec $CONTAINER_ID mongosh \
+       "mongodb://localhost:27017/${DB_NAME}" \
+       --eval "db.null_test.insertMany([{a:1},{a:null}])" 1>&2
+
 # The mongod docker container is kinda bad. The MONGO_INITDB_... environment vars
 # might look like the obvious solution, but they don't work as you would expect.
 #

diff --git a/testdata/sqllogictests_mongodb/read.slt b/testdata/sqllogictests_mongodb/read.slt
@@ -4,3 +4,14 @@ query I
 SELECT count(*) FROM read_mongodb('${MONGO_CONN_STRING}', 'test', 'bikeshare_stations');
 ----
 102
+
+query I
+SELECT count(*) FROM read_mongodb('${MONGO_CONN_STRING}', 'test', 'null_test');
+----
+2
+
+query I
+SELECT a FROM read_mongodb('${MONGO_CONN_STRING}', 'test', 'null_test');
+----
+1
+NULL
diff --git a/tests/tests/test_bson.py b/tests/tests/test_bson.py
@@ -84,3 +84,22 @@ def test_read_bson(
                 assert len(row) == 5
                 assert row["beatle_name"] in beatles
                 assert beatles.index(row["beatle_name"]) == row["beatle_idx"] - 1
+
+def test_null_handling(
+    glaredb_connection: psycopg2.extensions.connection,
+    tmp_path_factory: pytest.TempPathFactory,
+):
+    tmp_dir = tmp_path_factory.mktemp(basename="null_handling", numbered=True)
+    data_path = tmp_dir.joinpath("mixed.bson")
+
+    with open(data_path, "wb") as f:
+        for i in range(100):
+            f.write(bson.encode({"a": 1}))
+
+        for i in range(10):
+            f.write(bson.encode({"a": None}))
+
+    with glaredb_connection.cursor() as curr:
+        curr.execute(f"select count(*) from '{data_path}'")
+        r = curr.fetchone()
+        assert r[0] == 110