github.com/solvedata/migrate/v4@v4.8.7-0.20201127053940-c9fba4ce569f/database/ksql/ksql.go (about)

     1  package ksql
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"errors"
     7  	"fmt"
     8  	"io"
     9  	"io/ioutil"
    10  	"net/http"
    11  	"os"
    12  	"strconv"
    13  	"strings"
    14  	"time"
    15  
    16  	"github.com/solvedata/migrate/v4/database"
    17  )
    18  
    19  func init() {
    20  	database.Register("ksql", &Ksql{})
    21  }
    22  
    23  var CreateMigrationStreamSQL = `CREATE STREAM migrations
    24    (type VARCHAR,
    25    current_version INT,
    26    is_dirty BOOLEAN)
    27    WITH (KAFKA_TOPIC = 'ksql_migrations',
    28          VALUE_FORMAT='JSON',
    29          KEY = 'type',
    30          PARTITIONS = 1);`
    31  var CreateMigrationTableSQL = `CREATE TABLE schema_migrations
    32    WITH (KAFKA_TOPIC = 'ksql_schema_migrations',
    33          VALUE_FORMAT='JSON',
    34          PARTITIONS = 1)
    35    AS SELECT MAX(current_version) as current_version, type FROM migrations
    36    WHERE NOT is_dirty
    37    GROUP BY type;`
    38  var LatestSchemaMigrationSql = `SELECT current_version FROM schema_migrations WHERE type = 'schema' LIMIT 1;`
    39  
    40  type MigrationResult struct {
    41  	Row MigrationRow
    42  }
    43  
    44  type MigrationRow struct {
    45  	Columns []interface{}
    46  }
    47  
    48  type Ksql struct {
    49  	Url               string
    50  	HttpUrl           string
    51  	Instance          interface{}
    52  	CurrentVersion    int
    53  	MigrationSequence []string
    54  	LastRunMigration  []byte // todo: make []string
    55  	IsDirty           bool
    56  	FirstRun          bool
    57  	Client            *http.Client
    58  
    59  	Config *Config
    60  }
    61  
    62  func (s *Ksql) Open(url string) (database.Driver, error) {
    63  	fmt.Println("Opening at KSQL URL", url)
    64  	// Create HTTP client to use
    65  	timeout, err := strconv.ParseInt(os.Getenv("MIGRATE_KSQL_TIMEOUT"), 10, 64)
    66  	if err != nil {
    67  		fmt.Println("Unable to parse `MIGRATE_KSQL_TIMEOUT` environment variable. Defaulting to 10 seconds.")
    68  		timeout = 10
    69  	}
    70  	client := &http.Client{Timeout: time.Duration(timeout) * time.Second}
    71  	httpUrl := strings.Replace(url, "ksql://", "http://", 1)
    72  	fmt.Println("Setting HTTP URL with", httpUrl)
    73  
    74  	// We have a URL - can we connect?
    75  
    76  	ks := &Ksql{
    77  		Url:               url,
    78  		HttpUrl:           httpUrl,
    79  		Client:            client,
    80  		CurrentVersion:    -1,
    81  		FirstRun:          true,
    82  		MigrationSequence: make([]string, 0),
    83  		Config:            &Config{},
    84  	}
    85  
    86  	hasConnection := ks.ensureUrlConection()
    87  
    88  	if !hasConnection {
    89  		return nil, errors.New(fmt.Sprintf("Cannot connect to KSQL at %v", s.HttpUrl))
    90  	}
    91  
    92  	if err := ks.ensureVersionTable(); err != nil {
    93  		return nil, err
    94  	}
    95  
    96  	return ks, nil
    97  }
    98  
    99  type Config struct{}
   100  
   101  func (s *Ksql) Close() error {
   102  	return nil
   103  }
   104  
   105  func (s *Ksql) Lock() error {
   106  	return nil
   107  }
   108  
   109  func (s *Ksql) Unlock() error {
   110  	return nil
   111  }
   112  
   113  func (s *Ksql) Run(migration io.Reader) error {
   114  	m, err := ioutil.ReadAll(migration)
   115  	if err != nil {
   116  		return err
   117  	}
   118  
   119  	s.LastRunMigration = m
   120  	s.MigrationSequence = append(s.MigrationSequence, string(m[:]))
   121  
   122  	query := string(m[:])
   123  	// The migration is expecte to be valid KSQL. Send this to the KSQL server
   124  	resp, err := s.runKsql(query)
   125  	if err != nil {
   126  		return err
   127  	}
   128  
   129  	if resp.StatusCode != 200 {
   130  		// Something unexpected happened. Print out the response body and error out.
   131  		printResponseBody(resp)
   132  		return errors.New(fmt.Sprintf("Unexpected response code of %v", resp.Status))
   133  	}
   134  
   135  	return nil
   136  }
   137  
   138  // Adds a new record with the current migration version and it's dirty state
   139  func (s *Ksql) SetVersion(version int, dirty bool) error {
   140  	if version >= 0 {
   141  		query := fmt.Sprintf("INSERT INTO migrations VALUES ('schema', 'schema', %v, %v);", version, dirty)
   142  		_, err := s.runKsql(query)
   143  		if err != nil {
   144  			return nil
   145  		}
   146  
   147  		// Version updated in migration table successfully. Update instance
   148  		s.CurrentVersion = version
   149  		s.IsDirty = dirty
   150  	}
   151  	return nil
   152  }
   153  
   154  // Retrieves the current version of the KSQL migration state
   155  func (s *Ksql) Version() (version int, dirty bool, err error) {
   156  	if s.FirstRun {
   157  		// This is the first time _any_ migration has been run. No version to retrieve
   158  		// See .ensureVersionTable for where this is set
   159  		fmt.Println("First run, no version to set")
   160  		return -1, false, nil
   161  	}
   162  
   163  	currentVersion, isDirty, err := s.getLatestMigration()
   164  	if err != nil {
   165  		fmt.Println("Error getting latest migration version")
   166  		return -1, false, nil
   167  	}
   168  
   169  	return currentVersion, isDirty, nil
   170  }
   171  
   172  func (s *Ksql) Drop() error {
   173  	s.CurrentVersion = -1
   174  	s.LastRunMigration = nil
   175  	s.MigrationSequence = append(s.MigrationSequence, "DROP")
   176  	return nil
   177  }
   178  
   179  func (s *Ksql) ensureUrlConection() bool {
   180  	// Check that we can run a query with the given URL
   181  	query := "LIST TOPICS;"
   182  	resp, err := s.runKsql(query)
   183  	if err != nil {
   184  		fmt.Println("KSQL URL is not accepting requests")
   185  		return false
   186  	}
   187  
   188  	return resp.Status != "200"
   189  }
   190  
   191  // Makes sure that the schema migration state table is setup correctly
   192  func (s *Ksql) ensureVersionTable() (err error) {
   193  	stmt := "LIST TABLES;"
   194  	resp, err := s.runKsql(stmt)
   195  	if err != nil {
   196  		return err
   197  	}
   198  
   199  	body := resposeBodyText(resp)
   200  	lowerCaseBody := strings.ToLower(body)
   201  	// Simple check - does any text (i.e. table names) contain schema_migrations?
   202  	tableExists := strings.Contains(lowerCaseBody, "schema_migrations")
   203  
   204  	if tableExists {
   205  		fmt.Println("Schema migrations table already exists")
   206  		s.FirstRun = false
   207  		return nil
   208  	}
   209  
   210  	fmt.Println("Schema migrations table does not exist. Creating stream")
   211  	// First create the stream for the table to come off
   212  	resp, err = s.runKsql(CreateMigrationStreamSQL)
   213  	if err != nil {
   214  		return err
   215  	}
   216  
   217  	fmt.Println("Schema migrations table does not exist. Creating table")
   218  	// Now create the table itself
   219  	resp, err = s.runKsql(CreateMigrationTableSQL)
   220  	if err != nil {
   221  		return err
   222  	}
   223  
   224  	fmt.Println("Schema migrations table creation done!")
   225  	return nil
   226  }
   227  
   228  func (s *Ksql) runKsql(query string) (*http.Response, error) {
   229  	url := fmt.Sprintf(`%v/ksql`, s.HttpUrl)
   230  	return s.doQuery(url, query)
   231  }
   232  
   233  func (s *Ksql) runQuery(query string) (*http.Response, error) {
   234  	url := fmt.Sprintf(`%v/query`, s.HttpUrl)
   235  	return s.doQuery(url, query)
   236  }
   237  
   238  func (s *Ksql) doQuery(url string, query string) (*http.Response, error) {
   239  	formatted_query := fmt.Sprintf(`{"ksql":"%v","streamsProperties":{ "ksql.streams.auto.offset.reset": "earliest"}}`, strings.Replace(query, "\n", " ", -1))
   240  	req_body := []byte(formatted_query)
   241  	req, err := http.NewRequest("POST", url, bytes.NewBuffer(req_body))
   242  
   243  	if err != nil {
   244  		return nil, err
   245  	}
   246  
   247  	req.Header.Add("Content-Type", "application/vnd.ksql.v1+json; charset=utf-8")
   248  	resp, err := s.Client.Do(req)
   249  
   250  	if err != nil {
   251  		return nil, err
   252  	}
   253  
   254  	return resp, nil
   255  }
   256  
   257  // Does a request for the most recent event in the migration table
   258  func (s *Ksql) getLatestMigration() (int, bool, error) {
   259  	resp, err := s.runQuery(fmt.Sprintf(LatestSchemaMigrationSql))
   260  	if err != nil {
   261  		return -1, false, err
   262  	}
   263  	result, err := responseBodyMigrationResult(resp)
   264  	if err != nil {
   265  		return -1, false, err
   266  	}
   267  	currentVersion := int(result.Row.Columns[0].(float64))
   268  	fmt.Println("Current version:", currentVersion)
   269  
   270  	return currentVersion, false, nil
   271  }
   272  
   273  // Helper to grab the first line in a response body (while also removing whitespace etc)
   274  func responseBodyMigrationResult(resp *http.Response) (MigrationResult, error) {
   275  	body := strings.Trim(resposeBodyText(resp), "\n")
   276  	lines := strings.Split(body, "\n")
   277  
   278  	var result MigrationResult
   279  	err := json.Unmarshal([]byte(lines[0]), &result)
   280  	if err != nil {
   281  		return MigrationResult{}, err
   282  	}
   283  
   284  	return result, nil
   285  }
   286  
   287  // Helper to extract the HTTP response body
   288  func resposeBodyText(resp *http.Response) string {
   289  	bodyBytes, err := ioutil.ReadAll(resp.Body)
   290  	if err != nil {
   291  		return ""
   292  	}
   293  	return string(bodyBytes)
   294  }
   295  
   296  // Debuggering helper to print the HTTP response body
   297  func printResponseBody(resp *http.Response) {
   298  	bodyString := resposeBodyText(resp)
   299  	fmt.Println(bodyString)
   300  }