1 files changed, 79 insertions, 0 deletions
diff --git a/lib/sqlalchemy/dialects/mssql/pyodbc.py b/lib/sqlalchemy/dialects/mssql/pyodbc.py
new file mode 100644
index 000000000..9a2a9e4e7
--- /dev/null
+++ b/lib/sqlalchemy/dialects/mssql/pyodbc.py
@@ -0,0 +1,79 @@
+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