From 0f158d9c1b284a4d0bac348650ba706da783bc60 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Tue, 16 Jan 2024 01:55:00 +0100 Subject: [PATCH] Added optional `debug_truncate_bytes` parameter to the config and extend the default log truncation limit (#782) --- src/databricks/labs/ucx/config.py | 8 ++++++- tests/unit/test_install.py | 36 +++++++++++++++++++++++-------- 2 files changed, 34 insertions(+), 10 deletions(-) diff --git a/src/databricks/labs/ucx/config.py b/src/databricks/labs/ucx/config.py index 6a0f1e75fe..6b1055390a 100644 --- a/src/databricks/labs/ucx/config.py +++ b/src/databricks/labs/ucx/config.py @@ -27,6 +27,9 @@ class ConnectConfig: cluster_id: str | None = None profile: str | None = None debug_headers: bool | None = False + # Truncate JSON fields in HTTP requests and responses above this limit. + # If this occurs, the log message will include the text `... (XXX additional elements)` + debug_truncate_bytes: int | None = 250000 rate_limit: int | None = None max_connections_per_pool: int | None = None max_connection_pools: int | None = None @@ -45,6 +48,7 @@ def from_databricks_config(cfg: Config) -> "ConnectConfig": cluster_id=cfg.cluster_id, profile=cfg.profile, debug_headers=cfg.debug_headers, + debug_truncate_bytes=cfg.debug_truncate_bytes, rate_limit=cfg.rate_limit, max_connection_pools=cfg.max_connection_pools, max_connections_per_pool=cfg.max_connections_per_pool, @@ -64,6 +68,7 @@ def to_databricks_config(self): cluster_id=self.cluster_id, profile=self.profile, debug_headers=self.debug_headers, + debug_truncate_bytes=self.debug_truncate_bytes, rate_limit=self.rate_limit, max_connection_pools=self.max_connection_pools, max_connections_per_pool=self.max_connections_per_pool, @@ -213,7 +218,8 @@ class WorkspaceConfig(_Config["WorkspaceConfig"]): @classmethod def from_dict(cls, raw: dict): raw = cls._migrate_from_v1(raw) - return cls(**raw) + connect = ConnectConfig.from_dict(raw.pop("connect", {})) + return cls(connect=connect, **raw) def to_workspace_client(self) -> WorkspaceClient: return WorkspaceClient(config=self.to_databricks_config()) diff --git a/tests/unit/test_install.py b/tests/unit/test_install.py index 6de20050f9..da830d02bb 100644 --- a/tests/unit/test_install.py +++ b/tests/unit/test_install.py @@ -365,7 +365,9 @@ def not_found(_): ws.workspace.upload.assert_called_with( "/Users/me@example.com/.ucx/config.yml", - b"""default_catalog: ucx_default + b"""connect: + debug_truncate_bytes: 250000 +default_catalog: ucx_default inventory_database: ucx log_level: INFO num_threads: 8 @@ -387,7 +389,9 @@ def test_migrate_from_v1(ws, mocker): ws.current_user.me = lambda: iam.User(user_name="me@example.com", groups=[iam.ComplexValue(display="admins")]) ws.config.host = "https://foo" ws.config.is_aws = True - config_bytes = b"""default_catalog: ucx_default + config_bytes = b"""connect: + debug_truncate_bytes: 250000 +default_catalog: ucx_default groups: auto: true backup_group_prefix: db-temp- @@ -411,7 +415,9 @@ def test_migrate_from_v1(ws, mocker): ws.workspace.upload.assert_called_with( "/Users/me@example.com/.ucx/config.yml", - b"""default_catalog: ucx_default + b"""connect: + debug_truncate_bytes: 250000 +default_catalog: ucx_default inventory_database: ucx log_level: INFO num_threads: 8 @@ -433,7 +439,9 @@ def test_migrate_from_v1_selected_groups(ws, mocker): ws.current_user.me = lambda: iam.User(user_name="me@example.com", groups=[iam.ComplexValue(display="admins")]) ws.config.host = "https://foo" ws.config.is_aws = True - config_bytes = b"""default_catalog: ucx_default + config_bytes = b"""connect: + debug_truncate_bytes: 250000 +default_catalog: ucx_default groups: backup_group_prefix: 'backup_baguette_prefix' selected: @@ -459,7 +467,9 @@ def test_migrate_from_v1_selected_groups(ws, mocker): ws.workspace.upload.assert_called_with( "/Users/me@example.com/.ucx/config.yml", - b"""default_catalog: ucx_default + b"""connect: + debug_truncate_bytes: 250000 +default_catalog: ucx_default include_group_names: - '42' - '100' @@ -500,7 +510,9 @@ def not_found(_): ws.workspace.upload.assert_called_with( "/Users/me@example.com/.ucx/config.yml", - b"""default_catalog: ucx_default + b"""connect: + debug_truncate_bytes: 250000 +default_catalog: ucx_default inventory_database: ucx log_level: INFO num_threads: 8 @@ -538,7 +550,9 @@ def test_save_config_strip_group_names(ws): ws.workspace.upload.assert_called_with( "/Users/me@example.com/.ucx/config.yml", - b"""default_catalog: ucx_default + b"""connect: + debug_truncate_bytes: 250000 +default_catalog: ucx_default include_group_names: - g1 - g2 @@ -598,7 +612,9 @@ def test_save_config_with_custom_policy(ws): ws.workspace.upload.assert_called_with( "/Users/me@example.com/.ucx/config.yml", - b"""custom_cluster_policy_id: 0123456789ABCDEF + b"""connect: + debug_truncate_bytes: 250000 +custom_cluster_policy_id: 0123456789ABCDEF default_catalog: ucx_default inventory_database: ucx log_level: INFO @@ -655,7 +671,9 @@ def test_save_config_with_glue(ws): ws.workspace.upload.assert_called_with( "/Users/me@example.com/.ucx/config.yml", - b"""default_catalog: ucx_default + b"""connect: + debug_truncate_bytes: 250000 +default_catalog: ucx_default instance_profile: arn:aws:iam::111222333:instance-profile/foo-instance-profile inventory_database: ucx log_level: INFO