@@ -36,6 +36,7 @@ const (
3636 HeaderPanelPluginId = "X-Panel-Plugin-Id"
3737 HeaderQueryGroupID = "X-Query-Group-Id" // mainly useful for finding related queries with query chunking
3838 HeaderFromExpression = "X-Grafana-From-Expr" // used by datasources to identify expression queries
39+ headerCodeRabbitOrg = "X-CodeRabbit-Org-Id" // used by CodeRabbit Org Id use to set Row Level Security to scope queries to org
3940)
4041
4142func ProvideService (
@@ -259,6 +260,46 @@ func (s *ServiceImpl) handleQuerySingleDatasource(ctx context.Context, user iden
259260 if err != nil {
260261 return nil , err
261262 }
263+ 264+ codeRabbitOrgId := ""
265+ reqCtx := contexthandler .FromContext (ctx )
266+ if reqCtx != nil && reqCtx .Req != nil {
267+ codeRabbitOrgId = reqCtx .Req .Header .Get (headerCodeRabbitOrg )
268+ }
269+ 270+ if codeRabbitOrgId != "" {
271+ s .log .Info ("CodeRabbitOrgID found in header" )
272+ setQuery := s .createScopeToOrgQuery (codeRabbitOrgId , ds , true )
273+ resetQuery := s .createScopeToOrgQuery (codeRabbitOrgId , ds , false )
274+ 275+ setQueryReq := & backend.QueryDataRequest {
276+ PluginContext : pCtx ,
277+ Headers : map [string ]string {},
278+ Queries : []backend.DataQuery {},
279+ }
280+ setQueryReq .Queries = append (setQueryReq .Queries , setQuery )
281+ _ , err := s .pluginClient .QueryData (ctx , setQueryReq )
282+ if err != nil {
283+ s .log .Error ("Error querying data" , "error" , err )
284+ return nil , err
285+ } else {
286+ s .log .Info ("Applied RLS Query: " )
287+ }
288+ 289+ defer func () {
290+ resetQueryReq := & backend.QueryDataRequest {
291+ PluginContext : pCtx ,
292+ Headers : map [string ]string {},
293+ Queries : []backend.DataQuery {},
294+ }
295+ resetQueryReq .Queries = append (resetQueryReq .Queries , resetQuery )
296+ _ , err := s .pluginClient .QueryData (ctx , resetQueryReq )
297+ if err != nil {
298+ s .log .Error ("Error querying data" , "error" , err )
299+ }
300+ }()
301+ }
302+ 262303 req := & backend.QueryDataRequest {
263304 PluginContext : pCtx ,
264305 Headers : map [string ]string {},
@@ -341,6 +382,36 @@ func (s *ServiceImpl) parseMetricRequest(ctx context.Context, user identity.Requ
341382 return req , req .validateRequest (ctx )
342383}
343384
385+ func (s * ServiceImpl ) createScopeToOrgQuery (codeRabbitOrgId string , ds * datasources.DataSource , setOrg bool ) backend.DataQuery {
386+ setOrgID := fmt .Sprintf ("SET app.current_org_id = '%s'" , codeRabbitOrgId )
387+ resetOrgID := "RESET app.current_org_id"
388+ 389+ rawSql := ""
390+ if setOrg {
391+ rawSql = setOrgID
392+ } else {
393+ rawSql = resetOrgID
394+ }
395+ 396+ return backend.DataQuery {
397+ TimeRange : backend.TimeRange {
398+ From : time .Now ().Add (- time .Hour ),
399+ To : time .Now (),
400+ },
401+ RefID : "rls_setup" ,
402+ MaxDataPoints : 100 ,
403+ Interval : 1000 * time .Millisecond ,
404+ QueryType : rawSql ,
405+ JSON : []byte (`{
406+ "datasource": {"uid": "` + ds .UID + `"},
407+ "intervalMs": 1000,
408+ "maxDataPoints": 100,
409+ "rawSql": "` + rawSql + `",
410+ "format": "table",
411+ "refId": "rls_setup"}` ),
412+ }
413+ }
414+ 344415func (s * ServiceImpl ) getDataSourceFromQuery (ctx context.Context , user identity.Requester , skipDSCache bool , query * simplejson.Json , history map [string ]* datasources.DataSource ) (* datasources.DataSource , error ) {
345416 var err error
346417 uid := query .Get ("datasource" ).Get ("uid" ).MustString ()
0 commit comments