cartodb4/observatory-extension
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Improve fixtures generator

Browse Source
This commit is contained in:
Mario de Frutos 2018-04-12 18:01:25 +02:00
parent b891034146
commit e9dbc97772
1 changed files with 28 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
scripts/generate_fixtures.py
View File

@ -2,12 +2,18 @@ import os
import psycopg2 import psycopg2
import subprocess import subprocess
PGUSER = os.environ.get('PGUSER', 'postgres')
PGPASSWORD = os.environ.get('PGPASSWORD', '')
PGHOST=os.environ.get('PGHOST', 'localhost')
PGPORT=os.environ.get('PGPORT', '5432')
PGDATABASE=os.environ.get('PGDATABASE', 'postgres')
DB_CONN = psycopg2.connect('postgres://{user}:{password}@{host}:{port}/{database}'.format( DB_CONN = psycopg2.connect('postgres://{user}:{password}@{host}:{port}/{database}'.format(
user=os.environ.get('PGUSER', 'postgres'), user=PGUSER,
password=os.environ.get('PGPASSWORD', ''), password=PGPASSWORD,
host=os.environ.get('PGHOST', 'localhost'), host=PGHOST,
port=os.environ.get('PGPORT', '5432'), port=PGPORT,
database=os.environ.get('PGDATABASE', 'postgres'), database=PGDATABASE
)) ))
CURSOR = DB_CONN.cursor() CURSOR = DB_CONN.cursor()
@ -207,15 +213,12 @@ FIXTURES = [
('us.census.spielman_singleton_segments.X55', 'us.census.tiger.census_tract', '2010 - 2014'), ('us.census.spielman_singleton_segments.X55', 'us.census.tiger.census_tract', '2010 - 2014'),
('us.zillow.AllHomes_Zhvi', 'us.census.tiger.zcta5', '2014-01'), ('us.zillow.AllHomes_Zhvi', 'us.census.tiger.zcta5', '2014-01'),
('us.zillow.AllHomes_Zhvi', 'us.census.tiger.zcta5', '2016-06'), ('us.zillow.AllHomes_Zhvi', 'us.census.tiger.zcta5', '2016-06'),
('whosonfirst.wof_country_name', 'whosonfirst.wof_country_geom', '2016'), ('us.census.acs.B01003001', 'us.census.tiger.zcta5', '2010 - 2014'),
('us.census.acs.B01003001', 'us.census.tiger.zcta5_clipped', '2010 - 2014'), ('us.census.acs.B01003001', 'us.census.tiger.block_group', '2010 - 2014'),
('us.census.acs.B01003001', 'us.census.tiger.block_group_clipped', '2010 - 2014'), ('us.census.acs.B01003001', 'us.census.tiger.census_tract', '2010 - 2014'),
('us.census.acs.B01003001', 'us.census.tiger.census_tract_clipped', '2010 - 2014'), ('us.census.tiger.place_geoname', 'us.census.tiger.place_clipped', '2015'),
('us.census.tiger.fullname', 'us.census.tiger.pointlm_geom', '2016'), ('us.census.tiger.county_geoname', 'us.census.tiger.county_clipped', '2015'),
('us.census.tiger.fullname', 'us.census.tiger.prisecroads_geom', '2016'), ('us.census.tiger.block_group_geoname', 'us.census.tiger.block_group', '2015'),
('us.census.tiger.name', 'us.census.tiger.county', '2015'),
('us.census.tiger.name', 'us.census.tiger.county_clipped', '2015'),
('us.census.tiger.name', 'us.census.tiger.block_group', '2015'),
] ]
OUTFILE_PATH = os.path.join(os.path.dirname(__file__), '..', OUTFILE_PATH = os.path.join(os.path.dirname(__file__), '..',
@ -230,7 +233,8 @@ def dump(cols, tablename, where=''):
tablename=tablename, tablename=tablename,
)) ))
subprocess.check_call('pg_dump -x --section=pre-data -t observatory.{tablename} ' subprocess.check_call('PGPASSWORD={pgpassword} PGUSER={pguser} PGHOST={pghost} PGDATABASE={pgdb} '
'pg_dump -x --section=pre-data -t observatory.{tablename} '
' | sed "s:SET search_path.*::" ' ' | sed "s:SET search_path.*::" '
' | sed "s:CREATE TABLE :CREATE TABLE observatory.:" ' ' | sed "s:CREATE TABLE :CREATE TABLE observatory.:" '
' | sed "s:ALTER TABLE.*OWNER.*::" ' ' | sed "s:ALTER TABLE.*OWNER.*::" '
@ -238,19 +242,27 @@ def dump(cols, tablename, where=''):
' >> {outfile}'.format( ' >> {outfile}'.format(
tablename=tablename, tablename=tablename,
outfile=OUTFILE_PATH, outfile=OUTFILE_PATH,
pgpassword=PGPASSWORD,
pghost=PGHOST,
pgdb=PGDATABASE,
pguser=PGUSER
), shell=True) ), shell=True)
with open(OUTFILE_PATH, 'a') as outfile: with open(OUTFILE_PATH, 'a') as outfile:
outfile.write('COPY observatory."{}" FROM stdin WITH CSV HEADER;\n'.format(tablename)) outfile.write('COPY observatory."{}" FROM stdin WITH CSV HEADER;\n'.format(tablename))
subprocess.check_call(''' subprocess.check_call('''
psql -c "COPY (SELECT {cols} \ PGPASSWORD={pgpassword} psql -U {pguser} -d {pgdb} -h {pghost} -c "COPY (SELECT {cols} \
FROM observatory.{tablename} {where}) \ FROM observatory.{tablename} {where}) \
TO STDOUT WITH CSV HEADER" >> {outfile}'''.format( TO STDOUT WITH CSV HEADER" >> {outfile}'''.format(
cols=cols, cols=cols,
tablename=tablename, tablename=tablename,
where=where, where=where,
outfile=OUTFILE_PATH, outfile=OUTFILE_PATH,
pgpassword=PGPASSWORD,
pghost=PGHOST,
pgdb=PGDATABASE,
pguser=PGUSER
), shell=True) ), shell=True)
with open(OUTFILE_PATH, 'a') as outfile: with open(OUTFILE_PATH, 'a') as outfile: