Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
Use the new QgsAbstractDatabaseProviderConnection interface
  • Loading branch information
strk committed Jan 9, 2020
1 parent 3892fae commit f034317
Show file tree
Hide file tree
Showing 3 changed files with 81 additions and 73 deletions.
141 changes: 73 additions & 68 deletions python/plugins/db_manager/db_plugins/postgis/connector.py
Expand Up @@ -27,7 +27,14 @@
from functools import cmp_to_key

from qgis.PyQt.QtCore import QRegExp, QFile, QCoreApplication
from qgis.core import Qgis, QgsCredentials, QgsDataSourceUri, QgsCoordinateReferenceSystem
from qgis.core import (
Qgis,
QgsCredentials,
QgsDataSourceUri,
QgsProviderRegistry,
QgsAbstractDatabaseProviderConnection,
QgsProviderConnectionException,
)

from ..connector import DBConnector
from ..plugin import ConnectionError, DbError, Table
Expand All @@ -44,60 +51,69 @@ def classFactory():
return PostGisDBConnector


class CursorProxy():
def __init__(self, connection, sql):
self.connection = connection
self.sql = sql
self.result = None
self.closed = False
self.description = None
self._execute()

def _execute(self):
if self.result != None:
return
self.result = self.connection._executeSql(self.sql)
self.description = []
if len(self.result):
for i in range(len(self.result)):
self.description.append([
'column' + str(i), # name
str, # type_code
10, # display_size
10, # internal_size
0, # precision
None, # scale
True # null_ok
])

def fetchone(self):
self._execute()
return self.result[0]

def fetchall(self):
self._execute()
return self.result

def close(self):
self.result = None
self.closed = True


class PostGisDBConnector(DBConnector):

def __init__(self, uri):
DBConnector.__init__(self, uri)
def __init__(self, uri, dbplugin):
"""Creates a new PostgreSQL connector
self.host = uri.host() or os.environ.get('PGHOST')
self.port = uri.port() or os.environ.get('PGPORT')
:param uri: data source URI
:type uri: QgsDataSourceUri
:param dbplugin: the PostGisDBPlugin parent instance
:type dbplugin: PostGisDbPlugin
"""
DBConnector.__init__(self, uri)

username = uri.username() or os.environ.get('PGUSER')
password = uri.password() or os.environ.get('PGPASSWORD')
self.dbname = uri.database()

# Do not get db and user names from the env if service is used
if not uri.service():
if username is None:
username = os.environ.get('USER')
self.dbname = uri.database() or os.environ.get('PGDATABASE') or username
uri.setDatabase(self.dbname)
#self.connName = connName
#self.user = uri.username() or os.environ.get('USER')
#self.passwd = uri.password()
self.host = uri.host()

expandedConnInfo = self._connectionInfo()
try:
self.connection = psycopg2.connect(expandedConnInfo)
except self.connection_error_types() as e:
# get credentials if cached or asking to the user no more than 3 times
err = str(e)
uri = self.uri()
conninfo = uri.connectionInfo(False)

for i in range(3):
(ok, username, password) = QgsCredentials.instance().get(conninfo, username, password, err)
if not ok:
raise ConnectionError(QCoreApplication.translate('db_manager', 'Could not connect to database as user {user}').format(user=username))

if username:
uri.setUsername(username)

if password:
uri.setPassword(password)

newExpandedConnInfo = uri.connectionInfo(True)
try:
self.connection = psycopg2.connect(newExpandedConnInfo)
QgsCredentials.instance().put(conninfo, username, password)
except self.connection_error_types() as e:
if i == 2:
raise ConnectionError(e)
err = str(e)
finally:
# clear certs for each time trying to connect
self._clearSslTempCertsIfAny(newExpandedConnInfo)
finally:
# clear certs of the first connection try
self._clearSslTempCertsIfAny(expandedConnInfo)

self.connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
self.dbplugin = dbplugin
md = QgsProviderRegistry.instance().providerMetadata(dbplugin.providerName())
self.core_connection = md.findConnection(dbplugin.connectionName())
if self.core_connection is None:
self.core_connection = md.createConnection(dbplugin.connectionName(), uri.database())

c = self._execute(None, u"SELECT current_user,current_database()")
self.user, self.dbname = self._fetchone(c)
Expand Down Expand Up @@ -1034,33 +1050,22 @@ def execution_error_types(self):
def connection_error_types(self):
return psycopg2.InterfaceError, psycopg2.OperationalError

# moved into the parent class: DbConnector._execute()
# def _execute(self, cursor, sql):
# pass
def _execute(self, cursor, sql):
return CursorProxy(self, sql)

# moved into the parent class: DbConnector._execute_and_commit()
# def _execute_and_commit(self, sql):
# pass
def _executeSql(self, sql):
return self.core_connection.executeSql(sql)

# moved into the parent class: DbConnector._get_cursor()
# def _get_cursor(self, name=None):
# pass

# moved into the parent class: DbConnector._fetchall()
# def _fetchall(self, c):
# pass

# moved into the parent class: DbConnector._fetchone()
# def _fetchone(self, c):
# pass

# moved into the parent class: DbConnector._commit()
# def _commit(self):
# pass
def _commit(self):
pass

# moved into the parent class: DbConnector._rollback()
# def _rollback(self):
# pass
def _rollback(self):
pass

# moved into the parent class: DbConnector._get_cursor_columns()
# def _get_cursor_columns(self, c):
Expand Down
5 changes: 4 additions & 1 deletion python/plugins/db_manager/db_plugins/postgis/plugin.py
Expand Up @@ -112,7 +112,7 @@ def __init__(self, connection, uri):
Database.__init__(self, connection, uri)

def connectorsFactory(self, uri):
return PostGisDBConnector(uri)
return PostGisDBConnector(uri, self.connection())

def dataTablesFactory(self, row, db, schema=None):
return PGTable(row, db, schema)
Expand Down Expand Up @@ -184,6 +184,9 @@ def hasLowercaseFieldNamesOption(self):
def supportsComment(self):
return True

def executeSql(self, sql):
return self.connector._executeSql(sql)


class PGSchema(Schema):

Expand Down
Expand Up @@ -44,10 +44,10 @@ def load(db, mainwindow):
sql = u"""SELECT count(*)
FROM pg_class AS cls JOIN pg_namespace AS nsp ON nsp.oid = cls.relnamespace
WHERE cls.relname = 'topology' AND nsp.nspname = 'topology'"""
c = db.connector._get_cursor()
db.connector._execute(c, sql)
res = db.connector._fetchone(c)
if res is None or int(res[0]) <= 0:
#c = db.connector._get_cursor()
#db.connector._execute(c, sql)
res = db.executeSql(sql)
if res is None or len(res) < 1 or int(res[0][0]) <= 0:
return

# add the action to the DBManager menu
Expand Down

0 comments on commit f034317

Please sign in to comment.