1 import time

     2 import gobject, gtk, pango, cairo

     3 import gtksourceview2 as gtksourceview

     4 

     5 import psycopg2

     6 import psycopg2.extensions

     7 import psycopg2.extras

     8 

     9 from pgconsole.config import cfg

    10 from pgconsole.editor import Editor

    11 from pgconsole.dataview import DataView

    12 from pgconsole.database import Database, BadConnectionError, DatabaseError

    13 from pgconsole.settings import Settings

    14 from pgconsole.panedext import HPanedExt, VPanedExt

    15 

    16 

    17 class PgConsoleApp:

    18     def __init__(self):

    19         self.db = Database()

    20         self.conn = None

    21 

    22         win = gtk.Window(gtk.WINDOW_TOPLEVEL)

    23         self.win = win

    24         win.set_title('PostgreSQL Console')

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

    26         win.set_default_size(800, 600)

    27         self.restore_window_size()

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

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

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

    31 

    32         # toolbar

    33         toolbar = gtk.Toolbar()

    34         toolbar.set_style(gtk.TOOLBAR_ICONS)

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

    36 

    37         tb = gtk.ToolButton(gtk.STOCK_PREFERENCES)

    38         tb.set_label('Settings')

    39         tb.set_tooltip_text('Settings')

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

    41         self.tb_settings = tb

    42         toolbar.add(tb)

    43 

    44         self.cb_server = gtk.combo_box_entry_new_text()

    45         self.cb_server.set_tooltip_text('Server')

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

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

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

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

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

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

    52         ti = gtk.ToolItem()

    53         ti.add(self.cb_server)

    54         toolbar.add(ti)

    55 

    56         self.cb_dbname = gtk.combo_box_entry_new_text()

    57         self.cb_dbname.set_tooltip_text('Database')

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

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

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

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

    62         ti = gtk.ToolItem()

    63         ti.add(self.cb_dbname)

    64         toolbar.add(ti)

    65 

    66         tb = gtk.ToolButton(gtk.STOCK_CONNECT)

    67         tb.set_label('Connect')

    68         tb.set_tooltip_text('Connect')

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

    70         toolbar.add(tb)

    71         self.tb_connect = tb

    72 

    73         sep = gtk.SeparatorToolItem()

    74         toolbar.add(sep)

    75 

    76         tb = gtk.ToolButton(gtk.STOCK_EXECUTE)

    77         tb.set_label('Execute')

    78         tb.set_tooltip_text('Execute')

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

    80         toolbar.add(tb)

    81         tb.set_sensitive(False)

    82         self.tb_execute = tb

    83         tb = gtk.ToolButton(gtk.STOCK_NEW)

    84         tb.set_label('Begin transaction')

    85         tb.set_tooltip_text('Begin transaction')

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

    87         tb.set_sensitive(False)

    88         toolbar.add(tb)

    89         self.tb_begin = tb

    90         tb = gtk.ToolButton(gtk.STOCK_APPLY)

    91         tb.set_label('Commit')

    92         tb.set_tooltip_text('Commit')

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

    94         tb.set_sensitive(False)

    95         self.tb_commit = tb

    96         toolbar.add(tb)

    97         tb = gtk.ToolButton(gtk.STOCK_CANCEL)

    98         tb.set_label('Rollback')

    99         tb.set_tooltip_text('Rollback')

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

   101         tb.set_sensitive(False)

   102         self.tb_rollback = tb

   103         toolbar.add(tb)

   104 

   105         sep = gtk.SeparatorToolItem()

   106         toolbar.add(sep)

   107 

   108         # editor

   109         self.editor = Editor()

   110 

   111         # data view

   112         self.dataview = DataView()

   113 

   114 

   115         vbox = gtk.VBox(False, 2)

   116 

   117         sep = gtk.SeparatorToolItem()

   118         sep.set_expand(True)

   119         sep.set_draw(False)

   120         toolbar.add(sep)

   121 

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

   123         self.throbber = gtk.Image()

   124         tb = gtk.ToolItem()

   125         tb.add(self.throbber)

   126         toolbar.add(tb)

   127         sep = gtk.SeparatorToolItem()

   128         sep.set_draw(False)

   129         toolbar.add(sep)

   130 

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

   132 

   133         vpaned = VPanedExt()

   134         vpaned.set_border_width(5)

   135 

   136         hpaned = HPanedExt()

   137         hpaned.set_border_width(0)

   138         hpaned.add1(self.editor)

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

   140         hpaned.set_snap2(80)

   141         hpaned.set_property('position', 500)

   142 

   143         vpaned.add1(hpaned)

   144         vpaned.set_snap1(80)

   145         vpaned.set_property('position', 300)

   146 

   147         vpaned.add2(self.dataview)

   148         vpaned.set_snap2(80)

   149 

   150         self.vpaned = vpaned

   151         self.hpaned = hpaned

   152 

   153         # log

   154         self.logbuf = gtk.TextBuffer()

   155         view = gtk.TextView(self.logbuf)

   156         view.set_editable(False)

   157         font_desc = pango.FontDescription('monospace')

   158         if font_desc:

   159             view.modify_font(font_desc)

   160 

   161         sw = gtk.ScrolledWindow()

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

   163         sw.set_shadow_type(gtk.SHADOW_ETCHED_IN)

   164         sw.add(view)

   165         hpaned.add2(sw)

   166 

   167         vbox.pack_start(vpaned, padding=0)

   168 

   169         win.add(vbox)

   170 

   171         self.editor.view.grab_focus()

   172 

   173         self.reload_server_list(0)

   174 

   175         win.show_all()

   176 

   177         self.restore_win_state()

   178 

   179 

   180     def main(self):

   181         gtk.main()

   182 

   183 

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

   185         self.save_win_state()

   186         cfg.save()

   187         gtk.main_quit()

   188 

   189 

   190     def on_configure(self, w, ev):

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

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

   193 

   194 

   195     def restore_window_size(self):

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

   197 

   198 

   199     def save_win_state(self):

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

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

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

   203 

   204 

   205     def restore_win_state(self):

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

   207         if pos >= 0:

   208             self.vpaned.set_position(pos)

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

   210         if pos >= 0:

   211             self.hpaned.set_position(pos)

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

   213         if pos >= 0:

   214             self.editor.set_position(pos)

   215 

   216 

   217     def get_typename(self, oid, size):

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

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

   220         row = self.curs.fetchone()

   221 

   222         typname = None

   223         if row:

   224             typname = row[0]

   225             if typname == 'int4':

   226                 return 'integer'

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

   228                 return typname

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

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

   231         return typname

   232 

   233 

   234     def get_conninfo(self, nodb=False):

   235         sel = self.cb_server.get_active()

   236         srv = cfg.servers.server[sel]

   237         if nodb:

   238             dbname = 'postgres'

   239         else:

   240             dbname = self.cb_dbname.get_active_text()

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

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

   243         return conninfo

   244 

   245 

   246     def connect(self, w):

   247         conninfo = self.get_conninfo()

   248         if self.conn:

   249             # disconnect

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

   251             self.conn = None

   252             self.cb_server.set_sensitive(True)

   253             self.cb_dbname.set_sensitive(True)

   254             self.tb_connect.set_stock_id(gtk.STOCK_CONNECT)

   255             self.tb_connect.set_label('Connect')

   256             self.tb_connect.set_tooltip_text('Connect')

   257             self.tb_execute.set_sensitive(False)

   258             self.tb_begin.set_sensitive(False)

   259         else:

   260             # connect

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

   262             try:

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

   264             except DatabaseError as e:

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

   266                 return

   267             self.cb_server.set_sensitive(False)

   268             self.cb_dbname.set_sensitive(False)

   269             self.tb_connect.set_stock_id(gtk.STOCK_DISCONNECT)

   270             self.tb_connect.set_label('Disconnect')

   271             self.tb_connect.set_tooltip_text('Disconnect')

   272             self.tb_execute.set_sensitive(True)

   273             self.tb_begin.set_sensitive(True)

   274 

   275 

   276     def begin(self, w):

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

   278         curs = self.conn.cursor()

   279         curs.execute('BEGIN')

   280         psycopg2.extras.wait_select(curs.connection)

   281 

   282         self.tb_connect.set_sensitive(False)

   283         self.tb_begin.set_sensitive(False)

   284         self.tb_commit.set_sensitive(True)

   285         self.tb_rollback.set_sensitive(True)

   286 

   287 

   288     def commit(self, w):

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

   290         curs = self.conn.cursor()

   291         curs.execute('COMMIT')

   292         psycopg2.extras.wait_select(curs.connection)

   293 

   294         self.tb_connect.set_sensitive(True)

   295         self.tb_begin.set_sensitive(True)

   296         self.tb_commit.set_sensitive(False)

   297         self.tb_rollback.set_sensitive(False)

   298 

   299 

   300     def rollback(self, w):

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

   302         curs = self.conn.cursor()

   303         curs.execute('ROLLBACK')

   304         psycopg2.extras.wait_select(curs.connection)

   305 

   306         self.tb_connect.set_sensitive(True)

   307         self.tb_begin.set_sensitive(True)

   308         self.tb_commit.set_sensitive(False)

   309         self.tb_rollback.set_sensitive(False)

   310 

   311 

   312     def execute(self, widget):

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

   314 

   315         self.tb_connect.set_sensitive(False)

   316         self.tb_execute.set_sensitive(False)

   317         self.tb_begin.set_sensitive(False)

   318         self.tb_commit.set_sensitive(False)

   319         self.tb_rollback.set_sensitive(False)

   320         self.throbber.set_from_animation(self.throbber_anim)

   321 

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

   323 

   324         self.t1 = time.time()

   325         try:

   326             self.curs.execute(query)

   327         except (psycopg2.OperationalError, psycopg2.DatabaseError) as e:

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

   329             return

   330 

   331         self.execute_poll()

   332 

   333 

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

   335         try:

   336             state = self.conn.poll()

   337         except (psycopg2.OperationalError, psycopg2.DatabaseError) as e:

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

   339             return

   340 

   341         if state == psycopg2.extensions.POLL_OK:

   342             self.execute_finish()

   343         elif state == psycopg2.extensions.POLL_WRITE:

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

   345         elif state == psycopg2.extensions.POLL_READ:

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

   347         else:

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

   349             return

   350 

   351 

   352     def execute_finish(self):

   353         t2 = time.time()

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

   355 

   356         self.throbber.clear()

   357 

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

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

   360 

   361         # notices

   362         for n in self.conn.notices:

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

   364 

   365         if self.curs.rowcount >= 0:

   366             rows = self.curs.fetchall()

   367 

   368             names = []

   369             for c in self.curs.description:

   370                 name = c[0]

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

   372                 names += [(name, typename)]

   373 

   374             self.dataview.load_data(names, rows)

   375 

   376         self.tb_execute.set_sensitive(True)

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

   378             self.tb_commit.set_sensitive(True)

   379             self.tb_rollback.set_sensitive(True)

   380         else:

   381             self.tb_connect.set_sensitive(True)

   382             self.tb_begin.set_sensitive(True)

   383 

   384 

   385     def simulate_click(self, tb):

   386         if tb.get_property('sensitive'):

   387             tb.get_child().activate()

   388 

   389 

   390     def keypress(self, w, event):

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

   392         if keyname == 'F1':

   393             self.simulate_click(self.tb_settings)

   394             return True

   395         if keyname == 'F2':

   396             print(self.cb_server.popup())

   397             return True

   398         if keyname == 'F3':

   399             print(self.cb_dbname.popup())

   400             return True

   401         if keyname == 'F4':

   402             self.simulate_click(self.tb_connect)

   403             return True

   404         if keyname == 'F5':

   405             self.simulate_click(self.tb_execute)

   406             return True

   407         if keyname == 'F6':

   408             self.simulate_click(self.tb_begin)

   409             return True

   410         if keyname == 'F7':

   411             self.simulate_click(self.tb_commit)

   412             return True

   413         if keyname == 'F8':

   414             self.simulate_click(self.tb_rollback)

   415             return True

   416 

   417         return False

   418 

   419 

   420     def toolbar_server_keypress(self, w, event):

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

   422         if keyname == 'Tab':

   423             self.populate_db_list()

   424             self.cb_dbname.grab_focus()

   425             return True

   426 

   427 

   428     def on_server_changed(self, w):

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

   430             self.populate_db_list()

   431 

   432 

   433     def populate_db_list(self):

   434         conninfo = self.get_conninfo(True)

   435         try:

   436             conn = self.db.get_conn(conninfo)

   437         except DatabaseError as e:

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

   439             return

   440 

   441         curs = conn.cursor()

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

   443         psycopg2.extras.wait_select(conn)

   444         rows = curs.fetchall()

   445 

   446         self.db.put_conn(conninfo, conn)

   447 

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

   449             self.cb_dbname.remove_text(0)

   450 

   451         for row in rows:

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

   453 

   454         self.cb_dbname.set_active(0)

   455 

   456 

   457     def reload_server_list(self, sel=None):

   458         # clean

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

   460             self.cb_server.remove_text(0)

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

   462 

   463         try:

   464             # populate

   465             for server in cfg.servers.server:

   466                 if str(server.name):

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

   468                 else:

   469                     title = str(server.host)

   470                 self.cb_server.append_text(title)

   471         except AttributeError:

   472             pass

   473 

   474         if not sel is None:

   475             self.cb_server.set_active(sel)

   476 

   477 

   478     def settings(self, w):

   479         Settings(self)

   480