Skip to content

Commit

Permalink
..
Browse files Browse the repository at this point in the history
  • Loading branch information
nfx committed Mar 23, 2024
1 parent 819c2d1 commit af99cca
Show file tree
Hide file tree
Showing 4 changed files with 97 additions and 1 deletion.
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ classifiers = [
"Programming Language :: Python :: Implementation :: PyPy",
]
dependencies = [
"databricks-labs-blueprint>=0.4.2",
"databricks-labs-blueprint[yaml]>=0.4.2",
"databricks-sdk>=0.22.0",
"sqlglot>=22.3.1,<22.5.0"
]
Expand Down
82 changes: 82 additions & 0 deletions src/databricks/labs/lsql/dashboards.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
import json
from pathlib import Path
from typing import Protocol, ClassVar, runtime_checkable

Check warning on line 3 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L1-L3

Added lines #L1 - L3 were not covered by tests

import sqlglot
import yaml
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.workspace import ExportFormat

Check warning on line 8 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L5-L8

Added lines #L5 - L8 were not covered by tests

from databricks.labs.lsql.lakeview import Dashboard, Page, Query, NamedQuery, ControlFieldEncoding

Check warning on line 10 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L10

Added line #L10 was not covered by tests

@runtime_checkable
class _DataclassInstance(Protocol):
__dataclass_fields__: ClassVar[dict]

Check warning on line 14 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L13-L14

Added lines #L13 - L14 were not covered by tests

class Dashboards:
def __init__(self, ws: WorkspaceClient):
self._ws = ws

Check warning on line 18 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L16-L18

Added lines #L16 - L18 were not covered by tests

def get_dashboard(self, dashboard_path: str):

Check warning on line 20 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L20

Added line #L20 was not covered by tests
with self._ws.workspace.download(dashboard_path, format=ExportFormat.SOURCE) as f:
raw = f.read().decode('utf-8')
as_dict = json.loads(raw)
return Dashboard.from_dict(as_dict)

Check warning on line 24 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L22-L24

Added lines #L22 - L24 were not covered by tests

def save_to_folder(self, dashboard_path: str, local_path: Path):
local_path.mkdir(parents=True, exist_ok=True)
dashboard = self.get_dashboard(dashboard_path)
better_names = {}

Check warning on line 29 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L26-L29

Added lines #L26 - L29 were not covered by tests
for dataset in dashboard.datasets:
name = dataset.display_name
better_names[dataset.name] = name
query_path = local_path / f"{name}.sql"
sql_query = dataset.query
self._format_sql_file(sql_query, query_path)
lvdash_yml = local_path / "lvdash.yml"

Check warning on line 36 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L31-L36

Added lines #L31 - L36 were not covered by tests
with lvdash_yml.open('w') as f:
first_page = dashboard.pages[0]
self._replace_names(first_page, better_names)
page = first_page.as_dict()
yaml.safe_dump(page, f)
assert True

Check warning on line 42 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L38-L42

Added lines #L38 - L42 were not covered by tests

def _format_sql_file(self, sql_query, query_path):

Check warning on line 44 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L44

Added line #L44 was not covered by tests
with query_path.open('w') as f:
try:

Check warning on line 46 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L46

Added line #L46 was not covered by tests
for statement in sqlglot.parse(sql_query):
# see https://sqlglot.com/sqlglot/generator.html#Generator
pretty = statement.sql(

Check warning on line 49 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L49

Added line #L49 was not covered by tests
dialect='databricks',
normalize=True, # normalize identifiers to lowercase
pretty=True, # format the produced SQL string
normalize_functions='upper', # normalize function names to uppercase
max_text_width=80, # wrap text at 120 characters
)
f.write(f"{pretty};\n")
except sqlglot.ParseError:
f.write(sql_query)

Check warning on line 58 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L56-L58

Added lines #L56 - L58 were not covered by tests

def _replace_names(self, node: _DataclassInstance, better_names: dict[str, str]):

Check warning on line 60 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L60

Added line #L60 was not covered by tests
# walk evely dataclass instance recursively and replace names
if isinstance(node, _DataclassInstance):
for field in node.__dataclass_fields__.values():
value = getattr(node, field.name)

Check warning on line 64 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L64

Added line #L64 was not covered by tests
if isinstance(value, list):
setattr(node, field.name, [self._replace_names(item, better_names) for item in value])
elif isinstance(value, _DataclassInstance):
setattr(node, field.name, self._replace_names(value, better_names))

Check warning on line 68 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L68

Added line #L68 was not covered by tests
if isinstance(node, Query):
node.dataset_name = better_names.get(node.dataset_name, node.dataset_name)

Check warning on line 70 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L70

Added line #L70 was not covered by tests
elif isinstance(node, NamedQuery) and node.query:
# 'dashboards/01eeb077e38c17e6ba3511036985960c/datasets/01eeb081882017f6a116991d124d3068_...'
if node.name.startswith('dashboards/'):
parts = [node.query.dataset_name]

Check warning on line 74 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L74

Added line #L74 was not covered by tests
for field in node.query.fields:
parts.append(field.name)
new_name = '_'.join(parts)
better_names[node.name] = new_name
node.name = better_names.get(node.name, node.name)

Check warning on line 79 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L76-L79

Added lines #L76 - L79 were not covered by tests
elif isinstance(node, ControlFieldEncoding):
node.query_name = better_names.get(node.query_name, node.query_name)
return node

Check warning on line 82 in src/databricks/labs/lsql/dashboards.py

View check run for this annotation

Codecov / codecov/patch

src/databricks/labs/lsql/dashboards.py#L81-L82

Added lines #L81 - L82 were not covered by tests
1 change: 1 addition & 0 deletions tests/integration/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
sample/*
13 changes: 13 additions & 0 deletions tests/integration/test_dashboards.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
from pathlib import Path

from databricks.sdk import WorkspaceClient

from databricks.labs.lsql.dashboards import Dashboards


def test_load_dashboard():
ws = WorkspaceClient(profile='logfood-master')
dashboards = Dashboards(ws)
src = "/Workspace/Users/[email protected]/Databricks Labs GitHub telemetry.lvdash.json"
dst = Path(__file__).parent / "sample"
dashboards.save_to_folder(src, dst)

0 comments on commit af99cca

Please sign in to comment.