import_us_buildings.py

import os
import sys
import argparse
try:
    from urllib2 import urlopen
except ImportError:
    from urllib.request import urlopen

try:
    from StringIO import BytesIO
except ImportError:
    from io import BytesIO

from carto.auth import APIKeyAuthClient
from carto.sql import SQLClient, CopySQLClient

from ijson import items
from shapely.geometry import shape
from zipfile import ZipFile

parser = argparse.ArgumentParser(description=(
    'US buildings to CARTO map.'
))

parser.add_argument('--base_url', type=str, dest='CARTO_BASE_URL',
                    default=os.environ.get('CARTO_API_URL', ''),
                    help=('Set the base URL. For example:'
                          ' https://username.carto.com/'
                          ' (defaults to env variable CARTO_API_URL)'))

parser.add_argument('--api_key', dest='CARTO_API_KEY',
                    default=os.environ.get('CARTO_API_KEY', ''),
                    help=('Api key of the account'
                          ' (defaults to env variable CARTO_API_KEY)'))

# see available datasets names here -> https://github.com/Microsoft/USBuildingFootprints
parser.add_argument('--datasets', dest='DATASETS',
                    default=os.environ.get('DATASETS', ''),
                    help=('Name of the datasets to download'
                          ' (defaults to None)'))

args = parser.parse_args()

if not args.CARTO_BASE_URL or not args.CARTO_API_KEY or not args.DATASETS:
    sys.exit(parser.print_usage())

auth_client = APIKeyAuthClient(args.CARTO_BASE_URL, args.CARTO_API_KEY)
sql_client = SQLClient(auth_client)
copy_client = CopySQLClient(auth_client)
datasets = args.DATASETS.split(',')

sql_client.send("""CREATE TABLE IF NOT EXISTS buildings (
  the_geom geometry(Geometry,4326)
)""")

def rows(datasets):
    for dataset in datasets:
        resp = urlopen('https://usbuildingdata.blob.core.windows.net/usbuildings-v1-1/{dataset}.zip'.format(dataset=dataset))
        zipfile = ZipFile(BytesIO(resp.read()))
        files = zipfile.namelist()
        objects = items(zipfile.open(files[0]), 'features.item')
        features = (o for o in objects if o['type'] == 'Feature')
        for feature in features:
            g2 = shape(feature['geometry'])
            row = u'SRID=4326;{polygon}\n'.format(polygon=g2.wkt)
            yield row.encode()

try:
    result = copy_client.copyfrom(
        ('COPY buildings(the_geom)'
         ' FROM stdin WITH (FORMAT csv, QUOTE \'"\', DELIMITER \'|\')'),
        rows(datasets))

    sql_client.send(
        "SELECT CDB_CartodbfyTable(current_schema, 'buildings')")
except Exception as e:
    print(e)