From b0e2e797272a5dafce009225612d5c98ff03ea12 Mon Sep 17 00:00:00 2001
From: Gregory Eremin <g@erem.in>
Date: Mon, 12 Nov 2018 13:33:47 +0100
Subject: [PATCH] Use regexp to determine changed table

---
 reader/schema/manager.go      | 35 +++++++++++++++++++++--------------
 reader/schema/manager_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 14 deletions(-)
 create mode 100644 reader/schema/manager_test.go

diff --git a/reader/schema/manager.go b/reader/schema/manager.go
index 09ef8b2..bf9c338 100644
--- a/reader/schema/manager.go
+++ b/reader/schema/manager.go
@@ -2,25 +2,26 @@ package schema
 
 import (
 	"database/sql"
+	"regexp"
 	"strings"
 )
 
-// SchemaManager maintains table schemas.
-type SchemaManager struct {
+// Manager maintains table schemas.
+type Manager struct {
 	Schema *Schema
 	db     *sql.DB
 }
 
 // NewManager creates a new schema manager.
-func NewManager(db *sql.DB) *SchemaManager {
-	return &SchemaManager{
+func NewManager(db *sql.DB) *Manager {
+	return &Manager{
 		Schema: NewSchema(),
 		db:     db,
 	}
 }
 
 // Manage adds given tables to a list of managed tables and updates its details.
-func (m *SchemaManager) Manage(database, table string) error {
+func (m *Manager) Manage(database, table string) error {
 	cols, err := m.tableColumns(database, table)
 	if err != nil {
 		return err
@@ -31,20 +32,16 @@ func (m *SchemaManager) Manage(database, table string) error {
 }
 
 // ProcessQuery accepts an SQL query and updates schema if required.
-func (m *SchemaManager) ProcessQuery(query string) error {
-	if strings.HasPrefix(query, "ALTER TABLE") {
-		for database, tables := range m.Schema.tables {
-			for table := range tables {
-				if err := m.Manage(database, table); err != nil {
-					return err
-				}
-			}
+func (m *Manager) ProcessQuery(database, query string) error {
+	if tableName, ok := changedTable(query); !ok {
+		if tbl := m.Schema.Table(database, tableName); tbl != nil {
+			return m.Manage(database, tableName)
 		}
 	}
 	return nil
 }
 
-func (m *SchemaManager) tableColumns(database, table string) ([]Column, error) {
+func (m *Manager) tableColumns(database, table string) ([]Column, error) {
 	rows, err := m.db.Query(`
 		SELECT COLUMN_NAME, COLUMN_TYPE 
 		FROM INFORMATION_SCHEMA.COLUMNS 
@@ -71,3 +68,13 @@ func (m *SchemaManager) tableColumns(database, table string) ([]Column, error) {
 	}
 	return cols, nil
 }
+
+var alterRegexp = regexp.MustCompile(`(?im)^alter[\s\t\n]+table[\s\t\n]+` + "`" + `?([a-z0-9_]+)`)
+
+func changedTable(query string) (string, bool) {
+	m := alterRegexp.FindAllStringSubmatch(query, -1)
+	if len(m) > 0 {
+		return m[0][1], true
+	}
+	return "", false
+}
diff --git a/reader/schema/manager_test.go b/reader/schema/manager_test.go
new file mode 100644
index 0000000..3e9dacc
--- /dev/null
+++ b/reader/schema/manager_test.go
@@ -0,0 +1,34 @@
+package schema
+
+import "testing"
+
+func TestChangedTable(t *testing.T) {
+	inputs := []struct {
+		query, tbl string
+		ok         bool
+	}{
+		{"alter\ttable foobar add column", "foobar", true},
+		{"ALTER   TABLE foobar ADD COLUMN", "foobar", true},
+		{"alter table    `foobar` add column", "foobar", true},
+		{"ALTER TABLE `foobar` ADD COLUMN", "foobar", true},
+		{"alter\ntable     \n\tfoobar\nadd column", "foobar", true},
+		{"ALTER TABLE Foo_Bar111 ADD COLUMN", "Foo_Bar111", true},
+		{"SELECT * FROM foobar", "", false},
+		{"SELECT * FROM `foobar`", "", false},
+	}
+
+	for _, in := range inputs {
+		out, ok := changedTable(in.query)
+		if ok != in.ok {
+			if ok {
+				t.Errorf("Didn't expect to match table %q in query %q", out, in.query)
+			} else {
+				t.Errorf("Expected to match table %q in query %q", in.tbl, in.query)
+			}
+			continue
+		}
+		if out != in.tbl {
+			t.Errorf("Expected to match table %q, got %q in query %q", in.tbl, out, in.query)
+		}
+	}
+}