from sqlalchemy.dialects.mssql.base import MSExecutionContext, MSDialect
from sqlalchemy.connectors.pyodbc import PyODBCConnector
from sqlalchemy import types as sqltypes
import re
import sys

class MSExecutionContext_pyodbc(MSExecutionContext):
    _embedded_scope_identity = False
    
    def pre_exec(self):
        """where appropriate, issue "select scope_identity()" in the same statement.
        
        Background on why "scope_identity()" is preferable to "@@identity":
        http://msdn.microsoft.com/en-us/library/ms190315.aspx
        
        Background on why we attempt to embed "scope_identity()" into the same
        statement as the INSERT:
        http://code.google.com/p/pyodbc/wiki/FAQs#How_do_I_retrieve_autogenerated/identity_values?
        
        """
        
        super(MSExecutionContext_pyodbc, self).pre_exec()

        # don't embed the scope_identity select into an "INSERT .. DEFAULT VALUES"
        if self._select_lastrowid and \
                self.dialect.use_scope_identity and \
                len(self.parameters[0]):
            self._embedded_scope_identity = True
            
            self.statement += "; select scope_identity()"

    def post_exec(self):
        if self._embedded_scope_identity:
            # Fetch the last inserted id from the manipulated statement
            # We may have to skip over a number of result sets with no data (due to triggers, etc.)
            while True:
                try:
                    # fetchall() ensures the cursor is consumed without closing it (FreeTDS particularly)
                    row = self.cursor.fetchall()[0]  
                    break
                except self.dialect.dbapi.Error, e:
                    # no way around this - nextset() consumes the previous set
                    # so we need to just keep flipping
                    self.cursor.nextset()
                    
            self._lastrowid = int(row[0])
        else:
            super(MSExecutionContext_pyodbc, self).post_exec()


class MSDialect_pyodbc(PyODBCConnector, MSDialect):
    supports_sane_rowcount = True
    supports_sane_multi_rowcount = False

    execution_ctx_cls = MSExecutionContext_pyodbc

    pyodbc_driver_name = 'SQL Server'

    def __init__(self, description_encoding='latin-1', **params):
        super(MSDialect_pyodbc, self).__init__(**params)
        self.description_encoding = description_encoding
        self.use_scope_identity = self.dbapi and hasattr(self.dbapi.Cursor, 'nextset')
        
    def initialize(self, connection):
        super(MSDialect_pyodbc, self).initialize(connection)
        pyodbc = self.dbapi
        
        dbapi_con = connection.connection
        
        self._free_tds = re.match(r".*libtdsodbc.*\.so",  dbapi_con.getinfo(pyodbc.SQL_DRIVER_NAME))
    
        # the "Py2K only" part here is theoretical.
        # have not tried pyodbc + python3.1 yet.
        # Py2K
        self.supports_unicode_statements = not self._free_tds
        self.supports_unicode_binds = not self._free_tds
        # end Py2K
        
dialect = MSDialect_pyodbc
