package db_client import ( "context" "os" "strings" "sync" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/turbot/steampipe/v2/pkg/db/db_common" ) // TestDbClient_SessionRegistration verifies session registration in sessions map func TestDbClient_SessionRegistration(t *testing.T) { client := &DbClient{ sessions: make(map[uint32]*db_common.DatabaseSession), sessionsMutex: &sync.Mutex{}, } // Simulate session registration backendPid := uint32(12345) session := db_common.NewDBSession(backendPid) client.sessionsMutex.Lock() client.sessions[backendPid] = session client.sessionsMutex.Unlock() // Verify session is registered client.sessionsMutex.Lock() registeredSession, found := client.sessions[backendPid] client.sessionsMutex.Unlock() assert.True(t, found, "Session should be registered") assert.Equal(t, backendPid, registeredSession.BackendPid, "Backend PID should match") } // TestDbClient_SessionUnregistration verifies session cleanup via BeforeClose func TestDbClient_SessionUnregistration(t *testing.T) { client := &DbClient{ sessions: make(map[uint32]*db_common.DatabaseSession), sessionsMutex: &sync.Mutex{}, } // Add sessions backendPid1 := uint32(100) backendPid2 := uint32(200) client.sessionsMutex.Lock() client.sessions[backendPid1] = db_common.NewDBSession(backendPid1) client.sessions[backendPid2] = db_common.NewDBSession(backendPid2) client.sessionsMutex.Unlock() assert.Len(t, client.sessions, 2, "Should have 2 sessions") // Simulate BeforeClose callback for one session client.sessionsMutex.Lock() delete(client.sessions, backendPid1) client.sessionsMutex.Unlock() // Verify only one session remains client.sessionsMutex.Lock() _, found1 := client.sessions[backendPid1] _, found2 := client.sessions[backendPid2] client.sessionsMutex.Unlock() assert.False(t, found1, "First session should be removed") assert.True(t, found2, "Second session should still exist") assert.Len(t, client.sessions, 1, "Should have 1 session remaining") } // TestDbClient_ConcurrentSessionRegistration tests concurrent session additions func TestDbClient_ConcurrentSessionRegistration(t *testing.T) { if testing.Short() { t.Skip("Skipping concurrent test in short mode") } client := &DbClient{ sessions: make(map[uint32]*db_common.DatabaseSession), sessionsMutex: &sync.Mutex{}, } var wg sync.WaitGroup numGoroutines := 100 // Concurrently add sessions for i := 0; i < numGoroutines; i++ { wg.Add(1) go func(id uint32) { defer wg.Done() backendPid := id session := db_common.NewDBSession(backendPid) client.sessionsMutex.Lock() client.sessions[backendPid] = session client.sessionsMutex.Unlock() }(uint32(i)) } wg.Wait() // Verify all sessions were added assert.Len(t, client.sessions, numGoroutines, "All sessions should be registered") } // TestDbClient_SessionMapGrowthUnbounded tests for potential memory leaks // This verifies that sessions don't accumulate indefinitely func TestDbClient_SessionMapGrowthUnbounded(t *testing.T) { if testing.Short() { t.Skip("Skipping large dataset test in short mode") } client := &DbClient{ sessions: make(map[uint32]*db_common.DatabaseSession), sessionsMutex: &sync.Mutex{}, } // Simulate many connections numSessions := 10000 for i := 0; i < numSessions; i++ { backendPid := uint32(i) session := db_common.NewDBSession(backendPid) client.sessionsMutex.Lock() client.sessions[backendPid] = session client.sessionsMutex.Unlock() } assert.Len(t, client.sessions, numSessions, "Should have all sessions") // Simulate cleanup (BeforeClose callbacks) for i := 0; i < numSessions; i++ { backendPid := uint32(i) client.sessionsMutex.Lock() delete(client.sessions, backendPid) client.sessionsMutex.Unlock() } // Verify all sessions are cleaned up assert.Len(t, client.sessions, 0, "All sessions should be cleaned up") } // TestDbClient_SearchPathUpdates verifies session search path management func TestDbClient_SearchPathUpdates(t *testing.T) { client := &DbClient{ sessions: make(map[uint32]*db_common.DatabaseSession), sessionsMutex: &sync.Mutex{}, customSearchPath: []string{"schema1", "schema2"}, } // Add a session backendPid := uint32(12345) session := db_common.NewDBSession(backendPid) client.sessionsMutex.Lock() client.sessions[backendPid] = session client.sessionsMutex.Unlock() // Verify custom search path is set assert.NotNil(t, client.customSearchPath, "Custom search path should be set") assert.Len(t, client.customSearchPath, 2, "Should have 2 schemas in search path") } // TestSearchPathAccessShouldUseReadLocks checks that search path access does not block other goroutines unnecessarily. // // Holding an exclusive mutex during search-path reads in concurrent query setup can deadlock when // another goroutine is setting the path. The current code uses Lock/Unlock; this test documents // the expectation to move to a read/non-blocking lock so concurrent reads are safe. func TestSearchPathAccessShouldUseReadLocks(t *testing.T) { content, err := os.ReadFile("db_client_search_path.go") require.NoError(t, err, "should be able to read db_client_search_path.go") source := string(content) assert.Contains(t, source, "GetRequiredSessionSearchPath", "getter must exist") assert.Contains(t, source, "searchPathMutex", "getter must guard access to searchPath state") // Expect a read or non-blocking lock in getters; fail if only full Lock/Unlock is present. hasRLock := strings.Contains(source, "RLock") hasTry := strings.Contains(source, "TryLock") || strings.Contains(source, "tryLock") if !hasRLock && !hasTry { t.Fatalf("GetRequiredSessionSearchPath should avoid exclusive Lock/Unlock to prevent deadlocks under concurrent query setup") } } // TestDbClient_SessionConnectionNilSafety verifies handling of nil connections func TestDbClient_SessionConnectionNilSafety(t *testing.T) { session := db_common.NewDBSession(12345) // Session is created with nil connection initially assert.Nil(t, session.Connection, "New session should have nil connection initially") } // TestDbClient_SessionSearchPathUpdatesThreadSafe verifies that concurrent access // to customSearchPath does not cause data races. // Reference: https://github.com/turbot/steampipe/issues/4792 // // This test simulates concurrent goroutines accessing and modifying the customSearchPath // slice. Without proper synchronization, this causes a data race. // // Run with: go test -race -run TestDbClient_SessionSearchPathUpdatesThreadSafe func TestDbClient_SessionSearchPathUpdatesThreadSafe(t *testing.T) { // Create a DbClient with the fields we need for testing client := &DbClient{ customSearchPath: []string{"public", "internal"}, userSearchPath: []string{"public"}, searchPathMutex: &sync.RWMutex{}, } // Number of concurrent operations to test const numGoroutines = 100 var wg sync.WaitGroup wg.Add(numGoroutines * 3) // Simulate concurrent readers calling GetRequiredSessionSearchPath for i := 0; i < numGoroutines; i++ { go func() { defer wg.Done() _ = client.GetRequiredSessionSearchPath() }() } // Simulate concurrent readers calling GetCustomSearchPath for i := 0; i < numGoroutines; i++ { go func() { defer wg.Done() _ = client.GetCustomSearchPath() }() } // Simulate concurrent writers calling SetRequiredSessionSearchPath // This is the most dangerous operation as it modifies the slice for i := 0; i < numGoroutines; i++ { go func() { defer wg.Done() ctx := context.Background() // This will write to customSearchPath _ = client.SetRequiredSessionSearchPath(ctx) }() } wg.Wait() }