diff --git a/core/src/main/java/google/registry/batch/WipeOutCloudSqlAction.java b/core/src/main/java/google/registry/batch/WipeOutCloudSqlAction.java index 70ad115ff..bbbec41fe 100644 --- a/core/src/main/java/google/registry/batch/WipeOutCloudSqlAction.java +++ b/core/src/main/java/google/registry/batch/WipeOutCloudSqlAction.java @@ -22,7 +22,7 @@ import static javax.servlet.http.HttpServletResponse.SC_OK; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.flogger.FluentLogger; -import google.registry.config.RegistryConfig.Config; +import google.registry.config.RegistryEnvironment; import google.registry.persistence.PersistenceModule.SchemaManagerConnection; import google.registry.request.Action; import google.registry.request.Response; @@ -48,22 +48,18 @@ import javax.inject.Inject; public class WipeOutCloudSqlAction implements Runnable { private static final FluentLogger logger = FluentLogger.forEnclosingClass(); - // As a short-lived class, hardcode allowed projects here instead of using config files. - private static final ImmutableSet ALLOWED_PROJECTS = - ImmutableSet.of("domain-registry-qa"); + private static final ImmutableSet FORBIDDEN_ENVIRONMENTS = + ImmutableSet.of(RegistryEnvironment.PRODUCTION, RegistryEnvironment.SANDBOX); - private final String projectId; private final Supplier connectionSupplier; private final Response response; private final Retrier retrier; @Inject WipeOutCloudSqlAction( - @Config("projectId") String projectId, @SchemaManagerConnection Supplier connectionSupplier, Response response, Retrier retrier) { - this.projectId = projectId; this.connectionSupplier = connectionSupplier; this.response = response; this.retrier = retrier; @@ -73,9 +69,9 @@ public class WipeOutCloudSqlAction implements Runnable { public void run() { response.setContentType(PLAIN_TEXT_UTF_8); - if (!ALLOWED_PROJECTS.contains(projectId)) { + if (FORBIDDEN_ENVIRONMENTS.contains(RegistryEnvironment.get())) { response.setStatus(SC_FORBIDDEN); - response.setPayload("Wipeout is not allowed in " + projectId); + response.setPayload("Wipeout is not allowed in " + RegistryEnvironment.get()); return; } @@ -90,11 +86,11 @@ public class WipeOutCloudSqlAction implements Runnable { }, e -> !(e instanceof SQLException)); response.setStatus(SC_OK); - response.setPayload("Wiped out Cloud SQL in " + projectId); + response.setPayload("Wiped out Cloud SQL in " + RegistryEnvironment.get()); } catch (RuntimeException e) { logger.atSevere().withCause(e).log("Failed to wipe out Cloud SQL data."); response.setStatus(SC_INTERNAL_SERVER_ERROR); - response.setPayload("Failed to wipe out Cloud SQL in " + projectId); + response.setPayload("Failed to wipe out Cloud SQL in " + RegistryEnvironment.get()); } } diff --git a/core/src/main/java/google/registry/batch/WipeoutDatastoreAction.java b/core/src/main/java/google/registry/batch/WipeoutDatastoreAction.java index a2adcfa09..6d89f00fa 100644 --- a/core/src/main/java/google/registry/batch/WipeoutDatastoreAction.java +++ b/core/src/main/java/google/registry/batch/WipeoutDatastoreAction.java @@ -28,6 +28,7 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.flogger.FluentLogger; import google.registry.config.RegistryConfig.Config; +import google.registry.config.RegistryEnvironment; import google.registry.request.Action; import google.registry.request.Response; import google.registry.request.auth.Auth; @@ -49,9 +50,8 @@ public class WipeoutDatastoreAction implements Runnable { private static final String PIPELINE_NAME = "bulk_delete_datastore_pipeline"; - // As a short-lived class, hardcode allowed projects here instead of using config files. - private static final ImmutableSet ALLOWED_PROJECTS = - ImmutableSet.of("domain-registry-qa"); + private static final ImmutableSet FORBIDDEN_ENVIRONMENTS = + ImmutableSet.of(RegistryEnvironment.PRODUCTION, RegistryEnvironment.SANDBOX); private final String projectId; private final String jobRegion; @@ -80,9 +80,9 @@ public class WipeoutDatastoreAction implements Runnable { public void run() { response.setContentType(PLAIN_TEXT_UTF_8); - if (!ALLOWED_PROJECTS.contains(projectId)) { + if (FORBIDDEN_ENVIRONMENTS.contains(RegistryEnvironment.get())) { response.setStatus(SC_FORBIDDEN); - response.setPayload("Wipeout is not allowed in " + projectId); + response.setPayload("Wipeout is not allowed in " + RegistryEnvironment.get()); return; } diff --git a/core/src/test/java/google/registry/batch/WipeOutCloudSqlActionTest.java b/core/src/test/java/google/registry/batch/WipeOutCloudSqlActionTest.java index 484bbff09..509944871 100644 --- a/core/src/test/java/google/registry/batch/WipeOutCloudSqlActionTest.java +++ b/core/src/test/java/google/registry/batch/WipeOutCloudSqlActionTest.java @@ -28,6 +28,7 @@ import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; +import google.registry.config.RegistryEnvironment; import google.registry.testing.FakeClock; import google.registry.testing.FakeResponse; import google.registry.testing.FakeSleeper; @@ -73,8 +74,7 @@ public class WipeOutCloudSqlActionTest { @Test void run_projectAllowed() throws Exception { - WipeOutCloudSqlAction action = - new WipeOutCloudSqlAction("domain-registry-qa", () -> conn, response, retrier); + WipeOutCloudSqlAction action = new WipeOutCloudSqlAction(() -> conn, response, retrier); action.run(); assertThat(response.getStatus()).isEqualTo(SC_OK); verify(stmt, times(1)).executeQuery(anyString()); @@ -84,18 +84,21 @@ public class WipeOutCloudSqlActionTest { @Test void run_projectNotAllowed() { - WipeOutCloudSqlAction action = - new WipeOutCloudSqlAction("domain-registry", () -> conn, response, retrier); - action.run(); - assertThat(response.getStatus()).isEqualTo(SC_FORBIDDEN); - verifyNoInteractions(stmt); + try { + RegistryEnvironment.SANDBOX.setup(); + WipeOutCloudSqlAction action = new WipeOutCloudSqlAction(() -> conn, response, retrier); + action.run(); + assertThat(response.getStatus()).isEqualTo(SC_FORBIDDEN); + verifyNoInteractions(stmt); + } finally { + RegistryEnvironment.UNITTEST.setup(); + } } @Test void run_nonRetrieableFailure() throws Exception { doThrow(new SQLException()).when(conn).getMetaData(); - WipeOutCloudSqlAction action = - new WipeOutCloudSqlAction("domain-registry-qa", () -> conn, response, retrier); + WipeOutCloudSqlAction action = new WipeOutCloudSqlAction(() -> conn, response, retrier); action.run(); assertThat(response.getStatus()).isEqualTo(SC_INTERNAL_SERVER_ERROR); verifyNoInteractions(stmt); @@ -104,8 +107,7 @@ public class WipeOutCloudSqlActionTest { @Test void run_retrieableFailure() throws Exception { when(conn.getMetaData()).thenThrow(new RuntimeException()).thenReturn(metaData); - WipeOutCloudSqlAction action = - new WipeOutCloudSqlAction("domain-registry-qa", () -> conn, response, retrier); + WipeOutCloudSqlAction action = new WipeOutCloudSqlAction(() -> conn, response, retrier); action.run(); assertThat(response.getStatus()).isEqualTo(SC_OK); verify(stmt, times(1)).executeQuery(anyString()); diff --git a/core/src/test/java/google/registry/batch/WipeOutDatastoreActionTest.java b/core/src/test/java/google/registry/batch/WipeOutDatastoreActionTest.java index ce3f780be..2fcd94ec6 100644 --- a/core/src/test/java/google/registry/batch/WipeOutDatastoreActionTest.java +++ b/core/src/test/java/google/registry/batch/WipeOutDatastoreActionTest.java @@ -25,6 +25,7 @@ import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; import google.registry.beam.BeamActionTestBase; +import google.registry.config.RegistryEnvironment; import google.registry.testing.FakeClock; import org.junit.jupiter.api.Test; @@ -35,12 +36,22 @@ class WipeOutDatastoreActionTest extends BeamActionTestBase { @Test void run_projectNotAllowed() { - WipeoutDatastoreAction action = - new WipeoutDatastoreAction( - "domain-registry", "us-central1", "gs://some-bucket", clock, response, dataflow); - action.run(); - assertThat(response.getStatus()).isEqualTo(SC_FORBIDDEN); - verifyNoInteractions(dataflow); + try { + RegistryEnvironment.SANDBOX.setup(); + WipeoutDatastoreAction action = + new WipeoutDatastoreAction( + "domain-registry-sandbox", + "us-central1", + "gs://some-bucket", + clock, + response, + dataflow); + action.run(); + assertThat(response.getStatus()).isEqualTo(SC_FORBIDDEN); + verifyNoInteractions(dataflow); + } finally { + RegistryEnvironment.UNITTEST.setup(); + } } @Test