github.com/getsynq/migrate/v4@v4.15.3-0.20220615182648-8e72daaa5ed9/database/clickhouse/clickhouse.go (about)

     1  package clickhouse
     2  
     3  import (
     4  	"database/sql"
     5  	"fmt"
     6  	"io"
     7  	"io/ioutil"
     8  	"net/url"
     9  	"strconv"
    10  	"strings"
    11  	"time"
    12  
    13  	"go.uber.org/atomic"
    14  
    15  	"github.com/getsynq/migrate/v4"
    16  	"github.com/getsynq/migrate/v4/database"
    17  	"github.com/getsynq/migrate/v4/database/multistmt"
    18  	"github.com/hashicorp/go-multierror"
    19  )
    20  
    21  var (
    22  	multiStmtDelimiter = []byte(";")
    23  
    24  	DefaultMigrationsTable       = "schema_migrations"
    25  	DefaultMigrationsTableEngine = "TinyLog"
    26  	DefaultMultiStatementMaxSize = 10 * 1 << 20 // 10 MB
    27  
    28  	ErrNilConfig = fmt.Errorf("no config")
    29  )
    30  
    31  type Config struct {
    32  	DatabaseName          string
    33  	ClusterName           string
    34  	MigrationsTable       string
    35  	MigrationsTableEngine string
    36  	MultiStatementEnabled bool
    37  	MultiStatementMaxSize int
    38  }
    39  
    40  func init() {
    41  	database.Register("clickhouse", &ClickHouse{})
    42  }
    43  
    44  func WithInstance(conn *sql.DB, config *Config) (database.Driver, error) {
    45  	if config == nil {
    46  		return nil, ErrNilConfig
    47  	}
    48  
    49  	if err := conn.Ping(); err != nil {
    50  		return nil, err
    51  	}
    52  
    53  	ch := &ClickHouse{
    54  		conn:   conn,
    55  		config: config,
    56  	}
    57  
    58  	if err := ch.init(); err != nil {
    59  		return nil, err
    60  	}
    61  
    62  	return ch, nil
    63  }
    64  
    65  type ClickHouse struct {
    66  	conn     *sql.DB
    67  	config   *Config
    68  	isLocked atomic.Bool
    69  }
    70  
    71  func (ch *ClickHouse) Open(dsn string) (database.Driver, error) {
    72  	purl, err := url.Parse(dsn)
    73  	if err != nil {
    74  		return nil, err
    75  	}
    76  	q := migrate.FilterCustomQuery(purl)
    77  	q.Scheme = "tcp"
    78  	conn, err := sql.Open("clickhouse", q.String())
    79  	if err != nil {
    80  		return nil, err
    81  	}
    82  
    83  	multiStatementMaxSize := DefaultMultiStatementMaxSize
    84  	if s := purl.Query().Get("x-multi-statement-max-size"); len(s) > 0 {
    85  		multiStatementMaxSize, err = strconv.Atoi(s)
    86  		if err != nil {
    87  			return nil, err
    88  		}
    89  	}
    90  
    91  	migrationsTableEngine := DefaultMigrationsTableEngine
    92  	if s := purl.Query().Get("x-migrations-table-engine"); len(s) > 0 {
    93  		migrationsTableEngine = s
    94  	}
    95  
    96  	ch = &ClickHouse{
    97  		conn: conn,
    98  		config: &Config{
    99  			MigrationsTable:       purl.Query().Get("x-migrations-table"),
   100  			MigrationsTableEngine: migrationsTableEngine,
   101  			DatabaseName:          purl.Query().Get("database"),
   102  			ClusterName:           purl.Query().Get("x-cluster-name"),
   103  			MultiStatementEnabled: purl.Query().Get("x-multi-statement") == "true",
   104  			MultiStatementMaxSize: multiStatementMaxSize,
   105  		},
   106  	}
   107  
   108  	if err := ch.init(); err != nil {
   109  		return nil, err
   110  	}
   111  
   112  	return ch, nil
   113  }
   114  
   115  func (ch *ClickHouse) init() error {
   116  	if len(ch.config.DatabaseName) == 0 {
   117  		if err := ch.conn.QueryRow("SELECT currentDatabase()").Scan(&ch.config.DatabaseName); err != nil {
   118  			return err
   119  		}
   120  	}
   121  
   122  	if len(ch.config.MigrationsTable) == 0 {
   123  		ch.config.MigrationsTable = DefaultMigrationsTable
   124  	}
   125  
   126  	if ch.config.MultiStatementMaxSize <= 0 {
   127  		ch.config.MultiStatementMaxSize = DefaultMultiStatementMaxSize
   128  	}
   129  
   130  	if len(ch.config.MigrationsTableEngine) == 0 {
   131  		ch.config.MigrationsTableEngine = DefaultMigrationsTableEngine
   132  	}
   133  
   134  	return ch.ensureVersionTable()
   135  }
   136  
   137  func (ch *ClickHouse) Run(r io.Reader) error {
   138  	if ch.config.MultiStatementEnabled {
   139  		var err error
   140  		if e := multistmt.Parse(r, multiStmtDelimiter, ch.config.MultiStatementMaxSize, func(m []byte) bool {
   141  			tq := strings.TrimSpace(string(m))
   142  			if tq == "" {
   143  				return true
   144  			}
   145  			if _, e := ch.conn.Exec(string(m)); e != nil {
   146  				err = database.Error{OrigErr: e, Err: "migration failed", Query: m}
   147  				return false
   148  			}
   149  			return true
   150  		}); e != nil {
   151  			return e
   152  		}
   153  		return err
   154  	}
   155  
   156  	migration, err := ioutil.ReadAll(r)
   157  	if err != nil {
   158  		return err
   159  	}
   160  
   161  	if _, err := ch.conn.Exec(string(migration)); err != nil {
   162  		return database.Error{OrigErr: err, Err: "migration failed", Query: migration}
   163  	}
   164  
   165  	return nil
   166  }
   167  func (ch *ClickHouse) Version() (int, bool, error) {
   168  	var (
   169  		version int
   170  		dirty   uint8
   171  		query   = "SELECT version, dirty FROM `" + ch.config.MigrationsTable + "` ORDER BY sequence DESC LIMIT 1"
   172  	)
   173  	if err := ch.conn.QueryRow(query).Scan(&version, &dirty); err != nil {
   174  		if err == sql.ErrNoRows {
   175  			return database.NilVersion, false, nil
   176  		}
   177  		return 0, false, &database.Error{OrigErr: err, Query: []byte(query)}
   178  	}
   179  	return version, dirty == 1, nil
   180  }
   181  
   182  func (ch *ClickHouse) SetVersion(version int, dirty bool) error {
   183  	var (
   184  		bool = func(v bool) uint8 {
   185  			if v {
   186  				return 1
   187  			}
   188  			return 0
   189  		}
   190  		tx, err = ch.conn.Begin()
   191  	)
   192  	if err != nil {
   193  		return err
   194  	}
   195  
   196  	query := "INSERT INTO " + ch.config.MigrationsTable + " (version, dirty, sequence) VALUES (?, ?, ?)"
   197  	stmt, err := tx.Prepare(query)
   198  	if err != nil {
   199  		if rollbackErr := tx.Rollback(); rollbackErr != nil {
   200  			return fmt.Errorf("error during prepare statement %w and rollback %s", err, rollbackErr)
   201  		}
   202  
   203  		return err
   204  	}
   205  
   206  	if _, err := stmt.Exec(int64(version), bool(dirty), uint64(time.Now().UnixNano())); err != nil {
   207  		return &database.Error{OrigErr: err, Query: []byte(query)}
   208  	}
   209  
   210  	return tx.Commit()
   211  }
   212  
   213  // ensureVersionTable checks if versions table exists and, if not, creates it.
   214  // Note that this function locks the database, which deviates from the usual
   215  // convention of "caller locks" in the ClickHouse type.
   216  func (ch *ClickHouse) ensureVersionTable() (err error) {
   217  	if err = ch.Lock(); err != nil {
   218  		return err
   219  	}
   220  
   221  	defer func() {
   222  		if e := ch.Unlock(); e != nil {
   223  			if err == nil {
   224  				err = e
   225  			} else {
   226  				err = multierror.Append(err, e)
   227  			}
   228  		}
   229  	}()
   230  
   231  	var (
   232  		table string
   233  		query = "SHOW TABLES FROM " + ch.config.DatabaseName + " LIKE '" + ch.config.MigrationsTable + "'"
   234  	)
   235  	// check if migration table exists
   236  	if err := ch.conn.QueryRow(query).Scan(&table); err != nil {
   237  		if err != sql.ErrNoRows {
   238  			return &database.Error{OrigErr: err, Query: []byte(query)}
   239  		}
   240  	} else {
   241  		return nil
   242  	}
   243  
   244  	// if not, create the empty migration table
   245  	if len(ch.config.ClusterName) > 0 {
   246  		query = fmt.Sprintf(`
   247  			CREATE TABLE %s ON CLUSTER %s (
   248  				version    Int64,
   249  				dirty      UInt8,
   250  				sequence   UInt64
   251  			) Engine=%s`, ch.config.MigrationsTable, ch.config.ClusterName, ch.config.MigrationsTableEngine)
   252  	} else {
   253  		query = fmt.Sprintf(`
   254  			CREATE TABLE %s (
   255  				version    Int64,
   256  				dirty      UInt8,
   257  				sequence   UInt64
   258  			) Engine=%s`, ch.config.MigrationsTable, ch.config.MigrationsTableEngine)
   259  	}
   260  
   261  	if strings.HasSuffix(ch.config.MigrationsTableEngine, "Tree") {
   262  		query = fmt.Sprintf(`%s ORDER BY sequence`, query)
   263  	}
   264  
   265  	if _, err := ch.conn.Exec(query); err != nil {
   266  		return &database.Error{OrigErr: err, Query: []byte(query)}
   267  	}
   268  	return nil
   269  }
   270  
   271  func (ch *ClickHouse) Drop() (err error) {
   272  	query := "SHOW TABLES FROM " + ch.config.DatabaseName
   273  	tables, err := ch.conn.Query(query)
   274  
   275  	if err != nil {
   276  		return &database.Error{OrigErr: err, Query: []byte(query)}
   277  	}
   278  	defer func() {
   279  		if errClose := tables.Close(); errClose != nil {
   280  			err = multierror.Append(err, errClose)
   281  		}
   282  	}()
   283  
   284  	for tables.Next() {
   285  		var table string
   286  		if err := tables.Scan(&table); err != nil {
   287  			return err
   288  		}
   289  
   290  		query = "DROP TABLE IF EXISTS " + ch.config.DatabaseName + "." + table
   291  
   292  		if _, err := ch.conn.Exec(query); err != nil {
   293  			return &database.Error{OrigErr: err, Query: []byte(query)}
   294  		}
   295  	}
   296  	if err := tables.Err(); err != nil {
   297  		return &database.Error{OrigErr: err, Query: []byte(query)}
   298  	}
   299  
   300  	return nil
   301  }
   302  
   303  func (ch *ClickHouse) Lock() error {
   304  	if !ch.isLocked.CAS(false, true) {
   305  		return database.ErrLocked
   306  	}
   307  
   308  	return nil
   309  }
   310  func (ch *ClickHouse) Unlock() error {
   311  	if !ch.isLocked.CAS(true, false) {
   312  		return database.ErrNotLocked
   313  	}
   314  
   315  	return nil
   316  }
   317  func (ch *ClickHouse) Close() error { return ch.conn.Close() }