#!/usr/bin/env python
from elixir import *
from sqlalchemy import *
from datetime import datetime

class Project(Entity):
    has_field('name', Unicode, primary_key=True)

import threading
class ProjectIdMapperExtension(MapperExtension):
    def before_insert(self, mapper, connection, instance):
        # We just want to set the id to the next logical id
        # for a given project.
        if not instance.id:
            # Grab the underlying SQLAlchemy table
            ticket_table = instance._descriptor.entity.table

            # Build a custom query
            if not hasattr(self, '_thread_local'):
                self._thread_local = threading.local()
            if not hasattr(self._thread_local, 'projects'):
                self._thread_local.projects = {}

            project_contexts = self._thread_local.projects

            # Check if we already have an entry for this project
            if project_contexts.has_key(instance.project_name):
                project_contexts[instance.project_name] += 1
            else:
                qry = select([func.max(instance.c.id)], instance.c.project_name == instance.project_name)
                rproxy = qry.execute()
                last_id  = rproxy.fetchone()[0]
                if not last_id: last_id = 0
                project_contexts[instance.project_name] = last_id + 1

            instance.id = project_contexts[instance.project_name]

    def after_insert(self, mapper, connection, instance):
        if hasattr(self, '_thread_local'):
            if hasattr(self._thread_local, 'projects'):
                del self._thread_local.projects

from elixir.fields import Statement, HasField
from elixir.relationships import BelongsTo
from elixir.options import UsingMapperOptions
class ProjectId(object):
    def __init__(self, entity):
        BelongsTo(entity, 'project', of_kind='Project', column_kwargs={"primary_key": True})
        HasField(entity, 'id', Integer, default=0, primary_key=True, autoincrement=False)
        UsingMapperOptions(entity, extension=ProjectIdMapperExtension())

has_project_and_id = Statement(ProjectId)

class Milestone(Entity):
    has_project_and_id()
    has_field('name', Unicode)

metadata.connect("sqlite://in_memory")
metadata.engine.echo = True
create_all()

y = Project(name='fizzle')
x = Milestone(project=y, name='stuff0')
n = Milestone(project=y, name='stuff1')
n = Milestone(project=y, name='stuff2')
objectstore.flush()