|
| 1 | +from libs.utils import ( |
| 2 | + pull_validated_params, |
| 3 | +) |
| 4 | +from airflow.providers.postgres.hooks.postgres import PostgresHook |
| 5 | +import logging |
| 6 | + |
| 7 | +# PostgreSQL connection hook |
| 8 | +pg_hook = PostgresHook(postgres_conn_id="postgres_db_conn") |
| 9 | + |
| 10 | + |
| 11 | +def process_search_recommendations(**kwargs) -> None: |
| 12 | + """ |
| 13 | + Process search recommendations for a given scan report table. |
| 14 | +
|
| 15 | + This function: |
| 16 | + 1. Retrieves scan report value strings and searches the OMOP concept table using ILIKE queries |
| 17 | + 2. Selects the top 3 matches for each value |
| 18 | + 3. Inserts them as MappingRecommendations into the database |
| 19 | +
|
| 20 | + Parameters are validated upstream in validate_params_auto_mapping task. |
| 21 | + """ |
| 22 | + # Get validated parameters from XCom |
| 23 | + validated_params = pull_validated_params(kwargs, "validate_params_auto_mapping") |
| 24 | + |
| 25 | + table_id = validated_params["table_id"] |
| 26 | + |
| 27 | + try: |
| 28 | + # Get content type for ScanReportValue |
| 29 | + content_type_query = """ |
| 30 | + SELECT id FROM django_content_type |
| 31 | + WHERE app_label = 'mapping' AND model = 'scanreportvalue' |
| 32 | + LIMIT 1 |
| 33 | + """ |
| 34 | + content_type_result = pg_hook.get_first(content_type_query) |
| 35 | + if not content_type_result: |
| 36 | + raise ValueError("Could not find content type for ScanReportValue") |
| 37 | + content_type_id = content_type_result[0] |
| 38 | + |
| 39 | + # Get all scan report values for the table |
| 40 | + get_values_query = """ |
| 41 | + SELECT |
| 42 | + sr_value.id, |
| 43 | + sr_value.value |
| 44 | + FROM mapping_scanreportvalue sr_value |
| 45 | + JOIN mapping_scanreportfield sr_field ON sr_value.scan_report_field_id = sr_field.id |
| 46 | + WHERE sr_field.scan_report_table_id = %(table_id)s |
| 47 | + AND sr_value.value IS NOT NULL |
| 48 | + AND TRIM(sr_value.value) <> '' |
| 49 | + """ |
| 50 | + |
| 51 | + logging.info(f"Getting scan report values for table_id: {table_id}") |
| 52 | + values = pg_hook.get_records( |
| 53 | + get_values_query, parameters={"table_id": table_id} |
| 54 | + ) |
| 55 | + |
| 56 | + if not values: |
| 57 | + logging.info(f"No values found for table {table_id}") |
| 58 | + return |
| 59 | + |
| 60 | + # Process each value with individual queries |
| 61 | + # This is more efficient for large concept tables because: |
| 62 | + # 1. Each query has LIMIT 3, stopping early |
| 63 | + # 2. No cross-join with 1.9M concept records |
| 64 | + # 3. Database can optimize each focused query |
| 65 | + recommendations_created = 0 |
| 66 | + for value_id, value_string in values: |
| 67 | + try: |
| 68 | + # Search OMOP concepts using ILIKE with LIMIT 3 |
| 69 | + # This is more efficient than cross-join approach |
| 70 | + search_query = """ |
| 71 | + SELECT concept_id |
| 72 | + FROM omop.concept |
| 73 | + WHERE concept_name ILIKE %(search_term)s |
| 74 | + AND invalid_reason IS NULL |
| 75 | + AND standard_concept = 'S' |
| 76 | + ORDER BY concept_name |
| 77 | + LIMIT 3 |
| 78 | + """ |
| 79 | + |
| 80 | + matches = pg_hook.get_records( |
| 81 | + search_query, parameters={"search_term": f"%{value_string}%"} |
| 82 | + ) |
| 83 | + |
| 84 | + if matches: |
| 85 | + # Insert recommendations for each match |
| 86 | + for (concept_id,) in matches: |
| 87 | + insert_query = """ |
| 88 | + INSERT INTO mapping_mappingrecommendation ( |
| 89 | + content_type_id, |
| 90 | + object_id, |
| 91 | + concept_id, |
| 92 | + score, |
| 93 | + tool_name, |
| 94 | + tool_version, |
| 95 | + created_at, |
| 96 | + updated_at |
| 97 | + ) VALUES ( |
| 98 | + %(content_type_id)s, |
| 99 | + %(object_id)s, |
| 100 | + %(concept_id)s, |
| 101 | + %(score)s, |
| 102 | + %(tool_name)s, |
| 103 | + %(tool_version)s, |
| 104 | + NOW(), |
| 105 | + NOW() |
| 106 | + ) |
| 107 | + """ |
| 108 | + |
| 109 | + pg_hook.run( |
| 110 | + insert_query, |
| 111 | + parameters={ |
| 112 | + "content_type_id": content_type_id, |
| 113 | + "object_id": value_id, |
| 114 | + "concept_id": concept_id, |
| 115 | + "score": 0.5, # Base score for ILIKE matches |
| 116 | + "tool_name": "string-search", |
| 117 | + "tool_version": "1.0.0", |
| 118 | + }, |
| 119 | + ) |
| 120 | + |
| 121 | + recommendations_created += 1 |
| 122 | + |
| 123 | + logging.info( |
| 124 | + f"Created {len(matches)} recommendations for value '{value_string}'" |
| 125 | + ) |
| 126 | + else: |
| 127 | + logging.info(f"No matches found for value '{value_string}'") |
| 128 | + |
| 129 | + except Exception as e: |
| 130 | + logging.error(f"Error processing value '{value_string}': {str(e)}") |
| 131 | + # Continue with next value instead of failing the entire process |
| 132 | + continue |
| 133 | + |
| 134 | + logging.info( |
| 135 | + f"Successfully created {recommendations_created} search recommendations" |
| 136 | + ) |
| 137 | + |
| 138 | + except Exception as e: |
| 139 | + logging.error(f"Error in process_search_recommendations: {str(e)}") |
| 140 | + raise |
0 commit comments