|
41 | 41 | import org.apache.gobblin.metastore.testing.TestMetastoreDatabaseFactory; |
42 | 42 | import org.apache.gobblin.metrics.event.TimingEvent; |
43 | 43 | import org.apache.gobblin.service.ExecutionStatus; |
| 44 | +import org.apache.gobblin.service.ServiceConfigKeys; |
44 | 45 | import org.apache.gobblin.service.modules.flowgraph.Dag; |
45 | 46 | import org.apache.gobblin.service.modules.orchestration.DagManagerTest; |
46 | 47 | import org.apache.gobblin.service.modules.orchestration.DagProcessingEngine; |
47 | 48 | import org.apache.gobblin.service.modules.orchestration.DagTestUtils; |
48 | 49 | import org.apache.gobblin.service.modules.orchestration.MySqlDagManagementStateStore; |
49 | 50 | import org.apache.gobblin.service.modules.orchestration.MySqlDagManagementStateStoreTest; |
| 51 | +import org.apache.gobblin.service.modules.orchestration.UserQuotaManager; |
50 | 52 | import org.apache.gobblin.service.modules.orchestration.proc.LaunchDagProcTest; |
51 | 53 | import org.apache.gobblin.service.modules.spec.JobExecutionPlan; |
52 | 54 | import org.apache.gobblin.service.monitoring.FlowStatus; |
53 | 55 | import org.apache.gobblin.service.monitoring.JobStatus; |
54 | 56 | import org.apache.gobblin.service.monitoring.JobStatusRetriever; |
55 | 57 |
|
56 | 58 | import static org.mockito.ArgumentMatchers.anyString; |
| 59 | +import static org.mockito.Mockito.mock; |
57 | 60 | import static org.mockito.Mockito.spy; |
58 | 61 | import static org.mockito.Mockito.when; |
59 | 62 |
|
@@ -207,6 +210,115 @@ public void testSameFlowExecAlreadyCompiledWithinJobStartDeadline() throws IOExc |
207 | 210 | flowStartTime, this.dagManagementStateStore)); |
208 | 211 | } |
209 | 212 |
|
| 213 | + /** |
| 214 | + * Helper to build a {@link FlowCompilationValidationHelper} with the given service-level config. |
| 215 | + */ |
| 216 | + private FlowCompilationValidationHelper buildHelper(Config serviceConfig) throws Exception { |
| 217 | + ITestMetastoreDatabase db = TestMetastoreDatabaseFactory.get(); |
| 218 | + MySqlDagManagementStateStore dmss = spy(MySqlDagManagementStateStoreTest.getDummyDMSS(db)); |
| 219 | + LaunchDagProcTest.mockDMSSCommonBehavior(dmss); |
| 220 | + SharedFlowMetricsSingleton sharedFlowMetricsSingleton = new SharedFlowMetricsSingleton(serviceConfig); |
| 221 | + return new FlowCompilationValidationHelper(serviceConfig, sharedFlowMetricsSingleton, |
| 222 | + mock(UserQuotaManager.class), dmss); |
| 223 | + } |
| 224 | + |
| 225 | + @Test |
| 226 | + public void testResolveAllowConcurrentExecution_explicitFlowConfigTrue() throws Exception { |
| 227 | + // Per-flow setting should take precedence even when service-level is false and no prefix match |
| 228 | + Config serviceConfig = ConfigFactory.empty() |
| 229 | + .withValue(ServiceConfigKeys.FLOW_CONCURRENCY_ALLOWED, ConfigValueFactory.fromAnyRef(false)); |
| 230 | + FlowCompilationValidationHelper helper = buildHelper(serviceConfig); |
| 231 | + |
| 232 | + Config flowConfig = ConfigFactory.empty() |
| 233 | + .withValue(ConfigurationKeys.FLOW_ALLOW_CONCURRENT_EXECUTION, ConfigValueFactory.fromAnyRef(true)); |
| 234 | + Assert.assertTrue(helper.resolveAllowConcurrentExecution(flowConfig, "unmatched-group")); |
| 235 | + } |
| 236 | + |
| 237 | + @Test |
| 238 | + public void testResolveAllowConcurrentExecution_explicitFlowConfigFalse() throws Exception { |
| 239 | + // Per-flow setting of false should take precedence even when service-level is true |
| 240 | + Config serviceConfig = ConfigFactory.empty() |
| 241 | + .withValue(ServiceConfigKeys.FLOW_CONCURRENCY_ALLOWED, ConfigValueFactory.fromAnyRef(true)); |
| 242 | + FlowCompilationValidationHelper helper = buildHelper(serviceConfig); |
| 243 | + |
| 244 | + Config flowConfig = ConfigFactory.empty() |
| 245 | + .withValue(ConfigurationKeys.FLOW_ALLOW_CONCURRENT_EXECUTION, ConfigValueFactory.fromAnyRef(false)); |
| 246 | + Assert.assertFalse(helper.resolveAllowConcurrentExecution(flowConfig, "any-group")); |
| 247 | + } |
| 248 | + |
| 249 | + @Test |
| 250 | + public void testResolveAllowConcurrentExecution_noFlowConfig_serviceLevelTrue() throws Exception { |
| 251 | + // When flow config is absent and service-level is true, should allow without checking prefixes |
| 252 | + Config serviceConfig = ConfigFactory.empty() |
| 253 | + .withValue(ServiceConfigKeys.FLOW_CONCURRENCY_ALLOWED, ConfigValueFactory.fromAnyRef(true)); |
| 254 | + FlowCompilationValidationHelper helper = buildHelper(serviceConfig); |
| 255 | + |
| 256 | + Assert.assertTrue(helper.resolveAllowConcurrentExecution(ConfigFactory.empty(), "unmatched-group")); |
| 257 | + } |
| 258 | + |
| 259 | + @Test |
| 260 | + public void testResolveAllowConcurrentExecution_noFlowConfig_serviceLevelFalse_prefixMatch() throws Exception { |
| 261 | + // When flow config is absent and service-level is false, should fall back to prefix matching |
| 262 | + Config serviceConfig = ConfigFactory.empty() |
| 263 | + .withValue(ServiceConfigKeys.FLOW_CONCURRENCY_ALLOWED, ConfigValueFactory.fromAnyRef(false)) |
| 264 | + .withValue(ServiceConfigKeys.CONCURRENCY_ALLOWED_FLOWGROUP_PREFIXES, |
| 265 | + ConfigValueFactory.fromAnyRef("teamA,teamB-prod")); |
| 266 | + FlowCompilationValidationHelper helper = buildHelper(serviceConfig); |
| 267 | + |
| 268 | + Assert.assertTrue(helper.resolveAllowConcurrentExecution(ConfigFactory.empty(), "teamA-pipeline")); |
| 269 | + Assert.assertTrue(helper.resolveAllowConcurrentExecution(ConfigFactory.empty(), "teamB-prod-etl")); |
| 270 | + } |
| 271 | + |
| 272 | + @Test |
| 273 | + public void testResolveAllowConcurrentExecution_noFlowConfig_serviceLevelFalse_noPrefixMatch() throws Exception { |
| 274 | + // When flow config is absent, service-level is false, and no prefix matches, should disallow |
| 275 | + Config serviceConfig = ConfigFactory.empty() |
| 276 | + .withValue(ServiceConfigKeys.FLOW_CONCURRENCY_ALLOWED, ConfigValueFactory.fromAnyRef(false)) |
| 277 | + .withValue(ServiceConfigKeys.CONCURRENCY_ALLOWED_FLOWGROUP_PREFIXES, |
| 278 | + ConfigValueFactory.fromAnyRef("teamA,teamB-prod")); |
| 279 | + FlowCompilationValidationHelper helper = buildHelper(serviceConfig); |
| 280 | + |
| 281 | + Assert.assertFalse(helper.resolveAllowConcurrentExecution(ConfigFactory.empty(), "teamC-pipeline")); |
| 282 | + } |
| 283 | + |
| 284 | + @Test |
| 285 | + public void testResolveAllowConcurrentExecution_noFlowConfig_serviceLevelFalse_noPrefixesConfigured() throws Exception { |
| 286 | + // When no prefixes are configured at all, should disallow |
| 287 | + Config serviceConfig = ConfigFactory.empty() |
| 288 | + .withValue(ServiceConfigKeys.FLOW_CONCURRENCY_ALLOWED, ConfigValueFactory.fromAnyRef(false)); |
| 289 | + FlowCompilationValidationHelper helper = buildHelper(serviceConfig); |
| 290 | + |
| 291 | + Assert.assertFalse(helper.resolveAllowConcurrentExecution(ConfigFactory.empty(), "any-group")); |
| 292 | + } |
| 293 | + |
| 294 | + @Test |
| 295 | + public void testResolveAllowConcurrentExecution_noFlowConfig_serviceLevelFalse_blankPrefixesIgnored() throws Exception { |
| 296 | + // Blank entries from trailing/double commas should not match all flow groups |
| 297 | + Config serviceConfig = ConfigFactory.empty() |
| 298 | + .withValue(ServiceConfigKeys.FLOW_CONCURRENCY_ALLOWED, ConfigValueFactory.fromAnyRef(false)) |
| 299 | + .withValue(ServiceConfigKeys.CONCURRENCY_ALLOWED_FLOWGROUP_PREFIXES, |
| 300 | + ConfigValueFactory.fromAnyRef("teamA,,teamB, ")); |
| 301 | + FlowCompilationValidationHelper helper = buildHelper(serviceConfig); |
| 302 | + |
| 303 | + Assert.assertTrue(helper.resolveAllowConcurrentExecution(ConfigFactory.empty(), "teamA-pipeline")); |
| 304 | + Assert.assertTrue(helper.resolveAllowConcurrentExecution(ConfigFactory.empty(), "teamB-pipeline")); |
| 305 | + Assert.assertFalse(helper.resolveAllowConcurrentExecution(ConfigFactory.empty(), "teamC-pipeline")); |
| 306 | + } |
| 307 | + |
| 308 | + @Test |
| 309 | + public void testResolveAllowConcurrentExecution_explicitFlowConfigFalse_overridesPrefixMatch() throws Exception { |
| 310 | + // Explicit per-flow false should take precedence even when prefix matches |
| 311 | + Config serviceConfig = ConfigFactory.empty() |
| 312 | + .withValue(ServiceConfigKeys.FLOW_CONCURRENCY_ALLOWED, ConfigValueFactory.fromAnyRef(false)) |
| 313 | + .withValue(ServiceConfigKeys.CONCURRENCY_ALLOWED_FLOWGROUP_PREFIXES, |
| 314 | + ConfigValueFactory.fromAnyRef("teamA")); |
| 315 | + FlowCompilationValidationHelper helper = buildHelper(serviceConfig); |
| 316 | + |
| 317 | + Config flowConfig = ConfigFactory.empty() |
| 318 | + .withValue(ConfigurationKeys.FLOW_ALLOW_CONCURRENT_EXECUTION, ConfigValueFactory.fromAnyRef(false)); |
| 319 | + Assert.assertFalse(helper.resolveAllowConcurrentExecution(flowConfig, "teamA-pipeline")); |
| 320 | + } |
| 321 | + |
210 | 322 | private void insertFlowIntoDMSSMock(String flowGroup, String flowName, long flowStartTime, ExecutionStatus executionStatus, Config config) |
211 | 323 | throws URISyntaxException, IOException { |
212 | 324 | List<FlowStatus> list = new ArrayList<>(); |
|
0 commit comments