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

If you want add trait to every generated table you need override parents method in TableClass and EntityType traits of code generator:

class ExtSourceCodeGenerator(model: m.Model) extends SourceCodeGenerator(model: m.Model) {

  //Import packages with your traits
  override def code: String = "import models._\n" + super.code

  override def Table = new Table(_) {

    override def TableClass = new TableClass {
      //Add custom traits to Table classes
      override def parents: Seq[String] = {
        Seq("IdentityTable")
      }
    }

    override def EntityType = new EntityType {
      //Add custom traits to rows
      override def parents: Seq[String] = {
        Seq("Entity")
      }
    }
  }
}

If you need to filter some tables you can use model field to get current table name:

model.name.table

Hope this will help you.

Answer 2

Using the solutions outlined in the How to mix-in a trait to instance, answers this question (and my question too which is almost exactly as yours). We want to extend the Slick model generated code without modifying it (or for that matter the generator)

Adapted here for self-containment, convenience and to the specific use-case I had. This is my Slick code-gen generated code:

/** Table description of table user. Objects of this class serve as prototypes for rows in queries. */
class User(_tableTag: Tag) extends Table[UserRow](_tableTag, "user") { ... }

but I need my User Slick data model to implement the be.objectify.deadbolt.java.models.Subject Java interface so I can use Deadbolt2 as part of my Scala Play Web Application.

Therefore I'd do:

import be.objectify.deadbolt.java.models.Subject     

/**
 * Mixin framework or infrastructural code
 */
trait DynamicMixinCompanion[TT] {
  implicit def baseObject[OT](o: Mixin[OT]): OT = o.obj

  def ::[OT](o: OT): Mixin[OT] with TT
  class Mixin[OT] protected[DynamicMixinCompanion](val obj: OT)
}

/**
 * Subject Mixin implementation
 */
object SubjectMixinHelper extends DynamicMixinCompanion[Subject] {
  def ::[T](o: T) = new Mixin(o) with Subject {
     def getPermissions = ...  
     def getRoles = ...
  }
}

and finally:

import SubjectMixinHelper._
withSession{ implicit session =>
  val user = Query(User).where(_.id === id).firstOption :: Subject

  // then use user as a Subject too
  user.getPermissions
  user.getRoles
}

Note that I haven't tested it yet (but I will soon)

Answer 3

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.