match input/output types

censoredplanet · Jun 21, 2023 · ccf1ff1 · ccf1ff1
1 parent d7f4d1c
commit ccf1ff1
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 18 deletions.
diff --git a/pipeline/beam_tables.py b/pipeline/beam_tables.py
@@ -590,7 +590,8 @@ def derive_dashboard_rows(self, rows: beam.PCollection[schema.HyperquackRow]) ->
     pprint(dir(coder_rows))
     pprint(coder_rows.element_type)
 
-    dash_rows = (coder_rows | 'derive dashboard rows' >> SqlTransform(sql_query).with_output_types(BigqueryOutputRow))
+    dash_rows = (coder_rows | 'derive dashboard rows' >> SqlTransform(sql_query))
+                 #.with_output_types(BigqueryOutputRow))
 
     pprint("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
     pprint(dash_rows)

diff --git a/pipeline/metadata/schema.py b/pipeline/metadata/schema.py
@@ -225,15 +225,9 @@ class HyperquackRow(BigqueryRow):
    ('category', str),
    ('ip', str),
    ('ddate', str),
-   ('state_time', str),
-   ('end_time', str),
-   ('retry', int),
-   ('error', str),
-   ('anomaly', bool),
-   ('success', bool),
+   ('error', Optional[str]),
    ('is_control', bool),
    ('controls_failed', bool),
-   ('stateful_block', bool),
    ('measurement_id', str),
    ('source', str),
    ('outcome', str),
@@ -253,22 +247,16 @@ def convert_byperquack_row_to_bq_row_format(row: HyperquackRow) -> BigqueryInput
     row.category,
     row.ip,
     row.date,
-    row.start_time,
-    row.end_time,
-    row.retry,
     row.error,
-    row.anomaly,
-    row.success,
     row.is_control,
     row.controls_failed,
-    row.stateful_block,
     row.measurement_id,
     row.source,
     row.outcome,
     row.ip_metadata.country,
     row.ip_metadata.as_full_name,
-    row.ip_metadata.server_asn,
-    row.ip_metadata.server_organization,
+    row.ip_metadata.asn,
+    row.ip_metadata.organization,
     row.category
   )
 
@@ -284,7 +272,7 @@ def convert_byperquack_row_to_bq_row_format(row: HyperquackRow) -> BigqueryInput
    ('outcome', str),
    ('subnetwork', str),
    ('category', str),
-   ('count', int),
+   ('c', int),
    ('unexpected_count', int)
    ]
   )

diff --git a/table/queries/fake.sql b/table/queries/fake.sql
@@ -7,7 +7,7 @@ SELECT
         outcome AS outcome,
         server_asn AS subnetwork,
         COALESCE(domain_category, 'Uncategorized') AS category,
-        COUNT(*),
+        COUNT(*) AS c,
         1 AS unexpected_count
     FROM PCOLLECTION
     WHERE NOT controls_failed