Merge lp:~chromakode/boots/postgresql into lp:boots

[Max Goodhart]

Initial PostgreSQL support added. Table output is wonky, because psycopg2 does not calculate display_size for columns, so we'd need to calculate the maximum sizes ourselves.

We can merge this code now, but it won't be very useful until DSN support lands, or the user is able to specify the server type some other way.

In a future patch, we should probably update gen_table() in plain.py to calculate the display size if the value we got is None.

[lp:~chromakode/boots/postgresql]

163. By Max Goodhart on 2010-03-29

Add slow manual display_size calculation in output ascii table generation, in case display_size is not provided by the database driver.

[Max Goodhart (chromakode)]

And there's display_size calculation in gen_table(). Now PostgreSQL connections should be pretty indistinguishable from other connections.

[lp:~chromakode/boots/postgresql]

164. By Max Goodhart on 2010-03-29

Stylistic improvement to query code for PostgreSQL server_version.

Unmerged revisions

164. By Max Goodhart on 2010-03-29

Stylistic improvement to query code for PostgreSQL server_version.

163. By Max Goodhart on 2010-03-29

Add slow manual display_size calculation in output ascii table generation, in case display_size is not provided by the database driver.

162. By Max Goodhart on 2010-03-29

Initial internal PostgreSQL support, with the required tweaks and fixes.

161. By Max Goodhart on 2010-03-29

Refactor boots.api.api into boots.api.server.

=== modified file 'boots/api/constructors.py'
--- boots/api/constructors.py	2010-03-24 07:33:04 +0000
+++ boots/api/constructors.py	2010-03-29 08:42:30 +0000
@@ -23,7 +23,7 @@
 
 import re
 
-from boots.api.api import Rows, ResultInfo
+from boots.api.server import Rows, ResultInfo
 from boots.api.errors import BootsError
 from boots.api.nodes.node import NodeGraph, SyncNode, IteratorNode
 from boots.api import _, n_
 
=== modified file 'boots/api/errors.py'
--- boots/api/errors.py	2010-03-28 08:27:19 +0000
+++ boots/api/errors.py	2010-03-29 08:42:30 +0000
@@ -85,7 +85,10 @@
         self.code = code
         
     def __str__(self):
-        return "{msg} (#{code})".format(msg=self.msg, code=self.code)
+        # Add the error code to the end of the first line of the error message.
+        lines = self.msg.split("\n", 1)
+        lines[0] ="{msg} (#{code})".format(msg=lines[0], code=self.code)
+        return "\n".join(lines)
     
 class ConnectionError(ServerError, CausedException):
     """Errors related to server connections."""
 
=== renamed file 'boots/api/api.py' => 'boots/api/server.py'
--- boots/api/api.py	2010-03-27 10:24:03 +0000
+++ boots/api/server.py	2010-03-29 08:42:30 +0000
@@ -39,11 +39,12 @@
 ColumnDescription = namedtuple("ColumnDescription", 
                                ("name", "type_code", "display_size",
                                 "internal_size", "precision", "scale",
-                                "null_ok"))
+                                "null_ok", "index"))
 
 class RowDescription:
     def __init__(self, description):
-        self.description = [ColumnDescription(*column) for column in description]
+        self.description = [ColumnDescription(*column, index=index) 
+                            for index, column in enumerate(description)]
         self.column_index = dict(zip((column[0] for column in description),
                                      range(len(description))))
     
@@ -265,11 +266,6 @@
             Server.__init__(self, MySQLdb, hostname, port, database, auth)
 
         def _connect(self):
-            """Establishes database connection.
-
-            Connect to the database server whose details were provided through
-            the __init__ method."""
-
             converter = self.db.converters.conversions.copy()
             for x in range(0,256):
                 converter.pop(x, None)
@@ -304,6 +300,62 @@
 else:
     _server_classes.append(MySQLdbServer)
 
+try:
+    import psycopg2
+    import psycopg2.extensions
+    
+    # Disable psycopg2 type conversions
+    raw = psycopg2.extensions.new_type(tuple(psycopg2.extensions.string_types.keys()),
+                                       "raw",
+                                       lambda value, cur:value if value is not None else None)
+    psycopg2.extensions.register_type(raw)
+    psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
+    
+    class PsycopgServer(Server):
+        def __init__(self, hostname, port, database, auth):
+            Server.__init__(self, psycopg2, hostname, port, database, auth)
+
+        def _connect(self):
+            connection = self.db.connect(
+                host     =  self.hostname,  
+                port     =  self.port,
+                database =  self.database,
+                user     =  self.auth.get("username") or "",
+                password =  self.auth.get("password") or "")
+            connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+            return connection
+
+        def _wrap_exception(self, e):
+            if isinstance(e, psycopg2.Error):
+                msg = e.pgerror
+                if msg is not None and msg.startswith("ERROR:"):
+                    # Remove the ERROR: title.
+                    msg = msg.replace("ERROR:", "", 1).strip()
+                
+                if e.pgcode:
+                    return errors.CodedDatabaseError(e.pgcode, msg, e)
+                else:
+                    return errors.DatabaseError(msg, e)
+            else:
+                return e
+
+        @property
+        def server_version(self):
+            self._check_connected()
+            if hasattr(self._connection, "server_version"):
+                version = self._connection.server_version
+                return "{major}.{minor}.{revision}".format(major=version/10000,
+                                                           minor=version/100%100,
+                                                           revision=version%10000)
+            else:
+                # Versions of psycopg2 before 2.0.12 don't have the server_version property.
+                return self._execute("show server_version;").fetchone()[0]
+        
+except ImportError:
+    PsycopgServer = None
+else:
+    _server_classes.append(PsycopgServer)
+
 def DrizzleServer(hostname, port, database, auth):
     if LibDrizzleServer is not None:
         return LibDrizzleServer(hostname, port, database, auth)
 
=== modified file 'boots/lib/console.py'
--- boots/lib/console.py	2010-03-28 23:55:27 +0000
+++ boots/lib/console.py	2010-03-29 08:42:30 +0000
@@ -28,7 +28,7 @@
 import os.path
 import signal
 
-from boots.api import api
+from boots.api.server import DrizzleServer
 from boots.api.errors import BootsError, BootsWarning, CausedException, ConnectionError
 from boots.api.nodes.node import NodeGraph, SyncNode, IteratorNode, LinkStackNode
 from boots.lib.ui.plain import PlainUI
@@ -201,8 +201,8 @@
         return result_data
 
     def connect(self, host, port, database=None, username=None, password=None):
-        server = api.DrizzleServer(host, port, database,
-                                   {"username": username, "password": password})
+        server = DrizzleServer(host, port, database,
+                               {"username": username, "password": password})
         server.connect()
         self.servers.append(server)
         return True
 
=== modified file 'boots/lib/ui/plain.py'
--- boots/lib/ui/plain.py	2010-03-28 23:57:18 +0000
+++ boots/lib/ui/plain.py	2010-03-29 08:42:30 +0000
@@ -30,9 +30,9 @@
 import readline
 import subprocess
 import threading
-from itertools import izip
+from itertools import imap
 
-from boots.api.api import Rows, ResultInfo, DurationCounter
+from boots.api.server import Rows, ResultInfo, DurationCounter
 from boots.api.errors import CausedException
 from boots.api.nodes.node import Status, NodeGraph, SyncNode, QueueNode, filter_none
 from boots.lib.ui.components.help import HelpTopic, CommandHelpTopic
@@ -184,13 +184,20 @@
             conversion."""
             return value if value is not None else "NULL"
 
+        def display_size(column):
+            if column.display_size is not None:
+                return column.display_size
+            else:
+                # Welcome to slowville, population you.
+                return max(imap(len, (row[column.index] for row in self.buffer)))
+
         def _gen_table(info):
             with DurationCounter() as table_elapsed:
                 if self.buffer:
                     # If a column can contain the value NULL, we leave at least 4
                     # characters, since display_size does not account for NULL.
                     max_widths = map(max, [(len(column.name),
-                                            column.display_size,
+                                            display_size(column),
                                             4 if column.null_ok else 0)
                                      for column in info["description"]])
                     dashes = map(lambda x: "-"*(x+2), max_widths)
@@ -210,9 +217,12 @@
                                    "{count} rows in set",
                                    info["row_count"])
                 else:
-                    set_count = n_("{count} row affected",
-                                   "{count} rows affected",
-                                   info["row_count"])
+                    if info["row_count"] != -1:
+                        set_count = n_("{count} row affected",
+                                       "{count} rows affected",
+                                       info["row_count"])
+                    else:
+                        set_count = _("Operation completed")
 
                 nodes_elapsed = info["end_time"] - info["begin_time"] - info["server_elapsed"]                
                 timings = [_("{0:.2f}s server").format(info["server_elapsed"]),
 
=== renamed file 'tests/boots/api/api_tests.py' => 'tests/boots/api/server_tests.py'
--- tests/boots/api/api_tests.py	2010-03-07 10:59:51 +0000
+++ tests/boots/api/server_tests.py	2010-03-29 08:42:30 +0000
@@ -25,12 +25,13 @@
 import unittest
 from itertools import ifilter
 
-from boots.api import errors, api
+from boots.api import errors
+from boots.api.server import DrizzleServer
 from boots.api.nodes.node import filter_default
 import boots_unit_test
 
 class TestServerConnections(boots_unit_test.BootsQueryTester):
-    """Class that provides unit tests that test the api.DrizzleServer class' ability
+    """Class that provides unit tests that test the DrizzleServer class' ability
     to connect to servers as well as provide proper errors when connections
     should fail."""
     #FIXME should use the configuration provided.
@@ -40,54 +41,54 @@
         pass
 
     def test_valid_connect_fqdn(self):
-        server = api.DrizzleServer("capstonedd.cs.pdx.edu", 9306, {})
+        server = DrizzleServer("capstonedd.cs.pdx.edu", 9306, {})
         server.connect()
         self.assert_(server.is_connected)
         server.disconnect()
 
     def test_valid_connect_partial(self):
-        server = api.DrizzleServer("capstonedd", 9306, {})
+        server = DrizzleServer("capstonedd", 9306, {})
         server.connect()
         self.assert_(server.is_connected)
         server.disconnect()
 
     def test_valid_connect_ip(self):
-        server = api.DrizzleServer("131.252.214.178", 9306, {})
+        server = DrizzleServer("131.252.214.178", 9306, {})
         server.connect()
         self.assert_(server.is_connected)
         server.disconnect()
 
     def test_valid_connect_fqdn_with_db(self):
         boots_unit_test.BootsQueryTester.setUp(self)
-        server = api.DrizzleServer("capstonedd.cs.pdx.edu", 9306, {"database": "fec"})
+        server = DrizzleServer("capstonedd.cs.pdx.edu", 9306, {"database": "fec"})
         server.connect()
         self.assert_(server.is_connected)
         server.disconnect()
 
     def test_invalid_connect(self):
-        server = api.DrizzleServer("kororaa.cs.pdx.edu", 9306, {})
+        server = DrizzleServer("kororaa.cs.pdx.edu", 9306, {})
         self.assertRaises(errors.ConnectionError, server.connect)
         self.assert_(not server.is_connected)
 
     def test_valid_disconnect(self):
-        server = api.DrizzleServer("capstonedd.cs.pdx.edu", 9306, {})
+        server = DrizzleServer("capstonedd.cs.pdx.edu", 9306, {})
         server.connect()
         server.disconnect()
         self.assert_(not server.is_connected)
 
     def test_invalid_disconnect(self):
-        server = api.DrizzleServer("kororaa.cs.pdx.edu", 9306, {})
+        server = DrizzleServer("kororaa.cs.pdx.edu", 9306, {})
         self.assertRaises(errors.ConnectionError, server.connect)
         self.assertRaises(errors.ConnectionError, server.disconnect)
         self.assert_(not server.is_connected)
 
 class TestServerExecute(boots_unit_test.BootsQueryTester):
-    """Class that provides unit tests that tests api.DrizzleServer class' ability
+    """Class that provides unit tests that tests DrizzleServer class' ability
     to execute SQL statements on a server."""
     #FIXME needs to test more than just SELECT statements.
     def setUp(self):
         boots_unit_test.BootsQueryTester.setUp(self)
-        self.server = api.DrizzleServer(self.config["host"], self.config["port"], {"database": self.config["database"]})
+        self.server = DrizzleServer(self.config["host"], self.config["port"], {"database": self.config["database"]})
         self.server.connect()
 
     def tearDown(self):
 
=== modified file 'tests/boots/boots_unit_test.py'
--- tests/boots/boots_unit_test.py	2010-03-07 10:59:51 +0000
+++ tests/boots/boots_unit_test.py	2010-03-29 08:42:30 +0000
@@ -59,7 +59,7 @@
         fec_sql_path = os.path.join(base_tests_path, "fec.sql")
         boots_path = os.path.join(base_tests_path, "../../bin/boots")
         
-        cmd = "{boots} -H {host} -p {port} -f {file}".format(boots=boots_path,
+        cmd = "{boots} -h {host} -p {port} -f {file}".format(boots=boots_path,
                                                              file=fec_sql_path, 
                                                              **self.config)
         p = Popen(cmd, shell=True)
 
=== modified file 'tests/boots/tests.py'
--- tests/boots/tests.py	2010-03-07 10:59:51 +0000
+++ tests/boots/tests.py	2010-03-29 08:42:30 +0000
@@ -35,7 +35,7 @@
 # Path mangling to import boots files properly.
 new_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../..")
 sys.path.insert(0, new_path)
-from api import api_tests
+from api import server_tests
 from api.nodes import node_tests
 from app import app_tests
 from lib import console_tests
@@ -59,7 +59,7 @@
     (options, args) = parser.parse_args()
     BootsBaseTester.config.update(vars(options))
     if not args:
-        modules = [api_tests, node_tests, app_tests,
+        modules = [server_tests, node_tests, app_tests,
                    console_tests, lingo_tests, ui_tests]
     else:
         modules = map(globals().__getitem__, args)