Fix query for NodeExecutionsRepo.Count (flyteorg#472)

Signed-off-by: Andrew Dye <[email protected]> Signed-off-by: Andrew Dye <[email protected]>
eapolinario · Sep 7, 2022 · 2e2e8f7 · 2e2e8f7
1 parent 7570147
commit 2e2e8f7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/flyteadmin/pkg/repositories/gormimpl/node_execution_repo.go b/flyteadmin/pkg/repositories/gormimpl/node_execution_repo.go
@@ -164,12 +164,14 @@ func (r *NodeExecutionRepo) Exists(ctx context.Context, input interfaces.NodeExe
 
 func (r *NodeExecutionRepo) Count(ctx context.Context, input interfaces.CountResourceInput) (int64, error) {
 	var err error
-	tx := r.db.Model(&models.NodeExecution{})
+	tx := r.db.Model(&models.NodeExecution{}).Preload("ChildNodeExecutions")
 
 	// Add join condition (joining multiple tables is fine even we only filter on a subset of table attributes).
 	// (this query isn't called for deletes).
-	tx = tx.Joins(innerJoinNodeExecToNodeEvents)
-	tx = tx.Joins(innerJoinExecToNodeExec)
+	tx = tx.Joins(fmt.Sprintf("INNER JOIN %s ON %s.execution_project = %s.execution_project AND "+
+		"%s.execution_domain = %s.execution_domain AND %s.execution_name = %s.execution_name",
+		executionTableName, nodeExecutionTableName, executionTableName,
+		nodeExecutionTableName, executionTableName, nodeExecutionTableName, executionTableName))
 
 	// Apply filters
 	tx, err = applyScopedFilters(tx, input.InlineFilters, input.MapFilters)

diff --git a/flyteadmin/pkg/repositories/gormimpl/node_execution_repo_test.go b/flyteadmin/pkg/repositories/gormimpl/node_execution_repo_test.go
@@ -390,15 +390,15 @@ func TestCountNodeExecutions_Filters(t *testing.T) {
 
 	GlobalMock := mocket.Catcher.Reset()
 	GlobalMock.NewMock().WithQuery(
-		`SELECT count(*) FROM "node_executions" INNER JOIN node_executions ON node_event_executions.node_execution_id = node_executions.id INNER JOIN executions ON node_executions.execution_project = executions.execution_project AND node_executions.execution_domain = executions.execution_domain AND node_executions.execution_name = executions.execution_name WHERE node_executions.phase = $1 AND "error_code" IS NULL`).WithReply([]map[string]interface{}{{"rows": 3}})
+		`SELECT count(*) FROM "node_executions" INNER JOIN executions ON node_executions.execution_project = executions.execution_project AND node_executions.execution_domain = executions.execution_domain AND node_executions.execution_name = executions.execution_name WHERE node_executions.phase = $1 AND "node_executions"."error_code" IS NULL`).WithReply([]map[string]interface{}{{"rows": 3}})
 
 	count, err := nodeExecutionRepo.Count(context.Background(), interfaces.CountResourceInput{
 		InlineFilters: []common.InlineFilter{
 			getEqualityFilter(common.NodeExecution, "phase", core.NodeExecution_FAILED.String()),
 		},
 		MapFilters: []common.MapFilter{
 			common.NewMapFilter(map[string]interface{}{
-				"error_code": nil,
+				"\"node_executions\".\"error_code\"": nil,
 			}),
 		},
 	})