diff --git a/_integration-tests/tests/sql/sql.go b/_integration-tests/tests/sql/sql.go index b65e5224..ca044691 100644 --- a/_integration-tests/tests/sql/sql.go +++ b/_integration-tests/tests/sql/sql.go @@ -20,26 +20,38 @@ import ( type TestCase struct { *sql.DB + // untraced is the un-traced database connection, used to set up the test case + // without producing unnecessary spans. It is retained as a [TestCase] field + // and cleaned up using [t.Cleanup] to ensure the shared cache is not lost, as + // it stops existing once the last DB connection using it is closed. + untraced *sql.DB } func (tc *TestCase) Setup(t *testing.T) { + const ( + dn = "sqlite3" + dsn = "file::memory:?cache=shared" + ) + var err error - tc.DB, err = sql.Open("sqlite3", "file::memory:") + + //orchestrion:ignore + tc.untraced, err = sql.Open(dn, dsn) require.NoError(t, err) - t.Cleanup(func() { - assert.NoError(t, tc.DB.Close()) - }) + t.Cleanup(func() { require.NoError(t, tc.untraced.Close()) }) + + ctx := context.Background() - _, err = tc.DB.ExecContext(context.Background(), + _, err = tc.untraced.ExecContext(ctx, `CREATE TABLE IF NOT EXISTS notes ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - userid INTEGER, - content STRING, - created STRING -)`) + id INTEGER PRIMARY KEY AUTOINCREMENT, + userid INTEGER, + content STRING, + created STRING + )`) require.NoError(t, err) - _, err = tc.DB.ExecContext(context.Background(), + _, err = tc.untraced.ExecContext(ctx, `INSERT OR REPLACE INTO notes(userid, content, created) VALUES (1, 'Hello, John. This is John. You are leaving a note for yourself. You are welcome and thank you.', datetime('now')), (1, 'Hey, remember to mow the lawn.', datetime('now')), @@ -47,8 +59,14 @@ func (tc *TestCase) Setup(t *testing.T) { (2, 'Opportunities don''t happen, you create them.', datetime('now')), (3, 'Pick up cabbage from the store on the way home.', datetime('now')), (3, 'Review PR #1138', datetime('now') - );`) + )`) require.NoError(t, err) + + tc.DB, err = sql.Open(dn, dsn) + require.NoError(t, err) + t.Cleanup(func() { + assert.NoError(t, tc.DB.Close()) + }) } func (tc *TestCase) Run(t *testing.T) { diff --git a/internal/injector/builtin/generated.go b/internal/injector/builtin/generated.go index efd27118..fa4f8e3f 100644 --- a/internal/injector/builtin/generated.go +++ b/internal/injector/builtin/generated.go @@ -1564,7 +1564,14 @@ var Aspects = [...]aspect.Aspect{ { JoinPoint: join.FunctionCall("database/sql", "Register"), Advice: []advice.Advice{ - advice.ReplaceFunction("gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql", "Register"), + advice.WrapExpression(code.MustTemplate( + "func(driverName string, driver driver.Driver) {\n sql.Register(driverName, driver)\n sqltrace.Register(driverName, driver)\n}({{ index .AST.Args 0 }}, {{ index .AST.Args 1 }})", + map[string]string{ + "sql": "database/sql", + "sqltrace": "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql", + }, + context.GoLangVersion{}, + )), }, }, { @@ -1815,4 +1822,4 @@ var InjectedPaths = [...]string{ } // Checksum is a checksum of the built-in configuration which can be used to invalidate caches. -const Checksum = "sha512:MJ4+HJMaOH/ACkbYHyW+erCE7e7moqEflMCs1vOIMiUQv4nkKGASjiZOukrnylyCtDrtAk/InP376zusJcmapw==" +const Checksum = "sha512:6M7isTPZdLAR2kLE936vCfN6doV2jpSSzCYkLZTAtwd3kXfbPk8IWcCqjDeKMRM4em9YzcIApY3rEMmsxMDODA==" diff --git a/internal/injector/builtin/testdata/server/database.go.snap b/internal/injector/builtin/testdata/server/database.go.snap index 5ec346ed..c9fae181 100644 --- a/internal/injector/builtin/testdata/server/database.go.snap +++ b/internal/injector/builtin/testdata/server/database.go.snap @@ -18,7 +18,13 @@ import ( //line samples/server/database.go:16 func init() { - __orchestrion_sql.Register("test", &testDriver{}) +//line + func(driverName string, driver driver.Driver) { + sql.Register(driverName, driver) + __orchestrion_sql.Register(driverName, driver) + }( +//line samples/server/database.go:17 + "test", &testDriver{}) } type testDriver struct{} diff --git a/internal/injector/builtin/yaml/stdlib/database-sql.yml b/internal/injector/builtin/yaml/stdlib/database-sql.yml index 27ecf5cb..86ce2402 100644 --- a/internal/injector/builtin/yaml/stdlib/database-sql.yml +++ b/internal/injector/builtin/yaml/stdlib/database-sql.yml @@ -14,7 +14,15 @@ aspects: join-point: function-call: database/sql.Register advice: - - replace-function: gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql.Register + - wrap-expression: + imports: + sqltrace: gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql + sql: database/sql + template: |- + func(driverName string, driver driver.Driver) { + sql.Register(driverName, driver) + sqltrace.Register(driverName, driver) + }({{ index .AST.Args 0 }}, {{ index .AST.Args 1 }}) - id: Open join-point: