From 780fc962c1223237abcb89f3299f04c1a855f5b6 Mon Sep 17 00:00:00 2001
From: Tianzhou <t@bytebase.com>
Date: Thu, 16 Oct 2025 01:08:34 +0800
Subject: [PATCH] fix: generate_column_not_null

---
 internal/diff/table.go                                    | 4 ++--
 testdata/diff/create_table/add_column_generated/diff.sql  | 4 +++-
 testdata/diff/create_table/add_column_generated/new.sql   | 1 +
 testdata/diff/create_table/add_column_generated/plan.json | 6 ++++++
 testdata/diff/create_table/add_column_generated/plan.sql  | 2 ++
 testdata/diff/create_table/add_column_generated/plan.txt  | 3 +++
 6 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/internal/diff/table.go b/internal/diff/table.go
index b9e1e18a..f8f8fbd8 100644
--- a/internal/diff/table.go
+++ b/internal/diff/table.go
@@ -1116,8 +1116,8 @@ func buildColumnClauses(column *ir.Column, isPartOfAnyPK bool, tableSchema strin
 		parts = append(parts, fmt.Sprintf("GENERATED ALWAYS AS (%s) STORED", *column.GeneratedExpr))
 	}
 
-	// 4. NOT NULL (skip for PK including multi-column PKs, identity, SERIAL, or generated columns)
-	if !column.IsNullable && column.Identity == nil && !isSerialColumn(column) && !isPartOfAnyPK && !column.IsGenerated {
+	// 4. NOT NULL (skip for PK including multi-column PKs, identity, and SERIAL)
+	if !column.IsNullable && column.Identity == nil && !isSerialColumn(column) && !isPartOfAnyPK {
 		parts = append(parts, "NOT NULL")
 	}
 
diff --git a/testdata/diff/create_table/add_column_generated/diff.sql b/testdata/diff/create_table/add_column_generated/diff.sql
index 237798bb..6c442993 100644
--- a/testdata/diff/create_table/add_column_generated/diff.sql
+++ b/testdata/diff/create_table/add_column_generated/diff.sql
@@ -1,4 +1,6 @@
 ALTER TABLE merge_request
 ADD COLUMN iid integer GENERATED ALWAYS AS (CAST(data ->> 'iid' AS int)) STORED CONSTRAINT pk_merge_request_iid PRIMARY KEY;
 
-ALTER TABLE merge_request ADD COLUMN title text GENERATED ALWAYS AS (data ->> 'title') STORED;
\ No newline at end of file
+ALTER TABLE merge_request ADD COLUMN title text GENERATED ALWAYS AS (data ->> 'title') STORED;
+
+ALTER TABLE merge_request ADD COLUMN cleaned_title varchar(255) GENERATED ALWAYS AS (lower(data ->> 'title')) STORED NOT NULL;
\ No newline at end of file
diff --git a/testdata/diff/create_table/add_column_generated/new.sql b/testdata/diff/create_table/add_column_generated/new.sql
index f278c718..eca02abd 100644
--- a/testdata/diff/create_table/add_column_generated/new.sql
+++ b/testdata/diff/create_table/add_column_generated/new.sql
@@ -2,5 +2,6 @@ CREATE TABLE public.merge_request (
     data jsonb NOT NULL,
     iid integer GENERATED ALWAYS AS ((data ->> 'iid')::integer) STORED,
     title text GENERATED ALWAYS AS (data ->> 'title') STORED,
+    cleaned_title varchar(255) GENERATED ALWAYS AS (lower(data ->> 'title')) STORED NOT NULL,
     CONSTRAINT pk_merge_request_iid PRIMARY KEY (iid)
 );
\ No newline at end of file
diff --git a/testdata/diff/create_table/add_column_generated/plan.json b/testdata/diff/create_table/add_column_generated/plan.json
index e5a533a2..d3059dc6 100644
--- a/testdata/diff/create_table/add_column_generated/plan.json
+++ b/testdata/diff/create_table/add_column_generated/plan.json
@@ -19,6 +19,12 @@
           "type": "table.column",
           "operation": "create",
           "path": "public.merge_request.title"
+        },
+        {
+          "sql": "ALTER TABLE merge_request ADD COLUMN cleaned_title varchar(255) GENERATED ALWAYS AS (lower(data ->> 'title')) STORED NOT NULL;",
+          "type": "table.column",
+          "operation": "create",
+          "path": "public.merge_request.cleaned_title"
         }
       ]
     }
diff --git a/testdata/diff/create_table/add_column_generated/plan.sql b/testdata/diff/create_table/add_column_generated/plan.sql
index efc1fa5f..6ce6ff56 100644
--- a/testdata/diff/create_table/add_column_generated/plan.sql
+++ b/testdata/diff/create_table/add_column_generated/plan.sql
@@ -2,3 +2,5 @@ ALTER TABLE merge_request
 ADD COLUMN iid integer GENERATED ALWAYS AS (CAST(data ->> 'iid' AS int)) STORED CONSTRAINT pk_merge_request_iid PRIMARY KEY;
 
 ALTER TABLE merge_request ADD COLUMN title text GENERATED ALWAYS AS (data ->> 'title') STORED;
+
+ALTER TABLE merge_request ADD COLUMN cleaned_title varchar(255) GENERATED ALWAYS AS (lower(data ->> 'title')) STORED NOT NULL;
diff --git a/testdata/diff/create_table/add_column_generated/plan.txt b/testdata/diff/create_table/add_column_generated/plan.txt
index 5317def3..bcd23d8f 100644
--- a/testdata/diff/create_table/add_column_generated/plan.txt
+++ b/testdata/diff/create_table/add_column_generated/plan.txt
@@ -5,6 +5,7 @@ Summary by type:
 
 Tables:
   ~ merge_request
+    + cleaned_title (column)
     + iid (column)
     + title (column)
 
@@ -15,3 +16,5 @@ ALTER TABLE merge_request
 ADD COLUMN iid integer GENERATED ALWAYS AS (CAST(data ->> 'iid' AS int)) STORED CONSTRAINT pk_merge_request_iid PRIMARY KEY;
 
 ALTER TABLE merge_request ADD COLUMN title text GENERATED ALWAYS AS (data ->> 'title') STORED;
+
+ALTER TABLE merge_request ADD COLUMN cleaned_title varchar(255) GENERATED ALWAYS AS (lower(data ->> 'title')) STORED NOT NULL;