Add non-nullability support for LIMIT queries

jschaf · Feb 5, 2022 · 1266d15 · 1266d15
1 parent 43bad01
commit 1266d15
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/internal/pginfer/explain.go b/internal/pginfer/explain.go
@@ -12,6 +12,7 @@ type PlanType string
 
 const (
 	PlanResult      PlanType = "Result"      // select statement
+	PlanLimit       PlanType = "Limit"       // select statement with a limit
 	PlanModifyTable PlanType = "ModifyTable" // update, insert, or delete statement
 )
 

diff --git a/internal/pginfer/nullability.go b/internal/pginfer/nullability.go
@@ -24,10 +24,10 @@ func isColNullable(query *ast.SourceQuery, plan Plan, out string, column pg.Colu
 		// try below
 	}
 
-	// A plain select query with no joins where the column comes from a table and
-	// has a not-null constraint. Not full proof because of cross-join with comma
-	// syntax.
-	if plan.Type == PlanResult &&
+	// A plain select query (possibly with a LIMIT clause) with no joins where
+	// the column comes from a table and has a not-null constraint. Not full
+	// proof because of cross-join with comma syntax.
+	if (plan.Type == PlanResult || plan.Type == PlanLimit) &&
 		!strings.Contains(strings.ToLower(query.PreparedSQL), "join") &&
 		!column.Null {
 		return false