oids rows insert sort orders galore

1 files changed, 10 insertions, 13 deletions

diff --git a/lib/sqlalchemy/databases/postgres.py b/lib/sqlalchemy/databases/postgres.py
index cfaf63b57..74c2b5366 100644
--- a/lib/sqlalchemy/databases/postgres.py
+++ b/lib/sqlalchemy/databases/postgres.py
@@ -92,8 +92,13 @@ class PGSQLEngine(ansisql.ANSISQLEngine):
         raise "not implemented"
         
     def last_inserted_ids(self):
-        return self.context.last_inserted_ids
-
+        table = self.context.last_inserted_table
+        if self.context.lastrowid is not None and table is not None and len(table.primary_keys):
+            row = sql.select(table.primary_keys, table.rowid_column == self.context.lastrowid).execute().fetchone()
+            return [v for v in row]
+        else:
+            return None
+            
     def pre_exec(self, connection, cursor, statement, parameters, echo = None, compiled = None, **kwargs):
         if True: return
         # if a sequence was explicitly defined we do it here
@@ -123,18 +128,10 @@ class PGSQLEngine(ansisql.ANSISQLEngine):
     def post_exec(self, connection, cursor, statement, parameters, echo = None, compiled = None, **kwargs):
         if compiled is None: return
         if getattr(compiled, "isinsert", False):
-            # psycopg wants to return internal rowids, which I guess is what DBAPI2 really 
-            # specifies.
-            # well then post exec to get the row.  I guess this could be genericised to 
-            # be for all inserts somehow if the "rowid" col could be gotten off a table.
             table = compiled.statement.table
-            if len(table.primary_keys):
-                # TODO: cache this statement against the table to avoid multiple re-compiles
-                # TODO: instead of "oid" have the Table object have a "rowid_col" property
-                # that gives this col generically
-                row = sql.select(table.primary_keys, sql.ColumnClause("oid",table) == bindparam('oid', cursor.lastrowid) ).execute().fetchone()
-                self.context.last_inserted_ids = [v for v in row]
-
+            self.context.last_inserted_table = table
+            self.context.lastrowid = cursor.lastrowid
+            
     def dbapi(self):
         return self.module