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