Skip to content

Commit

Permalink
feat: support more ops in sql translate ast walker
Browse files Browse the repository at this point in the history
Added support for some more operators in the SQL translate AST walker:
- bitwise operators
- plus and minus operators
  • Loading branch information
tanmaykm committed Apr 29, 2024
1 parent 33c9a81 commit 1e3bffd
Show file tree
Hide file tree
Showing 4 changed files with 118 additions and 6 deletions.
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "OpenPolicyAgent"
uuid = "8f257efb-743c-4ebc-8197-d291a1f743b4"
authors = ["JuliaHub Inc.", "Tanmay Mohapatra <tanmaykm@gmail.com>"]
version = "0.3.0"
version = "0.3.1"

[deps]
Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
Expand Down
25 changes: 24 additions & 1 deletion src/utils/sql.jl
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,14 @@ const SQL_OP_MAP = Dict{String,String}(
"lte" => "<=",
"equal" => "=",
"internal.member_2" => "in",
"bits.and" => "&",
"bits.or" => "|",
"bits.xor" => "#",
"bits.negate" => "~",
"bits.lsh" => "<<",
"bits.rsh" => ">>",
"plus" => "+",
"minus" => "-",
)

const VALID_SQL_OPS = Set(keys(SQL_OP_MAP))
Expand Down Expand Up @@ -156,6 +164,21 @@ function visit(visitor::SQLVisitor, arr::AST.OPAArray)
return nothing
end

function visit(visitor::SQLVisitor, call::AST.OPACall)
op_name = walk(visitor, call.operator)
if !(op_name in VALID_SQL_OPS)
error("Invalid SQL operator: $op_name")
end
op = SQL_OP_MAP[op_name]

op_operands = call.operands
@assert length(op_operands) == 2
op_lhs = walk(visitor, op_operands[1])
op_rhs = walk(visitor, op_operands[2])

push!(visitor.result_stack, join(["(", op_lhs, op, op_rhs, ")"], " "))
end

function visit(visitor::SQLVisitor, expr::AST.OPAExpr)
@assert AST.is_call(expr)

Expand All @@ -171,7 +194,7 @@ function visit(visitor::SQLVisitor, expr::AST.OPAExpr)
op_rhs = walk(visitor, op_operands[2])
op = SQL_OP_MAP[op_name]

push!(visitor.result_stack, join([op_lhs, op, op_rhs], " "))
push!(visitor.result_stack, join(["(", op_lhs, op, op_rhs, ")"], " "))
end

end # module SQL
23 changes: 22 additions & 1 deletion test/sql_translate.jl
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,14 @@ const SQL_OP_MAP = Dict{String,String}(
"lte" => "<=",
"equal" => "=",
"internal.member_2" => "in",
"bits.and" => "&",
"bits.or" => "|",
"bits.xor" => "#",
"bits.negate" => "~",
"bits.lsh" => "<<",
"bits.rsh" => ">>",
"plus" => "+",
"minus" => "-",
)

const VALID_SQL_OPS = Set(keys(SQL_OP_MAP))
Expand Down Expand Up @@ -229,7 +237,20 @@ function to_sql(expr::OPAExpr)
op_lhs = to_sql(op_operands[1])
op_rhs = to_sql(op_operands[2])

return join([op_lhs, op, op_rhs], " ")
return join(["(", op_lhs, op, op_rhs, ")"], " ")
end

function to_sql(call::OPACall)
op_name = to_sql(call.operator)
if !(op_name in VALID_SQL_OPS)
error("Invalid SQL operator: $op_name")
end

op = SQL_OP_MAP[op_name]
op_lhs = to_sql(call.operands[1])
op_rhs = to_sql(call.operands[2])

return join(["(", op_lhs, op, op_rhs, ")"], " ")
end

function to_sql(query::OPAQuery)
Expand Down
74 changes: 71 additions & 3 deletions test/test_data.jl
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ const PARTIAL_COMPILE_CASES = [
"disableInlining" => []
),
unknowns = ["data.reports"],
sql = "4 >= public.juliahub_reports.clearance_level",
sql = "( 4 >= public.juliahub_reports.clearance_level )",
),
(
policy = """package example
Expand Down Expand Up @@ -66,7 +66,7 @@ const PARTIAL_COMPILE_CASES = [
"disableInlining" => []
),
unknowns = ["data.reports"],
sql = "public.juliahub_reports.public = true or\n4 >= public.juliahub_reports.clearance_level and 'bob' = public.juliahub_reports.owner",
sql = "( public.juliahub_reports.public = true ) or\n( 4 >= public.juliahub_reports.clearance_level ) and ( 'bob' = public.juliahub_reports.owner )",
),
(
# always allowed if the policy is fully satisfied with the given input for any one condition
Expand Down Expand Up @@ -163,7 +163,75 @@ const PARTIAL_COMPILE_CASES = [
"disableInlining" => []
),
unknowns = ["data.reports"],
sql = "public.juliahub_reports.category in ('public', 'pinned') or\n4 >= public.juliahub_reports.clearance_level",
sql = "( public.juliahub_reports.category in ('public', 'pinned') ) or\n( 4 >= public.juliahub_reports.clearance_level )",
),
(
policy = """package example
allow {
bits.and(data.reports[_].clearance_level, input.subject.clearance_level) >= input.subject.clearance_level
}""",
query = "data.example.allow == true",
input = Dict{String,Any}(
"subject" => Dict{String,Any}(
"clearance_level" => 4
)
),
options = Dict{String,Any}(
"disableInlining" => []
),
unknowns = ["data.reports"],
sql = "( ( public.juliahub_reports.clearance_level & 4 ) >= 4 )",
),
(
policy = """package example
allow {
bits.or(data.reports[_].clearance_level, input.subject.clearance_level) >= input.subject.clearance_level
}""",
query = "data.example.allow == true",
input = Dict{String,Any}(
"subject" => Dict{String,Any}(
"clearance_level" => 4
)
),
options = Dict{String,Any}(
"disableInlining" => []
),
unknowns = ["data.reports"],
sql = "( ( public.juliahub_reports.clearance_level | 4 ) >= 4 )",
),
(
policy = """package example
allow {
(data.reports[_].clearance_level + input.subject.clearance_level) >= input.subject.clearance_level
}""",
query = "data.example.allow == true",
input = Dict{String,Any}(
"subject" => Dict{String,Any}(
"clearance_level" => 4
)
),
options = Dict{String,Any}(
"disableInlining" => []
),
unknowns = ["data.reports"],
sql = "( ( public.juliahub_reports.clearance_level + 4 ) >= 4 )",
),
(
policy = """package example
allow {
(data.reports[_].clearance_level - input.subject.clearance_level) >= 0
}""",
query = "data.example.allow == true",
input = Dict{String,Any}(
"subject" => Dict{String,Any}(
"clearance_level" => 4
)
),
options = Dict{String,Any}(
"disableInlining" => []
),
unknowns = ["data.reports"],
sql = "( ( public.juliahub_reports.clearance_level - 4 ) >= 0 )",
),
]

Expand Down

0 comments on commit 1e3bffd

Please sign in to comment.