QGIS/python/plugins/db_manager/dlg_sql_window.py

# -*- coding: utf-8 -*-

"""
/***************************************************************************
Name                 : DB Manager
Description          : Database manager plugin for QuantumGIS
Date                 : May 23, 2011
copyright            : (C) 2011 by Giuseppe Sucameli
email                : brush.tyler@gmail.com

The content of this file is based on
- PG_Manager by Martin Dobias (GPLv2 license)
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""

from PyQt4.QtCore import *
from PyQt4.QtGui import *

from .db_plugins.plugin import BaseError
from .dlg_db_error import DlgDbError

from .ui.ui_DlgSqlWindow import Ui_DbManagerDlgSqlWindow as Ui_Dialog

from .highlighter import SqlHighlighter
from .completer import SqlCompleter

class DlgSqlWindow(QDialog, Ui_Dialog):

  def __init__(self, iface, db, parent=None):
    QDialog.__init__(self, parent)
    self.iface = iface
    self.db = db
    self.setupUi(self)
    self.setWindowTitle( u"%s - %s [%s]" % (self.windowTitle(), db.connection().connectionName(), db.connection().typeNameString()) )

    self.defaultLayerName = 'QueryLayer'

    settings = QSettings()
    self.restoreGeometry(settings.value("/DB_Manager/sqlWindow/geometry").toByteArray())

    self.editSql.setAcceptRichText(False)
    SqlCompleter(self.editSql, self.db)
    SqlHighlighter(self.editSql, self.db)

    # allow to copy results
    copyAction = QAction("copy", self)
    self.viewResult.addAction( copyAction )
    copyAction.setShortcuts(QKeySequence.Copy)
    QObject.connect(copyAction, SIGNAL("triggered()"), self.copySelectedResults)

    self.connect(self.btnExecute, SIGNAL("clicked()"), self.executeSql)
    self.connect(self.btnClear, SIGNAL("clicked()"), self.clearSql)
    self.connect(self.buttonBox.button(QDialogButtonBox.Close), SIGNAL("clicked()"), self.close)

    # hide the load query as layer if feature is not supported
    self._loadAsLayerAvailable = self.db.connector.hasCustomQuerySupport()
    self.loadAsLayerGroup.setVisible( self._loadAsLayerAvailable )
    if self._loadAsLayerAvailable:
      self.layerTypeWidget.hide() # show if load as raster is supported
      self.connect(self.loadLayerBtn, SIGNAL("clicked()"), self.loadSqlLayer)
      self.connect(self.getColumnsBtn, SIGNAL("clicked()"), self.fillColumnCombos)
      self.connect(self.loadAsLayerGroup, SIGNAL("toggled(bool)"), self.loadAsLayerToggled)
      self.loadAsLayerToggled(False)


  def closeEvent(self, e):
    """ save window state """
    settings = QSettings()
    settings.setValue("/DB_Manager/sqlWindow/geometry", QVariant(self.saveGeometry()))

    QDialog.closeEvent(self, e)

  def loadAsLayerToggled(self, checked):
    self.loadAsLayerGroup.setChecked( checked )
    self.loadAsLayerWidget.setVisible( checked )

  def getSql(self):
    # If the selection obtained from an editor spans a line break,
    # the text will contain a Unicode U+2029 paragraph separator
    # character instead of a newline \n character
    # (see https://qt-project.org/doc/qt-4.8/qtextcursor.html#selectedText)
    sql = self.editSql.textCursor().selectedText().replace(unichr(0x2029), "\n")
    if sql.isEmpty():
      sql = self.editSql.toPlainText()
    # try to sanitize query
    sql = sql.replace( QRegExp( ";\\s*$" ), "" )
    return sql

  def clearSql(self):
    self.editSql.clear()

  def executeSql(self):
    sql = self.getSql()
    if sql.isEmpty(): return

    QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))

    # delete the old model
    old_model = self.viewResult.model()
    self.viewResult.setModel(None)
    if old_model: old_model.deleteLater()

    self.uniqueCombo.clear()
    self.geomCombo.clear()

    try:
      # set the new model
      model = self.db.sqlResultModel( sql, self )
      self.viewResult.setModel( model )
      self.lblResult.setText("%d rows, %.1f seconds" % (model.affectedRows(), model.secs()))

    except BaseError, e:
      QApplication.restoreOverrideCursor()
      DlgDbError.showError(e, self)
      return

    cols = self.viewResult.model().columnNames()
    cols.sort()
    self.uniqueCombo.addItems( cols )
    self.geomCombo.addItems( cols )

    self.update()
    QApplication.restoreOverrideCursor()

  def loadSqlLayer(self):
    uniqueFieldName = self.uniqueCombo.currentText()
    geomFieldName = self.geomCombo.currentText()

    if geomFieldName.isEmpty() or uniqueFieldName.isEmpty():
      QMessageBox.warning(self, self.tr( "Sorry" ), self.tr( "You must fill the required fields: \ngeometry column - column with unique integer values" ) )
      return

    query = self.getSql()
    if query.isEmpty():
      return

    QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))

    from qgis.core import QgsMapLayer, QgsMapLayerRegistry
    layerType = QgsMapLayer.VectorLayer if self.vectorRadio.isChecked() else QgsMapLayer.RasterLayer

    # get a new layer name
    names = []
    for layer in QgsMapLayerRegistry.instance().mapLayers().values():
      names.append( layer.name() )

    layerName = self.layerNameEdit.text()
    if layerName.isEmpty():
      layerName = self.defaultLayerName
    newLayerName = layerName
    index = 1
    while newLayerName in names:
      index += 1
      newLayerName = u"%s_%d" % (layerName, index)

    # create the layer
    layer = self.db.toSqlLayer(query, geomFieldName, uniqueFieldName, newLayerName, layerType, self.avoidSelectById.isChecked())
    if layer.isValid():
      QgsMapLayerRegistry.instance().addMapLayers([layer], True)

    QApplication.restoreOverrideCursor()

  def fillColumnCombos(self):
    query = self.getSql()
    if query.isEmpty(): return

    QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))
    self.uniqueCombo.clear()
    self.geomCombo.clear()

    # get a new alias
    aliasIndex = 0
    while True:
      alias = "_%s__%d" % ("subQuery", aliasIndex)
      escaped = '\\b("?)' + QRegExp.escape(alias) + '\\1\\b'
      if not query.contains( QRegExp(escaped, Qt.CaseInsensitive) ):
        break
      aliasIndex += 1

    # get all the columns
    cols = []
    connector = self.db.connector
    sql = u"SELECT * FROM (%s\n) AS %s LIMIT 0" % ( unicode(query), connector.quoteId(alias) )

    c = None
    try:
      c = connector._execute(None, sql)
      cols = connector._get_cursor_columns(c)

    except BaseError as e:
      QApplication.restoreOverrideCursor()
      DlgDbError.showError(e, self)
      return

    finally:
      if c:
        c.close()
        del c

    cols.sort()
    self.uniqueCombo.addItems( cols )
    self.geomCombo.addItems( cols )

    QApplication.restoreOverrideCursor()

  def copySelectedResults(self):
    if len(self.viewResult.selectedIndexes()) <= 0:
      return
    model = self.viewResult.model()

    # convert to string using tab as separator
    text = model.headerToString( "\t" )
    for idx in self.viewResult.selectionModel().selectedRows():
      text += "\n" + model.rowToString( idx.row(), "\t" )

    QApplication.clipboard().setText( text, QClipboard.Selection )
    QApplication.clipboard().setText( text, QClipboard.Clipboard )