diff --git a/pkg/table/asserty/README.md b/pkg/table/asserty/README.md new file mode 100644 index 0000000..886f82a --- /dev/null +++ b/pkg/table/asserty/README.md @@ -0,0 +1,5 @@ +# What is this? + +Asserty is a package that provides a set of functions to help ensure that a database schema conforms to your expectations. + +Sure, you could probably use it in tests. But our goal for this is to use it as part of service readiness checks. We want to make sure that the database schema is what we expect it to be before we start to serve traffic. \ No newline at end of file diff --git a/pkg/table/asserty/asserty.go b/pkg/table/asserty/asserty.go new file mode 100644 index 0000000..5f6d3bc --- /dev/null +++ b/pkg/table/asserty/asserty.go @@ -0,0 +1,60 @@ +// Package asserty offers functionality to assert for certain DB properties. +package asserty + +import ( + "context" + "database/sql" + "errors" + "slices" + + "github.com/cashapp/spirit/pkg/table" + _ "github.com/go-sql-driver/mysql" +) + +type Table struct { + ti *table.TableInfo +} + +func LoadTable(db *sql.DB, schema, tableName string) (*Table, error) { + ti := table.NewTableInfo(db, schema, tableName) + if err := ti.SetInfo(context.TODO()); err != nil { + return nil, err + } + return &Table{ti: ti}, nil +} + +func (t *Table) ContainsColumns(columnNames ...string) error { + for _, col := range columnNames { + if !slices.Contains(t.ti.Columns, col) { + return errors.New("missing column " + col + " on table " + t.ti.QuotedName) + } + } + return nil +} + +func (t *Table) NotContainsColumns(columnNames ...string) error { + for _, col := range columnNames { + if slices.Contains(t.ti.Columns, col) { + return errors.New("unexpected column " + col + " on table " + t.ti.QuotedName) + } + } + return nil +} + +func (t *Table) ContainsIndexes(indexNames ...string) error { + for _, idx := range indexNames { + if !slices.Contains(t.ti.Indexes, idx) { + return errors.New("missing index " + idx + " on table " + t.ti.QuotedName) + } + } + return nil +} + +func (t *Table) NotContainsIndexes(indexNames ...string) error { + for _, idx := range indexNames { + if slices.Contains(t.ti.Indexes, idx) { + return errors.New("unexpected index " + idx + " on table " + t.ti.QuotedName) + } + } + return nil +} diff --git a/pkg/table/asserty/asserty_test.go b/pkg/table/asserty/asserty_test.go new file mode 100644 index 0000000..fca9f0f --- /dev/null +++ b/pkg/table/asserty/asserty_test.go @@ -0,0 +1,39 @@ +package asserty + +import ( + "database/sql" + "testing" + + "github.com/cashapp/spirit/pkg/testutils" + "github.com/stretchr/testify/require" +) + +func TestBasicUsage(t *testing.T) { + db, err := sql.Open("mysql", testutils.DSN()) + require.NoError(t, err) + defer db.Close() + + testutils.RunSQL(t, `DROP TABLE IF EXISTS asserty_test`) + table := `CREATE TABLE asserty_test ( + id int(11) unsigned NOT NULL AUTO_INCREMENT, + name varchar(115), + age int(11) NOT NULL, + PRIMARY KEY (id, age), + INDEX idx_name (name), + INDEX idx_age (age) + )` + testutils.RunSQL(t, table) + + tbl, err := LoadTable(db, "test", "asserty_test") + require.NoError(t, err) + require.NoError(t, tbl.ContainsColumns("name", "age")) + require.Error(t, tbl.ContainsColumns("name", "somethingelse")) + require.NoError(t, tbl.NotContainsColumns("col1", "col2")) + require.Error(t, tbl.NotContainsColumns("col1", "name", "col2")) + + require.NoError(t, tbl.ContainsIndexes("idx_name", "idx_age")) + require.Error(t, tbl.ContainsIndexes("idx_name", "idx_age", "idx_somethingelse")) + require.NoError(t, tbl.NotContainsIndexes("idx_col1", "idx_col2")) + require.Error(t, tbl.NotContainsIndexes("idx_col1", "idx_col2", "idx_name")) + require.Error(t, tbl.NotContainsIndexes("idx_name")) +} diff --git a/pkg/table/tableinfo.go b/pkg/table/tableinfo.go index 5aae5db..b804d7b 100644 --- a/pkg/table/tableinfo.go +++ b/pkg/table/tableinfo.go @@ -32,6 +32,7 @@ type TableInfo struct { TableName string QuotedName string Columns []string // all the column names + Indexes []string // all the index names columnsMySQLTps map[string]string // map from column name to MySQL type KeyColumns []string // the column names of the primaryKey keyColumnsMySQLTp []string // the MySQL types of the primaryKey @@ -86,6 +87,9 @@ func (t *TableInfo) SetInfo(ctx context.Context) error { if err := t.setPrimaryKey(ctx); err != nil { return err } + if err := t.setIndexes(ctx); err != nil { + return err + } return t.setMinMax(ctx) } @@ -106,6 +110,29 @@ func (t *TableInfo) setRowEstimate(ctx context.Context) error { return nil } +func (t *TableInfo) setIndexes(ctx context.Context) error { + rows, err := t.db.QueryContext(ctx, "SELECT DISTINCT INDEX_NAME FROM INFORMATION_SCHEMA.STATISTICS WHERE table_schema=? AND table_name=? AND index_name != 'PRIMARY'", + t.SchemaName, + t.TableName, + ) + if err != nil { + return err + } + defer rows.Close() + t.Indexes = []string{} + for rows.Next() { + var name string + if err := rows.Scan(&name); err != nil { + return err + } + t.Indexes = append(t.Indexes, name) + } + if rows.Err() != nil { + return rows.Err() + } + return nil +} + func (t *TableInfo) setColumns(ctx context.Context) error { rows, err := t.db.QueryContext(ctx, "SELECT column_name, column_type FROM information_schema.columns WHERE table_schema=? AND table_name=? ORDER BY ORDINAL_POSITION", t.SchemaName,