diff --git a/app/agent/tools/bigquery.py b/app/agent/tools/bigquery.py index 82b9b45..ae04ccb 100644 --- a/app/agent/tools/bigquery.py +++ b/app/agent/tools/bigquery.py @@ -111,11 +111,18 @@ def decode_table_values( search_query = f""" SELECT nome_coluna, chave, valor FROM {dict_table_id} - WHERE id_tabela = '{table_name}' + WHERE id_tabela = @table_name """ + query_params = [ + bq.ScalarQueryParameter("table_name", "STRING", table_name), + ] + if column_name is not None: - search_query += f"AND nome_coluna = '{column_name}'" + search_query += "AND nome_coluna = @column_name\n" + query_params.append( + bq.ScalarQueryParameter("column_name", "STRING", column_name), + ) search_query += "ORDER BY nome_coluna, chave" @@ -127,7 +134,10 @@ def decode_table_values( try: client = _get_client() - job = client.query(search_query, job_config=bq.QueryJobConfig(labels=labels)) + job = client.query( + search_query, + job_config=bq.QueryJobConfig(query_parameters=query_params, labels=labels), + ) results = [dict(row) for row in job.result()] except GoogleAPICallError as e: reason = e.errors[0].get("reason") if getattr(e, "errors", None) else None diff --git a/tests/app/agent/tools/test_bigquery.py b/tests/app/agent/tools/test_bigquery.py index 1d71563..f64e7b2 100644 --- a/tests/app/agent/tools/test_bigquery.py +++ b/tests/app/agent/tools/test_bigquery.py @@ -157,6 +157,15 @@ def test_decode_all_columns(self, mocker: MockerFixture, mock_config: dict): assert len(output) == 2 + # Verify parameterized table filter was added to query + call_args = mock_bigquery_client.query.call_args[0][0] + assert "id_tabela = @table_name" in call_args + + # Verify query parameters include table_name + job_config = mock_bigquery_client.query.call_args[1]["job_config"] + param_names = {p.name for p in job_config.query_parameters} + assert "table_name" in param_names + def test_decode_specific_column(self, mocker: MockerFixture, mock_config: dict): """Test decoding a specific column.""" mock_query_job = MagicMock() @@ -183,9 +192,15 @@ def test_decode_specific_column(self, mocker: MockerFixture, mock_config: dict): output = json.loads(result) assert len(output) == 2 - # Verify column filter was added to query + + # Verify parameterized column filter was added to query call_args = mock_bigquery_client.query.call_args[0][0] - assert "nome_coluna = 'col1'" in call_args + assert "nome_coluna = @column_name" in call_args + + # Verify query parameters include column_name + job_config = mock_bigquery_client.query.call_args[1]["job_config"] + param_names = {p.name for p in job_config.query_parameters} + assert "column_name" in param_names def test_dictionary_not_found(self, mocker: MockerFixture, mock_config: dict): """Test error when dictionary table doesn't exist."""