Fix extraction of ogr LayerName from multi-layer dataset URIs

Adds supports for "layerid" when present. Drop special handling for "table=" portions found in URI, making the code more generic. Includes testcase. Fixes #15698 - import geodatabase to postgis via processing
qgis · Oct 15, 2016 · 6c53641 · gacarrillor · Oct 15, 2016 · strk
1 parent 60b4b4d
commit 6c53641
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 14 deletions.
diff --git a/python/plugins/processing/tests/ToolsTest.py b/python/plugins/processing/tests/ToolsTest.py
@@ -31,11 +31,67 @@
 from qgis.core import (QgsVectorLayer, QgsFeatureRequest)
 from processing.core.ProcessingConfig import ProcessingConfig
 
+import os.path
+import errno
+import shutil
+
+dataFolder = os.path.join(os.path.dirname(__file__), '../../../../tests/testdata/')
+tmpBaseFolder = os.path.join(os.sep, 'tmp', 'qgis_test', str(os.getpid()))
+
+
+def mkDirP(path):
+    try:
+        os.makedirs(path)
+    except OSError as exc:
+        if exc.errno == errno.EEXIST and os.path.isdir(path):
+            pass
+        else:
+            raise
+
 start_app()
 
 
 class VectorTest(unittest.TestCase):
 
+    @classmethod
+    def setUpClass(cls):
+        mkDirP(tmpBaseFolder)
+
+    @classmethod
+    def tearDownClass(cls):
+        shutil.rmtree(tmpBaseFolder)
+        pass
+
+    # See http://hub.qgis.org/issues/15698
+    def test_ogrLayerName(self):
+        tmpdir = os.path.join(tmpBaseFolder, 'ogrLayerName')
+        os.mkdir(tmpdir)
+
+        def linkTestfile(f, t):
+            os.link(os.path.join(dataFolder, f), os.path.join(tmpdir, t))
+
+        linkTestfile('geom_data.csv', 'a.csv')
+        name = vector.ogrLayerName(tmpdir)
+        self.assertEqual(name, 'a')
+
+        linkTestfile('wkt_data.csv', 'b.csv')
+        name = vector.ogrLayerName(tmpdir + '|layerid=0')
+        self.assertEqual(name, 'a')
+        name = vector.ogrLayerName(tmpdir + '|layerid=1')
+        self.assertEqual(name, 'b')
+
+        name = vector.ogrLayerName(tmpdir + '|layerid=2')
+        self.assertEqual(name, 'invalid-layerid')
+
+        name = vector.ogrLayerName(tmpdir + '|layername=f')
+        self.assertEqual(name, 'f') # layername takes precedence
+
+        name = vector.ogrLayerName(tmpdir + '|layerid=0|layername=f2')
+        self.assertEqual(name, 'f2') # layername takes precedence
+
+        name = vector.ogrLayerName(tmpdir + '|layername=f2|layerid=0')
+        self.assertEqual(name, 'f2') # layername takes precedence
+
     def testFeatures(self):
         ProcessingConfig.initialize()
 

diff --git a/python/plugins/processing/tools/vector.py b/python/plugins/processing/tools/vector.py
@@ -41,6 +41,8 @@
 
 import psycopg2
 
+from osgeo import ogr
+
 from qgis.PyQt.QtCore import QVariant, QSettings
 from qgis.core import (Qgis, QgsFields, QgsField, QgsGeometry, QgsRectangle, QgsWkbTypes,
                        QgsSpatialIndex, QgsMapLayerRegistry, QgsMapLayer, QgsVectorLayer,
@@ -531,20 +533,26 @@ def ogrConnectionString(uri):
 
 
 def ogrLayerName(uri):
-    if 'host' in uri:
-        regex = re.compile('(table=")(.+?)(\.)(.+?)"')
-        r = regex.search(uri)
-        return '"' + r.groups()[1] + '.' + r.groups()[3] + '"'
-    elif 'dbname' in uri:
-        regex = re.compile('(table=")(.+?)"')
-        r = regex.search(uri)
-        return r.groups()[1]
-    elif 'layername' in uri:
-        regex = re.compile('(layername=)(.*)')
-        r = regex.search(uri)
-        return r.groups()[1]
-    else:
-        return os.path.basename(os.path.splitext(uri)[0])
+    fields = uri.split('|')
+    ogruri = fields[0]
+    fields = fields[1:]
+    layerid = 0
+    for f in fields:
+        if f.startswith('layername='):
+            # Name encoded in uri, nothing more needed
+            return f.split('=')[1]
+        if f.startswith('layerid='):
+            layerid = int(f.split('=')[1])
+            # Last layerid= takes precedence, to allow of layername to
+            # take precedence
+    ds = ogr.Open(ogruri)
+    if not ds:
+        return "invalid-uri"
+    ly = ds.GetLayer(layerid)
+    if not ly:
+        return "invalid-layerid"
+    name = ly.GetName()
+    return name
 
 
 class VectorWriter(object):