Merge pull request #69 from uc-cdis/add_timeout_to_queries

Adding a default timeout to all queries

Author: pieterlukasse

models/cohortdata.go

@@ -47,12 +47,14 @@ func (h CohortData) RetrieveDataByOriginalCohortAndNewCohort(sourceId int, origi
 	resultsDataSource := dataSourceModel.GetDataSource(sourceId, Results)
 	var personData []*PersonIdAndCohort
 
-	meta_result := resultsDataSource.Db.Model(&Cohort{}).
+	query := resultsDataSource.Db.Model(&Cohort{}).
 		Select("cohort.subject_id as person_id, cohort.cohort_definition_id as cohort_id").
 		Joins("INNER JOIN "+resultsDataSource.Schema+".cohort as original_cohort ON cohort.subject_id = original_cohort.subject_id").
 		Where("cohort.cohort_definition_id = ?", cohortDefinitionId).
-		Where("original_cohort.cohort_definition_id = ?", originalCohortDefinitionId).
-		Scan(&personData)
+		Where("original_cohort.cohort_definition_id = ?", originalCohortDefinitionId)
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	meta_result := query.Scan(&personData)
 	return personData, meta_result.Error
 }
 
@@ -68,14 +70,16 @@ func (h CohortData) RetrieveDataBySourceIdAndCohortIdAndConceptIdsOrderedByPerso
 
 	// get the observations for the subjects and the concepts, to build up the data rows to return:
 	var cohortData []*PersonConceptAndValue
-	meta_result := omopDataSource.Db.Table(omopDataSource.Schema+".observation_continuous as observation"+omopDataSource.GetViewDirective()).
+	query := omopDataSource.Db.Table(omopDataSource.Schema+".observation_continuous as observation"+omopDataSource.GetViewDirective()).
 		Select("observation.person_id, observation.observation_concept_id as concept_id, concept.concept_class_id, observation.value_as_string as concept_value_as_string, observation.value_as_number as concept_value_as_number, observation.value_as_concept_id as concept_value_as_concept_id").
 		Joins("INNER JOIN "+resultsDataSource.Schema+".cohort as cohort ON cohort.subject_id = observation.person_id").
 		Joins("INNER JOIN "+omopDataSource.Schema+".concept as concept ON concept.concept_id = observation.observation_concept_id").
 		Where("cohort.cohort_definition_id = ?", cohortDefinitionId).
 		Where("observation.observation_concept_id in (?)", conceptIds).
-		Order("observation.person_id asc"). // this order is important!
-		Scan(&cohortData)
+		Order("observation.person_id asc") // this order is important!
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	meta_result := query.Scan(&cohortData)
 	return cohortData, meta_result.Error
 }
 
@@ -95,6 +99,8 @@ func (h CohortData) RetrieveHistogramDataBySourceIdAndCohortIdAndConceptIdsAndCo
 
 	query = QueryFilterByConceptIdsHelper(query, sourceId, filterConceptIds, omopDataSource, resultsDataSource.Schema, "observation")
 
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
 	meta_result := query.Scan(&cohortData)
 	return cohortData, meta_result.Error
 }
@@ -118,6 +124,8 @@ func (h CohortData) RetrieveCohortOverlapStatsWithoutFilteringOnConceptValue(sou
 		query = QueryFilterByConceptIdsHelper(query, sourceId, otherFilterConceptIds, omopDataSource, resultsDataSource.Schema, "observation")
 	}
 	query = query.Where("control_cohort.cohort_definition_id = ?", controlCohortId)
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
 	meta_result := query.Scan(&cohortOverlapStats)
 	return cohortOverlapStats, meta_result.Error
 }
@@ -152,6 +160,8 @@ func (h CohortData) ValidateObservationData(observationConceptIdsToCheck []int64
 			Group("observation.person_id, observation.observation_concept_id").
 			Having("count(*) > 1")
 
+		query, cancel := utils.AddTimeoutToQuery(query)
+		defer cancel()
 		meta_result := query.Scan(&personConceptAndCount)
 		if meta_result.Error != nil {
 			return -1, meta_result.Error

models/cohortdefinition.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/uc-cdis/cohort-middleware/db"
+	"github.com/uc-cdis/cohort-middleware/utils"
 )
 
 type CohortDefinitionI interface {
@@ -32,29 +33,35 @@ type CohortDefinitionStats struct {
 func (h CohortDefinition) GetCohortDefinitionById(id int) (*CohortDefinition, error) {
 	db2 := db.GetAtlasDB().Db
 	var cohortDefinition *CohortDefinition
-	meta_result := db2.Model(&CohortDefinition{}).
+	query := db2.Model(&CohortDefinition{}).
 		Select("id, name, description").
-		Where("id = ?", id).
-		Scan(&cohortDefinition)
+		Where("id = ?", id)
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	meta_result := query.Scan(&cohortDefinition)
 	return cohortDefinition, meta_result.Error
 }
 
 func (h CohortDefinition) GetCohortDefinitionByName(name string) (*CohortDefinition, error) {
 	db2 := db.GetAtlasDB().Db
 	var cohortDefinition *CohortDefinition
-	meta_result := db2.Model(&CohortDefinition{}).
+	query := db2.Model(&CohortDefinition{}).
 		Select("id, name, description").
-		Where("name = ?", name).
-		Scan(&cohortDefinition)
+		Where("name = ?", name)
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	meta_result := query.Scan(&cohortDefinition)
 	return cohortDefinition, meta_result.Error
 }
 
 func (h CohortDefinition) GetAllCohortDefinitions() ([]*CohortDefinition, error) {
 	db2 := db.GetAtlasDB().Db
 	var cohortDefinition []*CohortDefinition
-	meta_result := db2.Model(&CohortDefinition{}).
-		Select("id, name, description").
-		Scan(&cohortDefinition)
+	query := db2.Model(&CohortDefinition{}).
+		Select("id, name, description")
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	meta_result := query.Scan(&cohortDefinition)
 	return cohortDefinition, meta_result.Error
 }
 
@@ -64,11 +71,13 @@ func (h CohortDefinition) GetAllCohortDefinitionsAndStatsOrderBySizeDesc(sourceI
 	var dataSourceModel = new(Source)
 	resultsDataSource := dataSourceModel.GetDataSource(sourceId, Results)
 	var cohortDefinitionStats []*CohortDefinitionStats
-	meta_result := resultsDataSource.Db.Model(&Cohort{}).
+	query := resultsDataSource.Db.Model(&Cohort{}).
 		Select("cohort_definition_id as id, '' as name, count(*) as cohort_size").
 		Group("cohort_definition_id").
-		Order("count(*) desc").
-		Scan(&cohortDefinitionStats)
+		Order("count(*) desc")
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	meta_result := query.Scan(&cohortDefinitionStats)
 
 	// add name details:
 	for _, cohortDefinitionStat := range cohortDefinitionStats {

models/concept.go

@@ -58,10 +58,12 @@ func (h Concept) RetriveAllBySourceId(sourceId int) ([]*Concept, error) {
 	omopDataSource := dataSourceModel.GetDataSource(sourceId, Omop)
 
 	var concepts []*Concept
-	meta_result := omopDataSource.Db.Model(&Concept{}).
+	query := omopDataSource.Db.Model(&Concept{}).
 		Select("concept_id, concept_name, concept_class_id as concept_type").
-		Order("concept_name").
-		Scan(&concepts)
+		Order("concept_name")
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	meta_result := query.Scan(&concepts)
 	return concepts, meta_result.Error
 }
 
@@ -95,11 +97,13 @@ func (h Concept) RetrieveInfoBySourceIdAndConceptIds(sourceId int, conceptIds []
 	omopDataSource := dataSourceModel.GetDataSource(sourceId, Omop)
 
 	var conceptItems []*ConceptSimple
-	meta_result := omopDataSource.Db.Model(&Concept{}).
+	query := omopDataSource.Db.Model(&Concept{}).
 		Select("concept_id, concept_name, concept_code, concept_class_id as concept_type").
 		Where("concept_id in (?)", conceptIds).
-		Order("concept_name").
-		Scan(&conceptItems)
+		Order("concept_name")
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	meta_result := query.Scan(&conceptItems)
 	if meta_result.Error != nil {
 		return nil, meta_result.Error
 	}
@@ -118,11 +122,14 @@ func (h Concept) RetrieveInfoBySourceIdAndConceptTypes(sourceId int, conceptType
 	omopDataSource := dataSourceModel.GetDataSource(sourceId, Omop)
 
 	var conceptItems []*ConceptSimple
-	meta_result := omopDataSource.Db.Model(&Concept{}).
+	query := omopDataSource.Db.Model(&Concept{}).
 		Select("concept_id, concept_name, concept_class_id as concept_type").
 		Where("concept_class_id in (?)", conceptTypes).
-		Order("concept_name").
-		Scan(&conceptItems)
+		Order("concept_name")
+
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	meta_result := query.Scan(&conceptItems)
 	if meta_result.Error != nil {
 		return nil, meta_result.Error
 	}
@@ -135,6 +142,7 @@ func (h Concept) RetrieveInfoBySourceIdAndConceptTypes(sourceId int, conceptType
 
 // Retrieve concept name, type and missing ratio statistics for given list of conceptIds.
 // Assumption is that both OMOP and RESULTS schemas are on same DB.
+// TODO - remove this code as it is NOT used anymore by the frontend
 func (h Concept) RetrieveStatsBySourceIdAndCohortIdAndConceptIds(sourceId int, cohortDefinitionId int, conceptIds []int64) ([]*ConceptStats, error) {
 	var dataSourceModel = new(Source)
 	omopDataSource := dataSourceModel.GetDataSource(sourceId, Omop)
@@ -227,6 +235,8 @@ func (h Concept) RetrieveBreakdownStatsBySourceIdAndCohortIdAndConceptIdsAndCoho
 	// which is a better performing SQL in this particular scenario:
 	query = QueryFilterByConceptIdsHelper(query, sourceId, filterConceptIds, omopDataSource, resultsDataSource.Schema, "observation")
 
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
 	meta_result := query.Group("observation.value_as_concept_id").
 		Scan(&conceptBreakdownList)
 

models/source.go

@@ -17,20 +17,24 @@ type Source struct {
 func (h Source) GetSourceById(id int) (*Source, error) {
 	db2 := db.GetAtlasDB().Db
 	var dataSource *Source
-	db2.Model(&Source{}).
+	query := db2.Model(&Source{}).
 		Select("source_id, source_name").
-		Where("source_id = ?", id).
-		Scan(&dataSource)
+		Where("source_id = ?", id)
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	query.Scan(&dataSource)
 	return dataSource, nil
 }
 
 func (h Source) GetSourceByIdWithConnection(id int) (*Source, error) {
 	db2 := db.GetAtlasDB().Db
 	var dataSource *Source
-	db2.Model(&Source{}).
+	query := db2.Model(&Source{}).
 		Select("source_id, source_name, source_connection, source_dialect, username, password").
-		Where("source_id = ?", id).
-		Scan(&dataSource)
+		Where("source_id = ?", id)
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	query.Scan(&dataSource)
 	return dataSource, nil
 }
 
@@ -42,12 +46,14 @@ func (h Source) GetSourceSchemaNameBySourceIdAndSourceType(id int, sourceType So
 	atlasDb := db.GetAtlasDB()
 	db2 := atlasDb.Db
 	var sourceSchema *SourceSchema
-	db2.Model(&Source{}).
+	query := db2.Model(&Source{}).
 		Select("source_daimon.table_qualifier as schema_name").
 		Joins("INNER JOIN "+atlasDb.Schema+".source_daimon ON source.source_id = source_daimon.source_id").
 		Where("source.source_id = ?", id).
-		Where("source_daimon.daimon_type = ?", sourceType).
-		Scan(&sourceSchema)
+		Where("source_daimon.daimon_type = ?", sourceType)
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	query.Scan(&sourceSchema)
 	return sourceSchema, nil
 }
 
@@ -74,18 +80,22 @@ func (h Source) GetDataSource(sourceId int, sourceType SourceType) *utils.DbAndS
 func (h Source) GetSourceByName(name string) (*Source, error) {
 	db2 := db.GetAtlasDB().Db
 	var dataSource *Source
-	db2.Model(&Source{}).
+	query := db2.Model(&Source{}).
 		Select("source_id, source_name").
-		Where("source_name = ?", name).
-		Scan(&dataSource)
+		Where("source_name = ?", name)
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	query.Scan(&dataSource)
 	return dataSource, nil
 }
 
 func (h Source) GetAllSources() ([]*Source, error) {
 	db2 := db.GetAtlasDB().Db
 	var dataSource []*Source
-	db2.Model(&Source{}).
-		Select("source_id, source_name").
-		Scan(&dataSource)
+	query := db2.Model(&Source{}).
+		Select("source_id, source_name")
+	query, cancel := utils.AddTimeoutToQuery(query)
+	defer cancel()
+	query.Scan(&dataSource)
 	return dataSource, nil
 }

tests/models_tests/models_test.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/uc-cdis/cohort-middleware/config"
 	"github.com/uc-cdis/cohort-middleware/db"
@@ -889,3 +890,28 @@ func TestRetrieveDataByOriginalCohortAndNewCohort(t *testing.T) {
 		}
 	}
 }
+
+func TestAddTimeoutToQuery(t *testing.T) {
+	setUp(t)
+
+	// take a simple query, run with short timeout, and expect error:
+	db2 := db.GetAtlasDB().Db
+	var dataSource []*models.Source
+	query := db2.Model(&models.Source{}).
+		Select("source_id, source_name")
+	query, cancel := utils.AddSpecificTimeoutToQuery(query, 2*time.Nanosecond)
+	defer cancel()
+	meta_result := query.Scan(&dataSource)
+	if meta_result.Error == nil || len(dataSource) > 0 {
+		t.Errorf("Expected timeout error and NO data")
+	}
+
+	// then switch to default (longer) timeout and expect a result:
+	query2, cancel2 := utils.AddTimeoutToQuery(query)
+	defer cancel2()
+	meta_result2 := query2.Scan(&dataSource)
+
+	if meta_result2.Error != nil || len(dataSource) == 0 {
+		t.Errorf("Expected result and NO error")
+	}
+}

utils/db.go

@@ -1,8 +1,10 @@
 package utils
 
 import (
+	"context"
 	"log"
 	"strings"
+	"time"
 
 	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlserver"
@@ -56,6 +58,20 @@ func GetDataSourceDB(sourceConnectionString string, dbSchema string) *DbAndSchem
 	return dataSourceDb
 }
 
+// Adds a default timeout to a query
+func AddTimeoutToQuery(query *gorm.DB) (*gorm.DB, context.CancelFunc) {
+	// default timeout of 3 minutes:
+	query, cancel := AddSpecificTimeoutToQuery(query, 180*time.Second)
+	return query, cancel
+}
+
+// Adds a specific timeout to a query
+func AddSpecificTimeoutToQuery(query *gorm.DB, timeout time.Duration) (*gorm.DB, context.CancelFunc) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	query = query.WithContext(ctx)
+	return query, cancel
+}
+
 // Returns extra DB dialect specific directives to optimize performance when using views:
 func (h DbAndSchema) GetViewDirective() string {
 	if h.Vendor == "sqlserver" {