diff --git a/xray/sqlcontext_test.go b/xray/sqlcontext_test.go index 1d9d1861..3f144fec 100644 --- a/xray/sqlcontext_test.go +++ b/xray/sqlcontext_test.go @@ -9,6 +9,8 @@ package xray import ( + "context" + "database/sql/driver" "encoding/json" "errors" "testing" @@ -44,6 +46,24 @@ func mockOracle(mock sqlmock.Sqlmock, err error) { mock.ExpectPrepare(`SELECT version FROM v\$instance UNION SELECT user, ora_database_name FROM dual`).ExpectQuery().WillReturnRows(row) } +func mockSnowflakeDB(mock sqlmock.Sqlmock, err error) { + row := sqlmock.NewRows([]string{"current_version()", "current_user()", "current_database()"}). + AddRow("test version", "test user", "test database"). + RowError(0, err) + mock.ExpectPrepare(`SELECT current_version\(\), current_user\(\), current_database\(\)`).ExpectQuery().WillReturnRows(row) +} + +func registerSnowflakeDriver() { + RegisterSQLDetector("snowflake", func(ctx context.Context, conn driver.Conn, attr *DBAttribute) error { + attr.DatabaseType = "Snowflake" + return queryRow( + ctx, conn, + "SELECT current_version(), current_user(), current_database()", + &attr.DatabaseVersion, &attr.User, &attr.DBname, + ) + }) +} + func capturePing(dsn string) (*Segment, error) { ctx, td := NewTestDaemon() defer td.Close()