feat: add where clause (#2299)

opengovern · Dec 22, 2024 · 03f0096 · 03f0096
1 parent 6133ab5
commit 03f0096
Showing 1 changed file with 38 additions and 2 deletions.
diff --git a/services/rego/internal/rego.go b/services/rego/internal/rego.go
@@ -9,6 +9,7 @@ import (
 	"github.com/open-policy-agent/opa/types"
 	steampipesdk "github.com/opengovern/og-util/pkg/steampipe"
 	"go.uber.org/zap"
+	"strings"
 	"time"
 )
 
@@ -70,11 +71,46 @@ func (r *RegoEngine) getRegoFunctionForTables(ctx context.Context) ([]func(*rego
 		f := rego.FunctionDyn(&rego.Function{
 			Name:             fmt.Sprintf("opencomply.%s", tableName),
 			Description:      "",
-			Decl:             types.NewFunction(nil, types.Any{}),
+			Decl:             types.NewFunction([]types.Type{types.NewObject(nil, &types.DynamicProperty{Key: types.S, Value: types.NewArray(nil, types.A)})}, types.NewArray(nil, types.A)),
 			Memoize:          true,
 			Nondeterministic: true,
 		}, func(bctx rego.BuiltinContext, terms []*ast.Term) (*ast.Term, error) {
-			rows, err := r.steampipe.Conn().Query(bctx.Context, fmt.Sprintf("SELECT * FROM %s", tableName))
+			var args []any
+			var whereClause string
+			if len(terms) > 0 {
+				whereObject, err := ast.ValueToInterface(terms[0].Value, nil)
+				if err != nil {
+					r.logger.Error("Unable to convert to interface", zap.Error(err))
+					r.logger.Sync()
+					return nil, err
+				}
+				if whereMap, ok := whereObject.(map[string]any); ok && len(whereMap) > 0 {
+					whereClauseBuilder := strings.Builder{}
+					whereClauseBuilder.WriteString("WHERE ")
+					counter := 1
+					for key, value := range whereMap {
+						if list, ok := value.([]interface{}); ok {
+							whereClauseBuilder.WriteString(fmt.Sprintf("%s IN (", key))
+							for i, v := range list {
+								whereClauseBuilder.WriteString(fmt.Sprintf("$%d", counter))
+								if i < len(list)-1 {
+									whereClauseBuilder.WriteString(", ")
+								}
+								args = append(args, v)
+								counter++
+							}
+							whereClauseBuilder.WriteString(") AND ")
+						} else {
+							return nil, fmt.Errorf("invalid where clause: %v", whereObject)
+						}
+					}
+					whereClause = strings.TrimSuffix(whereClauseBuilder.String(), "AND ")
+				} else {
+					return nil, fmt.Errorf("invalid where clause: %v", whereObject)
+				}
+			}
+
+			rows, err := r.steampipe.Conn().Query(bctx.Context, fmt.Sprintf("SELECT * FROM %s %s", tableName, whereClause), args...)
 			if err != nil {
 				r.logger.Error("Unable to query database", zap.Error(err), zap.String("table", tableName))
 				r.logger.Sync()