Fix BigQuery connector variable iteration for protobuf changes

- Update BigQueryConnector to iterate over VariableEntry list instead of dict
- Fix test fixtures to use VariableEntry with key/value attributes
- Add comprehensive test for new variable iteration pattern
- Update Snowflake tests for same protobuf structure change
- Fix workflow and Makefile for plugin testing

Signed-off-by: Kevin Su <pingsutw@apache.org>

Author: pingsutw

.github/workflows/tests.yml

@@ -60,5 +60,5 @@ jobs:
 
       - name: Run tests
         env:
-          FLYTE_PLUGINS: ${{ matrix.plugin }}
+          FLYTE_PLUGINS: plugins/${{ matrix.plugin }}
         run: make unit_test_plugins

Makefile

@@ -75,7 +75,7 @@ unit_test_plugins:
 	@for plugin in $${FLYTE_PLUGIN:-plugins/*}; do \
 		if [ -d "$$plugin/tests" ]; then \
 			echo "🚀 Testing plugin: $$plugin..."; \
-			cd "$$plugin" && uv run python -m pytest tests/ && cd ../../..; \
+			cd "$$plugin" && uv run --active python -m pytest tests/ && cd ../../..; \
 		fi \
 	done
 

plugins/bigquery/src/flyteplugins/bigquery/connector.py

@@ -69,8 +69,8 @@ async def create(
         job_config = None
         if inputs:
             python_interface_inputs = {
-                name: TypeEngine.guess_python_type(lt.type)
-                for name, lt in task_template.interface.inputs.variables.items()
+                variable.key: TypeEngine.guess_python_type(variable.value.type)
+                for variable in task_template.interface.inputs.variables
             }
             job_config = bigquery.QueryJobConfig(
                 query_parameters=[

plugins/bigquery/tests/test_connector.py

@@ -5,7 +5,7 @@
 import pytest
 from flyte.io import DataFrame
 from flyteidl2.core.execution_pb2 import TaskExecution
-from flyteidl2.core.interface_pb2 import Variable, VariableMap
+from flyteidl2.core.interface_pb2 import Variable, VariableEntry, VariableMap
 from flyteidl2.core.tasks_pb2 import Sql, TaskTemplate
 from flyteidl2.core.types_pb2 import LiteralType, SimpleType
 from google.protobuf import struct_pb2
@@ -72,13 +72,14 @@ def task_template_with_inputs(self):
         template.sql.CopyFrom(Sql(statement="SELECT * FROM table WHERE id = @user_id", dialect=Sql.Dialect.ANSI))
         template.metadata.runtime.version = "1.0.0"
 
-        # Add input variables
+        # Add input variables using the new list-based structure
         int_type = LiteralType()
         int_type.simple = SimpleType.INTEGER
         user_id_var = Variable(type=int_type)
 
         variables = VariableMap()
-        variables.variables["user_id"].CopyFrom(user_id_var)
+        var_entry = VariableEntry(key="user_id", value=user_id_var)
+        variables.variables.append(var_entry)
         template.interface.inputs.CopyFrom(variables)
 
         custom = struct_pb2.Struct()
@@ -329,7 +330,7 @@ async def test_create_with_multiple_input_types(self, connector):
         )
         template.metadata.runtime.version = "1.0.0"
 
-        # Add multiple input variables with different types
+        # Add multiple input variables with different types using the new list-based structure
         int_type = LiteralType()
         int_type.simple = SimpleType.INTEGER
         str_type = LiteralType()
@@ -338,9 +339,9 @@ async def test_create_with_multiple_input_types(self, connector):
         bool_type.simple = SimpleType.BOOLEAN
 
         variables = VariableMap()
-        variables.variables["user_id"].CopyFrom(Variable(type=int_type))
-        variables.variables["name"].CopyFrom(Variable(type=str_type))
-        variables.variables["active"].CopyFrom(Variable(type=bool_type))
+        variables.variables.append(VariableEntry(key="user_id", value=Variable(type=int_type)))
+        variables.variables.append(VariableEntry(key="name", value=Variable(type=str_type)))
+        variables.variables.append(VariableEntry(key="active", value=Variable(type=bool_type)))
         template.interface.inputs.CopyFrom(variables)
 
         custom = struct_pb2.Struct()
@@ -447,3 +448,71 @@ async def test_create_with_google_application_credentials(self, connector, task_
                 # Verify the credentials were passed to the client
                 mock_client_class.assert_called_once()
                 assert mock_client_class.call_args[1]["credentials"] == mock_credentials
+
+    @pytest.mark.asyncio
+    async def test_create_iterates_variables_with_new_structure(self, connector):
+        """Test that the connector correctly iterates over variables using the new iteration pattern.
+
+        This test verifies the change from:
+            for name, lt in task_template.interface.inputs.variables.items()
+        To:
+            for variable in task_template.interface.inputs.variables
+
+        The variables field changed from a map to a repeated field (list), so we now
+        iterate directly over the list of Variable objects which have key and value attributes.
+        """
+        template = TaskTemplate()
+        template.sql.CopyFrom(
+            Sql(
+                statement="SELECT * FROM table WHERE user_id = @user_id AND email = @email",
+                dialect=Sql.Dialect.ANSI,
+            )
+        )
+        template.metadata.runtime.version = "2.0.0"
+
+        # Create variables using the new list-based VariableMap structure
+        int_type = LiteralType()
+        int_type.simple = SimpleType.INTEGER
+        str_type = LiteralType()
+        str_type.simple = SimpleType.STRING
+
+        variables = VariableMap()
+        variables.variables.append(VariableEntry(key="user_id", value=Variable(type=int_type)))
+        variables.variables.append(VariableEntry(key="email", value=Variable(type=str_type)))
+        template.interface.inputs.CopyFrom(variables)
+
+        custom = struct_pb2.Struct()
+        custom["ProjectID"] = "test-project"
+        custom["Location"] = "US"
+        custom["Domain"] = "test-domain"
+        template.custom.CopyFrom(custom)
+
+        with patch("flyteplugins.bigquery.connector.bigquery.Client") as mock_client_class:
+            mock_client = MagicMock()
+            mock_client_class.return_value = mock_client
+
+            mock_query_job = MagicMock()
+            mock_query_job.job_id = "job-iteration-test"
+            mock_client.query.return_value = mock_query_job
+
+            inputs = {"user_id": 42, "email": "test@example.com"}
+            metadata = await connector.create(template, inputs=inputs)
+
+            assert metadata.job_id == "job-iteration-test"
+
+            # Verify that the query was called with proper parameters
+            call_args = mock_client.query.call_args
+            job_config = call_args[1]["job_config"]
+
+            # The new iteration pattern should successfully create query parameters
+            assert len(job_config.query_parameters) == 2
+
+            param_dict = {p.name: p for p in job_config.query_parameters}
+            assert "user_id" in param_dict
+            assert "email" in param_dict
+            assert param_dict["user_id"].value == 42
+            assert param_dict["email"].value == "test@example.com"
+
+            # Verify parameter types are correctly mapped
+            assert param_dict["user_id"].type_ == "INT64"
+            assert param_dict["email"].type_ == "STRING"

plugins/snowflake/tests/test_connector.py

@@ -3,7 +3,7 @@
 import pytest
 from flyte.io import DataFrame
 from flyteidl2.core.execution_pb2 import TaskExecution
-from flyteidl2.core.interface_pb2 import Variable, VariableMap
+from flyteidl2.core.interface_pb2 import Variable, VariableEntry, VariableMap
 from flyteidl2.core.tasks_pb2 import Sql, TaskTemplate
 from flyteidl2.core.types_pb2 import LiteralType, StructuredDatasetType
 from google.protobuf import struct_pb2
@@ -104,11 +104,14 @@ def task_template_with_output(self):
         template.custom.CopyFrom(custom)
 
         # Set output variables so has_output is True
-        template.interface.outputs.CopyFrom(
-            VariableMap(
-                variables={"results": Variable(type=LiteralType(structured_dataset_type=StructuredDatasetType()))}
+        output_vars = VariableMap()
+        output_vars.variables.append(
+            VariableEntry(
+                key="results",
+                value=Variable(type=LiteralType(structured_dataset_type=StructuredDatasetType()))
             )
         )
+        template.interface.outputs.CopyFrom(output_vars)
 
         return template