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

     1  package db_common
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"log"
     7  	"sync"
     8  
     9  	"github.com/turbot/steampipe-plugin-sdk/v5/sperr"
    10  
    11  	"github.com/jackc/pgx/v5"
    12  	"github.com/jackc/pgx/v5/pgconn"
    13  	"github.com/turbot/steampipe/pkg/constants"
    14  	"github.com/turbot/steampipe/pkg/error_helpers"
    15  )
    16  
    17  type NotificationListener struct {
    18  	notifications []*pgconn.Notification
    19  	conn          *pgx.Conn
    20  
    21  	onNotification func(*pgconn.Notification)
    22  	mut            sync.Mutex
    23  	cancel         context.CancelFunc
    24  }
    25  
    26  func NewNotificationListener(ctx context.Context, conn *pgx.Conn) (*NotificationListener, error) {
    27  	if conn == nil {
    28  		return nil, sperr.New("nil connection passed to NewNotificationListener")
    29  	}
    30  
    31  	listener := &NotificationListener{conn: conn}
    32  
    33  	// tell the connection to listen to notifications
    34  	listenSql := fmt.Sprintf("listen %s", constants.PostgresNotificationChannel)
    35  	_, err := conn.Exec(ctx, listenSql)
    36  	if err != nil {
    37  		log.Printf("[INFO] Error listening to notification channel: %s", err)
    38  		conn.Close(ctx)
    39  		return nil, err
    40  	}
    41  
    42  	// create cancel context to shutdown the listener
    43  	cancelCtx, cancel := context.WithCancel(ctx)
    44  	listener.cancel = cancel
    45  
    46  	// start the goroutine to listen
    47  	listener.listenToPgNotificationsAsync(cancelCtx)
    48  
    49  	return listener, nil
    50  }
    51  
    52  func (c *NotificationListener) Stop(ctx context.Context) {
    53  	c.conn.Close(ctx)
    54  	// stop the listener goroutine
    55  	c.cancel()
    56  }
    57  
    58  func (c *NotificationListener) RegisterListener(onNotification func(*pgconn.Notification)) {
    59  	c.mut.Lock()
    60  	defer c.mut.Unlock()
    61  
    62  	c.onNotification = onNotification
    63  	// send any notifications we have already collected
    64  	for _, n := range c.notifications {
    65  		onNotification(n)
    66  	}
    67  	// clear notifications
    68  	c.notifications = nil
    69  }
    70  
    71  func (c *NotificationListener) listenToPgNotificationsAsync(ctx context.Context) {
    72  	log.Printf("[INFO] notificationListener listenToPgNotificationsAsync")
    73  
    74  	go func() {
    75  		for ctx.Err() == nil {
    76  			log.Printf("[INFO] Wait for notification")
    77  			notification, err := c.conn.WaitForNotification(ctx)
    78  			if err != nil && !error_helpers.IsContextCancelledError(err) {
    79  				log.Printf("[WARN] Error waiting for notification: %s", err)
    80  				return
    81  			}
    82  
    83  			if notification != nil {
    84  				log.Printf("[INFO] got notification")
    85  				c.mut.Lock()
    86  				// if we have a callback, call it
    87  				if c.onNotification != nil {
    88  					log.Printf("[INFO] call notification handler")
    89  					c.onNotification(notification)
    90  				} else {
    91  					// otherwise cache the notification
    92  					log.Printf("[INFO] cache notification")
    93  					c.notifications = append(c.notifications, notification)
    94  				}
    95  				c.mut.Unlock()
    96  				log.Printf("[INFO] Handled notification")
    97  			}
    98  		}
    99  	}()
   100  
   101  	log.Printf("[TRACE] InteractiveClient listenToPgNotificationsAsync DONE")
   102  }