Is it possible to customize Slick code generation to have generated classes extend custom traits?

Question

I'm currently using Slick codegen (version 3.2.0-M1) to generate Slick code for a database. Many of my tables contain the same columns (with the same name and type), and so I would like to have some methods that can perform operations on these tables in a generic way, e.g., a generic method that can select rows from any of these tables based on a particular shared field.

To do this, I could create a trait that would contain these shared fields and then have the Slick table classes extend them or mix them in. Ideally, I'd like to have the code generator add extends <trait> or with <trait> to these classes for me.

I see that there's an overrideable code method in the generator, but I would like to avoid having to mess with the code directly, e.g. via regular expressions.

I haven't found anything online or in the Slick documentation that points towards an easy solution using code generator customization, so I was wondering if anyone out there knows if this is even possible.

Answer 1

I've managed to override several of the source code generator's customizable methods using code modified from slick.codegen.AbstractSourceCodeGenerator:

/* `TableUtils` contains the definitions of the `GenericRow`
      and `GenericTable` traits. */
  override def code = "import data.TableUtils._\n" + super.code

  override def Table = new Table(_) {

    override def EntityType = new EntityTypeDef {

      /* This code is adapted from the `EntityTypeDef` trait's `code` method
         within `AbstractSourceCodeGenerator`.
         All code is identical except for those lines which have a corresponding
         comment above them. */
      override def code = {
        val args = columns.map(c=>
          c.default.map( v =>
            s"${c.name}: ${c.exposedType} = $v"
          ).getOrElse(
            s"${c.name}: ${c.exposedType}"
          )
        ).mkString(", ")
        if(classEnabled){
          /* `rowList` contains the names of the generated "Row" case classes we
              wish to have extend our `GenericRow` trait. */
          val newParents = if (rowList.contains(name)) parents :+ "GenericRow" else parents
          /* Use our modified parent class sequence in place of the old one. */
          val prns = (newParents.take(1).map(" extends "+_) ++ newParents.drop(1).map(" with "+_)).mkString("")
          s"""case class $name($args)$prns"""
        } else {
          s"""type $name = $types
              /** Constructor for $name providing default values if available in the database schema. */
              def $name($args): $name = {
              ${compoundValue(columns.map(_.name))}
              }
            """.trim
        }
      }
    }

    override def TableClass = new TableClassDef {

      /* This code is adapted from the `TableClassDef` trait's `code` method
         within `AbstractSourceCodeGenerator`.
         All code is identical except for those lines which have a corresponding
         comment above them. */
      override def code = {
        /* `tableList` contains the names of the generated table classes we
            wish to have extend our `GenericTable` trait. */
        val newParents = if (tableList.contains(name)) parents :+ "GenericTable" else parents
        /* Use our modified parent class sequence in place of the old one. */
        val prns = newParents.map(" with " + _).mkString("")
        val args = model.name.schema.map(n => s"""Some("$n")""") ++ Seq("\""+model.name.table+"\"")
        s"""class $name(_tableTag: Tag) extends profile.api.Table[$elementType](_tableTag, ${args.mkString(", ")})$prns {
            ${indent(body.map(_.mkString("\n")).mkString("\n\n"))}
            }
          """.trim()
      }
    }
  }
}

This solution works for my purposes, but copying and modifying source code feels a little inelegant. If anyone knows of a nicer way to do this, I'd love to see what you've come up with.