diff --git a/scripts/python/idc_index_data_manager.py b/scripts/python/idc_index_data_manager.py index d4914af..6464569 100644 --- a/scripts/python/idc_index_data_manager.py +++ b/scripts/python/idc_index_data_manager.py @@ -222,21 +222,21 @@ def _extract_column_name(column_def: str) -> str | None: column_def = column_def.strip().rstrip(",").strip() # Look for the last AS clause (to handle nested AS in CAST expressions) - # Use a regex that finds the rightmost AS followed by a word - as_matches = list(re.finditer(r"\bAS\b\s+(\w+)", column_def, re.IGNORECASE)) + # Use a regex that finds the rightmost AS followed by a word (optionally backtick-quoted) + as_matches = list(re.finditer(r"\bAS\b\s+`?(\w+)`?", column_def, re.IGNORECASE)) if as_matches: # Return the last match (rightmost AS clause) return as_matches[-1].group(1) # If no AS clause, try to get the column name # Remove function calls and get the last word before comma - # Handle cases like: column_name, or just column_name + # Handle cases like: column_name, or `column_name` (backtick-quoted reserved words) parts = column_def.split() if parts: # Get the last word that looks like an identifier for original_part in reversed(parts): - # Remove trailing punctuation - part = original_part.rstrip(",").strip() + # Remove trailing punctuation and backtick quoting + part = original_part.rstrip(",").strip().strip("`") # Check if it's a valid identifier (word characters only) if re.match(r"^\w+$", part): return part diff --git a/tests/test_column_description_parser.py b/tests/test_column_description_parser.py index 9232024..2d047a6 100644 --- a/tests/test_column_description_parser.py +++ b/tests/test_column_description_parser.py @@ -125,6 +125,21 @@ def test_extract_column_name_with_as(self): ) assert IDCIndexDataManager._extract_column_name("column AS alias,") == "alias" + def test_extract_column_name_backtick_quoted(self): + """Test extracting column name from backtick-quoted reserved words (e.g. Rows, Columns).""" + assert IDCIndexDataManager._extract_column_name("`Rows`,") == "Rows" + assert IDCIndexDataManager._extract_column_name("`Columns`,") == "Columns" + assert ( + IDCIndexDataManager._extract_column_name("ANY_VALUE(`Rows`) AS `Rows`,") + == "Rows" + ) + assert ( + IDCIndexDataManager._extract_column_name( + "ANY_VALUE(`Columns`) AS `Columns`," + ) + == "Columns" + ) + def test_extract_column_name_complex(self): """Test extracting column name from complex expressions.""" assert ( @@ -151,6 +166,24 @@ def test_complex_multiline_select(self): assert "series_size_MB" in descriptions assert descriptions["series_size_MB"] == "total size of the series in megabytes" + def test_backtick_quoted_columns(self): + """Test parsing descriptions for backtick-quoted reserved-word columns (Rows, Columns).""" + sql_query = """ +SELECT + # description: + # number of pixel rows per image slice + ANY_VALUE(`Rows`) AS `Rows`, + # description: + # number of pixel columns per image slice + ANY_VALUE(`Columns`) AS `Columns`, +FROM table +""" + descriptions = IDCIndexDataManager.parse_column_descriptions(sql_query) + assert "Rows" in descriptions + assert descriptions["Rows"] == "number of pixel rows per image slice" + assert "Columns" in descriptions + assert descriptions["Columns"] == "number of pixel columns per image slice" + def test_no_descriptions(self): """Test SQL query with no descriptions.""" sql_query = """