Use QgsVectorLayer to retrive field names from queries

qgis · Jan 10, 2020 · bf4ed9a · elpaso · Jan 10, 2020 · bf4ed9a
1 parent 8e5ff28
commit bf4ed9a
Showing 1 changed file with 24 additions and 11 deletions.
diff --git a/python/plugins/db_manager/db_plugins/postgis/connector.py b/python/plugins/db_manager/db_plugins/postgis/connector.py
@@ -35,6 +35,7 @@
 from qgis.core import (
     Qgis,
     QgsCredentials,
+    QgsVectorLayer,
     QgsDataSourceUri,
     QgsProviderRegistry,
     QgsProviderConnectionException,
@@ -97,18 +98,30 @@ def _execute(self, sql=None):
         self.result = self._toStrResultSet(self.connection.executeSql(self.sql))
         self._debug("execute returned " + str(len(self.result)) + " rows")
         self.cursor = 0
+
         self.description = []
-        if len(self.result):
-            for i in range(len(self.result[0])):
-                self.description.append([
-                    'column' + str(i),                         # name
-                    str,                                      # type_code
-                    10,                                       # display_size
-                    10,                                       # internal_size
-                    0,                                        # precision
-                    None,                                     # scale
-                    True                                      # null_ok
-                ])
+
+        uri = QgsDataSourceUri(self.connection.uri())
+
+        # TODO: make this part provider-agnostic
+        uri.setTable('(SELECT row_number() OVER () AS __rid__, * FROM (' + self.sql + ') as foo)')
+        uri.setKeyColumn('__rid__')
+        # TODO: fetch provider name from connection (QgsAbstractConnectionProvider)
+        # TODO: re-use the VectorLayer for fetching rows in batch mode
+        vl = QgsVectorLayer(uri.uri(False), 'dbmanager_cursor', 'postgres')
+
+        fields = vl.fields()
+        for i in range(1, len(fields)): # skip first field (__rid__)
+            f = fields[i]
+            self.description.append([
+                f.name(),                         # name
+                str,                              # type_code
+                10,                               # display_size
+                10,                               # internal_size
+                0,                                # precision
+                None,                             # scale
+                True                              # null_ok
+            ])
         self._debug("execute returned " + str(len(self.description)) + " cols")
 
     def fetchone(self):