fix for inconsistent v4 api responses

Author: OriHoch

airflow/knesset_data_pipelines/kns_odata.py

@@ -1,5 +1,7 @@
+import json
 import os
 from textwrap import dedent
+import xml.etree.ElementTree as ET
 
 import requests
 
@@ -19,9 +21,12 @@ def get_table_res_field(field):
         'Edm.Int32': 'integer',  # 123
         'Edm.Int16': 'integer',  # 123
         'Edm.Byte': 'integer',  # 123
-        'Edm.DateTimeOffset': 'datetime',  # "2016-02-28T10:22:10.843+02:00"
+        'Edm.DateTimeOffset': 'datetime',  # "2016-02-28T10:22:10.843+02:00",
+        'Edm.DateTime': 'datetime',
         None: 'string',  # "string",
         'Edm.Boolean': 'boolean',
+        'Edm.String': 'string',
+        'Edm.Decimal': 'string',
     }
     assert field.get("$Type") in field_types, f"unknown field type {field.get('$Type')}"
     return {
@@ -30,8 +35,8 @@ def get_table_res_field(field):
 
 
 def get_parliamentinfo_tables():
-    data = requests.get('https://knesset.gov.il/OdataV4/ParliamentInfo/$metadata?$format=json').json()
     res = {}
+    data = requests.get('https://knesset.gov.il/OdataV4/ParliamentInfo/$metadata?$format=json').json()
     for table_name, fields in data['OdataService.DAL.ParliamentInfo'].items():
         assert fields.pop("$Kind") == 'EntityType'
         keys = fields.pop("$Key") or []
@@ -42,11 +47,38 @@ def get_parliamentinfo_tables():
                 field_name: get_table_res_field(field) for field_name, field in fields.items() if field.get('$Kind') is None
             }
         }
+    for service_name in (
+            'Lobbyists',
+            'Votes',
+    ):
+        xmlstr = requests.get(f'https://knesset.gov.il/Odata/{service_name}.svc/$metadata').text
+        root = ET.fromstring(xmlstr)
+        ns = {
+            "edmx": "http://schemas.microsoft.com/ado/2007/06/edmx",
+            "edm": "http://schemas.microsoft.com/ado/2009/11/edm",
+        }
+        for schema in root.findall(".//edm:Schema", ns):
+            for entity_type in schema.findall("edm:EntityType", ns):
+                table_name = f'{service_name.lower()}_{entity_type.get("Name")}'
+                fields = {}
+                for prop in entity_type.findall("edm:Property", ns):
+                    field_name = prop.get("Name")
+                    fields[field_name] = get_table_res_field({
+                        "$Type": prop.get("Type"),
+                    })
+                edm_keys = entity_type.findall("edm:Key", ns)
+                assert len(edm_keys) == 1
+                edm_key = edm_keys[0]
+                keys = [k.get("Name") for k in edm_key.findall("edm:PropertyRef", ns)]
+                res[table_name] = {
+                    'primary_key': keys[0] if len(keys) == 1 else keys,
+                    'fields': fields,
+                }
     return res
 
 
-def get_pipelines_res_field(field):
-    assert field['source'] == '{name}', f'unexpected source {field["source"]}'
+def get_pipelines_res_field(field, pipeline_id=None, field_name=None):
+    assert field['source'] == '{name}', f'unexpected source {field["source"]} in field {field_name} of pipeline {pipeline_id}'
     return {
         'type': field['type'],
         'primary_key': bool(field.get('primaryKey')),
@@ -56,88 +88,126 @@ def get_pipelines_res_field(field):
 def get_parliamentinfo_pipelines():
     res = {}
     for pipeline in list_pipelines(full=True):
-        if pipeline.get('dataservice_params') and pipeline['dataservice_params'].get('service-name') == 'api':
-            res[pipeline['dataservice_params']['method-name']] = {}
-            res[pipeline['dataservice_params']['method-name']]['fields'] = {
-                name: get_pipelines_res_field(field) for name, field in pipeline['dataservice_params']['fields'].items()
+        pipeline_id = pipeline['pipeline_id']
+        if pipeline.get('dataservice_params'):
+            dataservice_params = pipeline['dataservice_params']
+            odata_v4 = dataservice_params.get('odata-v4', True)
+            if odata_v4:
+                assert dataservice_params.get('service-name') == 'api', f'{pipeline_id}: api service should have service-name set to api'
+                method_name = dataservice_params['method-name']
+            else:
+                method_name = f'{dataservice_params["service-name"]}_{dataservice_params["method-name"]}'
+            res[method_name] = {}
+            res[method_name]['fields'] = {
+                name: get_pipelines_res_field(field, pipeline_id=pipeline_id, field_name=name)
+                for name, field in dataservice_params['fields'].items()
             }
-            primary_keys = [name for name, field in res[pipeline['dataservice_params']['method-name']]['fields'].items() if field.get('primary_key')]
-            assert len(primary_keys) <= 1, f"unexpected number of primary keys in pipeline {pipeline['pipeline_id']}: {primary_keys}"
-            res[pipeline['dataservice_params']['method-name']]['primary_key'] = primary_keys[0] if primary_keys else None
+            primary_keys = [name for name, field in res[method_name]['fields'].items() if field.get('primary_key')]
+            if len(primary_keys) == 1:
+                res[method_name]['primary_key'] = primary_keys[0]
+            elif len(primary_keys) > 1:
+                res[method_name]['primary_key'] = primary_keys
+            else:
+                raise Exception(f'no primary key found in pipeline {pipeline_id} for method {method_name}')
     return res
 
 
-class MissingPipeline:
+class MissingObject:
 
-    def __init__(self, table_name, table):
+    def __init__(self, dst_obj_type, table_name, table=None):
+        self.dst_obj_type = dst_obj_type
         self.table_name = table_name
         self.table = table
 
     def __str__(self):
-        res = f'table {self.table_name} is missing from pipelines\n'
-        filename = f'knesset/{self.table_name.lower()}.yaml'
-        filecontent = dedent(f'''
-        pipeline-type: knesset dataservice
-        dataservice-parameters:
-          service-name: api
-          method-name: "{self.table_name}"
-          fields:
-        ''')
-        for field_name, field in self.table['fields'].items():
-            filecontent += f'    {field_name}:\n'
-            filecontent += f'      source: "{{name}}"\n'
-            filecontent += f'      type: "{field["type"]}"\n'
-            if field_name == self.table['primary_key']:
-                filecontent += f'      primaryKey: true\n'
-        with open(os.path.join(os.path.dirname(__file__), '..', 'pipelines', filename), 'w') as f:
-            f.write(filecontent)
-        return res
-
-
-class IncompletePipeline:
-
-    def __init__(self, table_name, table, pipeline):
+        if self.dst_obj_type == 'pipelines':
+            filename = f'knesset/{self.table_name.lower()}.yaml'
+            filecontent = dedent(f'''
+            pipeline-type: knesset dataservice
+            dataservice-parameters:
+              service-name: api
+              method-name: "{self.table_name}"
+              fields:
+            ''')
+            for field_name, field in self.table['fields'].items():
+                filecontent += f'    {field_name}:\n'
+                filecontent += f'      source: "{{name}}"\n'
+                filecontent += f'      type: "{field["type"]}"\n'
+                if field_name == self.table['primary_key']:
+                    filecontent += f'      primaryKey: true\n'
+            with open(os.path.join(os.path.dirname(__file__), '..', 'pipelines', filename), 'w') as f:
+                f.write(filecontent)
+            return f'{self.table_name} is missing from {self.dst_obj_type}, wrote to filename {filename}\n'
+        else:
+            return f'{self.table_name} exists in pipelines but missing in tables'
+
+
+class InvalidPipeline:
+
+    def __init__(self, dst_obj_type, table_name, table):
+        self.dst_obj_type = dst_obj_type
         self.table_name = table_name
         self.table = table
-        self.pipeline = pipeline
         self.missing_fields = []
         self.wrong_type_fields = []
+        self.extra_fields = []
+        self.primary_key_mismatch = False
 
     def __str__(self):
-        res = f'pipeline for table {self.table_name} is incomplete:\n'
+        res = f'mismatch in pipeline for table {self.table_name}:\n' if self.dst_obj_type == 'pipelines' else f'mismatch in table for pipeline {self.table_name}:\n'
         if len(self.missing_fields) > 0:
-            res += ' - missing fields:\n'
+            res += ' - missing fields in pipeline:\n' if self.dst_obj_type == 'pipelines' else ' - extra fields in pipeline do not appear in table:\n'
             for field_name in self.missing_fields:
-                res += f'{field_name}: {{source: "{{name}}", type: "{self.table["fields"][field_name]["type"]}"}}\n'
+                res += f'{field_name}: {{source: "{{name}}", type: "{self.table["fields"].get(field_name, {}).get("type")}"}}\n'
         if len(self.wrong_type_fields) > 0:
-            res += ' - wrong type fields:\n'
+            res += ' - wrong type fields:\n' if self.dst_obj_type == 'pipelines' else ' - wrong type fields (type in table shown below does not match type in pipeline):\n'
             for field_name in self.wrong_type_fields:
                 res += f'{field_name}: {{source: "{{name}}", type: "{self.table["fields"][field_name]["type"]}"}}\n'
+        if len(self.extra_fields) > 0:
+            res += ' - extra fields in pipeline which are not in table:\n' if self.dst_obj_type == 'pipelines' else ' - missing  fields in pipeline which appear in table:\n'
+            for field_name in self.extra_fields:
+                res += f'{field_name}\n'
+        if self.primary_key_mismatch:
+            res += ' - primary key mismatch:\n'
+            res += f'table primary key: {self.table["primary_key"]}\n'
         return res
 
+    def has_errors(self):
+        return sum((
+            len(self.missing_fields),
+            len(self.wrong_type_fields),
+            len(self.extra_fields),
+        )) > 0 or self.primary_key_mismatch
+
 
 def compare_parliamentinfo_tables_pipelines():
     fix_objects = []
     tables = get_parliamentinfo_tables()
     pipelines = get_parliamentinfo_pipelines()
-    for table_name, table in tables.items():
-        if table_name not in pipelines:
-            fix_objects.append(MissingPipeline(table_name, table))
-            continue
-        pipeline = pipelines[table_name]
-        incomplete_pipeline = IncompletePipeline(table_name, table, pipeline)
-        for field_name, field in table['fields'].items():
-            if field_name not in pipeline['fields']:
-                if table['primary_key'] == field_name and pipeline['primary_key']:
-                    field_name = pipeline['primary_key']
-                else:
-                    incomplete_pipeline.missing_fields.append(field_name)
-                    continue
-            pipeline_field = pipeline['fields'][field_name]
-            if field['type'] != pipeline_field['type'] and field_name not in IGNORE_FIELD_TYPE_ERRORS.get(table_name, []):
-                incomplete_pipeline.wrong_type_fields.append(field_name)
-        if len(incomplete_pipeline.missing_fields) > 0 or len(incomplete_pipeline.wrong_type_fields) > 0:
-            fix_objects.append(incomplete_pipeline)
+    for dst_obj_type, src, dst in (('pipelines', tables, pipelines), ('tables', pipelines, tables)):
+        for src_name, src_obj in src.items():
+            if src_name not in dst:
+                fix_objects.append(MissingObject(dst_obj_type, src_name, src_obj if dst_obj_type == 'pipelines' else None))
+            else:
+                dst_obj = dst[src_name]
+                invalid_pipeline = InvalidPipeline(dst_obj_type, src_name, tables[src_name])
+                for field_name, field in src_obj['fields'].items():
+                    if field_name not in dst_obj['fields']:
+                        if dst_obj_type == 'pipelines' and src_obj['primary_key'] == field_name and dst_obj['primary_key']:
+                            # handle the special case of primary key for v4 being renamed to Id
+                            field_name = dst_obj['primary_key']
+                        elif dst_obj_type == 'tables' and src_obj['primary_key'] == field_name and dst_obj['primary_key'] == 'Id':
+                            # handle the special case of primary key for v4 being renamed to Id
+                            field_name = 'Id'
+                        else:
+                            invalid_pipeline.missing_fields.append(field_name)
+                            continue
+                    if field['type'] != dst_obj['fields'][field_name]['type'] and field_name not in IGNORE_FIELD_TYPE_ERRORS.get(src_name, []):
+                        invalid_pipeline.wrong_type_fields.append(field_name)
+                if set(src_obj['primary_key']) != set(dst_obj['primary_key']) and tables[src_name]['primary_key'] != 'Id':
+                    invalid_pipeline.primary_key_mismatch = True
+                if invalid_pipeline.has_errors():
+                    fix_objects.append(invalid_pipeline)
     if len(fix_objects) > 0:
         for fix_object in fix_objects:
             print(fix_object)
@@ -148,3 +218,7 @@ def compare_parliamentinfo_tables_pipelines():
 
 if __name__ == '__main__':
     compare_parliamentinfo_tables_pipelines()
+    # print(json.dumps(list(get_parliamentinfo_pipelines().keys()), indent=2))
+    # print(json.dumps(get_parliamentinfo_pipelines()["lobbyists_V_Lobbyists"], indent=2))
+    # print(json.dumps(list(get_parliamentinfo_tables().keys()), indent=2))
+    # print(json.dumps(get_parliamentinfo_tables()["lobbyists_V_Lobbyists"], indent=2))

airflow/knesset_data_pipelines/run_pipeline.py

@@ -13,6 +13,8 @@
 from pprint import pprint
 from textwrap import dedent
 from collections import defaultdict
+import pickle
+import functools
 
 import requests
 import dataflows as DF
@@ -25,7 +27,7 @@
 
 
 RECOVERABLE_SERVER_ERRORS = [500, 504]
-RECOVERABLE_THROTTLE_ERRORS = [503, 403, 504]
+RECOVERABLE_THROTTLE_ERRORS = [500, 503, 403, 504]
 
 # these errors indicate that the pipeline cannot run using the standard dataservice flow but have to run via Docker
 UNSUPPORTED_PIPELINE_PARAMS_ERRORS = ['pipeline dependencies are not supported', 'unknown pipeline-type: None', 'additional-steps is not supported']
@@ -215,6 +217,20 @@ def get_row_from_entry(params, entry, v4=False):
     return {fieldname: get_field_from_entry(fieldname, field, data, v4=v4) for fieldname, field in params['fields'].items()}
 
 
+def get_key_row_from_entry(params, entry, v4=False):
+    row = get_row_from_entry(params, entry, v4)
+    primary_keys = set()
+    for field_name, field in params['fields'].items():
+        if field.get('primaryKey'):
+            primary_keys.add(field_name)
+    key = ';'.join([str(row[k]) for k in sorted(list(primary_keys))])
+    assert len(key) > 0, f'unexpected key ({key}): {params} {entry}'
+    return {
+        'key': key,
+        'row': row,
+    }
+
+
 def get_next_skiptoken(first_skiptoken, i, processed_entry_ids):
     if len(processed_entry_ids) > i:
         skiptoken = processed_entry_ids[-(i+1)]
@@ -262,7 +278,7 @@ def add_dataservice_collection_resource_odata_v4(params, proxies, stats, limit_r
         num_entries = 0
         for entry in (res if isinstance(res, list) else res['value']):
             stats['rows'] += 1
-            yield get_row_from_entry(params, entry, v4=True)
+            yield get_key_row_from_entry(params, entry, v4=True)
             num_entries += 1
         skip = num_entries if skip is None else skip + num_entries
         if odata_count is None and isinstance(res, dict) and '@odata.count' in res:
@@ -341,7 +357,7 @@ def add_dataservice_collection_resource(params, proxies=None, stats=None, limit_
                     has_valid_entry_ids = False
             if not has_valid_entry_ids or entry_id not in processed_entry_ids:
                 stats['rows'] += 1
-                yield get_row_from_entry(params, entry)
+                yield get_key_row_from_entry(params, entry)
                 if has_valid_entry_ids:
                     processed_entry_ids.append(entry_id)
         try:
@@ -385,6 +401,42 @@ def get_schema_set_type_kwargs(dataservice_params):
     return res
 
 
+def validate_items_iterator(iterator_func, stats, num_validations=2, max_iterations=5):
+    print(f'Starting validate_items_iterator (num_validations={num_validations}, max_iterations={max_iterations})')
+    ok, last_keys = False, None
+    with tempfile.TemporaryDirectory() as tmpdir:
+        filename = os.path.join(tmpdir, 'pickle')
+        for i in range(max_iterations):
+            stats.clear()
+            print(f'validate_items_iterator iteration #{i+1}')
+            if i > 0:
+                os.unlink(filename)
+            with open(filename, 'wb') as f:
+                pickler = pickle.Pickler(f)
+                cur_keys = set()
+                for data in iterator_func():
+                    pickler.dump(data['row'])
+                    key = data['key']
+                    assert key not in cur_keys, f'duplicate key {key}'
+                    cur_keys.add(key)
+            if last_keys is not None and cur_keys == last_keys:
+                num_validations -= 1
+                print(f'Validated, remaining validations: {num_validations}')
+            if num_validations <= 0:
+                ok = True
+                break
+            last_keys = cur_keys
+        assert ok, f'Failed to validate items iterator after {max_iterations} iterations'
+        print('Validated all required validations')
+        with open(filename, 'rb') as f:
+            unpickler = pickle.Unpickler(f)
+            while True:
+                try:
+                    yield unpickler.load()
+                except EOFError:
+                    break
+
+
 def _run_pipeline(table_name, storage_url, pipeline_id, storage_path, dataservice_params, limit_rows, pipeline_name, dump_to_path, dump_to_db, dump_to_storage,
                   stop_on_throttled_error=False, start_url=None, load_from=None):
     if stop_on_throttled_error:
@@ -395,7 +447,14 @@ def _run_pipeline(table_name, storage_url, pipeline_id, storage_path, dataservic
     print(f'pipeline_id: {pipeline_id}\nstorage_url: {storage_url}\nstorage_path: {storage_path}\ntable_name: {table_name}')
     stats = defaultdict(int)
     temp_table_name = f'__temp__{table_name}'
-    rows_iterator = add_dataservice_collection_resource(dataservice_params, stats=stats, limit_rows=limit_rows, stop_on_throttled_error=stop_on_throttled_error, start_url=start_url, load_from=load_from)
+    rows_iterator = validate_items_iterator(
+        functools.partial(
+            add_dataservice_collection_resource,
+            dataservice_params,
+            stats=stats, limit_rows=limit_rows, stop_on_throttled_error=stop_on_throttled_error, start_url=start_url, load_from=load_from
+        ),
+        stats,
+    )
     first_row = next(rows_iterator, None)
     if first_row is not None:
         DF.Flow(