@@ -2921,13 +2921,18 @@ func (ex *connExecutor) execCopyOut(
29212921 }
29222922
29232923 if copyErr := ex .execWithProfiling (ctx , cmd .Stmt , nil , func (ctx context.Context ) error {
2924- ex .mu .Lock ()
2925- queryMeta , ok := ex .mu .ActiveQueries [queryID ]
2926- if ! ok {
2927- return errors .AssertionFailedf ("query %d not in registry" , queryID )
2924+ if err := func () error {
2925+ ex .mu .Lock ()
2926+ defer ex .mu .Unlock ()
2927+ queryMeta , ok := ex .mu .ActiveQueries [queryID ]
2928+ if ! ok {
2929+ return errors .AssertionFailedf ("query %d not in registry" , queryID )
2930+ }
2931+ queryMeta .phase = executing
2932+ return nil
2933+ }(); err != nil {
2934+ return err
29282935 }
2929- queryMeta .phase = executing
2930- ex .mu .Unlock ()
29312936
29322937 // We'll always have a txn on the planner since we called resetPlanner
29332938 // above.
@@ -3183,13 +3188,18 @@ func (ex *connExecutor) execCopyIn(
31833188 }
31843189
31853190 if copyErr = ex .execWithProfiling (ctx , cmd .Stmt , nil , func (ctx context.Context ) error {
3186- ex .mu .Lock ()
3187- queryMeta , ok := ex .mu .ActiveQueries [queryID ]
3188- if ! ok {
3189- return errors .AssertionFailedf ("query %d not in registry" , queryID )
3191+ if err := func () error {
3192+ ex .mu .Lock ()
3193+ defer ex .mu .Unlock ()
3194+ queryMeta , ok := ex .mu .ActiveQueries [queryID ]
3195+ if ! ok {
3196+ return errors .AssertionFailedf ("query %d not in registry" , queryID )
3197+ }
3198+ queryMeta .phase = executing
3199+ return nil
3200+ }(); err != nil {
3201+ return err
31903202 }
3191- queryMeta .phase = executing
3192- ex .mu .Unlock ()
31933203 return cm .run (ctx )
31943204 }); copyErr != nil {
31953205 // TODO(andrei): We don't have a full retriable error story for the copy machine.
0 commit comments