wizzat · October 23, 2013 04:29 · Oct 23, 2013
diff --git a/psycopg2_performance.py b/psycopg2_performance.py
@@ -0,0 +1,135 @@
+try:
+    from psycopg2cffi import compat
+    compat.register()
+except ImportError:
+    pass
+
+from pyutil.decorators import *
+import tempfile, psycopg2, psycopg2.extras
+
+def setup_test_table(conn):
+    with conn.cursor() as cur:
+        cur.execute("""
+            create table if not exists foobar (
+                some_long_column1 integer,
+                some_long_column2 integer,
+                some_long_column3 integer,
+                some_long_column4 integer,
+                some_long_column5 integer,
+                some_long_column6 integer,
+                some_long_column7 integer,
+                some_long_column8 integer
+            )
+        """)
+
+    with conn.cursor() as cur:
+        cur.execute("""
+            truncate table foobar
+        """)
+
+    conn.commit()
+
+    with tempfile.NamedTemporaryFile() as fp:
+        for start in xrange(10000):
+            fp.write('\t'.join([ str(x) for x in xrange(start, start+8) ] ))
+            fp.write("\n")
+        fp.seek(0)
+        with conn.cursor() as cur:
+            cur.copy_from(fp, 'foobar', columns = [
+                'some_long_column1',
+                'some_long_column2',
+                'some_long_column3',
+                'some_long_column4',
+                'some_long_column5',
+                'some_long_column6',
+                'some_long_column7',
+                'some_long_column8',
+            ])
+        conn.commit()
+
+@benchmark
+def execute(conn, sql, **bind_params):
+    """
+    Executes a SQL command against the connection with optional bind params.
+    """
+    with conn.cursor() as cur:
+        cur.execute(sql, bind_params)
+
+def iter_results(conn, sql, **bind_params):
+    """
+    Delays fetching the SQL results into memory until iteration
+    Keeps memory footprint low
+    """
+    with conn.cursor() as cur:
+        cur.execute(sql, bind_params)
+        for row in cur:
+            yield row
+
+@benchmark
+def fetch_iter_results(conn, sql, **bind_params):
+    return [ x for x in iter_results(conn, sql, **bind_params) ]
+
+@benchmark
+def fetch_results(conn, sql, **bind_params):
+    """
+    Immediatly fetches the SQL results into memory
+    Trades memory for the ability to immediately execute another query
+    """
+    with conn.cursor() as cur:
+        cur.execute(sql, bind_params)
+        return cur.fetchall()
+
+def test_conn(name, conn):
+    test_query = "select * from foobar"
+    for x in xrange(1000):
+        fetch_iter_results(conn_default, test_query)
+        fetch_results(conn_default, test_query)
+        conn.rollback()
+
+    print name
+    print BenchResults.format_stats()
+    BenchResults.clear()
+
+if __name__ == '__main__':
+    db_info = {
+        'host'     : 'localhost',
+        'port'     : 5432,
+        'user'     : 'pyutil',
+        'password' : 'pyutil',
+        'database' : 'pyutil_testdb',
+    }
+
+    conn_default = psycopg2.connect(**db_info)
+    conn_dict    = psycopg2.connect(cursor_factory = psycopg2.extras.DictCursor, **db_info)
+    conn_named   = psycopg2.connect(cursor_factory = psycopg2.extras.NamedTupleCursor, **db_info)
+
+    setup_test_table(conn_default)
+    test_conn('Default', conn_default)
+    test_conn('DictCursor', conn_dict)
+    test_conn('Named', conn_named)
+
+# 1k calls, cume duration
+# 10k rows fetched
+# +--------------------+----------------+--------------------+-------------------------+
+# |   Default Cursor   | psycopg2/2.7.3 | psycopg2cffi/2.7.3 | psycopg2cffi/pypy-2.0.2 |
+# +====================+================+====================+=========================+
+# | fetch_results      | 18.072         | 18.076             | 32.817                  |
+# +--------------------+----------------+--------------------+-------------------------+
+# | fetch_iter_results | 20.560         | 20.691             | 33.817                  |
+# +--------------------+----------------+--------------------+-------------------------+
+# 
+# +--------------------+----------------+--------------------+-------------------------+
+# |     DictCursor     | psycopg2/2.7.3 | psycopg2cffi/2.7.3 | psycopg2cffi/pypy-2.0.2 |
+# +====================+================+====================+=========================+
+# | fetch_results      | 18.405         | 18.377             | 32.434                  |
+# +--------------------+----------------+--------------------+-------------------------+
+# | fetch_iter_results | 19.563         | 19.674             | 33.265                  |
+# +--------------------+----------------+--------------------+-------------------------+
+# 
+# +--------------------+----------------+--------------------+-------------------------+
+# |  NamedTupleCursor  | psycopg2/2.7.3 | psycopg2cffi/2.7.3 | psycopg2cffi/pypy-2.0.2 |
+# +====================+================+====================+=========================+
+# | fetch_results      | 18.296         | 18.291             | 32.158                  |
+# +--------------------+----------------+--------------------+-------------------------+
+# | fetch_iter_results | 19.599         | 19.650             | 32.999                  |
+# +--------------------+----------------+--------------------+-------------------------+