github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/libraries/doltcore/sqle/statspro/initdbhook.go (about)

     1  // Copyright 2024 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package statspro
    16  
    17  import (
    18  	"context"
    19  	"strings"
    20  
    21  	"github.com/dolthub/go-mysql-server/sql"
    22  
    23  	"github.com/dolthub/dolt/go/libraries/doltcore/env"
    24  	"github.com/dolthub/dolt/go/libraries/doltcore/sqle"
    25  	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
    26  )
    27  
    28  func NewStatsInitDatabaseHook(
    29  	statsProv *Provider,
    30  	ctxFactory func(ctx context.Context) (*sql.Context, error),
    31  	bThreads *sql.BackgroundThreads,
    32  ) sqle.InitDatabaseHook {
    33  	return func(
    34  		ctx *sql.Context,
    35  		pro *sqle.DoltDatabaseProvider,
    36  		name string,
    37  		denv *env.DoltEnv,
    38  		db dsess.SqlDatabase,
    39  	) error {
    40  		statsDb, err := statsProv.sf.Init(ctx, db, statsProv.pro, denv.FS, env.GetCurrentUserHomeDir)
    41  		if err != nil {
    42  			ctx.GetLogger().Debugf("statistics load error: %s", err.Error())
    43  			return nil
    44  		}
    45  		statsProv.mu.Lock()
    46  		statsProv.setStatDb(strings.ToLower(db.Name()), statsDb)
    47  		statsProv.mu.Unlock()
    48  
    49  		ctx.GetLogger().Debugf("statistics refresh: initialize %s", name)
    50  		return statsProv.InitAutoRefresh(ctxFactory, name, bThreads)
    51  	}
    52  }
    53  
    54  func NewStatsDropDatabaseHook(statsProv *Provider) sqle.DropDatabaseHook {
    55  	return func(ctx *sql.Context, name string) {
    56  		statsProv.CancelRefreshThread(name)
    57  		statsProv.DropDbStats(ctx, name, false)
    58  
    59  		if db, ok := statsProv.getStatDb(name); ok {
    60  			if err := db.Close(); err != nil {
    61  				ctx.GetLogger().Debugf("failed to close stats database: %s", err)
    62  			}
    63  		}
    64  	}
    65  }