From 331a5a6ab575f16ea935005a18ce9c2af976696b Mon Sep 17 00:00:00 2001
From: Daylon Wilkins <daylon@liquidata.co>
Date: Wed, 9 Dec 2020 02:04:13 -0800
Subject: [PATCH] Fixed NULLs in pk on table import

---
 bats/import-create-tables.bats             | 21 +++++++++++++++++++++
 go/libraries/doltcore/mvdata/data_mover.go | 12 ++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/bats/import-create-tables.bats b/bats/import-create-tables.bats
index dd2522fe12..3de3660f58 100755
--- a/bats/import-create-tables.bats
+++ b/bats/import-create-tables.bats
@@ -29,10 +29,12 @@ DELIM
     "four":"c3"
 }
 JSON
+
     cat <<DELIM > name-map-data.csv
 one,two,three,four
 0,1,2,3
 DELIM
+
     cat <<SQL > name-map-sch.sql
 CREATE TABLE test (
     pk int not null,
@@ -459,6 +461,25 @@ SQL
     [ "$status" -eq 1 ]
 }
 
+@test "fail on import table creation when defined pk has a NULL value" {
+    cat <<DELIM > null-pk-1.csv
+pk,v1
+"a",1
+,2
+DELIM
+    cat <<DELIM > null-pk-2.csv
+pk1,pk2,v1
+0,0,0
+1,,1
+DELIM
+    run dolt table import -c --pk=pk test null-pk-1.csv
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "pk" ]]
+    run dolt table import -c --pk=pk1,pk2 test null-pk-2.csv
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "pk2" ]]
+}
+
 @test "table import -c infers types from data" {
     cat <<DELIM > types.csv
 pk,str,int,bool,float, date, time, datetime
diff --git a/go/libraries/doltcore/mvdata/data_mover.go b/go/libraries/doltcore/mvdata/data_mover.go
index 099ac374ed..b710c28014 100644
--- a/go/libraries/doltcore/mvdata/data_mover.go
+++ b/go/libraries/doltcore/mvdata/data_mover.go
@@ -254,6 +254,18 @@ func InferSchema(ctx context.Context, root *doltdb.RootValue, rd table.TableRead
 	pkSet := set.NewStrSet(pks)
 	newCols, _ := schema.MapColCollection(infCols, func(col schema.Column) (schema.Column, error) {
 		col.IsPartOfPK = pkSet.Contains(col.Name)
+		if col.IsPartOfPK {
+			hasNotNull := false
+			for _, constraint := range col.Constraints {
+				if _, ok := constraint.(schema.NotNullConstraint); ok {
+					hasNotNull = true
+					break
+				}
+			}
+			if !hasNotNull {
+				col.Constraints = append(col.Constraints, schema.NotNullConstraint{})
+			}
+		}
 		return col, nil
 	})