GetPrimaryKey returns pointer to the ColumnDefinition instead of the name of the column; thus we can change the property of the column

xwb1989 · xwb1989 · commit a1939ab915c6 · 2015-04-20T11:45:36.000-04:00
diff --git a/analyzer.go b/analyzer.go
@@ -25,19 +25,19 @@ func GetTableName(node SimpleTableExpr) string {
 }
 
 // Get the primary key of the table, sqlNode must be a CreateTable struct
-func GetPrimaryKey(sqlNode SQLNode) (string, error) {
+func GetPrimaryKey(sqlNode SQLNode) (*ColumnDefinition, error) {
 	node, ok := sqlNode.(*CreateTable)
 	if !ok {
-		return "", errors.New("fail to convert interface SQLNode to struct CreateTable")
+		return nil, errors.New("fail to convert interface SQLNode to struct CreateTable")
 	}
 	for _, col := range node.ColumnDefinitions {
 		for _, att := range col.ColumnAtts {
 			if att == AST_PRIMARY_KEY {
-				return col.ColName, nil
+				return col, nil
 			}
 		}
 	}
-	return "", errors.New("unable to find primary key")
+	return nil, errors.New("unable to find primary key")
 }
 
 // GetColName returns the column name, only if
diff --git a/analyzer_test.go b/analyzer_test.go
@@ -18,15 +18,41 @@ func TestPrimaryKey(t *testing.T) {
 	assert.Nil(t, err)
 	primary_key, err := GetPrimaryKey(tree)
 	assert.Nil(t, err)
-	assert.Equal(t, "ID", primary_key)
+	assert.Equal(t, "ID", primary_key.ColName)
 
-	sql = `create table t1 (
+}
+func TestNoPrimaryKey(t *testing.T) {
+	sql := `create table t1 (
 	LastName varchar(255),
 	FirstName varchar(255),
 	ID int unique key
 )`
-	tree, err = Parse(sql)
+	tree, err := Parse(sql)
 	assert.Nil(t, err)
 	_, err = GetPrimaryKey(tree)
 	assert.NotNil(t, err)
 }
+
+func TestChangePrimaryKeyProperty(t *testing.T) {
+	sql := `create table t1 (
+	LastName varchar(255),
+	FirstName varchar(255),
+	ID int primary key
+)`
+
+	tree, err := Parse(sql)
+	assert.Nil(t, err)
+	primary_key, err := GetPrimaryKey(tree)
+	assert.Nil(t, err)
+
+	primary_key.ColumnAtts = append(primary_key.ColumnAtts, "auto_increment")
+
+	sql_actual := String(tree)
+
+	sql_expected := `create table t1 (
+	LastName varchar(255),
+	FirstName varchar(255),
+	ID int primary key auto_increment
+)`
+	assert.Equal(t, sql_expected, sql_actual)
+}

Original file line number	Diff line number	Diff line change
`@@ -25,19 +25,19 @@ func GetTableName(node SimpleTableExpr) string {`
`25`	`25`	`}`
`26`	`26`
`27`	`27`	`// Get the primary key of the table, sqlNode must be a CreateTable struct`
`28`		`-func GetPrimaryKey(sqlNode SQLNode) (string, error) {`
	`28`	`+func GetPrimaryKey(sqlNode SQLNode) (*ColumnDefinition, error) {`
`29`	`29`	`node, ok := sqlNode.(*CreateTable)`
`30`	`30`	`if !ok {`
`31`		`- return "", errors.New("fail to convert interface SQLNode to struct CreateTable")`
	`31`	`+ return nil, errors.New("fail to convert interface SQLNode to struct CreateTable")`
`32`	`32`	`}`
`33`	`33`	`for _, col := range node.ColumnDefinitions {`
`34`	`34`	`for _, att := range col.ColumnAtts {`
`35`	`35`	`if att == AST_PRIMARY_KEY {`
`36`		`- return col.ColName, nil`
	`36`	`+ return col, nil`
`37`	`37`	`}`
`38`	`38`	`}`
`39`	`39`	`}`
`40`		`- return "", errors.New("unable to find primary key")`
	`40`	`+ return nil, errors.New("unable to find primary key")`
`41`	`41`	`}`
`42`	`42`
`43`	`43`	`// GetColName returns the column name, only if`