Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added optional debug_truncate_bytes parameter to the config and extend the default log truncation limit #782

Merged
merged 3 commits into from
Jan 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion src/databricks/labs/ucx/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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())
Expand Down
36 changes: 27 additions & 9 deletions tests/unit/test_install.py
Original file line number Diff line number Diff line change
Expand Up @@ -365,7 +365,9 @@ def not_found(_):

ws.workspace.upload.assert_called_with(
"/Users/[email protected]/.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
Expand All @@ -387,7 +389,9 @@ def test_migrate_from_v1(ws, mocker):
ws.current_user.me = lambda: iam.User(user_name="[email protected]", 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-
Expand All @@ -411,7 +415,9 @@ def test_migrate_from_v1(ws, mocker):

ws.workspace.upload.assert_called_with(
"/Users/[email protected]/.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
Expand All @@ -433,7 +439,9 @@ def test_migrate_from_v1_selected_groups(ws, mocker):
ws.current_user.me = lambda: iam.User(user_name="[email protected]", 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:
Expand All @@ -459,7 +467,9 @@ def test_migrate_from_v1_selected_groups(ws, mocker):

ws.workspace.upload.assert_called_with(
"/Users/[email protected]/.ucx/config.yml",
b"""default_catalog: ucx_default
b"""connect:
debug_truncate_bytes: 250000
default_catalog: ucx_default
include_group_names:
- '42'
- '100'
Expand Down Expand Up @@ -500,7 +510,9 @@ def not_found(_):

ws.workspace.upload.assert_called_with(
"/Users/[email protected]/.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
Expand Down Expand Up @@ -538,7 +550,9 @@ def test_save_config_strip_group_names(ws):

ws.workspace.upload.assert_called_with(
"/Users/[email protected]/.ucx/config.yml",
b"""default_catalog: ucx_default
b"""connect:
debug_truncate_bytes: 250000
default_catalog: ucx_default
include_group_names:
- g1
- g2
Expand Down Expand Up @@ -598,7 +612,9 @@ def test_save_config_with_custom_policy(ws):

ws.workspace.upload.assert_called_with(
"/Users/[email protected]/.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
Expand Down Expand Up @@ -655,7 +671,9 @@ def test_save_config_with_glue(ws):

ws.workspace.upload.assert_called_with(
"/Users/[email protected]/.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
Expand Down
Loading