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

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

class ProjectIdMapperExtension(MapperExtension):
    def before_insert(self, mapper, connection, instance):
        # If we don't have an ID already,
        # grab the next ID in this project
        if not instance.id:
            # Grab the underlying SQLAlchemy table
            ticket_table = instance._descriptor.entity.table

            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:
                # No entry -- query for the next ID
                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):
        # After this insert is over, trash the counter (it may well expire)
        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()


"""
# ./test.py
2008-02-05 13:18:41,121 INFO sqlalchemy.engine.base.Engine.0x..b4 PRAGMA table_info("__main___project")
2008-02-05 13:18:41,122 INFO sqlalchemy.engine.base.Engine.0x..b4 {}
2008-02-05 13:18:41,123 INFO sqlalchemy.engine.base.Engine.0x..b4 PRAGMA table_info("__main___milestone")
2008-02-05 13:18:41,123 INFO sqlalchemy.engine.base.Engine.0x..b4 {}
2008-02-05 13:18:41,124 INFO sqlalchemy.engine.base.Engine.0x..b4
CREATE TABLE __main___project (
        name TEXT NOT NULL,
        PRIMARY KEY (name)
)


2008-02-05 13:18:41,124 INFO sqlalchemy.engine.base.Engine.0x..b4 None
2008-02-05 13:18:41,125 INFO sqlalchemy.engine.base.Engine.0x..b4 COMMIT
2008-02-05 13:18:41,126 INFO sqlalchemy.engine.base.Engine.0x..b4
CREATE TABLE __main___milestone (
        id INTEGER NOT NULL,
        name TEXT,
        project_name TEXT NOT NULL,
        PRIMARY KEY (id, project_name),
         CONSTRAINT __main___milestone_project_name_fk FOREIGN KEY(project_name) REFERENCES __main___project (name)
)


2008-02-05 13:18:41,126 INFO sqlalchemy.engine.base.Engine.0x..b4 None
2008-02-05 13:18:41,127 INFO sqlalchemy.engine.base.Engine.0x..b4 COMMIT
2008-02-05 13:18:41,127 INFO sqlalchemy.engine.base.Engine.0x..b4 CREATE INDEX ix___main___milestone_project_name ON __main___milestone (project_name)
2008-02-05 13:18:41,128 INFO sqlalchemy.engine.base.Engine.0x..b4 None
2008-02-05 13:18:41,128 INFO sqlalchemy.engine.base.Engine.0x..b4 COMMIT
2008-02-05 13:18:41,134 INFO sqlalchemy.engine.base.Engine.0x..b4 BEGIN
2008-02-05 13:18:41,136 INFO sqlalchemy.engine.base.Engine.0x..b4 INSERT INTO __main___project (name) VALUES (?)
2008-02-05 13:18:41,136 INFO sqlalchemy.engine.base.Engine.0x..b4 ['fizzle']
2008-02-05 13:18:41,140 INFO sqlalchemy.engine.base.Engine.0x..b4 SELECT max(__main___milestone.id)
FROM __main___milestone
WHERE __main___milestone.project_name = ?
2008-02-05 13:18:41,140 INFO sqlalchemy.engine.base.Engine.0x..b4 ['fizzle']
2008-02-05 13:18:41,144 INFO sqlalchemy.engine.base.Engine.0x..b4 INSERT INTO __main___milestone (id, name, project_name) VALUES (?, ?, ?)
2008-02-05 13:18:41,144 INFO sqlalchemy.engine.base.Engine.0x..b4 [2, 'stuff0', 'fizzle']
2008-02-05 13:18:41,145 INFO sqlalchemy.engine.base.Engine.0x..b4 INSERT INTO __main___milestone (id, name, project_name) VALUES (?, ?, ?)
2008-02-05 13:18:41,146 INFO sqlalchemy.engine.base.Engine.0x..b4 [3, 'stuff1', 'fizzle']
2008-02-05 13:18:41,148 INFO sqlalchemy.engine.base.Engine.0x..b4 INSERT INTO __main___milestone (id, name, project_name) VALUES (?, ?, ?)
2008-02-05 13:18:41,148 INFO sqlalchemy.engine.base.Engine.0x..b4 [1, 'stuff2', 'fizzle']
2008-02-05 13:18:41,150 INFO sqlalchemy.engine.base.Engine.0x..b4 COMMIT
"""