def extract_class(query):
  """Extract original class object from a SQLAlchemy query.

  Args:
    query (query): SQLAlchemy query

  Returns:
    class: base class use when setting up the SQL query
  """
  first_expression = query.column_descriptions[0]['expr']

  try:
    # query returns subset of columns as tuples
    return first_expression.class_

  except AttributeError:
    # query returns a full class ORM object
    return first_expression