Skip to content

Instantly share code, notes, and snippets.

@cynici
Last active July 12, 2016 09:45
Show Gist options
  • Save cynici/221383c28a21387a771c3080cb6e06ad to your computer and use it in GitHub Desktop.
Save cynici/221383c28a21387a771c3080cb6e06ad to your computer and use it in GitHub Desktop.
Example using Python coroutine to insert large dataset into PostgreSQL in batches
#
# CAVEAT: untested code intended to demonstrate idea only
#
from osgeo import ogr, osr
import psycopg2
# Read gory details http://www.dabeaz.com/coroutines/Coroutines.pdf
def coroutine(func):
def start(*args,**kwargs):
cr = func(*args,**kwargs)
cr.next()
return cr
return start
@coroutine
def batch_sql(conn=None, sql=None, batch_size=1000):
"""Use executemany() in batches of batch_size"""
def execute_sql(conn, sql, namedicts):
if len(namedicts) <= 0:
return
cur = conn.cursor()
cur.executemany(sql, namedicts)
conn.commit()
cur.close()
del namedicts[:]
pending = []
try:
while True:
namedict = (yield)
pending.append(namedict)
if len(pending) == batch_size:
execute_sql(conn, sql, pending)
except GeneratorExit:
execute_sql(conn, sql, pending)
def load_wgs84shape2database(connection_dict,table,shapefile,tileset,yyyymm):
connect_string = "user=%(user)s password=%(password)s host=%(host)s dbname=%(dbname)s" %connection_dict
connection = psycopg2.connect(connect_string)
inserter = batch_sql(
conn=connection,
sql="INSERT INTO "+table+"(tileset,yyyymm,burnday,the_geom) VALUES ('%(tileset)s',%(yyyymm)s,%(burnday)s,st_multi(st_GeomFromText('%(the_geom)s',4326)));"
)
# Add data to a spatial table
src_shapefile = ogr.Open(shapefile)
layer = src_shapefile.GetLayer(0)
for feature in layer:
geom = feature.GetGeometryRef()
wkt = geom.ExportToWkt()
params = {
'tileset':feature.GetField('tileset'),
'yyyymm':feature.GetField('YYYYMM'),
'burnday':feature.GetField('BurnDay'),
'the_geom':wkt
}
inserter.send(params)
inserter.close()
print 'Loading shapefile %s was successful' %shapefile
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment