From adaf91ec9d58073d42cd1e8c962b186f49ecec69 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Thu, 25 Oct 2018 14:18:14 -0700
Subject: [PATCH] ZIP codes based on Polygon

(cherry picked from commit 0dbaaf9d293d2f08e542491ba484a91596a6c4b6)
(cherry picked from commit 046736cbd429c01d673b6fc75fb9eb2e21b1077b)
(cherry picked from commit 49a33981c23dbefbe7420cd8828b78140b46e961)
(cherry picked from commit 9c10547f19b628e81cbcd6e1fbac86a70ea510be)
---
 .../components/controls/SpatialControl.jsx    | 25 +++++++
 superset/assets/src/explore/controls.jsx      |  1 +
 .../deckgl/layers/Polygon/Polygon.jsx         |  1 -
 superset/viz.py                               | 72 ++++++++++++++++++-
 4 files changed, 97 insertions(+), 2 deletions(-)

diff --git a/superset/assets/src/explore/components/controls/SpatialControl.jsx b/superset/assets/src/explore/components/controls/SpatialControl.jsx
index 592f12a780ca1..9d38fca36a151 100644
--- a/superset/assets/src/explore/components/controls/SpatialControl.jsx
+++ b/superset/assets/src/explore/components/controls/SpatialControl.jsx
@@ -32,6 +32,7 @@ const spatialTypes = {
   latlong: 'latlong',
   delimited: 'delimited',
   geohash: 'geohash',
+  zipcode: 'zipcode',
 };
 
 const propTypes = {
@@ -63,6 +64,7 @@ export default class SpatialControl extends React.Component {
       lonlatCol: v.lonlatCol || defaultCol,
       reverseCheckbox: v.reverseCheckbox || false,
       geohashCol: v.geohashCol || defaultCol,
+      zipcodeCol: v.zipcodeCol || defaultCol,
       value: null,
       errors: [],
     };
@@ -97,6 +99,12 @@ export default class SpatialControl extends React.Component {
       if (!value.geohashCol) {
         errors.push(errMsg);
       }
+    } else if (type === spatialTypes.zipcode) {
+      value.zipcodeCol = this.state.zipcodeCol;
+      value.reverseCheckbox = this.state.reverseCheckbox;
+      if (!value.zipcodeCol) {
+        errors.push(errMsg);
+      }
     }
     this.setState({ value, errors });
     this.props.onChange(value, errors);
@@ -120,6 +128,8 @@ export default class SpatialControl extends React.Component {
       return `${this.state.lonlatCol}`;
     } else if (this.state.type === spatialTypes.geohash) {
       return `${this.state.geohashCol}`;
+    } else if (this.state.type === spatialTypes.zipcode) {
+      return `${this.state.zipcodeCol}`;
     }
     return null;
   }
@@ -199,6 +209,21 @@ export default class SpatialControl extends React.Component {
               </Col>
             </Row>
           </PopoverSection>
+          <PopoverSection
+            title={t('ZIP code')}
+            isSelected={this.state.type === spatialTypes.zipcode}
+            onSelect={this.setType.bind(this, spatialTypes.zipcode)}
+          >
+            <Row>
+              <Col md={6}>
+                Column
+                {this.renderSelect('zipcodeCol', spatialTypes.zipcode)}
+              </Col>
+              <Col md={6}>
+                {this.renderReverseCheckbox()}
+              </Col>
+            </Row>
+          </PopoverSection>
           <div className="clearfix">
             <Button
               bsSize="small"
diff --git a/superset/assets/src/explore/controls.jsx b/superset/assets/src/explore/controls.jsx
index c8d8b7bed9699..de493341cea41 100644
--- a/superset/assets/src/explore/controls.jsx
+++ b/superset/assets/src/explore/controls.jsx
@@ -2183,6 +2183,7 @@ export const controls = {
       ['polyline', 'Polyline'],
       ['json', 'JSON'],
       ['geohash', 'geohash (square)'],
+      ['zipcode', 'ZIP code'],
     ],
   },
 
diff --git a/superset/assets/src/visualizations/deckgl/layers/Polygon/Polygon.jsx b/superset/assets/src/visualizations/deckgl/layers/Polygon/Polygon.jsx
index 8a3c2f4028af1..cffae2f1482bc 100644
--- a/superset/assets/src/visualizations/deckgl/layers/Polygon/Polygon.jsx
+++ b/superset/assets/src/visualizations/deckgl/layers/Polygon/Polygon.jsx
@@ -84,7 +84,6 @@ export function getLayer(formData, payload, setTooltip, selected, onSelect, filt
   return new PolygonLayer({
     id: `path-layer-${fd.slice_id}`,
     data,
-    pickable: true,
     filled: fd.filled,
     stroked: fd.stroked,
     getPolygon: d => d.polygon,
diff --git a/superset/viz.py b/superset/viz.py
index dc84630e4f8a5..5678b9b57c9d5 100644
--- a/superset/viz.py
+++ b/superset/viz.py
@@ -29,6 +29,7 @@
 from itertools import product
 import logging
 import math
+import os
 import pickle as pkl
 import re
 import traceback
@@ -45,6 +46,7 @@
 from pandas.tseries.frequencies import to_offset
 import polyline
 import simplejson as json
+import sqlalchemy
 
 from superset import app, cache, get_css_manifest_files
 from superset.exceptions import NullValueException, SpatialException
@@ -2110,6 +2112,8 @@ def get_spatial_columns(self, key):
             return [spatial.get('lonlatCol')]
         elif spatial.get('type') == 'geohash':
             return [spatial.get('geohashCol')]
+        elif spatial.get('type') == 'zipcode':
+            return [spatial.get('zipcodeCol')]
 
     @staticmethod
     def parse_coordinates(s):
@@ -2127,6 +2131,11 @@ def reverse_geohash_decode(geohash_code):
         lat, lng = geohash.decode(geohash_code)
         return (lng, lat)
 
+    @staticmethod
+    def reverse_zipcode_decode(zipcode):
+        raise NotImplementedError(
+            'No mapping from ZIP code to single latitude/longitude')
+
     @staticmethod
     def reverse_latlong(df, key):
         df[key] = [
@@ -2152,6 +2161,9 @@ def process_spatial_data_obj(self, key, df):
         elif spatial.get('type') == 'geohash':
             df[key] = df[spatial.get('geohashCol')].map(self.reverse_geohash_decode)
             del df[spatial.get('geohashCol')]
+        elif spatial.get('type') == 'zipcode':
+            df[key] = df[spatial.get('zipcodeCol')].map(self.reverse_zipcode_decode)
+            del df[spatial.get('zipcodeCol')]
 
         if spatial.get('reverseCheckbox'):
             self.reverse_latlong(df, key)
@@ -2342,6 +2354,57 @@ def geohash_to_json(geohash_code):
     ]
 
 
+def zipcode_deser(zipcodes):
+    geojson = zipcodes_to_json(zipcodes)
+
+    def deser(zipcode):
+        return geojson[str(zipcode)]['coordinates'][0]
+    return deser
+
+
+def zipcodes_to_json(zipcodes):
+    user = os.environ.get('CREDENTIALS_LYFTPG_USER', '')
+    password = os.environ.get('CREDENTIALS_LYFTPG_PASSWORD', '')
+    url = (
+        'postgresql+psycopg2://'
+        '{user}:{password}'
+        '@analytics-platform-vpc.c067nfzisc99.us-east-1.rds.amazonaws.com:5432'
+        '/platform'.format(user=user, password=password)
+    )
+
+    out = {}
+    missing = set()
+    for zipcode in zipcodes:
+        cache_key = 'zipcode_geojson_{}'.format(zipcode)
+        geojson = cache and cache.get(cache_key)
+        if geojson:
+            out[zipcode] = geojson
+        else:
+            missing.add(str(zipcode))
+
+    if not missing:
+        return out
+
+    # fetch missing geojson from lyftpg
+    in_clause = ', '.join(['%s'] * len(missing))
+    query = (
+        'SELECT zipcode, geojson FROM zip_codes WHERE zipcode IN ({0})'
+        .format(in_clause))
+    conn = sqlalchemy.create_engine(url, client_encoding='utf8')
+    results = conn.execute(query, tuple(missing)).fetchall()
+
+    for zipcode, geojson in results:
+        out[zipcode] = geojson
+        if cache and len(results) < 10000:  # avoid storing too much
+            cache_key = 'zipcode_geojson_{}'.format(zipcode)
+            try:
+                cache.set(cache_key, geojson, timeout=86400)
+            except Exception:
+                pass
+
+    return out
+
+
 class DeckPathViz(BaseDeckGLViz):
 
     """deck.gl's PathLayer"""
@@ -2354,6 +2417,7 @@ class DeckPathViz(BaseDeckGLViz):
         'json': json.loads,
         'polyline': polyline.decode,
         'geohash': geohash_to_json,
+        'zipcode': None,  # per request
     }
 
     def query_obj(self):
@@ -2379,13 +2443,19 @@ def get_properties(self, d):
         if fd.get('reverse_long_lat'):
             path = [(o[1], o[0]) for o in path]
         d[self.deck_viz_key] = path
-        if line_type != 'geohash':
+        if line_type not in ['geohash', 'zipcode']:
             del d[line_column]
         d['__timestamp'] = d.get(DTTM_ALIAS) or d.get('__time')
         return d
 
     def get_data(self, df):
         self.metric_label = utils.get_metric_name(self.metric)
+        fd = self.form_data
+        line_type = fd.get('line_type')
+        if line_type == 'zipcode':
+            zipcodes = df[fd['line_column']].unique()
+            self.deser_map['zipcode'] = zipcode_deser(zipcodes)
+
         return super(DeckPathViz, self).get_data(df)