From 241a7fb342dc235133f0e5d49a64037b9630126b Mon Sep 17 00:00:00 2001 From: Zach Musgrave Date: Thu, 8 Apr 2021 15:10:31 -0700 Subject: [PATCH] Check constraint, working except for DROP CONSTRAINT. Also needs bats tests --- go/libraries/doltcore/schema/check_coll.go | 94 +++++++++++++++++++ .../schema/encoding/schema_marshaling.go | 34 ++++++- go/libraries/doltcore/schema/index.go | 1 + go/libraries/doltcore/schema/schema.go | 3 + go/libraries/doltcore/schema/schema_impl.go | 8 ++ .../sqle/enginetest/dolt_engine_test.go | 17 ++++ go/libraries/doltcore/sqle/tables.go | 88 +++++++++++++++++ 7 files changed, 244 insertions(+), 1 deletion(-) create mode 100755 go/libraries/doltcore/schema/check_coll.go diff --git a/go/libraries/doltcore/schema/check_coll.go b/go/libraries/doltcore/schema/check_coll.go new file mode 100755 index 0000000000..5546b8eb09 --- /dev/null +++ b/go/libraries/doltcore/schema/check_coll.go @@ -0,0 +1,94 @@ +// Copyright 2021 Dolthub, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package schema + +import ( + "fmt" +) + +type Check interface { + Name() string + Expression() string + Enforced() bool +} + +// CheckCollection is the set of `check` constraints on a table's schema +type CheckCollection interface { + // AddCheck adds a check to this collection and returns it + AddCheck(name, expression string, enforce bool) (Check, error) + DropCheck(name string) error + AllChecks() []Check + Count() int +} + +type check struct { + name string + expression string + enforced bool +} + +func (c check) Name() string { + return c.name +} + +func (c check) Expression() string { + return c.expression +} + +func (c check) Enforced() bool { + return c.enforced +} + +type checkCollection struct { + checks map[string]check +} + +func (c checkCollection) AddCheck(name, expression string, enforce bool) (Check, error) { + if _, ok := c.checks[name]; ok { + return nil, fmt.Errorf("name %s in use", name) + } + c.checks[name] = check{ + name: name, + expression: expression, + enforced: enforce, + } + + return c.checks[name], nil +} + +func (c checkCollection) DropCheck(name string) error { + delete(c.checks, name) + return nil +} + +func (c checkCollection) AllChecks() []Check { + checks := make([]Check, len(c.checks)) + i := 0 + for _, check := range c.checks { + checks[i] = check + i++ + } + return checks +} + +func (c checkCollection) Count() int { + return len(c.checks) +} + +func NewCheckCollection() CheckCollection { + return checkCollection{ + checks: make(map[string]check), + } +} \ No newline at end of file diff --git a/go/libraries/doltcore/schema/encoding/schema_marshaling.go b/go/libraries/doltcore/schema/encoding/schema_marshaling.go index b5c1cd244c..a9cae0dc3a 100644 --- a/go/libraries/doltcore/schema/encoding/schema_marshaling.go +++ b/go/libraries/doltcore/schema/encoding/schema_marshaling.go @@ -146,9 +146,16 @@ type encodedIndex struct { IsSystemDefined bool `noms:"hidden,omitempty" json:"hidden,omitempty"` // Was previously named Hidden, do not change noms name } +type encodedCheck struct { + Name string `noms:"name" json:"name"` + Expression string `noms:"expression" json:"expression"` + Enforced bool `noms:"enforced" json:"enforced"` +} + type schemaData struct { Columns []encodedColumn `noms:"columns" json:"columns"` IndexCollection []encodedIndex `noms:"idxColl,omitempty" json:"idxColl,omitempty"` + CheckConstraints []encodedCheck `noms:"checks,omitempty" json:"checks,omitempty"` } func toSchemaData(sch schema.Schema) (schemaData, error) { @@ -178,7 +185,21 @@ func toSchemaData(sch schema.Schema) (schemaData, error) { } } - return schemaData{encCols, encodedIndexes}, nil + encodedChecks := make([]encodedCheck, sch.Checks().Count()) + checks := sch.Checks() + for i, check := range checks.AllChecks() { + encodedChecks[i] = encodedCheck{ + Name: check.Name(), + Expression: check.Expression(), + Enforced: check.Enforced(), + } + } + + return schemaData{ + Columns: encCols, + IndexCollection: encodedIndexes, + CheckConstraints: encodedChecks, + }, nil } func (sd schemaData) decodeSchema() (schema.Schema, error) { @@ -215,6 +236,17 @@ func (sd schemaData) decodeSchema() (schema.Schema, error) { } } + for _, encodedCheck := range sd.CheckConstraints { + _, err = sch.Checks().AddCheck( + encodedCheck.Name, + encodedCheck.Expression, + encodedCheck.Enforced, + ) + if err != nil { + return nil, err + } + } + return sch, nil } diff --git a/go/libraries/doltcore/schema/index.go b/go/libraries/doltcore/schema/index.go index 76ce3acd41..aaabe72266 100644 --- a/go/libraries/doltcore/schema/index.go +++ b/go/libraries/doltcore/schema/index.go @@ -171,6 +171,7 @@ func (ix *indexImpl) Schema() Schema { nonPKCols: nonPkCols, allCols: allCols, indexCollection: NewIndexCollection(nil), + checkCollection: NewCheckCollection(), } } diff --git a/go/libraries/doltcore/schema/schema.go b/go/libraries/doltcore/schema/schema.go index be8ed796e8..901a2c9bd7 100644 --- a/go/libraries/doltcore/schema/schema.go +++ b/go/libraries/doltcore/schema/schema.go @@ -27,6 +27,9 @@ type Schema interface { // Indexes returns a collection of all indexes on the table that this schema belongs to. Indexes() IndexCollection + + // Checks returns a collection of all check constraints on the table that this schema belongs to. + Checks() CheckCollection } // ColFromTag returns a schema.Column from a schema and a tag diff --git a/go/libraries/doltcore/schema/schema_impl.go b/go/libraries/doltcore/schema/schema_impl.go index 29f540d27a..5a7c521144 100644 --- a/go/libraries/doltcore/schema/schema_impl.go +++ b/go/libraries/doltcore/schema/schema_impl.go @@ -32,6 +32,7 @@ var EmptySchema = &schemaImpl{ type schemaImpl struct { pkCols, nonPKCols, allCols *ColCollection indexCollection IndexCollection + checkCollection CheckCollection } // SchemaFromCols creates a Schema from a collection of columns @@ -59,6 +60,7 @@ func SchemaFromCols(allCols *ColCollection) (Schema, error) { nonPKCols: nonPKColColl, allCols: allCols, indexCollection: NewIndexCollection(allCols), + checkCollection: NewCheckCollection(), }, nil } @@ -123,6 +125,7 @@ func UnkeyedSchemaFromCols(allCols *ColCollection) Schema { nonPKCols: nonPKColColl, allCols: nonPKColColl, indexCollection: NewIndexCollection(nil), + checkCollection: NewCheckCollection(), } } @@ -156,6 +159,7 @@ func SchemaFromPKAndNonPKCols(pkCols, nonPKCols *ColCollection) (Schema, error) nonPKCols: nonPKCols, allCols: allColColl, indexCollection: NewIndexCollection(allColColl), + checkCollection: NewCheckCollection(), }, nil } @@ -207,3 +211,7 @@ func (si *schemaImpl) String() string { func (si *schemaImpl) Indexes() IndexCollection { return si.indexCollection } + +func (si *schemaImpl) Checks() CheckCollection { + return si.checkCollection +} diff --git a/go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go b/go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go index e4b0fd9454..3907e0cbe7 100644 --- a/go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go +++ b/go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go @@ -220,6 +220,23 @@ func TestDropForeignKeys(t *testing.T) { enginetest.TestDropForeignKeys(t, newDoltHarness(t)) } +func TestCreateCheckConstraints(t *testing.T) { + enginetest.TestCreateCheckConstraints(t, newDoltHarness(t)) +} + +func TestChecksOnInsert(t *testing.T) { + enginetest.TestChecksOnInsert(t, newDoltHarness(t)) +} + +func TestTestDisallowedCheckConstraints(t *testing.T) { + enginetest.TestDisallowedCheckConstraints(t, newDoltHarness(t)) +} + +func TestDropCheckConstraints(t *testing.T) { + enginetest.TestDropCheckConstraints(t, newDoltHarness(t)) +} + + func TestExplode(t *testing.T) { t.Skipf("Unsupported types") enginetest.TestExplode(t, newDoltHarness(t)) diff --git a/go/libraries/doltcore/sqle/tables.go b/go/libraries/doltcore/sqle/tables.go index 79aedf68ce..ad7e66adbd 100644 --- a/go/libraries/doltcore/sqle/tables.go +++ b/go/libraries/doltcore/sqle/tables.go @@ -388,6 +388,7 @@ var _ sql.InsertableTable = (*WritableDoltTable)(nil) var _ sql.ReplaceableTable = (*WritableDoltTable)(nil) var _ sql.AutoIncrementTable = (*WritableDoltTable)(nil) var _ sql.TruncateableTable = (*WritableDoltTable)(nil) +var _ sql.CheckTable = (*WritableDoltTable)(nil) func (t *WritableDoltTable) WithIndexLookup(lookup sql.IndexLookup) sql.Table { dil, ok := lookup.(*doltIndexLookup) @@ -526,6 +527,24 @@ func (t *WritableDoltTable) GetAutoIncrementValue(ctx *sql.Context) (interface{} return t.DoltTable.GetAutoIncrementValue(ctx) } +func (t *WritableDoltTable) GetChecks(ctx *sql.Context) ([]sql.CheckDefinition, error) { + sch, err := t.table.GetSchema(ctx) + if err != nil { + return nil, err + } + + checks := make([]sql.CheckDefinition, sch.Checks().Count()) + for i, check := range sch.Checks().AllChecks() { + checks[i] = sql.CheckDefinition{ + Name: check.Name(), + CheckExpression: check.Expression(), + Enforced: check.Enforced(), + } + } + + return checks, nil +} + // GetForeignKeys implements sql.ForeignKeyTable func (t *DoltTable) GetForeignKeys(ctx *sql.Context) ([]sql.ForeignKeyConstraint, error) { root, err := t.db.GetRoot(ctx) @@ -677,6 +696,7 @@ var _ sql.AlterableTable = (*AlterableDoltTable)(nil) var _ sql.IndexAlterableTable = (*AlterableDoltTable)(nil) var _ sql.ForeignKeyAlterableTable = (*AlterableDoltTable)(nil) var _ sql.ForeignKeyTable = (*AlterableDoltTable)(nil) +var _ sql.CheckAlterableTable = (*AlterableDoltTable)(nil) // AddColumn implements sql.AlterableTable func (t *AlterableDoltTable) AddColumn(ctx *sql.Context, column *sql.Column, order *sql.ColumnOrder) error { @@ -1422,3 +1442,71 @@ func (t *AlterableDoltTable) updateFromRoot(ctx *sql.Context, root *doltdb.RootV t.WritableDoltTable.DoltTable = updatedTable.WritableDoltTable.DoltTable return nil } + +func (t *AlterableDoltTable) CreateCheck(ctx *sql.Context, check *sql.CheckDefinition) error { + sch, err := t.table.GetSchema(ctx) + if err != nil { + return err + } + + _, err = sch.Checks().AddCheck(check.Name, check.CheckExpression, check.Enforced) + if err != nil { + return err + } + + newTable, err := t.table.UpdateSchema(ctx, sch) + if err != nil { + return err + } + + root, err := t.db.GetRoot(ctx) + if err != nil { + return err + } + + newRoot, err := root.PutTable(ctx, t.name, newTable) + if err != nil { + return err + } + + err = t.db.SetRoot(ctx, newRoot) + if err != nil { + return err + } + + return t.updateFromRoot(ctx, newRoot) +} + +func (t *AlterableDoltTable) DropCheck(ctx *sql.Context, chName string) error { + sch, err := t.table.GetSchema(ctx) + if err != nil { + return err + } + + err = sch.Checks().DropCheck(chName) + if err != nil { + return err + } + + newTable, err := t.table.UpdateSchema(ctx, sch) + if err != nil { + return err + } + + root, err := t.db.GetRoot(ctx) + if err != nil { + return err + } + + newRoot, err := root.PutTable(ctx, t.name, newTable) + if err != nil { + return err + } + + err = t.db.SetRoot(ctx, newRoot) + if err != nil { + return err + } + + return t.updateFromRoot(ctx, newRoot) +}