Skip to content

Commit 0d5a7ea

Browse files
kyleconroyclaude
kyleconroy
and
claude
committed
Add Always Encrypted column encryption key support
- Add AST types for CREATE/ALTER/DROP COLUMN ENCRYPTION KEY statements - Add ColumnEncryptionKeyValue and parameter types - Add parsing for column encryption key statements - Add ALTER COLUMN ENCRYPTION KEY ADD/DROP VALUE parsing - Add JSON marshaling for new statement types - Add TokenMaster, TokenKey, TokenEncryption to GRANT/DENY permission parsing - Add IsIfExists field to DropColumnMasterKeyStatement Enables AlwaysEncryptedTests130 and Baselines130_AlwaysEncryptedTests130. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent 934b947 commit 0d5a7ea

7 files changed

Lines changed: 410 additions & 11 deletions

File tree

‎ast/column_encryption.go‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,3 +35,70 @@ type ColumnEncryptionAlgorithmParameter struct {
3535
}
3636

3737
func (c *ColumnEncryptionAlgorithmParameter) columnEncryptionParameter() {}
38+
39+
// ColumnEncryptionKeyValueParameter represents a parameter in column encryption key values
40+
type ColumnEncryptionKeyValueParameter interface {
41+
columnEncryptionKeyValueParameter()
42+
}
43+
44+
// ColumnMasterKeyNameParameter represents COLUMN_MASTER_KEY parameter in CEK
45+
type ColumnMasterKeyNameParameter struct {
46+
Name *Identifier
47+
ParameterKind string // "ColumnMasterKeyName"
48+
}
49+
50+
func (c *ColumnMasterKeyNameParameter) node() {}
51+
func (c *ColumnMasterKeyNameParameter) columnEncryptionKeyValueParameter() {}
52+
53+
// ColumnEncryptionAlgorithmNameParameter represents ALGORITHM parameter in CEK
54+
type ColumnEncryptionAlgorithmNameParameter struct {
55+
Algorithm ScalarExpression
56+
ParameterKind string // "EncryptionAlgorithmName"
57+
}
58+
59+
func (c *ColumnEncryptionAlgorithmNameParameter) node() {}
60+
func (c *ColumnEncryptionAlgorithmNameParameter) columnEncryptionKeyValueParameter() {}
61+
62+
// EncryptedValueParameter represents ENCRYPTED_VALUE parameter
63+
type EncryptedValueParameter struct {
64+
Value ScalarExpression
65+
ParameterKind string // "EncryptedValue"
66+
}
67+
68+
func (e *EncryptedValueParameter) node() {}
69+
func (e *EncryptedValueParameter) columnEncryptionKeyValueParameter() {}
70+
71+
// ColumnEncryptionKeyValue represents a value in CREATE/ALTER COLUMN ENCRYPTION KEY
72+
type ColumnEncryptionKeyValue struct {
73+
Parameters []ColumnEncryptionKeyValueParameter
74+
}
75+
76+
func (c *ColumnEncryptionKeyValue) node() {}
77+
78+
// CreateColumnEncryptionKeyStatement represents CREATE COLUMN ENCRYPTION KEY statement
79+
type CreateColumnEncryptionKeyStatement struct {
80+
Name *Identifier
81+
ColumnEncryptionKeyValues []*ColumnEncryptionKeyValue
82+
}
83+
84+
func (c *CreateColumnEncryptionKeyStatement) node() {}
85+
func (c *CreateColumnEncryptionKeyStatement) statement() {}
86+
87+
// AlterColumnEncryptionKeyStatement represents ALTER COLUMN ENCRYPTION KEY statement
88+
type AlterColumnEncryptionKeyStatement struct {
89+
Name *Identifier
90+
AlterType string // "Add" or "Drop"
91+
ColumnEncryptionKeyValues []*ColumnEncryptionKeyValue
92+
}
93+
94+
func (a *AlterColumnEncryptionKeyStatement) node() {}
95+
func (a *AlterColumnEncryptionKeyStatement) statement() {}
96+
97+
// DropColumnEncryptionKeyStatement represents DROP COLUMN ENCRYPTION KEY statement
98+
type DropColumnEncryptionKeyStatement struct {
99+
Name *Identifier
100+
IsIfExists bool
101+
}
102+
103+
func (d *DropColumnEncryptionKeyStatement) node() {}
104+
func (d *DropColumnEncryptionKeyStatement) statement() {}

‎ast/column_master_key_statement.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,8 @@ func (c *ColumnMasterKeyEnclaveComputationsParameter) columnMasterKeyParameter()
4444

4545
// DropColumnMasterKeyStatement represents a DROP COLUMN MASTER KEY statement.
4646
type DropColumnMasterKeyStatement struct {
47-
Name *Identifier
47+
Name *Identifier
48+
IsIfExists bool
4849
}
4950

5051
func (d *DropColumnMasterKeyStatement) node() {}

‎parser/marshal.go‎

Lines changed: 107 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,12 @@ func statementToJSON(stmt ast.Statement) jsonNode {
175175
return createColumnMasterKeyStatementToJSON(s)
176176
case *ast.DropColumnMasterKeyStatement:
177177
return dropColumnMasterKeyStatementToJSON(s)
178+
case *ast.CreateColumnEncryptionKeyStatement:
179+
return createColumnEncryptionKeyStatementToJSON(s)
180+
case *ast.AlterColumnEncryptionKeyStatement:
181+
return alterColumnEncryptionKeyStatementToJSON(s)
182+
case *ast.DropColumnEncryptionKeyStatement:
183+
return dropColumnEncryptionKeyStatementToJSON(s)
178184
case *ast.AlterCryptographicProviderStatement:
179185
return alterCryptographicProviderStatementToJSON(s)
180186
case *ast.DropCryptographicProviderStatement:
@@ -8466,7 +8472,8 @@ func (p *Parser) parseGrantStatement() (*ast.GrantStatement, error) {
84668472
p.curTok.Type == TokenDatabase || p.curTok.Type == TokenTable ||
84678473
p.curTok.Type == TokenFunction || p.curTok.Type == TokenBackup ||
84688474
p.curTok.Type == TokenDefault || p.curTok.Type == TokenTrigger ||
8469-
p.curTok.Type == TokenSchema {
8475+
p.curTok.Type == TokenSchema || p.curTok.Type == TokenMaster ||
8476+
p.curTok.Type == TokenKey || p.curTok.Type == TokenEncryption {
84708477
perm.Identifiers = append(perm.Identifiers, &ast.Identifier{
84718478
Value: p.curTok.Literal,
84728479
QuoteType: "NotQuoted",
@@ -9040,7 +9047,8 @@ func (p *Parser) parseDenyStatement() (*ast.DenyStatement, error) {
90409047
p.curTok.Type == TokenDatabase || p.curTok.Type == TokenTable ||
90419048
p.curTok.Type == TokenFunction || p.curTok.Type == TokenBackup ||
90429049
p.curTok.Type == TokenDefault || p.curTok.Type == TokenTrigger ||
9043-
p.curTok.Type == TokenSchema {
9050+
p.curTok.Type == TokenSchema || p.curTok.Type == TokenMaster ||
9051+
p.curTok.Type == TokenKey || p.curTok.Type == TokenEncryption {
90449052
perm.Identifiers = append(perm.Identifiers, &ast.Identifier{
90459053
Value: p.curTok.Literal,
90469054
QuoteType: "NotQuoted",
@@ -20827,9 +20835,106 @@ func dropColumnMasterKeyStatementToJSON(s *ast.DropColumnMasterKeyStatement) jso
2082720835
if s.Name != nil {
2082820836
node["Name"] = identifierToJSON(s.Name)
2082920837
}
20838+
node["IsIfExists"] = s.IsIfExists
20839+
return node
20840+
}
20841+
20842+
func createColumnEncryptionKeyStatementToJSON(s *ast.CreateColumnEncryptionKeyStatement) jsonNode {
20843+
node := jsonNode{
20844+
"$type": "CreateColumnEncryptionKeyStatement",
20845+
}
20846+
if s.Name != nil {
20847+
node["Name"] = identifierToJSON(s.Name)
20848+
}
20849+
if len(s.ColumnEncryptionKeyValues) > 0 {
20850+
values := make([]jsonNode, len(s.ColumnEncryptionKeyValues))
20851+
for i, v := range s.ColumnEncryptionKeyValues {
20852+
values[i] = columnEncryptionKeyValueToJSON(v)
20853+
}
20854+
node["ColumnEncryptionKeyValues"] = values
20855+
}
20856+
return node
20857+
}
20858+
20859+
func alterColumnEncryptionKeyStatementToJSON(s *ast.AlterColumnEncryptionKeyStatement) jsonNode {
20860+
node := jsonNode{
20861+
"$type": "AlterColumnEncryptionKeyStatement",
20862+
}
20863+
if s.AlterType != "" {
20864+
node["AlterType"] = s.AlterType
20865+
}
20866+
if s.Name != nil {
20867+
node["Name"] = identifierToJSON(s.Name)
20868+
}
20869+
if len(s.ColumnEncryptionKeyValues) > 0 {
20870+
values := make([]jsonNode, len(s.ColumnEncryptionKeyValues))
20871+
for i, v := range s.ColumnEncryptionKeyValues {
20872+
values[i] = columnEncryptionKeyValueToJSON(v)
20873+
}
20874+
node["ColumnEncryptionKeyValues"] = values
20875+
}
2083020876
return node
2083120877
}
2083220878

20879+
func dropColumnEncryptionKeyStatementToJSON(s *ast.DropColumnEncryptionKeyStatement) jsonNode {
20880+
node := jsonNode{
20881+
"$type": "DropColumnEncryptionKeyStatement",
20882+
}
20883+
if s.Name != nil {
20884+
node["Name"] = identifierToJSON(s.Name)
20885+
}
20886+
node["IsIfExists"] = s.IsIfExists
20887+
return node
20888+
}
20889+
20890+
func columnEncryptionKeyValueToJSON(v *ast.ColumnEncryptionKeyValue) jsonNode {
20891+
node := jsonNode{
20892+
"$type": "ColumnEncryptionKeyValue",
20893+
}
20894+
if len(v.Parameters) > 0 {
20895+
params := make([]jsonNode, len(v.Parameters))
20896+
for i, p := range v.Parameters {
20897+
params[i] = columnEncryptionKeyValueParameterToJSON(p)
20898+
}
20899+
node["Parameters"] = params
20900+
}
20901+
return node
20902+
}
20903+
20904+
func columnEncryptionKeyValueParameterToJSON(p ast.ColumnEncryptionKeyValueParameter) jsonNode {
20905+
switch param := p.(type) {
20906+
case *ast.ColumnMasterKeyNameParameter:
20907+
node := jsonNode{
20908+
"$type": "ColumnMasterKeyNameParameter",
20909+
}
20910+
if param.Name != nil {
20911+
node["Name"] = identifierToJSON(param.Name)
20912+
}
20913+
node["ParameterKind"] = param.ParameterKind
20914+
return node
20915+
case *ast.ColumnEncryptionAlgorithmNameParameter:
20916+
node := jsonNode{
20917+
"$type": "ColumnEncryptionAlgorithmNameParameter",
20918+
}
20919+
if param.Algorithm != nil {
20920+
node["Algorithm"] = scalarExpressionToJSON(param.Algorithm)
20921+
}
20922+
node["ParameterKind"] = param.ParameterKind
20923+
return node
20924+
case *ast.EncryptedValueParameter:
20925+
node := jsonNode{
20926+
"$type": "EncryptedValueParameter",
20927+
}
20928+
if param.Value != nil {
20929+
node["Value"] = scalarExpressionToJSON(param.Value)
20930+
}
20931+
node["ParameterKind"] = param.ParameterKind
20932+
return node
20933+
default:
20934+
return jsonNode{"$type": "UnknownColumnEncryptionKeyValueParameter"}
20935+
}
20936+
}
20937+
2083320938
func alterCryptographicProviderStatementToJSON(s *ast.AlterCryptographicProviderStatement) jsonNode {
2083420939
node := jsonNode{
2083520940
"$type": "AlterCryptographicProviderStatement",

‎parser/parse_ddl.go‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,8 @@ func (p *Parser) parseDropStatement() (ast.Statement, error) {
160160
return p.parseDropServiceStatement()
161161
case "EVENT":
162162
return p.parseDropEventNotificationStatement()
163+
case "COLUMN":
164+
return p.parseDropColumnStatement()
163165
}
164166

165167
// Handle LOGIN token explicitly
@@ -2275,6 +2277,42 @@ func (p *Parser) parseDropMasterKeyStatement() (*ast.DropMasterKeyStatement, err
22752277
return stmt, nil
22762278
}
22772279

2280+
func (p *Parser) parseDropColumnStatement() (ast.Statement, error) {
2281+
// Consume COLUMN
2282+
p.nextToken()
2283+
2284+
keyword := strings.ToUpper(p.curTok.Literal)
2285+
p.nextToken() // consume MASTER or ENCRYPTION
2286+
2287+
if keyword == "MASTER" {
2288+
// DROP COLUMN MASTER KEY
2289+
if strings.ToUpper(p.curTok.Literal) == "KEY" {
2290+
p.nextToken() // consume KEY
2291+
}
2292+
stmt := &ast.DropColumnMasterKeyStatement{
2293+
Name: p.parseIdentifier(),
2294+
}
2295+
if p.curTok.Type == TokenSemicolon {
2296+
p.nextToken()
2297+
}
2298+
return stmt, nil
2299+
} else if keyword == "ENCRYPTION" {
2300+
// DROP COLUMN ENCRYPTION KEY
2301+
if strings.ToUpper(p.curTok.Literal) == "KEY" {
2302+
p.nextToken() // consume KEY
2303+
}
2304+
stmt := &ast.DropColumnEncryptionKeyStatement{
2305+
Name: p.parseIdentifier(),
2306+
}
2307+
if p.curTok.Type == TokenSemicolon {
2308+
p.nextToken()
2309+
}
2310+
return stmt, nil
2311+
}
2312+
2313+
return nil, fmt.Errorf("unexpected token after DROP COLUMN: expected MASTER or ENCRYPTION, got %s", keyword)
2314+
}
2315+
22782316
func (p *Parser) parseDropXmlSchemaCollectionStatement() (*ast.DropXmlSchemaCollectionStatement, error) {
22792317
// Consume XML
22802318
p.nextToken()
@@ -2621,6 +2659,8 @@ func (p *Parser) parseAlterStatement() (ast.Statement, error) {
26212659
return p.parseAlterEventSessionStatement()
26222660
case "SECURITY":
26232661
return p.parseAlterSecurityPolicyStatement()
2662+
case "COLUMN":
2663+
return p.parseAlterColumnEncryptionKeyStatement()
26242664
}
26252665
return nil, fmt.Errorf("unexpected token after ALTER: %s", p.curTok.Literal)
26262666
default:
@@ -12704,3 +12744,97 @@ func (p *Parser) parseAlterAuthorizationStatement() (*ast.AlterAuthorizationStat
1270412744

1270512745
return stmt, nil
1270612746
}
12747+
12748+
func (p *Parser) parseAlterColumnEncryptionKeyStatement() (ast.Statement, error) {
12749+
// ALTER COLUMN ENCRYPTION KEY name ADD|DROP VALUE (...)
12750+
// Currently on COLUMN
12751+
p.nextToken() // consume COLUMN
12752+
12753+
if strings.ToUpper(p.curTok.Literal) != "ENCRYPTION" {
12754+
return nil, fmt.Errorf("expected ENCRYPTION after COLUMN, got %s", p.curTok.Literal)
12755+
}
12756+
p.nextToken() // consume ENCRYPTION
12757+
12758+
if strings.ToUpper(p.curTok.Literal) != "KEY" {
12759+
return nil, fmt.Errorf("expected KEY after ENCRYPTION, got %s", p.curTok.Literal)
12760+
}
12761+
p.nextToken() // consume KEY
12762+
12763+
stmt := &ast.AlterColumnEncryptionKeyStatement{}
12764+
12765+
// Parse key name
12766+
stmt.Name = p.parseIdentifier()
12767+
12768+
// Parse ADD VALUE or DROP VALUE
12769+
keyword := strings.ToUpper(p.curTok.Literal)
12770+
if keyword == "ADD" {
12771+
stmt.AlterType = "Add"
12772+
p.nextToken() // consume ADD
12773+
} else if keyword == "DROP" {
12774+
stmt.AlterType = "Drop"
12775+
p.nextToken() // consume DROP
12776+
} else {
12777+
return nil, fmt.Errorf("expected ADD or DROP, got %s", p.curTok.Literal)
12778+
}
12779+
12780+
if strings.ToUpper(p.curTok.Literal) == "VALUE" {
12781+
p.nextToken() // consume VALUE
12782+
}
12783+
12784+
// Parse the value - enclosed in ( ... )
12785+
if p.curTok.Type == TokenLParen {
12786+
value := &ast.ColumnEncryptionKeyValue{}
12787+
p.nextToken() // consume (
12788+
12789+
// Parse parameters
12790+
for p.curTok.Type != TokenRParen && p.curTok.Type != TokenEOF {
12791+
paramName := strings.ToUpper(p.curTok.Literal)
12792+
p.nextToken() // consume parameter name
12793+
12794+
if p.curTok.Type == TokenEquals {
12795+
p.nextToken() // consume =
12796+
}
12797+
12798+
switch paramName {
12799+
case "COLUMN_MASTER_KEY":
12800+
value.Parameters = append(value.Parameters, &ast.ColumnMasterKeyNameParameter{
12801+
Name: p.parseIdentifier(),
12802+
ParameterKind: "ColumnMasterKeyName",
12803+
})
12804+
case "ALGORITHM":
12805+
expr, _ := p.parseScalarExpression()
12806+
value.Parameters = append(value.Parameters, &ast.ColumnEncryptionAlgorithmNameParameter{
12807+
Algorithm: expr,
12808+
ParameterKind: "EncryptionAlgorithmName",
12809+
})
12810+
case "ENCRYPTED_VALUE":
12811+
expr, _ := p.parseScalarExpression()
12812+
value.Parameters = append(value.Parameters, &ast.EncryptedValueParameter{
12813+
Value: expr,
12814+
ParameterKind: "EncryptedValue",
12815+
})
12816+
default:
12817+
// Skip unknown parameter
12818+
p.nextToken()
12819+
}
12820+
12821+
// Skip comma if present
12822+
if p.curTok.Type == TokenComma {
12823+
p.nextToken()
12824+
}
12825+
}
12826+
12827+
// Consume closing )
12828+
if p.curTok.Type == TokenRParen {
12829+
p.nextToken()
12830+
}
12831+
12832+
stmt.ColumnEncryptionKeyValues = append(stmt.ColumnEncryptionKeyValues, value)
12833+
}
12834+
12835+
if p.curTok.Type == TokenSemicolon {
12836+
p.nextToken()
12837+
}
12838+
12839+
return stmt, nil
12840+
}

0 commit comments

Comments
 (0)