Merge pull request #29 from reidab/pg-connection-string

Support PostgreSQL connections for source data

ogr2osm.py

@@ -46,6 +46,7 @@
 import os
 import optparse
 import logging as l
+import re
 l.basicConfig(level=l.DEBUG, format="%(message)s")
 
 from osgeo import ogr
@@ -85,7 +86,10 @@
 
 
 # Setup program usage
-usage = "usage: %prog SRCFILE"
+usage = """%prog SRCFILE
+
+SRCFILE can be a file path or a org PostgreSQL connection string such as:
+"PG:dbname=pdx_bldgs user=emma host=localhost" (including the quotes)"""
 parser = optparse.OptionParser(usage=usage)
 parser.add_option("-t", "--translation", dest="translationMethod",
                   metavar="TRANSLATION",
@@ -143,6 +147,9 @@
 parser.add_option("--add-timestamp", dest="addTimestamp", action="store_true",
                     help=optparse.SUPPRESS_HELP)
 
+parser.add_option("--sql", dest="sqlQuery", type=str, default=None,
+                     help="SQL query to execute on a PostgreSQL source")
+
 parser.set_defaults(sourceEPSG=None, sourcePROJ4=None, verbose=False,
                     debugTags=False,
                     translationMethod=None, outputFile=None,
@@ -169,15 +176,23 @@
 
 # Input and output file
 # if no output file given, use the basename of the source but with .osm
-sourceFile = args[0]
+source = args[0]
+sourceIsDatabase = bool(re.match('^PG:', source))
+
 if options.outputFile is not None:
     options.outputFile = os.path.realpath(options.outputFile)
+elif sourceIsDatabase:
+    parser.error("ERROR: An output file must be explicitly specified when using a database source")
 else:
-    (base, ext) = os.path.splitext(os.path.basename(sourceFile))
+    (base, ext) = os.path.splitext(os.path.basename(source))
     options.outputFile = os.path.join(os.getcwd(), base + ".osm")
+
+if options.sqlQuery and not sourceIsDatabase:
+    parser.error("ERROR: You must use a database source when specifying a query with --sql")
+
 if not options.forceOverwrite and os.path.exists(options.outputFile):
     parser.error("ERROR: output file '%s' exists" % (options.outputFile))
-l.info("Preparing to convert file '%s' to '%s'." % (sourceFile, options.outputFile))
+l.info("Preparing to convert '%s' to '%s'." % (source, options.outputFile))
 
 # Projection
 if not options.sourcePROJ4 and not options.sourceEPSG:
@@ -249,6 +264,20 @@
 if options.positiveID:
     Geometry.elementIdCounterIncr = 1 # default is -1
 
+def openData(source):
+    if re.match('^PG:', source):
+        return openDatabaseSource(source)
+    else:
+        return getFileData(source)
+
+def openDatabaseSource(source):
+    dataSource = ogr.Open(source, 0)  # 0 means read-only
+    if dataSource is None:
+        l.error('OGR failed to open connection to' + source)
+        sys.exit(1)
+    else:
+        return dataSource
+
 def getFileData(filename):
     ogr_accessmethods = [ "/vsicurl/", "/vsicurl_streaming/", "/vsisubfile/",
         "/vsistdin/" ]
@@ -289,10 +318,15 @@ def getFileData(filename):
 def parseData(dataSource):
     l.debug("Parsing data")
     global translations
-    for i in range(dataSource.GetLayerCount()):
-        layer = dataSource.GetLayer(i)
+    if options.sqlQuery:
+        layer = dataSource.ExecuteSQL(options.sqlQuery)
         layer.ResetReading()
         parseLayer(translations.filterLayer(layer))
+    else:
+        for i in range(dataSource.GetLayerCount()):
+            layer = dataSource.GetLayer(i)
+            layer.ResetReading()
+            parseLayer(translations.filterLayer(layer))
 
 def getTransform(layer):
     global options
@@ -608,7 +642,7 @@ def output():
 
 
 # Main flow
-data = getFileData(sourceFile)
+data = openData(source)
 parseData(data)
 mergePoints()
 mergeWayPoints()

testfiles/basic.t

@@ -8,6 +8,9 @@ usage:
   running with lxml.etree
   Usage: ogr2osm.py SRCFILE
 
+  SRCFILE can be a file path or a org PostgreSQL connection string such as:
+  "PG:dbname=pdx_bldgs user=emma host=localhost" (including the quotes)
+
   Options:
     -h, --help            show this help message and exit
     -t TRANSLATION, --translation=TRANSLATION
@@ -38,12 +41,13 @@ usage:
                           Defaults to 0.
     --idfile=IDFILE       Read ID to start counting from from a file.
     --saveid=SAVEID       Save last ID after execution to a file.
+    --sql=SQLQUERY        SQL query to execute on a PostgreSQL source
 
 test1:
   $ rm -f test1.osm
   $ ogr2osm $TESTDIR/shapefiles/test1.shp
   running with lxml.etree
-  Preparing to convert file .* (re)
+  Preparing to convert .* (re)
   Will try to detect projection from source metadata, or fall back to EPSG:4326
   Using default translations
   Using default filterLayer
@@ -78,13 +82,17 @@ duplicatefile:
   running with lxml.etree
   Usage: ogr2osm.py SRCFILE
 
+  SRCFILE can be a file path or a org PostgreSQL connection string such as:
+  "PG:dbname=pdx_bldgs user=emma host=localhost" (including the quotes)
+
   ogr2osm.py: error: ERROR: output file .*test1.osm' exists (re)
   [2]
 
+
 force:
   $ ogr2osm -f $TESTDIR/shapefiles/test1.shp
   running with lxml.etree
-  Preparing to convert file .* (re)
+  Preparing to convert .* (re)
   Will try to detect projection from source metadata, or fall back to EPSG:4326
   Using default translations
   Using default filterLayer
@@ -117,7 +125,7 @@ force:
 nomemorycopy:
   $ ogr2osm -f --no-memory-copy $TESTDIR/shapefiles/test1.shp
   running with lxml.etree
-  Preparing to convert file .* (re)
+  Preparing to convert .* (re)
   Will try to detect projection from source metadata, or fall back to EPSG:4326
   Using default translations
   Using default filterLayer
@@ -150,7 +158,7 @@ nomemorycopy:
 positiveid:
   $ ogr2osm -f --positive-id $TESTDIR/shapefiles/test1.shp
   running with lxml.etree
-  Preparing to convert file .* (re)
+  Preparing to convert .* (re)
   Will try to detect projection from source metadata, or fall back to EPSG:4326
   Using default translations
   Using default filterLayer
@@ -183,7 +191,7 @@ positiveid:
 version:
   $ ogr2osm -f --add-version $TESTDIR/shapefiles/test1.shp
   running with lxml.etree
-  Preparing to convert file .* (re)
+  Preparing to convert .* (re)
   Will try to detect projection from source metadata, or fall back to EPSG:4326
   Using default translations
   Using default filterLayer
@@ -216,7 +224,7 @@ version:
 timestamp:
   $ ogr2osm -f --add-timestamp $TESTDIR/shapefiles/test1.shp
   running with lxml.etree
-  Preparing to convert file .* (re)
+  Preparing to convert .* (re)
   Will try to detect projection from source metadata, or fall back to EPSG:4326
   Using default translations
   Using default filterLayer
@@ -248,7 +256,7 @@ timestamp:
 utf8:
   $ ogr2osm -f $TESTDIR/shapefiles/sp_usinas.shp
   running with lxml.etree
-  Preparing to convert file .* (re)
+  Preparing to convert .* (re)
   Will try to detect projection from source metadata, or fall back to EPSG:4326
   Using default translations
   Using default filterLayer
@@ -273,7 +281,7 @@ utf8:
 japanese:
   $ ogr2osm --encoding shift_jis -f $TESTDIR/shapefiles/japanese.shp
   running with lxml.etree
-  Preparing to convert file .* (re)
+  Preparing to convert .* (re)
   Will try to detect projection from source metadata, or fall back to EPSG:4326
   Using default translations
   Using default filterLayer
@@ -293,7 +301,7 @@ japanese:
 duplicatewaynodes:
   $ ogr2osm -f $TESTDIR/duplicate-way-nodes.gml
   running with lxml.etree
-  Preparing to convert file .* (re)
+  Preparing to convert .* (re)
   Will try to detect projection from source metadata, or fall back to EPSG:4326
   Using default translations
   Using default filterLayer
@@ -361,3 +369,27 @@ duplicatewaynodes:
   Merging duplicate points in ways
   Outputting XML
   $ xmllint --format duplicate-way-nodes.osm | diff -uNr - $TESTDIR/duplicate-way-nodes.xml
+
+require_output_file_when_using_db_source:
+
+  $ ogr2osm "PG:dbname=test"
+  running with lxml.etree
+  Usage: ogr2osm.py SRCFILE
+  
+  SRCFILE can be a file path or a org PostgreSQL connection string such as:
+  "PG:dbname=pdx_bldgs user=emma host=localhost" (including the quotes)
+  
+  ogr2osm.py: error: ERROR: An output file must be explicitly specified when using a database source
+  [2]
+
+require_db_source_for_sql_query:
+
+  $ ogr2osm $TESTDIR/shapefiles/test1.shp --sql="SELECT * FROM wombats"
+  running with lxml.etree
+  Usage: ogr2osm.py SRCFILE
+  
+  SRCFILE can be a file path or a org PostgreSQL connection string such as:
+  "PG:dbname=pdx_bldgs user=emma host=localhost" (including the quotes)
+  
+  ogr2osm.py: error: ERROR: You must use a database source when specifying a query with --sql
+  [2]