Repository: aurora Updated Branches: refs/heads/master be75c36f6 -> bf7f9b7f9
http://git-wip-us.apache.org/repos/asf/aurora/blob/bf7f9b7f/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.xml ---------------------------------------------------------------------- diff --git a/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.xml b/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.xml new file mode 100644 index 0000000..8258fb1 --- /dev/null +++ b/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.xml @@ -0,0 +1,323 @@ +<?xml version="1.0" encoding="UTF-8" ?> +<!-- + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + --> + +<!DOCTYPE mapper + PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN" + "http://mybatis.org/dtd/mybatis-3-mapper.dtd";> +<mapper namespace="org.apache.aurora.scheduler.storage.db.TaskConfigMapper"> + <cache size="1000" readOnly="true" /> + <insert id="insert" useGeneratedKeys="true" keyColumn="id" keyProperty="result.id"> + INSERT INTO task_configs ( + job_key_id, + creator_user, + service, + num_cpus, + ram_mb, + disk_mb, + priority, + max_task_failures, + production, + contact_email, + executor_name, + executor_data, + ) VALUES ( + ( + SELECT ID + FROM job_keys + WHERE role = #{config.job.role} + AND environment = #{config.job.environment} + AND name = #{config.job.name} + ), + #{config.owner.user}, + #{config.isService}, + #{config.numCpus}, + #{config.ramMb}, + #{config.diskMb}, + #{config.priority}, + #{config.maxTaskFailures}, + #{config.production}, + #{config.contactEmail}, + #{config.executorConfig.name}, + #{config.executorConfig.data}, + ) + </insert> + + <resultMap id="limitConstraintMap" type="org.apache.aurora.gen.LimitConstraint"> + <id column="id"/> + </resultMap> + + <resultMap id="valueConstraintMap" type="org.apache.aurora.gen.ValueConstraint"> + <id column="id"/> + + <collection property="values" ofType="String" columnPrefix="v_"> + <result column="value" /> + </collection> + </resultMap> + + <!-- + A subclass of TaskConstraint is used to work around trouble using TaskConstraint. See docs in + TaskConstraintShim for details. + --> + <resultMap id="taskConstraintMap" type="org.apache.aurora.scheduler.storage.db.shims.TaskConstraintShim"> + <id column="id"/> + <association property="value" resultMap="valueConstraintMap" columnPrefix="v_"/> + <association property="limit" resultMap="limitConstraintMap" columnPrefix="l_"/> + </resultMap> + + <resultMap id="constraintMap" type="org.apache.aurora.gen.Constraint"> + <id column="id"/> + <association property="constraint" resultMap="taskConstraintMap"/> + </resultMap> + + <select id="selectConstraints" resultMap="constraintMap"> + SELECT + tc.id AS id, + tc.name AS name, + tlc.id AS l_id, + tlc.value AS l_limit, + tvc.id AS v_id, + tvc.negated AS v_negated, + tvcv.value AS v_v_value + FROM task_constraints AS tc + LEFT OUTER JOIN limit_constraints AS tlc ON tlc.constraint_id = tc.id + LEFT OUTER JOIN value_constraints AS tvc ON tvc.constraint_id = tc.id + LEFT OUTER JOIN value_constraint_values AS tvcv ON tvcv.value_constraint_id = tvc.id + WHERE tc.task_config_id = #{id} + </select> + + <resultMap id="dockerContainerMap" type="org.apache.aurora.gen.DockerContainer"> + <result column="image" property="image"/> + </resultMap> + + <resultMap id="containerMap" type="org.apache.aurora.scheduler.storage.db.shims.ContainerShim"> + <association property="docker" resultMap="dockerContainerMap"/> + </resultMap> + + <resultMap id="metadataMap" type="org.apache.aurora.gen.Metadata"> + <id column="id" /> + </resultMap> + + <resultMap id="taskConfigMap" type="org.apache.aurora.gen.TaskConfig"> + <id column="id" /> + <result column="j_role" property="owner.role"/> + <result column="j_environment" property="environment"/> + <result column="j_name" property="jobName"/> + <result column="creator_user" property="owner.user"/> + <result column="executor_name" property="executorConfig.name"/> + <result column="executor_data" property="executorConfig.data"/> + <association + property="job" + resultMap="org.apache.aurora.scheduler.storage.db.JobKeyMapper.jobKeyMap" + columnPrefix="j_"/> + <association property="container" resultMap="containerMap" columnPrefix="c_"/> + <collection + property="constraints" + select="selectConstraints" + column="id" + foreignColumn="task_config_id"/> + <collection property="requestedPorts" ofType="String" columnPrefix="p_"> + <result column="port_name" /> + </collection> + <collection property="metadata" resultMap="metadataMap" columnPrefix="m_"/> + </resultMap> + + <sql id="unscopedConfigSelect"> + SELECT + c.id AS id, + c.creator_user AS creator_user, + c.service AS is_service, + c.num_cpus AS num_cpus, + c.ram_mb AS ram_mb, + c.disk_mb AS disk_mb, + c.priority AS priority, + c.max_task_failures AS max_task_failures, + c.production AS production, + c.contact_email AS contact_email, + c.executor_name AS executor_name, + c.executor_data AS executor_data, + j.role AS j_role, + j.environment AS j_environment, + j.name AS j_name, + p.port_name AS p_port_name, + d.id AS c_id, + d.image AS c_image, + m.id AS m_id, + m.key AS m_key, + m.value AS m_value + FROM task_configs AS c + INNER JOIN job_keys AS j ON j.id = c.job_key_id + LEFT OUTER JOIN task_config_requested_ports AS p ON p.task_config_id = c.id + LEFT OUTER JOIN task_config_docker_containers AS d ON d.task_config_id = c.id + LEFT OUTER JOIN task_config_metadata AS m ON m.task_config_id = c.id + </sql> + + <select id="selectConfig" resultMap="taskConfigMap"> + <include refid="unscopedConfigSelect"/> + WHERE c.id = #{id} + </select> + + <resultMap + id="taskConfigRowMap" + type="org.apache.aurora.scheduler.storage.db.views.TaskConfigRow"> + + <id column="id" property="id" /> + <association property="config" resultMap="taskConfigMap"/> + </resultMap> + + <select id="selectConfigsByJob" resultMap="taskConfigRowMap"> + <include refid="unscopedConfigSelect"/> + WHERE j.role = #{role} + AND j.environment = #{environment} + AND j.name = #{name} + </select> + + <select id="selectConfigsByTaskId" resultType="long"> + SELECT + c.id AS id + FROM task_configs AS c + INNER JOIN tasks AS t ON t.task_config_row_id = c.id + WHERE t.task_id IN ( + <foreach item="taskId" collection="taskIds" separator=","> + #{taskId} + </foreach> + ) + </select> + + <select id="selectTasksByConfigId" resultType="long"> + SELECT + t.id AS id + FROM tasks AS t + INNER JOIN task_configs AS c ON c.id = t.task_config_row_id + WHERE c.id IN ( + <foreach item="configId" collection="configIds" separator=","> + #{configId} + </foreach> + ) + </select> + + <insert id="insertConstraint" useGeneratedKeys="true" keyColumn="id" keyProperty="result.id"> + INSERT INTO task_constraints ( + task_config_id, + name + ) VALUES ( + #{configId}, + #{constraint.name} + ) + </insert> + + <insert id="insertLimitConstraint"> + INSERT INTO limit_constraints ( + constraint_id, + value + ) VALUES ( + #{constraintId}, + #{constraint.limit} + ) + </insert> + + <insert id="insertValueConstraint" useGeneratedKeys="true" keyColumn="id" keyProperty="result.id"> + INSERT INTO value_constraints ( + constraint_id, + negated + ) VALUES ( + #{constraintId}, + #{constraint.negated} + ) + </insert> + + <insert id="insertValueConstraintValues"> + INSERT INTO value_constraint_values ( + value_constraint_id, + value + ) VALUES ( + <foreach item="value" collection="values" separator="),("> + #{valueConstraintId}, + #{value} + </foreach> + ) + </insert> + + <insert id="insertRequestedPorts"> + INSERT INTO task_config_requested_ports ( + task_config_id, + port_name + ) VALUES ( + <foreach item="port" collection="ports" separator="),("> + #{configId}, + #{port} + </foreach> + ) + </insert> + + <insert id="insertTaskLinks" > + INSERT INTO task_config_task_links ( + task_config_id, + label, + url + ) VALUES ( + <foreach item="link" collection="links" separator="),("> + #{configId}, + #{link.label}, + #{link.url} + </foreach> + ) + </insert> + + <resultMap id="taskLinkMap" type="org.apache.aurora.scheduler.storage.db.views.TaskLink"> + <id column="id"/> + </resultMap> + + <select id="selectTaskLinks" resultMap="taskLinkMap"> + SELECT + id, + label, + url + FROM task_config_task_links + WHERE task_config_id = #{configId} + </select> + + <insert id="insertContainer"> + INSERT INTO task_config_docker_containers ( + task_config_id, + image + ) VALUES ( + #{configId}, + #{container.image} + ) + </insert> + + <insert id="insertMetadata"> + INSERT INTO task_config_metadata ( + task_config_id, + key, + value + ) VALUES ( + <foreach item="entry" collection="metadata" separator="),("> + #{configId}, + #{entry.key}, + #{entry.value} + </foreach> + ) + </insert> + + <delete id="delete"> + DELETE FROM task_configs + WHERE id IN ( + <foreach item="configId" collection="configIds" separator=","> + #{configId} + </foreach> + ) + </delete> +</mapper> http://git-wip-us.apache.org/repos/asf/aurora/blob/bf7f9b7f/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml ---------------------------------------------------------------------- diff --git a/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml b/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml new file mode 100644 index 0000000..c3ab3eb --- /dev/null +++ b/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml @@ -0,0 +1,228 @@ +<?xml version="1.0" encoding="UTF-8" ?> +<!-- + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + --> + +<!DOCTYPE mapper + PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN" + "http://mybatis.org/dtd/mybatis-3-mapper.dtd";> +<mapper namespace="org.apache.aurora.scheduler.storage.db.TaskMapper"> + <cache size="10000" readOnly="true" /> + <insert id="insertScheduledTask" useGeneratedKeys="true" keyColumn="id" keyProperty="taskRowId"> + INSERT INTO tasks ( + task_id, + slave_row_id, + instance_id, + status, + failure_count, + ancestor_task_id, + task_config_row_id, + ) VALUES ( + #{task.assignedTask.taskId}, + ( + SELECT ID + FROM host_attributes + WHERE slave_id = #{task.assignedTask.slaveId} + AND host = #{task.assignedTask.slaveHost} + ), + #{task.assignedTask.instanceId}, + #{task.status, typeHandler=org.apache.aurora.scheduler.storage.db.typehandlers.ScheduleStatusTypeHandler}, + #{task.failureCount}, + #{task.ancestorId}, + #{taskConfigRowId} + ) + </insert> + + <resultMap id="taskEventMap" type="org.apache.aurora.gen.TaskEvent"> + <id column="id"/> + <result property="status" + column="status" + typeHandler="org.apache.aurora.scheduler.storage.db.typehandlers.ScheduleStatusTypeHandler" /> + </resultMap> + + <select id="selectTaskEvents" resultMap="taskEventMap"> + SELECT + e.id AS id, + e.timestamp_ms AS timestamp, + e.status AS status, + e.message AS message, + e.scheduler_host AS scheduler + FROM task_events AS e + WHERE e.task_row_id = #{id} + ORDER BY e.timestamp_ms ASC + </select> + + <resultMap id="scheduledTaskMap" type="org.apache.aurora.gen.ScheduledTask"> + <id column="row_id"/> + <result property="status" + column="status" + typeHandler="org.apache.aurora.scheduler.storage.db.typehandlers.ScheduleStatusTypeHandler" /> + <result column="task_id" property="assignedTask.taskId"/> + <result column="slave_id" property="assignedTask.slaveId"/> + <result column="slave_host" property="assignedTask.slaveHost"/> + <result column="instance_id" property="assignedTask.instanceId"/> + <association + property="assignedTask.task" + select="org.apache.aurora.scheduler.storage.db.TaskConfigMapper.selectConfig" + column="task_config_row_id" + foreignColumn="id"/> + <collection + property="taskEvents" + select="selectTaskEvents" + column="row_id" + foreignColumn="task_row_id"/> + </resultMap> + + <resultMap id="taskWrapperMap" type="org.apache.aurora.scheduler.storage.db.views.ScheduledTaskWrapper"> + <id column="row_id" property="taskRowId"/> + <result column="c_id" property="taskConfigRowId"/> + <association property="task" resultMap="scheduledTaskMap"/> + </resultMap> + + <!-- TODO(wfarner): Consider adding selectById and/or selectByIds methods. The dynamic SQL here + seems to come with a ~3x performance hit. + --> + <select id="select" resultMap="taskWrapperMap"> + SELECT + t.id AS row_id, + t.task_config_row_id AS task_config_row_id, + t.task_id AS task_id, + t.instance_id AS instance_id, + t.status AS status, + t.failure_count AS failure_count, + t.ancestor_task_id AS ancestor_id, + j.role AS c_j_role, + j.environment AS c_j_environment, + j.name AS c_j_name, + h.slave_id AS slave_id, + h.host AS slave_host + FROM tasks AS t + INNER JOIN task_configs as c ON c.id = t.task_config_row_id + INNER JOIN job_keys AS j ON j.id = c.job_key_id + LEFT OUTER JOIN host_attributes AS h ON h.id = t.slave_row_id + <where> + <if test="role != null"> + j.role = #{role} + </if> + <if test="owner != null and owner.role != null"> + AND j.role = #{owner.role} + </if> + <if test="environment != null"> + AND j.environment = #{environment} + </if> + <if test="jobName != null"> + AND j.name = #{jobName} + </if> + <if test="taskIds != null"> + AND t.task_id IN ( + <foreach item="task_id" collection="taskIds" separator=","> + #{task_id} + </foreach> + ) + </if> + <if test="statuses != null and !statuses.isEmpty()"> + AND t.status IN ( + <foreach item="element" collection="statuses" separator=","> + #{element, typeHandler=org.apache.aurora.scheduler.storage.db.typehandlers.ScheduleStatusTypeHandler} + </foreach> + ) + </if> + <if test="instanceIds != null and !instanceIds.isEmpty()"> + AND t.instance_id IN ( + <foreach item="instance_id" collection="instanceIds" separator=","> + #{instance_id} + </foreach> + ) + </if> + <if test="slaveHosts != null"> + AND h.host IN ( + <foreach item="host" collection="slaveHosts" separator=","> + #{host} + </foreach> + ) + </if> + <if test="jobKeys != null"> + AND ( + <foreach item="jobKey" collection="jobKeys" open="(" separator=") OR (" close=")"> + j.role = #{jobKey.role} + AND j.name = #{jobKey.name} + AND j.environment = #{jobKey.environment} + </foreach> + ) + </if> + </where> + </select> + + <insert id="insertTaskEvents"> + INSERT INTO task_events( + task_row_id, + timestamp_ms, + status, + message, + scheduler_host + ) VALUES ( + <foreach item="event" collection="events" separator="),("> + #{taskRowId}, + #{event.timestamp}, + #{event.status, typeHandler=org.apache.aurora.scheduler.storage.db.typehandlers.ScheduleStatusTypeHandler}, + #{event.message}, + #{event.scheduler} + </foreach> + ) + </insert> + + <insert id="insertPorts"> + INSERT INTO task_ports( + task_row_id, + name, + port + ) VALUES ( + <foreach item="port" collection="ports" separator="),("> + #{taskRowId}, + #{port.name}, + #{port.port} + </foreach> + ) + </insert> + + <resultMap id="portMap" type="org.apache.aurora.scheduler.storage.db.views.AssignedPort"> + <id column="id"/> + </resultMap> + + <select id="selectPorts" resultMap="portMap"> + SELECT + id, + name, + port + FROM task_ports + WHERE task_row_id = #{taskRowId} + </select> + + <delete id="truncate"> + <!-- + This assumes cascading deletes will clean up all references. Also, once the job store is + migrated, there will be a clash between deletes on the two that needs to be resolved. At that + point it probably makes sense to remove all of the store-specific truncate verbs and use a + single control. + --> + DELETE FROM tasks + </delete> + + <delete id="deleteTasks"> + DELETE FROM tasks WHERE task_id IN ( + <foreach item="task_id" collection="taskIds" separator=","> + #{task_id} + </foreach> + ) + </delete> +</mapper> http://git-wip-us.apache.org/repos/asf/aurora/blob/bf7f9b7f/src/test/java/org/apache/aurora/scheduler/state/StateManagerImplTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/state/StateManagerImplTest.java b/src/test/java/org/apache/aurora/scheduler/state/StateManagerImplTest.java index afb7db8..15e4d38 100644 --- a/src/test/java/org/apache/aurora/scheduler/state/StateManagerImplTest.java +++ b/src/test/java/org/apache/aurora/scheduler/state/StateManagerImplTest.java @@ -197,6 +197,7 @@ public class StateManagerImplTest extends EasyMockTest { .setScheduler(StateManagerImpl.LOCAL_HOST_SUPPLIER.get()) .setStatus(PENDING))) .setAssignedTask(new AssignedTask() + .setAssignedPorts(ImmutableMap.<String, Integer>of()) .setInstanceId(3) .setTaskId(taskId) .setTask(NON_SERVICE_CONFIG.newBuilder())); http://git-wip-us.apache.org/repos/asf/aurora/blob/bf7f9b7f/src/test/java/org/apache/aurora/scheduler/storage/AbstractTaskStoreTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/AbstractTaskStoreTest.java b/src/test/java/org/apache/aurora/scheduler/storage/AbstractTaskStoreTest.java index 6a6ff27..0f3a1a0 100644 --- a/src/test/java/org/apache/aurora/scheduler/storage/AbstractTaskStoreTest.java +++ b/src/test/java/org/apache/aurora/scheduler/storage/AbstractTaskStoreTest.java @@ -26,6 +26,7 @@ import com.google.common.collect.Iterables; import com.google.common.collect.Maps; import com.google.common.util.concurrent.ThreadFactoryBuilder; import com.google.inject.Guice; +import com.google.inject.Injector; import com.google.inject.Module; import com.twitter.common.quantity.Amount; import com.twitter.common.quantity.Time; @@ -60,31 +61,33 @@ import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; public abstract class AbstractTaskStoreTest { - private static final IHostAttributes HOST_A = IHostAttributes.build( + protected static final IHostAttributes HOST_A = IHostAttributes.build( new HostAttributes( "hostA", ImmutableSet.of(new Attribute("zone", ImmutableSet.of("1a")))) .setSlaveId("slaveIdA") .setMode(MaintenanceMode.NONE)); - private static final IHostAttributes HOST_B = IHostAttributes.build( + protected static final IHostAttributes HOST_B = IHostAttributes.build( new HostAttributes( "hostB", ImmutableSet.of(new Attribute("zone", ImmutableSet.of("1a")))) .setSlaveId("slaveIdB") .setMode(MaintenanceMode.NONE)); protected static final IScheduledTask TASK_A = createTask("a"); - private static final IScheduledTask TASK_B = + protected static final IScheduledTask TASK_B = setContainer(createTask("b"), Container.mesos(new MesosContainer())); - private static final IScheduledTask TASK_C = createTask("c"); - private static final IScheduledTask TASK_D = createTask("d"); + protected static final IScheduledTask TASK_C = createTask("c"); + protected static final IScheduledTask TASK_D = createTask("d"); + protected Injector injector; protected Storage storage; protected abstract Module getStorageModule(); @Before public void baseSetUp() { - storage = Guice.createInjector(getStorageModule()).getInstance(Storage.class); + injector = Guice.createInjector(getStorageModule()); + storage = injector.getInstance(Storage.class); storage.prepare(); storage.write(new Storage.MutateWork.NoResult.Quiet() { @@ -106,7 +109,7 @@ public abstract class AbstractTaskStoreTest { }); } - private void saveTasks(final IScheduledTask... tasks) { + protected void saveTasks(final IScheduledTask... tasks) { saveTasks(ImmutableSet.copyOf(tasks)); } @@ -140,7 +143,7 @@ public abstract class AbstractTaskStoreTest { }); } - private void deleteTasks(final String... taskIds) { + protected void deleteTasks(final String... taskIds) { storage.write(new Storage.MutateWork.NoResult.Quiet() { @Override protected void execute(Storage.MutableStoreProvider storeProvider) { http://git-wip-us.apache.org/repos/asf/aurora/blob/bf7f9b7f/src/test/java/org/apache/aurora/scheduler/storage/db/DbJobUpdateStoreTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/db/DbJobUpdateStoreTest.java b/src/test/java/org/apache/aurora/scheduler/storage/db/DbJobUpdateStoreTest.java index 7d856d0..177d720 100644 --- a/src/test/java/org/apache/aurora/scheduler/storage/db/DbJobUpdateStoreTest.java +++ b/src/test/java/org/apache/aurora/scheduler/storage/db/DbJobUpdateStoreTest.java @@ -278,7 +278,7 @@ public class DbJobUpdateStoreTest { for (Map.Entry<JobUpdateStatus, T> entry : expected.entrySet()) { assertEquals( entry.getValue().longValue(), - stats.getLongValue(DBJobUpdateStore.statName(entry.getKey()))); + stats.getLongValue(DbJobUpdateStore.statName(entry.getKey()))); } } http://git-wip-us.apache.org/repos/asf/aurora/blob/bf7f9b7f/src/test/java/org/apache/aurora/scheduler/storage/db/DbTaskStoreTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/db/DbTaskStoreTest.java b/src/test/java/org/apache/aurora/scheduler/storage/db/DbTaskStoreTest.java new file mode 100644 index 0000000..dda988d --- /dev/null +++ b/src/test/java/org/apache/aurora/scheduler/storage/db/DbTaskStoreTest.java @@ -0,0 +1,66 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.aurora.scheduler.storage.db; + +import com.google.common.collect.ImmutableList; +import com.google.inject.AbstractModule; +import com.google.inject.Module; +import com.google.inject.util.Modules; +import com.twitter.common.stats.StatsProvider; +import com.twitter.common.util.Clock; +import com.twitter.common.util.testing.FakeClock; + +import org.apache.aurora.scheduler.base.Tasks; +import org.apache.aurora.scheduler.storage.AbstractTaskStoreTest; +import org.apache.aurora.scheduler.storage.db.views.TaskConfigRow; +import org.apache.aurora.scheduler.testing.FakeStatsProvider; +import org.junit.Before; +import org.junit.Test; + +import static org.junit.Assert.assertEquals; + +public class DbTaskStoreTest extends AbstractTaskStoreTest { + + private TaskConfigManager configManager; + + @Before + public void setUp() { + configManager = injector.getInstance(TaskConfigManager.class); + } + + @Override + protected Module getStorageModule() { + return Modules.combine( + DbModule.testModule(), + new AbstractModule() { + @Override + protected void configure() { + bind(StatsProvider.class).toInstance(new FakeStatsProvider()); + bind(Clock.class).toInstance(new FakeClock()); + } + }); + } + + @Test + public void testRelationsRemoved() { + // When there are no remaining references to a task config, it should be removed. + saveTasks(TASK_A); + deleteTasks(Tasks.id(TASK_A)); + assertEquals( + ImmutableList.<TaskConfigRow>of(), + configManager.getConfigs(TASK_A.getAssignedTask().getTask().getJob())); + + // TODO(wfarner): Check that the job key was removed. + } +} http://git-wip-us.apache.org/repos/asf/aurora/blob/bf7f9b7f/src/test/java/org/apache/aurora/scheduler/storage/mem/InMemTaskStoreTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/mem/InMemTaskStoreTest.java b/src/test/java/org/apache/aurora/scheduler/storage/mem/InMemTaskStoreTest.java index 8f139fc..999d5e8 100644 --- a/src/test/java/org/apache/aurora/scheduler/storage/mem/InMemTaskStoreTest.java +++ b/src/test/java/org/apache/aurora/scheduler/storage/mem/InMemTaskStoreTest.java @@ -17,6 +17,7 @@ import com.google.common.collect.ImmutableSet; import com.google.inject.AbstractModule; import com.google.inject.Module; import com.google.inject.util.Modules; +import com.twitter.common.inject.Bindings.KeyFactory; import com.twitter.common.stats.StatsProvider; import org.apache.aurora.scheduler.base.Tasks; @@ -37,7 +38,7 @@ public class InMemTaskStoreTest extends AbstractTaskStoreTest { protected Module getStorageModule() { statsProvider = new FakeStatsProvider(); return Modules.combine( - DbModule.testModule(), + DbModule.testModule(new InMemStoresModule.TaskStoreModule(KeyFactory.PLAIN)), new AbstractModule() { @Override protected void configure() { http://git-wip-us.apache.org/repos/asf/aurora/blob/bf7f9b7f/src/test/java/org/apache/aurora/scheduler/storage/mem/StorageTransactionTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/mem/StorageTransactionTest.java b/src/test/java/org/apache/aurora/scheduler/storage/mem/StorageTransactionTest.java index bad9eb5..e5e889a 100644 --- a/src/test/java/org/apache/aurora/scheduler/storage/mem/StorageTransactionTest.java +++ b/src/test/java/org/apache/aurora/scheduler/storage/mem/StorageTransactionTest.java @@ -23,7 +23,6 @@ import java.util.concurrent.Future; import com.google.common.collect.FluentIterable; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; -import com.google.common.collect.Iterables; import com.google.common.testing.TearDown; import com.google.common.testing.junit4.TearDownTestCase; import com.google.common.util.concurrent.ThreadFactoryBuilder; @@ -31,12 +30,9 @@ import com.twitter.common.quantity.Amount; import com.twitter.common.quantity.Time; import com.twitter.common.util.concurrent.ExecutorServiceShutdown; -import org.apache.aurora.gen.AssignedTask; -import org.apache.aurora.gen.Identity; import org.apache.aurora.gen.ResourceAggregate; -import org.apache.aurora.gen.ScheduledTask; -import org.apache.aurora.gen.TaskConfig; import org.apache.aurora.scheduler.base.Query; +import org.apache.aurora.scheduler.base.TaskTestUtil; import org.apache.aurora.scheduler.base.Tasks; import org.apache.aurora.scheduler.storage.Storage; import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider; @@ -113,13 +109,7 @@ public class StorageTransactionTest extends TearDownTestCase { } private IScheduledTask makeTask(String taskId) { - return IScheduledTask.build(new ScheduledTask().setAssignedTask( - new AssignedTask() - .setTaskId(taskId) - .setTask(new TaskConfig() - .setOwner(new Identity().setRole("owner-" + taskId)) - .setJobName("job-" + taskId) - .setEnvironment("env-" + taskId)))); + return TaskTestUtil.makeTask(taskId, TaskTestUtil.JOB); } private static class CustomException extends RuntimeException { @@ -186,7 +176,7 @@ public class StorageTransactionTest extends TearDownTestCase { throw new CustomException(); } }); - expectTasks("a", "b"); + expectTasks(); storage.write(new MutateWork.NoResult.Quiet() { @Override @@ -203,7 +193,14 @@ public class StorageTransactionTest extends TearDownTestCase { throw new CustomException(); } }); - expectTasks(); + expectTasks("a", "b"); + + storage.write(new MutateWork.NoResult.Quiet() { + @Override + protected void execute(MutableStoreProvider storeProvider) { + storeProvider.getUnsafeTaskStore().deleteAllTasks(); + } + }); expectWriteFail(new MutateWork.NoResult.Quiet() { @Override @@ -212,15 +209,12 @@ public class StorageTransactionTest extends TearDownTestCase { throw new CustomException(); } }); - expectTasks("a"); - storage.read(new Work.Quiet<Void>() { + expectTasks(); + + storage.write(new MutateWork.NoResult.Quiet() { @Override - public Void apply(StoreProvider storeProvider) { - assertEquals( - makeTask("a"), - Iterables.getOnlyElement(storeProvider.getTaskStore().fetchTasks( - Query.taskScoped("a")))); - return null; + protected void execute(MutableStoreProvider storeProvider) { + storeProvider.getUnsafeTaskStore().saveTasks(ImmutableSet.of(makeTask("a"))); } }); @@ -238,7 +232,7 @@ public class StorageTransactionTest extends TearDownTestCase { }); } }); - expectTasks("a", "c", "d"); + expectTasks("a"); // Nested transaction where outer transaction fails. expectWriteFail(new MutateWork.NoResult.Quiet() { @@ -254,6 +248,6 @@ public class StorageTransactionTest extends TearDownTestCase { throw new CustomException(); } }); - expectTasks("a", "c", "d"); + expectTasks("a"); } }
