Postgres/MySQL/MSSQL: Add setting to limit maximum amount of rows processed (#38986)

Adds a new setting dataproxy.row_limit that allows an operator to limit the 
amount of rows being processed/accepted in response to database queries 
originating from SQL data sources.

Closes #38975
Ref #39095

Co-authored-by: achatterjee-grafana <70489351+achatterjee-grafana@users.noreply.github.com>
This commit is contained in:
Marcus Efraimsson
2021-09-13 15:27:51 +02:00
committed by GitHub
parent 5b7dc16c06
commit 1c892a2fc4
12 changed files with 279 additions and 21 deletions

View File

@ -26,10 +26,9 @@ var logger = log.New("tsdb.postgres")
func ProvideService(cfg *setting.Cfg, manager backendplugin.Manager) (*Service, error) {
s := &Service{
Cfg: cfg,
tlsManager: newTLSManager(logger, cfg.DataPath),
}
s.im = datasource.NewInstanceManager(s.newInstanceSettings())
s.im = datasource.NewInstanceManager(s.newInstanceSettings(cfg))
factory := coreplugin.New(backend.ServeOpts{
QueryDataHandler: s,
})
@ -41,7 +40,6 @@ func ProvideService(cfg *setting.Cfg, manager backendplugin.Manager) (*Service,
}
type Service struct {
Cfg *setting.Cfg
tlsManager tlsSettingsProvider
im instancemgmt.InstanceManager
}
@ -63,7 +61,7 @@ func (s *Service) QueryData(ctx context.Context, req *backend.QueryDataRequest)
return dsInfo.QueryData(ctx, req)
}
func (s *Service) newInstanceSettings() datasource.InstanceFactoryFunc {
func (s *Service) newInstanceSettings(cfg *setting.Cfg) datasource.InstanceFactoryFunc {
return func(settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
logger.Debug("Creating Postgres query endpoint")
jsonData := sqleng.JsonData{
@ -94,7 +92,7 @@ func (s *Service) newInstanceSettings() datasource.InstanceFactoryFunc {
return nil, err
}
if s.Cfg.Env == setting.Dev {
if cfg.Env == setting.Dev {
logger.Debug("getEngine", "connection", cnnstr)
}
@ -103,6 +101,7 @@ func (s *Service) newInstanceSettings() datasource.InstanceFactoryFunc {
ConnectionString: cnnstr,
DSInfo: dsInfo,
MetricColumnTypes: []string{"UNKNOWN", "TEXT", "VARCHAR", "CHAR"},
RowLimit: cfg.DataProxyRowLimit,
}
queryResultTransformer := postgresQueryResultTransformer{

View File

@ -112,7 +112,6 @@ func TestGenerateConnectionString(t *testing.T) {
for _, tt := range testCases {
t.Run(tt.desc, func(t *testing.T) {
svc := Service{
Cfg: cfg,
tlsManager: &tlsTestManager{settings: tt.tlsSettings},
}
@ -190,6 +189,7 @@ func TestPostgres(t *testing.T) {
ConnectionString: "",
DSInfo: dsInfo,
MetricColumnTypes: []string{"UNKNOWN", "TEXT", "VARCHAR", "CHAR"},
RowLimit: 1000000,
}
queryResultTransformer := postgresQueryResultTransformer{
@ -1225,6 +1225,84 @@ func TestPostgres(t *testing.T) {
require.Equal(t, data.FieldTypeNullableTime, frames[0].Fields[0].Type())
require.Equal(t, data.FieldTypeNullableTime, frames[0].Fields[1].Type())
})
t.Run("When row limit set to 1", func(t *testing.T) {
dsInfo := sqleng.DataSourceInfo{}
config := sqleng.DataPluginConfiguration{
DriverName: "postgres",
ConnectionString: "",
DSInfo: dsInfo,
MetricColumnTypes: []string{"UNKNOWN", "TEXT", "VARCHAR", "CHAR"},
RowLimit: 1,
}
queryResultTransformer := postgresQueryResultTransformer{
log: logger,
}
handler, err := sqleng.NewQueryDataHandler(config, &queryResultTransformer, newPostgresMacroEngine(false), logger)
require.NoError(t, err)
t.Run("When doing a table query that returns 2 rows should limit the result to 1 row", func(t *testing.T) {
query := &backend.QueryDataRequest{
Queries: []backend.DataQuery{
{
JSON: []byte(`{
"rawSql": "SELECT 1 as value UNION ALL select 2 as value",
"format": "table"
}`),
RefID: "A",
TimeRange: backend.TimeRange{
From: time.Now(),
To: time.Now(),
},
},
},
}
resp, err := handler.QueryData(context.Background(), query)
require.NoError(t, err)
queryResult := resp.Responses["A"]
require.NoError(t, queryResult.Error)
frames := queryResult.Frames
require.NoError(t, err)
require.Equal(t, 1, len(frames))
require.Equal(t, 1, len(frames[0].Fields))
require.Equal(t, 1, frames[0].Rows())
require.Len(t, frames[0].Meta.Notices, 1)
require.Equal(t, data.NoticeSeverityWarning, frames[0].Meta.Notices[0].Severity)
})
t.Run("When doing a time series query that returns 2 rows should limit the result to 1 row", func(t *testing.T) {
query := &backend.QueryDataRequest{
Queries: []backend.DataQuery{
{
JSON: []byte(`{
"rawSql": "SELECT 1 as time, 1 as value UNION ALL select 2 as time, 2 as value",
"format": "time_series"
}`),
RefID: "A",
TimeRange: backend.TimeRange{
From: time.Now(),
To: time.Now(),
},
},
},
}
resp, err := handler.QueryData(context.Background(), query)
require.NoError(t, err)
queryResult := resp.Responses["A"]
require.NoError(t, queryResult.Error)
frames := queryResult.Frames
require.NoError(t, err)
require.Equal(t, 1, len(frames))
require.Equal(t, 2, len(frames[0].Fields))
require.Equal(t, 1, frames[0].Rows())
require.Len(t, frames[0].Meta.Notices, 1)
require.Equal(t, data.NoticeSeverityWarning, frames[0].Meta.Notices[0].Severity)
})
})
})
}