diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index 3216fe0c0e..c0cbdeb1d1 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -49,12 +49,15 @@ def __init__(self, cfg: PlatformConfig, **kwargs): # StreamRemoved exception. # https://github.com/flyteorg/flyte/blob/e8588f3a04995a420559327e78c3f95fbf64dc01/flyteadmin/pkg/common/constants.go#L14 # 32KB for error messages, 20MB for actual messages. - options = (("grpc.max_metadata_size", 32 * 1024), ("grpc.max_receive_message_length", 20 * 1024 * 1024)) + options = [ + ("grpc.max_metadata_size", 32 * 1024), + ("grpc.max_receive_message_length", 20 * 1024 * 1024), + ] + kwargs.pop("options", []) self._cfg = cfg self._channel = wrap_exceptions_channel( cfg, upgrade_channel_to_authenticated( - cfg, upgrade_channel_to_proxy_authenticated(cfg, get_channel(cfg, options=options)) + cfg, upgrade_channel_to_proxy_authenticated(cfg, get_channel(cfg, options=options, **kwargs)) ), ) self._stub = _admin_service.AdminServiceStub(self._channel) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index 00ca38b807..e86b50a9ca 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -20,3 +20,22 @@ def test_list_projects_paginated(mock_channel, mock_admin): project_list_request = _project_pb2.ProjectListRequest(limit=100, token="", filters=None, sort_by=None) client.list_projects(project_list_request) mock_admin.AdminServiceStub().ListProjects.assert_called_with(project_list_request, metadata=None) + + +@mock.patch("flytekit.clients.raw.grpc.insecure_channel") +def test_kwargs_passed_to_channel(mock_channel): + RawSynchronousFlyteClient( + PlatformConfig(endpoint="a.b.com", insecure=True), + options=[("grpc.default_authority", "foo.b.net")], + ) + + mock_channel.assert_called_with( + "a.b.com", + options=[ + # ensure we're always passing these defaults + ("grpc.max_metadata_size", mock.ANY), + ("grpc.max_receive_message_length", mock.ANY), + # as well as the user-provided options + ("grpc.default_authority", "foo.b.net"), + ], + )