github.com/turbot/steampipe@v1.7.0-rc.0.0.20240517123944-7cef272d4458/pkg/connection/refresh_connections_state.go (about) 1 package connection 2 3 import ( 4 "context" 5 "fmt" 6 "log" 7 "os" 8 "strconv" 9 "strings" 10 "sync" 11 "time" 12 13 "github.com/jackc/pgx/v5/pgxpool" 14 "github.com/spf13/cobra" 15 "github.com/spf13/viper" 16 "github.com/turbot/go-kit/helpers" 17 "github.com/turbot/steampipe-plugin-sdk/v5/grpc/proto" 18 "github.com/turbot/steampipe-plugin-sdk/v5/plugin" 19 "github.com/turbot/steampipe-plugin-sdk/v5/sperr" 20 "github.com/turbot/steampipe/pkg/constants" 21 "github.com/turbot/steampipe/pkg/db/db_common" 22 "github.com/turbot/steampipe/pkg/db/db_local" 23 "github.com/turbot/steampipe/pkg/error_helpers" 24 "github.com/turbot/steampipe/pkg/introspection" 25 "github.com/turbot/steampipe/pkg/steampipeconfig" 26 "github.com/turbot/steampipe/pkg/utils" 27 "golang.org/x/exp/maps" 28 "golang.org/x/sync/semaphore" 29 ) 30 31 type connectionError struct { 32 name string 33 err error 34 } 35 36 type refreshConnectionState struct { 37 // a connection pool to the DB service which uses the server appname 38 pool *pgxpool.Pool 39 searchPath []string 40 connectionUpdates *steampipeconfig.ConnectionUpdates 41 tableUpdater *connectionStateTableUpdater 42 res *steampipeconfig.RefreshConnectionResult 43 forceUpdateConnectionNames []string 44 // properties for schema/comment cloning 45 exemplarSchemaMapMut sync.Mutex 46 47 // maps keyed by plugin which gives an exemplar connection name, 48 // if a plugin has an entry in this map, all connections schemas can be cloned from the exemplar schema 49 exemplarSchemaMap map[string]string 50 // if a plugin has an entry in this map, all connections schemas can be cloned from the exemplar schema 51 exemplarCommentsMap map[string]string 52 pluginManager pluginManager 53 } 54 55 func newRefreshConnectionState(ctx context.Context, pluginManager pluginManager, forceUpdateConnectionNames []string) (*refreshConnectionState, error) { 56 log.Println("[DEBUG] newRefreshConnectionState start") 57 defer log.Println("[DEBUG] newRefreshConnectionState end") 58 59 pool := pluginManager.Pool() 60 // set user search path first 61 log.Printf("[INFO] setting up search path") 62 searchPath, err := db_local.SetUserSearchPath(ctx, pool) 63 if err != nil { 64 return nil, err 65 } 66 67 res := &refreshConnectionState{ 68 pool: pool, 69 searchPath: searchPath, 70 forceUpdateConnectionNames: forceUpdateConnectionNames, 71 pluginManager: pluginManager, 72 } 73 74 return res, nil 75 } 76 77 // RefreshConnections loads required connections from config 78 // and update the database schema and search path to reflect the required connections 79 // return whether any changes have been made 80 func (s *refreshConnectionState) refreshConnections(ctx context.Context) { 81 log.Println("[DEBUG] refreshConnectionState.refreshConnections start") 82 defer log.Println("[DEBUG] refreshConnectionState.refreshConnections end") 83 // if there was an error (other than a connection error, which will NOT have been assigned to res), 84 // set state of all incomplete connections to error 85 defer func() { 86 if s.res != nil { 87 if s.res.Error != nil { 88 s.setIncompleteConnectionStateToError(ctx, sperr.WrapWithMessage(s.res.Error, "refreshConnections failed before connection update was complete")) 89 } 90 if !s.res.ErrorAndWarnings.Empty() { 91 log.Printf("[INFO] refreshConnections completed with errors, sending notification") 92 s.pluginManager.SendPostgresErrorsAndWarningsNotification(ctx, s.res.ErrorAndWarnings) 93 } 94 95 } 96 }() 97 log.Printf("[INFO] building connectionUpdates") 98 99 var opts []steampipeconfig.ConnectionUpdatesOption 100 if len(s.forceUpdateConnectionNames) > 0 { 101 opts = append(opts, steampipeconfig.WithForceUpdate(s.forceUpdateConnectionNames)) 102 } 103 104 // build a ConnectionUpdates struct 105 // this determines any necessary connection updates and starts any necessary plugins 106 s.connectionUpdates, s.res = steampipeconfig.NewConnectionUpdates(ctx, s.pool, s.pluginManager, opts...) 107 108 defer s.logRefreshConnectionResults() 109 // were we successful? 110 if s.res.Error != nil { 111 return 112 } 113 114 // if any connections in the final state are in error, that may mean we failed to start them 115 // - update the connection state table 116 if err := s.setFailedConnectionsToError(ctx); err != nil { 117 s.res.Error = err 118 return 119 } 120 121 log.Printf("[INFO] created connectionUpdates") 122 123 // reload plugin rate limiter definitions for all plugins which are updated - the plugin will already be loaded 124 // also repopulate the plugin column table 125 if err := s.updateRateLimiterDefinitions(ctx); err != nil { 126 s.res.Error = err 127 return 128 } 129 130 // update the plugin column table, based on connection updates and plugins with updated binaries 131 if err := s.updatePluginColumnTable(ctx); err != nil { 132 s.res.Error = err 133 return 134 } 135 136 // delete the connection state file - it will be rewritten when we are complete 137 log.Printf("[INFO] deleting connections state file") 138 steampipeconfig.DeleteConnectionStateFile() 139 defer func() { 140 if s.res.Error == nil { 141 log.Printf("[INFO] saving connections state file") 142 steampipeconfig.SaveConnectionStateFile(s.res, s.connectionUpdates) 143 } 144 }() 145 146 // warn about missing plugins 147 s.addMissingPluginWarnings() 148 149 // create object to update the connection state table and notify of state changes 150 s.tableUpdater = newConnectionStateTableUpdater(s.connectionUpdates, s.pool) 151 152 // NOTE: delete any DYNAMIC plugin connections which will be updated 153 // to avoid them being accessed before they are updated 154 if err := s.executeDeleteQueries(ctx, s.connectionUpdates.DynamicUpdates()); err != nil { 155 s.res.Error = err 156 return 157 } 158 159 // update connectionState table to reflect the updates (i.e. set connections to updating/deleting/ready as appropriate) 160 // also this will update the schema hashes of plugins 161 if err := s.tableUpdater.start(ctx); err != nil { 162 s.res.Error = err 163 return 164 } 165 166 // if there are no updates, just return 167 if !s.connectionUpdates.HasUpdates() { 168 log.Println("[INFO] no updates required") 169 return 170 } 171 172 log.Printf("[INFO] execute connection queries") 173 174 // execute any necessary queries 175 s.executeConnectionQueries(ctx) 176 if s.res.Error != nil { 177 log.Printf("[WARN] refreshConnections failed with err %s", s.res.Error.Error()) 178 return 179 } 180 181 s.res.UpdatedConnections = true 182 } 183 184 func (s *refreshConnectionState) setFailedConnectionsToError(ctx context.Context) error { 185 conn, err := s.pool.Acquire(ctx) 186 if err != nil { 187 return sperr.WrapWithMessage(err, "failed to update connection state table") 188 } 189 defer conn.Release() 190 191 for _, c := range s.connectionUpdates.FinalConnectionState { 192 if c.State == constants.ConnectionStateError { 193 if err := s.tableUpdater.onConnectionError(ctx, conn.Conn(), c.ConnectionName, fmt.Errorf(c.Error())); err != nil { 194 return sperr.WrapWithMessage(err, "failed to update connection state table") 195 } 196 } 197 } 198 return nil 199 } 200 201 // if any plugin binaries have changed update the rate limiter definitions 202 func (s *refreshConnectionState) updateRateLimiterDefinitions(ctx context.Context) error { 203 if len(s.connectionUpdates.PluginsWithUpdatedBinary) == 0 { 204 return nil 205 } 206 207 updatedPluginLimiters, err := s.pluginManager.LoadPluginRateLimiters(s.connectionUpdates.PluginsWithUpdatedBinary) 208 209 if err != nil { 210 return err 211 } 212 213 if len(updatedPluginLimiters) > 0 { 214 err := s.pluginManager.HandlePluginLimiterChanges(updatedPluginLimiters) 215 if err != nil { 216 s.pluginManager.SendPostgresErrorsAndWarningsNotification(ctx, error_helpers.NewErrorsAndWarning(err)) 217 } 218 } 219 return nil 220 } 221 222 // if any plugin binaries have changed update the plugin column table 223 func (s *refreshConnectionState) updatePluginColumnTable(ctx context.Context) error { 224 var deletedPlugins []string 225 var updatedPlugins = map[string]*proto.Schema{} 226 227 currentPluginConnectionMap := s.connectionUpdates.CurrentConnectionState.GetPluginToConnectionMap() 228 finalPluginConnectionMap := s.connectionUpdates.FinalConnectionState.GetPluginToConnectionMap() 229 230 // add into plugin column table any plugins which have connections for the first time 231 for _, connectionState := range s.connectionUpdates.Update { 232 connectionName := connectionState.ConnectionName 233 if connectionState.SchemaMode == plugin.SchemaModeDynamic { 234 // plugin column table only supports static for now 235 continue 236 } 237 p := connectionState.Plugin 238 if _, ok := currentPluginConnectionMap[p]; !ok { 239 updatedPlugins[p] = s.connectionUpdates.ConnectionPlugins[connectionName].ConnectionMap[connectionName].Schema 240 } 241 } 242 243 // remove from plugin column table any plugins which have no connections 244 for connectionName := range s.connectionUpdates.Delete { 245 // get plugin for this connection 246 connectionState, ok := s.connectionUpdates.CurrentConnectionState[connectionName] 247 if !ok { 248 continue 249 } 250 251 p := connectionState.Plugin 252 if _, ok := finalPluginConnectionMap[p]; !ok { 253 deletedPlugins = append(deletedPlugins, p) 254 } 255 } 256 257 // update plugin column table for any plugins which have updated binaries 258 for p, connectionName := range s.connectionUpdates.PluginsWithUpdatedBinary { 259 // do we actually have a connection plugin for this plugin? 260 if connectionPlugin, ok := s.connectionUpdates.ConnectionPlugins[connectionName]; ok { 261 updatedPlugins[p] = connectionPlugin.ConnectionMap[connectionName].Schema 262 } 263 } 264 265 return s.pluginManager.UpdatePluginColumnsTable(ctx, updatedPlugins, deletedPlugins) 266 267 } 268 269 func (s *refreshConnectionState) addMissingPluginWarnings() { 270 log.Printf("[INFO] refreshConnections: identify missing plugins") 271 272 var connectionNames []string 273 // add warning if there are connections left over, from missing plugins 274 if len(s.connectionUpdates.MissingPlugins) > 0 { 275 // warning 276 for _, conns := range s.connectionUpdates.MissingPlugins { 277 for _, con := range conns { 278 connectionNames = append(connectionNames, con.Name) 279 } 280 281 } 282 pluginNames := maps.Keys(s.connectionUpdates.MissingPlugins) 283 284 s.res.AddWarning(fmt.Sprintf("%d %s required by %d %s %s missing. To install, please run: %s", 285 len(pluginNames), 286 utils.Pluralize("plugin", len(pluginNames)), 287 len(connectionNames), 288 utils.Pluralize("connection", len(connectionNames)), 289 utils.Pluralize("is", len(pluginNames)), 290 constants.Bold(fmt.Sprintf("steampipe plugin install %s", strings.Join(pluginNames, " "))))) 291 } 292 } 293 294 func (s *refreshConnectionState) logRefreshConnectionResults() { 295 var cmdName = viper.Get(constants.ConfigKeyActiveCommand).(*cobra.Command).Name() 296 if cmdName != "plugin-manager" { 297 return 298 } 299 300 var op strings.Builder 301 if s.connectionUpdates != nil { 302 op.WriteString(s.connectionUpdates.String()) 303 } 304 if s.res != nil { 305 op.WriteString(fmt.Sprintf("%s\n", s.res.String())) 306 } 307 308 log.Printf("[TRACE] refresh connections: \n%s\n", helpers.Tabify(op.String(), " ")) 309 } 310 311 func (s *refreshConnectionState) executeConnectionQueries(ctx context.Context) { 312 log.Println("[DEBUG] refreshConnectionState.executeConnectionQueries start") 313 defer log.Println("[DEBUG] refreshConnectionState.executeConnectionQueries end") 314 315 // execute deletions 316 if err := s.executeDeleteQueries(ctx, s.connectionUpdates.GetConnectionsToDelete()); err != nil { 317 // just log 318 log.Printf("[WARN] failed to delete all unused schemas: %s", err.Error()) 319 } 320 321 // execute updates 322 numUpdates := len(s.connectionUpdates.Update) 323 numMissingComments := len(s.connectionUpdates.MissingComments) 324 log.Printf("[INFO] executeConnectionQueries: num updates: %d, connections missing comments: %d", numUpdates, numMissingComments) 325 326 if numUpdates+numMissingComments > 0 { 327 // get schema queries - this updates schemas for validated plugins and drops schemas for unvalidated plugins 328 s.executeUpdateQueries(ctx) 329 // done 330 return 331 } 332 333 if len(s.connectionUpdates.Delete) > 0 { 334 log.Printf("[INFO] deleted all unnecessary schemas - sending notification") 335 336 // if there are no updates and there ARE deletes, notify 337 // (is there are updates, deletes will be notified by executeUpdateQueries) 338 if err := s.pluginManager.SendPostgresSchemaNotification(ctx); err != nil { 339 // just log 340 log.Printf("[WARN] failed to send schema deletion Postgres notification: %s", err.Error()) 341 } 342 } 343 } 344 345 // execute all update queries 346 // NOTE: this only sets res.Error if there is a failure to set update the connection state table 347 // - all other connection based failures are recorded in the connection state table 348 func (s *refreshConnectionState) executeUpdateQueries(ctx context.Context) { 349 log.Println("[DEBUG] refreshConnectionState.executeUpdateQueries start") 350 defer log.Println("[DEBUG] refreshConnectionState.executeUpdateQueries end") 351 352 defer func() { 353 if s.res.Error != nil { 354 log.Printf("[INFO] executeUpdateQueries returned error: %v", s.res.Error) 355 } 356 }() 357 358 connectionUpdates := s.connectionUpdates 359 connectionPlugins := connectionUpdates.ConnectionPlugins 360 numUpdates := len(connectionUpdates.Update) 361 362 // we need to execute the updates in search path order 363 // i.e. we first need to update the first search path connection for each plugin (this can be done in parallel) 364 // then we can update the remaining connections in parallel 365 initialUpdates, remainingUpdates, dynamicUpdates := s.getInitialAndRemainingUpdates() 366 367 // dynamic plugins must be updated for each plugin in search path order 368 // dynamicUpdates is a map keyed by plugin with all the updates for that plugin 369 370 // create exemplar maps 371 s.exemplarSchemaMap = make(map[string]string) 372 s.exemplarCommentsMap = make(map[string]string) 373 log.Printf("[INFO] executing %d update %s", numUpdates, utils.Pluralize("query", numUpdates)) 374 375 // execute initial updates 376 log.Printf("[INFO] executing initial updates") 377 var errors []error 378 moreErrors := s.executeUpdatesInParallel(ctx, initialUpdates) 379 errors = append(errors, moreErrors...) 380 381 // execute dynamic updates (note, we update all connections in search path order, 382 // so must call executeUpdateSetsInParallel) 383 log.Printf("[INFO] executing dynamic updates") 384 moreErrors = s.executeUpdateSetsInParallel(ctx, dynamicUpdates) 385 errors = append(errors, moreErrors...) 386 387 // if any of the initial schemas failed, do not proceed - these schemas are required to ensure we correctly 388 // resolve unqualified queries/tables 389 if len(errors) > 0 { 390 s.res.Error = error_helpers.CombineErrors(errors...) 391 log.Printf("[WARN] initial updates failed: %s", s.res.Error.Error()) 392 return 393 } 394 395 log.Printf("[INFO] set comments for initial updates") 396 // now set comments for initial updates and dynamic connections 397 // note errors will be empty to get here 398 s.UpdateCommentsInParallel(ctx, maps.Values(initialUpdates), connectionPlugins) 399 400 log.Printf("[INFO] set comments for dynamic updates") 401 // convert dynamicUpdates to an array of connection states 402 var dynamicUpdateArray = updateSetMapToArray(dynamicUpdates) 403 s.UpdateCommentsInParallel(ctx, dynamicUpdateArray, connectionPlugins) 404 405 log.Printf("[INFO] updated all exemplar schemas - sending notification") 406 // now that we have updated all exemplar schemars, send postgres notification 407 // this gives any attached interactive clients a chance to update their inspect data and autocomplete 408 if err := s.pluginManager.SendPostgresSchemaNotification(ctx); err != nil { 409 // just log 410 log.Printf("[WARN] failed to send schem update Postgres notification: %s", err.Error()) 411 } 412 413 log.Printf("[INFO] Execute %d remaining %s", 414 len(remainingUpdates), 415 utils.Pluralize("updates", len(remainingUpdates))) 416 // now execute remaining updates 417 moreErrors = s.executeUpdatesInParallel(ctx, remainingUpdates) 418 errors = append(errors, moreErrors...) 419 420 log.Printf("[INFO] Set comments for %d remaining %s and %d %s missing comments", 421 len(remainingUpdates), 422 utils.Pluralize("updates", len(remainingUpdates)), 423 len(connectionUpdates.MissingComments), 424 utils.Pluralize("updates", len(connectionUpdates.MissingComments)), 425 ) 426 // set comments for remaining updates 427 s.UpdateCommentsInParallel(ctx, maps.Values(remainingUpdates), connectionPlugins) 428 // set comments for any other connection without comment set 429 s.UpdateCommentsInParallel(ctx, maps.Values(s.connectionUpdates.MissingComments), connectionPlugins) 430 431 if len(errors) > 0 { 432 s.res.Error = error_helpers.CombineErrors(errors...) 433 } 434 435 log.Printf("[INFO] all update queries executed") 436 437 for _, failure := range connectionUpdates.InvalidConnections { 438 log.Printf("[TRACE] remove schema for connection failing validation connection %s, plugin Name %s\n ", failure.ConnectionName, failure.Plugin) 439 if failure.ShouldDropIfExists { 440 _, err := s.pool.Exec(ctx, db_common.GetDeleteConnectionQuery(failure.ConnectionName)) 441 if err != nil { 442 // NOTE: do not return an error if we fail to remove an invalid connection - just log it 443 log.Printf("[WARN] failed to delete invalid connection '%s' (%s) : %s", failure.ConnectionName, failure.Message, err.Error()) 444 } 445 } 446 } 447 log.Printf("[INFO] executeUpdateQueries complete") 448 return 449 } 450 451 // convert map update sets (used for dynamic schemas) to an array of the underlying connection states 452 func updateSetMapToArray(updateSetMap map[string][]*steampipeconfig.ConnectionState) []*steampipeconfig.ConnectionState { 453 var res []*steampipeconfig.ConnectionState 454 for _, updates := range updateSetMap { 455 res = append(res, updates...) 456 } 457 return res 458 } 459 460 // create/update connections 461 462 func (s *refreshConnectionState) executeUpdatesInParallel(ctx context.Context, updates map[string]*steampipeconfig.ConnectionState) (errors []error) { 463 log.Println("[DEBUG] refreshConnectionState.executeUpdatesInParallel start") 464 defer log.Println("[DEBUG] refreshConnectionState.executeUpdatesInParallel end") 465 466 // convert updates to update sets 467 updatesAsSets := make(map[string][]*steampipeconfig.ConnectionState, len(updates)) 468 for k, v := range updates { 469 updatesAsSets[k] = []*steampipeconfig.ConnectionState{v} 470 } 471 // just call executeUpdateSetsInParallel 472 return s.executeUpdateSetsInParallel(ctx, updatesAsSets) 473 } 474 475 // execute sets of updates in parallel - this is required as for dynamic plugins, we must update all connections in 476 // search path order 477 // - for convenience we also use this function for static connections by mapping the input data 478 // from map[string]*steampipeconfig.ConnectionState to map[string][]*steampipeconfig.ConnectionState 479 func (s *refreshConnectionState) executeUpdateSetsInParallel(ctx context.Context, updates map[string][]*steampipeconfig.ConnectionState) (errors []error) { 480 log.Println("[DEBUG] refreshConnectionState.executeUpdateSetsInParallel start") 481 defer log.Println("[DEBUG] refreshConnectionState.executeUpdateSetsInParallel end") 482 483 var wg sync.WaitGroup 484 var errChan = make(chan *connectionError) 485 486 // default to running a single update at a time 487 var maxParallel = int64(1) 488 // allow override of this behaviour vis env var 489 if envMaxStr, ok := os.LookupEnv("STEAMPIPE_UPDATE_SCHEMA_MAX_PARALLEL"); ok { 490 envMax, err := strconv.Atoi(envMaxStr) 491 if err == nil { 492 maxParallel = int64(envMax) 493 } 494 } 495 log.Printf("[INFO] executeUpdateSetsInParallel - maxParallel= %d", maxParallel) 496 497 sem := semaphore.NewWeighted(maxParallel) 498 499 go func() { 500 for { 501 select { 502 case connectionError := <-errChan: 503 if connectionError == nil { 504 return 505 } 506 errors = append(errors, connectionError.err) 507 conn, poolErr := s.pool.Acquire(ctx) 508 if poolErr == nil { 509 if err := s.tableUpdater.onConnectionError(ctx, conn.Conn(), connectionError.name, connectionError.err); err != nil { 510 log.Println("[WARN] failed to update connection state table", err.Error()) 511 } 512 conn.Release() 513 } 514 } 515 } 516 }() 517 518 // allow disabling of schema clone via env var 519 var cloneSchemaEnabled = true 520 if envClone, ok := os.LookupEnv("STEAMPIPE_CLONE_SCHEMA"); ok { 521 cloneSchemaEnabled = strings.ToLower(envClone) == "true" 522 } 523 log.Printf("[INFO] executeUpdateForConnections - cloneSchema=%v", cloneSchemaEnabled) 524 525 // each update may be multiple connections, to execute in order 526 for _, states := range updates { 527 wg.Add(1) 528 // use semaphore to limit goroutines 529 if err := sem.Acquire(ctx, 1); err != nil { 530 errors = append(errors, err) 531 // if we fail to acquire semaphore, just give up 532 return errors 533 } 534 go func(connectionStates []*steampipeconfig.ConnectionState) { 535 defer func() { 536 wg.Done() 537 sem.Release(1) 538 }() 539 540 s.executeUpdateForConnections(ctx, errChan, cloneSchemaEnabled, connectionStates...) 541 }(states) 542 543 } 544 545 wg.Wait() 546 close(errChan) 547 548 return errors 549 } 550 551 // syncronously execute the update queries for one or more connections 552 func (s *refreshConnectionState) executeUpdateForConnections(ctx context.Context, errChan chan *connectionError, cloneSchemaEnabled bool, connectionStates ...*steampipeconfig.ConnectionState) { 553 log.Println("[DEBUG] refreshConnectionState.executeUpdateForConnections start") 554 defer log.Println("[DEBUG] refreshConnectionState.executeUpdateForConnections end") 555 556 for _, connectionState := range connectionStates { 557 connectionName := connectionState.ConnectionName 558 pluginSchemaName := utils.PluginFQNToSchemaName(connectionState.Plugin) 559 var sql string 560 561 s.exemplarSchemaMapMut.Lock() 562 // is this plugin in the exemplarSchemaMap 563 exemplarSchemaName, haveExemplarSchema := s.exemplarSchemaMap[connectionState.Plugin] 564 if haveExemplarSchema && cloneSchemaEnabled { 565 // we can clone! 566 sql = getCloneSchemaQuery(exemplarSchemaName, connectionState) 567 } else { 568 // just get sql to execute update query, and update the connection state table, in a transaction 569 sql = db_common.GetUpdateConnectionQuery(connectionName, pluginSchemaName) 570 } 571 s.exemplarSchemaMapMut.Unlock() 572 573 // the only error this will return is the failure to update the state table 574 // - all other errors are written to the state table 575 if err := s.executeUpdateQuery(ctx, sql, connectionName); err != nil { 576 errChan <- &connectionError{connectionName, err} 577 } else { 578 // we can clone this plugin, add to exemplarSchemaMap 579 // (AFTER executing the update query) 580 if !haveExemplarSchema && connectionState.CanCloneSchema() { 581 s.exemplarSchemaMap[connectionState.Plugin] = connectionName 582 } 583 } 584 } 585 } 586 587 func (s *refreshConnectionState) executeUpdateQuery(ctx context.Context, sql, connectionName string) (err error) { 588 log.Println("[DEBUG] refreshConnectionState.executeUpdateQuery start") 589 defer log.Println("[DEBUG] refreshConnectionState.executeUpdateQuery end") 590 591 // create a transaction 592 tx, err := s.pool.Begin(ctx) 593 if err != nil { 594 return sperr.WrapWithMessage(err, "failed to create transaction to perform update query") 595 } 596 defer func() { 597 if err != nil { 598 tx.Rollback(ctx) 599 } else { 600 tx.Commit(ctx) 601 } 602 }() 603 604 // execute update sql 605 _, err = tx.Exec(ctx, sql) 606 if err != nil { 607 // update failed connections in result 608 s.res.AddFailedConnection(connectionName, err.Error()) 609 610 // update the state table 611 //(the transaction will be aborted - create a connection for the update) 612 if conn, poolErr := s.pool.Acquire(ctx); poolErr == nil { 613 defer conn.Release() 614 if statusErr := s.tableUpdater.onConnectionError(ctx, conn.Conn(), connectionName, err); statusErr != nil { 615 // NOTE: do not return the error - unless we failed to update the connection state table 616 return error_helpers.CombineErrorsWithPrefix(fmt.Sprintf("failed to update connection %s and failed to update connection_state table", connectionName), err, statusErr) 617 } 618 } 619 return nil 620 } 621 622 // update state table (inside transaction) 623 err = s.tableUpdater.onConnectionReady(ctx, tx.Conn(), connectionName) 624 if err != nil { 625 return sperr.WrapWithMessage(err, "failed to update connection state table") 626 } 627 return nil 628 } 629 630 // set connection comments 631 632 func (s *refreshConnectionState) UpdateCommentsInParallel(ctx context.Context, updates []*steampipeconfig.ConnectionState, plugins map[string]*steampipeconfig.ConnectionPlugin) (errors []error) { 633 if !viper.GetBool(constants.ArgSchemaComments) { 634 return nil 635 } 636 637 var wg sync.WaitGroup 638 var errChan = make(chan *connectionError) 639 640 // use as many goroutines as we have connections 641 var maxUpdateThreads = int64(s.pool.Config().MaxConns) 642 sem := semaphore.NewWeighted(maxUpdateThreads) 643 644 go func() { 645 for { 646 select { 647 case connectionError := <-errChan: 648 if connectionError == nil { 649 return 650 } 651 errors = append(errors, connectionError.err) 652 // TODO just log errors 653 } 654 } 655 }() 656 657 // each update may be multiple connections, to execute in order 658 for _, connectionState := range updates { 659 wg.Add(1) 660 // use semaphore to limit goroutines 661 if err := sem.Acquire(ctx, 1); err != nil { 662 errors = append(errors, err) 663 // if we fail to acquire semaphore, just give up 664 return errors 665 } 666 go func(connectionState *steampipeconfig.ConnectionState) { 667 defer func() { 668 wg.Done() 669 sem.Release(1) 670 }() 671 672 s.updateCommentsForConnection(ctx, errChan, plugins, connectionState) 673 }(connectionState) 674 675 } 676 677 wg.Wait() 678 close(errChan) 679 680 return errors 681 } 682 683 // syncronously execute the comments queries for one or more connections 684 func (s *refreshConnectionState) updateCommentsForConnection(ctx context.Context, errChan chan *connectionError, connectionPluginMap map[string]*steampipeconfig.ConnectionPlugin, connectionState *steampipeconfig.ConnectionState) { 685 connectionName := connectionState.ConnectionName 686 687 var sql string 688 689 // we should have a connectionPlugin loaded for this connection 690 connectionPlugin, ok := connectionPluginMap[connectionName] 691 if !ok { 692 log.Printf("[WARN] no connection plugin loaded for connection '%s', which needs comments updating", connectionName) 693 return 694 } 695 696 schema := connectionPlugin.ConnectionMap[connectionName].Schema.Schema 697 // just get sql to execute update query, and update the connection state table, in a transaction 698 sql = db_common.GetCommentsQueryForPlugin(connectionName, schema) 699 700 // comment cloning disabled for now 701 //// if this schema is static, add to the exemplar map 702 //state.exemplarSchemaMapMut.Lock() 703 //// is this plugin in the exemplarSchemaMap 704 //exemplarSchemaName, haveExemplarSchema := state.exemplarCommentsMap[connectionState.Plugin] 705 //if haveExemplarSchema { 706 //// we can clone! 707 // sql = getCloneCommentsQuery(sql, exemplarSchemaName, connectionState) 708 //} else { 709 // // get the schema from the connection plugin 710 // schema := connectionPluginMap[connectionName].ConnectionMap[connectionName].Schema.Schema 711 // // just get sql to execute update query, and update the connection state table, in a transaction 712 // sql = db_common.GetCommentsQueryForPlugin(connectionName, schema) 713 //} 714 //state.exemplarSchemaMapMut.Unlock() 715 716 // the only error this will return is the failure to update the state table 717 // - all other errors are written to the state table 718 if err := s.executeCommentQuery(ctx, sql, connectionName); err != nil { 719 errChan <- &connectionError{connectionName, err} 720 } //else { 721 // // we can clone this plugin, add to exemplarCommentsMap 722 // // (AFTER executing the update query) 723 // if !haveExemplarSchema && connectionState.CanCloneSchema() { 724 // state.exemplarCommentsMap[connectionState.Plugin] = connectionName 725 // } 726 //} 727 } 728 729 func (s *refreshConnectionState) executeCommentQuery(ctx context.Context, sql, connectionName string) error { 730 // create a transaction 731 tx, err := s.pool.Begin(ctx) 732 if err != nil { 733 return sperr.WrapWithMessage(err, "failed to create transaction to perform update query") 734 } 735 defer func() { 736 if err != nil { 737 tx.Rollback(ctx) 738 } else { 739 tx.Commit(ctx) 740 } 741 }() 742 743 // execute update sql 744 _, err = tx.Exec(ctx, sql) 745 if err != nil { 746 // update the state table 747 //(the transaction will be aborted - create a connection for the update) 748 if conn, poolErr := s.pool.Acquire(ctx); poolErr == nil { 749 defer conn.Release() 750 if statusErr := s.tableUpdater.onConnectionError(ctx, conn.Conn(), connectionName, err); statusErr != nil { 751 // NOTE: do not return the error - unless we failed to update the connection state table 752 return error_helpers.CombineErrorsWithPrefix(fmt.Sprintf("failed to update connection %s and failed to update connection_state table", connectionName), err, statusErr) 753 } 754 } 755 756 return nil 757 } 758 759 // update state table (inside transaction) 760 // ignore error 761 if err := s.tableUpdater.onConnectionCommentsLoaded(ctx, tx.Conn(), connectionName); err != nil { 762 log.Printf("[WARN] failed to set 'comments_set' for connection '%s': %s", connectionName, err.Error()) 763 } 764 765 return nil 766 } 767 768 func getCloneSchemaQuery(exemplarSchemaName string, connectionState *steampipeconfig.ConnectionState) string { 769 return fmt.Sprintf("select clone_foreign_schema('%s', '%s', '%s');", exemplarSchemaName, connectionState.ConnectionName, connectionState.Plugin) 770 } 771 772 func (s *refreshConnectionState) getInitialAndRemainingUpdates() (initialUpdates, remainingUpdates map[string]*steampipeconfig.ConnectionState, dynamicUpdates map[string][]*steampipeconfig.ConnectionState) { 773 updates := s.connectionUpdates.Update 774 searchPathConnections := s.connectionUpdates.FinalConnectionState.GetFirstSearchPathConnectionForPlugins(s.searchPath) 775 776 initialUpdates = make(map[string]*steampipeconfig.ConnectionState) 777 remainingUpdates = make(map[string]*steampipeconfig.ConnectionState) 778 // dynamic plugins must be updated for each plugin in search path order 779 // build a map keyed by plugin, with the value the ordered updates for that plugin 780 dynamicUpdates = make(map[string][]*steampipeconfig.ConnectionState) 781 782 // convert this into a lookup of initial updates to execute 783 for _, connectionName := range searchPathConnections { 784 if connectionState, updateRequired := updates[connectionName]; updateRequired { 785 if connectionState.SchemaMode == plugin.SchemaModeDynamic { 786 pluginInstance := *connectionState.PluginInstance 787 dynamicUpdates[pluginInstance] = append(dynamicUpdates[pluginInstance], connectionState) 788 } else { 789 initialUpdates[connectionName] = connectionState 790 } 791 } 792 } 793 // now add remaining updates to remainingUpdates 794 for connectionName, connectionState := range updates { 795 _, isInitialUpdate := initialUpdates[connectionName] 796 if connectionState.SchemaMode == plugin.SchemaModeStatic && !isInitialUpdate { 797 remainingUpdates[connectionName] = connectionState 798 } 799 800 } 801 return initialUpdates, remainingUpdates, dynamicUpdates 802 } 803 804 func (s *refreshConnectionState) executeDeleteQueries(ctx context.Context, deletions []string) error { 805 t := time.Now() 806 log.Printf("[INFO] execute %d delete %s", len(deletions), utils.Pluralize("query", len(deletions))) 807 defer func() { 808 log.Printf("[INFO] completed execute delete queries (%fs)", time.Since(t).Seconds()) 809 }() 810 811 var errors []error 812 813 for _, c := range deletions { 814 err := s.executeDeleteQuery(ctx, c) 815 if err != nil { 816 errors = append(errors, err) 817 } 818 } 819 return error_helpers.CombineErrors(errors...) 820 } 821 822 // delete the schema and update remove the connection from the state table 823 // NOTE: this only returns an error if we fail to update the state table 824 func (s *refreshConnectionState) executeDeleteQuery(ctx context.Context, connectionName string) error { 825 // create a transaction 826 tx, err := s.pool.Begin(ctx) 827 if err != nil { 828 return sperr.WrapWithMessage(err, "failed to create transaction to perform delete query") 829 } 830 defer func() { 831 if err != nil { 832 _ = tx.Rollback(ctx) 833 } else { 834 err = tx.Commit(ctx) 835 } 836 }() 837 838 sql := db_common.GetDeleteConnectionQuery(connectionName) 839 840 // execute delete sql 841 _, err = tx.Exec(ctx, sql) 842 if err != nil { 843 // update the state table 844 //(the transaction will be aborted - create a connection for the update) 845 if conn, poolErr := s.pool.Acquire(ctx); poolErr == nil { 846 defer conn.Release() 847 if statusErr := s.tableUpdater.onConnectionError(ctx, conn.Conn(), connectionName, err); statusErr != nil { 848 // NOTE: do not return the error - unless we failed to update the connection state table 849 return error_helpers.CombineErrorsWithPrefix(fmt.Sprintf("failed to update connection %s and failed to update connection_state table", connectionName), err, statusErr) 850 } 851 } 852 853 return nil 854 } 855 856 // delete state table entry (inside transaction) 857 err = s.tableUpdater.onConnectionDeleted(ctx, tx.Conn(), connectionName) 858 if err != nil { 859 return sperr.WrapWithMessage(err, "failed to delete connection state table entry for '%s'", connectionName) 860 } 861 return nil 862 } 863 864 // set the state of any incomplete connections to error 865 func (s *refreshConnectionState) setIncompleteConnectionStateToError(ctx context.Context, err error) { 866 // create wrapped error 867 connectionStateError := sperr.WrapWithMessage(err, "failed to update Steampipe connections") 868 // load connection state 869 conn, err := s.pool.Acquire(ctx) 870 if err != nil { 871 log.Printf("[WARN] setAllConnectionStateToError failed to acquire connection from pool: %s", err.Error()) 872 return 873 } 874 defer conn.Release() 875 876 queries := introspection.GetIncompleteConnectionStateErrorSql(connectionStateError) 877 878 if _, err = db_local.ExecuteSqlWithArgsInTransaction(ctx, conn.Conn(), queries...); err != nil { 879 log.Printf("[WARN] setAllConnectionStateToError failed to set connection states to error: %s", err.Error()) 880 return 881 } 882 }