1 #!/usr/bin/env python

     2 

     3 import time

     4 import gobject, gtk, pango, cairo

     5 import gtksourceview2 as gtksourceview

     6 

     7 import psycopg2

     8 import psycopg2.extensions

     9 import psycopg2.extras

    10 

    11 from pgconsole.config import cfg

    12 from pgconsole.editor import Editor

    13 from pgconsole.dataview import DataView

    14 from pgconsole.database import Database, BadConnectionError, DatabaseError

    15 from pgconsole.settings import Settings

    16 from pgconsole.panedext import HPanedExt, VPanedExt

    17 

    18 

    19 class PgConsoleApp:

    20     def __init__(self):

    21         self.db = Database()

    22         self.conn = None

    23 

    24         win = gtk.Window(gtk.WINDOW_TOPLEVEL)

    25         self.win = win

    26         win.set_title('PostgreSQL Console')

    27         win.set_size_request(300, 200)  # minimal size

    28         win.set_default_size(800, 600)

    29         self.restore_window_size()

    30         win.connect("destroy", self.destroy)

    31         win.connect("key_press_event", self.keypress)

    32         win.connect('configure-event', self.on_configure)

    33 

    34         # toolbar

    35         toolbar = gtk.Toolbar()

    36         toolbar.set_style(gtk.TOOLBAR_ICONS)

    37         toolbar.set_property("icon-size", gtk.ICON_SIZE_SMALL_TOOLBAR)

    38 

    39         tb = gtk.ToolButton(gtk.STOCK_PREFERENCES)

    40         tb.set_label('Settings')

    41         tb.set_tooltip_text('Settings')

    42         tb.connect('clicked', self.settings)

    43         self.tb_settings = tb

    44         toolbar.add(tb)

    45 

    46         self.cb_server = gtk.combo_box_entry_new_text()

    47         self.cb_server.set_tooltip_text('Server')

    48         self.cb_server.get_child().set_property("editable", False)

    49         self.cb_server.set_property("add_tearoffs", True)

    50         #self.cb_server.set_property("focus-on-click", False)

    51         self.cb_server.set_property("can-focus", True)

    52         self.cb_server.connect('changed', self.on_server_changed)

    53         self.cb_server.connect("key_press_event", self.toolbar_server_keypress)

    54         ti = gtk.ToolItem()

    55         ti.add(self.cb_server)

    56         toolbar.add(ti)

    57 

    58         self.cb_dbname = gtk.combo_box_entry_new_text()

    59         self.cb_dbname.set_tooltip_text('Database')

    60         self.cb_dbname.get_child().set_property("editable", False)

    61         self.cb_dbname.set_property("add_tearoffs", True)

    62         #self.cb_dbname.set_property("focus-on-click", False)

    63         self.cb_dbname.set_property("can-focus", True)

    64         ti = gtk.ToolItem()

    65         ti.add(self.cb_dbname)

    66         toolbar.add(ti)

    67 

    68         tb = gtk.ToolButton(gtk.STOCK_CONNECT)

    69         tb.set_label('Connect')

    70         tb.set_tooltip_text('Connect')

    71         tb.connect('clicked', self.connect)

    72         toolbar.add(tb)

    73         self.tb_connect = tb

    74 

    75         sep = gtk.SeparatorToolItem()

    76         toolbar.add(sep)

    77 

    78         tb = gtk.ToolButton(gtk.STOCK_EXECUTE)

    79         tb.set_label('Execute')

    80         tb.set_tooltip_text('Execute')

    81         tb.connect('clicked', self.execute)

    82         toolbar.add(tb)

    83         tb.set_sensitive(False)

    84         self.tb_execute = tb

    85         tb = gtk.ToolButton(gtk.STOCK_NEW)

    86         tb.set_label('Begin transaction')

    87         tb.set_tooltip_text('Begin transaction')

    88         tb.connect('clicked', self.begin)

    89         tb.set_sensitive(False)

    90         toolbar.add(tb)

    91         self.tb_begin = tb

    92         tb = gtk.ToolButton(gtk.STOCK_APPLY)

    93         tb.set_label('Commit')

    94         tb.set_tooltip_text('Commit')

    95         tb.connect('clicked', self.commit)

    96         tb.set_sensitive(False)

    97         self.tb_commit = tb

    98         toolbar.add(tb)

    99         tb = gtk.ToolButton(gtk.STOCK_CANCEL)

   100         tb.set_label('Rollback')

   101         tb.set_tooltip_text('Rollback')

   102         tb.connect('clicked', self.rollback)

   103         tb.set_sensitive(False)

   104         self.tb_rollback = tb

   105         toolbar.add(tb)

   106 

   107         sep = gtk.SeparatorToolItem()

   108         toolbar.add(sep)

   109 

   110         # editor

   111         self.editor = Editor()

   112 

   113         # data view

   114         self.dataview = DataView()

   115 

   116 

   117         vbox = gtk.VBox(False, 2)

   118 

   119         sep = gtk.SeparatorToolItem()

   120         sep.set_expand(True)

   121         sep.set_draw(False)

   122         toolbar.add(sep)

   123 

   124         self.throbber_anim = gtk.gdk.PixbufAnimation('data/Throbber.gif')

   125         self.throbber = gtk.Image()

   126         tb = gtk.ToolItem()

   127         tb.add(self.throbber)

   128         toolbar.add(tb)

   129         sep = gtk.SeparatorToolItem()

   130         sep.set_draw(False)

   131         toolbar.add(sep)

   132 

   133         vbox.pack_start(toolbar, False, False, 0)

   134 

   135         vpaned = VPanedExt()

   136         vpaned.set_border_width(5)

   137 

   138         hpaned = HPanedExt()

   139         hpaned.set_border_width(0)

   140         hpaned.add1(self.editor)

   141         hpaned.child_set_property(self.editor, 'shrink', False)

   142         hpaned.set_snap2(80)

   143         hpaned.set_property('position', 500)

   144 

   145         vpaned.add1(hpaned)

   146         vpaned.set_snap1(80)

   147         vpaned.set_property('position', 300)

   148 

   149         vpaned.add2(self.dataview)

   150         vpaned.set_snap2(80)

   151 

   152         self.vpaned = vpaned

   153         self.hpaned = hpaned

   154 

   155         # log

   156         self.logbuf = gtk.TextBuffer()

   157         view = gtk.TextView(self.logbuf)

   158         view.set_editable(False)

   159         font_desc = pango.FontDescription('monospace')

   160         if font_desc:

   161             view.modify_font(font_desc)

   162 

   163         sw = gtk.ScrolledWindow()

   164         sw.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)

   165         sw.set_shadow_type(gtk.SHADOW_ETCHED_IN)

   166         sw.add(view)

   167         hpaned.add2(sw)

   168 

   169         vbox.pack_start(vpaned, padding=0)

   170 

   171         win.add(vbox)

   172 

   173         self.editor.view.grab_focus()

   174 

   175         self.reload_server_list(0)

   176 

   177         win.show_all()

   178 

   179         self.restore_win_state()

   180 

   181 

   182     def main(self):

   183         gtk.main()

   184 

   185 

   186     def destroy(self, widget, data=None):

   187         self.save_win_state()

   188         cfg.save()

   189         gtk.main_quit()

   190 

   191 

   192     def on_configure(self, w, ev):

   193         cfg.root.window.size.width = ev.width

   194         cfg.root.window.size.height = ev.height

   195 

   196 

   197     def restore_window_size(self):

   198         self.win.resize(cfg.root.window.size.width, cfg.root.window.size.height)

   199 

   200 

   201     def save_win_state(self):

   202         cfg.root.window.dividers.verticaldivider = self.vpaned.get_position()

   203         cfg.root.window.dividers.horizontaldivider = self.hpaned.get_position()

   204         cfg.root.window.dividers.editordivider = self.editor.get_position()

   205 

   206 

   207     def restore_win_state(self):

   208         pos = cfg.root.window.dividers.verticaldivider

   209         if pos >= 0:

   210             self.vpaned.set_position(pos)

   211         pos = cfg.root.window.dividers.horizontaldivider

   212         if pos >= 0:

   213             self.hpaned.set_position(pos)

   214         pos = cfg.root.window.dividers.editordivider

   215         if pos >= 0:

   216             self.editor.set_position(pos)

   217 

   218 

   219     def get_typename(self, oid, size):

   220         self.curs.execute('SELECT typname FROM pg_type WHERE oid=%s', [oid])

   221         psycopg2.extras.wait_select(self.curs.connection)

   222         row = self.curs.fetchone()

   223 

   224         typname = None

   225         if row:

   226             typname = row[0]

   227             if typname == 'int4':

   228                 return 'integer'

   229             if typname in ('timestamp', 'interval', 'date'):

   230                 return typname

   231             if size and size > 0 and size < 65535:

   232                 typname += '(%s)' % size

   233         return typname

   234 

   235 

   236     def get_conninfo(self, nodb=False):

   237         sel = self.cb_server.get_active()

   238         srv = cfg.servers.server[sel]

   239         if nodb:

   240             dbname = 'postgres'

   241         else:

   242             dbname = self.cb_dbname.get_active_text()

   243         conninfo = 'host=%s port=%s dbname=%s user=%s password=%s' \

   244             % (srv.host, srv.port, dbname, srv.user, srv.password)

   245         return conninfo

   246 

   247 

   248     def connect(self, w):

   249         conninfo = self.get_conninfo()

   250         if self.conn:

   251             # disconnect

   252             self.db.put_conn(conninfo, self.conn)

   253             self.conn = None

   254             self.cb_server.set_sensitive(True)

   255             self.cb_dbname.set_sensitive(True)

   256             self.tb_connect.set_stock_id(gtk.STOCK_CONNECT)

   257             self.tb_connect.set_label('Connect')

   258             self.tb_connect.set_tooltip_text('Connect')

   259             self.tb_execute.set_sensitive(False)

   260             self.tb_begin.set_sensitive(False)

   261         else:

   262             # connect

   263             self.logbuf.insert(self.logbuf.get_end_iter(), 'Connect %s\n' % conninfo)

   264             try:

   265                 self.conn = self.db.get_conn(conninfo)

   266             except DatabaseError, e:

   267                 self.logbuf.insert(self.logbuf.get_end_iter(), 'Error:\n%s\n' % e)

   268                 return

   269             self.cb_server.set_sensitive(False)

   270             self.cb_dbname.set_sensitive(False)

   271             self.tb_connect.set_stock_id(gtk.STOCK_DISCONNECT)

   272             self.tb_connect.set_label('Disconnect')

   273             self.tb_connect.set_tooltip_text('Disconnect')

   274             self.tb_execute.set_sensitive(True)

   275             self.tb_begin.set_sensitive(True)

   276 

   277 

   278     def begin(self, w):

   279         self.logbuf.insert(self.logbuf.get_end_iter(), 'Begin transaction\n')

   280         curs = self.conn.cursor()

   281         curs.execute('BEGIN')

   282         psycopg2.extras.wait_select(curs.connection)

   283 

   284         self.tb_connect.set_sensitive(False)

   285         self.tb_begin.set_sensitive(False)

   286         self.tb_commit.set_sensitive(True)

   287         self.tb_rollback.set_sensitive(True)

   288 

   289 

   290     def commit(self, w):

   291         self.logbuf.insert(self.logbuf.get_end_iter(), 'Commit\n')

   292         curs = self.conn.cursor()

   293         curs.execute('COMMIT')

   294         psycopg2.extras.wait_select(curs.connection)

   295 

   296         self.tb_connect.set_sensitive(True)

   297         self.tb_begin.set_sensitive(True)

   298         self.tb_commit.set_sensitive(False)

   299         self.tb_rollback.set_sensitive(False)

   300 

   301 

   302     def rollback(self, w):

   303         self.logbuf.insert(self.logbuf.get_end_iter(), 'Rollback\n')

   304         curs = self.conn.cursor()

   305         curs.execute('ROLLBACK')

   306         psycopg2.extras.wait_select(curs.connection)

   307 

   308         self.tb_connect.set_sensitive(True)

   309         self.tb_begin.set_sensitive(True)

   310         self.tb_commit.set_sensitive(False)

   311         self.tb_rollback.set_sensitive(False)

   312 

   313 

   314     def execute(self, widget):

   315         query = self.editor.get_selection() or self.editor.get_text()

   316 

   317         self.tb_connect.set_sensitive(False)

   318         self.tb_execute.set_sensitive(False)

   319         self.tb_begin.set_sensitive(False)

   320         self.tb_commit.set_sensitive(False)

   321         self.tb_rollback.set_sensitive(False)

   322         self.throbber.set_from_animation(self.throbber_anim)

   323 

   324         self.curs = self.conn.cursor()

   325 

   326         self.t1 = time.time()

   327         try:

   328             self.curs.execute(query)

   329         except (psycopg2.OperationalError, psycopg2.DatabaseError), e:

   330             self.logbuf.insert(self.logbuf.get_end_iter(), 'Error:\n' + str(e))

   331             return

   332 

   333         self.execute_poll()

   334 

   335 

   336     def execute_poll(self, source=None, cond=None):

   337         try:

   338             state = self.conn.poll()

   339         except (psycopg2.OperationalError, psycopg2.DatabaseError), e:

   340             self.logbuf.insert(self.logbuf.get_end_iter(), 'Error:\n' + str(e))

   341             return

   342 

   343         if state == psycopg2.extensions.POLL_OK:

   344             self.execute_finish()

   345         elif state == psycopg2.extensions.POLL_WRITE:

   346             gobject.io_add_watch(self.conn.fileno(), gobject.IO_OUT, self.execute_poll)

   347         elif state == psycopg2.extensions.POLL_READ:

   348             gobject.io_add_watch(self.conn.fileno(), gobject.IO_IN, self.execute_poll)

   349         else:

   350             self.logbuf.insert(self.logbuf.get_end_iter(), "poll() returned %s" % state)

   351             return

   352 

   353 

   354     def execute_finish(self):

   355         t2 = time.time()

   356         t = (t2 - self.t1)*1000

   357 

   358         self.throbber.clear()

   359 

   360         self.logbuf.insert(self.logbuf.get_end_iter(),

   361             'Query successful (%d ms, %d rows)\n' % (t, self.curs.rowcount))

   362 

   363         # notices

   364         for n in self.conn.notices:

   365             self.logbuf.insert(self.logbuf.get_end_iter(), n)

   366 

   367         if self.curs.rowcount >= 0:

   368             rows = self.curs.fetchall()

   369 

   370             names = []

   371             for c in self.curs.description:

   372                 name = c[0]

   373                 typename = self.get_typename(c[1], c[3])

   374                 names += [(name, typename)]

   375 

   376             self.dataview.load_data(names, rows)

   377 

   378         self.tb_execute.set_sensitive(True)

   379         if self.conn.get_transaction_status() == psycopg2.extensions.TRANSACTION_STATUS_INTRANS:

   380             self.tb_commit.set_sensitive(True)

   381             self.tb_rollback.set_sensitive(True)

   382         else:

   383             self.tb_connect.set_sensitive(True)

   384             self.tb_begin.set_sensitive(True)

   385 

   386 

   387     def simulate_click(self, tb):

   388         if tb.get_property('sensitive'):

   389             tb.get_child().activate()

   390 

   391 

   392     def keypress(self, w, event):

   393         keyname = gtk.gdk.keyval_name(event.keyval)

   394         if keyname == 'F1':

   395             self.simulate_click(self.tb_settings)

   396             return True

   397         if keyname == 'F2':

   398             print self.cb_server.popup()

   399             return True

   400         if keyname == 'F3':

   401             print self.cb_dbname.popup()

   402             return True

   403         if keyname == 'F4':

   404             self.simulate_click(self.tb_connect)

   405             return True

   406         if keyname == 'F5':

   407             self.simulate_click(self.tb_execute)

   408             return True

   409         if keyname == 'F6':

   410             self.simulate_click(self.tb_begin)

   411             return True

   412         if keyname == 'F7':

   413             self.simulate_click(self.tb_commit)

   414             return True

   415         if keyname == 'F8':

   416             self.simulate_click(self.tb_rollback)

   417             return True

   418 

   419         return False

   420 

   421 

   422     def toolbar_server_keypress(self, w, event):

   423         keyname = gtk.gdk.keyval_name(event.keyval)

   424         if keyname == 'Tab':

   425             self.populate_db_list()

   426             self.cb_dbname.grab_focus()

   427             return True

   428 

   429 

   430     def on_server_changed(self, w):

   431         if self.cb_server.get_active() != -1:

   432             self.populate_db_list()

   433 

   434 

   435     def populate_db_list(self):

   436         conninfo = self.get_conninfo(True)

   437         try:

   438             conn = self.db.get_conn(conninfo)

   439         except DatabaseError, e:

   440             self.logbuf.insert(self.logbuf.get_end_iter(), 'Error:\n%s\n' % e)

   441             return

   442 

   443         curs = conn.cursor()

   444         curs.execute('SELECT * FROM pg_catalog.pg_database WHERE NOT datistemplate ORDER BY 1')

   445         psycopg2.extras.wait_select(conn)

   446         rows = curs.fetchall()

   447 

   448         self.db.put_conn(conninfo, conn)

   449 

   450         for i in range(self.cb_dbname.get_model().iter_n_children(None)):

   451             self.cb_dbname.remove_text(0)

   452 

   453         for row in rows:

   454             self.cb_dbname.append_text(row[0])

   455 

   456         self.cb_dbname.set_active(0)

   457 

   458 

   459     def reload_server_list(self, sel=None):

   460         # clean

   461         for i in range(self.cb_server.get_model().iter_n_children(None)):

   462             self.cb_server.remove_text(0)

   463         self.cb_server.get_child().set_text('')

   464 

   465         try:

   466             # populate

   467             for server in cfg.servers.server:

   468                 if str(server.name):

   469                     title = '%s (%s)' % (str(server.name), str(server.host))

   470                 else:

   471                     title = str(server.host)

   472                 self.cb_server.append_text(title)

   473         except AttributeError:

   474             pass

   475 

   476         if not sel is None:

   477             self.cb_server.set_active(sel)

   478 

   479 

   480     def settings(self, w):

   481         Settings(self)

   482 

   483 

   484 

   485 if __name__ == '__main__':

   486     try:

   487         cfg.load('pgconsole.xml.gz')

   488     except IOError:

   489         cfg.new('pgconsole.xml.gz')

   490 

   491     try:

   492         gtkrc = cfg.root.style.gtkrc.text

   493     except AttributeError:

   494         pass

   495     else:

   496         gtk.rc_parse_string(gtkrc)

   497 

   498     app = PgConsoleApp()

   499     app.main()

   500