Skip to content

Commit

Permalink
MINOR - Fix lineage GET for names with / and standardize quote calls (
Browse files Browse the repository at this point in the history
#17748)

* MINOR - Fix lineage GET for names with `/` and standardize quote calls

* format

* fix import
  • Loading branch information
pmbrull authored Sep 6, 2024
1 parent a5dc937 commit 1a27645
Show file tree
Hide file tree
Showing 9 changed files with 100 additions and 15 deletions.
2 changes: 1 addition & 1 deletion ingestion/src/metadata/ingestion/ometa/mixins/es_mixin.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,11 @@
from typing import Generic, Iterable, List, Optional, Set, Type, TypeVar

from pydantic import BaseModel
from requests.utils import quote

from metadata.generated.schema.entity.data.container import Container
from metadata.generated.schema.entity.data.query import Query
from metadata.ingestion.ometa.client import REST, APIError
from metadata.ingestion.ometa.utils import quote
from metadata.utils.elasticsearch import ES_INDEX_MAP
from metadata.utils.logger import ometa_logger

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
from metadata.ingestion.lineage.parser import LINEAGE_PARSING_TIMEOUT
from metadata.ingestion.models.patch_request import build_patch
from metadata.ingestion.ometa.client import REST, APIError
from metadata.ingestion.ometa.utils import get_entity_type
from metadata.ingestion.ometa.utils import get_entity_type, quote
from metadata.utils.logger import ometa_logger
from metadata.utils.lru_cache import LRU_CACHE_SIZE, LRUCache

Expand Down Expand Up @@ -279,7 +279,7 @@ def get_lineage_by_name(
"""
return self._get_lineage(
entity=entity,
path=f"name/{fqn}",
path=f"name/{quote(fqn)}",
up_depth=up_depth,
down_depth=down_depth,
)
Expand Down
5 changes: 2 additions & 3 deletions ingestion/src/metadata/ingestion/ometa/mixins/table_mixin.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
from typing import List, Optional, Type, TypeVar

from pydantic import BaseModel, validate_call
from requests.utils import quote

from metadata.generated.schema.api.data.createTableProfile import (
CreateTableProfileRequest,
Expand All @@ -39,7 +38,7 @@
from metadata.generated.schema.type.usageRequest import UsageRequest
from metadata.ingestion.ometa.client import REST
from metadata.ingestion.ometa.models import EntityList
from metadata.ingestion.ometa.utils import model_str
from metadata.ingestion.ometa.utils import model_str, quote
from metadata.utils.logger import ometa_logger

logger = ometa_logger()
Expand Down Expand Up @@ -290,7 +289,7 @@ def get_latest_table_profile(
Returns:
Optional[Table]: OM table object
"""
return self._get(Table, f"{quote(model_str(fqn), safe='')}/tableProfile/latest")
return self._get(Table, f"{quote(fqn)}/tableProfile/latest")

def create_or_update_custom_metric(
self, custom_metric: CreateCustomMetricRequest, table_id: str
Expand Down
5 changes: 2 additions & 3 deletions ingestion/src/metadata/ingestion/ometa/mixins/tests_mixin.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
import traceback
from datetime import datetime
from typing import List, Optional, Type, Union
from urllib.parse import quote
from uuid import UUID

from metadata.generated.schema.api.tests.createLogicalTestCases import (
Expand Down Expand Up @@ -46,7 +45,7 @@
from metadata.generated.schema.tests.testSuite import TestSuite
from metadata.generated.schema.type.entityReference import EntityReference
from metadata.ingestion.ometa.client import REST
from metadata.ingestion.ometa.utils import model_str
from metadata.ingestion.ometa.utils import model_str, quote
from metadata.utils.logger import ometa_logger

logger = ometa_logger()
Expand Down Expand Up @@ -76,7 +75,7 @@ def add_test_case_results(
_type_: _description_
"""
resp = self.client.put(
f"{self.get_suffix(TestCase)}/{quote(test_case_fqn,safe='')}/testCaseResult",
f"{self.get_suffix(TestCase)}/{quote(test_case_fqn)}/testCaseResult",
test_results.model_dump_json(),
)

Expand Down
5 changes: 2 additions & 3 deletions ingestion/src/metadata/ingestion/ometa/ometa_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
from typing import Dict, Generic, Iterable, List, Optional, Type, TypeVar, Union

from pydantic import BaseModel
from requests.utils import quote

from metadata.generated.schema.api.services.ingestionPipelines.createIngestionPipeline import (
CreateIngestionPipelineRequest,
Expand Down Expand Up @@ -57,7 +56,7 @@
from metadata.ingestion.ometa.mixins.version_mixin import OMetaVersionMixin
from metadata.ingestion.ometa.models import EntityList
from metadata.ingestion.ometa.routes import ROUTES
from metadata.ingestion.ometa.utils import get_entity_type, model_str
from metadata.ingestion.ometa.utils import get_entity_type, model_str, quote
from metadata.utils.logger import ometa_logger
from metadata.utils.secrets.secrets_manager_factory import SecretsManagerFactory
from metadata.utils.ssl_registry import get_verify_ssl_fn
Expand Down Expand Up @@ -293,7 +292,7 @@ def get_by_name(

return self._get(
entity=entity,
path=f"name/{quote(model_str(fqn), safe='')}",
path=f"name/{quote(fqn)}",
fields=fields,
nullable=nullable,
)
Expand Down
11 changes: 11 additions & 0 deletions ingestion/src/metadata/ingestion/ometa/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@
from typing import Any, Type, TypeVar, Union

from pydantic import BaseModel
from requests.utils import quote as url_quote

from metadata.generated.schema.type.basic import FullyQualifiedEntityName

T = TypeVar("T", bound=BaseModel)

Expand Down Expand Up @@ -74,3 +77,11 @@ def model_str(arg: Any) -> str:
return str(arg.root)

return str(arg)


def quote(fqn: Union[FullyQualifiedEntityName, str]) -> str:
"""
Quote the FQN so that it's safe to pass to the API.
E.g., `"foo.bar/baz"` -> `%22foo.bar%2Fbaz%22`
"""
return url_quote(model_str(fqn), safe="")
54 changes: 51 additions & 3 deletions ingestion/tests/integration/ometa/test_ometa_lineage_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
from metadata.generated.schema.entity.services.dashboardService import DashboardService
from metadata.generated.schema.entity.services.databaseService import DatabaseService
from metadata.generated.schema.entity.services.pipelineService import PipelineService
from metadata.generated.schema.type.basic import EntityName
from metadata.generated.schema.type.entityLineage import (
ColumnLineage,
EntitiesEdge,
Expand Down Expand Up @@ -85,7 +86,7 @@ def setUpClass(cls) -> None:
)
)

create_schema_entity = cls.metadata.create_or_update(
cls.create_schema_entity = cls.metadata.create_or_update(
data=get_create_entity(
entity=DatabaseSchema,
reference=create_db_entity.fullyQualifiedName,
Expand All @@ -96,14 +97,14 @@ def setUpClass(cls) -> None:
cls.table1 = get_create_entity(
name=generate_name(),
entity=Table,
reference=create_schema_entity.fullyQualifiedName,
reference=cls.create_schema_entity.fullyQualifiedName,
)

cls.table1_entity = cls.metadata.create_or_update(data=cls.table1)
cls.table2 = get_create_entity(
name=generate_name(),
entity=Table,
reference=create_schema_entity.fullyQualifiedName,
reference=cls.create_schema_entity.fullyQualifiedName,
)

cls.table2_entity = cls.metadata.create_or_update(data=cls.table2)
Expand Down Expand Up @@ -337,3 +338,50 @@ def test_table_datamodel_lineage(self):
)
entity_lineage = EntityLineage.model_validate(datamodel_lineage)
self.assertEqual(from_id, str(entity_lineage.upstreamEdges[0].fromEntity.root))

def test_table_with_slash_in_name(self):
"""E.g., `foo.bar/baz`"""
name = EntityName("foo.bar/baz")
new_table: Table = self.metadata.create_or_update(
data=get_create_entity(
entity=Table,
name=name,
reference=self.create_schema_entity.fullyQualifiedName,
)
)

res: Table = self.metadata.get_by_name(
entity=Table, fqn=new_table.fullyQualifiedName
)

assert res.name == name

self.metadata.add_lineage(
data=AddLineageRequest(
edge=EntitiesEdge(
fromEntity=EntityReference(id=self.table1_entity.id, type="table"),
toEntity=EntityReference(id=new_table.id, type="table"),
lineageDetails=LineageDetails(
columnsLineage=[
ColumnLineage(
fromColumns=[
self.table1_entity.columns[0].fullyQualifiedName
],
toColumn=new_table.columns[0].fullyQualifiedName,
)
]
),
),
)
)

# use the SDK to get the lineage
lineage = self.metadata.get_lineage_by_name(
entity=Table,
fqn=new_table.fullyQualifiedName.root,
)
entity_lineage = EntityLineage.model_validate(lineage)
assert (
entity_lineage.upstreamEdges[0].fromEntity.root
== self.table1_entity.id.root
)
19 changes: 19 additions & 0 deletions ingestion/tests/integration/ometa/test_ometa_table_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,8 @@
from metadata.generated.schema.type.usageRequest import UsageRequest
from metadata.ingestion.ometa.client import REST

from ..integration_base import get_create_entity

BAD_RESPONSE = {
"data": [
{
Expand Down Expand Up @@ -643,3 +645,20 @@ def test_list_all_w_skip_on_failure(self):

# We should have 2 tables, the 3rd one is broken and should be skipped
assert len(list(res)) == 2

def test_table_with_slash_in_name(self):
"""E.g., `foo.bar/baz`"""
name = EntityName("foo.bar/baz")
new_table: Table = self.metadata.create_or_update(
data=get_create_entity(
entity=Table,
name=name,
reference=self.create_schema_entity.fullyQualifiedName,
)
)

res: Table = self.metadata.get_by_name(
entity=Table, fqn=new_table.fullyQualifiedName
)

assert res.name == name
10 changes: 10 additions & 0 deletions ingestion/tests/unit/test_fqn.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
import pytest

from metadata.generated.schema.entity.data.table import Table
from metadata.generated.schema.type.basic import FullyQualifiedEntityName
from metadata.ingestion.ometa.utils import quote
from metadata.utils import fqn


Expand Down Expand Up @@ -145,3 +147,11 @@ def test_split_test_case_fqn(self):

with pytest.raises(ValueError):
fqn.split_test_case_fqn("local_redshift.dev.dbt_jaffle.customers")

def test_quote_fqns(self):
"""We can properly quote FQNs for URL usage"""
assert quote(FullyQualifiedEntityName("a.b.c")) == "a.b.c"
# Works with strings directly
assert quote("a.b.c") == "a.b.c"
assert quote(FullyQualifiedEntityName('"foo.bar".baz')) == "%22foo.bar%22.baz"
assert quote('"foo.bar/baz".hello') == "%22foo.bar%2Fbaz%22.hello"

0 comments on commit 1a27645

Please sign in to comment.