From bf4a5349e1904cd2adb26f8e33b1eec54e7a2e6a Mon Sep 17 00:00:00 2001
From: Simon Binder <oss@simonbinder.eu>
Date: Fri, 22 Nov 2024 17:56:49 +0100
Subject: [PATCH] Propagate types with known nullability first

Closes https://github.com/simolus3/drift/issues/3351
---
 sqlparser/CHANGELOG.md                            |  2 ++
 .../lib/src/analysis/types/graph/type_graph.dart  |  4 ++++
 sqlparser/lib/src/analysis/types/types.dart       |  1 +
 .../test/analysis/types2/misc_cases_test.dart     | 15 +++++++++++++++
 4 files changed, 22 insertions(+)

diff --git a/sqlparser/CHANGELOG.md b/sqlparser/CHANGELOG.md
index 6c45e10e1..df2838c76 100644
--- a/sqlparser/CHANGELOG.md
+++ b/sqlparser/CHANGELOG.md
@@ -1,6 +1,8 @@
 ## 0.40.0-dev
 
 - Add support for the `dbstat` module.
+- Prioritize null propagation in type resolver, leading to more accurate
+  analysis on which columns are nullable.
 
 ## 0.39.2
 
diff --git a/sqlparser/lib/src/analysis/types/graph/type_graph.dart b/sqlparser/lib/src/analysis/types/graph/type_graph.dart
index f3b848a37..636a74835 100644
--- a/sqlparser/lib/src/analysis/types/graph/type_graph.dart
+++ b/sqlparser/lib/src/analysis/types/graph/type_graph.dart
@@ -64,6 +64,9 @@ class TypeGraph {
     _indexRelations();
 
     var queue = List.of(_knownTypes.keys);
+    // Apply type propagation to known types for which we know the nullability
+    // too first.
+    queue.sortBy<num>((e) => _knownNullability.containsKey(e) ? 1 : 0);
     while (queue.isNotEmpty) {
       _propagateTypeInfo(queue, queue.removeLast());
     }
@@ -153,6 +156,7 @@ class TypeGraph {
           .any((nullable) => nullable == true);
 
       _knownNullability[edge.target] = nullable;
+      resolved.add(edge.target);
     }
   }
 
diff --git a/sqlparser/lib/src/analysis/types/types.dart b/sqlparser/lib/src/analysis/types/types.dart
index 18a8afdec..c003f0270 100644
--- a/sqlparser/lib/src/analysis/types/types.dart
+++ b/sqlparser/lib/src/analysis/types/types.dart
@@ -1,5 +1,6 @@
 import 'dart:math';
 
+import 'package:collection/collection.dart';
 import 'package:sqlparser/sqlparser.dart';
 
 part 'expectation.dart';
diff --git a/sqlparser/test/analysis/types2/misc_cases_test.dart b/sqlparser/test/analysis/types2/misc_cases_test.dart
index c99632329..122f1b147 100644
--- a/sqlparser/test/analysis/types2/misc_cases_test.dart
+++ b/sqlparser/test/analysis/types2/misc_cases_test.dart
@@ -176,4 +176,19 @@ WITH RECURSIVE
       ResolvedType(type: BasicType.int),
     ]);
   });
+
+  test('keeps nullability information across unions', () {
+    // https://github.com/simolus3/drift/issues/3351
+    final engine = SqlEngine();
+    final ctx = engine.analyze('''
+      SELECT CAST(NULL AS INT) AS value
+      UNION
+      SELECT CAST(NULL AS INT);
+    ''');
+
+    final select = ctx.root as CompoundSelectStatement;
+    final column = select.resolvedColumns!.first;
+    expect(ctx.typeOf(column),
+        ResolveResult(ResolvedType(nullable: true, type: BasicType.int)));
+  });
 }