Skip to content

Commit

Permalink
Fix issue when creating view using WITH clause (#1809)
Browse files Browse the repository at this point in the history
## Changes
Fix issue

### Linked issues
<!-- DOC: Link issue with a keyword: close, closes, closed, fix, fixes,
fixed, resolve, resolves, resolved. See
https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword
-->

Resolves #1798

### Functionality 

- [ ] added relevant user documentation
- [ ] added new CLI command
- [ ] modified existing command: `databricks labs ucx ...`
- [ ] added a new workflow
- [ ] modified existing workflow: `...`
- [ ] added a new table
- [ ] modified existing table: `...`

### Tests

- [ ] manually tested
- [x] added unit tests
- [ ] added integration tests
- [ ] verified on staging environment (screenshot attached)

Co-authored-by: Eric Vergnaud <eric.vergnaud@databricks.com>
  • Loading branch information
ericvergnaud and ericvergnaud authored May 30, 2024
1 parent 612d3e2 commit bcdf452
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 0 deletions.
11 changes: 11 additions & 0 deletions src/databricks/labs/ucx/hive_metastore/view_migrate.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,10 @@ def _view_dependencies(self):
if len(statements) != 1 or statements[0] is None:
raise ValueError(f"Could not analyze view SQL: {self.src.view_text}")
statement = statements[0]
aliases = self._read_aliases(statement)
for old_table in statement.find_all(expressions.Table):
if old_table.name in aliases:
continue
if old_table.catalog and old_table.catalog != 'hive_metastore':
continue
src_db = old_table.db if old_table.db else self.src.database
Expand All @@ -41,6 +44,14 @@ def _view_dependencies(self):
continue
yield TableView("hive_metastore", src_db, old_table.name)

def _read_aliases(self, statement: expressions.Expression):
aliases = set()
for with_clause in statement.find_all(expressions.With):
for expression in with_clause.expressions:
if isinstance(expression, expressions.CTE):
aliases.add(expression.alias_or_name)
return aliases

def sql_migrate_view(self, index: MigrationIndex) -> str:
from_table = FromTable(index, CurrentSessionState(self.src.database))
assert self.src.view_text is not None, 'Expected a view text'
Expand Down
12 changes: 12 additions & 0 deletions tests/unit/hive_metastore/test_table_migrate.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
MigrationStatusRefresher,
MigrationIndex,
MigrationStatus,
TableView,
)
from databricks.labs.ucx.hive_metastore.tables import (
AclMigrationWhat,
Expand All @@ -31,6 +32,7 @@
What,
)
from databricks.labs.ucx.hive_metastore.udfs import UdfsCrawler
from databricks.labs.ucx.hive_metastore.view_migrate import ViewToMigrate
from databricks.labs.ucx.workspace_access.groups import GroupManager

from .. import GROUPS, mock_table_mapping, mock_workspace_client
Expand Down Expand Up @@ -512,6 +514,16 @@ def test_migrate_view_should_produce_proper_queries(ws):
principal_grants.get_interactive_cluster_grants.assert_not_called()


def test_migrate_view_with_local_dataset_should_be_skipped(ws):
ddl = "CREATE VIEW v AS WITH t(a) AS (SELECT 1) SELECT * FROM t;"
to_migrate = ViewToMigrate(
Table("hive_metastore", "database", "", "EXTERNAL", "VIEW", None, ddl),
Rule("workspace", "catalog", "", "", "", ""),
)
dependencies = to_migrate.dependencies
assert dependencies == [TableView(catalog="hive_metastore", schema="database", name="v")]


def test_migrate_view_with_columns(ws):
errors = {}
create = "CREATE OR REPLACE VIEW hive_metastore.db1_src.view_src (a,b) AS SELECT * FROM db1_src.managed_dbfs"
Expand Down

0 comments on commit bcdf452

Please sign in to comment.