From ce7ba67e2cd4da8396732946feecdbe39e194443 Mon Sep 17 00:00:00 2001 From: Shang Ma Date: Tue, 20 May 2025 10:39:50 -0700 Subject: [PATCH] Add support for adding custom codecs with pom file --- .../utils/DataSizeToBytesThriftCodec.java | 11 +++- .../utils/DurationToMillisThriftCodec.java | 11 +++- .../JodaDateTimeToEpochMillisThriftCodec.java | 11 +++- .../codec/utils/LocaleToLanguageTagCodec.java | 11 +++- .../facebook/drift/codec/CodecThriftType.java | 28 +++++++++ .../builtin/OptionalDoubleThriftCodec.java | 11 +++- .../builtin/OptionalIntThriftCodec.java | 11 +++- .../builtin/OptionalLongThriftCodec.java | 11 +++- .../AbstractThriftMetadataBuilder.java | 2 +- .../idl/generator/ThriftIdlGenerator.java | 63 ++++++++++++++++++- .../generator/ThriftIdlGeneratorConfig.java | 22 ++++++- .../drift/maven/IdlGeneratorMojo.java | 4 ++ 12 files changed, 184 insertions(+), 12 deletions(-) create mode 100644 drift-codec/src/main/java/com/facebook/drift/codec/CodecThriftType.java diff --git a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DataSizeToBytesThriftCodec.java b/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DataSizeToBytesThriftCodec.java index 1e546e74a..d65c9dd40 100644 --- a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DataSizeToBytesThriftCodec.java +++ b/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DataSizeToBytesThriftCodec.java @@ -15,6 +15,7 @@ */ package com.facebook.drift.codec.utils; +import com.facebook.drift.codec.CodecThriftType; import com.facebook.drift.codec.ThriftCodec; import com.facebook.drift.codec.internal.coercion.FromThrift; import com.facebook.drift.codec.internal.coercion.ToThrift; @@ -32,16 +33,24 @@ public class DataSizeToBytesThriftCodec implements ThriftCodec { + private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.DOUBLE, DataSize.class); + @Inject public DataSizeToBytesThriftCodec(ThriftCatalog thriftCatalog) { thriftCatalog.addDefaultCoercions(getClass()); } + @CodecThriftType + public static ThriftType getThriftType() + { + return THRIFT_TYPE; + } + @Override public ThriftType getType() { - return new ThriftType(ThriftType.DOUBLE, DataSize.class); + return THRIFT_TYPE; } @Override diff --git a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DurationToMillisThriftCodec.java b/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DurationToMillisThriftCodec.java index 55086cac3..af7d1d9fa 100644 --- a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DurationToMillisThriftCodec.java +++ b/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DurationToMillisThriftCodec.java @@ -15,6 +15,7 @@ */ package com.facebook.drift.codec.utils; +import com.facebook.drift.codec.CodecThriftType; import com.facebook.drift.codec.ThriftCodec; import com.facebook.drift.codec.internal.coercion.FromThrift; import com.facebook.drift.codec.internal.coercion.ToThrift; @@ -32,16 +33,24 @@ public class DurationToMillisThriftCodec implements ThriftCodec { + private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.DOUBLE, Duration.class); + @Inject public DurationToMillisThriftCodec(ThriftCatalog thriftCatalog) { thriftCatalog.addDefaultCoercions(getClass()); } + @CodecThriftType + public static ThriftType getThriftType() + { + return THRIFT_TYPE; + } + @Override public ThriftType getType() { - return new ThriftType(ThriftType.DOUBLE, Duration.class); + return THRIFT_TYPE; } @Override diff --git a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/JodaDateTimeToEpochMillisThriftCodec.java b/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/JodaDateTimeToEpochMillisThriftCodec.java index 4aa01b809..91a15c022 100644 --- a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/JodaDateTimeToEpochMillisThriftCodec.java +++ b/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/JodaDateTimeToEpochMillisThriftCodec.java @@ -15,6 +15,7 @@ */ package com.facebook.drift.codec.utils; +import com.facebook.drift.codec.CodecThriftType; import com.facebook.drift.codec.ThriftCodec; import com.facebook.drift.codec.internal.coercion.FromThrift; import com.facebook.drift.codec.internal.coercion.ToThrift; @@ -31,16 +32,24 @@ public class JodaDateTimeToEpochMillisThriftCodec implements ThriftCodec { + private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.I64, DateTime.class); + @Inject public JodaDateTimeToEpochMillisThriftCodec(ThriftCatalog thriftCatalog) { thriftCatalog.addDefaultCoercions(getClass()); } + @CodecThriftType + public static ThriftType getThriftType() + { + return THRIFT_TYPE; + } + @Override public ThriftType getType() { - return new ThriftType(ThriftType.I64, DateTime.class); + return THRIFT_TYPE; } @Override diff --git a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/LocaleToLanguageTagCodec.java b/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/LocaleToLanguageTagCodec.java index 58c33a153..3ad6bc368 100644 --- a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/LocaleToLanguageTagCodec.java +++ b/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/LocaleToLanguageTagCodec.java @@ -15,6 +15,7 @@ */ package com.facebook.drift.codec.utils; +import com.facebook.drift.codec.CodecThriftType; import com.facebook.drift.codec.ThriftCodec; import com.facebook.drift.codec.internal.coercion.FromThrift; import com.facebook.drift.codec.internal.coercion.ToThrift; @@ -32,16 +33,24 @@ public class LocaleToLanguageTagCodec implements ThriftCodec { + private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.STRING, Locale.class); + @Inject public LocaleToLanguageTagCodec(ThriftCatalog thriftCatalog) { thriftCatalog.addDefaultCoercions(getClass()); } + @CodecThriftType + public static ThriftType getThriftType() + { + return THRIFT_TYPE; + } + @Override public ThriftType getType() { - return new ThriftType(ThriftType.STRING, Locale.class); + return THRIFT_TYPE; } @Override diff --git a/drift-codec/src/main/java/com/facebook/drift/codec/CodecThriftType.java b/drift-codec/src/main/java/com/facebook/drift/codec/CodecThriftType.java new file mode 100644 index 000000000..92679a37a --- /dev/null +++ b/drift-codec/src/main/java/com/facebook/drift/codec/CodecThriftType.java @@ -0,0 +1,28 @@ +/* + * Copyright (C) 2012 Facebook, Inc. + * + * 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 com.facebook.drift.codec; + +import java.lang.annotation.Retention; +import java.lang.annotation.Target; + +import static java.lang.annotation.ElementType.METHOD; +import static java.lang.annotation.RetentionPolicy.RUNTIME; + +@Retention(RUNTIME) +@Target(METHOD) +public @interface CodecThriftType +{ +} diff --git a/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalDoubleThriftCodec.java b/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalDoubleThriftCodec.java index 764f76d89..2657c0633 100644 --- a/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalDoubleThriftCodec.java +++ b/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalDoubleThriftCodec.java @@ -15,6 +15,7 @@ */ package com.facebook.drift.codec.internal.builtin; +import com.facebook.drift.codec.CodecThriftType; import com.facebook.drift.codec.ThriftCodec; import com.facebook.drift.codec.metadata.ThriftType; import com.facebook.drift.protocol.TProtocolReader; @@ -30,10 +31,18 @@ public class OptionalDoubleThriftCodec implements ThriftCodec { + private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.DOUBLE, OptionalDouble.class, OptionalDouble.empty()); + + @CodecThriftType + public static ThriftType getThriftType() + { + return THRIFT_TYPE; + } + @Override public ThriftType getType() { - return new ThriftType(ThriftType.DOUBLE, OptionalDouble.class, OptionalDouble.empty()); + return THRIFT_TYPE; } @Override diff --git a/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalIntThriftCodec.java b/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalIntThriftCodec.java index 9b242b32b..68096a744 100644 --- a/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalIntThriftCodec.java +++ b/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalIntThriftCodec.java @@ -15,6 +15,7 @@ */ package com.facebook.drift.codec.internal.builtin; +import com.facebook.drift.codec.CodecThriftType; import com.facebook.drift.codec.ThriftCodec; import com.facebook.drift.codec.metadata.ThriftType; import com.facebook.drift.protocol.TProtocolReader; @@ -30,10 +31,18 @@ public class OptionalIntThriftCodec implements ThriftCodec { + private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.I32, OptionalInt.class, OptionalInt.empty()); + @Override public ThriftType getType() { - return new ThriftType(ThriftType.I32, OptionalInt.class, OptionalInt.empty()); + return THRIFT_TYPE; + } + + @CodecThriftType + public static ThriftType getThriftType() + { + return THRIFT_TYPE; } @Override diff --git a/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalLongThriftCodec.java b/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalLongThriftCodec.java index 118ec7b81..8e8c17f65 100644 --- a/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalLongThriftCodec.java +++ b/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalLongThriftCodec.java @@ -15,6 +15,7 @@ */ package com.facebook.drift.codec.internal.builtin; +import com.facebook.drift.codec.CodecThriftType; import com.facebook.drift.codec.ThriftCodec; import com.facebook.drift.codec.metadata.ThriftType; import com.facebook.drift.protocol.TProtocolReader; @@ -30,10 +31,18 @@ public class OptionalLongThriftCodec implements ThriftCodec { + private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.I64, OptionalLong.class, OptionalLong.empty()); + + @CodecThriftType + public static ThriftType getThriftType() + { + return THRIFT_TYPE; + } + @Override public ThriftType getType() { - return new ThriftType(ThriftType.I64, OptionalLong.class, OptionalLong.empty()); + return THRIFT_TYPE; } @Override diff --git a/drift-codec/src/main/java/com/facebook/drift/codec/metadata/AbstractThriftMetadataBuilder.java b/drift-codec/src/main/java/com/facebook/drift/codec/metadata/AbstractThriftMetadataBuilder.java index 1b8586e92..1d2ddfe81 100644 --- a/drift-codec/src/main/java/com/facebook/drift/codec/metadata/AbstractThriftMetadataBuilder.java +++ b/drift-codec/src/main/java/com/facebook/drift/codec/metadata/AbstractThriftMetadataBuilder.java @@ -704,7 +704,7 @@ protected final Requiredness extractFieldRequiredness(short fieldId, String fiel } /** - * Verifies that the the fields all have a supported Java type and that all fields map to the + * Verifies that the fields all have a supported Java type and that all fields map to the * exact same ThriftType. */ protected final void verifyFieldType(short id, String name, Collection fields, ThriftCatalog catalog) diff --git a/drift-idl-generator/src/main/java/com/facebook/drift/idl/generator/ThriftIdlGenerator.java b/drift-idl-generator/src/main/java/com/facebook/drift/idl/generator/ThriftIdlGenerator.java index e7bf83fe3..ab23c1d05 100644 --- a/drift-idl-generator/src/main/java/com/facebook/drift/idl/generator/ThriftIdlGenerator.java +++ b/drift-idl-generator/src/main/java/com/facebook/drift/idl/generator/ThriftIdlGenerator.java @@ -16,6 +16,8 @@ package com.facebook.drift.idl.generator; import com.facebook.drift.annotations.ThriftService; +import com.facebook.drift.codec.CodecThriftType; +import com.facebook.drift.codec.ThriftCodec; import com.facebook.drift.codec.ThriftCodecManager; import com.facebook.drift.codec.ThriftProtocolType; import com.facebook.drift.codec.metadata.FieldKind; @@ -36,6 +38,9 @@ import com.google.common.collect.ImmutableSet; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.lang.reflect.Modifier; import java.lang.reflect.Type; import java.net.URI; import java.util.ArrayList; @@ -129,11 +134,65 @@ public ThriftIdlGenerator(ThriftIdlGeneratorConfig config, ClassLoader classLoad this.namespaces = config.getNamespaces(); this.recursive = config.isRecursive(); + + for (String customCodecClassName : config.getCustomCodecs()) { + loadCustomCodec(customCodecClassName); + } + } + + private void loadCustomCodec(String customCodecClassName) + { + Class codecClass = load(customCodecClassName); + if (codecClass == null) { + throw new ThriftIdlGeneratorException("Failed to load codec class: " + customCodecClassName); + } + if (!ThriftCodec.class.isAssignableFrom(codecClass)) { + throw new ThriftIdlGeneratorException("Class " + customCodecClassName + " does not implement ThriftCodec"); + } + + // Look for a static method annotated with @CodecThriftType + ThriftType thriftType = null; + for (Method method : codecClass.getDeclaredMethods()) { + if (method.isAnnotationPresent(CodecThriftType.class)) { + if (!Modifier.isPublic(method.getModifiers())) { + throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must be public: " + customCodecClassName + "#" + method.getName()); + } + + if (!Modifier.isStatic(method.getModifiers())) { + throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must be static: " + customCodecClassName + "#" + method.getName()); + } + + if (method.getParameterCount() != 0) { + throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must have no parameters: " + customCodecClassName + "#" + method.getName()); + } + + if (!ThriftType.class.isAssignableFrom(method.getReturnType())) { + throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must return ThriftType: " + customCodecClassName + "#" + method.getName()); + } + + try { + thriftType = (ThriftType) method.invoke(null); + if (thriftType == null) { + throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType returns null: " + customCodecClassName + "#" + method.getName()); + } + break; + } + catch (IllegalAccessException | InvocationTargetException e) { + throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType can not be invoked: " + customCodecClassName + "#" + method.getName()); + } + } + } + if (thriftType == null) { + throw new ThriftIdlGeneratorException("Codec must have a public static method annotated with @CodecThriftType and return ThriftType: " + customCodecClassName); + } + + addCustomType(thriftType); + catalog.addThriftType(thriftType); } - public void addCustomType(ThriftType... type) + public void addCustomType(ThriftType... types) { - for (ThriftType t : type) { + for (ThriftType t : types) { customTypes.add(t); knownTypes.add(t); } diff --git a/drift-idl-generator/src/main/java/com/facebook/drift/idl/generator/ThriftIdlGeneratorConfig.java b/drift-idl-generator/src/main/java/com/facebook/drift/idl/generator/ThriftIdlGeneratorConfig.java index 5d032b3dd..1367e2925 100644 --- a/drift-idl-generator/src/main/java/com/facebook/drift/idl/generator/ThriftIdlGeneratorConfig.java +++ b/drift-idl-generator/src/main/java/com/facebook/drift/idl/generator/ThriftIdlGeneratorConfig.java @@ -15,8 +15,10 @@ */ package com.facebook.drift.idl.generator; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import java.util.List; import java.util.Map; import java.util.function.Consumer; @@ -31,6 +33,7 @@ public class ThriftIdlGeneratorConfig private final Consumer errorLogger; private final Consumer warningLogger; private final Consumer verboseLogger; + private final List customCodecs; private ThriftIdlGeneratorConfig( String defaultPackage, @@ -39,7 +42,8 @@ private ThriftIdlGeneratorConfig( boolean recursive, Consumer errorLogger, Consumer warningLogger, - Consumer verboseLogger) + Consumer verboseLogger, + List customCodecs) { this.defaultPackage = firstNonNull(defaultPackage, ""); this.namespaces = ImmutableMap.copyOf(firstNonNull(namespaces, ImmutableMap.of())); @@ -48,6 +52,7 @@ private ThriftIdlGeneratorConfig( this.errorLogger = firstNonNull(errorLogger, ignored -> {}); this.warningLogger = firstNonNull(warningLogger, ignored -> {}); this.verboseLogger = firstNonNull(verboseLogger, ignored -> {}); + this.customCodecs = ImmutableList.copyOf(firstNonNull(customCodecs, ImmutableList.of())); } public static Builder builder() @@ -90,6 +95,11 @@ public Consumer getVerboseLogger() return verboseLogger; } + public List getCustomCodecs() + { + return customCodecs; + } + public static class Builder { private String defaultPackage; @@ -99,6 +109,7 @@ public static class Builder private Consumer errorLogger; private Consumer warningLogger; private Consumer verboseLogger; + private List customCodecs; private Builder() {} @@ -144,6 +155,12 @@ public Builder verboseLogger(Consumer verboseLogger) return this; } + public Builder customCodecs(List customCodecs) + { + this.customCodecs = ImmutableList.copyOf(firstNonNull(customCodecs, ImmutableList.of())); + return this; + } + public ThriftIdlGeneratorConfig build() { return new ThriftIdlGeneratorConfig( @@ -153,7 +170,8 @@ public ThriftIdlGeneratorConfig build() recursive, errorLogger, warningLogger, - verboseLogger); + verboseLogger, + customCodecs); } } } diff --git a/drift-maven-plugin/src/main/java/com/facebook/drift/maven/IdlGeneratorMojo.java b/drift-maven-plugin/src/main/java/com/facebook/drift/maven/IdlGeneratorMojo.java index db3957bf7..bb951161f 100644 --- a/drift-maven-plugin/src/main/java/com/facebook/drift/maven/IdlGeneratorMojo.java +++ b/drift-maven-plugin/src/main/java/com/facebook/drift/maven/IdlGeneratorMojo.java @@ -99,6 +99,9 @@ public class IdlGeneratorMojo @Parameter private boolean quiet; + @Parameter + private List customCodecs; + @Override public void execute() throws MojoExecutionException @@ -113,6 +116,7 @@ public void execute() .errorLogger(message -> getLog().error(message)) .warningLogger(message -> getLog().warn(message)) .verboseLogger(this::verbose) + .customCodecs(customCodecs) .build(); String idl;