trinodb · ebyhr · Feb 3, 2025 · Jan 31, 2025
diff --git a/plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiRecordSet.java b/plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiRecordSet.java
@@ -49,7 +49,7 @@ public LokiRecordSet(LokiClient lokiClient, LokiSplit split, List<LokiColumnHand
 
         // Execute the query
         try {
-            this.result = lokiClient.rangeQuery(split.query(), split.start(), split.end());
+            this.result = lokiClient.rangeQuery(split.query(), split.start(), split.end(), split.step());
         }
         catch (LokiClientException e) {
             throw new TrinoException(LOKI_CLIENT_ERROR, e);

diff --git a/plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiSplit.java b/plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiSplit.java
@@ -19,7 +19,7 @@
 
 import static java.util.Objects.requireNonNull;
 
-public record LokiSplit(String query, Instant start, Instant end)
+public record LokiSplit(String query, Instant start, Instant end, int step)
         implements ConnectorSplit
 {
     public LokiSplit

diff --git a/plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiSplitManager.java b/plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiSplitManager.java
@@ -42,7 +42,7 @@ public ConnectorSplitSource getSplits(
     {
         final LokiTableHandle table = (LokiTableHandle) connectorTableHandle;
 
-        List<ConnectorSplit> splits = ImmutableList.of(new LokiSplit(table.query(), table.start(), table.end()));
+        List<ConnectorSplit> splits = ImmutableList.of(new LokiSplit(table.query(), table.start(), table.end(), table.step()));
 
         log.debug("created %d splits", splits.size());
         return new FixedSplitSource(splits);

diff --git a/plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiTableHandle.java b/plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiTableHandle.java
@@ -22,7 +22,7 @@
 
 import static java.util.Objects.requireNonNull;
 
-public record LokiTableHandle(String query, Instant start, Instant end, List<ColumnHandle> columnHandles)
+public record LokiTableHandle(String query, Instant start, Instant end, int step, List<ColumnHandle> columnHandles)
         implements ConnectorTableHandle
 {
     public LokiTableHandle

diff --git a/plugin/trino-loki/src/main/java/io/trino/plugin/loki/QueryRangeTableFunction.java b/plugin/trino-loki/src/main/java/io/trino/plugin/loki/QueryRangeTableFunction.java
@@ -41,6 +41,7 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
 import static io.trino.spi.function.table.ReturnTypeSpecification.GenericTable.GENERIC_TABLE;
+import static io.trino.spi.type.IntegerType.INTEGER;
 import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_NANOS;
 import static io.trino.spi.type.Timestamps.MILLISECONDS_PER_SECOND;
 import static io.trino.spi.type.Timestamps.NANOSECONDS_PER_MILLISECOND;
@@ -74,6 +75,11 @@ public QueryRangeTableFunction(LokiMetadata metadata)
                         ScalarArgumentSpecification.builder()
                                 .name("END")
                                 .type(TIMESTAMP_TZ_NANOS)
+                                .build(),
+                        ScalarArgumentSpecification.builder()
+                                .name("STEP")
+                                .type(INTEGER)
+                                .defaultValue(0L)
                                 .build()),
                 GENERIC_TABLE);
 
@@ -89,6 +95,11 @@ public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransact
         LongTimestampWithTimeZone startArgument = (LongTimestampWithTimeZone) ((ScalarArgument) arguments.get("START")).getValue();
         LongTimestampWithTimeZone endArgument = (LongTimestampWithTimeZone) ((ScalarArgument) arguments.get("END")).getValue();
 
+        Long step = (Long) ((ScalarArgument) arguments.get("STEP")).getValue();
+        if (step == null || step < 0L) {
+            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "step must be positive");
+        }
+
         if (Strings.isNullOrEmpty(query)) {
             throw new TrinoException(INVALID_FUNCTION_ARGUMENT, query);
         }
@@ -118,6 +129,7 @@ public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransact
                 query,
                 start,
                 end,
+                step.intValue(),
                 columnHandles);
 
         return TableFunctionAnalysis.builder()

diff --git a/plugin/trino-loki/src/test/java/io/trino/plugin/loki/TestLokiIntegration.java b/plugin/trino-loki/src/test/java/io/trino/plugin/loki/TestLokiIntegration.java
@@ -21,13 +21,11 @@
 
 import java.time.Duration;
 import java.time.Instant;
-import java.time.LocalDate;
 import java.time.ZoneId;
 import java.time.ZoneOffset;
 import java.time.format.DateTimeFormatter;
 import java.time.temporal.ChronoUnit;
 
-import static io.trino.type.DateTimes.MILLISECONDS_PER_DAY;
 import static java.lang.String.format;
 
 final class TestLokiIntegration
@@ -37,6 +35,7 @@ final class TestLokiIntegration
 
     private static final DateTimeFormatter timestampFormatter = DateTimeFormatter.ofPattern("uuuu-MM-dd HH:mm:ss.SSS'Z'").withZone(ZoneOffset.UTC);
     private static final DateTimeFormatter timestampFormatterAtEasternTime = DateTimeFormatter.ofPattern("uuuu-MM-dd HH:mm:ss'-05:00'").withZone(ZoneId.of("US/Eastern"));
+    private static final DateTimeFormatter isoTimestampFormatter = DateTimeFormatter.ofPattern("uuuu-MM-dd'T'HH:mm:ss.SSS'Z'").withZone(ZoneOffset.UTC);
 
     @Override
     protected QueryRunner createQueryRunner()
@@ -155,7 +154,6 @@ void testLabelsComplex()
     void testSelectTimestampLogsQuery()
             throws Exception
     {
-        DateTimeFormatter isoTimestampFormatter = DateTimeFormatter.ofPattern("uuuu-MM-dd'T'HH:mm:ss.SSS'Z'").withZone(ZoneOffset.UTC);
         Instant start = Instant.now().truncatedTo(ChronoUnit.DAYS).minus(Duration.ofHours(12));
         Instant end = start.plus(Duration.ofHours(4));
         Instant firstLineTimestamp = start.truncatedTo(ChronoUnit.MILLIS);
@@ -184,29 +182,58 @@ void testSelectTimestampLogsQuery()
     void testTimestampMetricsQuery()
             throws Exception
     {
-        LocalDate baseLineDate = LocalDate.now();
-        Instant start = Instant.ofEpochMilli(baseLineDate.toEpochDay() * MILLISECONDS_PER_DAY);
-        Instant end = start.plus(Duration.ofHours(4));
+        Instant start = Instant.now().truncatedTo(ChronoUnit.HOURS).minus(Duration.ofHours(4));
+        Instant end = start.plus(Duration.ofHours(3));
 
-        this.client.pushLogLine("line 1", start.plus(Duration.ofHours(1)), ImmutableMap.of("test", "timestamp_metrics_query"));
+        this.client.pushLogLine("line 1", start.plus(Duration.ofMinutes(4)), ImmutableMap.of("test", "timestamp_metrics_query"));
         this.client.pushLogLine("line 2", start.plus(Duration.ofHours(2)), ImmutableMap.of("test", "timestamp_metrics_query"));
         this.client.pushLogLine("line 3", start.plus(Duration.ofHours(3)), ImmutableMap.of("test", "timestamp_metrics_query"));
         this.client.flush();
         assertQuery(format("""
-                        SELECT CAST(timestamp AS DATE) FROM
+                        SELECT to_iso8601(timestamp), value FROM
                         TABLE(system.query_range(
                          'count_over_time({test="timestamp_metrics_query"}[5m])',
                          TIMESTAMP '%s',
-                         TIMESTAMP '%s'
+                         TIMESTAMP '%s',
+                         300
                         ))
                         LIMIT 1
                         """, timestampFormatter.format(start), timestampFormatter.format(end)),
-                "VALUES DATE '%s'".formatted(baseLineDate));
+                "VALUES ('%s', 1.0)".formatted(isoTimestampFormatter.format(start.plus(Duration.ofMinutes(5)))));
     }
 
     @Test
     void testSelectFromTableFails()
     {
         assertQueryFails("SELECT * FROM default", "Loki connector does not support querying tables directly. Use the TABLE function instead.");
     }
+
+    @Test
+    void testQueryRangeInvalidArguments()
+    {
+        assertQueryFails(
+                """
+                SELECT to_iso8601(timestamp), value FROM
+                TABLE(system.query_range(
+                 'count_over_time({test="timestamp_metrics_query"}[5m])',
+                 TIMESTAMP '2012-08-08',
+                 TIMESTAMP '2012-08-09',
+                 -300
+                ))
+                LIMIT 1
+                """,
+                "step must be positive");
+        assertQueryFails(
+                """
+                SELECT to_iso8601(timestamp), value FROM
+                TABLE(system.query_range(
+                 'count_over_time({test="timestamp_metrics_query"}[5m])',
+                 TIMESTAMP '2012-08-08',
+                 TIMESTAMP '2012-08-09',
+                 NULL
+                ))
+                LIMIT 1
+                """,
+                "step must be positive");
+    }
 }