diff --git a/postgres-persistence/src/main/java/com/netflix/conductor/postgres/util/PostgresIndexQueryBuilder.java b/postgres-persistence/src/main/java/com/netflix/conductor/postgres/util/PostgresIndexQueryBuilder.java index 5dce1c4d9..6a88d6f97 100644 --- a/postgres-persistence/src/main/java/com/netflix/conductor/postgres/util/PostgresIndexQueryBuilder.java +++ b/postgres-persistence/src/main/java/com/netflix/conductor/postgres/util/PostgresIndexQueryBuilder.java @@ -187,7 +187,7 @@ public String getCountQuery() { .map(c -> c.getQueryFragment()) .collect(Collectors.toList())); } - return "SELECT COUNT(json_data::TEXT) FROM " + table + queryString; + return "SELECT COUNT(json_data) FROM " + table + queryString; } public void addParameters(Query q) throws SQLException { diff --git a/postgres-persistence/src/test/java/com/netflix/conductor/postgres/util/PostgresIndexQueryBuilderTest.java b/postgres-persistence/src/test/java/com/netflix/conductor/postgres/util/PostgresIndexQueryBuilderTest.java index c2ea72d22..759d92d81 100644 --- a/postgres-persistence/src/test/java/com/netflix/conductor/postgres/util/PostgresIndexQueryBuilderTest.java +++ b/postgres-persistence/src/test/java/com/netflix/conductor/postgres/util/PostgresIndexQueryBuilderTest.java @@ -48,6 +48,19 @@ void shouldGenerateQueryForEmptyString() throws SQLException { verifyNoMoreInteractions(mockQuery); } + @Test + void shouldGenerateCountQueryForEmptyString() throws SQLException { + String inputQuery = ""; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String generatedQuery = builder.getCountQuery(); + assertEquals("SELECT COUNT(json_data) FROM table_name", generatedQuery); + Query mockQuery = mock(Query.class); + builder.addParameters(mockQuery); + verifyNoMoreInteractions(mockQuery); + } + @Test void shouldGenerateQueryForNull() throws SQLException { String inputQuery = null; @@ -65,6 +78,19 @@ void shouldGenerateQueryForNull() throws SQLException { verifyNoMoreInteractions(mockQuery); } + @Test + void shouldGenerateCountQueryForNull() throws SQLException { + String inputQuery = null; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String generatedQuery = builder.getCountQuery(); + assertEquals("SELECT COUNT(json_data) FROM table_name", generatedQuery); + Query mockQuery = mock(Query.class); + builder.addParameters(mockQuery); + verifyNoMoreInteractions(mockQuery); + } + @Test void shouldGenerateQueryForWorkflowId() throws SQLException { String inputQuery = "workflowId=\"abc123\""; @@ -85,6 +111,22 @@ void shouldGenerateQueryForWorkflowId() throws SQLException { verifyNoMoreInteractions(mockQuery); } + @Test + void shouldGenerateCountQueryForWorkflowId() throws SQLException { + String inputQuery = "workflowId=\"abc123\""; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String generatedQuery = builder.getCountQuery(); + assertEquals( + "SELECT COUNT(json_data) FROM table_name WHERE workflow_id = ?", generatedQuery); + Query mockQuery = mock(Query.class); + builder.addParameters(mockQuery); + InOrder inOrder = Mockito.inOrder(mockQuery); + inOrder.verify(mockQuery).addParameter("abc123"); + verifyNoMoreInteractions(mockQuery); + } + @Test void shouldGenerateQueryForMultipleInClause() throws SQLException { String inputQuery = "status IN (COMPLETED,RUNNING)"; @@ -105,6 +147,22 @@ void shouldGenerateQueryForMultipleInClause() throws SQLException { verifyNoMoreInteractions(mockQuery); } + @Test + void shouldGenerateCountQueryForMultipleInClause() throws SQLException { + String inputQuery = "status IN (COMPLETED,RUNNING)"; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String generatedQuery = builder.getCountQuery(); + assertEquals( + "SELECT COUNT(json_data) FROM table_name WHERE status = ANY(?)", generatedQuery); + Query mockQuery = mock(Query.class); + builder.addParameters(mockQuery); + InOrder inOrder = Mockito.inOrder(mockQuery); + inOrder.verify(mockQuery).addParameter(new ArrayList<>(List.of("COMPLETED", "RUNNING"))); + verifyNoMoreInteractions(mockQuery); + } + @Test void shouldGenerateQueryForSingleInClause() throws SQLException { String inputQuery = "status IN (COMPLETED)"; @@ -125,6 +183,21 @@ void shouldGenerateQueryForSingleInClause() throws SQLException { verifyNoMoreInteractions(mockQuery); } + @Test + void shouldGenerateCountQueryForSingleInClause() throws SQLException { + String inputQuery = "status IN (COMPLETED)"; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String generatedQuery = builder.getCountQuery(); + assertEquals("SELECT COUNT(json_data) FROM table_name WHERE status = ?", generatedQuery); + Query mockQuery = mock(Query.class); + builder.addParameters(mockQuery); + InOrder inOrder = Mockito.inOrder(mockQuery); + inOrder.verify(mockQuery).addParameter("COMPLETED"); + verifyNoMoreInteractions(mockQuery); + } + @Test void shouldGenerateQueryForStartTimeGt() throws SQLException { String inputQuery = "startTime>1675702498000"; @@ -145,6 +218,23 @@ void shouldGenerateQueryForStartTimeGt() throws SQLException { verifyNoMoreInteractions(mockQuery); } + @Test + void shouldGenerateCountQueryForStartTimeGt() throws SQLException { + String inputQuery = "startTime>1675702498000"; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String generatedQuery = builder.getCountQuery(); + assertEquals( + "SELECT COUNT(json_data) FROM table_name WHERE start_time > ?::TIMESTAMPTZ", + generatedQuery); + Query mockQuery = mock(Query.class); + builder.addParameters(mockQuery); + InOrder inOrder = Mockito.inOrder(mockQuery); + inOrder.verify(mockQuery).addParameter("2023-02-06T16:54:58Z"); + verifyNoMoreInteractions(mockQuery); + } + @Test void shouldGenerateQueryForStartTimeLt() throws SQLException { String inputQuery = "startTime<1675702498000"; @@ -165,6 +255,23 @@ void shouldGenerateQueryForStartTimeLt() throws SQLException { verifyNoMoreInteractions(mockQuery); } + @Test + void shouldGenerateCountQueryForStartTimeLt() throws SQLException { + String inputQuery = "startTime<1675702498000"; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String generatedQuery = builder.getCountQuery(); + assertEquals( + "SELECT COUNT(json_data) FROM table_name WHERE start_time < ?::TIMESTAMPTZ", + generatedQuery); + Query mockQuery = mock(Query.class); + builder.addParameters(mockQuery); + InOrder inOrder = Mockito.inOrder(mockQuery); + inOrder.verify(mockQuery).addParameter("2023-02-06T16:54:58Z"); + verifyNoMoreInteractions(mockQuery); + } + @Test void shouldGenerateQueryForUpdateTimeGt() throws SQLException { String inputQuery = "updateTime>1675702498000"; @@ -185,6 +292,23 @@ void shouldGenerateQueryForUpdateTimeGt() throws SQLException { verifyNoMoreInteractions(mockQuery); } + @Test + void shouldGenerateCountQueryForUpdateTimeGt() throws SQLException { + String inputQuery = "updateTime>1675702498000"; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String generatedQuery = builder.getCountQuery(); + assertEquals( + "SELECT COUNT(json_data) FROM table_name WHERE update_time > ?::TIMESTAMPTZ", + generatedQuery); + Query mockQuery = mock(Query.class); + builder.addParameters(mockQuery); + InOrder inOrder = Mockito.inOrder(mockQuery); + inOrder.verify(mockQuery).addParameter("2023-02-06T16:54:58Z"); + verifyNoMoreInteractions(mockQuery); + } + @Test void shouldGenerateQueryForUpdateTimeLt() throws SQLException { String inputQuery = "updateTime<1675702498000"; @@ -205,6 +329,23 @@ void shouldGenerateQueryForUpdateTimeLt() throws SQLException { verifyNoMoreInteractions(mockQuery); } + @Test + void shouldGenerateCountQueryForUpdateTimeLt() throws SQLException { + String inputQuery = "updateTime<1675702498000"; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String generatedQuery = builder.getCountQuery(); + assertEquals( + "SELECT COUNT(json_data) FROM table_name WHERE update_time < ?::TIMESTAMPTZ", + generatedQuery); + Query mockQuery = mock(Query.class); + builder.addParameters(mockQuery); + InOrder inOrder = Mockito.inOrder(mockQuery); + inOrder.verify(mockQuery).addParameter("2023-02-06T16:54:58Z"); + verifyNoMoreInteractions(mockQuery); + } + @Test void shouldGenerateQueryForMultipleConditions() throws SQLException { String inputQuery = @@ -230,6 +371,28 @@ void shouldGenerateQueryForMultipleConditions() throws SQLException { verifyNoMoreInteractions(mockQuery); } + @Test + void shouldGenerateCountQueryForMultipleConditions() throws SQLException { + String inputQuery = + "workflowId=\"abc123\" AND workflowType IN (one,two) AND status IN (COMPLETED,RUNNING) AND startTime>1675701498000 AND startTime<1675702498000"; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String generatedQuery = builder.getCountQuery(); + assertEquals( + "SELECT COUNT(json_data) FROM table_name WHERE start_time < ?::TIMESTAMPTZ AND start_time > ?::TIMESTAMPTZ AND status = ANY(?) AND workflow_id = ? AND workflow_type = ANY(?)", + generatedQuery); + Query mockQuery = mock(Query.class); + builder.addParameters(mockQuery); + InOrder inOrder = Mockito.inOrder(mockQuery); + inOrder.verify(mockQuery).addParameter("2023-02-06T16:54:58Z"); + inOrder.verify(mockQuery).addParameter("2023-02-06T16:38:18Z"); + inOrder.verify(mockQuery).addParameter(new ArrayList<>(List.of("COMPLETED", "RUNNING"))); + inOrder.verify(mockQuery).addParameter("abc123"); + inOrder.verify(mockQuery).addParameter(new ArrayList<>(List.of("one", "two"))); + verifyNoMoreInteractions(mockQuery); + } + @Test void shouldGenerateOrderBy() throws SQLException { String inputQuery = "updateTime<1675702498000"; @@ -264,6 +427,16 @@ void shouldNotAllowInvalidColumns() throws SQLException { assertEquals(expectedQuery, builder.getQuery()); } + @Test + void shouldNotAllowInvalidColumnsOnCountQuery() throws SQLException { + String inputQuery = "sqlInjection<1675702498000"; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, new ArrayList<>(), properties); + String expectedQuery = "SELECT COUNT(json_data) FROM table_name"; + assertEquals(expectedQuery, builder.getCountQuery()); + } + @Test void shouldNotAllowInvalidSortColumn() throws SQLException { String inputQuery = "updateTime<1675702498000"; @@ -276,6 +449,18 @@ void shouldNotAllowInvalidSortColumn() throws SQLException { assertEquals(expectedQuery, builder.getQuery()); } + @Test + void shouldNotAllowInvalidSortColumnOnCountQuery() throws SQLException { + String inputQuery = "updateTime<1675702498000"; + String[] query = {"sqlInjection:DESC"}; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", inputQuery, "", 0, 15, Arrays.asList(query), properties); + String expectedQuery = + "SELECT COUNT(json_data) FROM table_name WHERE update_time < ?::TIMESTAMPTZ"; + assertEquals(expectedQuery, builder.getCountQuery()); + } + @Test void shouldAllowFullTextSearch() throws SQLException { String freeText = "correlation-id"; @@ -288,6 +473,18 @@ void shouldAllowFullTextSearch() throws SQLException { assertEquals(expectedQuery, builder.getQuery()); } + @Test + void shouldAllowFullTextSearchOnCountQuery() throws SQLException { + String freeText = "correlation-id"; + String[] query = {"sqlInjection:DESC"}; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", "", freeText, 0, 15, Arrays.asList(query), properties); + String expectedQuery = + "SELECT COUNT(json_data) FROM table_name WHERE jsonb_to_tsvector('english', json_data, '[\"all\"]') @@ to_tsquery(?)"; + assertEquals(expectedQuery, builder.getCountQuery()); + } + @Test void shouldAllowJsonSearch() throws SQLException { String freeText = "{\"correlationId\":\"not-the-id\"}"; @@ -300,6 +497,18 @@ void shouldAllowJsonSearch() throws SQLException { assertEquals(expectedQuery, builder.getQuery()); } + @Test + void shouldAllowJsonSearchOnCountQuery() throws SQLException { + String freeText = "{\"correlationId\":\"not-the-id\"}"; + String[] query = {"sqlInjection:DESC"}; + PostgresIndexQueryBuilder builder = + new PostgresIndexQueryBuilder( + "table_name", "", freeText, 0, 15, Arrays.asList(query), properties); + String expectedQuery = + "SELECT COUNT(json_data) FROM table_name WHERE json_data @> ?::JSONB"; + assertEquals(expectedQuery, builder.getCountQuery()); + } + @Test() void shouldThrowIllegalArgumentExceptionWhenQueryStringIsInvalid() { String inputQuery =