public/grafana
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Spanner: support for lastInsertId (#101770)

Browse Source 
* Support for returning inserted ID in Spanner.
* Convert bytes to string when storing them into text column.
This commit is contained in:
Peter Štibraný
2025-03-10 15:59:06 +01:00
committed by GitHub
parent d5451f2e13
commit 0ce67dad74
8 changed files with 124 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/services/sqlstore/migrator/migrator.go
+3 -3
View File
@@ -240,19 +240,19 @@ func (mg *Migrator) run(ctx context.Context) (err error) {
migrationLogExists, err := mg.DBEngine.IsTableExist(mg.tableName)
if err != nil {
return fmt.Errorf("%v: %w", "failed to check table existence", err)
return fmt.Errorf("failed to check table existence: %w", err)
}
if !migrationLogExists {
// Check if dialect can initialize database from a snapshot.
err := mg.Dialect.CreateDatabaseFromSnapshot(ctx, mg.DBEngine, mg.tableName)
if err != nil {
return fmt.Errorf("%v: %w", "failed to create database from snapshot", err)
return fmt.Errorf("failed to create database from snapshot: %w", err)
}
migrationLogExists, err = mg.DBEngine.IsTableExist(mg.tableName)
if err != nil {
return fmt.Errorf("%v: %w", "failed to check table existence after applying snapshot", err)
return fmt.Errorf("failed to check table existence after applying snapshot: %w", err)
}
}
pkg/services/sqlstore/session.go
+7 -2
View File
@@ -11,6 +11,7 @@ import (
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
"go.opentelemetry.io/otel/trace/noop"
"xorm.io/xorm"
"github.com/grafana/grafana/pkg/apimachinery/errutil"
@@ -140,14 +141,18 @@ func (sess *DBSession) InsertId(bean any, dialect migrator.Dialect) error {
}
func (sess *DBSession) WithReturningID(driverName string, query string, args []any) (int64, error) {
supported := driverName != migrator.Postgres
var id int64
if !supported {
if driverName == migrator.Postgres {
query = fmt.Sprintf("%s RETURNING id", query)
if _, err := sess.SQL(query, args...).Get(&id); err != nil {
return id, err
}
} else {
if driverName == migrator.Spanner {
// Only works with INSERT statements.
query = fmt.Sprintf("%s THEN RETURN id", query)
}
sqlOrArgs := append([]any{query}, args...)
res, err := sess.Exec(sqlOrArgs...)
if err != nil {
pkg/util/xorm/go.mod
+3 -1
View File
@@ -3,7 +3,9 @@ module github.com/grafana/grafana/pkg/util/xorm
go 1.23.7
require (
cloud.google.com/go/spanner v1.75.0
github.com/googleapis/go-sql-spanner v1.11.1
github.com/grafana/grafana v5.4.5+incompatible
github.com/mattn/go-sqlite3 v1.14.22
github.com/stretchr/testify v1.10.0
xorm.io/builder v0.3.6
@@ -19,9 +21,9 @@ require (
cloud.google.com/go/iam v1.3.1 // indirect
cloud.google.com/go/longrunning v0.6.4 // indirect
cloud.google.com/go/monitoring v1.23.0 // indirect
cloud.google.com/go/spanner v1.75.0 // indirect
github.com/GoogleCloudPlatform/grpc-gcp-go/grpcgcp v1.5.2 // indirect
github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.25.0 // indirect
github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/cncf/xds/go v0.0.0-20240905190251-b4127c9b8d78 // indirect
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
pkg/util/xorm/go.sum
+4
View File
@@ -632,6 +632,8 @@ github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kd
github.com/apache/arrow/go/v10 v10.0.1/go.mod h1:YvhnlEePVnBS4+0z3fhPfUy7W1Ikj0Ih0vcRo/gZ1M0=
github.com/apache/arrow/go/v11 v11.0.0/go.mod h1:Eg5OsL5H+e299f7u5ssuXsuHQVEGC4xei5aX110hRiI=
github.com/apache/thrift v0.16.0/go.mod h1:PHK3hniurgQaNMZYaCLEqXKsYK8upmhPbmdP2FXSqgU=
github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4YnC6+E63dPcxHo2sUxDIu8g3QgEJdRY=
github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
github.com/boombuler/barcode v1.0.0/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
github.com/boombuler/barcode v1.0.1/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
@@ -832,6 +834,8 @@ github.com/googleapis/go-sql-spanner v1.11.1 h1:z3ThtKV5HFvaNv9UGc26+ggS+lS0dsCA
github.com/googleapis/go-sql-spanner v1.11.1/go.mod h1:fuA5q4yMS3SZiVfRr5bvksPNk7zUn/irbQW62H/ffZw=
github.com/googleapis/go-type-adapters v1.0.0/go.mod h1:zHW75FOG2aur7gAO2B+MLby+cLsWGBF62rFAi7WjWO4=
github.com/googleapis/google-cloud-go-testing v0.0.0-20200911160855-bcd43fbb19e8/go.mod h1:dvDLG8qkwmyD9a/MJJN3XJcT3xFxOKAvTZGvuZmac9g=
github.com/grafana/grafana v5.4.5+incompatible h1:xNuhSBxLgwDwesuQIAhQu1QCk6tD0TAghKHE36/hxrs=
github.com/grafana/grafana v5.4.5+incompatible/go.mod h1:U8QyUclJHj254BFcuw45p6sg7eeGYX44qn1ShYo5rGE=
github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
github.com/grpc-ecosystem/grpc-gateway/v2 v2.7.0/go.mod h1:hgWBS7lorOAVIJEQMi4ZsPv9hVvWI6+ch50m39Pf2Ks=
github.com/grpc-ecosystem/grpc-gateway/v2 v2.11.3/go.mod h1:o//XUCC/F+yRGJoPO/VU0GSB0f8Nhgmxx0VIRUvaC0w=
pkg/util/xorm/session_insert.go
+29 -46
View File
@@ -438,9 +438,11 @@ func (session *Session) innerInsert(bean any) (int64, error) {
}
if len(table.AutoIncrement) > 0 && session.engine.dialect.DBType() == core.POSTGRES {
if _, err := buf.WriteString(" RETURNING " + session.engine.Quote(table.AutoIncrement)); err != nil {
return 0, err
}
buf.WriteString(" RETURNING " + session.engine.Quote(table.AutoIncrement))
}
if len(table.AutoIncrement) > 0 && session.engine.dialect.DBType() == "spanner" {
buf.WriteString(" THEN RETURN " + session.engine.Quote(table.AutoIncrement))
}
sqlStr := buf.String()
@@ -476,6 +478,7 @@ func (session *Session) innerInsert(bean any) (int64, error) {
// for postgres, many of them didn't implement lastInsertId, so we should
// implemented it ourself.
var insertID, rowsAffected int64
if session.engine.dialect.DBType() == core.ORACLE && len(table.AutoIncrement) > 0 {
res, err := session.queryBytes("select seq_atable.currval from dual", args...)
if err != nil {
@@ -498,23 +501,11 @@ func (session *Session) innerInsert(bean any) (int64, error) {
}
idByte := res[0][table.AutoIncrement]
id, err := strconv.ParseInt(string(idByte), 10, 64)
if err != nil || id <= 0 {
insertID, err = strconv.ParseInt(string(idByte), 10, 64)
if err != nil || insertID <= 0 {
return 1, err
}
aiValue, err := table.AutoIncrColumn().ValueOf(bean)
if err != nil {
session.engine.logger.Error(err)
}
if aiValue == nil || !aiValue.IsValid() || !aiValue.CanSet() {
return 1, nil
}
aiValue.Set(int64ToIntValue(id, aiValue.Type()))
return 1, nil
rowsAffected = 1
} else if len(table.AutoIncrement) > 0 && (session.engine.dialect.DBType() == core.POSTGRES) {
res, err := session.queryBytes(sqlStr, args...)
@@ -537,23 +528,11 @@ func (session *Session) innerInsert(bean any) (int64, error) {
}
idByte := res[0][table.AutoIncrement]
id, err := strconv.ParseInt(string(idByte), 10, 64)
if err != nil || id <= 0 {
insertID, err = strconv.ParseInt(string(idByte), 10, 64)
if err != nil || insertID <= 0 {
return 1, err
}
aiValue, err := table.AutoIncrColumn().ValueOf(bean)
if err != nil {
session.engine.logger.Error(err)
}
if aiValue == nil || !aiValue.IsValid() || !aiValue.CanSet() {
return 1, nil
}
aiValue.Set(int64ToIntValue(id, aiValue.Type()))
return 1, nil
rowsAffected = 1
} else {
res, err := session.exec(sqlStr, args...)
if err != nil {
@@ -575,25 +554,29 @@ func (session *Session) innerInsert(bean any) (int64, error) {
return res.RowsAffected()
}
var id int64
id, err = res.LastInsertId()
if err != nil || id <= 0 {
insertID, err = res.LastInsertId()
if err != nil || insertID <= 0 {
return res.RowsAffected()
}
aiValue, err := table.AutoIncrColumn().ValueOf(bean)
rowsAffected, err = res.RowsAffected()
if err != nil {
session.engine.logger.Error(err)
return 0, err
}
if aiValue == nil || !aiValue.IsValid() || !aiValue.CanSet() {
return res.RowsAffected()
}
aiValue.Set(int64ToIntValue(id, aiValue.Type()))
return res.RowsAffected()
}
// Set insertID back to the bean.
aiValue, err := table.AutoIncrColumn().ValueOf(bean)
if err != nil {
session.engine.logger.Error(err)
}
if aiValue == nil || !aiValue.IsValid() || !aiValue.CanSet() {
return rowsAffected, nil
}
aiValue.Set(int64ToIntValue(insertID, aiValue.Type()))
return rowsAffected, nil
}
// InsertOne insert only one struct into database as a record.
pkg/util/xorm/statement.go
+10 -2
View File
@@ -320,7 +320,11 @@ func (statement *Statement) buildUpdates(bean any,
if err != nil {
engine.logger.Error(err)
} else {
val = data
if col.SQLType.IsText() {
val = string(data)
} else {
val = data
}
}
goto APPEND
}
@@ -331,7 +335,11 @@ func (statement *Statement) buildUpdates(bean any,
if err != nil {
engine.logger.Error(err)
} else {
val = data
if col.SQLType.IsText() {
val = string(data)
} else {
val = data
}
}
goto APPEND
}
pkg/util/xorm/xorm_spanner_test.go
+29
View File
@@ -0,0 +1,29 @@
//go:build enterprise || pro
package xorm
import (
"fmt"
"testing"
"cloud.google.com/go/spanner/spannertest"
_ "github.com/mattn/go-sqlite3"
"github.com/stretchr/testify/require"
)
func TestBasicOperationsWithSpanner(t *testing.T) {
span, err := spannertest.NewServer("localhost:0")
require.NoError(t, err)
defer span.Close()
eng, err := NewEngine("spanner", fmt.Sprintf("%s/projects/test/instances/test/databases/test;usePlainText=true", span.Addr))
require.NoError(t, err)
require.NotNil(t, eng)
require.Equal(t, "spanner", eng.DriverName())
_, err = eng.Exec("CREATE TABLE test_struct (id int64, comment string(max), json string(max)) primary key (id)")
require.NoError(t, err)
// Currently broken because simple INSERT into spannertest doesn't work: https://github.com/googleapis/go-sql-spanner/issues/392
// testBasicOperations(t, eng)
}
pkg/util/xorm/xorm_test.go
+39 -5
View File
@@ -5,13 +5,47 @@ import (
_ "github.com/mattn/go-sqlite3"
"github.com/stretchr/testify/require"
"github.com/grafana/grafana/pkg/components/simplejson"
)
func TestNewEngine(t *testing.T) {
t.Run("successfully create a new engine", func(t *testing.T) {
eng, err := NewEngine("sqlite3", "./test.db")
func TestBasicOperationsWithSqlite(t *testing.T) {
eng, err := NewEngine("sqlite3", ":memory:")
require.NoError(t, err)
require.NotNil(t, eng)
require.Equal(t, "sqlite3", eng.DriverName())
_, err = eng.Exec("CREATE TABLE test_struct (id int primary key, comment text, json text)")
require.NoError(t, err)
testBasicOperations(t, eng)
}
func testBasicOperations(t *testing.T, eng *Engine) {
t.Run("insert object", func(t *testing.T) {
obj := &TestStruct{Comment: "test comment"}
_, err := eng.Insert(obj)
require.NoError(t, err)
require.Equal(t, int64(1), obj.Id)
})
t.Run("update object with json field", func(t *testing.T) {
sess := eng.NewSession()
defer sess.Close()
obj := &TestStruct{Comment: "new comment"}
_, err := sess.Insert(obj)
require.NoError(t, err)
require.NotZero(t, obj.Id)
obj.Json = simplejson.MustJson([]byte(`{"test": "test", "key": null}`))
_, err = sess.Update(obj)
require.NoError(t, err)
require.NotNil(t, eng)
require.Equal(t, "sqlite3", eng.DriverName())
})
}
type TestStruct struct {
Id int64
Comment string
Json *simplejson.Json
}