feat: Add upsert operation mode (#400)

* feat: Add `upsert` operation mode

* test(push): fix upsert test

---------

Co-authored-by: Adrien Aury <[email protected]>

Author: youen

CHANGELOG.md

@@ -14,8 +14,9 @@ Types of changes
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
-## [3.5.1]
+## [3.6.0]
 
+- `Added` new push mode `upsert` to insert or update rows based on primary keys
 - `Fixed` relations extraction with Oracle database create duplicated relations
 
 ## [3.5.0]

README.md

@@ -449,6 +449,18 @@ lino push update source --table actor <<<'{"actor_id":998,"last_name":"CHASE","_
 
 The `__usingpk__` field can also be used with an ingress descriptor at any level in the data. The name of this field can be changed to another value with the `--using-pk-field` flag.
 
+### Upsert
+
+The `upsert` mode allows to insert new rows or update existing ones based on the primary key.
+
+```bash
+$ lino push upsert dest < data.jsonl
+```
+
+If a row with the same primary key already exists in the destination table, it will be updated with the new values. If it does not exist, it will be inserted.
+
+**Note:** This feature is currently supported for PostgreSQL and Oracle databases.
+
 ### How to recover from error
 
 Use options `lino pull --exclude-from-file` (shortcut `-X`) and `lino push --savepoint` combined to handle error recovery. The process will restart where it failed if an error has interrupted it in a previous run.

internal/app/push/cli.go

@@ -102,7 +102,7 @@ func NewCommand(fullName string, err *os.File, out *os.File, in *os.File) *cobra
 	)
 
 	cmd := &cobra.Command{
-		Use:     "push {<truncate>|<insert>|<update>|<delete>} [Data Connector Name]",
+		Use:     "push {<truncate>|<insert>|<update>|<delete>|<upsert>} [Data Connector Name]",
 		Short:   "Push data to a database with a pushing mode (insert by default)",
 		Long:    "",
 		Example: fmt.Sprintf("  %[1]s push truncate dstdatabase\n  %[1]s push dstdatabase", fullName),

internal/infra/push/datadestination_db2.go

@@ -91,6 +91,10 @@ func (d Db2Dialect) InsertStatement(tableName string, selectValues []ValueDescri
 	return sql.String(), selectValues
 }
 
+func (d Db2Dialect) UpsertStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
+	return "", nil, &push.Error{Description: "upsert not implemented for db2"}
+}
+
 // UpdateStatement
 func (d Db2Dialect) UpdateStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
 	sql := &strings.Builder{}

internal/infra/push/datadestination_db2_dummy.go

@@ -67,6 +67,11 @@ func (d Db2Dialect) InsertStatement(tableName string, selectValues []ValueDescri
 	panic(fmt.Errorf("Not implemented"))
 }
 
+// UpsertStatement
+func (d Db2Dialect) UpsertStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
+	panic(fmt.Errorf("Not implemented"))
+}
+
 // UpdateStatement
 func (d Db2Dialect) UpdateStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
 	panic(fmt.Errorf("Not implemented"))

internal/infra/push/datadestination_mariadb.go

@@ -89,6 +89,10 @@ func (d MariadbDialect) InsertStatement(tableName string, selectValues []ValueDe
 	return sql.String(), selectValues
 }
 
+func (d MariadbDialect) UpsertStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
+	return "", nil, &push.Error{Description: "upsert not implemented for mariadb"}
+}
+
 func (d MariadbDialect) UpdateStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
 	sql := &strings.Builder{}
 	sql.WriteString("UPDATE ")

internal/infra/push/datadestination_oracle.go

@@ -124,6 +124,63 @@ func (d OracleDialect) InsertStatement(tableName string, selectValues []ValueDes
 	return sql.String(), selectValues
 }
 
+// UpsertStatement
+func (d OracleDialect) UpsertStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
+	sql := &strings.Builder{}
+	sql.WriteString("MERGE INTO ")
+	sql.WriteString(tableName)
+	sql.WriteString(" target USING (SELECT ")
+
+	for i, col := range selectValues {
+		if i > 0 {
+			sql.WriteString(", ")
+		}
+		sql.WriteString(d.Placeholder(i + 1))
+		sql.WriteString(" AS \"")
+		sql.WriteString(col.name)
+		sql.WriteString("\"")
+	}
+	sql.WriteString(" FROM dual) source ON (")
+
+	for i, pk := range primaryKeys {
+		if i > 0 {
+			sql.WriteString(" AND ")
+		}
+		sql.WriteString(fmt.Sprintf("target.\"%s\" = source.\"%s\"", pk, pk))
+	}
+	sql.WriteString(") WHEN MATCHED THEN UPDATE SET ")
+
+	first := true
+	for _, col := range selectValues {
+		if isAPrimaryKey(col.name, primaryKeys) {
+			continue
+		}
+		if !first {
+			sql.WriteString(", ")
+		}
+		sql.WriteString(fmt.Sprintf("target.\"%s\" = source.\"%s\"", col.name, col.name))
+		first = false
+	}
+
+	sql.WriteString(" WHEN NOT MATCHED THEN INSERT (")
+	for i, col := range selectValues {
+		if i > 0 {
+			sql.WriteString(", ")
+		}
+		sql.WriteString(fmt.Sprintf("\"%s\"", col.name))
+	}
+	sql.WriteString(") VALUES (")
+	for i, col := range selectValues {
+		if i > 0 {
+			sql.WriteString(", ")
+		}
+		sql.WriteString(fmt.Sprintf("source.\"%s\"", col.name))
+	}
+	sql.WriteString(")")
+
+	return sql.String(), selectValues, nil
+}
+
 // UpdateStatement
 func (d OracleDialect) UpdateStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
 	sql := &strings.Builder{}

internal/infra/push/datadestination_postgres.go

@@ -100,6 +100,49 @@ func (d PostgresDialect) InsertStatement(tableName string, selectValues []ValueD
 	return sql.String(), selectValues
 }
 
+func (d PostgresDialect) UpsertStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
+	protectedColumns := []string{}
+	for _, c := range selectValues {
+		protectedColumns = append(protectedColumns, fmt.Sprintf("\"%s\"", c.name))
+	}
+
+	sql := &strings.Builder{}
+	sql.WriteString("INSERT INTO ")
+	sql.WriteString(tableName)
+	sql.WriteString("(")
+	sql.WriteString(strings.Join(protectedColumns, ","))
+	sql.WriteString(") VALUES (")
+	for i := 1; i <= len(selectValues); i++ {
+		sql.WriteString(d.Placeholder(i))
+		if i < len(selectValues) {
+			sql.WriteString(", ")
+		}
+	}
+
+	if len(primaryKeys) > 0 {
+		sql.WriteString(") ON CONFLICT (")
+		sql.WriteString(strings.Join(primaryKeys, ","))
+		sql.WriteString(") DO UPDATE SET ")
+
+		first := true
+		for _, column := range selectValues {
+			// Skip primary keys in update set
+			if isAPrimaryKey(column.name, primaryKeys) {
+				continue
+			}
+			if !first {
+				sql.WriteString(", ")
+			}
+			sql.WriteString(fmt.Sprintf("\"%s\" = EXCLUDED.\"%s\"", column.name, column.name))
+			first = false
+		}
+	} else {
+		sql.WriteString(")")
+	}
+
+	return sql.String(), selectValues, nil
+}
+
 func (d PostgresDialect) UpdateStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
 	sql := &strings.Builder{}
 	sql.WriteString("UPDATE ")

internal/infra/push/datadestination_sql.go

@@ -299,6 +299,12 @@ func (rw *SQLRowWriter) createStatement(row push.Row, where push.Row) *push.Erro
 			return pusherr
 		}
 
+	case rw.dd.mode == push.Upsert:
+		prepareStmt, rw.headers, pusherr = rw.dd.dialect.UpsertStatement(rw.tableName(), selectValues, whereValues, rw.table.PrimaryKey())
+		if pusherr != nil {
+			return pusherr
+		}
+
 	default: // Insert:
 		prepareStmt, rw.headers = rw.dd.dialect.InsertStatement(rw.tableName(), selectValues, rw.table.PrimaryKey())
 	}
@@ -376,6 +382,7 @@ func (rw *SQLRowWriter) Write(row push.Row, where push.Row) *push.Error {
 	rw.sqlLogger.Write(values)
 
 	_, err2 := rw.statement.Exec(values...)
+	log.Trace().AnErr("error", err2).Msg("push error")
 	if err2 != nil {
 		// reset statement after error
 		if err := rw.close(); err != nil {
@@ -483,6 +490,7 @@ type SQLDialect interface {
 	EnableConstraintsStatement(tableName string) string
 	TruncateStatement(tableName string) string
 	InsertStatement(tableName string, selectValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor)
+	UpsertStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error)
 	UpdateStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error)
 	IsDuplicateError(error) bool
 	ConvertValue(push.Value, ValueDescriptor) push.Value

internal/infra/push/datadestination_sqlserver.go

@@ -87,6 +87,10 @@ func (d SQLServerDialect) InsertStatement(tableName string, selectValues []Value
 	return sql.String(), selectValues
 }
 
+func (d SQLServerDialect) UpsertStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
+	return "", nil, &push.Error{Description: "upsert not implemented for sqlserver"}
+}
+
 func (d SQLServerDialect) UpdateStatement(tableName string, selectValues []ValueDescriptor, whereValues []ValueDescriptor, primaryKeys []string) (statement string, headers []ValueDescriptor, err *push.Error) {
 	sql := &strings.Builder{}
 	sql.WriteString("UPDATE ")