.NET: Creating data class with dynamic table attribute for DynamoDB

Context

For the last two weeks, I have been working on integrating DynamoDB in a .NET Core Lambda project. AWS SDK for .NET has the following APIs to interact with DynamoDB

LevelNameDescription
LowProtocol-level interfaceHTTP(S) based Request and Response
MidDocument ModelDocument model classes that wrap some of the low-level operations
HighObject Persistance ModelMap client-side classes to DynamoDB tables

The Object Persistance Model can be used with DynamoDBContext class, which acts as an entry point DynamoDB. To map the client-side classes, DynamoDBContext uses Attributes

[DynamoDBTable("Reply")]
public class Reply {
   [DynamoDBHashKey]
   public int ThreadId { get; set; }
   [DynamoDBRangeKey]
   public string Replenishment { get; set; }
   [DynamoDBProperty("Message")]
   public string Message { get; set; }
   // Additional properties go here.
}

DynamoDBContext will use the attribute value of DynamoDBTableAttribute to determine the table name. Additionally the constructor of DynamoDBContext accepts DynamoDBContextConfig class where we can set the table prefix with TableNamePrefix property. Then the table name will be determined by adding the prefix, ex. Prod-Reply, Stage-Reply, Dev-Reply, etc.

The Purpose

In general, writing a single class for each table makes perfect sense. Aside from that, we may want to use projection on multiple tables and if that projected data has the same schema, then we could create the same class for each table (with inheritance chaining or copying the whole class) to apply the table attribute. This would work when all possible table names can be defined at the compile-time, but it won’t work if the tables are generated dynamically and can’t be defined at the compile-time - ex. an external service provides the table names at runtime.

Possible Solution

Once types are compiled, they can not be changed. That’s why when we add an attribute to a type, it can not be changed.

One possible solution I can think of right now is to create types with Reflection API at runtime and add the DynamoDBTableAttribute. Why? because DynamoDBContext uses Generic Types and the metadata from the type to map table and properties. Once compiled to IL (Intermediate Language), a specialized generic type will be created for a generic parameter (method/class) and will be used for all reference types passed to that generic parameter. Read this C# documentation to know more about generics in the runtime.

So, if we can create another class at runtime with the target model as the base, we can create a generic method with that new class and invoke it. The returned object from that invoked method can be cast directly to the base class because of the inheritance. I’m getting a “Reflection-on-steroids” vibe from this idea (to be honest 😂).

The Code

As this will be highly experimental, I think it would be better to abstract away the wrapper implementation, so that we can revert to the original DynamoDBContext. I’ll start by creating a class that implements IDynamoDBContext and also takes the table name.

public class CustomDynamoDBContext : IDynamoDBContext, IDisposable
{
  private readonly IDynamoDBContext _context;

  public CustomDynamoDBContext(IAmazonDynamoDB client, string tableName) : this(client, tableName, null)
  { }

  public CustomDynamoDBContext(IAmazonDynamoDB client, string tableName, DynamoDBContextConfig config)
  {
    _tableName = tableName;
    _context = new DynamoDBContext(client, config);
  }
  // ...
  // Implement methods that will call the same method on _context
  // and return the result
}

To create our dynamic types, we need AssemblyBuilder and ModuleBuilder

public class CustomDynamoDBContext : IDynamoDBContext, IDisposable
{
  // ...
  // Other codes
  private static readonly AssemblyBuilder assemblyBuilder;
  private static readonly ModuleBuilder moduleBuilder;
  private static readonly Dictionary<string, Type> typeMap;

  static CustomDynamoDBContext()
  {
    assemblyBuilder = AssemblyBuilder.DefineDynamicAssembly(new AssemblyName
    {
      Name = "DynamoDbContainerAssembly"
    }, AssemblyBuilderAccess.Run);
    moduleBuilder = assemblyBuilder.DefineDynamicModule("DynamoDbContainerModule");
    typeMap = new Dictionary<string, Type>();
  }
  // ...
  // Other codes
}

The purpose of typeMap is to avoid creating the same type twice. So, types will be generated only once per table.

Now we can create a method that will create the type and attach the DynamoDBTableAttribute with our dynamic table name from the constructor.

public class CustomDynamoDBContext : IDynamoDBContext, IDisposable
{
  // ...
  // Other codes
  private Type AttachAttribute<T>()
  {
    Type targetType = typeof(T);

    string typeCacheKey = $"{targetType.FullName}_{_tableName}";
    if (typeMap.TryGetValue(typeCacheKey, out var existingType))
    {
        return existingType;
    }
    TypeBuilder typeBuilder = moduleBuilder.DefineType(targetType.Name, TypeAttributes.Public | TypeAttributes.Class, targetType);
    Type[] ctorParams = new Type[] { typeof(string) };
    ConstructorInfo classCtorInfo = typeof(DynamoDBTableAttribute).GetConstructor(ctorParams);
    CustomAttributeBuilder myCABuilder = new CustomAttributeBuilder(
                classCtorInfo,
                new object[] { _tableName });
    typeBuilder.SetCustomAttribute(myCABuilder);
    Type typeWithAttribute = typeBuilder.CreateType();
    typeMap.Add(typeCacheKey, typeWithAttribute);
    return typeWithAttribute;
  }
  // ...
  // Other codes
}

Now, we can intercept method calls and provide our newly created type. As the caller-provided type will be the base of the generated type, there won’t be any problem casting the base type to our generated type.

public class CustomDynamoDBContext : IDynamoDBContext, IDisposable
{
  // ...
  // Other codes
  public async Task<T> LoadAsync<T>(object hashKey, object rangeKey, CancellationToken cancellationToken = default)
  {
    Type TT = AttachAttribute<T>();
    MethodInfo methodLoadAsync = typeof(DynamoDBContext)
        .GetMethod("LoadAsync", 1, new Type[] { typeof(object), typeof(object), typeof(CancellationToken) })
        .MakeGenericMethod(TT);

    var returnedValue = methodLoadAsync.Invoke(_context, new object[] { hashKey, rangeKey, cancellationToken });
    return await (dynamic)returnedValue;
  }
  // ...
  // Other codes
}

Now we can replace public class MyRepository DynamoDBContext with CustomDynamoDBContext and we will be able to define the table name in the CustomDynamoDBContext constructor rather than using the DynamoDBTableAttribute to bind with the data class.

{
  // ...
  // Other codes

  private readonly IDynamoDBContext _dbContext;

  public MyRepository(AmazonDynamoDBClient amazonDynamoDBClient)
  {
-   _dbContext = new DynamoDBContext(amazonDynamoDBClient);
+   _dbContext = new CustomDynamoDBContext(amazonDynamoDBClient, tableName);
  }

  // ...
  // Other codes
}

Closing up

Even though it was kind of an experiment, it did work with the project that I was working on. One reason why I created a separate CustomDynamoDBContext class implementing the IDynamoDBContext interface was to separate this work-around so that I can incrementally add fixes for cases that may arise as the project matures. And if this doesn’t work out for some reason, I can swap with DynamoDBContext any time, or remove the usage Object Persistence Model altogether.

References