Add support for adding custom codecs with pom file

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<DataSize>
 {
+    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

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<Duration>
 {
+    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

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<DateTime>
 {
+    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

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<Locale>
 {
+    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

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
+{
+}

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<OptionalDouble>
 {
+    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

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<OptionalInt>
 {
+    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

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<OptionalLong>
 {
+    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

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<FieldMetadata> fields, ThriftCatalog catalog)

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);
         }