Added avg as a column for both SqlaTable and Druid

apache · Oct 25, 2016 · e3f64ba · e3f64ba
1 parent ba2a2e5
commit e3f64ba
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 2 deletions.
diff --git a/caravel/migrations/versions/ad4d656d92bc_add_avg_metric.py b/caravel/migrations/versions/ad4d656d92bc_add_avg_metric.py
@@ -0,0 +1,25 @@
+"""Add avg() to default metrics
+
+Revision ID: ad4d656d92bc
+Revises: b46fa1b0b39e
+Create Date: 2016-10-25 10:16:39.871078
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = 'ad4d656d92bc'
+down_revision = '7e3ddad2a00b'
+
+from alembic import op
+import sqlalchemy as sa
+
+
+def upgrade():
+    op.add_column('columns', sa.Column('avg', sa.Boolean(), nullable=True))
+    op.add_column('table_columns', sa.Column('avg', sa.Boolean(), nullable=True))
+
+def downgrade():
+  with op.batch_alter_table('columns') as batch_op:
+    batch_op.drop_column('avg')
+  with op.batch_alter_table('table_columns') as batch_op:
+    batch_op.drop_column('avg')
diff --git a/caravel/models.py b/caravel/models.py
@@ -1169,6 +1169,7 @@ def fetch_metadata(self):
                 dbcol.groupby = dbcol.is_string
                 dbcol.filterable = dbcol.is_string
                 dbcol.sum = dbcol.isnum
+                dbcol.avg = dbcol.isnum
                 dbcol.is_dttm = dbcol.is_time
 
             db.session.merge(self)
@@ -1186,7 +1187,7 @@ def fetch_metadata(self):
                     metric_type='sum',
                     expression="SUM({})".format(quoted)
                 ))
-            if dbcol.sum:
+            if dbcol.avg:
                 metrics.append(M(
                     metric_name='avg__' + dbcol.column_name,
                     verbose_name='avg__' + dbcol.column_name,
@@ -1373,6 +1374,7 @@ class TableColumn(Model, AuditMixinNullable, ImportMixin):
     groupby = Column(Boolean, default=False)
     count_distinct = Column(Boolean, default=False)
     sum = Column(Boolean, default=False)
+    avg = Column(Boolean, default=False)
     max = Column(Boolean, default=False)
     min = Column(Boolean, default=False)
     filterable = Column(Boolean, default=False)
@@ -1386,7 +1388,7 @@ class TableColumn(Model, AuditMixinNullable, ImportMixin):
     str_types = ('VARCHAR', 'STRING', 'CHAR')
     export_fields = (
         'table_id', 'column_name', 'verbose_name', 'is_dttm', 'is_active',
-        'type', 'groupby', 'count_distinct', 'sum', 'max', 'min',
+        'type', 'groupby', 'count_distinct', 'sum', 'avg', 'max', 'min',
         'filterable', 'expression', 'description', 'python_date_format',
         'database_expression'
     )
@@ -2144,6 +2146,7 @@ class DruidColumn(Model, AuditMixinNullable):
     groupby = Column(Boolean, default=False)
     count_distinct = Column(Boolean, default=False)
     sum = Column(Boolean, default=False)
+    avg = Column(Boolean, default=False)
     max = Column(Boolean, default=False)
     min = Column(Boolean, default=False)
     filterable = Column(Boolean, default=False)
@@ -2182,6 +2185,18 @@ def generate_metrics(self):
                 json=json.dumps({
                     'type': mt, 'name': name, 'fieldName': self.column_name})
             ))
+
+        if self.avg and self.isnum:
+            mt = corrected_type.lower() + 'Avg'
+            name = 'avg__' + self.column_name
+            metrics.append(DruidMetric(
+                metric_name=name,
+                metric_type='avg',
+                verbose_name='AVG({})'.format(self.column_name),
+                json=json.dumps({
+                    'type': mt, 'name': name, 'fieldName': self.column_name})
+            ))
+
         if self.min and self.isnum:
             mt = corrected_type.lower() + 'Min'
             name = 'min__' + self.column_name