Skip to content

Commit 3a99f94

Browse files
lalitiumlalitium
authored
feat: Add LogElements transform to Java SDK (#38533)
* feat: Add Java LogElements transform * Addressed LogElements review feedback * Add LogElements tests * Add LogElements level logging test * Refactor LogElements logging test
1 parent db1b852 commit 3a99f94

2 files changed

Lines changed: 342 additions & 0 deletions

File tree

‎sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/LogElements.java‎

Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.sdk.transforms;
19+
20+
import java.util.Objects;
21+
import org.apache.beam.sdk.transforms.display.DisplayData;
22+
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
23+
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
24+
import org.apache.beam.sdk.values.PCollection;
25+
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
26+
import org.checkerframework.checker.nullness.qual.Nullable;
27+
import org.joda.time.Instant;
28+
import org.slf4j.Logger;
29+
import org.slf4j.LoggerFactory;
30+
import org.slf4j.event.Level;
31+
32+
/**
33+
* {@link PTransform} for logging elements of a {@link PCollection}.
34+
*
35+
* <p>Each element is logged and then emitted unchanged.
36+
*
37+
* <pre>{@code
38+
* PCollection<String> words = ...;
39+
* PCollection<String> loggedWords = words.apply(LogElements.info().withPrefix("word: "));
40+
* }</pre>
41+
*
42+
* @param <T> the element type of the input {@link PCollection}
43+
*/
44+
public class LogElements<T> extends PTransform<PCollection<T>, PCollection<T>> {
45+
private static final Logger LOG = LoggerFactory.getLogger(LogElements.class);
46+
47+
private final Level level;
48+
private final String prefix;
49+
private final boolean withTimestamp;
50+
private final boolean withWindow;
51+
private final boolean withPaneInfo;
52+
53+
/** Returns a {@link LogElements} transform that logs elements at {@link Level#TRACE}. */
54+
public static <T> LogElements<T> trace() {
55+
return of(Level.TRACE);
56+
}
57+
58+
/** Returns a {@link LogElements} transform that logs elements at {@link Level#DEBUG}. */
59+
public static <T> LogElements<T> debug() {
60+
return of(Level.DEBUG);
61+
}
62+
63+
/** Returns a {@link LogElements} transform that logs elements at {@link Level#INFO}. */
64+
public static <T> LogElements<T> info() {
65+
return of(Level.INFO);
66+
}
67+
68+
/** Returns a {@link LogElements} transform that logs elements at {@link Level#WARN}. */
69+
public static <T> LogElements<T> warn() {
70+
return of(Level.WARN);
71+
}
72+
73+
/** Returns a {@link LogElements} transform that logs elements at {@link Level#ERROR}. */
74+
public static <T> LogElements<T> error() {
75+
return of(Level.ERROR);
76+
}
77+
78+
/** Returns a {@link LogElements} transform that logs elements at the given level. */
79+
public static <T> LogElements<T> of(Level level) {
80+
return new LogElements<>(level, "", false, false, false);
81+
}
82+
83+
private LogElements(
84+
Level level, String prefix, boolean withTimestamp, boolean withWindow, boolean withPaneInfo) {
85+
this.level = Objects.requireNonNull(level, "level");
86+
this.prefix = Objects.requireNonNull(prefix, "prefix");
87+
this.withTimestamp = withTimestamp;
88+
this.withWindow = withWindow;
89+
this.withPaneInfo = withPaneInfo;
90+
}
91+
92+
/** Returns a new {@link LogElements} transform with the given prefix before each element. */
93+
public LogElements<T> withPrefix(String prefix) {
94+
return new LogElements<>(level, prefix, withTimestamp, withWindow, withPaneInfo);
95+
}
96+
97+
/** Returns a new {@link LogElements} transform that logs each element's timestamp. */
98+
public LogElements<T> withTimestamp() {
99+
return new LogElements<>(level, prefix, true, withWindow, withPaneInfo);
100+
}
101+
102+
/** Returns a new {@link LogElements} transform that logs each element's window. */
103+
public LogElements<T> withWindow() {
104+
return new LogElements<>(level, prefix, withTimestamp, true, withPaneInfo);
105+
}
106+
107+
/** Returns a new {@link LogElements} transform that logs each element's pane info. */
108+
public LogElements<T> withPaneInfo() {
109+
return new LogElements<>(level, prefix, withTimestamp, withWindow, true);
110+
}
111+
112+
@Override
113+
public PCollection<T> expand(PCollection<T> input) {
114+
return input.apply("Log", ParDo.of(new LoggingFn<>(this)));
115+
}
116+
117+
@Override
118+
public void populateDisplayData(DisplayData.Builder builder) {
119+
super.populateDisplayData(builder);
120+
builder
121+
.add(DisplayData.item("level", level.name()).withLabel("Log Level"))
122+
.addIfNotDefault(DisplayData.item("prefix", prefix).withLabel("Prefix"), "")
123+
.addIfNotDefault(
124+
DisplayData.item("withTimestamp", withTimestamp).withLabel("Log Timestamp"), false)
125+
.addIfNotDefault(DisplayData.item("withWindow", withWindow).withLabel("Log Window"), false)
126+
.addIfNotDefault(
127+
DisplayData.item("withPaneInfo", withPaneInfo).withLabel("Log Pane Info"), false);
128+
}
129+
130+
static String formatForLogging(
131+
@Nullable Object element,
132+
String prefix,
133+
boolean withTimestamp,
134+
boolean withWindow,
135+
boolean withPaneInfo,
136+
Instant timestamp,
137+
BoundedWindow window,
138+
PaneInfo paneInfo) {
139+
StringBuilder builder = new StringBuilder(prefix).append(element);
140+
if (withTimestamp) {
141+
builder.append(", timestamp=").append(timestamp);
142+
}
143+
if (withWindow) {
144+
builder.append(", window=").append(window);
145+
}
146+
if (withPaneInfo) {
147+
builder.append(", paneInfo=").append(paneInfo);
148+
}
149+
return builder.toString();
150+
}
151+
152+
@VisibleForTesting
153+
static void log(Level level, String message) {
154+
switch (level) {
155+
case TRACE:
156+
LOG.trace("{}", message);
157+
break;
158+
case DEBUG:
159+
LOG.debug("{}", message);
160+
break;
161+
case INFO:
162+
LOG.info("{}", message);
163+
break;
164+
case WARN:
165+
LOG.warn("{}", message);
166+
break;
167+
case ERROR:
168+
LOG.error("{}", message);
169+
break;
170+
default:
171+
throw unsupportedLogLevel(level);
172+
}
173+
}
174+
175+
private static boolean isLoggingEnabled(Level level) {
176+
switch (level) {
177+
case TRACE:
178+
return LOG.isTraceEnabled();
179+
case DEBUG:
180+
return LOG.isDebugEnabled();
181+
case INFO:
182+
return LOG.isInfoEnabled();
183+
case WARN:
184+
return LOG.isWarnEnabled();
185+
case ERROR:
186+
return LOG.isErrorEnabled();
187+
default:
188+
throw unsupportedLogLevel(level);
189+
}
190+
}
191+
192+
private static IllegalArgumentException unsupportedLogLevel(Level level) {
193+
return new IllegalArgumentException("Unsupported log level: " + level);
194+
}
195+
196+
private static class LoggingFn<T> extends DoFn<T, T> {
197+
private final Level level;
198+
private final String prefix;
199+
private final boolean withTimestamp;
200+
private final boolean withWindow;
201+
private final boolean withPaneInfo;
202+
203+
private LoggingFn(LogElements<T> transform) {
204+
this.level = transform.level;
205+
this.prefix = transform.prefix;
206+
this.withTimestamp = transform.withTimestamp;
207+
this.withWindow = transform.withWindow;
208+
this.withPaneInfo = transform.withPaneInfo;
209+
}
210+
211+
@ProcessElement
212+
public void processElement(
213+
@Element T element,
214+
@DoFn.Timestamp Instant timestamp,
215+
BoundedWindow window,
216+
PaneInfo paneInfo,
217+
OutputReceiver<T> receiver) {
218+
if (isLoggingEnabled(level)) {
219+
log(
220+
level,
221+
formatForLogging(
222+
element,
223+
prefix,
224+
withTimestamp,
225+
withWindow,
226+
withPaneInfo,
227+
timestamp,
228+
window,
229+
paneInfo));
230+
}
231+
receiver.output(element);
232+
}
233+
}
234+
}

‎sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/LogElementsTest.java‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.sdk.transforms;
19+
20+
import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
21+
import static org.hamcrest.MatcherAssert.assertThat;
22+
import static org.hamcrest.Matchers.containsString;
23+
24+
import java.util.Arrays;
25+
import java.util.List;
26+
import org.apache.beam.sdk.testing.ExpectedLogs;
27+
import org.apache.beam.sdk.testing.NeedsRunner;
28+
import org.apache.beam.sdk.testing.PAssert;
29+
import org.apache.beam.sdk.testing.TestPipeline;
30+
import org.apache.beam.sdk.transforms.display.DisplayData;
31+
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
32+
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
33+
import org.apache.beam.sdk.values.PCollection;
34+
import org.joda.time.Duration;
35+
import org.joda.time.Instant;
36+
import org.junit.Rule;
37+
import org.junit.Test;
38+
import org.junit.experimental.categories.Category;
39+
import org.junit.runner.RunWith;
40+
import org.junit.runners.JUnit4;
41+
import org.slf4j.event.Level;
42+
43+
/** Tests for {@link LogElements}. */
44+
@RunWith(JUnit4.class)
45+
public class LogElementsTest {
46+
@Rule public final transient TestPipeline pipeline = TestPipeline.create();
47+
@Rule public ExpectedLogs expectedLogs = ExpectedLogs.none(LogElements.class);
48+
49+
@Test
50+
@Category(NeedsRunner.class)
51+
public void testLogElementsPreservesElements() {
52+
List<String> elements = Arrays.asList("a", "b", "c");
53+
54+
PCollection<String> output =
55+
pipeline.apply(Create.of(elements)).apply(LogElements.<String>info());
56+
57+
PAssert.that(output).containsInAnyOrder(elements);
58+
pipeline.run();
59+
}
60+
61+
@Test
62+
public void testFormatForLoggingIncludesConfiguredMetadata() {
63+
Instant timestamp = new Instant(0);
64+
IntervalWindow window = new IntervalWindow(timestamp, Duration.standardMinutes(1));
65+
66+
String message =
67+
LogElements.formatForLogging(
68+
"a", "row: ", true, true, true, timestamp, window, PaneInfo.NO_FIRING);
69+
70+
assertThat(message, containsString("row: a"));
71+
assertThat(message, containsString("timestamp=1970-01-01T00:00:00.000Z"));
72+
assertThat(
73+
message, containsString("window=[1970-01-01T00:00:00.000Z..1970-01-01T00:01:00.000Z)"));
74+
assertThat(message, containsString("paneInfo=PaneInfo.NO_FIRING"));
75+
}
76+
77+
@Test
78+
public void testLogElementsLogsAtConfiguredLevels() {
79+
LogElements.log(Level.TRACE, "trace: trace-element");
80+
LogElements.log(Level.DEBUG, "debug: debug-element");
81+
LogElements.log(Level.INFO, "info: info-element");
82+
LogElements.log(Level.WARN, "warn: warn-element");
83+
LogElements.log(Level.ERROR, "error: error-element");
84+
85+
expectedLogs.verifyTrace("trace: trace-element");
86+
expectedLogs.verifyDebug("debug: debug-element");
87+
expectedLogs.verifyInfo("info: info-element");
88+
expectedLogs.verifyWarn("warn: warn-element");
89+
expectedLogs.verifyError("error: error-element");
90+
}
91+
92+
@Test
93+
public void testDisplayData() {
94+
DisplayData displayData =
95+
DisplayData.from(
96+
LogElements.of(Level.WARN)
97+
.withPrefix("row: ")
98+
.withTimestamp()
99+
.withWindow()
100+
.withPaneInfo());
101+
102+
assertThat(displayData, hasDisplayItem("level", "WARN"));
103+
assertThat(displayData, hasDisplayItem("prefix", "row: "));
104+
assertThat(displayData, hasDisplayItem("withTimestamp", true));
105+
assertThat(displayData, hasDisplayItem("withWindow", true));
106+
assertThat(displayData, hasDisplayItem("withPaneInfo", true));
107+
}
108+
}

0 commit comments

Comments
 (0)