github.com/turbot/steampipe@v1.7.0-rc.0.0.20240517123944-7cef272d4458/pkg/db/db_common/db_session.go (about)

     1  package db_common
     2  
     3  import (
     4  	"log"
     5  	"time"
     6  
     7  	"github.com/jackc/pgx/v5/pgxpool"
     8  )
     9  
    10  // DatabaseSession wraps over the raw database connection
    11  // the purpose is to be able
    12  //   - to store the current search path of the connection without having to make a database round-trip
    13  //   - To store the last scan_metadata id used on this connection
    14  type DatabaseSession struct {
    15  	BackendPid uint32   `json:"backend_pid"`
    16  	SearchPath []string `json:"-"`
    17  
    18  	// this gets rewritten, since the database/sql gives back a new instance everytime
    19  	Connection *pgxpool.Conn `json:"-"`
    20  }
    21  
    22  func NewDBSession(backendPid uint32) *DatabaseSession {
    23  	return &DatabaseSession{
    24  		BackendPid: backendPid,
    25  	}
    26  }
    27  
    28  func (s *DatabaseSession) Close(waitForCleanup bool) {
    29  	if s.Connection != nil {
    30  		if waitForCleanup {
    31  			log.Printf("[TRACE] DatabaseSession.Close wait for connection cleanup")
    32  			select {
    33  			case <-time.After(5 * time.Second):
    34  				log.Printf("[TRACE] DatabaseSession.Close timed out waiting for connection cleanup")
    35  			case <-s.Connection.Conn().PgConn().CleanupDone():
    36  				log.Printf("[TRACE] DatabaseSession.Close connection cleanup complete")
    37  			}
    38  		}
    39  		s.Connection.Release()
    40  	}
    41  	s.Connection = nil
    42  
    43  }