code.vegaprotocol.io/vega@v0.79.0/blockexplorer/store/wipe.go (about)

     1  // Copyright (C) 2023 Gobalsky Labs Limited
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    15  
    16  package store
    17  
    18  import (
    19  	"context"
    20  	"fmt"
    21  
    22  	"code.vegaprotocol.io/vega/logging"
    23  
    24  	"github.com/jackc/pgx/v4"
    25  )
    26  
    27  func DropAllTablesAndViews(log *logging.Logger, config Config) error {
    28  	var err error
    29  	ctx := context.Background()
    30  
    31  	poolConfig, err := config.Postgres.ToPgxPoolConfig()
    32  	if err != nil {
    33  		return fmt.Errorf("determining database connection params: %w", err)
    34  	}
    35  
    36  	conn, err := pgx.ConnectConfig(ctx, poolConfig.ConnConfig)
    37  	if err != nil {
    38  		return fmt.Errorf("connecting to database: %w", err)
    39  	}
    40  
    41  	views, err := queryViews(ctx, conn)
    42  	if err != nil {
    43  		return fmt.Errorf("fetching view list: %w", err)
    44  	}
    45  
    46  	tables, err := queryTables(ctx, conn)
    47  	if err != nil {
    48  		return fmt.Errorf("fetching table list: %w", err)
    49  	}
    50  
    51  	if len(views) == 0 && len(tables) == 0 {
    52  		log.Info("database already empty")
    53  	}
    54  
    55  	for _, view := range views {
    56  		if err := dropView(ctx, conn, view); err != nil {
    57  			return err
    58  		}
    59  		log.Info("dropped view", logging.String("schema", view[0]), logging.String("name", view[1]))
    60  	}
    61  
    62  	for _, table := range tables {
    63  		if err := dropTable(ctx, conn, table); err != nil {
    64  			return err
    65  		}
    66  		log.Info("dropped table", logging.String("schema", table[0]), logging.String("name", table[1]))
    67  	}
    68  
    69  	return nil
    70  }
    71  
    72  func queryViews(ctx context.Context, conn *pgx.Conn) ([]pgx.Identifier, error) {
    73  	return queryIdentifierList(
    74  		ctx,
    75  		conn,
    76  		`select schemaname, viewname from pg_catalog.pg_views where schemaname=current_schema()`,
    77  	)
    78  }
    79  
    80  func queryTables(ctx context.Context, conn *pgx.Conn) ([]pgx.Identifier, error) {
    81  	return queryIdentifierList(
    82  		ctx,
    83  		conn,
    84  		`select schemaname, tablename from pg_catalog.pg_tables where schemaname=current_schema()`,
    85  	)
    86  }
    87  
    88  func dropView(ctx context.Context, conn *pgx.Conn, view pgx.Identifier) error {
    89  	statement := fmt.Sprintf("DROP VIEW IF EXISTS %s CASCADE", view.Sanitize())
    90  	_, err := conn.Exec(ctx, statement)
    91  	if err != nil {
    92  		return fmt.Errorf("dropping view %s: %w", view.Sanitize(), err)
    93  	}
    94  	return nil
    95  }
    96  
    97  func dropTable(ctx context.Context, conn *pgx.Conn, table pgx.Identifier) error {
    98  	statement := fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE", table.Sanitize())
    99  	_, err := conn.Exec(ctx, statement)
   100  	if err != nil {
   101  		return fmt.Errorf("dropping table %s: %w", table.Sanitize(), err)
   102  	}
   103  	return nil
   104  }
   105  
   106  func queryIdentifierList(ctx context.Context, conn *pgx.Conn, query string) ([]pgx.Identifier, error) {
   107  	rows, err := conn.Query(ctx, query)
   108  	if err != nil {
   109  		return nil, err
   110  	}
   111  
   112  	var identifiers []pgx.Identifier
   113  	for rows.Next() {
   114  		var schema, relation string
   115  		if err := rows.Scan(&schema, &relation); err != nil {
   116  			return nil, err
   117  		}
   118  		identifiers = append(identifiers, pgx.Identifier{schema, relation})
   119  	}
   120  
   121  	if rows.Err() != nil {
   122  		return nil, rows.Err()
   123  	}
   124  	return identifiers, nil
   125  }