-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.swift
310 lines (244 loc) · 11 KB
/
main.swift
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
import Foundation
import CoreGraphics
struct Extents {
let minX: Double
let minY: Double
let maxX: Double
let maxY: Double
}
enum FileFormat:String {
case png
case pdf
}
// Utility function to print progress
var lastPercent: Int = -1
func progress(_ current: Int, _ total: Int) {
var percent = Int(10 * round((Double(current) / Double(total)) * 10))
percent = percent >= 100 ? 100 : percent
if percent > lastPercent {
lastPercent = percent
print("\(percent)%")
}
}
// Utility function to run an SQL command that expects no returned values
func execSql(_ conn: OpaquePointer, _ sql: String) {
let res = PQexec(conn, sql)
if (PQresultStatus(res) != PGRES_COMMAND_OK) {
print("PQexec failed starting transaction: \(String(cString: PQerrorMessage(conn)))")
PQclear(res)
exit(1)
}
PQclear(res)
}
// Utility to generate extents from a WTK string
func getExtents(_ wkt: String) -> Extents {
let GEOS_HANDLE = GEOS_init_r()
let WKTReader = GEOSWKTReader_create_r(GEOS_HANDLE)
let GEOSGeom = GEOSWKTReader_read_r(GEOS_HANDLE, WKTReader, wkt)
var minX: Double = 0.0
var minY: Double = 0.0
var maxX: Double = 0.0
var maxY: Double = 0.0
GEOSGeom_getXMin_r(GEOS_HANDLE, GEOSGeom, &minX)
GEOSGeom_getYMin_r(GEOS_HANDLE, GEOSGeom, &minY)
GEOSGeom_getXMax_r(GEOS_HANDLE, GEOSGeom, &maxX)
GEOSGeom_getYMax_r(GEOS_HANDLE, GEOSGeom, &maxY)
GEOSWKTReader_destroy_r(GEOS_HANDLE, WKTReader)
GEOSGeom_destroy_r(GEOS_HANDLE, GEOSGeom)
return Extents(minX: minX, minY: minY, maxX: maxX, maxY: maxY)
}
// Utility to generate extents from an SQL query
func getExtents(_ conn: OpaquePointer, _ query: String) -> Extents {
print("Generating the extents")
let res = PQexec(conn, "with foo(geom) as ( \(query) ) select st_astext(st_extent(geom)) from foo")
if PQresultStatus(res) != PGRES_TUPLES_OK {
print("PQexec failed while getting extents: \(String(cString: PQerrorMessage(conn)))")
PQclear(res)
exit(1)
}
let data = String.init(cString: PQgetvalue(res, 0, 0))
PQclear(res)
return getExtents(data)
}
// Commandline argument utility
func getRequiredOpt(option: String, errorMessage: String) -> String {
guard let index = CommandLine.arguments.index(of: option), CommandLine.arguments.indices.contains(index + 1) else {
print(errorMessage)
exit(1)
}
return CommandLine.arguments[index+1]
}
// Commandline argument utility
func getOpt(option: String) -> String? {
guard let index = CommandLine.arguments.index(of: option), CommandLine.arguments.indices.contains(index + 1) else {
return nil
}
return CommandLine.arguments[index+1]
}
// Process commandline arguments
if CommandLine.argc <= 1 || CommandLine.arguments.index(of: "-help") != nil || CommandLine.arguments.index(of: "-h") != nil {
print("""
-f: Specify the filename to save the image
-pg: Set the database connection information in the format \"host=hostname.rds.amazonaws.com user=troutspotting dbname=xxx password=yyyy\"
-query: An SQL querty that returns the geometry in WKT format along with a numerical value representing the line width
-extents: Optional. Specify the extents in WKT format, otherwise this is dynamically calculated
-imageWidth: Defaults to 5000; height is auto-calculated to maintain the aspect ratio of the image
-scale: Defaults to 1.0
-totalrows: Optional, ignored if progress is not enabled. When set, it prevents a dynamic calculation of total rows of data to render, improving performance when the total row count calculation is time consuming.
Example:
./gen_image -pg \"host=hostname.rds.amazonaws.com user=username dbname=xxx password=yyyy\" -query \"with extents(geom) as (select st_geomfromtext('POLYGON((-129 23,-129 51,-62 51,-62 23,-129 23))')) select st_astext(shape) from streams join extents on shape && geom\" -extents \"POLYGON((-129 23,-129 51,-62 51,-62 23,-129 23))\" -f output.png
""")
exit(1)
}
// Get commandline arguments
let pq = getRequiredOpt(option: "-pg", errorMessage: "Need to provide -pg parameter")
let query = getRequiredOpt(option: "-query", errorMessage: "Need to provide a query param using -query")
let filePath = getRequiredOpt(option: "-f", errorMessage: "Need to specify an output file path using -f")
let url = NSURL(fileURLWithPath: filePath)
guard let fileSuffix = url.pathExtension, let format = FileFormat(rawValue: fileSuffix) else {
print("Invalid file format: \(filePath)")
exit(1)
}
let imageWidth = getOpt(option: "-width") == nil ? 5000 : Int(getOpt(option: "-width")!)!
let extentsProvided = getOpt(option: "-extents")
let printProgress = CommandLine.arguments.contains("-progress")
var totalRows = getOpt(option: "-totalrows") == nil ? 0 : Int(getOpt(option: "-totalrows")!)!
var scale = getOpt(option: "-scale") == nil ? 1.0 : Double(getOpt(option: "-scale")!)!
let increment = 1000 // the batch size of results fetched via a cursor
print("Connecting to database using: \"\(pq)\"")
print("Output format: \(format)")
print("Query: \"\(query)\"")
print("Image Width: \(imageWidth)")
print("File path: \(filePath)")
print("Image scale: \(scale)")
// Connect to the database
guard let conn = PQconnectdb(pq), PQstatus(conn) == CONNECTION_OK else {
print("Unable to connect to the database using \(pq)")
exit(1)
}
// If no estimate row count was provided using -totalrows, dynamically calculate it
// This may be a slow query, so its generally best to estimate the total rows
// for large data sets
if printProgress && totalRows == 0 {
let res = PQexec(conn, "with data as (\(query)) select count(*) from data")
if PQresultStatus(res) != PGRES_TUPLES_OK {
print("PQexec failed declaring cursor: \(String(cString: PQerrorMessage(conn)))")
PQclear(res)
exit(1)
}
totalRows = Int(String(cString: PQgetvalue(res, 0, 0)))!
print("Total rows calculated as \(totalRows)")
PQclear(res)
}
// Get the extents
let extents = extentsProvided == nil ? getExtents(conn, query) : getExtents(extentsProvided!)
let lineWidth = extents.maxX - extents.minX
let lineHeight = extents.maxY - extents.minY
let aspectRatio = lineWidth / lineHeight
let imageHeight = Int(Double(imageWidth) / aspectRatio)
print("Extents: \(extents)")
print("lineWidth: \(lineWidth) lineHeight: \(lineHeight) aspectRatio: \(aspectRatio) imageWidth: \(imageWidth) imageHeight: \(imageHeight)")
// Begin a transaction; this is required for a cursor
execSql(conn, "BEGIN")
// Using a cursor keeps memory usage to a minimim
execSql(conn, "DECLARE mycursor CURSOR FOR \(query)")
// Initialize the GEOS library. This is used to convert the PostGIS geometry into something
// that can be used in code
var GEOS_HANDLE = GEOS_init_r()
let WKTReader = GEOSWKTReader_create_r(GEOS_HANDLE)
// These ratios are used during the conversion of lat/lon to x/y screen coordinates
let imageToLineWidthRatio = Double(imageWidth) / lineWidth
let imageToLineHeightRatio = Double(imageHeight) / lineHeight
// Init CoreGraphics
let colorSpace = CGColorSpaceCreateDeviceRGB()
let size = CGSize(width: imageWidth, height: imageHeight)
var context: CGContext?
if format == .png {
let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.premultipliedLast.rawValue)
context = CGContext.init(data: nil, width: Int(size.width), height: Int(size.height), bitsPerComponent: 8, bytesPerRow: 0, space: colorSpace, bitmapInfo: bitmapInfo.rawValue)
} else {
var mediaBox: CGRect = CGRect(x: 0, y: 0, width: imageWidth, height: imageHeight)
context = CGContext(url, mediaBox: &mediaBox, nil)
let boxData = NSData(bytes: &mediaBox, length: MemoryLayout.size(ofValue: mediaBox))
let pageInfo = [ kCGPDFContextMediaBox as String: boxData ]
context?.beginPDFPage(pageInfo as NSDictionary)
}
context?.setShouldAntialias(true)
//context?.setFillColor(.black)
//context?.setStrokeColor(.white)
context?.setFillColor(.white)
let darkBlue = CGColor(red: 0.495, green: 0.615, blue: 0.905, alpha: 1.0)
let lightBlue = CGColor(red: 0.125, green: 0.278, blue: 0.615, alpha: 1.0)
context?.setStrokeColor(lightBlue)
context?.fill(CGRect(origin: CGPoint(x: 0, y: 0), size: size))
context?.setLineCap(.round)
context?.setLineJoin(.round)
var hasData = true
var count = 0
print("Start rendering...")
while hasData {
let res = PQexec(conn, "fetch forward \(increment) from mycursor")
if PQresultStatus(res) != PGRES_TUPLES_OK {
print("PQexec failed declaring cursor: \(String(cString: PQerrorMessage(conn)))")
PQclear(res)
exit(1)
}
let numRows = PQntuples(res)
hasData = numRows > 0
if !hasData {
print("Finished retrieving data")
}
count += increment
if printProgress {
progress(count, totalRows)
}
for row in 0..<numRows {
let data = String(cString: PQgetvalue(res, row, 0))
let length:CGFloat = CGFloat(Float(String(cString: PQgetvalue(res, row, 1)))!)
let GEOSGeom = GEOSWKTReader_read_r(GEOS_HANDLE, WKTReader, data)
let geometryTypeId = GEOSGeomTypeId_r(GEOS_HANDLE, GEOSGeom)
if geometryTypeId != GEOS_LINESTRING.rawValue {
print("Geometry must be a linestring, skipping")
continue
}
let sequence = GEOSGeom_getCoordSeq_r(GEOS_HANDLE, GEOSGeom)
var tmpNumCoordinates: UInt32 = 0
GEOSCoordSeq_getSize_r(GEOS_HANDLE, sequence, &tmpNumCoordinates)
let numCoordinates = Int(tmpNumCoordinates)
var xy = Array(repeating: Array(repeating: Double(0.0), count: 2), count: numCoordinates)
for i in 0 ..< numCoordinates {
var longitude: Double = 0
var latitude: Double = 0
GEOSCoordSeq_getX_r(GEOS_HANDLE, sequence, UInt32(i), &longitude)
GEOSCoordSeq_getY_r(GEOS_HANDLE, sequence, UInt32(i), &latitude)
xy[i][0] = (extents.minX - longitude) * -1 * imageToLineWidthRatio * scale
xy[i][1] = (latitude - extents.minY) * imageToLineHeightRatio * scale
}
// Drawing sequence
context?.beginPath()
context?.setLineWidth(length)
context?.setShadow(offset: CGSize(width: 0, height: 0), blur: length*3, color: darkBlue)
for i in 0 ..< numCoordinates {
let p = CGPoint(x: xy[i][0], y: xy[i][1])
i == 0 ? context?.move(to: p) : context?.addLine(to: p)
}
context?.strokePath()
GEOSGeom_destroy_r(GEOS_HANDLE, GEOSGeom)
}
PQclear(res)
}
print("Generating output to \(filePath)")
switch format {
case .png:
let image = context?.makeImage()
let destination = CGImageDestinationCreateWithURL(url, kUTTypePNG, 1, nil)
CGImageDestinationAddImage(destination!, image!, nil)
CGImageDestinationFinalize(destination!)
case .pdf:
context?.endPDFPage()
context?.closePDF()
}
execSql(conn, "CLOSE mycursor")
execSql(conn, "END")
PQfinish(conn)
GEOSWKTReader_destroy_r(GEOS_HANDLE, WKTReader)