From 303a2d0863d7a86e96dd3c32655f1a044dc6bffe Mon Sep 17 00:00:00 2001 From: Mayuri Nehate <33225191+mayurinehate@users.noreply.github.com> Date: Wed, 13 Sep 2023 00:00:24 +0530 Subject: [PATCH 001/156] build(ingest): upgrade to sqlalchemy 1.4, drop 1.3 support (#8810) Co-authored-by: Harshal Sheth --- docs/how/updating-datahub.md | 1 + metadata-ingestion/build.gradle | 3 -- .../scripts/install-sqlalchemy-stubs.sh | 28 --------------- metadata-ingestion/setup.py | 25 ++++++------- .../source/datahub/datahub_database_reader.py | 6 +--- .../source/snowflake/snowflake_usage_v2.py | 9 +---- .../ingestion/source/sql/clickhouse.py | 35 +------------------ .../source/usage/clickhouse_usage.py | 6 +--- .../ingestion/source/usage/redshift_usage.py | 4 +-- .../source/usage/starburst_trino_usage.py | 6 +--- 10 files changed, 17 insertions(+), 106 deletions(-) delete mode 100755 metadata-ingestion/scripts/install-sqlalchemy-stubs.sh diff --git a/docs/how/updating-datahub.md b/docs/how/updating-datahub.md index 1ef7413a88ebd..9b19291ee246a 100644 --- a/docs/how/updating-datahub.md +++ b/docs/how/updating-datahub.md @@ -5,6 +5,7 @@ This file documents any backwards-incompatible changes in DataHub and assists pe ## Next ### Breaking Changes +- #8810 - Removed support for SQLAlchemy 1.3.x. Only SQLAlchemy 1.4.x is supported now. ### Potential Downtime diff --git a/metadata-ingestion/build.gradle b/metadata-ingestion/build.gradle index 199ccc59c21e0..408ea771bc93f 100644 --- a/metadata-ingestion/build.gradle +++ b/metadata-ingestion/build.gradle @@ -71,7 +71,6 @@ task installDev(type: Exec, dependsOn: [install]) { commandLine 'bash', '-c', "source ${venv_name}/bin/activate && set -x && " + "${venv_name}/bin/pip install -e .[dev] ${extra_pip_requirements} && " + - "./scripts/install-sqlalchemy-stubs.sh && " + "touch ${sentinel_file}" } @@ -82,7 +81,6 @@ task installAll(type: Exec, dependsOn: [install]) { commandLine 'bash', '-c', "source ${venv_name}/bin/activate && set -x && " + "${venv_name}/bin/pip install -e .[all] ${extra_pip_requirements} && " + - "./scripts/install-sqlalchemy-stubs.sh && " + "touch ${sentinel_file}" } @@ -119,7 +117,6 @@ task lint(type: Exec, dependsOn: installDev) { task lintFix(type: Exec, dependsOn: installDev) { commandLine 'bash', '-c', "source ${venv_name}/bin/activate && set -x && " + - "./scripts/install-sqlalchemy-stubs.sh && " + "black src/ tests/ examples/ && " + "isort src/ tests/ examples/ && " + "flake8 src/ tests/ examples/ && " + diff --git a/metadata-ingestion/scripts/install-sqlalchemy-stubs.sh b/metadata-ingestion/scripts/install-sqlalchemy-stubs.sh deleted file mode 100755 index 7c14a06464f99..0000000000000 --- a/metadata-ingestion/scripts/install-sqlalchemy-stubs.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/bin/bash - -set -euo pipefail - -# ASSUMPTION: This assumes that we're running from inside the venv. - -SQLALCHEMY_VERSION=$(python -c 'import sqlalchemy; print(sqlalchemy.__version__)') - -if [[ $SQLALCHEMY_VERSION == 1.3.* ]]; then - ENSURE_NOT_INSTALLED=sqlalchemy2-stubs - ENSURE_INSTALLED=sqlalchemy-stubs -elif [[ $SQLALCHEMY_VERSION == 1.4.* ]]; then - ENSURE_NOT_INSTALLED=sqlalchemy-stubs - ENSURE_INSTALLED=sqlalchemy2-stubs -else - echo "Unsupported SQLAlchemy version: $SQLALCHEMY_VERSION" - exit 1 -fi - -FORCE_REINSTALL="" -if pip show $ENSURE_NOT_INSTALLED >/dev/null 2>&1 ; then - pip uninstall --yes $ENSURE_NOT_INSTALLED - FORCE_REINSTALL="--force-reinstall" -fi - -if [ -n "$FORCE_REINSTALL" ] || ! pip show $ENSURE_INSTALLED >/dev/null 2>&1 ; then - pip install $FORCE_REINSTALL $ENSURE_INSTALLED -fi diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py index d8668e8925546..09f71fa769fd3 100644 --- a/metadata-ingestion/setup.py +++ b/metadata-ingestion/setup.py @@ -112,7 +112,8 @@ def get_long_description(): sql_common = { # Required for all SQL sources. - "sqlalchemy>=1.3.24, <2", + # This is temporary lower bound that we're open to loosening/tightening as requirements show up + "sqlalchemy>=1.4.39, <2", # Required for SQL profiling. "great-expectations>=0.15.12, <=0.15.50", # scipy version restricted to reduce backtracking, used by great-expectations, @@ -172,13 +173,13 @@ def get_long_description(): } clickhouse_common = { - # Clickhouse 0.1.8 requires SQLAlchemy 1.3.x, while the newer versions - # allow SQLAlchemy 1.4.x. - "clickhouse-sqlalchemy>=0.1.8", + # Clickhouse 0.2.0 adds support for SQLAlchemy 1.4.x + "clickhouse-sqlalchemy>=0.2.0", } redshift_common = { - "sqlalchemy-redshift", + # Clickhouse 0.8.3 adds support for SQLAlchemy 1.4.x + "sqlalchemy-redshift>=0.8.3", "psycopg2-binary", "GeoAlchemy2", *sqllineage_lib, @@ -188,13 +189,8 @@ def get_long_description(): snowflake_common = { # Snowflake plugin utilizes sql common *sql_common, - # Required for all Snowflake sources. - # See https://github.com/snowflakedb/snowflake-sqlalchemy/issues/234 for why 1.2.5 is blocked. - "snowflake-sqlalchemy>=1.2.4, !=1.2.5", - # Because of https://github.com/snowflakedb/snowflake-sqlalchemy/issues/350 we need to restrict SQLAlchemy's max version. - # Eventually we should just require snowflake-sqlalchemy>=1.4.3, but I won't do that immediately - # because it may break Airflow users that need SQLAlchemy 1.3.x. - "SQLAlchemy<1.4.42", + # https://github.com/snowflakedb/snowflake-sqlalchemy/issues/350 + "snowflake-sqlalchemy>=1.4.3", # See https://github.com/snowflakedb/snowflake-connector-python/pull/1348 for why 2.8.2 is blocked "snowflake-connector-python!=2.8.2", "pandas", @@ -206,9 +202,7 @@ def get_long_description(): } trino = { - # Trino 0.317 broke compatibility with SQLAlchemy 1.3.24. - # See https://github.com/trinodb/trino-python-client/issues/250. - "trino[sqlalchemy]>=0.308, !=0.317", + "trino[sqlalchemy]>=0.308", } pyhive_common = { @@ -430,6 +424,7 @@ def get_long_description(): "types-Deprecated", "types-protobuf>=4.21.0.1", "types-tzlocal", + "sqlalchemy2-stubs", } diff --git a/metadata-ingestion/src/datahub/ingestion/source/datahub/datahub_database_reader.py b/metadata-ingestion/src/datahub/ingestion/source/datahub/datahub_database_reader.py index a5aadbd6e246b..96184d8d445e4 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/datahub/datahub_database_reader.py +++ b/metadata-ingestion/src/datahub/ingestion/source/datahub/datahub_database_reader.py @@ -69,11 +69,7 @@ def get_aspects( return for i, row in enumerate(rows): - # TODO: Replace with namedtuple usage once we drop sqlalchemy 1.3 - if hasattr(row, "_asdict"): - row_dict = row._asdict() - else: - row_dict = dict(row) + row_dict = row._asdict() mcp = self._parse_row(row_dict) if mcp: yield mcp, row_dict["createdon"] diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_usage_v2.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_usage_v2.py index d041d219c4bdd..1cbd4a3b3ea24 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_usage_v2.py +++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_usage_v2.py @@ -451,17 +451,10 @@ def _get_operation_aspect_work_unit( yield wu def _process_snowflake_history_row( - self, row: Any + self, event_dict: dict ) -> Iterable[SnowflakeJoinedAccessEvent]: try: # big hammer try block to ensure we don't fail on parsing events self.report.rows_processed += 1 - # Make some minor type conversions. - if hasattr(row, "_asdict"): - # Compat with SQLAlchemy 1.3 and 1.4 - # See https://docs.sqlalchemy.org/en/14/changelog/migration_14.html#rowproxy-is-no-longer-a-proxy-is-now-called-row-and-behaves-like-an-enhanced-named-tuple. - event_dict = row._asdict() - else: - event_dict = dict(row) # no use processing events that don't have a query text if not event_dict["QUERY_TEXT"]: diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py b/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py index 20130ef21e5e6..1626f86b92545 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py @@ -38,7 +38,6 @@ logger, register_custom_type, ) -from datahub.ingestion.source.sql.sql_config import make_sqlalchemy_uri from datahub.ingestion.source.sql.two_tier_sql_source import ( TwoTierSQLAlchemyConfig, TwoTierSQLAlchemySource, @@ -147,7 +146,6 @@ class ClickHouseConfig( include_materialized_views: Optional[bool] = Field(default=True, description="") def get_sql_alchemy_url(self, current_db=None): - url = make_url( super().get_sql_alchemy_url(uri_opts=self.uri_opts, current_db=current_db) ) @@ -158,42 +156,11 @@ def get_sql_alchemy_url(self, current_db=None): ) # We can setup clickhouse ingestion in sqlalchemy_uri form and config form. - - # If we use sqlalchemu_uri form then super().get_sql_alchemy_url doesn't - # update current_db because it return self.sqlalchemy_uri without any update. - # This code bellow needed for rewriting sqlalchemi_uri and replace database with current_db.from - # For the future without python3.7 and sqlalchemy 1.3 support we can use code - # url=url.set(db=current_db), but not now. - # Why we need to update database in uri at all? # Because we get database from sqlalchemy inspector and inspector we form from url inherited from # TwoTierSQLAlchemySource and SQLAlchemySource - if self.sqlalchemy_uri and current_db: - self.scheme = url.drivername - self.username = url.username - self.password = ( - pydantic.SecretStr(str(url.password)) - if url.password - else pydantic.SecretStr("") - ) - if url.host and url.port: - self.host_port = url.host + ":" + str(url.port) - elif url.host: - self.host_port = url.host - # untill released https://github.com/python/mypy/pull/15174 - self.uri_opts = {str(k): str(v) for (k, v) in url.query.items()} - - url = make_url( - make_sqlalchemy_uri( - self.scheme, - self.username, - self.password.get_secret_value() if self.password else None, - self.host_port, - current_db if current_db else self.database, - uri_opts=self.uri_opts, - ) - ) + url = url.set(database=current_db) return str(url) diff --git a/metadata-ingestion/src/datahub/ingestion/source/usage/clickhouse_usage.py b/metadata-ingestion/src/datahub/ingestion/source/usage/clickhouse_usage.py index 855958f0755e1..f659ea0c1c5c0 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/usage/clickhouse_usage.py +++ b/metadata-ingestion/src/datahub/ingestion/source/usage/clickhouse_usage.py @@ -143,11 +143,7 @@ def _get_clickhouse_history(self): results = engine.execute(query) events = [] for row in results: - # minor type conversion - if hasattr(row, "_asdict"): - event_dict = row._asdict() - else: - event_dict = dict(row) + event_dict = row._asdict() # stripping extra spaces caused by above _asdict() conversion for k, v in event_dict.items(): diff --git a/metadata-ingestion/src/datahub/ingestion/source/usage/redshift_usage.py b/metadata-ingestion/src/datahub/ingestion/source/usage/redshift_usage.py index 99a980b326e53..691eaa8211054 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/usage/redshift_usage.py +++ b/metadata-ingestion/src/datahub/ingestion/source/usage/redshift_usage.py @@ -298,9 +298,7 @@ def _gen_access_events_from_history_query( for row in results: if not self._should_process_row(row): continue - if hasattr(row, "_asdict"): - # Compatibility with sqlalchemy 1.4.x. - row = row._asdict() + row = row._asdict() access_event = RedshiftAccessEvent(**dict(row.items())) # Replace database name with the alias name if one is provided in the config. if self.config.database_alias: diff --git a/metadata-ingestion/src/datahub/ingestion/source/usage/starburst_trino_usage.py b/metadata-ingestion/src/datahub/ingestion/source/usage/starburst_trino_usage.py index 9394a8bba5e0b..c38800b3a6983 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/usage/starburst_trino_usage.py +++ b/metadata-ingestion/src/datahub/ingestion/source/usage/starburst_trino_usage.py @@ -162,11 +162,7 @@ def _get_trino_history(self): results = engine.execute(query) events = [] for row in results: - # minor type conversion - if hasattr(row, "_asdict"): - event_dict = row._asdict() - else: - event_dict = dict(row) + event_dict = row._asdict() # stripping extra spaces caused by above _asdict() conversion for k, v in event_dict.items(): From f7fee743bfddf27f072e5c56512ef905d942eab6 Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Tue, 12 Sep 2023 13:11:01 -0700 Subject: [PATCH 002/156] fix(ingest): use epoch 1 for dev build versions (#8824) --- docker/datahub-ingestion-base/smoke.Dockerfile | 2 +- docker/datahub-ingestion/Dockerfile | 4 ++-- docker/datahub-ingestion/Dockerfile-slim-only | 2 +- metadata-ingestion-modules/airflow-plugin/scripts/release.sh | 2 +- .../airflow-plugin/src/datahub_airflow_plugin/__init__.py | 2 +- metadata-ingestion/scripts/release.sh | 2 +- metadata-ingestion/src/datahub/__init__.py | 2 +- 7 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docker/datahub-ingestion-base/smoke.Dockerfile b/docker/datahub-ingestion-base/smoke.Dockerfile index 276f6dbc4436e..15dc46ae5b882 100644 --- a/docker/datahub-ingestion-base/smoke.Dockerfile +++ b/docker/datahub-ingestion-base/smoke.Dockerfile @@ -20,7 +20,7 @@ RUN DEBIAN_FRONTEND=noninteractive apt-get install -y openjdk-11-jdk COPY . /datahub-src ARG RELEASE_VERSION RUN cd /datahub-src/metadata-ingestion && \ - sed -i.bak "s/__version__ = \"0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" src/datahub/__init__.py && \ + sed -i.bak "s/__version__ = \"1!0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" src/datahub/__init__.py && \ cat src/datahub/__init__.py && \ cd ../ && \ ./gradlew :metadata-ingestion:installAll diff --git a/docker/datahub-ingestion/Dockerfile b/docker/datahub-ingestion/Dockerfile index 2ceff6a800ebb..8b726df5e8842 100644 --- a/docker/datahub-ingestion/Dockerfile +++ b/docker/datahub-ingestion/Dockerfile @@ -11,8 +11,8 @@ COPY ./metadata-ingestion-modules/airflow-plugin /datahub-ingestion/airflow-plug ARG RELEASE_VERSION WORKDIR /datahub-ingestion -RUN sed -i.bak "s/__version__ = \"0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" src/datahub/__init__.py && \ - sed -i.bak "s/__version__ = \"0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" airflow-plugin/src/datahub_airflow_plugin/__init__.py && \ +RUN sed -i.bak "s/__version__ = \"1!0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" src/datahub/__init__.py && \ + sed -i.bak "s/__version__ = \"1!0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" airflow-plugin/src/datahub_airflow_plugin/__init__.py && \ cat src/datahub/__init__.py && \ chown -R datahub /datahub-ingestion diff --git a/docker/datahub-ingestion/Dockerfile-slim-only b/docker/datahub-ingestion/Dockerfile-slim-only index 678bee7e306f6..9ae116f839aa0 100644 --- a/docker/datahub-ingestion/Dockerfile-slim-only +++ b/docker/datahub-ingestion/Dockerfile-slim-only @@ -9,7 +9,7 @@ COPY ./metadata-ingestion /datahub-ingestion ARG RELEASE_VERSION WORKDIR /datahub-ingestion -RUN sed -i.bak "s/__version__ = \"0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" src/datahub/__init__.py && \ +RUN sed -i.bak "s/__version__ = \"1!0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" src/datahub/__init__.py && \ cat src/datahub/__init__.py && \ chown -R datahub /datahub-ingestion diff --git a/metadata-ingestion-modules/airflow-plugin/scripts/release.sh b/metadata-ingestion-modules/airflow-plugin/scripts/release.sh index 7134187a45885..87157479f37d6 100755 --- a/metadata-ingestion-modules/airflow-plugin/scripts/release.sh +++ b/metadata-ingestion-modules/airflow-plugin/scripts/release.sh @@ -13,7 +13,7 @@ MODULE=datahub_airflow_plugin python -c 'import setuptools; where="./src"; assert setuptools.find_packages(where) == setuptools.find_namespace_packages(where), "you seem to be missing or have extra __init__.py files"' if [[ ${RELEASE_VERSION:-} ]]; then # Replace version with RELEASE_VERSION env variable - sed -i.bak "s/__version__ = \"0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" src/${MODULE}/__init__.py + sed -i.bak "s/__version__ = \"1!0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" src/${MODULE}/__init__.py else vim src/${MODULE}/__init__.py fi diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/__init__.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/__init__.py index ce98a0fc1fb60..b2c45d3a1e75d 100644 --- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/__init__.py +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/__init__.py @@ -1,6 +1,6 @@ # Published at https://pypi.org/project/acryl-datahub/. __package_name__ = "acryl-datahub-airflow-plugin" -__version__ = "0.0.0.dev0" +__version__ = "1!0.0.0.dev0" def is_dev_mode() -> bool: diff --git a/metadata-ingestion/scripts/release.sh b/metadata-ingestion/scripts/release.sh index 0a09c4e0307b3..eacaf1d920a8d 100755 --- a/metadata-ingestion/scripts/release.sh +++ b/metadata-ingestion/scripts/release.sh @@ -11,7 +11,7 @@ fi python -c 'import setuptools; where="./src"; assert setuptools.find_packages(where) == setuptools.find_namespace_packages(where), "you seem to be missing or have extra __init__.py files"' if [[ ${RELEASE_VERSION:-} ]]; then # Replace version with RELEASE_VERSION env variable - sed -i.bak "s/__version__ = \"0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" src/datahub/__init__.py + sed -i.bak "s/__version__ = \"1!0.0.0.dev0\"/__version__ = \"$RELEASE_VERSION\"/" src/datahub/__init__.py else vim src/datahub/__init__.py fi diff --git a/metadata-ingestion/src/datahub/__init__.py b/metadata-ingestion/src/datahub/__init__.py index 3ac3efefc14f0..a470de7b500be 100644 --- a/metadata-ingestion/src/datahub/__init__.py +++ b/metadata-ingestion/src/datahub/__init__.py @@ -3,7 +3,7 @@ # Published at https://pypi.org/project/acryl-datahub/. __package_name__ = "acryl-datahub" -__version__ = "0.0.0.dev0" +__version__ = "1!0.0.0.dev0" def is_dev_mode() -> bool: From 449cc9ba91bfc51bc8e5a66de7920340f164f272 Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Tue, 12 Sep 2023 13:15:05 -0700 Subject: [PATCH 003/156] ci: make wheel builds more robust (#8815) --- docs-website/sphinx/Makefile | 5 ++++- docs-website/sphinx/requirements.txt | 2 +- docs-website/yarn.lock | 18 +++++++++++------- .../airflow-plugin/build.gradle | 6 +++--- metadata-ingestion/build.gradle | 6 +++--- 5 files changed, 22 insertions(+), 15 deletions(-) diff --git a/docs-website/sphinx/Makefile b/docs-website/sphinx/Makefile index 00ece7ae25331..c01b45e322c67 100644 --- a/docs-website/sphinx/Makefile +++ b/docs-website/sphinx/Makefile @@ -22,7 +22,7 @@ $(VENV_SENTINEL): requirements.txt $(VENV_DIR)/bin/pip install -r requirements.txt touch $(VENV_SENTINEL) -.PHONY: help html doctest linkcheck clean serve md +.PHONY: help html doctest linkcheck clean clean_all serve md # Not using Python's http.server because it enables caching headers. serve: @@ -35,3 +35,6 @@ md: html # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). html doctest linkcheck clean: venv Makefile @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +clean_all: clean + -rm -rf $(VENV_DIR) diff --git a/docs-website/sphinx/requirements.txt b/docs-website/sphinx/requirements.txt index a63fd05853259..94ddd40579f0e 100644 --- a/docs-website/sphinx/requirements.txt +++ b/docs-website/sphinx/requirements.txt @@ -1,4 +1,4 @@ --e ../../metadata-ingestion[datahub-rest] +-e ../../metadata-ingestion[datahub-rest,sql-parsing] beautifulsoup4==4.11.2 Sphinx==6.1.3 sphinx-click==4.4.0 diff --git a/docs-website/yarn.lock b/docs-website/yarn.lock index 209a57a43dab0..0613fe71ef78e 100644 --- a/docs-website/yarn.lock +++ b/docs-website/yarn.lock @@ -2986,6 +2986,13 @@ dependencies: "@types/node" "*" +"@types/websocket@^1.0.3": + version "1.0.6" + resolved "https://registry.yarnpkg.com/@types/websocket/-/websocket-1.0.6.tgz#ec8dce5915741632ac3a4b1f951b6d4156e32d03" + integrity sha512-JXkliwz93B2cMWOI1ukElQBPN88vMg3CruvW4KVSKpflt3NyNCJImnhIuB/f97rG7kakqRJGFiwkA895Kn02Dg== + dependencies: + "@types/node" "*" + "@types/ws@^8.5.5": version "8.5.5" resolved "https://registry.yarnpkg.com/@types/ws/-/ws-8.5.5.tgz#af587964aa06682702ee6dcbc7be41a80e4b28eb" @@ -7053,7 +7060,6 @@ node-forge@^1: resolved "https://registry.yarnpkg.com/node-forge/-/node-forge-1.3.1.tgz#be8da2af243b2417d5f646a770663a92b7e9ded3" integrity sha512-dPEtOeMvF9VMcYV/1Wb8CPoVAXtp6MKMlcbAt4ddqmGqUJ6fQZFXkNZNkNlfevtNkGtaSoXf/vNNNSvgrdXwtA== - node-gyp-build@^4.3.0: version "4.6.1" resolved "https://registry.yarnpkg.com/node-gyp-build/-/node-gyp-build-4.6.1.tgz#24b6d075e5e391b8d5539d98c7fc5c210cac8a3e" @@ -9903,6 +9909,10 @@ use-sidecar@^1.1.2: detect-node-es "^1.1.0" tslib "^2.0.0" +use-sync-external-store@^1.2.0: + version "1.2.0" + resolved "https://registry.yarnpkg.com/use-sync-external-store/-/use-sync-external-store-1.2.0.tgz#7dbefd6ef3fe4e767a0cf5d7287aacfb5846928a" + integrity sha512-eEgnFxGQ1Ife9bzYs6VLi8/4X6CObHMw9Qr9tPY43iKwsPw8xE8+EFsf/2cFZ5S3esXgpWgtSCtLNS41F+sKPA== utf-8-validate@^5.0.2: version "5.0.10" @@ -9911,12 +9921,6 @@ utf-8-validate@^5.0.2: dependencies: node-gyp-build "^4.3.0" -use-sync-external-store@^1.2.0: - version "1.2.0" - resolved "https://registry.yarnpkg.com/use-sync-external-store/-/use-sync-external-store-1.2.0.tgz#7dbefd6ef3fe4e767a0cf5d7287aacfb5846928a" - integrity sha512-eEgnFxGQ1Ife9bzYs6VLi8/4X6CObHMw9Qr9tPY43iKwsPw8xE8+EFsf/2cFZ5S3esXgpWgtSCtLNS41F+sKPA== - - util-deprecate@^1.0.1, util-deprecate@^1.0.2, util-deprecate@~1.0.1: version "1.0.2" resolved "https://registry.yarnpkg.com/util-deprecate/-/util-deprecate-1.0.2.tgz#450d4dc9fa70de732762fbd2d4a28981419a0ccf" diff --git a/metadata-ingestion-modules/airflow-plugin/build.gradle b/metadata-ingestion-modules/airflow-plugin/build.gradle index d1e6f2f646491..58a2bc9e670e3 100644 --- a/metadata-ingestion-modules/airflow-plugin/build.gradle +++ b/metadata-ingestion-modules/airflow-plugin/build.gradle @@ -110,14 +110,14 @@ task testFull(type: Exec, dependsOn: [testQuick, installDevTest]) { commandLine 'bash', '-x', '-c', "source ${venv_name}/bin/activate && pytest -m 'not slow_integration' -vv --continue-on-collection-errors --junit-xml=junit.full.xml" } -task buildWheel(type: Exec, dependsOn: [install]) { - commandLine 'bash', '-c', "source ${venv_name}/bin/activate && " + 'pip install build && RELEASE_VERSION="\${RELEASE_VERSION:-0.0.0.dev1}" RELEASE_SKIP_TEST=1 RELEASE_SKIP_UPLOAD=1 ./scripts/release.sh' -} task cleanPythonCache(type: Exec) { commandLine 'bash', '-c', "find src -type f -name '*.py[co]' -delete -o -type d -name __pycache__ -delete -o -type d -empty -delete" } +task buildWheel(type: Exec, dependsOn: [install, cleanPythonCache]) { + commandLine 'bash', '-c', "source ${venv_name}/bin/activate && " + 'pip install build && RELEASE_VERSION="\${RELEASE_VERSION:-0.0.0.dev1}" RELEASE_SKIP_TEST=1 RELEASE_SKIP_UPLOAD=1 ./scripts/release.sh' +} build.dependsOn install check.dependsOn lint diff --git a/metadata-ingestion/build.gradle b/metadata-ingestion/build.gradle index 408ea771bc93f..c20d98cbcbb58 100644 --- a/metadata-ingestion/build.gradle +++ b/metadata-ingestion/build.gradle @@ -185,9 +185,6 @@ task specGen(type: Exec, dependsOn: [codegen, installDevTest]) { task docGen(type: Exec, dependsOn: [codegen, installDevTest, specGen]) { commandLine 'bash', '-c', "source ${venv_name}/bin/activate && ./scripts/docgen.sh" } -task buildWheel(type: Exec, dependsOn: [install, codegen]) { - commandLine 'bash', '-c', "source ${venv_name}/bin/activate && " + 'pip install build && RELEASE_VERSION="\${RELEASE_VERSION:-0.0.0.dev1}" RELEASE_SKIP_TEST=1 RELEASE_SKIP_UPLOAD=1 ./scripts/release.sh' -} @@ -195,6 +192,9 @@ task cleanPythonCache(type: Exec) { commandLine 'bash', '-c', "find src tests -type f -name '*.py[co]' -delete -o -type d -name __pycache__ -delete -o -type d -empty -delete" } +task buildWheel(type: Exec, dependsOn: [install, codegen, cleanPythonCache]) { + commandLine 'bash', '-c', "source ${venv_name}/bin/activate && " + 'pip install build && RELEASE_VERSION="\${RELEASE_VERSION:-0.0.0.dev1}" RELEASE_SKIP_TEST=1 RELEASE_SKIP_UPLOAD=1 ./scripts/release.sh' +} build.dependsOn install check.dependsOn lint From 138f6c0f74a4799d31560e9fde19ef6011089990 Mon Sep 17 00:00:00 2001 From: Pedro Silva Date: Tue, 12 Sep 2023 22:26:30 +0100 Subject: [PATCH 004/156] feat(cli): fix upload ingest cli endpoint (#8826) --- metadata-ingestion/src/datahub/cli/ingest_cli.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/metadata-ingestion/src/datahub/cli/ingest_cli.py b/metadata-ingestion/src/datahub/cli/ingest_cli.py index 42c0ea1601c74..5931bf89b010b 100644 --- a/metadata-ingestion/src/datahub/cli/ingest_cli.py +++ b/metadata-ingestion/src/datahub/cli/ingest_cli.py @@ -282,12 +282,14 @@ def deploy( "urn": urn, "name": name, "type": pipeline_config["source"]["type"], - "schedule": {"interval": schedule, "timezone": time_zone}, "recipe": json.dumps(pipeline_config), "executorId": executor_id, "version": cli_version, } + if schedule is not None: + variables["schedule"] = {"interval": schedule, "timezone": time_zone} + if urn: if not datahub_graph.exists(urn): logger.error(f"Could not find recipe for provided urn: {urn}") @@ -331,6 +333,7 @@ def deploy( $version: String) { createIngestionSource(input: { + name: $name, type: $type, schedule: $schedule, config: { From 3cc0f76d178f239acc018e06ec408eb6b38bfb5d Mon Sep 17 00:00:00 2001 From: Adriano Vega Llobell Date: Tue, 12 Sep 2023 23:34:24 +0200 Subject: [PATCH 005/156] docs(transformer): fix names in sample code of 'pattern_add_dataset_domain' (#8755) --- metadata-ingestion/docs/transformer/dataset_transformer.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/metadata-ingestion/docs/transformer/dataset_transformer.md b/metadata-ingestion/docs/transformer/dataset_transformer.md index cb06656940918..f0fa44687a109 100644 --- a/metadata-ingestion/docs/transformer/dataset_transformer.md +++ b/metadata-ingestion/docs/transformer/dataset_transformer.md @@ -909,7 +909,7 @@ in both of the cases domain should be provisioned on DataHub GMS - Add domains, however replace existing domains sent by ingestion source ```yaml transformers: - - type: "pattern_add_dataset_ownership" + - type: "pattern_add_dataset_domain" config: replace_existing: true # false is default behaviour domain_pattern: @@ -920,7 +920,7 @@ in both of the cases domain should be provisioned on DataHub GMS - Add domains, however overwrite the domains available for the dataset on DataHub GMS ```yaml transformers: - - type: "pattern_add_dataset_ownership" + - type: "pattern_add_dataset_domain" config: semantics: OVERWRITE # OVERWRITE is default behaviour domain_pattern: @@ -931,7 +931,7 @@ in both of the cases domain should be provisioned on DataHub GMS - Add domains, however keep the domains available for the dataset on DataHub GMS ```yaml transformers: - - type: "pattern_add_dataset_ownership" + - type: "pattern_add_dataset_domain" config: semantics: PATCH domain_pattern: From 785ab7718df8e4e46bdd612ed3deaafbda1d42cc Mon Sep 17 00:00:00 2001 From: ethan-cartwright Date: Wed, 13 Sep 2023 03:45:58 -0400 Subject: [PATCH 006/156] fix(siblingsHook): check number of dbtUpstreams instead of all upStreams (#8817) Co-authored-by: Ethan Cartwright --- .../hook/siblings/SiblingAssociationHook.java | 19 ++- .../siblings/SiblingAssociationHookTest.java | 112 ++++++++++++++---- 2 files changed, 100 insertions(+), 31 deletions(-) diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java index 2be719ed263ea..06545ef3525dd 100644 --- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java +++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java @@ -200,10 +200,19 @@ private void handleSourceDatasetEvent(MetadataChangeLog event, DatasetUrn source UpstreamLineage upstreamLineage = getUpstreamLineageFromEvent(event); if (upstreamLineage != null && upstreamLineage.hasUpstreams()) { UpstreamArray upstreams = upstreamLineage.getUpstreams(); - if ( - upstreams.size() == 1 - && upstreams.get(0).getDataset().getPlatformEntity().getPlatformNameEntity().equals(DBT_PLATFORM_NAME)) { - setSiblingsAndSoftDeleteSibling(upstreams.get(0).getDataset(), sourceUrn); + + // an entity can have merged lineage (eg. dbt + snowflake), but by default siblings are only between dbt <> non-dbt + UpstreamArray dbtUpstreams = new UpstreamArray( + upstreams.stream() + .filter(obj -> obj.getDataset().getPlatformEntity().getPlatformNameEntity().equals(DBT_PLATFORM_NAME)) + .collect(Collectors.toList()) + ); + // We're assuming a data asset (eg. snowflake table) will only ever be downstream of 1 dbt model + if (dbtUpstreams.size() == 1) { + setSiblingsAndSoftDeleteSibling(dbtUpstreams.get(0).getDataset(), sourceUrn); + } else { + log.error("{} has an unexpected number of dbt upstreams: {}. Not adding any as siblings.", sourceUrn.toString(), dbtUpstreams.size()); + } } } @@ -219,7 +228,7 @@ private void setSiblingsAndSoftDeleteSibling(Urn dbtUrn, Urn sourceUrn) { existingDbtSiblingAspect != null && existingSourceSiblingAspect != null && existingDbtSiblingAspect.getSiblings().contains(sourceUrn.toString()) - && existingDbtSiblingAspect.getSiblings().contains(dbtUrn.toString()) + && existingSourceSiblingAspect.getSiblings().contains(dbtUrn.toString()) ) { // we have already connected them- we can abort here return; diff --git a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java index 5fb2cfaaef2d1..78d304d67bfc0 100644 --- a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java +++ b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java @@ -36,6 +36,8 @@ import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import java.net.URISyntaxException; + import static com.linkedin.metadata.Constants.*; import static org.mockito.ArgumentMatchers.*; @@ -78,15 +80,12 @@ public void testInvokeWhenThereIsAPairWithDbtSourceNode() throws Exception { _mockAuthentication )).thenReturn(mockResponse); - MetadataChangeLog event = new MetadataChangeLog(); - event.setEntityType(DATASET_ENTITY_NAME); - event.setAspectName(UPSTREAM_LINEAGE_ASPECT_NAME); - event.setChangeType(ChangeType.UPSERT); + + MetadataChangeLog event = createEvent(DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, ChangeType.UPSERT); + + Upstream upstream = createUpstream("urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj.jaffle_shop.customers,PROD)", DatasetLineageType.TRANSFORMED); final UpstreamLineage upstreamLineage = new UpstreamLineage(); final UpstreamArray upstreamArray = new UpstreamArray(); - final Upstream upstream = new Upstream(); - upstream.setType(DatasetLineageType.TRANSFORMED); - upstream.setDataset(DatasetUrn.createFromString("urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj.jaffle_shop.customers,PROD)")); upstreamArray.add(upstream); upstreamLineage.setUpstreams(upstreamArray); @@ -151,15 +150,11 @@ public void testInvokeWhenThereIsNoPairWithDbtModel() throws Exception { _mockAuthentication )).thenReturn(mockResponse); - MetadataChangeLog event = new MetadataChangeLog(); - event.setEntityType(DATASET_ENTITY_NAME); - event.setAspectName(UPSTREAM_LINEAGE_ASPECT_NAME); - event.setChangeType(ChangeType.UPSERT); + MetadataChangeLog event = createEvent(DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, ChangeType.UPSERT); + Upstream upstream = createUpstream("urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj.jaffle_shop.customers,PROD)", DatasetLineageType.TRANSFORMED); + final UpstreamLineage upstreamLineage = new UpstreamLineage(); final UpstreamArray upstreamArray = new UpstreamArray(); - final Upstream upstream = new Upstream(); - upstream.setType(DatasetLineageType.TRANSFORMED); - upstream.setDataset(DatasetUrn.createFromString("urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj.jaffle_shop.customers,PROD)")); upstreamArray.add(upstream); upstreamLineage.setUpstreams(upstreamArray); @@ -189,15 +184,11 @@ public void testInvokeWhenThereIsNoPairWithDbtModel() throws Exception { public void testInvokeWhenThereIsAPairWithBigqueryDownstreamNode() throws Exception { Mockito.when(_mockEntityClient.exists(Mockito.any(), Mockito.any())).thenReturn(true); - MetadataChangeLog event = new MetadataChangeLog(); - event.setEntityType(DATASET_ENTITY_NAME); - event.setAspectName(UPSTREAM_LINEAGE_ASPECT_NAME); - event.setChangeType(ChangeType.UPSERT); + + MetadataChangeLog event = createEvent(DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, ChangeType.UPSERT); final UpstreamLineage upstreamLineage = new UpstreamLineage(); final UpstreamArray upstreamArray = new UpstreamArray(); - final Upstream upstream = new Upstream(); - upstream.setType(DatasetLineageType.TRANSFORMED); - upstream.setDataset(DatasetUrn.createFromString("urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)")); + Upstream upstream = createUpstream("urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)", DatasetLineageType.TRANSFORMED); upstreamArray.add(upstream); upstreamLineage.setUpstreams(upstreamArray); @@ -259,10 +250,7 @@ public void testInvokeWhenThereIsAKeyBeingReingested() throws Exception { .setSkipAggregates(true).setSkipHighlighting(true)) )).thenReturn(returnSearchResult); - MetadataChangeLog event = new MetadataChangeLog(); - event.setEntityType(DATASET_ENTITY_NAME); - event.setAspectName(DATASET_KEY_ASPECT_NAME); - event.setChangeType(ChangeType.UPSERT); + MetadataChangeLog event = createEvent(DATASET_ENTITY_NAME, DATASET_KEY_ASPECT_NAME, ChangeType.UPSERT); final DatasetKey datasetKey = new DatasetKey(); datasetKey.setName("my-proj.jaffle_shop.customers"); datasetKey.setOrigin(FabricType.PROD); @@ -304,4 +292,76 @@ public void testInvokeWhenThereIsAKeyBeingReingested() throws Exception { Mockito.eq(_mockAuthentication) ); } -} + @Test + public void testInvokeWhenSourceUrnHasTwoDbtUpstreams() throws Exception { + + MetadataChangeLog event = createEvent(DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, ChangeType.UPSERT); + final UpstreamLineage upstreamLineage = new UpstreamLineage(); + final UpstreamArray upstreamArray = new UpstreamArray(); + Upstream dbtUpstream1 = createUpstream("urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.source_entity1,PROD)", DatasetLineageType.TRANSFORMED); + Upstream dbtUpstream2 = createUpstream("urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.source_entity2,PROD)", DatasetLineageType.TRANSFORMED); + upstreamArray.add(dbtUpstream1); + upstreamArray.add(dbtUpstream2); + upstreamLineage.setUpstreams(upstreamArray); + + event.setAspect(GenericRecordUtils.serializeAspect(upstreamLineage)); + event.setEntityUrn(Urn.createFromString("urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj.jaffle_shop.customers,PROD)")); + _siblingAssociationHook.invoke(event); + + + Mockito.verify(_mockEntityClient, Mockito.times(0)).ingestProposal( + Mockito.any(), + Mockito.eq(_mockAuthentication) + ); + + + } + + @Test + public void testInvokeWhenSourceUrnHasTwoUpstreamsOneDbt() throws Exception { + + MetadataChangeLog event = createEvent(DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, ChangeType.UPSERT); + final UpstreamLineage upstreamLineage = new UpstreamLineage(); + final UpstreamArray upstreamArray = new UpstreamArray(); + Upstream dbtUpstream = createUpstream("urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.source_entity1,PROD)", DatasetLineageType.TRANSFORMED); + Upstream snowflakeUpstream = + createUpstream("urn:li:dataset:(urn:li:dataPlatform:snowflake,my-proj.jaffle_shop.customers,PROD)", DatasetLineageType.TRANSFORMED); + upstreamArray.add(dbtUpstream); + upstreamArray.add(snowflakeUpstream); + upstreamLineage.setUpstreams(upstreamArray); + + event.setAspect(GenericRecordUtils.serializeAspect(upstreamLineage)); + event.setEntityUrn(Urn.createFromString("urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj.jaffle_shop.customers,PROD)")); + _siblingAssociationHook.invoke(event); + + + Mockito.verify(_mockEntityClient, Mockito.times(2)).ingestProposal( + Mockito.any(), + Mockito.eq(_mockAuthentication) + ); + + + } + + private MetadataChangeLog createEvent(String entityType, String aspectName, ChangeType changeType) { + MetadataChangeLog event = new MetadataChangeLog(); + event.setEntityType(entityType); + event.setAspectName(aspectName); + event.setChangeType(changeType); + return event; + } + private Upstream createUpstream(String urn, DatasetLineageType upstreamType) { + + final Upstream upstream = new Upstream(); + upstream.setType(upstreamType); + try { + upstream.setDataset(DatasetUrn.createFromString(urn)); + } catch (URISyntaxException e) { + throw new RuntimeException(e); + } + + return upstream; + } + + + } From e9b4727c8e270d22c80c4be7133a3315adbc5691 Mon Sep 17 00:00:00 2001 From: Chris Collins Date: Wed, 13 Sep 2023 11:18:52 -0400 Subject: [PATCH 007/156] fix(java) Update DataProductMapper to always return a name (#8832) --- .../types/dataproduct/mappers/DataProductMapper.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapper.java index 9cb6840067e7b..254b43ecb96cc 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapper.java @@ -50,7 +50,8 @@ public DataProduct apply(@Nonnull final EntityResponse entityResponse) { EnvelopedAspectMap aspectMap = entityResponse.getAspects(); MappingHelper mappingHelper = new MappingHelper<>(aspectMap, result); - mappingHelper.mapToResult(DATA_PRODUCT_PROPERTIES_ASPECT_NAME, this::mapDataProductProperties); + mappingHelper.mapToResult(DATA_PRODUCT_PROPERTIES_ASPECT_NAME, (dataProduct, dataMap) -> + mapDataProductProperties(dataProduct, dataMap, entityUrn)); mappingHelper.mapToResult(GLOBAL_TAGS_ASPECT_NAME, (dataProduct, dataMap) -> dataProduct.setTags(GlobalTagsMapper.map(new GlobalTags(dataMap), entityUrn))); mappingHelper.mapToResult(GLOSSARY_TERMS_ASPECT_NAME, (dataProduct, dataMap) -> @@ -65,11 +66,12 @@ public DataProduct apply(@Nonnull final EntityResponse entityResponse) { return result; } - private void mapDataProductProperties(@Nonnull DataProduct dataProduct, @Nonnull DataMap dataMap) { + private void mapDataProductProperties(@Nonnull DataProduct dataProduct, @Nonnull DataMap dataMap, @Nonnull Urn urn) { DataProductProperties dataProductProperties = new DataProductProperties(dataMap); com.linkedin.datahub.graphql.generated.DataProductProperties properties = new com.linkedin.datahub.graphql.generated.DataProductProperties(); - properties.setName(dataProductProperties.getName()); + final String name = dataProductProperties.hasName() ? dataProductProperties.getName() : urn.getId(); + properties.setName(name); properties.setDescription(dataProductProperties.getDescription()); if (dataProductProperties.hasExternalUrl()) { properties.setExternalUrl(dataProductProperties.getExternalUrl().toString()); From 1474ac01b19f47d1011dc836f0fceeb59bd1720d Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Wed, 13 Sep 2023 12:32:45 -0700 Subject: [PATCH 008/156] build(ingest): Bump jsonschema for Python >= 3.8 (#8836) --- metadata-ingestion/setup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py index 09f71fa769fd3..7a5fd355803cb 100644 --- a/metadata-ingestion/setup.py +++ b/metadata-ingestion/setup.py @@ -58,7 +58,8 @@ def get_long_description(): "requests_file", "jsonref", # jsonschema drops python 3.7 support in v4.18.0 - "jsonschema<=4.17.3", + "jsonschema<=4.17.3 ; python_version < '3.8'", + "jsonschema>=4.18.0 ; python_version >= '3.8'", "ruamel.yaml", } From 493d31531a1ed829adc106ea7722c88c50b70270 Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Wed, 13 Sep 2023 14:00:58 -0700 Subject: [PATCH 009/156] feat(ingest/rest-emitter): Do not raise error on retry failure to get better error messages (#8837) --- metadata-ingestion/src/datahub/emitter/rest_emitter.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/metadata-ingestion/src/datahub/emitter/rest_emitter.py b/metadata-ingestion/src/datahub/emitter/rest_emitter.py index acb5763280905..937e0902d6d8c 100644 --- a/metadata-ingestion/src/datahub/emitter/rest_emitter.py +++ b/metadata-ingestion/src/datahub/emitter/rest_emitter.py @@ -120,11 +120,15 @@ def __init__( self._retry_max_times = retry_max_times try: + # Set raise_on_status to False to propagate errors: + # https://stackoverflow.com/questions/70189330/determine-status-code-from-python-retry-exception + # Must call `raise_for_status` after making a request, which we do retry_strategy = Retry( total=self._retry_max_times, status_forcelist=self._retry_status_codes, backoff_factor=2, allowed_methods=self._retry_methods, + raise_on_status=False, ) except TypeError: # Prior to urllib3 1.26, the Retry class used `method_whitelist` instead of `allowed_methods`. @@ -133,6 +137,7 @@ def __init__( status_forcelist=self._retry_status_codes, backoff_factor=2, method_whitelist=self._retry_methods, + raise_on_status=False, ) adapter = HTTPAdapter( From 31abf383d13538cdb2fdb3b89ca3ca1fe6b1590f Mon Sep 17 00:00:00 2001 From: Hyejin Yoon <0327jane@gmail.com> Date: Thu, 14 Sep 2023 11:34:21 +0900 Subject: [PATCH 010/156] ci: add markdown-link-check (#8771) --- README.md | 6 +- docs-website/build.gradle | 2 +- docs-website/markdown-link-check-config.json | 50 +++++++ docs-website/package.json | 3 +- docs-website/yarn.lock | 122 ++++++++++++++++-- docs/advanced/no-code-modeling.md | 7 +- docs/api/graphql/how-to-set-up-graphql.md | 2 +- docs/architecture/architecture.md | 2 +- docs/authentication/guides/add-users.md | 8 +- .../guides/sso/configure-oidc-react.md | 2 +- docs/cli.md | 2 +- docs/domains.md | 19 ++- docs/how/add-new-aspect.md | 10 +- docs/modeling/extending-the-metadata-model.md | 10 +- docs/modeling/metadata-model.md | 4 +- docs/tags.md | 10 +- docs/townhall-history.md | 2 +- docs/what/gms.md | 4 +- docs/what/mxe.md | 2 +- docs/what/relationship.md | 3 - docs/what/search-document.md | 1 - .../add_stateful_ingestion_to_source.md | 13 +- .../docs/dev_guides/reporting_telemetry.md | 2 +- .../docs/dev_guides/stateful.md | 16 +-- metadata-ingestion/docs/sources/gcs/README.md | 4 +- .../docs/sources/kafka-connect/README.md | 10 +- metadata-ingestion/docs/sources/s3/README.md | 4 +- .../examples/transforms/README.md | 2 +- .../source/usage/starburst_trino_usage.py | 3 - metadata-jobs/README.md | 4 +- metadata-models/docs/entities/dataPlatform.md | 4 +- 31 files changed, 236 insertions(+), 97 deletions(-) create mode 100644 docs-website/markdown-link-check-config.json diff --git a/README.md b/README.md index 951dcebad6498..79f85433fbc18 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ export const Logo = (props) => {
DataHub Logo
@@ -18,7 +18,7 @@ export const Logo = (props) => {

-DataHub +DataHub

@@ -156,7 +156,7 @@ Here are the companies that have officially adopted DataHub. Please feel free to - [DataHub Blog](https://blog.datahubproject.io/) - [DataHub YouTube Channel](https://www.youtube.com/channel/UC3qFQC5IiwR5fvWEqi_tJ5w) -- [Optum: Data Mesh via DataHub](https://optum.github.io/blog/2022/03/23/data-mesh-via-datahub/) +- [Optum: Data Mesh via DataHub](https://opensource.optum.com/blog/2022/03/23/data-mesh-via-datahub) - [Saxo Bank: Enabling Data Discovery in Data Mesh](https://medium.com/datahub-project/enabling-data-discovery-in-a-data-mesh-the-saxo-journey-451b06969c8f) - [Bringing The Power Of The DataHub Real-Time Metadata Graph To Everyone At Acryl Data](https://www.dataengineeringpodcast.com/acryl-data-datahub-metadata-graph-episode-230/) - [DataHub: Popular Metadata Architectures Explained](https://engineering.linkedin.com/blog/2020/datahub-popular-metadata-architectures-explained) diff --git a/docs-website/build.gradle b/docs-website/build.gradle index 370ae3eec9176..a213ec1ae8194 100644 --- a/docs-website/build.gradle +++ b/docs-website/build.gradle @@ -89,7 +89,7 @@ task fastReload(type: YarnTask) { args = ['run', 'generate-rsync'] } -task yarnLint(type: YarnTask, dependsOn: [yarnInstall]) { +task yarnLint(type: YarnTask, dependsOn: [yarnInstall, yarnGenerate]) { inputs.files(projectMdFiles) args = ['run', 'lint-check'] outputs.dir("dist") diff --git a/docs-website/markdown-link-check-config.json b/docs-website/markdown-link-check-config.json new file mode 100644 index 0000000000000..26e040edde6f7 --- /dev/null +++ b/docs-website/markdown-link-check-config.json @@ -0,0 +1,50 @@ +{ + "ignorePatterns": [ + { + "pattern": "^http://demo\\.datahubproject\\.io" + }, + { + "pattern": "^http://localhost" + }, + { + "pattern": "^http://www.famfamfam.com" + }, + { + "pattern": "^http://www.linkedin.com" + }, + { + "pattern": "\\.md$" + }, + { + "pattern":"\\.json$" + }, + { + "pattern":"\\.txt$" + }, + { + "pattern": "\\.java$" + }, + { + "pattern": "\\.md#.*$" + }, + { + "pattern": "^https://oauth2.googleapis.com/token" + }, + { + "pattern": "^https://login.microsoftonline.com/common/oauth2/na$" + }, + { + "pattern": "#v(\\d+)-(\\d+)-(\\d+)" + }, + { + "pattern": "^https://github.com/mohdsiddique$" + }, + { + "pattern": "^https://github.com/2x$" + }, + { + "pattern": "^https://github.com/datahub-project/datahub/assets/15873986/2f47d033-6c2b-483a-951d-e6d6b807f0d0%22%3E$" + } + ], + "aliveStatusCodes": [200, 206, 0, 999, 400, 401, 403] +} \ No newline at end of file diff --git a/docs-website/package.json b/docs-website/package.json index 400ef4143c786..1722f92169692 100644 --- a/docs-website/package.json +++ b/docs-website/package.json @@ -17,7 +17,7 @@ "generate": "rm -rf genDocs genStatic && mkdir genDocs genStatic && yarn _generate-docs && mv docs/* genDocs/ && rmdir docs", "generate-rsync": "mkdir -p genDocs genStatic && yarn _generate-docs && rsync -v --checksum -r -h -i --delete docs/ genDocs && rm -rf docs", "lint": "prettier -w generateDocsDir.ts sidebars.js src/pages/index.js", - "lint-check": "prettier -l generateDocsDir.ts sidebars.js src/pages/index.js", + "lint-check": "prettier -l generateDocsDir.ts sidebars.js src/pages/index.js && find ./genDocs -name \\*.md -not -path \"./genDocs/python-sdk/models.md\" -print0 | xargs -0 -n1 markdown-link-check -p -q -c markdown-link-check-config.json", "lint-fix": "prettier --write generateDocsDir.ts sidebars.js src/pages/index.js" }, "dependencies": { @@ -37,6 +37,7 @@ "docusaurus-graphql-plugin": "0.5.0", "docusaurus-plugin-sass": "^0.2.1", "dotenv": "^16.0.1", + "markdown-link-check": "^3.11.2", "markprompt": "^0.1.7", "react": "^18.2.0", "react-dom": "18.2.0", diff --git a/docs-website/yarn.lock b/docs-website/yarn.lock index 0613fe71ef78e..5698029bff70a 100644 --- a/docs-website/yarn.lock +++ b/docs-website/yarn.lock @@ -3414,6 +3414,11 @@ async-validator@^4.1.0: resolved "https://registry.yarnpkg.com/async-validator/-/async-validator-4.2.5.tgz#c96ea3332a521699d0afaaceed510a54656c6339" integrity sha512-7HhHjtERjqlNbZtqNqy2rckN/SpOOlmDliet+lP7k+eKZEjPk3DgyeU9lIXLdeLz0uBbbVp+9Qdow9wJWgwwfg== +async@^3.2.4: + version "3.2.4" + resolved "https://registry.yarnpkg.com/async/-/async-3.2.4.tgz#2d22e00f8cddeb5fde5dd33522b56d1cf569a81c" + integrity sha512-iAB+JbDEGXhyIUavoDl9WP/Jj106Kz9DEn1DPgYw5ruDn0e3Wgi3sKFm55sASdGBNOQB8F59d9qQ7deqrHA8wQ== + asynckit@^0.4.0: version "0.4.0" resolved "https://registry.yarnpkg.com/asynckit/-/asynckit-0.4.0.tgz#c79ed97f7f34cb8f2ba1bc9790bcc366474b4b79" @@ -3765,6 +3770,11 @@ chalk@^4.0.0, chalk@^4.1.0, chalk@^4.1.2: ansi-styles "^4.1.0" supports-color "^7.1.0" +chalk@^5.2.0: + version "5.3.0" + resolved "https://registry.yarnpkg.com/chalk/-/chalk-5.3.0.tgz#67c20a7ebef70e7f3970a01f90fa210cb6860385" + integrity sha512-dLitG79d+GV1Nb/VYcCDFivJeK1hiukt9QjRNVOsUtTy1rR1YJsmpGGTZ3qJos+uw7WmWF4wUwBd9jxjocFC2w== + character-entities-legacy@^1.0.0: version "1.1.4" resolved "https://registry.yarnpkg.com/character-entities-legacy/-/character-entities-legacy-1.1.4.tgz#94bc1845dce70a5bb9d2ecc748725661293d8fc1" @@ -3797,7 +3807,7 @@ cheerio-select@^2.1.0: domhandler "^5.0.3" domutils "^3.0.1" -cheerio@^1.0.0-rc.12: +cheerio@^1.0.0-rc.10, cheerio@^1.0.0-rc.12: version "1.0.0-rc.12" resolved "https://registry.yarnpkg.com/cheerio/-/cheerio-1.0.0-rc.12.tgz#788bf7466506b1c6bf5fae51d24a2c4d62e47683" integrity sha512-VqR8m68vM46BNnuZ5NtnGBKIE/DfN0cRIzg9n40EIq9NOv90ayxLBXA8fXC5gquFRGJSTRqBq25Jt2ECLR431Q== @@ -3984,6 +3994,11 @@ comma-separated-tokens@^2.0.0: resolved "https://registry.yarnpkg.com/comma-separated-tokens/-/comma-separated-tokens-2.0.3.tgz#4e89c9458acb61bc8fef19f4529973b2392839ee" integrity sha512-Fu4hJdvzeylCfQPp9SGWidpzrMs7tTrlu6Vb8XGaRGck8QSNZJJp538Wrb60Lax4fPwR64ViY468OIUTbRlGZg== +commander@^10.0.1: + version "10.0.1" + resolved "https://registry.yarnpkg.com/commander/-/commander-10.0.1.tgz#881ee46b4f77d1c1dccc5823433aa39b022cbe06" + integrity sha512-y4Mg2tXshplEbSGzx7amzPwKKOCGuoSRP/CjEdwwk0FOGlUbq6lKuoyDZTNZkmxHdJtp54hdfY/JUrdL7Xfdug== + commander@^2.20.0: version "2.20.3" resolved "https://registry.yarnpkg.com/commander/-/commander-2.20.3.tgz#fd485e84c03eb4881c20722ba48035e8531aeb33" @@ -4385,6 +4400,13 @@ debug@4, debug@^4.0.0, debug@^4.1.0, debug@^4.1.1: dependencies: ms "2.1.2" +debug@^3.2.6: + version "3.2.7" + resolved "https://registry.yarnpkg.com/debug/-/debug-3.2.7.tgz#72580b7e9145fb39b6676f9c5e5fb100b934179a" + integrity sha512-CFjzYYAi4ThfiQvizrFQevTTXHtnCqWfe7x1AhgEscTz6ZbLbfoLRLPugTQyBth6f8ZERVUSyWHFD/7Wu4t1XQ== + dependencies: + ms "^2.1.1" + decode-named-character-reference@^1.0.0: version "1.0.2" resolved "https://registry.yarnpkg.com/decode-named-character-reference/-/decode-named-character-reference-1.0.2.tgz#daabac9690874c394c81e4162a0304b35d824f0e" @@ -5551,6 +5573,13 @@ html-entities@^2.3.2: resolved "https://registry.yarnpkg.com/html-entities/-/html-entities-2.4.0.tgz#edd0cee70402584c8c76cc2c0556db09d1f45061" integrity sha512-igBTJcNNNhvZFRtm8uA6xMY6xYleeDwn3PeBCkDz7tHttv4F2hsDI2aPgNERWzvRcNYHNT3ymRaQzllmXj4YsQ== +html-link-extractor@^1.0.5: + version "1.0.5" + resolved "https://registry.yarnpkg.com/html-link-extractor/-/html-link-extractor-1.0.5.tgz#a4be345cb13b8c3352d82b28c8b124bb7bf5dd6f" + integrity sha512-ADd49pudM157uWHwHQPUSX4ssMsvR/yHIswOR5CUfBdK9g9ZYGMhVSE6KZVHJ6kCkR0gH4htsfzU6zECDNVwyw== + dependencies: + cheerio "^1.0.0-rc.10" + html-minifier-terser@^6.0.2, html-minifier-terser@^6.1.0: version "6.1.0" resolved "https://registry.yarnpkg.com/html-minifier-terser/-/html-minifier-terser-6.1.0.tgz#bfc818934cc07918f6b3669f5774ecdfd48f32ab" @@ -5673,6 +5702,13 @@ iconv-lite@0.4.24: dependencies: safer-buffer ">= 2.1.2 < 3" +iconv-lite@^0.6.3: + version "0.6.3" + resolved "https://registry.yarnpkg.com/iconv-lite/-/iconv-lite-0.6.3.tgz#a52f80bf38da1952eb5c681790719871a1a72501" + integrity sha512-4fCk79wshMdzMp2rH06qWrJE4iolqLhCUH+OiuIgU++RB0+94NlDL81atO7GX55uUKueo0txHNtvEyI6D7WdMw== + dependencies: + safer-buffer ">= 2.1.2 < 3.0.0" + icss-utils@^5.0.0, icss-utils@^5.1.0: version "5.1.0" resolved "https://registry.yarnpkg.com/icss-utils/-/icss-utils-5.1.0.tgz#c6be6858abd013d768e98366ae47e25d5887b1ae" @@ -5795,6 +5831,11 @@ ipaddr.js@^2.0.1: resolved "https://registry.yarnpkg.com/ipaddr.js/-/ipaddr.js-2.1.0.tgz#2119bc447ff8c257753b196fc5f1ce08a4cdf39f" integrity sha512-LlbxQ7xKzfBusov6UMi4MFpEg0m+mAm9xyNGEduwXMEDuf4WfzB/RZwMVYEd7IKGvh4IUkEXYxtAVu9T3OelJQ== +is-absolute-url@^4.0.1: + version "4.0.1" + resolved "https://registry.yarnpkg.com/is-absolute-url/-/is-absolute-url-4.0.1.tgz#16e4d487d4fded05cfe0685e53ec86804a5e94dc" + integrity sha512-/51/TKE88Lmm7Gc4/8btclNXWS+g50wXhYJq8HWIBAGUBnoAdRu1aXeh364t/O7wXDAcTJDP8PNuNKWUDWie+A== + is-alphabetical@1.0.4, is-alphabetical@^1.0.0: version "1.0.4" resolved "https://registry.yarnpkg.com/is-alphabetical/-/is-alphabetical-1.0.4.tgz#9e7d6b94916be22153745d184c298cbf986a686d" @@ -5963,6 +6004,13 @@ is-regexp@^1.0.0: resolved "https://registry.yarnpkg.com/is-regexp/-/is-regexp-1.0.0.tgz#fd2d883545c46bac5a633e7b9a09e87fa2cb5069" integrity sha512-7zjFAPO4/gwyQAAgRRmqeEeyIICSdmCqa3tsVHMdBzaXXRiqopZL4Cyghg/XulGWrtABTpbnYYzzIRffLkP4oA== +is-relative-url@^4.0.0: + version "4.0.0" + resolved "https://registry.yarnpkg.com/is-relative-url/-/is-relative-url-4.0.0.tgz#4d8371999ff6033b76e4d9972fb5bf496fddfa97" + integrity sha512-PkzoL1qKAYXNFct5IKdKRH/iBQou/oCC85QhXj6WKtUQBliZ4Yfd3Zk27RHu9KQG8r6zgvAA2AQKC9p+rqTszg== + dependencies: + is-absolute-url "^4.0.1" + is-root@^2.1.0: version "2.1.0" resolved "https://registry.yarnpkg.com/is-root/-/is-root-2.1.0.tgz#809e18129cf1129644302a4f8544035d51984a9c" @@ -6010,6 +6058,13 @@ isarray@~1.0.0: resolved "https://registry.yarnpkg.com/isarray/-/isarray-1.0.0.tgz#bb935d48582cba168c06834957a54a3e07124f11" integrity sha512-VLghIWNM6ELQzo7zwmcg0NmTVyWKYjvIeM83yjp0wRDTmUnrM678fQbcKBo6n2CJEF0szoG//ytg+TKla89ALQ== +isemail@^3.2.0: + version "3.2.0" + resolved "https://registry.yarnpkg.com/isemail/-/isemail-3.2.0.tgz#59310a021931a9fb06bbb51e155ce0b3f236832c" + integrity sha512-zKqkK+O+dGqevc93KNsbZ/TqTUFd46MwWjYOoMrjIMZ51eU7DtQG3Wmd9SQQT7i7RVnuTPEiYEWHU3MSbxC1Tg== + dependencies: + punycode "2.x.x" + isexe@^2.0.0: version "2.0.0" resolved "https://registry.yarnpkg.com/isexe/-/isexe-2.0.0.tgz#e8fbf374dc556ff8947a10dcb0572d633f2cfa10" @@ -6205,6 +6260,16 @@ lines-and-columns@^1.1.6: resolved "https://registry.yarnpkg.com/lines-and-columns/-/lines-and-columns-1.2.4.tgz#eca284f75d2965079309dc0ad9255abb2ebc1632" integrity sha512-7ylylesZQ/PV29jhEDl3Ufjo6ZX7gCqJr5F7PKrqc93v7fzSymt1BpwEU8nAUXs8qzzvqhbjhK5QZg6Mt/HkBg== +link-check@^5.2.0: + version "5.2.0" + resolved "https://registry.yarnpkg.com/link-check/-/link-check-5.2.0.tgz#595a339d305900bed8c1302f4342a29c366bf478" + integrity sha512-xRbhYLaGDw7eRDTibTAcl6fXtmUQ13vkezQiTqshHHdGueQeumgxxmQMIOmJYsh2p8BF08t8thhDQ++EAOOq3w== + dependencies: + is-relative-url "^4.0.0" + isemail "^3.2.0" + ms "^2.1.3" + needle "^3.1.0" + loader-runner@^4.2.0: version "4.3.0" resolved "https://registry.yarnpkg.com/loader-runner/-/loader-runner-4.3.0.tgz#c1b4a163b99f614830353b16755e7149ac2314e1" @@ -6366,6 +6431,28 @@ markdown-escapes@^1.0.0: resolved "https://registry.yarnpkg.com/markdown-escapes/-/markdown-escapes-1.0.4.tgz#c95415ef451499d7602b91095f3c8e8975f78535" integrity sha512-8z4efJYk43E0upd0NbVXwgSTQs6cT3T06etieCMEg7dRbzCbxUCK/GHlX8mhHRDcp+OLlHkPKsvqQTCvsRl2cg== +markdown-link-check@^3.11.2: + version "3.11.2" + resolved "https://registry.yarnpkg.com/markdown-link-check/-/markdown-link-check-3.11.2.tgz#303a8a03d4a34c42ef3158e0b245bced26b5d904" + integrity sha512-zave+vI4AMeLp0FlUllAwGbNytSKsS3R2Zgtf3ufVT892Z/L6Ro9osZwE9PNA7s0IkJ4onnuHqatpsaCiAShJw== + dependencies: + async "^3.2.4" + chalk "^5.2.0" + commander "^10.0.1" + link-check "^5.2.0" + lodash "^4.17.21" + markdown-link-extractor "^3.1.0" + needle "^3.2.0" + progress "^2.0.3" + +markdown-link-extractor@^3.1.0: + version "3.1.0" + resolved "https://registry.yarnpkg.com/markdown-link-extractor/-/markdown-link-extractor-3.1.0.tgz#0d5a703630d791a9e2017449e1a9b294f2d2b676" + integrity sha512-r0NEbP1dsM+IqB62Ru9TXLP/HDaTdBNIeylYXumuBi6Xv4ufjE1/g3TnslYL8VNqNcGAGbMptQFHrrdfoZ/Sug== + dependencies: + html-link-extractor "^1.0.5" + marked "^4.1.0" + markdown-table@^3.0.0: version "3.0.3" resolved "https://registry.yarnpkg.com/markdown-table/-/markdown-table-3.0.3.tgz#e6331d30e493127e031dd385488b5bd326e4a6bd" @@ -6376,6 +6463,11 @@ marked@^2.0.3: resolved "https://registry.yarnpkg.com/marked/-/marked-2.1.3.tgz#bd017cef6431724fd4b27e0657f5ceb14bff3753" integrity sha512-/Q+7MGzaETqifOMWYEA7HVMaZb4XbcRfaOzcSsHZEith83KGlvaSG33u0SKu89Mj5h+T8V2hM+8O45Qc5XTgwA== +marked@^4.1.0: + version "4.3.0" + resolved "https://registry.yarnpkg.com/marked/-/marked-4.3.0.tgz#796362821b019f734054582038b116481b456cf3" + integrity sha512-PRsaiG84bK+AMvxziE/lCFss8juXjNaWzVbN5tXAm4XjeaS9NAHhop+PjQxz2A9h8Q4M/xGmzP8vqNwy6JeK0A== + markprompt@^0.1.7: version "0.1.7" resolved "https://registry.yarnpkg.com/markprompt/-/markprompt-0.1.7.tgz#fa049e11109d93372c45c38b3ca40bd5fdf751ea" @@ -6978,7 +7070,7 @@ ms@2.1.2: resolved "https://registry.yarnpkg.com/ms/-/ms-2.1.2.tgz#d09d1f357b443f493382a8eb3ccd183872ae6009" integrity sha512-sGkPx+VjMtmA6MX27oA4FBFELFCZZ4S4XqeGOXCv68tT+jb3vk/RyaKWP0PTKyWtmLSM0b+adUTEvbs1PEaH2w== -ms@2.1.3: +ms@2.1.3, ms@^2.1.1, ms@^2.1.3: version "2.1.3" resolved "https://registry.yarnpkg.com/ms/-/ms-2.1.3.tgz#574c8138ce1d2b5861f0b44579dbadd60c6615b2" integrity sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA== @@ -7001,6 +7093,15 @@ napi-build-utils@^1.0.1: resolved "https://registry.yarnpkg.com/napi-build-utils/-/napi-build-utils-1.0.2.tgz#b1fddc0b2c46e380a0b7a76f984dd47c41a13806" integrity sha512-ONmRUqK7zj7DWX0D9ADe03wbwOBZxNAfF20PlGfCWQcD3+/MakShIHrMqx9YwPTfxDdF1zLeL+RGZiR9kGMLdg== +needle@^3.1.0, needle@^3.2.0: + version "3.2.0" + resolved "https://registry.yarnpkg.com/needle/-/needle-3.2.0.tgz#07d240ebcabfd65c76c03afae7f6defe6469df44" + integrity sha512-oUvzXnyLiVyVGoianLijF9O/RecZUf7TkBfimjGrLM4eQhXyeJwM6GeAWccwfQ9aa4gMCZKqhAOuLaMIcQxajQ== + dependencies: + debug "^3.2.6" + iconv-lite "^0.6.3" + sax "^1.2.4" + negotiator@0.6.3: version "0.6.3" resolved "https://registry.yarnpkg.com/negotiator/-/negotiator-0.6.3.tgz#58e323a72fedc0d6f9cd4d31fe49f51479590ccd" @@ -7753,6 +7854,11 @@ process-nextick-args@~2.0.0: resolved "https://registry.yarnpkg.com/process-nextick-args/-/process-nextick-args-2.0.1.tgz#7820d9b16120cc55ca9ae7792680ae7dba6d7fe2" integrity sha512-3ouUOpQhtgrbOa17J7+uxOTpITYWaGP7/AhoR3+A+/1e9skrzelGi/dXzEYyvbxubEF6Wn2ypscTKiKJFFn1ag== +progress@^2.0.3: + version "2.0.3" + resolved "https://registry.yarnpkg.com/progress/-/progress-2.0.3.tgz#7e8cf8d8f5b8f239c1bc68beb4eb78567d572ef8" + integrity sha512-7PiHtLll5LdnKIMw100I+8xJXR5gW2QwWYkT6iJva0bXitZKa/XMrSbdmg3r2Xnaidz9Qumd0VPaMrZlF9V9sA== + promise@^7.1.1: version "7.3.1" resolved "https://registry.yarnpkg.com/promise/-/promise-7.3.1.tgz#064b72602b18f90f29192b8b1bc418ffd1ebd3bf" @@ -7805,16 +7911,16 @@ pump@^3.0.0: end-of-stream "^1.1.0" once "^1.3.1" +punycode@2.x.x, punycode@^2.1.0: + version "2.3.0" + resolved "https://registry.yarnpkg.com/punycode/-/punycode-2.3.0.tgz#f67fa67c94da8f4d0cfff981aee4118064199b8f" + integrity sha512-rRV+zQD8tVFys26lAGR9WUuS4iUAngJScM+ZRSKtvl5tKeZ2t5bvdNFdNHBW9FWR4guGHlgmsZ1G7BSm2wTbuA== + punycode@^1.3.2: version "1.4.1" resolved "https://registry.yarnpkg.com/punycode/-/punycode-1.4.1.tgz#c0d5a63b2718800ad8e1eb0fa5269c84dd41845e" integrity sha512-jmYNElW7yvO7TV33CjSmvSiE2yco3bV2czu/OzDKdMNVZQWfxCblURLhf+47syQRBntjfLdd/H0egrzIG+oaFQ== -punycode@^2.1.0: - version "2.3.0" - resolved "https://registry.yarnpkg.com/punycode/-/punycode-2.3.0.tgz#f67fa67c94da8f4d0cfff981aee4118064199b8f" - integrity sha512-rRV+zQD8tVFys26lAGR9WUuS4iUAngJScM+ZRSKtvl5tKeZ2t5bvdNFdNHBW9FWR4guGHlgmsZ1G7BSm2wTbuA== - pupa@^2.1.1: version "2.1.1" resolved "https://registry.yarnpkg.com/pupa/-/pupa-2.1.1.tgz#f5e8fd4afc2c5d97828faa523549ed8744a20d62" @@ -8789,7 +8895,7 @@ safe-buffer@5.2.1, safe-buffer@>=5.1.0, safe-buffer@^5.0.1, safe-buffer@^5.1.0, resolved "https://registry.yarnpkg.com/safe-buffer/-/safe-buffer-5.2.1.tgz#1eaf9fa9bdb1fdd4ec75f58f9cdb4e6b7827eec6" integrity sha512-rp3So07KcdmmKbGvgaNxQSJr7bGVSVk5S9Eq1F+ppbRo70+YeaDxkw5Dd8NPN+GD6bjnYm2VuPuCXmpuYvmCXQ== -"safer-buffer@>= 2.1.2 < 3": +"safer-buffer@>= 2.1.2 < 3", "safer-buffer@>= 2.1.2 < 3.0.0": version "2.1.2" resolved "https://registry.yarnpkg.com/safer-buffer/-/safer-buffer-2.1.2.tgz#44fa161b0187b9549dd84bb91802f9bd8385cd6a" integrity sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg== diff --git a/docs/advanced/no-code-modeling.md b/docs/advanced/no-code-modeling.md index d76b776d3dddb..172e63f821eab 100644 --- a/docs/advanced/no-code-modeling.md +++ b/docs/advanced/no-code-modeling.md @@ -100,10 +100,9 @@ Currently, there are various models in GMS: 1. [Urn](https://github.com/datahub-project/datahub/blob/master/li-utils/src/main/pegasus/com/linkedin/common/DatasetUrn.pdl) - Structs composing primary keys 2. [Root] [Snapshots](https://github.com/datahub-project/datahub/blob/master/metadata-models/src/main/pegasus/com/linkedin/metadata/snapshot/Snapshot.pdl) - Container of aspects 3. [Aspects](https://github.com/datahub-project/datahub/blob/master/metadata-models/src/main/pegasus/com/linkedin/metadata/aspect/DashboardAspect.pdl) - Optional container of fields -4. [Values](https://github.com/datahub-project/datahub/blob/master/gms/api/src/main/pegasus/com/linkedin/dataset/Dataset.pdl), [Keys](https://github.com/datahub-project/datahub/blob/master/gms/api/src/main/pegasus/com/linkedin/dataset/DatasetKey.pdl) - Model returned by GMS [Rest.li](http://rest.li) API (public facing) -5. [Entities](https://github.com/datahub-project/datahub/blob/master/metadata-models/src/main/pegasus/com/linkedin/metadata/entity/DatasetEntity.pdl) - Records with fields derived from the URN. Used only in graph / relationships -6. [Relationships](https://github.com/datahub-project/datahub/blob/master/metadata-models/src/main/pegasus/com/linkedin/metadata/relationship/Relationship.pdl) - Edges between 2 entities with optional edge properties -7. [Search Documents](https://github.com/datahub-project/datahub/blob/master/metadata-models/src/main/pegasus/com/linkedin/metadata/search/ChartDocument.pdl) - Flat documents for indexing within Elastic index +4. [Keys](https://github.com/datahub-project/datahub/blob/master/metadata-models/src/main/pegasus/com/linkedin/metadata/key/DatasetKey.pdl) - Model returned by GMS [Rest.li](http://rest.li) API (public facing) +5. [Relationships](https://github.com/datahub-project/datahub/blob/master/metadata-models/src/main/pegasus/com/linkedin/common/EntityRelationship.pdl) - Edges between 2 entities with optional edge properties +6. Search Documents - Flat documents for indexing within Elastic index - And corresponding index [mappings.json](https://github.com/datahub-project/datahub/blob/master/gms/impl/src/main/resources/index/chart/mappings.json), [settings.json](https://github.com/datahub-project/datahub/blob/master/gms/impl/src/main/resources/index/chart/settings.json) Various components of GMS depend on / make assumptions about these model types: diff --git a/docs/api/graphql/how-to-set-up-graphql.md b/docs/api/graphql/how-to-set-up-graphql.md index 584bf34ad3f92..2be2f935b12b1 100644 --- a/docs/api/graphql/how-to-set-up-graphql.md +++ b/docs/api/graphql/how-to-set-up-graphql.md @@ -68,7 +68,7 @@ In the request body, select the `GraphQL` option and enter your GraphQL query in

-Please refer to [Querying with GraphQL](https://learning.postman.com/docs/sending-requests/graphql/graphql/) in the Postman documentation for more information. +Please refer to [Querying with GraphQL](https://learning.postman.com/docs/sending-requests/graphql/graphql-overview/) in the Postman documentation for more information. ### Authentication + Authorization diff --git a/docs/architecture/architecture.md b/docs/architecture/architecture.md index 6a9c1860d71b0..20f18f09d949b 100644 --- a/docs/architecture/architecture.md +++ b/docs/architecture/architecture.md @@ -17,7 +17,7 @@ The figures below describe the high-level architecture of DataHub.

- +

diff --git a/docs/authentication/guides/add-users.md b/docs/authentication/guides/add-users.md index f5dfc83201083..d380cacd6665e 100644 --- a/docs/authentication/guides/add-users.md +++ b/docs/authentication/guides/add-users.md @@ -19,13 +19,13 @@ To do so, navigate to the **Users & Groups** section inside of Settings page. He do not have the correct privileges to invite users, this button will be disabled.

- +

To invite new users, simply share the link with others inside your organization.

- +

When a new user visits the link, they will be directed to a sign up screen where they can create their DataHub account. @@ -37,13 +37,13 @@ and click **Reset user password** inside the menu dropdown on the right hand sid `Manage User Credentials` [Platform Privilege](../../authorization/access-policies-guide.md) in order to reset passwords.

- +

To reset the password, simply share the password reset link with the user who needs to change their password. Password reset links expire after 24 hours.

- +

# Configuring Single Sign-On with OpenID Connect diff --git a/docs/authentication/guides/sso/configure-oidc-react.md b/docs/authentication/guides/sso/configure-oidc-react.md index d27792ce3967b..512d6adbf916f 100644 --- a/docs/authentication/guides/sso/configure-oidc-react.md +++ b/docs/authentication/guides/sso/configure-oidc-react.md @@ -26,7 +26,7 @@ please see [this guide](../jaas.md) to mount a custom user.props file for a JAAS To configure OIDC in React, you will most often need to register yourself as a client with your identity provider (Google, Okta, etc). Each provider may have their own instructions. Provided below are links to examples for Okta, Google, Azure AD, & Keycloak. -- [Registering an App in Okta](https://developer.okta.com/docs/guides/add-an-external-idp/apple/register-app-in-okta/) +- [Registering an App in Okta](https://developer.okta.com/docs/guides/add-an-external-idp/openidconnect/main/) - [OpenID Connect in Google Identity](https://developers.google.com/identity/protocols/oauth2/openid-connect) - [OpenID Connect authentication with Azure Active Directory](https://docs.microsoft.com/en-us/azure/active-directory/fundamentals/auth-oidc) - [Keycloak - Securing Applications and Services Guide](https://www.keycloak.org/docs/latest/securing_apps/) diff --git a/docs/cli.md b/docs/cli.md index eb8bb406b0107..267f289d9f54a 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -547,7 +547,7 @@ Old Entities Migrated = {'urn:li:dataset:(urn:li:dataPlatform:hive,logging_event ### Using docker [![Docker Hub](https://img.shields.io/docker/pulls/acryldata/datahub-ingestion?style=plastic)](https://hub.docker.com/r/acryldata/datahub-ingestion) -[![datahub-ingestion docker](https://github.com/acryldata/datahub/actions/workflows/docker-ingestion.yml/badge.svg)](https://github.com/acryldata/datahub/actions/workflows/docker-ingestion.yml) +[![datahub-ingestion docker](https://github.com/acryldata/datahub/workflows/datahub-ingestion%20docker/badge.svg)](https://github.com/acryldata/datahub/actions/workflows/docker-ingestion.yml) If you don't want to install locally, you can alternatively run metadata ingestion within a Docker container. We have prebuilt images available on [Docker hub](https://hub.docker.com/r/acryldata/datahub-ingestion). All plugins will be installed and enabled automatically. diff --git a/docs/domains.md b/docs/domains.md index c846a753417c5..1b2ebc9d47f39 100644 --- a/docs/domains.md +++ b/docs/domains.md @@ -22,20 +22,20 @@ You can create this privileges by creating a new [Metadata Policy](./authorizati To create a Domain, first navigate to the **Domains** tab in the top-right menu of DataHub.

- +

Once you're on the Domains page, you'll see a list of all the Domains that have been created on DataHub. Additionally, you can view the number of entities inside each Domain.

- +

To create a new Domain, click '+ New Domain'.

- +

Inside the form, you can choose a name for your Domain. Most often, this will align with your business units or groups, for example @@ -48,7 +48,7 @@ for the Domain. This option is useful if you intend to refer to Domains by a com key to be human-readable. Proceed with caution: once you select a custom id, it cannot be easily changed.

- +

By default, you don't need to worry about this. DataHub will auto-generate a unique Domain id for you. @@ -64,7 +64,7 @@ To assign an asset to a Domain, simply navigate to the asset's profile page. At see a 'Domain' section. Click 'Set Domain', and then search for the Domain you'd like to add to. When you're done, click 'Add'.

- +

To remove an asset from a Domain, click the 'x' icon on the Domain tag. @@ -149,27 +149,27 @@ source: Once you've created a Domain, you can use the search bar to find it.

- +

Clicking on the search result will take you to the Domain's profile, where you can edit its description, add / remove owners, and view the assets inside the Domain.

- +

Once you've added assets to a Domain, you can filter search results to limit to those Assets within a particular Domain using the left-side search filters.

- +

On the homepage, you'll also find a list of the most popular Domains in your organization.

- +

## Additional Resources @@ -242,7 +242,6 @@ DataHub supports Tags, Glossary Terms, & Domains as distinct types of Metadata t - **Tags**: Informal, loosely controlled labels that serve as a tool for search & discovery. Assets may have multiple tags. No formal, central management. - **Glossary Terms**: A controlled vocabulary, with optional hierarchy. Terms are typically used to standardize types of leaf-level attributes (i.e. schema fields) for governance. E.g. (EMAIL_PLAINTEXT) - **Domains**: A set of top-level categories. Usually aligned to business units / disciplines to which the assets are most relevant. Central or distributed management. Single Domain assignment per data asset. - *Need more help? Join the conversation in [Slack](http://slack.datahubproject.io)!* ### Related Features diff --git a/docs/how/add-new-aspect.md b/docs/how/add-new-aspect.md index 6ea7256ed75cc..d1fe567018903 100644 --- a/docs/how/add-new-aspect.md +++ b/docs/how/add-new-aspect.md @@ -1,20 +1,20 @@ # How to add a new metadata aspect? Adding a new metadata [aspect](../what/aspect.md) is one of the most common ways to extend an existing [entity](../what/entity.md). -We'll use the [CorpUserEditableInfo](https://github.com/datahub-project/datahub/blob/master/metadata-models/src/main/pegasus/com/linkedin/identity/CorpUserEditableInfo.pdl) as an example here. +We'll use the CorpUserEditableInfo as an example here. 1. Add the aspect model to the corresponding namespace (e.g. [`com.linkedin.identity`](https://github.com/datahub-project/datahub/tree/master/metadata-models/src/main/pegasus/com/linkedin/identity)) -2. Extend the entity's aspect union to include the new aspect (e.g. [`CorpUserAspect`](https://github.com/datahub-project/datahub/blob/master/metadata-models/src/main/pegasus/com/linkedin/metadata/aspect/CorpUserAspect.pdl)) +2. Extend the entity's aspect union to include the new aspect. 3. Rebuild the rest.li [IDL & snapshot](https://linkedin.github.io/rest.li/modeling/compatibility_check) by running the following command from the project root ``` ./gradlew :metadata-service:restli-servlet-impl:build -Prest.model.compatibility=ignore ``` -4. To surface the new aspect at the top-level [resource endpoint](https://linkedin.github.io/rest.li/user_guide/restli_server#writing-resources), extend the resource data model (e.g. [`CorpUser`](https://github.com/datahub-project/datahub/blob/master/gms/api/src/main/pegasus/com/linkedin/identity/CorpUser.pdl)) with an optional field (e.g. [`editableInfo`](https://github.com/datahub-project/datahub/blob/master/gms/api/src/main/pegasus/com/linkedin/identity/CorpUser.pdl#L21)). You'll also need to extend the `toValue` & `toSnapshot` methods of the top-level resource (e.g. [`CorpUsers`](https://github.com/datahub-project/datahub/blob/master/gms/impl/src/main/java/com/linkedin/metadata/resources/identity/CorpUsers.java)) to convert between the snapshot & value models. +4. To surface the new aspect at the top-level [resource endpoint](https://linkedin.github.io/rest.li/user_guide/restli_server#writing-resources), extend the resource data model with an optional field. You'll also need to extend the `toValue` & `toSnapshot` methods of the top-level resource (e.g. [`CorpUsers`](https://github.com/datahub-project/datahub/blob/master/gms/impl/src/main/java/com/linkedin/metadata/resources/identity/CorpUsers.java)) to convert between the snapshot & value models. -5. (Optional) If there's need to update the aspect via API (instead of/in addition to MCE), add a [sub-resource](https://linkedin.github.io/rest.li/user_guide/restli_server#sub-resources) endpoint for the new aspect (e.g. [`CorpUsersEditableInfoResource`](https://github.com/datahub-project/datahub/blob/master/gms/impl/src/main/java/com/linkedin/metadata/resources/identity/CorpUsersEditableInfoResource.java)). The sub-resource endpiont also allows you to retrieve previous versions of the aspect as well as additional metadata such as the audit stamp. +5. (Optional) If there's need to update the aspect via API (instead of/in addition to MCE), add a [sub-resource](https://linkedin.github.io/rest.li/user_guide/restli_server#sub-resources) endpoint for the new aspect (e.g. `CorpUsersEditableInfoResource`). The sub-resource endpiont also allows you to retrieve previous versions of the aspect as well as additional metadata such as the audit stamp. -6. After rebuilding & restarting [gms](https://github.com/datahub-project/datahub/tree/master/gms), [mce-consumer-job](https://github.com/datahub-project/datahub/tree/master/metadata-jobs/mce-consumer-job) & [mae-consumer-job](https://github.com/datahub-project/datahub/tree/master/metadata-jobs/mae-consumer-job), +6. After rebuilding & restarting gms, [mce-consumer-job](https://github.com/datahub-project/datahub/tree/master/metadata-jobs/mce-consumer-job) & [mae-consumer-job](https://github.com/datahub-project/datahub/tree/master/metadata-jobs/mae-consumer-job),z you should be able to start emitting [MCE](../what/mxe.md) with the new aspect and have it automatically ingested & stored in DB. diff --git a/docs/modeling/extending-the-metadata-model.md b/docs/modeling/extending-the-metadata-model.md index 98f70f6d933e4..be2d7d795de70 100644 --- a/docs/modeling/extending-the-metadata-model.md +++ b/docs/modeling/extending-the-metadata-model.md @@ -24,7 +24,7 @@ We will refer to the two options as the **open-source fork** and **custom reposi ## This Guide This guide will outline what the experience of adding a new Entity should look like through a real example of adding the -Dashboard Entity. If you want to extend an existing Entity, you can skip directly to [Step 3](#step_3). +Dashboard Entity. If you want to extend an existing Entity, you can skip directly to [Step 3](#step-3-define-custom-aspects-or-attach-existing-aspects-to-your-entity). At a high level, an entity is made up of: @@ -82,14 +82,14 @@ Because they are aspects, keys need to be annotated with an @Aspect annotation, can be a part of. The key can also be annotated with the two index annotations: @Relationship and @Searchable. This instructs DataHub -infra to use the fields in the key to create relationships and index fields for search. See [Step 3](#step_3) for more details on +infra to use the fields in the key to create relationships and index fields for search. See [Step 3](#step-3-define-custom-aspects-or-attach-existing-aspects-to-your-entity) for more details on the annotation model. **Constraints**: Note that each field in a Key Aspect MUST be of String or Enum type. ### Step 2: Create the new entity with its key aspect -Define the entity within an `entity-registry.yml` file. Depending on your approach, the location of this file may vary. More on that in steps [4](#step_4) and [5](#step_5). +Define the entity within an `entity-registry.yml` file. Depending on your approach, the location of this file may vary. More on that in steps [4](#step-4-choose-a-place-to-store-your-model-extension) and [5](#step-5-attaching-your-non-key-aspects-to-the-entity). Example: ```yaml @@ -212,11 +212,11 @@ After you create your Aspect, you need to attach to all the entities that it app **Constraints**: Note that all aspects MUST be of type Record. -### Step 4: Choose a place to store your model extension +### Step 4: Choose a place to store your model extension At the beginning of this document, we walked you through a flow-chart that should help you decide whether you need to maintain a fork of the open source DataHub repo for your model extensions, or whether you can just use a model extension repository that can stay independent of the DataHub repo. Depending on what path you took, the place you store your aspect model files (the .pdl files) and the entity-registry files (the yaml file called `entity-registry.yaml` or `entity-registry.yml`) will vary. -- Open source Fork: Aspect files go under [`metadata-models`](../../metadata-models) module in the main repo, entity registry goes into [`metadata-models/src/main/resources/entity-registry.yml`](../../metadata-models/src/main/resources/entity-registry.yml). Read on for more details in [Step 5](#step_5). +- Open source Fork: Aspect files go under [`metadata-models`](../../metadata-models) module in the main repo, entity registry goes into [`metadata-models/src/main/resources/entity-registry.yml`](../../metadata-models/src/main/resources/entity-registry.yml). Read on for more details in [Step 5](#step-5-attaching-your-non-key-aspects-to-the-entity). - Custom repository: Read the [metadata-models-custom](../../metadata-models-custom/README.md) documentation to learn how to store and version your aspect models and registry. ### Step 5: Attaching your non-key Aspect(s) to the Entity diff --git a/docs/modeling/metadata-model.md b/docs/modeling/metadata-model.md index 037c9c7108a6e..a8958985a0a72 100644 --- a/docs/modeling/metadata-model.md +++ b/docs/modeling/metadata-model.md @@ -433,7 +433,7 @@ aggregation query against a timeseries aspect. The *@TimeseriesField* and the *@TimeseriesFieldCollection* are two new annotations that can be attached to a field of a *Timeseries aspect* that allows it to be part of an aggregatable query. The kinds of aggregations allowed on these annotated fields depends on the type of the field, as well as the kind of aggregation, as -described [here](#Performing-an-aggregation-on-a-Timeseries-aspect). +described [here](#performing-an-aggregation-on-a-timeseries-aspect). * `@TimeseriesField = {}` - this annotation can be used with any type of non-collection type field of the aspect such as primitive types and records (see the fields *stat*, *strStat* and *strArray* fields @@ -515,7 +515,7 @@ my_emitter = DatahubRestEmitter("http://localhost:8080") my_emitter.emit(mcpw) ``` -###### Performing an aggregation on a Timeseries aspect. +###### Performing an aggregation on a Timeseries aspect Aggreations on timeseries aspects can be performed by the GMS REST API for `/analytics?action=getTimeseriesStats` which accepts the following params. diff --git a/docs/tags.md b/docs/tags.md index 945b514dc7b47..cb08c9fafea49 100644 --- a/docs/tags.md +++ b/docs/tags.md @@ -27,25 +27,25 @@ You can create these privileges by creating a new [Metadata Policy](./authorizat To add a tag at the dataset or container level, simply navigate to the page for that entity and click on the **Add Tag** button.

- +

Type in the name of the tag you want to add. You can add a new tag, or add a tag that already exists (the autocomplete will pull up the tag if it already exists).

- +

Click on the "Add" button and you'll see the tag has been added!

- +

If you would like to add a tag at the schema level, hover over the "Tags" column for a schema until the "Add Tag" button shows up, and then follow the same flow as above.

- +

### Removing a Tag @@ -57,7 +57,7 @@ To remove a tag, simply click on the "X" button in the tag. Then click "Yes" whe You can search for a tag in the search bar, and even filter entities by the presence of a specific tag.

- +

## Additional Resources diff --git a/docs/townhall-history.md b/docs/townhall-history.md index e235a70c5d7b9..d92905af0cd72 100644 --- a/docs/townhall-history.md +++ b/docs/townhall-history.md @@ -328,7 +328,7 @@ November Town Hall (in December!) * Welcome - 5 mins * Latest React App Demo! ([video](https://www.youtube.com/watch?v=RQBEJhcen5E)) by John Joyce and Gabe Lyons - 5 mins -* Use-Case: DataHub at Geotab ([slides](https://docs.google.com/presentation/d/1qcgO3BW5NauuG0HnPqrxGcujsK-rJ1-EuU-7cbexkqE/edit?usp=sharing),[video](https://www.youtube.com/watch?v=boyjT2OrlU4)) by [John Yoon](https://www.linkedin.com/in/yhjyoon/) - 15 mins +* Use-Case: DataHub at Geotab ([video](https://www.youtube.com/watch?v=boyjT2OrlU4)) by [John Yoon](https://www.linkedin.com/in/yhjyoon/) - 15 mins * Tech Deep Dive: Tour of new pull-based Python Ingestion scripts ([slides](https://docs.google.com/presentation/d/15Xay596WDIhzkc5c8DEv6M-Bv1N4hP8quup1tkws6ms/edit#slide=id.gb478361595_0_10),[video](https://www.youtube.com/watch?v=u0IUQvG-_xI)) by [Harshal Sheth](https://www.linkedin.com/in/hsheth2/) - 15 mins * General Q&A from sign up sheet, slack, and participants - 15 mins * Closing remarks - 5 mins diff --git a/docs/what/gms.md b/docs/what/gms.md index 9e1cea1b9540e..a39450d28ae83 100644 --- a/docs/what/gms.md +++ b/docs/what/gms.md @@ -2,6 +2,4 @@ Metadata for [entities](entity.md) [onboarded](../modeling/metadata-model.md) to [GMA](gma.md) is served through microservices known as Generalized Metadata Service (GMS). GMS typically provides a [Rest.li](http://rest.li) API and must access the metadata using [GMA DAOs](../architecture/metadata-serving.md). -While a GMS is completely free to define its public APIs, we do provide a list of [resource base classes](https://github.com/datahub-project/datahub-gma/tree/master/restli-resources/src/main/java/com/linkedin/metadata/restli) to leverage for common patterns. - -GMA is designed to support a distributed fleet of GMS, each serving a subset of the [GMA graph](graph.md). However, for simplicity we include a single centralized GMS ([datahub-gms](../../gms)) that serves all entities. +GMA is designed to support a distributed fleet of GMS, each serving a subset of the [GMA graph](graph.md). However, for simplicity we include a single centralized GMS that serves all entities. diff --git a/docs/what/mxe.md b/docs/what/mxe.md index 8af96360858a3..25294e04ea3d9 100644 --- a/docs/what/mxe.md +++ b/docs/what/mxe.md @@ -266,7 +266,7 @@ A Metadata Change Event represents a request to change multiple aspects for the It leverages a deprecated concept of `Snapshot`, which is a strongly-typed list of aspects for the same entity. -A MCE is a "proposal" for a set of metadata changes, as opposed to [MAE](#metadata-audit-event), which is conveying a committed change. +A MCE is a "proposal" for a set of metadata changes, as opposed to [MAE](#metadata-audit-event-mae), which is conveying a committed change. Consequently, only successfully accepted and processed MCEs will lead to the emission of a corresponding MAE / MCLs. ### Emission diff --git a/docs/what/relationship.md b/docs/what/relationship.md index dcfe093a1b124..d5348dc04b3c0 100644 --- a/docs/what/relationship.md +++ b/docs/what/relationship.md @@ -102,9 +102,6 @@ For one, the actual direction doesn’t really impact the execution of graph que That being said, generally there’s a more "natural way" to specify the direction of a relationship, which closely relate to how the metadata is stored. For example, the membership information for an LDAP group is generally stored as a list in group’s metadata. As a result, it’s more natural to model a `HasMember` relationship that points from a group to a member, instead of a `IsMemberOf` relationship pointing from member to group. -Since all relationships are explicitly declared, it’s fairly easy for a user to discover what relationships are available and their directionality by inspecting -the [relationships directory](../../metadata-models/src/main/pegasus/com/linkedin/metadata/relationship). It’s also possible to provide a UI for the catalog of entities and relationships for analysts who are interested in building complex graph queries to gain insights into the metadata. - ## High Cardinality Relationships See [this doc](../advanced/high-cardinality.md) for suggestions on how to best model relationships with high cardinality. diff --git a/docs/what/search-document.md b/docs/what/search-document.md index 81359a55d0cae..bd27656e512c3 100644 --- a/docs/what/search-document.md +++ b/docs/what/search-document.md @@ -13,7 +13,6 @@ As a result, one may be tempted to add as many attributes as needed. This is acc Below shows an example schema for the `User` search document. Note that: 1. Each search document is required to have a type-specific `urn` field, generally maps to an entity in the [graph](graph.md). 2. Similar to `Entity`, each document has an optional `removed` field for "soft deletion". -This is captured in [BaseDocument](../../metadata-models/src/main/pegasus/com/linkedin/metadata/search/BaseDocument.pdl), which is expected to be included by all documents. 3. Similar to `Entity`, all remaining fields are made `optional` to support partial updates. 4. `management` shows an example of a string array field. 5. `ownedDataset` shows an example on how a field can be derived from metadata [aspects](aspect.md) associated with other types of entity (in this case, `Dataset`). diff --git a/metadata-ingestion/docs/dev_guides/add_stateful_ingestion_to_source.md b/metadata-ingestion/docs/dev_guides/add_stateful_ingestion_to_source.md index 6a1204fb0f2b3..9e39d24fb8578 100644 --- a/metadata-ingestion/docs/dev_guides/add_stateful_ingestion_to_source.md +++ b/metadata-ingestion/docs/dev_guides/add_stateful_ingestion_to_source.md @@ -60,16 +60,14 @@ class StaleEntityCheckpointStateBase(CheckpointStateBase, ABC, Generic[Derived]) ``` Examples: -1. [KafkaCheckpointState](https://github.com/datahub-project/datahub/blob/master/metadata-ingestion/src/datahub/ingestion/source/state/kafka_state.py#L11). -2. [DbtCheckpointState](https://github.com/datahub-project/datahub/blob/master/metadata-ingestion/src/datahub/ingestion/source/state/dbt_state.py#L16) -3. [BaseSQLAlchemyCheckpointState](https://github.com/datahub-project/datahub/blob/master/metadata-ingestion/src/datahub/ingestion/source/state/sql_common_state.py#L17) +* [BaseSQLAlchemyCheckpointState](https://github.com/datahub-project/datahub/blob/master/metadata-ingestion/src/datahub/ingestion/source/state/sql_common_state.py#L17) ### 2. Modifying the SourceConfig The source's config must inherit from `StatefulIngestionConfigBase`, and should declare a field named `stateful_ingestion` of type `Optional[StatefulStaleMetadataRemovalConfig]`. Examples: -1. The `KafkaSourceConfig` +- The `KafkaSourceConfig` ```python from typing import List, Optional import pydantic @@ -84,9 +82,6 @@ class KafkaSourceConfig(StatefulIngestionConfigBase): stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None ``` -2. The [DBTStatefulIngestionConfig](https://github.com/datahub-project/datahub/blob/master/metadata-ingestion/src/datahub/ingestion/source/dbt.py#L131) - and the [DBTConfig](https://github.com/datahub-project/datahub/blob/master/metadata-ingestion/src/datahub/ingestion/source/dbt.py#L317). - ### 3. Modifying the SourceReport The report class of the source should inherit from `StaleEntityRemovalSourceReport` whose definition is shown below. ```python @@ -102,7 +97,7 @@ class StaleEntityRemovalSourceReport(StatefulIngestionReport): ``` Examples: -1. The `KafkaSourceReport` +* The `KafkaSourceReport` ```python from dataclasses import dataclass from datahub.ingestion.source.state.stale_entity_removal_handler import StaleEntityRemovalSourceReport @@ -110,7 +105,7 @@ from datahub.ingestion.source.state.stale_entity_removal_handler import StaleEnt class KafkaSourceReport(StaleEntityRemovalSourceReport): # Date: Thu, 14 Sep 2023 11:40:38 +0530 Subject: [PATCH 011/156] docs(managed datahub): release notes 0.2.11 (#8830) --- docs-website/sidebars.js | 1 + .../managed-datahub/release-notes/v_0_2_11.md | 73 +++++++++++++++++++ 2 files changed, 74 insertions(+) create mode 100644 docs/managed-datahub/release-notes/v_0_2_11.md diff --git a/docs-website/sidebars.js b/docs-website/sidebars.js index fcf82b786a1b9..12691e9f8268a 100644 --- a/docs-website/sidebars.js +++ b/docs-website/sidebars.js @@ -597,6 +597,7 @@ module.exports = { }, { "Managed DataHub Release History": [ + "docs/managed-datahub/release-notes/v_0_2_11", "docs/managed-datahub/release-notes/v_0_2_10", "docs/managed-datahub/release-notes/v_0_2_9", "docs/managed-datahub/release-notes/v_0_2_8", diff --git a/docs/managed-datahub/release-notes/v_0_2_11.md b/docs/managed-datahub/release-notes/v_0_2_11.md new file mode 100644 index 0000000000000..1f42090848712 --- /dev/null +++ b/docs/managed-datahub/release-notes/v_0_2_11.md @@ -0,0 +1,73 @@ +# v0.2.11 +--- + +Release Availability Date +--- +14-Sep-2023 + +Recommended CLI/SDK +--- +- `v0.11.0` with release notes at https://github.com/acryldata/datahub/releases/tag/v0.10.5.5 +- [Deprecation] In LDAP ingestor, the manager_pagination_enabled changed to general pagination_enabled + +If you are using an older CLI/SDK version then please upgrade it. This applies for all CLI/SDK usages, if you are using it through your terminal, github actions, airflow, in python SDK somewhere, Java SKD etc. This is a strong recommendation to upgrade as we keep on pushing fixes in the CLI and it helps us support you better. + +Special Notes +--- +- Deployment process for this release is going to have a downtime when systme will be in a read only mode. A rough estimate 1 hour for every 2.3 million entities (includes soft-deleted entities). + + +## Release Changelog +--- +- Since `v0.2.10` these changes from OSS DataHub https://github.com/datahub-project/datahub/compare/2b0952195b7895df0a2bf92b28e71aac18217781...75252a3d9f6a576904be5a0790d644b9ae2df6ac have been pulled in. +- Misc fixes & features + - Proposals + - Group names shown correctly for proposal Inbox + - Metadata tests + - Deprecate/Un-deprecate actions available in Metadata tests + - Last Observed (in underlying sql) available as a filter in metadata tests + - [Breaking change] Renamed `__lastUpdated` -> `__created` as a filter to correctly represent what it was. This was not surfaced in the UI. But if you were using it then this needs to be renamed. Acryl Customer Success team will keep an eye out to pro-actively find and bring this up if you are affected by this. + - Robustness improvements to metadata test runs + - Copy urn for metadata tests to allow for easier filtering for iteration over metadata test results via our APIs. + - A lot more fixes to subscriptions, notifications and Observability (Beta). + - Some performance improvements to lineage queries + +## Some notable features in this SaaS release +- We now enable you to create and delete pinned announcements on your DataHub homepage! If you have the “Manage Home Page Posts” platform privilege you’ll see a new section in settings called “Home Page Posts” where you can create and delete text posts and link posts that your users see on the home page. +- Improvements to search experience +
+ -

- -**DataHub November 2022 Town Hall - Including Manual Lineage Demo** - -

- -

- -### GraphQL - -* [updateLineage](../../graphql/mutations.md#updatelineage) -* [searchAcrossLineage](../../graphql/queries.md#searchacrosslineage) -* [searchAcrossLineageInput](../../graphql/inputObjects.md#searchacrosslineageinput) - -#### Examples - -**Updating Lineage** - -```graphql -mutation updateLineage { - updateLineage(input: { - edgesToAdd: [ - { - downstreamUrn: "urn:li:dataset:(urn:li:dataPlatform:kafka,SampleKafkaDataset,PROD)", - upstreamUrn: "urn:li:dataset:(urn:li:dataPlatform:datahub,Dataset,PROD)" - } - ], - edgesToRemove: [ - { - downstreamUrn: "urn:li:dataset:(urn:li:dataPlatform:hdfs,SampleHdfsDataset,PROD)", - upstreamUrn: "urn:li:dataset:(urn:li:dataPlatform:kafka,SampleKafkaDataset,PROD)" - } - ] - }) -} -``` - -### DataHub Blog - -* [Acryl Data introduces lineage support and automated propagation of governance information for Snowflake in DataHub](https://blog.datahubproject.io/acryl-data-introduces-lineage-support-and-automated-propagation-of-governance-information-for-339c99536561) -* [Data in Context: Lineage Explorer in DataHub](https://blog.datahubproject.io/data-in-context-lineage-explorer-in-datahub-a53a9a476dc4) -* [Harnessing the Power of Data Lineage with DataHub](https://blog.datahubproject.io/harnessing-the-power-of-data-lineage-with-datahub-ad086358dec4) - -## FAQ and Troubleshooting - -**The Lineage Tab is greyed out - why can’t I click on it?** - -This means you have not yet ingested lineage metadata for that entity. Please ingest lineage to proceed. - -**Are there any recommended practices for emitting lineage?** - -We recommend emitting aspects as MetadataChangeProposalWrapper over emitting them via the MetadataChangeEvent. - -*Need more help? Join the conversation in [Slack](http://slack.datahubproject.io)!* - -### Related Features - -* [DataHub Lineage Impact Analysis](../act-on-metadata/impact-analysis.md) diff --git a/metadata-ingestion/scripts/docgen.py b/metadata-ingestion/scripts/docgen.py index b9f558011fc90..1a4db09e961ce 100644 --- a/metadata-ingestion/scripts/docgen.py +++ b/metadata-ingestion/scripts/docgen.py @@ -883,6 +883,150 @@ def generate( if metrics["plugins"].get("failed", 0) > 0: # type: ignore sys.exit(1) + ### Create Lineage doc + + source_dir = "../docs/generated/lineage" + os.makedirs(source_dir, exist_ok=True) + doc_file = f"{source_dir}/lineage-feature-guide.md" + with open(doc_file, "w+") as f: + f.write("import FeatureAvailability from '@site/src/components/FeatureAvailability';\n\n") + f.write(f"# About DataHub Lineage\n\n") + f.write("\n") + + f.write(""" +Lineage is used to capture data dependencies within an organization. It allows you to track the inputs from which a data asset is derived, along with the data assets that depend on it downstream. + +## Viewing Lineage + +You can view lineage under **Lineage** tab or **Lineage Visualization** screen. + +

+ +

+ +The UI shows the latest version of the lineage. The time picker can be used to filter out edges within the latest version to exclude those that were last updated outside of the time window. Selecting time windows in the patch will not show you historical lineages. It will only filter the view of the latest version of the lineage. + +

+ +

+ + +:::tip The Lineage Tab is greyed out - why can’t I click on it? +This means you have not yet ingested lineage metadata for that entity. Please ingest lineage to proceed. + +::: + +## Adding Lineage + +### Ingestion Source + +If you're using an ingestion source that supports extraction of Lineage (e.g. **Table Lineage Capability**), then lineage information can be extracted automatically. +For detailed instructions, refer to the [source documentation](https://datahubproject.io/integrations) for the source you are using. + +### UI + +As of `v0.9.5`, DataHub supports the manual editing of lineage between entities. Data experts are free to add or remove upstream and downstream lineage edges in both the Lineage Visualization screen as well as the Lineage tab on entity pages. Use this feature to supplement automatic lineage extraction or establish important entity relationships in sources that do not support automatic extraction. Editing lineage by hand is supported for Datasets, Charts, Dashboards, and Data Jobs. +Please refer to our [UI Guides on Lineage](../../features/feature-guides/ui-lineage.md) for more information. + +:::caution Recommendation on UI-based lineage + +Lineage added by hand and programmatically may conflict with one another to cause unwanted overwrites. +It is strongly recommend that lineage is edited manually in cases where lineage information is not also extracted in automated fashion, e.g. by running an ingestion source. + +::: + +### API + +If you are not using a Lineage-support ingestion source, you can programmatically emit lineage edges between entities via API. +Please refer to [API Guides on Lineage](../../api/tutorials/lineage.md) for more information. + + +## Lineage Support + +### Automatic Lineage Extraction Support + +This is a summary of automatic lineage extraciton support in our data source. Please refer to the **Important Capabilities** table in the source documentation. Note that even if the source does not support automatic extraction, you can still add lineage manually using our API & SDKs.\n""") + + f.write("\n| Source | Table-Level Lineage | Column-Level Lineage | Related Configs |\n") + f.write("| ---------- | ------ | ----- |----- |\n") + + for platform_id, platform_docs in sorted( + source_documentation.items(), + key=lambda x: (x[1]["name"].casefold(), x[1]["name"]) + if "name" in x[1] + else (x[0].casefold(), x[0]), + ): + for plugin, plugin_docs in sorted( + platform_docs["plugins"].items(), + key=lambda x: str(x[1].get("doc_order")) + if x[1].get("doc_order") + else x[0], + ): + platform_name = platform_docs['name'] + if len(platform_docs["plugins"].keys()) > 1: + # We only need to show this if there are multiple modules. + platform_name = f"{platform_name} `{plugin}`" + + # Initialize variables + table_level_supported = "❌" + column_level_supported = "❌" + config_names = '' + + if "capabilities" in plugin_docs: + plugin_capabilities = plugin_docs["capabilities"] + + for cap_setting in plugin_capabilities: + capability_text = get_capability_text(cap_setting.capability) + capability_supported = get_capability_supported_badge(cap_setting.supported) + + if capability_text == "Table-Level Lineage" and capability_supported == "✅": + table_level_supported = "✅" + + if capability_text == "Column-level Lineage" and capability_supported == "✅": + column_level_supported = "✅" + + if not (table_level_supported == "❌" and column_level_supported == "❌"): + if "config_schema" in plugin_docs: + config_properties = json.loads(plugin_docs['config_schema']).get('properties', {}) + config_names = '
'.join( + [f'- {property_name}' for property_name in config_properties if 'lineage' in property_name]) + lineage_not_applicable_sources = ['azure-ad', 'csv', 'demo-data', 'dynamodb', 'iceberg', 'json-schema', 'ldap', 'openapi', 'pulsar', 'sqlalchemy' ] + if platform_id not in lineage_not_applicable_sources : + f.write( + f"| [{platform_name}](../../generated/ingestion/sources/{platform_id}.md) | {table_level_supported} | {column_level_supported} | {config_names}|\n" + ) + + f.write(""" + +### Types of Lineage Connections + +Types of lineage connections supported in DataHub and the example codes are as follows. + +| Connection | Examples | A.K.A | +|---------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-----------------| +| Dataset to Dataset | - [lineage_emitter_mcpw_rest.py](../../../metadata-ingestion/examples/library/lineage_emitter_mcpw_rest.py)
- [lineage_emitter_rest.py](../../../metadata-ingestion/examples/library/lineage_emitter_rest.py)
- [lineage_emitter_kafka.py](../../../metadata-ingestion/examples/library/lineage_emitter_kafka.py)
- [lineage_emitter_dataset_finegrained.py](../../../metadata-ingestion/examples/library/lineage_emitter_dataset_finegrained.py)
- [Datahub BigQuery Lineage](https://github.com/datahub-project/datahub/blob/a1bf95307b040074c8d65ebb86b5eb177fdcd591/metadata-ingestion/src/datahub/ingestion/source/sql/bigquery.py#L229)
- [Datahub Snowflake Lineage](https://github.com/datahub-project/datahub/blob/master/metadata-ingestion/src/datahub/ingestion/source/sql/snowflake.py#L249) | +| DataJob to DataFlow | - [lineage_job_dataflow.py](../../../metadata-ingestion/examples/library/lineage_job_dataflow.py) | | +| DataJob to Dataset | - [lineage_dataset_job_dataset.py](../../../metadata-ingestion/examples/library/lineage_dataset_job_dataset.py)
| Pipeline Lineage | +| Chart to Dashboard | - [lineage_chart_dashboard.py](../../../metadata-ingestion/examples/library/lineage_chart_dashboard.py) | | +| Chart to Dataset | - [lineage_dataset_chart.py](../../../metadata-ingestion/examples/library/lineage_dataset_chart.py) | | + + +:::tip Our Roadmap +We're actively working on expanding lineage support for new data sources. +Visit our [Official Roadmap](https://feature-requests.datahubproject.io/roadmap) for upcoming updates! +::: + +## References + +- [DataHub Basics: Lineage 101](https://www.youtube.com/watch?v=rONGpsndzRw&t=1s) +- [DataHub November 2022 Town Hall](https://www.youtube.com/watch?v=BlCLhG8lGoY&t=1s) - Including Manual Lineage Demo +- [Acryl Data introduces lineage support and automated propagation of governance information for Snowflake in DataHub](https://blog.datahubproject.io/acryl-data-introduces-lineage-support-and-automated-propagation-of-governance-information-for-339c99536561) +- [Data in Context: Lineage Explorer in DataHub](https://blog.datahubproject.io/data-in-context-lineage-explorer-in-datahub-a53a9a476dc4) +- [Harnessing the Power of Data Lineage with DataHub](https://blog.datahubproject.io/harnessing-the-power-of-data-lineage-with-datahub-ad086358dec4) +- [DataHub Lineage Impact Analysis](https://datahubproject.io/docs/next/act-on-metadata/impact-analysis) + """) + + print("Lineage Documentation Generation Complete") if __name__ == "__main__": logger.setLevel("INFO") From c415d63ddae884de4e7a5d4ff3311f82057d3a78 Mon Sep 17 00:00:00 2001 From: siddiquebagwan-gslab Date: Wed, 4 Oct 2023 16:22:51 +0530 Subject: [PATCH 081/156] feat(ingestion/powerbi): column level lineage extraction for M-Query (#8796) --- .../docs/sources/powerbi/powerbi_pre.md | 2 +- .../ingestion/source/powerbi/config.py | 36 + .../powerbi/m_query/native_sql_parser.py | 6 +- .../source/powerbi/m_query/parser.py | 2 +- .../source/powerbi/m_query/resolver.py | 189 ++- .../ingestion/source/powerbi/powerbi.py | 102 +- .../integration/powerbi/golden_test_cll.json | 1357 +++++++++++++++++ .../integration/powerbi/test_m_parser.py | 155 +- .../tests/integration/powerbi/test_powerbi.py | 95 +- 9 files changed, 1804 insertions(+), 140 deletions(-) create mode 100644 metadata-ingestion/tests/integration/powerbi/golden_test_cll.json diff --git a/metadata-ingestion/docs/sources/powerbi/powerbi_pre.md b/metadata-ingestion/docs/sources/powerbi/powerbi_pre.md index 0323e214045ae..fcfae6cd1e6d7 100644 --- a/metadata-ingestion/docs/sources/powerbi/powerbi_pre.md +++ b/metadata-ingestion/docs/sources/powerbi/powerbi_pre.md @@ -40,7 +40,7 @@ PowerBI Source supports M-Query expression for below listed PowerBI Data Sources 4. Microsoft SQL Server 5. Google BigQuery -Native SQL query parsing is supported for `Snowflake` and `Amazon Redshift` data-sources and only first table from `FROM` clause will be ingested as upstream table. Advance SQL construct like JOIN and SUB-QUERIES in `FROM` clause are not supported. +Native SQL query parsing is supported for `Snowflake` and `Amazon Redshift` data-sources. For example refer below native SQL query. The table `OPERATIONS_ANALYTICS.TRANSFORMED_PROD.V_UNIT_TARGET` will be ingested as upstream table. diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py index ffa685fb25826..a8c7e48f3785c 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py @@ -397,6 +397,42 @@ class PowerBiDashboardSourceConfig( "as this option generates the upstream datasets URN in lowercase.", ) + # Enable CLL extraction + extract_column_level_lineage: bool = pydantic.Field( + default=False, + description="Whether to extract column level lineage. " + "Works only if configs `native_query_parsing`, `enable_advance_lineage_sql_construct` & `extract_lineage` are enabled. " + "Works for M-Query where native SQL is used for transformation.", + ) + + @root_validator + @classmethod + def validate_extract_column_level_lineage(cls, values: Dict) -> Dict: + flags = [ + "native_query_parsing", + "enable_advance_lineage_sql_construct", + "extract_lineage", + ] + + if ( + "extract_column_level_lineage" in values + and values["extract_column_level_lineage"] is False + ): + # Flag is not set. skip validation + return values + + logger.debug(f"Validating additional flags: {flags}") + + is_flag_enabled: bool = True + for flag in flags: + if flag not in values or values[flag] is False: + is_flag_enabled = False + + if not is_flag_enabled: + raise ValueError(f"Enable all these flags in recipe: {flags} ") + + return values + @validator("dataset_type_mapping") @classmethod def map_data_platform(cls, value): diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/native_sql_parser.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/native_sql_parser.py index 021c429c3c633..0afa8e7ff4564 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/native_sql_parser.py +++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/native_sql_parser.py @@ -9,7 +9,7 @@ SPECIAL_CHARACTERS = ["#(lf)", "(lf)"] -logger = logging.getLogger() +logger = logging.getLogger(__name__) def remove_special_characters(native_query: str) -> str: @@ -21,7 +21,7 @@ def remove_special_characters(native_query: str) -> str: def get_tables(native_query: str) -> List[str]: native_query = remove_special_characters(native_query) - logger.debug(f"Processing query = {native_query}") + logger.debug(f"Processing native query = {native_query}") tables: List[str] = [] parsed = sqlparse.parse(native_query)[0] tokens: List[sqlparse.sql.Token] = list(parsed.tokens) @@ -65,7 +65,7 @@ def parse_custom_sql( sql_query = remove_special_characters(query) - logger.debug(f"Parsing sql={sql_query}") + logger.debug(f"Processing native query = {sql_query}") return sqlglot_l.create_lineage_sql_parsed_result( query=sql_query, diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/parser.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/parser.py index 8cc38c366c42a..9134932c39fe0 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/parser.py +++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/parser.py @@ -56,7 +56,7 @@ def get_upstream_tables( ctx: PipelineContext, config: PowerBiDashboardSourceConfig, parameters: Dict[str, str] = {}, -) -> List[resolver.DataPlatformTable]: +) -> List[resolver.Lineage]: if table.expression is None: logger.debug(f"Expression is none for table {table.full_name}") return [] diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/resolver.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/resolver.py index 479f1decff903..e200ff41f71c2 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/resolver.py +++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/resolver.py @@ -27,7 +27,7 @@ IdentifierAccessor, ) from datahub.ingestion.source.powerbi.rest_api_wrapper.data_classes import Table -from datahub.utilities.sqlglot_lineage import SqlParsingResult +from datahub.utilities.sqlglot_lineage import ColumnLineageInfo, SqlParsingResult logger = logging.getLogger(__name__) @@ -38,6 +38,16 @@ class DataPlatformTable: urn: str +@dataclass +class Lineage: + upstreams: List[DataPlatformTable] + column_lineage: List[ColumnLineageInfo] + + @staticmethod + def empty() -> "Lineage": + return Lineage(upstreams=[], column_lineage=[]) + + def urn_to_lowercase(value: str, flag: bool) -> str: if flag is True: return value.lower() @@ -120,9 +130,9 @@ def __init__( self.platform_instance_resolver = platform_instance_resolver @abstractmethod - def create_dataplatform_tables( + def create_lineage( self, data_access_func_detail: DataAccessFunctionDetail - ) -> List[DataPlatformTable]: + ) -> Lineage: pass @abstractmethod @@ -147,7 +157,7 @@ def get_db_detail_from_argument( def parse_custom_sql( self, query: str, server: str, database: Optional[str], schema: Optional[str] - ) -> List[DataPlatformTable]: + ) -> Lineage: dataplatform_tables: List[DataPlatformTable] = [] @@ -174,7 +184,7 @@ def parse_custom_sql( if parsed_result is None: logger.debug("Failed to parse query") - return dataplatform_tables + return Lineage.empty() for urn in parsed_result.in_tables: dataplatform_tables.append( @@ -184,9 +194,15 @@ def parse_custom_sql( ) ) + logger.debug(f"Native Query parsed result={parsed_result}") logger.debug(f"Generated dataplatform_tables={dataplatform_tables}") - return dataplatform_tables + return Lineage( + upstreams=dataplatform_tables, + column_lineage=parsed_result.column_lineage + if parsed_result.column_lineage is not None + else [], + ) class AbstractDataAccessMQueryResolver(ABC): @@ -215,7 +231,7 @@ def resolve_to_data_platform_table_list( ctx: PipelineContext, config: PowerBiDashboardSourceConfig, platform_instance_resolver: AbstractDataPlatformInstanceResolver, - ) -> List[DataPlatformTable]: + ) -> List[Lineage]: pass @@ -471,8 +487,8 @@ def resolve_to_data_platform_table_list( ctx: PipelineContext, config: PowerBiDashboardSourceConfig, platform_instance_resolver: AbstractDataPlatformInstanceResolver, - ) -> List[DataPlatformTable]: - data_platform_tables: List[DataPlatformTable] = [] + ) -> List[Lineage]: + lineage: List[Lineage] = [] # Find out output variable as we are doing backtracking in M-Query output_variable: Optional[str] = tree_function.get_output_variable( @@ -484,7 +500,7 @@ def resolve_to_data_platform_table_list( f"{self.table.full_name}-output-variable", "output-variable not found in table expression", ) - return data_platform_tables + return lineage # Parse M-Query and use output_variable as root of tree and create instance of DataAccessFunctionDetail table_links: List[ @@ -509,7 +525,7 @@ def resolve_to_data_platform_table_list( # From supported_resolver enum get respective resolver like AmazonRedshift or Snowflake or Oracle or NativeQuery and create instance of it # & also pass additional information that will be need to generate urn - table_full_name_creator: AbstractDataPlatformTableCreator = ( + table_qualified_name_creator: AbstractDataPlatformTableCreator = ( supported_resolver.get_table_full_name_creator()( ctx=ctx, config=config, @@ -517,11 +533,9 @@ def resolve_to_data_platform_table_list( ) ) - data_platform_tables.extend( - table_full_name_creator.create_dataplatform_tables(f_detail) - ) + lineage.append(table_qualified_name_creator.create_lineage(f_detail)) - return data_platform_tables + return lineage class DefaultTwoStepDataAccessSources(AbstractDataPlatformTableCreator, ABC): @@ -536,7 +550,7 @@ class DefaultTwoStepDataAccessSources(AbstractDataPlatformTableCreator, ABC): def two_level_access_pattern( self, data_access_func_detail: DataAccessFunctionDetail - ) -> List[DataPlatformTable]: + ) -> Lineage: logger.debug( f"Processing {self.get_platform_pair().powerbi_data_platform_name} data-access function detail {data_access_func_detail}" ) @@ -545,7 +559,7 @@ def two_level_access_pattern( data_access_func_detail.arg_list ) if server is None or db_name is None: - return [] # Return empty list + return Lineage.empty() # Return empty list schema_name: str = cast( IdentifierAccessor, data_access_func_detail.identifier_accessor @@ -568,19 +582,21 @@ def two_level_access_pattern( server=server, qualified_table_name=qualified_table_name, ) - - return [ - DataPlatformTable( - data_platform_pair=self.get_platform_pair(), - urn=urn, - ) - ] + return Lineage( + upstreams=[ + DataPlatformTable( + data_platform_pair=self.get_platform_pair(), + urn=urn, + ) + ], + column_lineage=[], + ) class PostgresDataPlatformTableCreator(DefaultTwoStepDataAccessSources): - def create_dataplatform_tables( + def create_lineage( self, data_access_func_detail: DataAccessFunctionDetail - ) -> List[DataPlatformTable]: + ) -> Lineage: return self.two_level_access_pattern(data_access_func_detail) def get_platform_pair(self) -> DataPlatformPair: @@ -630,10 +646,10 @@ def create_urn_using_old_parser( return dataplatform_tables - def create_dataplatform_tables( + def create_lineage( self, data_access_func_detail: DataAccessFunctionDetail - ) -> List[DataPlatformTable]: - dataplatform_tables: List[DataPlatformTable] = [] + ) -> Lineage: + arguments: List[str] = tree_function.strip_char_from_list( values=tree_function.remove_whitespaces_from_list( tree_function.token_values(data_access_func_detail.arg_list) @@ -647,14 +663,17 @@ def create_dataplatform_tables( if len(arguments) >= 4 and arguments[2] != "Query": logger.debug("Unsupported case is found. Second index is not the Query") - return dataplatform_tables + return Lineage.empty() if self.config.enable_advance_lineage_sql_construct is False: # Use previous parser to generate URN to keep backward compatibility - return self.create_urn_using_old_parser( - query=arguments[3], - db_name=arguments[1], - server=arguments[0], + return Lineage( + upstreams=self.create_urn_using_old_parser( + query=arguments[3], + db_name=arguments[1], + server=arguments[0], + ), + column_lineage=[], ) return self.parse_custom_sql( @@ -684,9 +703,9 @@ def _get_server_and_db_name(value: str) -> Tuple[Optional[str], Optional[str]]: return tree_function.strip_char_from_list([splitter_result[0]])[0], db_name - def create_dataplatform_tables( + def create_lineage( self, data_access_func_detail: DataAccessFunctionDetail - ) -> List[DataPlatformTable]: + ) -> Lineage: logger.debug( f"Processing Oracle data-access function detail {data_access_func_detail}" ) @@ -698,7 +717,7 @@ def create_dataplatform_tables( server, db_name = self._get_server_and_db_name(arguments[0]) if db_name is None or server is None: - return [] + return Lineage.empty() schema_name: str = cast( IdentifierAccessor, data_access_func_detail.identifier_accessor @@ -719,18 +738,21 @@ def create_dataplatform_tables( qualified_table_name=qualified_table_name, ) - return [ - DataPlatformTable( - data_platform_pair=self.get_platform_pair(), - urn=urn, - ) - ] + return Lineage( + upstreams=[ + DataPlatformTable( + data_platform_pair=self.get_platform_pair(), + urn=urn, + ) + ], + column_lineage=[], + ) class DatabrickDataPlatformTableCreator(AbstractDataPlatformTableCreator): - def create_dataplatform_tables( + def create_lineage( self, data_access_func_detail: DataAccessFunctionDetail - ) -> List[DataPlatformTable]: + ) -> Lineage: logger.debug( f"Processing Databrick data-access function detail {data_access_func_detail}" ) @@ -749,7 +771,7 @@ def create_dataplatform_tables( logger.debug( "expecting instance to be IdentifierAccessor, please check if parsing is done properly" ) - return [] + return Lineage.empty() db_name: str = value_dict["Database"] schema_name: str = value_dict["Schema"] @@ -762,7 +784,7 @@ def create_dataplatform_tables( logger.info( f"server information is not available for {qualified_table_name}. Skipping upstream table" ) - return [] + return Lineage.empty() urn = urn_creator( config=self.config, @@ -772,12 +794,15 @@ def create_dataplatform_tables( qualified_table_name=qualified_table_name, ) - return [ - DataPlatformTable( - data_platform_pair=self.get_platform_pair(), - urn=urn, - ) - ] + return Lineage( + upstreams=[ + DataPlatformTable( + data_platform_pair=self.get_platform_pair(), + urn=urn, + ) + ], + column_lineage=[], + ) def get_platform_pair(self) -> DataPlatformPair: return SupportedDataPlatform.DATABRICK_SQL.value @@ -789,9 +814,9 @@ def get_datasource_server( ) -> str: return tree_function.strip_char_from_list([arguments[0]])[0] - def create_dataplatform_tables( + def create_lineage( self, data_access_func_detail: DataAccessFunctionDetail - ) -> List[DataPlatformTable]: + ) -> Lineage: logger.debug( f"Processing {self.get_platform_pair().datahub_data_platform_name} function detail {data_access_func_detail}" ) @@ -826,12 +851,15 @@ def create_dataplatform_tables( qualified_table_name=qualified_table_name, ) - return [ - DataPlatformTable( - data_platform_pair=self.get_platform_pair(), - urn=urn, - ) - ] + return Lineage( + upstreams=[ + DataPlatformTable( + data_platform_pair=self.get_platform_pair(), + urn=urn, + ) + ], + column_lineage=[], + ) class SnowflakeDataPlatformTableCreator(DefaultThreeStepDataAccessSources): @@ -859,9 +887,9 @@ class AmazonRedshiftDataPlatformTableCreator(AbstractDataPlatformTableCreator): def get_platform_pair(self) -> DataPlatformPair: return SupportedDataPlatform.AMAZON_REDSHIFT.value - def create_dataplatform_tables( + def create_lineage( self, data_access_func_detail: DataAccessFunctionDetail - ) -> List[DataPlatformTable]: + ) -> Lineage: logger.debug( f"Processing AmazonRedshift data-access function detail {data_access_func_detail}" ) @@ -870,7 +898,7 @@ def create_dataplatform_tables( data_access_func_detail.arg_list ) if db_name is None or server is None: - return [] # Return empty list + return Lineage.empty() # Return empty list schema_name: str = cast( IdentifierAccessor, data_access_func_detail.identifier_accessor @@ -891,12 +919,15 @@ def create_dataplatform_tables( qualified_table_name=qualified_table_name, ) - return [ - DataPlatformTable( - data_platform_pair=self.get_platform_pair(), - urn=urn, - ) - ] + return Lineage( + upstreams=[ + DataPlatformTable( + data_platform_pair=self.get_platform_pair(), + urn=urn, + ) + ], + column_lineage=[], + ) class NativeQueryDataPlatformTableCreator(AbstractDataPlatformTableCreator): @@ -916,9 +947,7 @@ def is_native_parsing_supported(data_access_function_name: str) -> bool: in NativeQueryDataPlatformTableCreator.SUPPORTED_NATIVE_QUERY_DATA_PLATFORM ) - def create_urn_using_old_parser( - self, query: str, server: str - ) -> List[DataPlatformTable]: + def create_urn_using_old_parser(self, query: str, server: str) -> Lineage: dataplatform_tables: List[DataPlatformTable] = [] tables: List[str] = native_sql_parser.get_tables(query) @@ -947,12 +976,14 @@ def create_urn_using_old_parser( logger.debug(f"Generated dataplatform_tables {dataplatform_tables}") - return dataplatform_tables + return Lineage( + upstreams=dataplatform_tables, + column_lineage=[], + ) - def create_dataplatform_tables( + def create_lineage( self, data_access_func_detail: DataAccessFunctionDetail - ) -> List[DataPlatformTable]: - dataplatform_tables: List[DataPlatformTable] = [] + ) -> Lineage: t1: Tree = cast( Tree, tree_function.first_arg_list_func(data_access_func_detail.arg_list) ) @@ -963,7 +994,7 @@ def create_dataplatform_tables( f"Expecting 2 argument, actual argument count is {len(flat_argument_list)}" ) logger.debug(f"Flat argument list = {flat_argument_list}") - return dataplatform_tables + return Lineage.empty() data_access_tokens: List[str] = tree_function.remove_whitespaces_from_list( tree_function.token_values(flat_argument_list[0]) ) @@ -981,7 +1012,7 @@ def create_dataplatform_tables( f"Server is not available in argument list for data-platform {data_access_tokens[0]}. Returning empty " "list" ) - return dataplatform_tables + return Lineage.empty() self.current_data_platform = self.SUPPORTED_NATIVE_QUERY_DATA_PLATFORM[ data_access_tokens[0] diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/powerbi.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/powerbi.py index 5d477ee090e7e..52bcef66658c8 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/powerbi.py +++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/powerbi.py @@ -44,6 +44,11 @@ StatefulIngestionSourceBase, ) from datahub.metadata.com.linkedin.pegasus2avro.common import ChangeAuditStamps +from datahub.metadata.com.linkedin.pegasus2avro.dataset import ( + FineGrainedLineage, + FineGrainedLineageDownstreamType, + FineGrainedLineageUpstreamType, +) from datahub.metadata.schema_classes import ( BrowsePathsClass, ChangeTypeClass, @@ -71,6 +76,7 @@ ViewPropertiesClass, ) from datahub.utilities.dedup_list import deduplicate_list +from datahub.utilities.sqlglot_lineage import ColumnLineageInfo # Logger instance logger = logging.getLogger(__name__) @@ -165,6 +171,48 @@ def extract_dataset_schema( ) return [schema_mcp] + def make_fine_grained_lineage_class( + self, lineage: resolver.Lineage, dataset_urn: str + ) -> List[FineGrainedLineage]: + fine_grained_lineages: List[FineGrainedLineage] = [] + + if ( + self.__config.extract_column_level_lineage is False + or self.__config.extract_lineage is False + ): + return fine_grained_lineages + + if lineage is None: + return fine_grained_lineages + + logger.info("Extracting column level lineage") + + cll: List[ColumnLineageInfo] = lineage.column_lineage + + for cll_info in cll: + downstream = ( + [builder.make_schema_field_urn(dataset_urn, cll_info.downstream.column)] + if cll_info.downstream is not None + and cll_info.downstream.column is not None + else [] + ) + + upstreams = [ + builder.make_schema_field_urn(column_ref.table, column_ref.column) + for column_ref in cll_info.upstreams + ] + + fine_grained_lineages.append( + FineGrainedLineage( + downstreamType=FineGrainedLineageDownstreamType.FIELD, + downstreams=downstream, + upstreamType=FineGrainedLineageUpstreamType.FIELD_SET, + upstreams=upstreams, + ) + ) + + return fine_grained_lineages + def extract_lineage( self, table: powerbi_data_classes.Table, ds_urn: str ) -> List[MetadataChangeProposalWrapper]: @@ -174,8 +222,9 @@ def extract_lineage( parameters = table.dataset.parameters if table.dataset else {} upstream: List[UpstreamClass] = [] + cll_lineage: List[FineGrainedLineage] = [] - upstream_dpts: List[resolver.DataPlatformTable] = parser.get_upstream_tables( + upstream_lineage: List[resolver.Lineage] = parser.get_upstream_tables( table=table, reporter=self.__reporter, platform_instance_resolver=self.__dataplatform_instance_resolver, @@ -185,34 +234,49 @@ def extract_lineage( ) logger.debug( - f"PowerBI virtual table {table.full_name} and it's upstream dataplatform tables = {upstream_dpts}" + f"PowerBI virtual table {table.full_name} and it's upstream dataplatform tables = {upstream_lineage}" ) - for upstream_dpt in upstream_dpts: - if ( - upstream_dpt.data_platform_pair.powerbi_data_platform_name - not in self.__config.dataset_type_mapping.keys() - ): - logger.debug( - f"Skipping upstream table for {ds_urn}. The platform {upstream_dpt.data_platform_pair.powerbi_data_platform_name} is not part of dataset_type_mapping", + for lineage in upstream_lineage: + for upstream_dpt in lineage.upstreams: + if ( + upstream_dpt.data_platform_pair.powerbi_data_platform_name + not in self.__config.dataset_type_mapping.keys() + ): + logger.debug( + f"Skipping upstream table for {ds_urn}. The platform {upstream_dpt.data_platform_pair.powerbi_data_platform_name} is not part of dataset_type_mapping", + ) + continue + + upstream_table_class = UpstreamClass( + upstream_dpt.urn, + DatasetLineageTypeClass.TRANSFORMED, ) - continue - upstream_table_class = UpstreamClass( - upstream_dpt.urn, - DatasetLineageTypeClass.TRANSFORMED, - ) + upstream.append(upstream_table_class) - upstream.append(upstream_table_class) + # Add column level lineage if any + cll_lineage.extend( + self.make_fine_grained_lineage_class( + lineage=lineage, + dataset_urn=ds_urn, + ) + ) if len(upstream) > 0: - upstream_lineage = UpstreamLineageClass(upstreams=upstream) + + upstream_lineage_class: UpstreamLineageClass = UpstreamLineageClass( + upstreams=upstream, + fineGrainedLineages=cll_lineage or None, + ) + logger.debug(f"Dataset urn = {ds_urn} and its lineage = {upstream_lineage}") + mcp = MetadataChangeProposalWrapper( entityType=Constant.DATASET, changeType=ChangeTypeClass.UPSERT, entityUrn=ds_urn, - aspect=upstream_lineage, + aspect=upstream_lineage_class, ) mcps.append(mcp) @@ -1075,6 +1139,10 @@ def report_to_datahub_work_units( SourceCapability.OWNERSHIP, "Disabled by default, configured using `extract_ownership`", ) +@capability( + SourceCapability.LINEAGE_FINE, + "Disabled by default, configured using `extract_column_level_lineage`. ", +) class PowerBiDashboardSource(StatefulIngestionSourceBase): """ This plugin extracts the following: diff --git a/metadata-ingestion/tests/integration/powerbi/golden_test_cll.json b/metadata-ingestion/tests/integration/powerbi/golden_test_cll.json new file mode 100644 index 0000000000000..5f92cdcfb5bde --- /dev/null +++ b/metadata-ingestion/tests/integration/powerbi/golden_test_cll.json @@ -0,0 +1,1357 @@ +[ +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.public_issue_history,DEV)", + "changeType": "UPSERT", + "aspectName": "viewProperties", + "aspect": { + "json": { + "materialized": false, + "viewLogic": "dummy", + "viewLanguage": "m_query" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.public_issue_history,DEV)", + "changeType": "UPSERT", + "aspectName": "datasetProperties", + "aspect": { + "json": { + "customProperties": { + "datasetId": "05169CD2-E713-41E6-9600-1D8066D95445" + }, + "externalUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/05169CD2-E713-41E6-9600-1D8066D95445/details", + "name": "public issue_history", + "description": "Library dataset description", + "tags": [] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.public_issue_history,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.public_issue_history,DEV)", + "changeType": "UPSERT", + "aspectName": "subTypes", + "aspect": { + "json": { + "typeNames": [ + "PowerBI Dataset Table", + "View" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.SNOWFLAKE_TESTTABLE,DEV)", + "changeType": "UPSERT", + "aspectName": "viewProperties", + "aspect": { + "json": { + "materialized": false, + "viewLogic": "let\n Source = Snowflake.Databases(\"hp123rt5.ap-southeast-2.fakecomputing.com\",\"PBI_TEST_WAREHOUSE_PROD\",[Role=\"PBI_TEST_MEMBER\"]),\n PBI_TEST_Database = Source{[Name=\"PBI_TEST\",Kind=\"Database\"]}[Data],\n TEST_Schema = PBI_TEST_Database{[Name=\"TEST\",Kind=\"Schema\"]}[Data],\n TESTTABLE_Table = TEST_Schema{[Name=\"TESTTABLE\",Kind=\"Table\"]}[Data]\nin\n TESTTABLE_Table", + "viewLanguage": "m_query" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.SNOWFLAKE_TESTTABLE,DEV)", + "changeType": "UPSERT", + "aspectName": "datasetProperties", + "aspect": { + "json": { + "customProperties": { + "datasetId": "05169CD2-E713-41E6-9600-1D8066D95445" + }, + "externalUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/05169CD2-E713-41E6-9600-1D8066D95445/details", + "name": "SNOWFLAKE_TESTTABLE", + "description": "Library dataset description", + "tags": [] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.SNOWFLAKE_TESTTABLE,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.SNOWFLAKE_TESTTABLE,DEV)", + "changeType": "UPSERT", + "aspectName": "subTypes", + "aspect": { + "json": { + "typeNames": [ + "PowerBI Dataset Table", + "View" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.SNOWFLAKE_TESTTABLE,DEV)", + "changeType": "UPSERT", + "aspectName": "upstreamLineage", + "aspect": { + "json": { + "upstreams": [ + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:snowflake,PBI_TEST.TEST.TESTTABLE,PROD)", + "type": "TRANSFORMED" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query,DEV)", + "changeType": "UPSERT", + "aspectName": "viewProperties", + "aspect": { + "json": { + "materialized": false, + "viewLogic": "let\n Source = Value.NativeQuery(Snowflake.Databases(\"bu20658.ap-southeast-2.snowflakecomputing.com\",\"operations_analytics_warehouse_prod\",[Role=\"OPERATIONS_ANALYTICS_MEMBER\"]){[Name=\"OPERATIONS_ANALYTICS\"]}[Data], \"SELECT#(lf)concat((UPPER(REPLACE(SELLER,'-',''))), MONTHID) as AGENT_KEY,#(lf)concat((UPPER(REPLACE(CLIENT_DIRECTOR,'-',''))), MONTHID) as CD_AGENT_KEY,#(lf) *#(lf)FROM#(lf)OPERATIONS_ANALYTICS.TRANSFORMED_PROD.V_APS_SME_UNITS_V4\", null, [EnableFolding=true]),\n #\"Added Conditional Column\" = Table.AddColumn(Source, \"SME Units ENT\", each if [DEAL_TYPE] = \"SME Unit\" then [UNIT] else 0),\n #\"Added Conditional Column1\" = Table.AddColumn(#\"Added Conditional Column\", \"Banklink Units\", each if [DEAL_TYPE] = \"Banklink\" then [UNIT] else 0),\n #\"Removed Columns\" = Table.RemoveColumns(#\"Added Conditional Column1\",{\"Banklink Units\"}),\n #\"Added Custom\" = Table.AddColumn(#\"Removed Columns\", \"Banklink Units\", each if [DEAL_TYPE] = \"Banklink\" and [SALES_TYPE] = \"3 - Upsell\"\nthen [UNIT]\n\nelse if [SALES_TYPE] = \"Adjusted BL Migration\"\nthen [UNIT]\n\nelse 0),\n #\"Added Custom1\" = Table.AddColumn(#\"Added Custom\", \"SME Units in $ (*$361)\", each if [DEAL_TYPE] = \"SME Unit\" \nand [SALES_TYPE] <> \"4 - Renewal\"\n then [UNIT] * 361\nelse 0),\n #\"Added Custom2\" = Table.AddColumn(#\"Added Custom1\", \"Banklink in $ (*$148)\", each [Banklink Units] * 148)\nin\n #\"Added Custom2\"", + "viewLanguage": "m_query" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query,DEV)", + "changeType": "UPSERT", + "aspectName": "datasetProperties", + "aspect": { + "json": { + "customProperties": { + "datasetId": "05169CD2-E713-41E6-9600-1D8066D95445" + }, + "externalUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/05169CD2-E713-41E6-9600-1D8066D95445/details", + "name": "snowflake native-query", + "description": "Library dataset description", + "tags": [] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query,DEV)", + "changeType": "UPSERT", + "aspectName": "subTypes", + "aspect": { + "json": { + "typeNames": [ + "PowerBI Dataset Table", + "View" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query,DEV)", + "changeType": "UPSERT", + "aspectName": "upstreamLineage", + "aspect": { + "json": { + "upstreams": [ + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:snowflake,operations_analytics.transformed_prod.v_aps_sme_units_v4,PROD)", + "type": "TRANSFORMED" + } + ], + "fineGrainedLineages": [ + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,operations_analytics.transformed_prod.v_aps_sme_units_v4,PROD),monthid)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,operations_analytics.transformed_prod.v_aps_sme_units_v4,PROD),seller)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query,DEV),agent_key)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,operations_analytics.transformed_prod.v_aps_sme_units_v4,PROD),client_director)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,operations_analytics.transformed_prod.v_aps_sme_units_v4,PROD),monthid)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query,DEV),cd_agent_key)" + ], + "confidenceScore": 1.0 + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.big-query-with-parameter,DEV)", + "changeType": "UPSERT", + "aspectName": "viewProperties", + "aspect": { + "json": { + "materialized": false, + "viewLogic": "let\n Source = GoogleBigQuery.Database([BillingProject = #\"Parameter - Source\"]),\n#\"gcp-project\" = Source{[Name=#\"Parameter - Source\"]}[Data],\nuniversal_Schema = #\"gcp-project\"{[Name=\"universal\",Kind=\"Schema\"]}[Data],\nD_WH_DATE_Table = universal_Schema{[Name=\"D_WH_DATE\",Kind=\"Table\"]}[Data],\n#\"Filtered Rows\" = Table.SelectRows(D_WH_DATE_Table, each [D_DATE] > #datetime(2019, 9, 10, 0, 0, 0)),\n#\"Filtered Rows1\" = Table.SelectRows(#\"Filtered Rows\", each DateTime.IsInPreviousNHours([D_DATE], 87600))\n in \n#\"Filtered Rows1\"", + "viewLanguage": "m_query" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.big-query-with-parameter,DEV)", + "changeType": "UPSERT", + "aspectName": "datasetProperties", + "aspect": { + "json": { + "customProperties": { + "datasetId": "05169CD2-E713-41E6-9600-1D8066D95445" + }, + "externalUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/05169CD2-E713-41E6-9600-1D8066D95445/details", + "name": "big-query-with-parameter", + "description": "Library dataset description", + "tags": [] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.big-query-with-parameter,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.big-query-with-parameter,DEV)", + "changeType": "UPSERT", + "aspectName": "subTypes", + "aspect": { + "json": { + "typeNames": [ + "PowerBI Dataset Table", + "View" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query-with-join,DEV)", + "changeType": "UPSERT", + "aspectName": "viewProperties", + "aspect": { + "json": { + "materialized": false, + "viewLogic": "let\n Source = Value.NativeQuery(Snowflake.Databases(\"xaa48144.snowflakecomputing.com\",\"GSL_TEST_WH\",[Role=\"ACCOUNTADMIN\"]){[Name=\"GSL_TEST_DB\"]}[Data], \"select A.name from GSL_TEST_DB.PUBLIC.SALES_ANALYST as A inner join GSL_TEST_DB.PUBLIC.SALES_FORECAST as B on A.name = B.name where startswith(A.name, 'mo')\", null, [EnableFolding=true])\nin\n Source", + "viewLanguage": "m_query" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query-with-join,DEV)", + "changeType": "UPSERT", + "aspectName": "datasetProperties", + "aspect": { + "json": { + "customProperties": { + "datasetId": "05169CD2-E713-41E6-9600-1D8066D95445" + }, + "externalUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/05169CD2-E713-41E6-9600-1D8066D95445/details", + "name": "snowflake native-query-with-join", + "description": "Library dataset description", + "tags": [] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.big-query-with-parameter,DEV)", + "changeType": "UPSERT", + "aspectName": "upstreamLineage", + "aspect": { + "json": { + "upstreams": [ + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:bigquery,my-test-project.universal.D_WH_DATE,PROD)", + "type": "TRANSFORMED" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query-with-join,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query-with-join,DEV)", + "changeType": "UPSERT", + "aspectName": "subTypes", + "aspect": { + "json": { + "typeNames": [ + "PowerBI Dataset Table", + "View" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.job-history,DEV)", + "changeType": "UPSERT", + "aspectName": "viewProperties", + "aspect": { + "json": { + "materialized": false, + "viewLogic": "let\n Source = Oracle.Database(\"localhost:1521/salesdb.GSLAB.COM\", [HierarchicalNavigation=true]), HR = Source{[Schema=\"HR\"]}[Data], EMPLOYEES1 = HR{[Name=\"EMPLOYEES\"]}[Data] \n in EMPLOYEES1", + "viewLanguage": "m_query" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query-with-join,DEV)", + "changeType": "UPSERT", + "aspectName": "upstreamLineage", + "aspect": { + "json": { + "upstreams": [ + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:snowflake,gsl_test_db.public.sales_analyst,PROD)", + "type": "TRANSFORMED" + }, + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:snowflake,gsl_test_db.public.sales_forecast,PROD)", + "type": "TRANSFORMED" + } + ], + "fineGrainedLineages": [ + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,gsl_test_db.public.sales_analyst,PROD),name)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query-with-join,DEV),name)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,gsl_test_db.public.sales_analyst,PROD),name)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query-with-join,DEV),name)" + ], + "confidenceScore": 1.0 + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.job-history,DEV)", + "changeType": "UPSERT", + "aspectName": "datasetProperties", + "aspect": { + "json": { + "customProperties": { + "datasetId": "05169CD2-E713-41E6-9600-1D8066D95445" + }, + "externalUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/05169CD2-E713-41E6-9600-1D8066D95445/details", + "name": "job-history", + "description": "Library dataset description", + "tags": [] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.job-history,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.job-history,DEV)", + "changeType": "UPSERT", + "aspectName": "subTypes", + "aspect": { + "json": { + "typeNames": [ + "PowerBI Dataset Table", + "View" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.job-history,DEV)", + "changeType": "UPSERT", + "aspectName": "upstreamLineage", + "aspect": { + "json": { + "upstreams": [ + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:oracle,salesdb.HR.EMPLOYEES,PROD)", + "type": "TRANSFORMED" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.postgres_test_table,DEV)", + "changeType": "UPSERT", + "aspectName": "viewProperties", + "aspect": { + "json": { + "materialized": false, + "viewLogic": "let\n Source = PostgreSQL.Database(\"localhost\" , \"mics\" ),\n public_order_date = Source{[Schema=\"public\",Item=\"order_date\"]}[Data] \n in \n public_order_date", + "viewLanguage": "m_query" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.postgres_test_table,DEV)", + "changeType": "UPSERT", + "aspectName": "datasetProperties", + "aspect": { + "json": { + "customProperties": { + "datasetId": "05169CD2-E713-41E6-9600-1D8066D95445" + }, + "externalUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/05169CD2-E713-41E6-9600-1D8066D95445/details", + "name": "postgres_test_table", + "description": "Library dataset description", + "tags": [] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.postgres_test_table,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.postgres_test_table,DEV)", + "changeType": "UPSERT", + "aspectName": "subTypes", + "aspect": { + "json": { + "typeNames": [ + "PowerBI Dataset Table", + "View" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.postgres_test_table,DEV)", + "changeType": "UPSERT", + "aspectName": "upstreamLineage", + "aspect": { + "json": { + "upstreams": [ + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:postgres,mics.public.order_date,PROD)", + "type": "TRANSFORMED" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.dbo_book_issue,DEV)", + "changeType": "UPSERT", + "aspectName": "viewProperties", + "aspect": { + "json": { + "materialized": false, + "viewLogic": "let\n Source = Sql.Database(\"localhost\", \"library\"),\n dbo_book_issue = Source{[Schema=\"dbo\",Item=\"book_issue\"]}[Data]\n in dbo_book_issue", + "viewLanguage": "m_query" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.dbo_book_issue,DEV)", + "changeType": "UPSERT", + "aspectName": "datasetProperties", + "aspect": { + "json": { + "customProperties": { + "datasetId": "ba0130a1-5b03-40de-9535-b34e778ea6ed" + }, + "externalUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/ba0130a1-5b03-40de-9535-b34e778ea6ed/details", + "name": "dbo_book_issue", + "description": "hr pbi test description", + "tags": [] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.dbo_book_issue,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.dbo_book_issue,DEV)", + "changeType": "UPSERT", + "aspectName": "subTypes", + "aspect": { + "json": { + "typeNames": [ + "PowerBI Dataset Table", + "View" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.ms_sql_native_table,DEV)", + "changeType": "UPSERT", + "aspectName": "viewProperties", + "aspect": { + "json": { + "materialized": false, + "viewLogic": "let\n Source = Sql.Database(\"AUPRDWHDB\", \"COMMOPSDB\", [Query=\"select *,#(lf)concat((UPPER(REPLACE(CLIENT_DIRECTOR,'-',''))), MONTH_WID) as CD_AGENT_KEY,#(lf)concat((UPPER(REPLACE(CLIENT_MANAGER_CLOSING_MONTH,'-',''))), MONTH_WID) as AGENT_KEY#(lf)#(lf)from V_PS_CD_RETENTION\", CommandTimeout=#duration(0, 1, 30, 0)]),\n #\"Changed Type\" = Table.TransformColumnTypes(Source,{{\"mth_date\", type date}}),\n #\"Added Custom\" = Table.AddColumn(#\"Changed Type\", \"Month\", each Date.Month([mth_date])),\n #\"Added Custom1\" = Table.AddColumn(#\"Added Custom\", \"TPV Opening\", each if [Month] = 1 then [TPV_AMV_OPENING]\nelse if [Month] = 2 then 0\nelse if [Month] = 3 then 0\nelse if [Month] = 4 then [TPV_AMV_OPENING]\nelse if [Month] = 5 then 0\nelse if [Month] = 6 then 0\nelse if [Month] = 7 then [TPV_AMV_OPENING]\nelse if [Month] = 8 then 0\nelse if [Month] = 9 then 0\nelse if [Month] = 10 then [TPV_AMV_OPENING]\nelse if [Month] = 11 then 0\nelse if [Month] = 12 then 0\n\nelse 0)\nin\n #\"Added Custom1\"", + "viewLanguage": "m_query" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.ms_sql_native_table,DEV)", + "changeType": "UPSERT", + "aspectName": "datasetProperties", + "aspect": { + "json": { + "customProperties": { + "datasetId": "ba0130a1-5b03-40de-9535-b34e778ea6ed" + }, + "externalUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/ba0130a1-5b03-40de-9535-b34e778ea6ed/details", + "name": "ms_sql_native_table", + "description": "hr pbi test description", + "tags": [] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.dbo_book_issue,DEV)", + "changeType": "UPSERT", + "aspectName": "upstreamLineage", + "aspect": { + "json": { + "upstreams": [ + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:mssql,library.dbo.book_issue,PROD)", + "type": "TRANSFORMED" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.ms_sql_native_table,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.ms_sql_native_table,DEV)", + "changeType": "UPSERT", + "aspectName": "subTypes", + "aspect": { + "json": { + "typeNames": [ + "PowerBI Dataset Table", + "View" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "corpuser", + "entityUrn": "urn:li:corpuser:users.User1@foo.com", + "changeType": "UPSERT", + "aspectName": "corpUserKey", + "aspect": { + "json": { + "username": "User1@foo.com" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "corpuser", + "entityUrn": "urn:li:corpuser:users.User2@foo.com", + "changeType": "UPSERT", + "aspectName": "corpUserKey", + "aspect": { + "json": { + "username": "User2@foo.com" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "chart", + "entityUrn": "urn:li:chart:(powerbi,charts.B8E293DC-0C83-4AA0-9BB9-0A8738DF24A0)", + "changeType": "UPSERT", + "aspectName": "chartInfo", + "aspect": { + "json": { + "customProperties": { + "createdFrom": "Dataset", + "datasetId": "05169CD2-E713-41E6-9600-1D8066D95445", + "datasetWebUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/05169CD2-E713-41E6-9600-1D8066D95445/details" + }, + "title": "test_tile", + "description": "test_tile", + "lastModified": { + "created": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + } + }, + "inputs": [ + { + "string": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.public_issue_history,DEV)" + }, + { + "string": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.SNOWFLAKE_TESTTABLE,DEV)" + }, + { + "string": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query,DEV)" + }, + { + "string": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.big-query-with-parameter,DEV)" + }, + { + "string": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.snowflake_native-query-with-join,DEV)" + }, + { + "string": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.job-history,DEV)" + }, + { + "string": "urn:li:dataset:(urn:li:dataPlatform:powerbi,library-dataset.postgres_test_table,DEV)" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "chart", + "entityUrn": "urn:li:chart:(powerbi,charts.B8E293DC-0C83-4AA0-9BB9-0A8738DF24A0)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "chart", + "entityUrn": "urn:li:chart:(powerbi,charts.B8E293DC-0C83-4AA0-9BB9-0A8738DF24A0)", + "changeType": "UPSERT", + "aspectName": "chartKey", + "aspect": { + "json": { + "dashboardTool": "powerbi", + "chartId": "powerbi.linkedin.com/charts/B8E293DC-0C83-4AA0-9BB9-0A8738DF24A0" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "chart", + "entityUrn": "urn:li:chart:(powerbi,charts.B8E293DC-0C83-4AA0-9BB9-0A8738DF24A0)", + "changeType": "UPSERT", + "aspectName": "browsePaths", + "aspect": { + "json": { + "paths": [ + "/powerbi/demo-workspace" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "chart", + "entityUrn": "urn:li:chart:(powerbi,charts.B8E293DC-0C83-4AA0-9BB9-0A8738DF24A0)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "demo-workspace" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "chart", + "entityUrn": "urn:li:chart:(powerbi,charts.23212598-23b5-4980-87cc-5fc0ecd84385)", + "changeType": "UPSERT", + "aspectName": "chartInfo", + "aspect": { + "json": { + "customProperties": { + "createdFrom": "Dataset", + "datasetId": "ba0130a1-5b03-40de-9535-b34e778ea6ed", + "datasetWebUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/ba0130a1-5b03-40de-9535-b34e778ea6ed/details" + }, + "title": "yearly_sales", + "description": "yearly_sales", + "lastModified": { + "created": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + } + }, + "inputs": [ + { + "string": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.dbo_book_issue,DEV)" + }, + { + "string": "urn:li:dataset:(urn:li:dataPlatform:powerbi,hr_pbi_test.ms_sql_native_table,DEV)" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "chart", + "entityUrn": "urn:li:chart:(powerbi,charts.23212598-23b5-4980-87cc-5fc0ecd84385)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "chart", + "entityUrn": "urn:li:chart:(powerbi,charts.23212598-23b5-4980-87cc-5fc0ecd84385)", + "changeType": "UPSERT", + "aspectName": "chartKey", + "aspect": { + "json": { + "dashboardTool": "powerbi", + "chartId": "powerbi.linkedin.com/charts/23212598-23b5-4980-87cc-5fc0ecd84385" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "chart", + "entityUrn": "urn:li:chart:(powerbi,charts.23212598-23b5-4980-87cc-5fc0ecd84385)", + "changeType": "UPSERT", + "aspectName": "browsePaths", + "aspect": { + "json": { + "paths": [ + "/powerbi/demo-workspace" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "chart", + "entityUrn": "urn:li:chart:(powerbi,charts.23212598-23b5-4980-87cc-5fc0ecd84385)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "demo-workspace" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dashboard", + "entityUrn": "urn:li:dashboard:(powerbi,dashboards.7D668CAD-7FFC-4505-9215-655BCA5BEBAE)", + "changeType": "UPSERT", + "aspectName": "browsePaths", + "aspect": { + "json": { + "paths": [ + "/powerbi/demo-workspace" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dashboard", + "entityUrn": "urn:li:dashboard:(powerbi,dashboards.7D668CAD-7FFC-4505-9215-655BCA5BEBAE)", + "changeType": "UPSERT", + "aspectName": "dashboardInfo", + "aspect": { + "json": { + "customProperties": { + "chartCount": "2", + "workspaceName": "demo-workspace", + "workspaceId": "64ED5CAD-7C10-4684-8180-826122881108" + }, + "title": "test_dashboard", + "description": "Description of test dashboard", + "charts": [ + "urn:li:chart:(powerbi,charts.B8E293DC-0C83-4AA0-9BB9-0A8738DF24A0)", + "urn:li:chart:(powerbi,charts.23212598-23b5-4980-87cc-5fc0ecd84385)" + ], + "datasets": [], + "lastModified": { + "created": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + } + }, + "dashboardUrl": "https://localhost/dashboards/web/1" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dashboard", + "entityUrn": "urn:li:dashboard:(powerbi,dashboards.7D668CAD-7FFC-4505-9215-655BCA5BEBAE)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dashboard", + "entityUrn": "urn:li:dashboard:(powerbi,dashboards.7D668CAD-7FFC-4505-9215-655BCA5BEBAE)", + "changeType": "UPSERT", + "aspectName": "dashboardKey", + "aspect": { + "json": { + "dashboardTool": "powerbi", + "dashboardId": "powerbi.linkedin.com/dashboards/7D668CAD-7FFC-4505-9215-655BCA5BEBAE" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dashboard", + "entityUrn": "urn:li:dashboard:(powerbi,dashboards.7D668CAD-7FFC-4505-9215-655BCA5BEBAE)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:users.User1@foo.com", + "type": "NONE" + }, + { + "owner": "urn:li:corpuser:users.User2@foo.com", + "type": "NONE" + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + } + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dashboard", + "entityUrn": "urn:li:dashboard:(powerbi,dashboards.7D668CAD-7FFC-4505-9215-655BCA5BEBAE)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "demo-workspace" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,employee-dataset.employee_ctc,DEV)", + "changeType": "UPSERT", + "aspectName": "viewProperties", + "aspect": { + "json": { + "materialized": false, + "viewLogic": "dummy", + "viewLanguage": "m_query" + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "corpuser", + "entityUrn": "urn:li:corpuser:users.User1@foo.com", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,employee-dataset.employee_ctc,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,employee-dataset.employee_ctc,DEV)", + "changeType": "UPSERT", + "aspectName": "subTypes", + "aspect": { + "json": { + "typeNames": [ + "PowerBI Dataset Table", + "View" + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:powerbi,employee-dataset.employee_ctc,DEV)", + "changeType": "UPSERT", + "aspectName": "datasetProperties", + "aspect": { + "json": { + "customProperties": { + "datasetId": "91580e0e-1680-4b1c-bbf9-4f6764d7a5ff" + }, + "externalUrl": "http://localhost/groups/64ED5CAD-7C10-4684-8180-826122881108/datasets/91580e0e-1680-4b1c-bbf9-4f6764d7a5ff/details", + "name": "employee_ctc", + "description": "Employee Management", + "tags": [] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +}, +{ + "entityType": "corpuser", + "entityUrn": "urn:li:corpuser:users.User2@foo.com", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "powerbi-test" + } +} +] \ No newline at end of file diff --git a/metadata-ingestion/tests/integration/powerbi/test_m_parser.py b/metadata-ingestion/tests/integration/powerbi/test_m_parser.py index 2fcbf5a0c0860..2e9c02ef759a5 100644 --- a/metadata-ingestion/tests/integration/powerbi/test_m_parser.py +++ b/metadata-ingestion/tests/integration/powerbi/test_m_parser.py @@ -15,8 +15,9 @@ AbstractDataPlatformInstanceResolver, create_dataplatform_instance_resolver, ) -from datahub.ingestion.source.powerbi.m_query import parser, tree_function -from datahub.ingestion.source.powerbi.m_query.resolver import DataPlatformTable +from datahub.ingestion.source.powerbi.m_query import parser, resolver, tree_function +from datahub.ingestion.source.powerbi.m_query.resolver import DataPlatformTable, Lineage +from datahub.utilities.sqlglot_lineage import ColumnLineageInfo, DownstreamColumnRef pytestmark = pytest.mark.slow @@ -70,6 +71,15 @@ def get_default_instances( return PipelineContext(run_id="fake"), config, platform_instance_resolver +def combine_upstreams_from_lineage(lineage: List[Lineage]) -> List[DataPlatformTable]: + data_platforms: List[DataPlatformTable] = [] + + for item in lineage: + data_platforms.extend(item.upstreams) + + return data_platforms + + @pytest.mark.integration def test_parse_m_query1(): expression: str = M_QUERIES[0] @@ -182,7 +192,7 @@ def test_snowflake_regular_case(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -212,7 +222,7 @@ def test_postgres_regular_case(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -242,7 +252,7 @@ def test_databricks_regular_case(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -272,7 +282,7 @@ def test_oracle_regular_case(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -302,7 +312,7 @@ def test_mssql_regular_case(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -348,7 +358,7 @@ def test_mssql_with_query(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert data_platform_tables[0].urn == expected_tables[index] @@ -388,7 +398,7 @@ def test_snowflake_native_query(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert data_platform_tables[0].urn == expected_tables[index] @@ -410,7 +420,7 @@ def test_google_bigquery_1(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -442,7 +452,7 @@ def test_google_bigquery_2(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -472,7 +482,7 @@ def test_for_each_expression_1(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -501,7 +511,7 @@ def test_for_each_expression_2(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -523,15 +533,15 @@ def test_native_query_disabled(): reporter = PowerBiDashboardSourceReport() ctx, config, platform_instance_resolver = get_default_instances() - config.native_query_parsing = False - data_platform_tables: List[DataPlatformTable] = parser.get_upstream_tables( + config.native_query_parsing = False # Disable native query parsing + lineage: List[Lineage] = parser.get_upstream_tables( table, reporter, ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, ) - assert len(data_platform_tables) == 0 + assert len(lineage) == 0 @pytest.mark.integration @@ -548,12 +558,14 @@ def test_multi_source_table(): ctx, config, platform_instance_resolver = get_default_instances() - data_platform_tables: List[DataPlatformTable] = parser.get_upstream_tables( - table, - reporter, - ctx=ctx, - config=config, - platform_instance_resolver=platform_instance_resolver, + data_platform_tables: List[DataPlatformTable] = combine_upstreams_from_lineage( + parser.get_upstream_tables( + table, + reporter, + ctx=ctx, + config=config, + platform_instance_resolver=platform_instance_resolver, + ) ) assert len(data_platform_tables) == 2 @@ -581,12 +593,14 @@ def test_table_combine(): ctx, config, platform_instance_resolver = get_default_instances() - data_platform_tables: List[DataPlatformTable] = parser.get_upstream_tables( - table, - reporter, - ctx=ctx, - config=config, - platform_instance_resolver=platform_instance_resolver, + data_platform_tables: List[DataPlatformTable] = combine_upstreams_from_lineage( + parser.get_upstream_tables( + table, + reporter, + ctx=ctx, + config=config, + platform_instance_resolver=platform_instance_resolver, + ) ) assert len(data_platform_tables) == 2 @@ -624,7 +638,7 @@ def test_expression_is_none(): ctx, config, platform_instance_resolver = get_default_instances() - data_platform_tables: List[DataPlatformTable] = parser.get_upstream_tables( + lineage: List[Lineage] = parser.get_upstream_tables( table, reporter, ctx=ctx, @@ -632,7 +646,7 @@ def test_expression_is_none(): platform_instance_resolver=platform_instance_resolver, ) - assert len(data_platform_tables) == 0 + assert len(lineage) == 0 def test_redshift_regular_case(): @@ -651,7 +665,7 @@ def test_redshift_regular_case(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -678,7 +692,7 @@ def test_redshift_native_query(): ctx=ctx, config=config, platform_instance_resolver=platform_instance_resolver, - ) + )[0].upstreams assert len(data_platform_tables) == 1 assert ( @@ -708,7 +722,7 @@ def test_sqlglot_parser(): } ) - data_platform_tables: List[DataPlatformTable] = parser.get_upstream_tables( + lineage: List[resolver.Lineage] = parser.get_upstream_tables( table, reporter, ctx=ctx, @@ -716,6 +730,8 @@ def test_sqlglot_parser(): platform_instance_resolver=platform_instance_resolver, ) + data_platform_tables: List[DataPlatformTable] = lineage[0].upstreams + assert len(data_platform_tables) == 2 assert ( data_platform_tables[0].urn @@ -725,3 +741,76 @@ def test_sqlglot_parser(): data_platform_tables[1].urn == "urn:li:dataset:(urn:li:dataPlatform:snowflake,sales_deployment.operations_analytics.transformed_prod.v_sme_unit_targets,PROD)" ) + + assert lineage[0].column_lineage == [ + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="client_director"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="tier"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column='upper("manager")'), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="team_type"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="date_target"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="monthid"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="target_team"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="seller_email"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="agent_key"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="sme_quota"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="revenue_quota"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="service_quota"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="bl_target"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="software_quota"), + upstreams=[], + logic=None, + ), + ] diff --git a/metadata-ingestion/tests/integration/powerbi/test_powerbi.py b/metadata-ingestion/tests/integration/powerbi/test_powerbi.py index 044532021a19c..b0695e3ea9954 100644 --- a/metadata-ingestion/tests/integration/powerbi/test_powerbi.py +++ b/metadata-ingestion/tests/integration/powerbi/test_powerbi.py @@ -1,4 +1,5 @@ import logging +import re import sys from typing import Any, Dict, List, cast from unittest import mock @@ -1127,7 +1128,7 @@ def test_dataset_type_mapping_error( """ register_mock_api(request_mock=requests_mock) - try: + with pytest.raises(Exception, match=r"dataset_type_mapping is deprecated"): Pipeline.create( { "run_id": "powerbi-test", @@ -1150,11 +1151,6 @@ def test_dataset_type_mapping_error( }, } ) - except Exception as e: - assert ( - "dataset_type_mapping is deprecated. Use server_to_platform_instance only." - in str(e) - ) @freeze_time(FROZEN_TIME) @@ -1506,3 +1502,90 @@ def test_independent_datasets_extraction( output_path=tmp_path / "powerbi_independent_mces.json", golden_path=f"{test_resources_dir}/{golden_file}", ) + + +@freeze_time(FROZEN_TIME) +@mock.patch("msal.ConfidentialClientApplication", side_effect=mock_msal_cca) +def test_cll_extraction(mock_msal, pytestconfig, tmp_path, mock_time, requests_mock): + + test_resources_dir = pytestconfig.rootpath / "tests/integration/powerbi" + + register_mock_api( + request_mock=requests_mock, + ) + + default_conf: dict = default_source_config() + + del default_conf[ + "dataset_type_mapping" + ] # delete this key so that connector set it to default (all dataplatform) + + pipeline = Pipeline.create( + { + "run_id": "powerbi-test", + "source": { + "type": "powerbi", + "config": { + **default_conf, + "extract_lineage": True, + "extract_column_level_lineage": True, + "enable_advance_lineage_sql_construct": True, + "native_query_parsing": True, + "extract_independent_datasets": True, + }, + }, + "sink": { + "type": "file", + "config": { + "filename": f"{tmp_path}/powerbi_cll_mces.json", + }, + }, + } + ) + + pipeline.run() + pipeline.raise_from_status() + golden_file = "golden_test_cll.json" + + mce_helpers.check_golden_file( + pytestconfig, + output_path=tmp_path / "powerbi_cll_mces.json", + golden_path=f"{test_resources_dir}/{golden_file}", + ) + + +@freeze_time(FROZEN_TIME) +@mock.patch("msal.ConfidentialClientApplication", side_effect=mock_msal_cca) +def test_cll_extraction_flags( + mock_msal, pytestconfig, tmp_path, mock_time, requests_mock +): + + register_mock_api( + request_mock=requests_mock, + ) + + default_conf: dict = default_source_config() + pattern: str = re.escape( + "Enable all these flags in recipe: ['native_query_parsing', 'enable_advance_lineage_sql_construct', 'extract_lineage']" + ) + + with pytest.raises(Exception, match=pattern): + + Pipeline.create( + { + "run_id": "powerbi-test", + "source": { + "type": "powerbi", + "config": { + **default_conf, + "extract_column_level_lineage": True, + }, + }, + "sink": { + "type": "file", + "config": { + "filename": f"{tmp_path}/powerbi_cll_mces.json", + }, + }, + } + ) From a300b39f15cd689b42b7c32ce9e5087ccf5a356e Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Wed, 4 Oct 2023 06:53:15 -0400 Subject: [PATCH 082/156] feat(ingest/airflow): airflow plugin v2 (#8853) --- .github/workflows/airflow-plugin.yml | 25 +- .github/workflows/build-and-test.yml | 9 +- docker/airflow/local_airflow.md | 2 +- docs/how/updating-datahub.md | 3 + docs/lineage/airflow.md | 251 ++- .../airflow-plugin/build.gradle | 23 +- .../airflow-plugin/pyproject.toml | 1 + .../airflow-plugin/setup.cfg | 30 +- .../airflow-plugin/setup.py | 71 +- .../datahub_airflow_plugin/_airflow_shims.py | 32 + .../src/datahub_airflow_plugin/_config.py | 80 + .../_datahub_listener_module.py | 7 + .../_datahub_ol_adapter.py | 23 + .../src/datahub_airflow_plugin/_extractors.py | 244 ++ .../client/airflow_generator.py | 69 +- .../datahub_listener.py | 494 +++++ .../datahub_airflow_plugin/datahub_plugin.py | 391 +--- .../datahub_plugin_v22.py | 336 +++ .../example_dags/lineage_emission_dag.py | 22 +- .../datahub_airflow_plugin/hooks/datahub.py | 115 +- .../{ => lineage}/_lineage_core.py | 30 +- .../datahub_airflow_plugin/lineage/datahub.py | 28 +- .../operators/datahub.py | 4 +- .../airflow-plugin/tests/conftest.py | 6 + .../tests/integration/dags/basic_iolets.py | 34 + .../tests/integration/dags/simple_dag.py | 34 + .../integration/dags/snowflake_operator.py | 32 + .../tests/integration/dags/sqlite_operator.py | 75 + .../integration/goldens/v1_basic_iolets.json | 533 +++++ .../integration/goldens/v1_simple_dag.json | 718 ++++++ .../integration/goldens/v2_basic_iolets.json | 535 +++++ .../v2_basic_iolets_no_dag_listener.json | 535 +++++ .../integration/goldens/v2_simple_dag.json | 666 ++++++ .../v2_simple_dag_no_dag_listener.json | 722 ++++++ .../goldens/v2_snowflake_operator.json | 507 +++++ .../goldens/v2_sqlite_operator.json | 1735 +++++++++++++++ .../v2_sqlite_operator_no_dag_listener.json | 1955 +++++++++++++++++ .../integration/integration_test_dummy.py | 2 - .../tests/integration/test_plugin.py | 392 ++++ .../airflow-plugin/tests/unit/test_airflow.py | 25 +- .../airflow-plugin/tests/unit/test_dummy.py | 2 - .../tests/unit/test_packaging.py | 8 + .../airflow-plugin/tox.ini | 39 +- metadata-ingestion/setup.py | 20 +- .../api/entities/corpgroup/corpgroup.py | 33 +- .../datahub/api/entities/corpuser/corpuser.py | 9 +- .../datahub/api/entities/datajob/dataflow.py | 19 +- .../datahub/api/entities/datajob/datajob.py | 39 +- .../dataprocess/dataprocess_instance.py | 21 +- .../api/entities/dataproduct/dataproduct.py | 22 +- .../src/datahub/emitter/generic_emitter.py | 31 + .../src/datahub/emitter/kafka_emitter.py | 3 +- .../src/datahub/emitter/rest_emitter.py | 16 +- .../emitter/synchronized_file_emitter.py | 60 + .../src/datahub/ingestion/graph/client.py | 17 + .../datahub/ingestion/source/kafka_connect.py | 4 +- .../ingestion/source/sql/sql_common.py | 48 - .../source/sql/sqlalchemy_uri_mapper.py | 47 + .../src/datahub/ingestion/source/superset.py | 6 +- .../src/datahub/ingestion/source/tableau.py | 11 +- .../integrations/great_expectations/action.py | 4 +- .../datahub/testing/compare_metadata_json.py | 22 +- .../src/datahub/utilities/sqlglot_lineage.py | 40 +- .../goldens/test_create_table_ddl.json | 8 + .../unit/sql_parsing/test_sqlglot_lineage.py | 15 + .../tests/unit/test_sql_common.py | 7 +- 66 files changed, 10457 insertions(+), 890 deletions(-) create mode 100644 metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_config.py create mode 100644 metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_listener_module.py create mode 100644 metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_ol_adapter.py create mode 100644 metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py create mode 100644 metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py create mode 100644 metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py rename metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/{ => lineage}/_lineage_core.py (72%) create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/conftest.py create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/dags/basic_iolets.py create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/dags/sqlite_operator.py create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_basic_iolets.json create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_simple_dag.json create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets.json create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets_no_dag_listener.json create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag.json create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag_no_dag_listener.json create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_snowflake_operator.json create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator_no_dag_listener.json delete mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/integration_test_dummy.py create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py delete mode 100644 metadata-ingestion-modules/airflow-plugin/tests/unit/test_dummy.py create mode 100644 metadata-ingestion-modules/airflow-plugin/tests/unit/test_packaging.py create mode 100644 metadata-ingestion/src/datahub/emitter/generic_emitter.py create mode 100644 metadata-ingestion/src/datahub/emitter/synchronized_file_emitter.py create mode 100644 metadata-ingestion/src/datahub/ingestion/source/sql/sqlalchemy_uri_mapper.py create mode 100644 metadata-ingestion/tests/unit/sql_parsing/goldens/test_create_table_ddl.json diff --git a/.github/workflows/airflow-plugin.yml b/.github/workflows/airflow-plugin.yml index 63bab821cc398..a250bddcc16d1 100644 --- a/.github/workflows/airflow-plugin.yml +++ b/.github/workflows/airflow-plugin.yml @@ -32,16 +32,21 @@ jobs: strategy: matrix: include: - - python-version: "3.7" - extraPythonRequirement: "apache-airflow~=2.1.0" - - python-version: "3.7" - extraPythonRequirement: "apache-airflow~=2.2.0" + - python-version: "3.8" + extra_pip_requirements: "apache-airflow~=2.1.4" + extra_pip_extras: plugin-v1 + - python-version: "3.8" + extra_pip_requirements: "apache-airflow~=2.2.4" + extra_pip_extras: plugin-v1 - python-version: "3.10" - extraPythonRequirement: "apache-airflow~=2.4.0" + extra_pip_requirements: "apache-airflow~=2.4.0" + extra_pip_extras: plugin-v2 - python-version: "3.10" - extraPythonRequirement: "apache-airflow~=2.6.0" + extra_pip_requirements: "apache-airflow~=2.6.0" + extra_pip_extras: plugin-v2 - python-version: "3.10" - extraPythonRequirement: "apache-airflow>2.6.0" + extra_pip_requirements: "apache-airflow>=2.7.0" + extra_pip_extras: plugin-v2 fail-fast: false steps: - uses: actions/checkout@v3 @@ -51,13 +56,13 @@ jobs: cache: "pip" - name: Install dependencies run: ./metadata-ingestion/scripts/install_deps.sh - - name: Install airflow package and test (extras ${{ matrix.extraPythonRequirement }}) - run: ./gradlew -Pextra_pip_requirements='${{ matrix.extraPythonRequirement }}' :metadata-ingestion-modules:airflow-plugin:lint :metadata-ingestion-modules:airflow-plugin:testQuick + - name: Install airflow package and test (extras ${{ matrix.extra_pip_requirements }}) + run: ./gradlew -Pextra_pip_requirements='${{ matrix.extra_pip_requirements }}' -Pextra_pip_extras='${{ matrix.extra_pip_extras }}' :metadata-ingestion-modules:airflow-plugin:lint :metadata-ingestion-modules:airflow-plugin:testQuick - name: pip freeze show list installed if: always() run: source metadata-ingestion-modules/airflow-plugin/venv/bin/activate && pip freeze - uses: actions/upload-artifact@v3 - if: ${{ always() && matrix.python-version == '3.10' && matrix.extraPythonRequirement == 'apache-airflow>2.6.0' }} + if: ${{ always() && matrix.python-version == '3.10' && matrix.extra_pip_requirements == 'apache-airflow>=2.7.0' }} with: name: Test Results (Airflow Plugin ${{ matrix.python-version}}) path: | diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml index f6320e1bd5c9f..3f409878b191f 100644 --- a/.github/workflows/build-and-test.yml +++ b/.github/workflows/build-and-test.yml @@ -26,9 +26,9 @@ jobs: matrix: command: [ - "./gradlew build -x :metadata-ingestion:build -x :metadata-ingestion:check -x docs-website:build -x :metadata-integration:java:spark-lineage:test -x :metadata-io:test -x :metadata-ingestion-modules:airflow-plugin:build -x :datahub-frontend:build -x :datahub-web-react:build --parallel", + # metadata-ingestion and airflow-plugin each have dedicated build jobs + "./gradlew build -x :metadata-ingestion:build -x :metadata-ingestion:check -x docs-website:build -x :metadata-integration:java:spark-lineage:test -x :metadata-io:test -x :metadata-ingestion-modules:airflow-plugin:build -x :metadata-ingestion-modules:airflow-plugin:check -x :datahub-frontend:build -x :datahub-web-react:build --parallel", "./gradlew :datahub-frontend:build :datahub-web-react:build --parallel", - "./gradlew :metadata-ingestion-modules:airflow-plugin:build --parallel" ] timezone: [ @@ -51,7 +51,8 @@ jobs: java-version: 11 - uses: actions/setup-python@v4 with: - python-version: "3.7" + python-version: "3.10" + cache: pip - name: Gradle build (and test) run: | ${{ matrix.command }} @@ -81,7 +82,7 @@ jobs: - uses: actions/checkout@v3 - uses: actions/setup-python@v4 with: - python-version: "3.7" + python-version: "3.10" - name: Download YQ uses: chrisdickinson/setup-yq@v1.0.1 with: diff --git a/docker/airflow/local_airflow.md b/docker/airflow/local_airflow.md index 55a64f5c122c5..fbfc1d17327c5 100644 --- a/docker/airflow/local_airflow.md +++ b/docker/airflow/local_airflow.md @@ -1,6 +1,6 @@ :::caution -This feature is currently unmaintained. As of 0.10.0 the container described is not published alongside the DataHub CLI. If you'd like to use it, please reach out to us on the [community slack.](docs/slack.md) +This guide is currently unmaintained. As of 0.10.0 the container described is not published alongside the DataHub CLI. If you'd like to use it, please reach out to us on the [community slack.](docs/slack.md) ::: diff --git a/docs/how/updating-datahub.md b/docs/how/updating-datahub.md index 9b19291ee246a..4df8d435cf1c4 100644 --- a/docs/how/updating-datahub.md +++ b/docs/how/updating-datahub.md @@ -5,7 +5,10 @@ This file documents any backwards-incompatible changes in DataHub and assists pe ## Next ### Breaking Changes + - #8810 - Removed support for SQLAlchemy 1.3.x. Only SQLAlchemy 1.4.x is supported now. +- #8853 - The Airflow plugin no longer supports Airflow 2.0.x or Python 3.7. See the docs for more details. +- #8853 - Introduced the Airflow plugin v2. If you're using Airflow 2.3+, the v2 plugin will be enabled by default, and so you'll need to switch your requirements to include `pip install 'acryl-datahub-airflow-plugin[plugin-v2]'`. To continue using the v1 plugin, set the `DATAHUB_AIRFLOW_PLUGIN_USE_V1_PLUGIN` environment variable to `true`. ### Potential Downtime diff --git a/docs/lineage/airflow.md b/docs/lineage/airflow.md index 49de5352f6d58..19ed1598d4c5a 100644 --- a/docs/lineage/airflow.md +++ b/docs/lineage/airflow.md @@ -1,74 +1,137 @@ # Airflow Integration -DataHub supports integration of +:::note -- Airflow Pipeline (DAG) metadata -- DAG and Task run information as well as -- Lineage information when present +If you're looking to schedule DataHub ingestion using Airflow, see the guide on [scheduling ingestion with Airflow](../../metadata-ingestion/schedule_docs/airflow.md). -You can use either the DataHub Airflow lineage plugin (recommended) or the Airflow lineage backend (deprecated). +::: -## Using Datahub's Airflow lineage plugin +The DataHub Airflow plugin supports: -:::note +- Automatic column-level lineage extraction from various operators e.g. `SqlOperator`s (including `MySqlOperator`, `PostgresOperator`, `SnowflakeOperator`, and more), `S3FileTransformOperator`, and a few others. +- Airflow DAG and tasks, including properties, ownership, and tags. +- Task run information, including task successes and failures. +- Manual lineage annotations using `inlets` and `outlets` on Airflow operators. -The Airflow lineage plugin is only supported with Airflow version >= 2.0.2 or on MWAA with an Airflow version >= 2.0.2. +There's two actively supported implementations of the plugin, with different Airflow version support. -If you're using Airflow 1.x, use the Airflow lineage plugin with acryl-datahub-airflow-plugin <= 0.9.1.0. +| Approach | Airflow Version | Notes | +| --------- | --------------- | --------------------------------------------------------------------------- | +| Plugin v2 | 2.3+ | Recommended. Requires Python 3.8+ | +| Plugin v1 | 2.1+ | No automatic lineage extraction; may not extract lineage if the task fails. | -::: +If you're using Airflow older than 2.1, it's possible to use the v1 plugin with older versions of `acryl-datahub-airflow-plugin`. See the [compatibility section](#compatibility) for more details. -This plugin registers a task success/failure callback on every task with a cluster policy and emits DataHub events from that. This allows this plugin to be able to register both task success as well as failures compared to the older Airflow Lineage Backend which could only support emitting task success. + + -### Setup +## DataHub Plugin v2 -1. You need to install the required dependency in your airflow. +### Installation + +The v2 plugin requires Airflow 2.3+ and Python 3.8+. If you don't meet these requirements, use the v1 plugin instead. ```shell -pip install acryl-datahub-airflow-plugin +pip install 'acryl-datahub-airflow-plugin[plugin-v2]' ``` -:::note +### Configuration -The [DataHub Rest](../../metadata-ingestion/sink_docs/datahub.md#datahub-rest) emitter is included in the plugin package by default. To use [DataHub Kafka](../../metadata-ingestion/sink_docs/datahub.md#datahub-kafka) install `pip install acryl-datahub-airflow-plugin[datahub-kafka]`. +Set up a DataHub connection in Airflow. -::: +```shell +airflow connections add --conn-type 'datahub-rest' 'datahub_rest_default' --conn-host 'http://datahub-gms:8080' --conn-password '' +``` + +No additional configuration is required to use the plugin. However, there are some optional configuration parameters that can be set in the `airflow.cfg` file. + +```ini title="airflow.cfg" +[datahub] +# Optional - additional config here. +enabled = True # default +``` + +| Name | Default value | Description | +| -------------------------- | -------------------- | ---------------------------------------------------------------------------------------- | +| enabled | true | If the plugin should be enabled. | +| conn_id | datahub_rest_default | The name of the datahub rest connection. | +| cluster | prod | name of the airflow cluster | +| capture_ownership_info | true | Extract DAG ownership. | +| capture_tags_info | true | Extract DAG tags. | +| capture_executions | true | Extract task runs and success/failure statuses. This will show up in DataHub "Runs" tab. | +| enable_extractors | true | Enable automatic lineage extraction. | +| disable_openlineage_plugin | true | Disable the OpenLineage plugin to avoid duplicative processing. | +| log_level | _no change_ | [debug] Set the log level for the plugin. | +| debug_emitter | false | [debug] If true, the plugin will log the emitted events. | + +### Automatic lineage extraction + +To automatically extract lineage information, the v2 plugin builds on top of Airflow's built-in [OpenLineage extractors](https://openlineage.io/docs/integrations/airflow/default-extractors). -2. Disable lazy plugin loading in your airflow.cfg. - On MWAA you should add this config to your [Apache Airflow configuration options](https://docs.aws.amazon.com/mwaa/latest/userguide/configuring-env-variables.html#configuring-2.0-airflow-override). +The SQL-related extractors have been updated to use DataHub's SQL parser, which is more robust than the built-in one and uses DataHub's metadata information to generate column-level lineage. We discussed the DataHub SQL parser, including why schema-aware parsing works better and how it performs on benchmarks, during the [June 2023 community town hall](https://youtu.be/1QVcUmRQK5E?si=U27zygR7Gi_KdkzE&t=2309). + +## DataHub Plugin v1 + +### Installation + +The v1 plugin requires Airflow 2.1+ and Python 3.8+. If you're on older versions, it's still possible to use an older version of the plugin. See the [compatibility section](#compatibility) for more details. + +If you're using Airflow 2.3+, we recommend using the v2 plugin instead. If you need to use the v1 plugin with Airflow 2.3+, you must also set the environment variable `DATAHUB_AIRFLOW_PLUGIN_USE_V1_PLUGIN=true`. + +```shell +pip install 'acryl-datahub-airflow-plugin[plugin-v1]' + +# The DataHub rest connection type is included by default. +# To use the DataHub Kafka connection type, install the plugin with the kafka extras. +pip install 'acryl-datahub-airflow-plugin[plugin-v1,datahub-kafka]' +``` + + + +### Configuration + +#### Disable lazy plugin loading ```ini title="airflow.cfg" [core] lazy_load_plugins = False ``` -3. You must configure an Airflow hook for Datahub. We support both a Datahub REST hook and a Kafka-based hook, but you only need one. +On MWAA you should add this config to your [Apache Airflow configuration options](https://docs.aws.amazon.com/mwaa/latest/userguide/configuring-env-variables.html#configuring-2.0-airflow-override). + +#### Setup a DataHub connection - ```shell - # For REST-based: - airflow connections add --conn-type 'datahub_rest' 'datahub_rest_default' --conn-host 'http://datahub-gms:8080' --conn-password '' - # For Kafka-based (standard Kafka sink config can be passed via extras): - airflow connections add --conn-type 'datahub_kafka' 'datahub_kafka_default' --conn-host 'broker:9092' --conn-extra '{}' - ``` +You must configure an Airflow connection for Datahub. We support both a Datahub REST and a Kafka-based connections, but you only need one. -4. Add your `datahub_conn_id` and/or `cluster` to your `airflow.cfg` file if it is not align with the default values. See configuration parameters below +```shell +# For REST-based: +airflow connections add --conn-type 'datahub_rest' 'datahub_rest_default' --conn-host 'http://datahub-gms:8080' --conn-password '' +# For Kafka-based (standard Kafka sink config can be passed via extras): +airflow connections add --conn-type 'datahub_kafka' 'datahub_kafka_default' --conn-host 'broker:9092' --conn-extra '{}' +``` - **Configuration options:** +#### Configure the plugin - | Name | Default value | Description | - | ------------------------------ | -------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | - | datahub.enabled | true | If the plugin should be enabled. | - | datahub.conn_id | datahub_rest_default | The name of the datahub connection you set in step 1. | - | datahub.cluster | prod | name of the airflow cluster | - | datahub.capture_ownership_info | true | If true, the owners field of the DAG will be capture as a DataHub corpuser. | - | datahub.capture_tags_info | true | If true, the tags field of the DAG will be captured as DataHub tags. | - | datahub.capture_executions | true | If true, we'll capture task runs in DataHub in addition to DAG definitions. | - | datahub.graceful_exceptions | true | If set to true, most runtime errors in the lineage backend will be suppressed and will not cause the overall task to fail. Note that configuration issues will still throw exceptions. | +If your config doesn't align with the default values, you can configure the plugin in your `airflow.cfg` file. + +```ini title="airflow.cfg" +[datahub] +enabled = true +conn_id = datahub_rest_default # or datahub_kafka_default +# etc. +``` -5. Configure `inlets` and `outlets` for your Airflow operators. For reference, look at the sample DAG in [`lineage_backend_demo.py`](../../metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_backend_demo.py), or reference [`lineage_backend_taskflow_demo.py`](../../metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_backend_taskflow_demo.py) if you're using the [TaskFlow API](https://airflow.apache.org/docs/apache-airflow/stable/concepts/taskflow.html). -6. [optional] Learn more about [Airflow lineage](https://airflow.apache.org/docs/apache-airflow/stable/lineage.html), including shorthand notation and some automation. +| Name | Default value | Description | +| ---------------------- | -------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| enabled | true | If the plugin should be enabled. | +| conn_id | datahub_rest_default | The name of the datahub connection you set in step 1. | +| cluster | prod | name of the airflow cluster | +| capture_ownership_info | true | If true, the owners field of the DAG will be capture as a DataHub corpuser. | +| capture_tags_info | true | If true, the tags field of the DAG will be captured as DataHub tags. | +| capture_executions | true | If true, we'll capture task runs in DataHub in addition to DAG definitions. | +| graceful_exceptions | true | If set to true, most runtime errors in the lineage backend will be suppressed and will not cause the overall task to fail. Note that configuration issues will still throw exceptions. | -### How to validate installation +#### Validate that the plugin is working 1. Go and check in Airflow at Admin -> Plugins menu if you can see the DataHub plugin 2. Run an Airflow DAG. In the task logs, you should see Datahub related log messages like: @@ -77,9 +140,22 @@ lazy_load_plugins = False Emitting DataHub ... ``` -### Emitting lineage via a custom operator to the Airflow Plugin +## Manual Lineage Annotation + +### Using `inlets` and `outlets` + +You can manually annotate lineage by setting `inlets` and `outlets` on your Airflow operators. This is useful if you're using an operator that doesn't support automatic lineage extraction, or if you want to override the automatic lineage extraction. + +We have a few code samples that demonstrate how to use `inlets` and `outlets`: -If you have created a custom Airflow operator [docs](https://airflow.apache.org/docs/apache-airflow/stable/howto/custom-operator.html) that inherits from the BaseOperator class, +- [`lineage_backend_demo.py`](../../metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_backend_demo.py) +- [`lineage_backend_taskflow_demo.py`](../../metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_backend_taskflow_demo.py) - uses the [TaskFlow API](https://airflow.apache.org/docs/apache-airflow/stable/concepts/taskflow.html) + +For more information, take a look at the [Airflow lineage docs](https://airflow.apache.org/docs/apache-airflow/stable/lineage.html). + +### Custom Operators + +If you have created a [custom Airflow operator](https://airflow.apache.org/docs/apache-airflow/stable/howto/custom-operator.html) that inherits from the BaseOperator class, when overriding the `execute` function, set inlets and outlets via `context['ti'].task.inlets` and `context['ti'].task.outlets`. The DataHub Airflow plugin will then pick up those inlets and outlets after the task runs. @@ -90,7 +166,7 @@ class DbtOperator(BaseOperator): def execute(self, context): # do something inlets, outlets = self._get_lineage() - # inlets/outlets are lists of either datahub_provider.entities.Dataset or datahub_provider.entities.Urn + # inlets/outlets are lists of either datahub_airflow_plugin.entities.Dataset or datahub_airflow_plugin.entities.Urn context['ti'].task.inlets = self.inlets context['ti'].task.outlets = self.outlets @@ -100,78 +176,25 @@ class DbtOperator(BaseOperator): return inlets, outlets ``` -If you override the `pre_execute` and `post_execute` function, ensure they include the `@prepare_lineage` and `@apply_lineage` decorators respectively. [source](https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/lineage.html#lineage) - -## Using DataHub's Airflow lineage backend (deprecated) - -:::caution - -The DataHub Airflow plugin (above) is the recommended way to integrate Airflow with DataHub. For managed services like MWAA, the lineage backend is not supported and so you must use the Airflow plugin. - -If you're using Airflow 1.x, we recommend using the Airflow lineage backend with acryl-datahub <= 0.9.1.0. - -::: - -:::note - -If you are looking to run Airflow and DataHub using docker locally, follow the guide [here](../../docker/airflow/local_airflow.md). Otherwise proceed to follow the instructions below. -::: - -### Setting up Airflow to use DataHub as Lineage Backend - -1. You need to install the required dependency in your airflow. See - -```shell -pip install acryl-datahub[airflow] -# If you need the Kafka-based emitter/hook: -pip install acryl-datahub[airflow,datahub-kafka] -``` - -2. You must configure an Airflow hook for Datahub. We support both a Datahub REST hook and a Kafka-based hook, but you only need one. - - ```shell - # For REST-based: - airflow connections add --conn-type 'datahub_rest' 'datahub_rest_default' --conn-host 'http://datahub-gms:8080' --conn-password '' - # For Kafka-based (standard Kafka sink config can be passed via extras): - airflow connections add --conn-type 'datahub_kafka' 'datahub_kafka_default' --conn-host 'broker:9092' --conn-extra '{}' - ``` +If you override the `pre_execute` and `post_execute` function, ensure they include the `@prepare_lineage` and `@apply_lineage` decorators respectively. Reference the [Airflow docs](https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/lineage.html#lineage) for more details. -3. Add the following lines to your `airflow.cfg` file. +## Emit Lineage Directly - ```ini title="airflow.cfg" - [lineage] - backend = datahub_provider.lineage.datahub.DatahubLineageBackend - datahub_kwargs = { - "enabled": true, - "datahub_conn_id": "datahub_rest_default", - "cluster": "prod", - "capture_ownership_info": true, - "capture_tags_info": true, - "graceful_exceptions": true } - # The above indentation is important! - ``` +If you can't use the plugin or annotate inlets/outlets, you can also emit lineage using the `DatahubEmitterOperator`. - **Configuration options:** +Reference [`lineage_emission_dag.py`](../../metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_emission_dag.py) for a full example. - - `datahub_conn_id` (required): Usually `datahub_rest_default` or `datahub_kafka_default`, depending on what you named the connection in step 1. - - `cluster` (defaults to "prod"): The "cluster" to associate Airflow DAGs and tasks with. - - `capture_ownership_info` (defaults to true): If true, the owners field of the DAG will be capture as a DataHub corpuser. - - `capture_tags_info` (defaults to true): If true, the tags field of the DAG will be captured as DataHub tags. - - `capture_executions` (defaults to false): If true, it captures task runs as DataHub DataProcessInstances. - - `graceful_exceptions` (defaults to true): If set to true, most runtime errors in the lineage backend will be suppressed and will not cause the overall task to fail. Note that configuration issues will still throw exceptions. +In order to use this example, you must first configure the Datahub hook. Like in ingestion, we support a Datahub REST hook and a Kafka-based hook. See the plugin configuration for examples. -4. Configure `inlets` and `outlets` for your Airflow operators. For reference, look at the sample DAG in [`lineage_backend_demo.py`](../../metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_backend_demo.py), or reference [`lineage_backend_taskflow_demo.py`](../../metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_backend_taskflow_demo.py) if you're using the [TaskFlow API](https://airflow.apache.org/docs/apache-airflow/stable/concepts/taskflow.html). -5. [optional] Learn more about [Airflow lineage](https://airflow.apache.org/docs/apache-airflow/stable/lineage.html), including shorthand notation and some automation. - -## Emitting lineage via a separate operator - -Take a look at this sample DAG: +## Debugging -- [`lineage_emission_dag.py`](../../metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_emission_dag.py) - emits lineage using the DatahubEmitterOperator. +### Missing lineage -In order to use this example, you must first configure the Datahub hook. Like in ingestion, we support a Datahub REST hook and a Kafka-based hook. See step 1 above for details. +If you're not seeing lineage in DataHub, check the following: -## Debugging +- Validate that the plugin is loaded in Airflow. Go to Admin -> Plugins and check that the DataHub plugin is listed. +- If using the v2 plugin's automatic lineage, ensure that the `enable_extractors` config is set to true and that automatic lineage is supported for your operator. +- If using manual lineage annotation, ensure that you're using the `datahub_airflow_plugin.entities.Dataset` or `datahub_airflow_plugin.entities.Urn` classes for your inlets and outlets. ### Incorrect URLs @@ -179,9 +202,21 @@ If your URLs aren't being generated correctly (usually they'll start with `http: ```ini title="airflow.cfg" [webserver] -base_url = http://airflow.example.com +base_url = http://airflow.mycorp.example.com ``` +## Compatibility + +We no longer officially support Airflow <2.1. However, you can use older versions of `acryl-datahub-airflow-plugin` with older versions of Airflow. +Both of these options support Python 3.7+. + +- Airflow 1.10.x, use DataHub plugin v1 with acryl-datahub-airflow-plugin <= 0.9.1.0. +- Airflow 2.0.x, use DataHub plugin v1 with acryl-datahub-airflow-plugin <= 0.11.0.1. + +DataHub also previously supported an Airflow [lineage backend](https://airflow.apache.org/docs/apache-airflow/2.2.0/lineage.html#lineage-backend) implementation. While the implementation is still in our codebase, it is deprecated and will be removed in a future release. +Note that the lineage backend did not support automatic lineage extraction, did not capture task failures, and did not work in AWS MWAA. +The [documentation for the lineage backend](https://docs-website-1wmaehubl-acryldata.vercel.app/docs/lineage/airflow/#using-datahubs-airflow-lineage-backend-deprecated) has already been archived. + ## Additional references Related Datahub videos: diff --git a/metadata-ingestion-modules/airflow-plugin/build.gradle b/metadata-ingestion-modules/airflow-plugin/build.gradle index 58a2bc9e670e3..dacf12dc020df 100644 --- a/metadata-ingestion-modules/airflow-plugin/build.gradle +++ b/metadata-ingestion-modules/airflow-plugin/build.gradle @@ -10,6 +10,13 @@ ext { if (!project.hasProperty("extra_pip_requirements")) { ext.extra_pip_requirements = "" } +if (!project.hasProperty("extra_pip_extras")) { + ext.extra_pip_extras = "plugin-v2" +} +// If extra_pip_extras is non-empty, we need to add a comma to the beginning of the string. +if (extra_pip_extras != "") { + ext.extra_pip_extras = "," + extra_pip_extras +} def pip_install_command = "${venv_name}/bin/pip install -e ../../metadata-ingestion" @@ -36,7 +43,7 @@ task installPackage(type: Exec, dependsOn: [environmentSetup, ':metadata-ingesti // and https://github.com/datahub-project/datahub/pull/8435. commandLine 'bash', '-x', '-c', "${pip_install_command} install 'Cython<3.0' 'PyYAML<6' --no-build-isolation && " + - "${pip_install_command} -e . ${extra_pip_requirements} &&" + + "${pip_install_command} -e .[ignore${extra_pip_extras}] ${extra_pip_requirements} &&" + "touch ${sentinel_file}" } @@ -47,7 +54,7 @@ task installDev(type: Exec, dependsOn: [install]) { inputs.file file('setup.py') outputs.file("${sentinel_file}") commandLine 'bash', '-x', '-c', - "${pip_install_command} -e .[dev] ${extra_pip_requirements} && " + + "${pip_install_command} -e .[dev${extra_pip_extras}] ${extra_pip_requirements} && " + "touch ${sentinel_file}" } @@ -79,7 +86,8 @@ task installDevTest(type: Exec, dependsOn: [installDev]) { outputs.dir("${venv_name}") outputs.file("${sentinel_file}") commandLine 'bash', '-x', '-c', - "${pip_install_command} -e .[dev,integration-tests] && touch ${sentinel_file}" + "${pip_install_command} -e .[dev,integration-tests${extra_pip_extras}] ${extra_pip_requirements} && " + + "touch ${sentinel_file}" } def testFile = hasProperty('testFile') ? testFile : 'unknown' @@ -97,20 +105,13 @@ task testSingle(dependsOn: [installDevTest]) { } task testQuick(type: Exec, dependsOn: installDevTest) { - // We can't enforce the coverage requirements if we run a subset of the tests. inputs.files(project.fileTree(dir: "src/", include: "**/*.py")) inputs.files(project.fileTree(dir: "tests/")) - outputs.dir("${venv_name}") commandLine 'bash', '-x', '-c', - "source ${venv_name}/bin/activate && pytest -vv --continue-on-collection-errors --junit-xml=junit.quick.xml" + "source ${venv_name}/bin/activate && pytest -vv --continue-on-collection-errors --junit-xml=junit.quick.xml" } -task testFull(type: Exec, dependsOn: [testQuick, installDevTest]) { - commandLine 'bash', '-x', '-c', - "source ${venv_name}/bin/activate && pytest -m 'not slow_integration' -vv --continue-on-collection-errors --junit-xml=junit.full.xml" -} - task cleanPythonCache(type: Exec) { commandLine 'bash', '-c', "find src -type f -name '*.py[co]' -delete -o -type d -name __pycache__ -delete -o -type d -empty -delete" diff --git a/metadata-ingestion-modules/airflow-plugin/pyproject.toml b/metadata-ingestion-modules/airflow-plugin/pyproject.toml index fba81486b9f67..648040c1951db 100644 --- a/metadata-ingestion-modules/airflow-plugin/pyproject.toml +++ b/metadata-ingestion-modules/airflow-plugin/pyproject.toml @@ -12,6 +12,7 @@ include = '\.pyi?$' [tool.isort] indent = ' ' +known_future_library = ['__future__', 'datahub.utilities._markupsafe_compat', 'datahub_provider._airflow_compat'] profile = 'black' sections = 'FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER' diff --git a/metadata-ingestion-modules/airflow-plugin/setup.cfg b/metadata-ingestion-modules/airflow-plugin/setup.cfg index 157bcce1c298d..c25256c5751b8 100644 --- a/metadata-ingestion-modules/airflow-plugin/setup.cfg +++ b/metadata-ingestion-modules/airflow-plugin/setup.cfg @@ -41,29 +41,29 @@ ignore_missing_imports = no [tool:pytest] asyncio_mode = auto -addopts = --cov=src --cov-report term-missing --cov-config setup.cfg --strict-markers +addopts = --cov=src --cov-report='' --cov-config setup.cfg --strict-markers -s -v +markers = + integration: marks tests to only run in integration (deselect with '-m "not integration"') testpaths = tests/unit tests/integration -[coverage:run] -# Because of some quirks in the way setup.cfg, coverage.py, pytest-cov, -# and tox interact, we should not uncomment the following line. -# See https://pytest-cov.readthedocs.io/en/latest/config.html and -# https://coverage.readthedocs.io/en/coverage-5.0/config.html. -# We also have some additional pytest/cov config options in tox.ini. -# source = src +# [coverage:run] +# # Because of some quirks in the way setup.cfg, coverage.py, pytest-cov, +# # and tox interact, we should not uncomment the following line. +# # See https://pytest-cov.readthedocs.io/en/latest/config.html and +# # https://coverage.readthedocs.io/en/coverage-5.0/config.html. +# # We also have some additional pytest/cov config options in tox.ini. +# # source = src -[coverage:paths] -# This is necessary for tox-based coverage to be counted properly. -source = - src - */site-packages +# [coverage:paths] +# # This is necessary for tox-based coverage to be counted properly. +# source = +# src +# */site-packages [coverage:report] -# The fail_under value ensures that at least some coverage data is collected. -# We override its value in the tox config. show_missing = true exclude_lines = pragma: no cover diff --git a/metadata-ingestion-modules/airflow-plugin/setup.py b/metadata-ingestion-modules/airflow-plugin/setup.py index 47069f59c314d..a5af881022d8c 100644 --- a/metadata-ingestion-modules/airflow-plugin/setup.py +++ b/metadata-ingestion-modules/airflow-plugin/setup.py @@ -1,5 +1,6 @@ import os import pathlib +from typing import Dict, Set import setuptools @@ -13,23 +14,43 @@ def get_long_description(): return pathlib.Path(os.path.join(root, "README.md")).read_text() +_version = package_metadata["__version__"] +_self_pin = f"=={_version}" if not _version.endswith("dev0") else "" + + rest_common = {"requests", "requests_file"} base_requirements = { # Compatibility. "dataclasses>=0.6; python_version < '3.7'", - # Typing extension should be >=3.10.0.2 ideally but we can't restrict due to Airflow 2.0.2 dependency conflict - "typing_extensions>=3.7.4.3 ; python_version < '3.8'", - "typing_extensions>=3.10.0.2,<4.6.0 ; python_version >= '3.8'", "mypy_extensions>=0.4.3", # Actual dependencies. - "typing-inspect", "pydantic>=1.5.1", "apache-airflow >= 2.0.2", *rest_common, - f"acryl-datahub == {package_metadata['__version__']}", } +plugins: Dict[str, Set[str]] = { + "datahub-rest": { + f"acryl-datahub[datahub-rest]{_self_pin}", + }, + "datahub-kafka": { + f"acryl-datahub[datahub-kafka]{_self_pin}", + }, + "datahub-file": { + f"acryl-datahub[sync-file-emitter]{_self_pin}", + }, + "plugin-v1": set(), + "plugin-v2": { + # The v2 plugin requires Python 3.8+. + f"acryl-datahub[sql-parser]{_self_pin}", + "openlineage-airflow==1.2.0; python_version >= '3.8'", + }, +} + +# Include datahub-rest in the base requirements. +base_requirements.update(plugins["datahub-rest"]) + mypy_stubs = { "types-dataclasses", @@ -45,11 +66,9 @@ def get_long_description(): # versions 0.1.13 and 0.1.14 seem to have issues "types-click==0.1.12", "types-tabulate", - # avrogen package requires this - "types-pytz", } -base_dev_requirements = { +dev_requirements = { *base_requirements, *mypy_stubs, "black==22.12.0", @@ -66,6 +85,7 @@ def get_long_description(): "pytest-cov>=2.8.1", "tox", "deepdiff", + "tenacity", "requests-mock", "freezegun", "jsonpickle", @@ -74,8 +94,24 @@ def get_long_description(): "packaging", } -dev_requirements = { - *base_dev_requirements, +integration_test_requirements = { + *dev_requirements, + *plugins["datahub-file"], + *plugins["datahub-kafka"], + f"acryl-datahub[testing-utils]{_self_pin}", + # Extra requirements for loading our test dags. + "apache-airflow[snowflake]>=2.0.2", + # https://github.com/snowflakedb/snowflake-sqlalchemy/issues/350 + # Eventually we want to set this to "snowflake-sqlalchemy>=1.4.3". + # However, that doesn't work with older versions of Airflow. Instead + # of splitting this into integration-test-old and integration-test-new, + # adding a bound to SQLAlchemy was the simplest solution. + "sqlalchemy<1.4.42", + # To avoid https://github.com/snowflakedb/snowflake-connector-python/issues/1188, + # we need https://github.com/snowflakedb/snowflake-connector-python/pull/1193 + "snowflake-connector-python>=2.7.10", + "virtualenv", # needed by PythonVirtualenvOperator + "apache-airflow-providers-sqlite", } @@ -88,7 +124,7 @@ def get_long_description(): setuptools.setup( # Package metadata. name=package_metadata["__package_name__"], - version=package_metadata["__version__"], + version=_version, url="https://datahubproject.io/", project_urls={ "Documentation": "https://datahubproject.io/docs/", @@ -131,17 +167,8 @@ def get_long_description(): # Dependencies. install_requires=list(base_requirements), extras_require={ + **{plugin: list(dependencies) for plugin, dependencies in plugins.items()}, "dev": list(dev_requirements), - "datahub-kafka": [ - f"acryl-datahub[datahub-kafka] == {package_metadata['__version__']}" - ], - "integration-tests": [ - f"acryl-datahub[datahub-kafka] == {package_metadata['__version__']}", - # Extra requirements for Airflow. - "apache-airflow[snowflake]>=2.0.2", # snowflake is used in example dags - # Because of https://github.com/snowflakedb/snowflake-sqlalchemy/issues/350 we need to restrict SQLAlchemy's max version. - "SQLAlchemy<1.4.42", - "virtualenv", # needed by PythonVirtualenvOperator - ], + "integration-tests": list(integration_test_requirements), }, ) diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_airflow_shims.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_airflow_shims.py index 5ad20e1f72551..10f014fbd586f 100644 --- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_airflow_shims.py +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_airflow_shims.py @@ -1,3 +1,7 @@ +from typing import List + +import airflow.version +import packaging.version from airflow.models.baseoperator import BaseOperator from datahub_airflow_plugin._airflow_compat import AIRFLOW_PATCHED @@ -21,7 +25,35 @@ assert AIRFLOW_PATCHED +# Approach suggested by https://stackoverflow.com/a/11887885/5004662. +AIRFLOW_VERSION = packaging.version.parse(airflow.version.version) +HAS_AIRFLOW_STANDALONE_CMD = AIRFLOW_VERSION >= packaging.version.parse("2.2.0.dev0") +HAS_AIRFLOW_LISTENER_API = AIRFLOW_VERSION >= packaging.version.parse("2.3.0.dev0") +HAS_AIRFLOW_DAG_LISTENER_API = AIRFLOW_VERSION >= packaging.version.parse("2.5.0.dev0") + + +def get_task_inlets(operator: "Operator") -> List: + # From Airflow 2.4 _inlets is dropped and inlets used consistently. Earlier it was not the case, so we have to stick there to _inlets + if hasattr(operator, "_inlets"): + return operator._inlets # type: ignore[attr-defined, union-attr] + if hasattr(operator, "get_inlet_defs"): + return operator.get_inlet_defs() # type: ignore[attr-defined] + return operator.inlets + + +def get_task_outlets(operator: "Operator") -> List: + # From Airflow 2.4 _outlets is dropped and inlets used consistently. Earlier it was not the case, so we have to stick there to _outlets + # We have to use _outlets because outlets is empty in Airflow < 2.4.0 + if hasattr(operator, "_outlets"): + return operator._outlets # type: ignore[attr-defined, union-attr] + if hasattr(operator, "get_outlet_defs"): + return operator.get_outlet_defs() + return operator.outlets + + __all__ = [ + "AIRFLOW_VERSION", + "HAS_AIRFLOW_LISTENER_API", "Operator", "MappedOperator", "EmptyOperator", diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_config.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_config.py new file mode 100644 index 0000000000000..67843da2ba995 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_config.py @@ -0,0 +1,80 @@ +from typing import TYPE_CHECKING, Optional + +import datahub.emitter.mce_builder as builder +from airflow.configuration import conf +from datahub.configuration.common import ConfigModel + +if TYPE_CHECKING: + from datahub_airflow_plugin.hooks.datahub import DatahubGenericHook + + +class DatahubLineageConfig(ConfigModel): + # This class is shared between the lineage backend and the Airflow plugin. + # The defaults listed here are only relevant for the lineage backend. + # The Airflow plugin's default values come from the fallback values in + # the get_lineage_config() function below. + + enabled: bool = True + + # DataHub hook connection ID. + datahub_conn_id: str + + # Cluster to associate with the pipelines and tasks. Defaults to "prod". + cluster: str = builder.DEFAULT_FLOW_CLUSTER + + # If true, the owners field of the DAG will be capture as a DataHub corpuser. + capture_ownership_info: bool = True + + # If true, the tags field of the DAG will be captured as DataHub tags. + capture_tags_info: bool = True + + capture_executions: bool = False + + enable_extractors: bool = True + + log_level: Optional[str] = None + debug_emitter: bool = False + + disable_openlineage_plugin: bool = True + + # Note that this field is only respected by the lineage backend. + # The Airflow plugin behaves as if it were set to True. + graceful_exceptions: bool = True + + def make_emitter_hook(self) -> "DatahubGenericHook": + # This is necessary to avoid issues with circular imports. + from datahub_airflow_plugin.hooks.datahub import DatahubGenericHook + + return DatahubGenericHook(self.datahub_conn_id) + + +def get_lineage_config() -> DatahubLineageConfig: + """Load the DataHub plugin config from airflow.cfg.""" + + enabled = conf.get("datahub", "enabled", fallback=True) + datahub_conn_id = conf.get("datahub", "conn_id", fallback="datahub_rest_default") + cluster = conf.get("datahub", "cluster", fallback=builder.DEFAULT_FLOW_CLUSTER) + capture_tags_info = conf.get("datahub", "capture_tags_info", fallback=True) + capture_ownership_info = conf.get( + "datahub", "capture_ownership_info", fallback=True + ) + capture_executions = conf.get("datahub", "capture_executions", fallback=True) + enable_extractors = conf.get("datahub", "enable_extractors", fallback=True) + log_level = conf.get("datahub", "log_level", fallback=None) + debug_emitter = conf.get("datahub", "debug_emitter", fallback=False) + disable_openlineage_plugin = conf.get( + "datahub", "disable_openlineage_plugin", fallback=True + ) + + return DatahubLineageConfig( + enabled=enabled, + datahub_conn_id=datahub_conn_id, + cluster=cluster, + capture_ownership_info=capture_ownership_info, + capture_tags_info=capture_tags_info, + capture_executions=capture_executions, + enable_extractors=enable_extractors, + log_level=log_level, + debug_emitter=debug_emitter, + disable_openlineage_plugin=disable_openlineage_plugin, + ) diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_listener_module.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_listener_module.py new file mode 100644 index 0000000000000..f39d37b122228 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_listener_module.py @@ -0,0 +1,7 @@ +from datahub_airflow_plugin.datahub_listener import get_airflow_plugin_listener + +_listener = get_airflow_plugin_listener() +if _listener: + on_task_instance_running = _listener.on_task_instance_running + on_task_instance_success = _listener.on_task_instance_success + on_task_instance_failed = _listener.on_task_instance_failed diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_ol_adapter.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_ol_adapter.py new file mode 100644 index 0000000000000..7d35791bf1db4 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_ol_adapter.py @@ -0,0 +1,23 @@ +import logging + +import datahub.emitter.mce_builder as builder +from openlineage.client.run import Dataset as OpenLineageDataset + +logger = logging.getLogger(__name__) + + +OL_SCHEME_TWEAKS = { + "sqlserver": "mssql", + "trino": "presto", + "awsathena": "athena", +} + + +def translate_ol_to_datahub_urn(ol_uri: OpenLineageDataset) -> str: + namespace = ol_uri.namespace + name = ol_uri.name + + scheme, *rest = namespace.split("://", maxsplit=1) + + platform = OL_SCHEME_TWEAKS.get(scheme, scheme) + return builder.make_dataset_urn(platform=platform, name=name) diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py new file mode 100644 index 0000000000000..f84b7b56f6119 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py @@ -0,0 +1,244 @@ +import contextlib +import logging +import unittest.mock +from typing import TYPE_CHECKING, Optional + +import datahub.emitter.mce_builder as builder +from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import ( + get_platform_from_sqlalchemy_uri, +) +from datahub.utilities.sqlglot_lineage import ( + SqlParsingResult, + create_lineage_sql_parsed_result, +) +from openlineage.airflow.extractors import BaseExtractor +from openlineage.airflow.extractors import ExtractorManager as OLExtractorManager +from openlineage.airflow.extractors import TaskMetadata +from openlineage.airflow.extractors.snowflake_extractor import SnowflakeExtractor +from openlineage.airflow.extractors.sql_extractor import SqlExtractor +from openlineage.airflow.utils import get_operator_class, try_import_from_string +from openlineage.client.facet import ( + ExtractionError, + ExtractionErrorRunFacet, + SqlJobFacet, +) + +from datahub_airflow_plugin._airflow_shims import Operator +from datahub_airflow_plugin._datahub_ol_adapter import OL_SCHEME_TWEAKS + +if TYPE_CHECKING: + from airflow.models import DagRun, TaskInstance + from datahub.ingestion.graph.client import DataHubGraph + +logger = logging.getLogger(__name__) +_DATAHUB_GRAPH_CONTEXT_KEY = "datahub_graph" +SQL_PARSING_RESULT_KEY = "datahub_sql" + + +class ExtractorManager(OLExtractorManager): + # TODO: On Airflow 2.7, the OLExtractorManager is part of the built-in Airflow API. + # When available, we should use that instead. The same goe for most of the OL + # extractors. + + def __init__(self): + super().__init__() + + _sql_operator_overrides = [ + # The OL BigQuery extractor has some complex logic to fetch detect + # the BigQuery job_id and fetch lineage from there. However, it can't + # generate CLL, so we disable it and use our own extractor instead. + "BigQueryOperator", + "BigQueryExecuteQueryOperator", + # Athena also does something similar. + "AthenaOperator", + "AWSAthenaOperator", + # Additional types that OL doesn't support. This is only necessary because + # on older versions of Airflow, these operators don't inherit from SQLExecuteQueryOperator. + "SqliteOperator", + ] + for operator in _sql_operator_overrides: + self.task_to_extractor.extractors[operator] = GenericSqlExtractor + + self._graph: Optional["DataHubGraph"] = None + + @contextlib.contextmanager + def _patch_extractors(self): + with contextlib.ExitStack() as stack: + # Patch the SqlExtractor.extract() method. + stack.enter_context( + unittest.mock.patch.object( + SqlExtractor, + "extract", + _sql_extractor_extract, + ) + ) + + # Patch the SnowflakeExtractor.default_schema property. + stack.enter_context( + unittest.mock.patch.object( + SnowflakeExtractor, + "default_schema", + property(snowflake_default_schema), + ) + ) + + # TODO: Override the BigQuery extractor to use the DataHub SQL parser. + # self.extractor_manager.add_extractor() + + # TODO: Override the Athena extractor to use the DataHub SQL parser. + + yield + + def extract_metadata( + self, + dagrun: "DagRun", + task: "Operator", + complete: bool = False, + task_instance: Optional["TaskInstance"] = None, + task_uuid: Optional[str] = None, + graph: Optional["DataHubGraph"] = None, + ) -> TaskMetadata: + self._graph = graph + with self._patch_extractors(): + return super().extract_metadata( + dagrun, task, complete, task_instance, task_uuid + ) + + def _get_extractor(self, task: "Operator") -> Optional[BaseExtractor]: + # By adding this, we can use the generic extractor as a fallback for + # any operator that inherits from SQLExecuteQueryOperator. + clazz = get_operator_class(task) + SQLExecuteQueryOperator = try_import_from_string( + "airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator" + ) + if SQLExecuteQueryOperator and issubclass(clazz, SQLExecuteQueryOperator): + self.task_to_extractor.extractors.setdefault( + clazz.__name__, GenericSqlExtractor + ) + + extractor = super()._get_extractor(task) + if extractor: + extractor.set_context(_DATAHUB_GRAPH_CONTEXT_KEY, self._graph) + return extractor + + +class GenericSqlExtractor(SqlExtractor): + # Note that the extract() method is patched elsewhere. + + @property + def default_schema(self): + return super().default_schema + + def _get_scheme(self) -> Optional[str]: + # Best effort conversion to DataHub platform names. + + with contextlib.suppress(Exception): + if self.hook: + if hasattr(self.hook, "get_uri"): + uri = self.hook.get_uri() + return get_platform_from_sqlalchemy_uri(uri) + + return self.conn.conn_type or super().dialect + + def _get_database(self) -> Optional[str]: + if self.conn: + # For BigQuery, the "database" is the project name. + if hasattr(self.conn, "project_id"): + return self.conn.project_id + + return self.conn.schema + return None + + +def _sql_extractor_extract(self: "SqlExtractor") -> TaskMetadata: + # Why not override the OL sql_parse method directly, instead of overriding + # extract()? A few reasons: + # + # 1. We would want to pass the default_db and graph instance into our sql parser + # method. The OL code doesn't pass the default_db (despite having it available), + # and it's not clear how to get the graph instance into that method. + # 2. OL has some janky logic to fetch table schemas as part of the sql extractor. + # We don't want that behavior and this lets us disable it. + # 3. Our SqlParsingResult already has DataHub urns, whereas using SqlMeta would + # require us to convert those urns to OL uris, just for them to get converted + # back to urns later on in our processing. + + task_name = f"{self.operator.dag_id}.{self.operator.task_id}" + sql = self.operator.sql + + run_facets = {} + job_facets = {"sql": SqlJobFacet(query=self._normalize_sql(sql))} + + # Prepare to run the SQL parser. + graph = self.context.get(_DATAHUB_GRAPH_CONTEXT_KEY, None) + + default_database = getattr(self.operator, "database", None) + if not default_database: + default_database = self.database + default_schema = self.default_schema + + # TODO: Add better handling for sql being a list of statements. + if isinstance(sql, list): + logger.info(f"Got list of SQL statements for {task_name}. Using first one.") + sql = sql[0] + + # Run the SQL parser. + scheme = self.scheme + platform = OL_SCHEME_TWEAKS.get(scheme, scheme) + self.log.debug( + "Running the SQL parser %s (platform=%s, default db=%s, schema=%s): %s", + "with graph client" if graph else "in offline mode", + platform, + default_database, + default_schema, + sql, + ) + sql_parsing_result: SqlParsingResult = create_lineage_sql_parsed_result( + query=sql, + graph=graph, + platform=platform, + platform_instance=None, + env=builder.DEFAULT_ENV, + database=default_database, + schema=default_schema, + ) + self.log.debug(f"Got sql lineage {sql_parsing_result}") + + if sql_parsing_result.debug_info.error: + error = sql_parsing_result.debug_info.error + run_facets["extractionError"] = ExtractionErrorRunFacet( + totalTasks=1, + failedTasks=1, + errors=[ + ExtractionError( + errorMessage=str(error), + stackTrace=None, + task="datahub_sql_parser", + taskNumber=None, + ) + ], + ) + + # Save sql_parsing_result to the facets dict. It is removed from the + # facet dict in the extractor's processing logic. + run_facets[SQL_PARSING_RESULT_KEY] = sql_parsing_result # type: ignore + + return TaskMetadata( + name=task_name, + inputs=[], + outputs=[], + run_facets=run_facets, + job_facets=job_facets, + ) + + +def snowflake_default_schema(self: "SnowflakeExtractor") -> Optional[str]: + if hasattr(self.operator, "schema") and self.operator.schema is not None: + return self.operator.schema + return ( + self.conn.extra_dejson.get("extra__snowflake__schema", "") + or self.conn.extra_dejson.get("schema", "") + or self.conn.schema + ) + # TODO: Should we try a fallback of: + # execute_query_on_hook(self.hook, "SELECT current_schema();")[0][0] diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py index b5e86e14d85d0..16585f70e820b 100644 --- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py @@ -1,4 +1,5 @@ -from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union, cast +from datetime import datetime +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union, cast from airflow.configuration import conf from datahub.api.entities.datajob import DataFlow, DataJob @@ -6,6 +7,7 @@ DataProcessInstance, InstanceRunResult, ) +from datahub.emitter.generic_emitter import Emitter from datahub.metadata.schema_classes import DataProcessTypeClass from datahub.utilities.urns.data_flow_urn import DataFlowUrn from datahub.utilities.urns.data_job_urn import DataJobUrn @@ -17,8 +19,6 @@ if TYPE_CHECKING: from airflow import DAG from airflow.models import DagRun, TaskInstance - from datahub.emitter.kafka_emitter import DatahubKafkaEmitter - from datahub.emitter.rest_emitter import DatahubRestEmitter from datahub_airflow_plugin._airflow_shims import Operator @@ -91,7 +91,7 @@ def _get_dependencies( ) # if the task triggers the subdag, link it to this node in the subdag - if subdag_task_id in _task_downstream_task_ids(upstream_task): + if subdag_task_id in sorted(_task_downstream_task_ids(upstream_task)): upstream_subdag_triggers.append(upstream_task_urn) # If the operator is an ExternalTaskSensor then we set the remote task as upstream. @@ -143,7 +143,7 @@ def generate_dataflow( """ id = dag.dag_id orchestrator = "airflow" - description = f"{dag.description}\n\n{dag.doc_md or ''}" + description = "\n\n".join(filter(None, [dag.description, dag.doc_md])) or None data_flow = DataFlow( env=cluster, id=id, orchestrator=orchestrator, description=description ) @@ -153,7 +153,7 @@ def generate_dataflow( allowed_flow_keys = [ "_access_control", "_concurrency", - "_default_view", + # "_default_view", "catchup", "fileloc", "is_paused_upon_creation", @@ -171,7 +171,7 @@ def generate_dataflow( data_flow.url = f"{base_url}/tree?dag_id={dag.dag_id}" if capture_owner and dag.owner: - data_flow.owners.add(dag.owner) + data_flow.owners.update(owner.strip() for owner in dag.owner.split(",")) if capture_tags and dag.tags: data_flow.tags.update(dag.tags) @@ -227,10 +227,7 @@ def generate_datajob( job_property_bag: Dict[str, str] = {} - allowed_task_keys = [ - "_downstream_task_ids", - "_inlets", - "_outlets", + allowed_task_keys: List[Union[str, Tuple[str, ...]]] = [ "_task_type", "_task_module", "depends_on_past", @@ -243,15 +240,28 @@ def generate_datajob( "trigger_rule", "wait_for_downstream", # In Airflow 2.3, _downstream_task_ids was renamed to downstream_task_ids - "downstream_task_ids", + ("downstream_task_ids", "_downstream_task_ids"), # In Airflow 2.4, _inlets and _outlets were removed in favor of non-private versions. - "inlets", - "outlets", + ("inlets", "_inlets"), + ("outlets", "_outlets"), ] for key in allowed_task_keys: - if hasattr(task, key): - job_property_bag[key] = repr(getattr(task, key)) + if isinstance(key, tuple): + out_key: str = key[0] + try_keys = key + else: + out_key = key + try_keys = (key,) + + for k in try_keys: + if hasattr(task, k): + v = getattr(task, k) + if out_key == "downstream_task_ids": + # Generate these in a consistent order. + v = list(sorted(v)) + job_property_bag[out_key] = repr(v) + break datajob.properties = job_property_bag base_url = conf.get("webserver", "base_url") @@ -288,7 +298,7 @@ def create_datajob_instance( @staticmethod def run_dataflow( - emitter: Union["DatahubRestEmitter", "DatahubKafkaEmitter"], + emitter: Emitter, cluster: str, dag_run: "DagRun", start_timestamp_millis: Optional[int] = None, @@ -340,7 +350,7 @@ def run_dataflow( @staticmethod def complete_dataflow( - emitter: Union["DatahubRestEmitter", "DatahubKafkaEmitter"], + emitter: Emitter, cluster: str, dag_run: "DagRun", end_timestamp_millis: Optional[int] = None, @@ -348,7 +358,7 @@ def complete_dataflow( ) -> None: """ - :param emitter: DatahubRestEmitter - the datahub rest emitter to emit the generated mcps + :param emitter: Emitter - the datahub emitter to emit the generated mcps :param cluster: str - name of the cluster :param dag_run: DagRun :param end_timestamp_millis: Optional[int] - the completion time in milliseconds if not set the current time will be used. @@ -386,7 +396,7 @@ def complete_dataflow( @staticmethod def run_datajob( - emitter: Union["DatahubRestEmitter", "DatahubKafkaEmitter"], + emitter: Emitter, cluster: str, ti: "TaskInstance", dag: "DAG", @@ -413,16 +423,13 @@ def run_datajob( job_property_bag["end_date"] = str(ti.end_date) job_property_bag["execution_date"] = str(ti.execution_date) job_property_bag["try_number"] = str(ti.try_number - 1) - job_property_bag["hostname"] = str(ti.hostname) job_property_bag["max_tries"] = str(ti.max_tries) # Not compatible with Airflow 1 if hasattr(ti, "external_executor_id"): job_property_bag["external_executor_id"] = str(ti.external_executor_id) - job_property_bag["pid"] = str(ti.pid) job_property_bag["state"] = str(ti.state) job_property_bag["operator"] = str(ti.operator) job_property_bag["priority_weight"] = str(ti.priority_weight) - job_property_bag["unixname"] = str(ti.unixname) job_property_bag["log_url"] = ti.log_url dpi.properties.update(job_property_bag) dpi.url = ti.log_url @@ -442,8 +449,10 @@ def run_datajob( dpi.type = DataProcessTypeClass.BATCH_AD_HOC if start_timestamp_millis is None: - assert ti.start_date - start_timestamp_millis = int(ti.start_date.timestamp() * 1000) + if ti.start_date: + start_timestamp_millis = int(ti.start_date.timestamp() * 1000) + else: + start_timestamp_millis = int(datetime.now().timestamp() * 1000) if attempt is None: attempt = ti.try_number @@ -458,7 +467,7 @@ def run_datajob( @staticmethod def complete_datajob( - emitter: Union["DatahubRestEmitter", "DatahubKafkaEmitter"], + emitter: Emitter, cluster: str, ti: "TaskInstance", dag: "DAG", @@ -469,7 +478,7 @@ def complete_datajob( ) -> DataProcessInstance: """ - :param emitter: DatahubRestEmitter + :param emitter: Emitter - the datahub emitter to emit the generated mcps :param cluster: str :param ti: TaskInstance :param dag: DAG @@ -483,8 +492,10 @@ def complete_datajob( datajob = AirflowGenerator.generate_datajob(cluster, ti.task, dag) if end_timestamp_millis is None: - assert ti.end_date - end_timestamp_millis = int(ti.end_date.timestamp() * 1000) + if ti.end_date: + end_timestamp_millis = int(ti.end_date.timestamp() * 1000) + else: + end_timestamp_millis = int(datetime.now().timestamp() * 1000) if result is None: # We should use TaskInstanceState but it is not available in Airflow 1 diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py new file mode 100644 index 0000000000000..a3f5cb489e29f --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py @@ -0,0 +1,494 @@ +import copy +import functools +import logging +import threading +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, TypeVar, cast + +import airflow +import datahub.emitter.mce_builder as builder +from datahub.api.entities.datajob import DataJob +from datahub.api.entities.dataprocess.dataprocess_instance import InstanceRunResult +from datahub.emitter.rest_emitter import DatahubRestEmitter +from datahub.ingestion.graph.client import DataHubGraph +from datahub.metadata.schema_classes import ( + FineGrainedLineageClass, + FineGrainedLineageDownstreamTypeClass, + FineGrainedLineageUpstreamTypeClass, +) +from datahub.telemetry import telemetry +from datahub.utilities.sqlglot_lineage import SqlParsingResult +from datahub.utilities.urns.dataset_urn import DatasetUrn +from openlineage.airflow.listener import TaskHolder +from openlineage.airflow.utils import redact_with_exclusions +from openlineage.client.serde import Serde + +from datahub_airflow_plugin._airflow_shims import ( + HAS_AIRFLOW_DAG_LISTENER_API, + Operator, + get_task_inlets, + get_task_outlets, +) +from datahub_airflow_plugin._config import DatahubLineageConfig, get_lineage_config +from datahub_airflow_plugin._datahub_ol_adapter import translate_ol_to_datahub_urn +from datahub_airflow_plugin._extractors import SQL_PARSING_RESULT_KEY, ExtractorManager +from datahub_airflow_plugin.client.airflow_generator import AirflowGenerator +from datahub_airflow_plugin.entities import _Entity + +_F = TypeVar("_F", bound=Callable[..., None]) +if TYPE_CHECKING: + from airflow.models import DAG, DagRun, TaskInstance + from sqlalchemy.orm import Session + + # To placate mypy on Airflow versions that don't have the listener API, + # we define a dummy hookimpl that's an identity function. + + def hookimpl(f: _F) -> _F: # type: ignore[misc] # noqa: F811 + return f + +else: + from airflow.listeners import hookimpl + +logger = logging.getLogger(__name__) + +_airflow_listener_initialized = False +_airflow_listener: Optional["DataHubListener"] = None +_RUN_IN_THREAD = True +_RUN_IN_THREAD_TIMEOUT = 30 + + +def get_airflow_plugin_listener() -> Optional["DataHubListener"]: + # Using globals instead of functools.lru_cache to make testing easier. + global _airflow_listener_initialized + global _airflow_listener + + if not _airflow_listener_initialized: + _airflow_listener_initialized = True + + plugin_config = get_lineage_config() + + if plugin_config.enabled: + _airflow_listener = DataHubListener(config=plugin_config) + + if plugin_config.disable_openlineage_plugin: + # Deactivate the OpenLineagePlugin listener to avoid conflicts. + from openlineage.airflow.plugin import OpenLineagePlugin + + OpenLineagePlugin.listeners = [] + + telemetry.telemetry_instance.ping( + "airflow-plugin-init", + { + "airflow-version": airflow.__version__, + "datahub-airflow-plugin": "v2", + "datahub-airflow-plugin-dag-events": HAS_AIRFLOW_DAG_LISTENER_API, + "capture_executions": plugin_config.capture_executions, + "capture_tags": plugin_config.capture_tags_info, + "capture_ownership": plugin_config.capture_ownership_info, + "enable_extractors": plugin_config.enable_extractors, + "disable_openlineage_plugin": plugin_config.disable_openlineage_plugin, + }, + ) + return _airflow_listener + + +def run_in_thread(f: _F) -> _F: + # This is also responsible for catching exceptions and logging them. + + @functools.wraps(f) + def wrapper(*args, **kwargs): + try: + if _RUN_IN_THREAD: + # A poor-man's timeout mechanism. + # This ensures that we don't hang the task if the extractors + # are slow or the DataHub API is slow to respond. + + thread = threading.Thread( + target=f, args=args, kwargs=kwargs, daemon=True + ) + thread.start() + + thread.join(timeout=_RUN_IN_THREAD_TIMEOUT) + if thread.is_alive(): + logger.warning( + f"Thread for {f.__name__} is still running after {_RUN_IN_THREAD_TIMEOUT} seconds. " + "Continuing without waiting for it to finish." + ) + else: + f(*args, **kwargs) + except Exception as e: + logger.exception(e) + + return cast(_F, wrapper) + + +class DataHubListener: + __name__ = "DataHubListener" + + def __init__(self, config: DatahubLineageConfig): + self.config = config + self._set_log_level() + + self._emitter = config.make_emitter_hook().make_emitter() + self._graph: Optional[DataHubGraph] = None + logger.info(f"DataHub plugin using {repr(self._emitter)}") + + # See discussion here https://github.com/OpenLineage/OpenLineage/pull/508 for + # why we need to keep track of tasks ourselves. + self._task_holder = TaskHolder() + + # In our case, we also want to cache the initial datajob object + # so that we can add to it when the task completes. + self._datajob_holder: Dict[str, DataJob] = {} + + self.extractor_manager = ExtractorManager() + + # This "inherits" from types.ModuleType to avoid issues with Airflow's listener plugin loader. + # It previously (v2.4.x and likely other versions too) would throw errors if it was not a module. + # https://github.com/apache/airflow/blob/e99a518970b2d349a75b1647f6b738c8510fa40e/airflow/listeners/listener.py#L56 + # self.__class__ = types.ModuleType + + @property + def emitter(self): + return self._emitter + + @property + def graph(self) -> Optional[DataHubGraph]: + if self._graph: + return self._graph + + if isinstance(self._emitter, DatahubRestEmitter) and not isinstance( + self._emitter, DataHubGraph + ): + # This is lazy initialized to avoid throwing errors on plugin load. + self._graph = self._emitter.to_graph() + self._emitter = self._graph + + return self._graph + + def _set_log_level(self) -> None: + """Set the log level for the plugin and its dependencies. + + This may need to be called multiple times, since Airflow sometimes + messes with the logging configuration after the plugin is loaded. + In particular, the loggers may get changed when the worker starts + executing a task. + """ + + if self.config.log_level: + logging.getLogger(__name__.split(".")[0]).setLevel(self.config.log_level) + if self.config.debug_emitter: + logging.getLogger("datahub.emitter").setLevel(logging.DEBUG) + + def _make_emit_callback(self) -> Callable[[Optional[Exception], str], None]: + def emit_callback(err: Optional[Exception], msg: str) -> None: + if err: + logger.error(f"Error sending metadata to datahub: {msg}", exc_info=err) + + return emit_callback + + def _extract_lineage( + self, + datajob: DataJob, + dagrun: "DagRun", + task: "Operator", + task_instance: "TaskInstance", + complete: bool = False, + ) -> None: + """ + Combine lineage (including column lineage) from task inlets/outlets and + extractor-generated task_metadata and write it to the datajob. This + routine is also responsible for converting the lineage to DataHub URNs. + """ + + input_urns: List[str] = [] + output_urns: List[str] = [] + fine_grained_lineages: List[FineGrainedLineageClass] = [] + + task_metadata = None + if self.config.enable_extractors: + task_metadata = self.extractor_manager.extract_metadata( + dagrun, + task, + complete=complete, + task_instance=task_instance, + task_uuid=str(datajob.urn), + graph=self.graph, + ) + logger.debug(f"Got task metadata: {task_metadata}") + + # Translate task_metadata.inputs/outputs to DataHub URNs. + input_urns.extend( + translate_ol_to_datahub_urn(dataset) for dataset in task_metadata.inputs + ) + output_urns.extend( + translate_ol_to_datahub_urn(dataset) + for dataset in task_metadata.outputs + ) + + # Add DataHub-native SQL parser results. + sql_parsing_result: Optional[SqlParsingResult] = None + if task_metadata: + sql_parsing_result = task_metadata.run_facets.pop( + SQL_PARSING_RESULT_KEY, None + ) + if sql_parsing_result: + if sql_parsing_result.debug_info.error: + datajob.properties["datahub_sql_parser_error"] = str( + sql_parsing_result.debug_info.error + ) + if not sql_parsing_result.debug_info.table_error: + input_urns.extend(sql_parsing_result.in_tables) + output_urns.extend(sql_parsing_result.out_tables) + + if sql_parsing_result.column_lineage: + fine_grained_lineages.extend( + FineGrainedLineageClass( + upstreamType=FineGrainedLineageUpstreamTypeClass.FIELD_SET, + downstreamType=FineGrainedLineageDownstreamTypeClass.FIELD, + upstreams=[ + builder.make_schema_field_urn( + upstream.table, upstream.column + ) + for upstream in column_lineage.upstreams + ], + downstreams=[ + builder.make_schema_field_urn( + downstream.table, downstream.column + ) + for downstream in [column_lineage.downstream] + if downstream.table + ], + ) + for column_lineage in sql_parsing_result.column_lineage + ) + + # Add DataHub-native inlets/outlets. + # These are filtered out by the extractor, so we need to add them manually. + input_urns.extend( + iolet.urn for iolet in get_task_inlets(task) if isinstance(iolet, _Entity) + ) + output_urns.extend( + iolet.urn for iolet in get_task_outlets(task) if isinstance(iolet, _Entity) + ) + + # Write the lineage to the datajob object. + datajob.inlets.extend(DatasetUrn.create_from_string(urn) for urn in input_urns) + datajob.outlets.extend( + DatasetUrn.create_from_string(urn) for urn in output_urns + ) + datajob.fine_grained_lineages.extend(fine_grained_lineages) + + # Merge in extra stuff that was present in the DataJob we constructed + # at the start of the task. + if complete: + original_datajob = self._datajob_holder.get(str(datajob.urn), None) + else: + self._datajob_holder[str(datajob.urn)] = datajob + original_datajob = None + + if original_datajob: + logger.debug("Merging start datajob into finish datajob") + datajob.inlets.extend(original_datajob.inlets) + datajob.outlets.extend(original_datajob.outlets) + datajob.fine_grained_lineages.extend(original_datajob.fine_grained_lineages) + + for k, v in original_datajob.properties.items(): + datajob.properties.setdefault(k, v) + + # Deduplicate inlets/outlets. + datajob.inlets = list(sorted(set(datajob.inlets), key=lambda x: str(x))) + datajob.outlets = list(sorted(set(datajob.outlets), key=lambda x: str(x))) + + # Write all other OL facets as DataHub properties. + if task_metadata: + for k, v in task_metadata.job_facets.items(): + datajob.properties[f"openlineage_job_facet_{k}"] = Serde.to_json( + redact_with_exclusions(v) + ) + + for k, v in task_metadata.run_facets.items(): + datajob.properties[f"openlineage_run_facet_{k}"] = Serde.to_json( + redact_with_exclusions(v) + ) + + @hookimpl + @run_in_thread + def on_task_instance_running( + self, + previous_state: None, + task_instance: "TaskInstance", + session: "Session", # This will always be QUEUED + ) -> None: + self._set_log_level() + + # This if statement mirrors the logic in https://github.com/OpenLineage/OpenLineage/pull/508. + if not hasattr(task_instance, "task"): + # The type ignore is to placate mypy on Airflow 2.1.x. + logger.warning( + f"No task set for task_id: {task_instance.task_id} - " # type: ignore[attr-defined] + f"dag_id: {task_instance.dag_id} - run_id {task_instance.run_id}" # type: ignore[attr-defined] + ) + return + + logger.debug( + f"DataHub listener got notification about task instance start for {task_instance.task_id}" + ) + + # Render templates in a copy of the task instance. + # This is necessary to get the correct operator args in the extractors. + task_instance = copy.deepcopy(task_instance) + task_instance.render_templates() + + # The type ignore is to placate mypy on Airflow 2.1.x. + dagrun: "DagRun" = task_instance.dag_run # type: ignore[attr-defined] + task = task_instance.task + dag: "DAG" = task.dag # type: ignore[assignment] + + self._task_holder.set_task(task_instance) + + # Handle async operators in Airflow 2.3 by skipping deferred state. + # Inspired by https://github.com/OpenLineage/OpenLineage/pull/1601 + if task_instance.next_method is not None: # type: ignore[attr-defined] + return + + # If we don't have the DAG listener API, we just pretend that + # the start of the task is the start of the DAG. + # This generates duplicate events, but it's better than not + # generating anything. + if not HAS_AIRFLOW_DAG_LISTENER_API: + self.on_dag_start(dagrun) + + datajob = AirflowGenerator.generate_datajob( + cluster=self.config.cluster, + task=task, + dag=dag, + capture_tags=self.config.capture_tags_info, + capture_owner=self.config.capture_ownership_info, + ) + + # TODO: Make use of get_task_location to extract github urls. + + # Add lineage info. + self._extract_lineage(datajob, dagrun, task, task_instance) + + # TODO: Add handling for Airflow mapped tasks using task_instance.map_index + + datajob.emit(self.emitter, callback=self._make_emit_callback()) + logger.debug(f"Emitted DataHub Datajob start: {datajob}") + + if self.config.capture_executions: + dpi = AirflowGenerator.run_datajob( + emitter=self.emitter, + cluster=self.config.cluster, + ti=task_instance, + dag=dag, + dag_run=dagrun, + datajob=datajob, + emit_templates=False, + ) + logger.debug(f"Emitted DataHub DataProcess Instance start: {dpi}") + + self.emitter.flush() + + logger.debug( + f"DataHub listener finished processing notification about task instance start for {task_instance.task_id}" + ) + + def on_task_instance_finish( + self, task_instance: "TaskInstance", status: InstanceRunResult + ) -> None: + dagrun: "DagRun" = task_instance.dag_run # type: ignore[attr-defined] + task = self._task_holder.get_task(task_instance) or task_instance.task + dag: "DAG" = task.dag # type: ignore[assignment] + + datajob = AirflowGenerator.generate_datajob( + cluster=self.config.cluster, + task=task, + dag=dag, + capture_tags=self.config.capture_tags_info, + capture_owner=self.config.capture_ownership_info, + ) + + # Add lineage info. + self._extract_lineage(datajob, dagrun, task, task_instance, complete=True) + + datajob.emit(self.emitter, callback=self._make_emit_callback()) + logger.debug(f"Emitted DataHub Datajob finish w/ status {status}: {datajob}") + + if self.config.capture_executions: + dpi = AirflowGenerator.complete_datajob( + emitter=self.emitter, + cluster=self.config.cluster, + ti=task_instance, + dag=dag, + dag_run=dagrun, + datajob=datajob, + result=status, + ) + logger.debug( + f"Emitted DataHub DataProcess Instance with status {status}: {dpi}" + ) + + self.emitter.flush() + + @hookimpl + @run_in_thread + def on_task_instance_success( + self, previous_state: None, task_instance: "TaskInstance", session: "Session" + ) -> None: + self._set_log_level() + + logger.debug( + f"DataHub listener got notification about task instance success for {task_instance.task_id}" + ) + self.on_task_instance_finish(task_instance, status=InstanceRunResult.SUCCESS) + logger.debug( + f"DataHub listener finished processing task instance success for {task_instance.task_id}" + ) + + @hookimpl + @run_in_thread + def on_task_instance_failed( + self, previous_state: None, task_instance: "TaskInstance", session: "Session" + ) -> None: + self._set_log_level() + + logger.debug( + f"DataHub listener got notification about task instance failure for {task_instance.task_id}" + ) + + # TODO: Handle UP_FOR_RETRY state. + self.on_task_instance_finish(task_instance, status=InstanceRunResult.FAILURE) + logger.debug( + f"DataHub listener finished processing task instance failure for {task_instance.task_id}" + ) + + def on_dag_start(self, dag_run: "DagRun") -> None: + dag = dag_run.dag + if not dag: + return + + dataflow = AirflowGenerator.generate_dataflow( + cluster=self.config.cluster, + dag=dag, + capture_tags=self.config.capture_tags_info, + capture_owner=self.config.capture_ownership_info, + ) + dataflow.emit(self.emitter, callback=self._make_emit_callback()) + + if HAS_AIRFLOW_DAG_LISTENER_API: + + @hookimpl + @run_in_thread + def on_dag_run_running(self, dag_run: "DagRun", msg: str) -> None: + self._set_log_level() + + logger.debug( + f"DataHub listener got notification about dag run start for {dag_run.dag_id}" + ) + + self.on_dag_start(dag_run) + + self.emitter.flush() + + # TODO: Add hooks for on_dag_run_success, on_dag_run_failed -> call AirflowGenerator.complete_dataflow diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin.py index d1cec9e5c1b54..c96fab31647f5 100644 --- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin.py +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin.py @@ -1,367 +1,74 @@ import contextlib import logging -import traceback -from typing import Any, Callable, Iterable, List, Optional, Union +import os -from airflow.configuration import conf -from airflow.lineage import PIPELINE_OUTLETS -from airflow.models.baseoperator import BaseOperator from airflow.plugins_manager import AirflowPlugin -from airflow.utils.module_loading import import_string -from cattr import structure -from datahub.api.entities.dataprocess.dataprocess_instance import InstanceRunResult from datahub_airflow_plugin._airflow_compat import AIRFLOW_PATCHED -from datahub_airflow_plugin._airflow_shims import MappedOperator, Operator -from datahub_airflow_plugin.client.airflow_generator import AirflowGenerator -from datahub_airflow_plugin.hooks.datahub import DatahubGenericHook -from datahub_airflow_plugin.lineage.datahub import DatahubLineageConfig +from datahub_airflow_plugin._airflow_shims import ( + HAS_AIRFLOW_DAG_LISTENER_API, + HAS_AIRFLOW_LISTENER_API, +) assert AIRFLOW_PATCHED logger = logging.getLogger(__name__) -TASK_ON_FAILURE_CALLBACK = "on_failure_callback" -TASK_ON_SUCCESS_CALLBACK = "on_success_callback" +_USE_AIRFLOW_LISTENER_INTERFACE = HAS_AIRFLOW_LISTENER_API and not os.getenv( + "DATAHUB_AIRFLOW_PLUGIN_USE_V1_PLUGIN", "false" +).lower() in ("true", "1") -def get_lineage_config() -> DatahubLineageConfig: - """Load the lineage config from airflow.cfg.""" +if _USE_AIRFLOW_LISTENER_INTERFACE: + try: + from openlineage.airflow.utils import try_import_from_string # noqa: F401 + except ImportError: + # If v2 plugin dependencies are not installed, we fall back to v1. + logger.debug("Falling back to v1 plugin due to missing dependencies.") + _USE_AIRFLOW_LISTENER_INTERFACE = False - enabled = conf.get("datahub", "enabled", fallback=True) - datahub_conn_id = conf.get("datahub", "conn_id", fallback="datahub_rest_default") - cluster = conf.get("datahub", "cluster", fallback="prod") - graceful_exceptions = conf.get("datahub", "graceful_exceptions", fallback=True) - capture_tags_info = conf.get("datahub", "capture_tags_info", fallback=True) - capture_ownership_info = conf.get( - "datahub", "capture_ownership_info", fallback=True - ) - capture_executions = conf.get("datahub", "capture_executions", fallback=True) - return DatahubLineageConfig( - enabled=enabled, - datahub_conn_id=datahub_conn_id, - cluster=cluster, - graceful_exceptions=graceful_exceptions, - capture_ownership_info=capture_ownership_info, - capture_tags_info=capture_tags_info, - capture_executions=capture_executions, - ) +with contextlib.suppress(Exception): + if not os.getenv("DATAHUB_AIRFLOW_PLUGIN_SKIP_FORK_PATCH", "false").lower() in ( + "true", + "1", + ): + # From https://github.com/apache/airflow/discussions/24463#discussioncomment-4404542 + # I'm not exactly sure why this fixes it, but I suspect it's that this + # forces the proxy settings to get cached before the fork happens. + # + # For more details, see https://github.com/python/cpython/issues/58037 + # and https://wefearchange.org/2018/11/forkmacos.rst.html + # and https://bugs.python.org/issue30385#msg293958 + # An alternative fix is to set NO_PROXY='*' -def _task_inlets(operator: "Operator") -> List: - # From Airflow 2.4 _inlets is dropped and inlets used consistently. Earlier it was not the case, so we have to stick there to _inlets - if hasattr(operator, "_inlets"): - return operator._inlets # type: ignore[attr-defined, union-attr] - return operator.inlets + from _scproxy import _get_proxy_settings + _get_proxy_settings() -def _task_outlets(operator: "Operator") -> List: - # From Airflow 2.4 _outlets is dropped and inlets used consistently. Earlier it was not the case, so we have to stick there to _outlets - # We have to use _outlets because outlets is empty in Airflow < 2.4.0 - if hasattr(operator, "_outlets"): - return operator._outlets # type: ignore[attr-defined, union-attr] - return operator.outlets +class DatahubPlugin(AirflowPlugin): + name = "datahub_plugin" -def get_inlets_from_task(task: BaseOperator, context: Any) -> Iterable[Any]: - # TODO: Fix for https://github.com/apache/airflow/commit/1b1f3fabc5909a447a6277cafef3a0d4ef1f01ae - # in Airflow 2.4. - # TODO: ignore/handle airflow's dataset type in our lineage - - inlets: List[Any] = [] - task_inlets = _task_inlets(task) - # From Airflow 2.3 this should be AbstractOperator but due to compatibility reason lets use BaseOperator - if isinstance(task_inlets, (str, BaseOperator)): - inlets = [ - task_inlets, - ] - - if task_inlets and isinstance(task_inlets, list): - inlets = [] - task_ids = ( - {o for o in task_inlets if isinstance(o, str)} - .union(op.task_id for op in task_inlets if isinstance(op, BaseOperator)) - .intersection(task.get_flat_relative_ids(upstream=True)) - ) - - from airflow.lineage import AUTO - - # pick up unique direct upstream task_ids if AUTO is specified - if AUTO.upper() in task_inlets or AUTO.lower() in task_inlets: - print("Picking up unique direct upstream task_ids as AUTO is specified") - task_ids = task_ids.union( - task_ids.symmetric_difference(task.upstream_task_ids) - ) - - inlets = task.xcom_pull( - context, task_ids=list(task_ids), dag_id=task.dag_id, key=PIPELINE_OUTLETS - ) - - # re-instantiate the obtained inlets - inlets = [ - structure(item["data"], import_string(item["type_name"])) - # _get_instance(structure(item, Metadata)) - for sublist in inlets - if sublist - for item in sublist - ] - - for inlet in task_inlets: - if not isinstance(inlet, str): - inlets.append(inlet) - - return inlets - - -def _make_emit_callback( - logger: logging.Logger, -) -> Callable[[Optional[Exception], str], None]: - def emit_callback(err: Optional[Exception], msg: str) -> None: - if err: - logger.error(f"Error sending metadata to datahub: {msg}", exc_info=err) - - return emit_callback - - -def datahub_task_status_callback(context, status): - ti = context["ti"] - task: "BaseOperator" = ti.task - dag = context["dag"] - - # This code is from the original airflow lineage code -> - # https://github.com/apache/airflow/blob/main/airflow/lineage/__init__.py - inlets = get_inlets_from_task(task, context) - - emitter = ( - DatahubGenericHook(context["_datahub_config"].datahub_conn_id) - .get_underlying_hook() - .make_emitter() - ) - - dataflow = AirflowGenerator.generate_dataflow( - cluster=context["_datahub_config"].cluster, - dag=dag, - capture_tags=context["_datahub_config"].capture_tags_info, - capture_owner=context["_datahub_config"].capture_ownership_info, - ) - task.log.info(f"Emitting Datahub Dataflow: {dataflow}") - dataflow.emit(emitter, callback=_make_emit_callback(task.log)) - - datajob = AirflowGenerator.generate_datajob( - cluster=context["_datahub_config"].cluster, - task=task, - dag=dag, - capture_tags=context["_datahub_config"].capture_tags_info, - capture_owner=context["_datahub_config"].capture_ownership_info, - ) - - for inlet in inlets: - datajob.inlets.append(inlet.urn) - - task_outlets = _task_outlets(task) - for outlet in task_outlets: - datajob.outlets.append(outlet.urn) - - task.log.info(f"Emitting Datahub Datajob: {datajob}") - datajob.emit(emitter, callback=_make_emit_callback(task.log)) - - if context["_datahub_config"].capture_executions: - dpi = AirflowGenerator.run_datajob( - emitter=emitter, - cluster=context["_datahub_config"].cluster, - ti=context["ti"], - dag=dag, - dag_run=context["dag_run"], - datajob=datajob, - start_timestamp_millis=int(ti.start_date.timestamp() * 1000), - ) - - task.log.info(f"Emitted Start Datahub Dataprocess Instance: {dpi}") - - dpi = AirflowGenerator.complete_datajob( - emitter=emitter, - cluster=context["_datahub_config"].cluster, - ti=context["ti"], - dag_run=context["dag_run"], - result=status, - dag=dag, - datajob=datajob, - end_timestamp_millis=int(ti.end_date.timestamp() * 1000), - ) - task.log.info(f"Emitted Completed Data Process Instance: {dpi}") - - emitter.flush() - - -def datahub_pre_execution(context): - ti = context["ti"] - task: "BaseOperator" = ti.task - dag = context["dag"] - - task.log.info("Running Datahub pre_execute method") - - emitter = ( - DatahubGenericHook(context["_datahub_config"].datahub_conn_id) - .get_underlying_hook() - .make_emitter() - ) - - # This code is from the original airflow lineage code -> - # https://github.com/apache/airflow/blob/main/airflow/lineage/__init__.py - inlets = get_inlets_from_task(task, context) - - datajob = AirflowGenerator.generate_datajob( - cluster=context["_datahub_config"].cluster, - task=context["ti"].task, - dag=dag, - capture_tags=context["_datahub_config"].capture_tags_info, - capture_owner=context["_datahub_config"].capture_ownership_info, - ) - - for inlet in inlets: - datajob.inlets.append(inlet.urn) - - task_outlets = _task_outlets(task) - - for outlet in task_outlets: - datajob.outlets.append(outlet.urn) - - task.log.info(f"Emitting Datahub dataJob {datajob}") - datajob.emit(emitter, callback=_make_emit_callback(task.log)) - - if context["_datahub_config"].capture_executions: - dpi = AirflowGenerator.run_datajob( - emitter=emitter, - cluster=context["_datahub_config"].cluster, - ti=context["ti"], - dag=dag, - dag_run=context["dag_run"], - datajob=datajob, - start_timestamp_millis=int(ti.start_date.timestamp() * 1000), - ) - - task.log.info(f"Emitting Datahub Dataprocess Instance: {dpi}") - - emitter.flush() - - -def _wrap_pre_execution(pre_execution): - def custom_pre_execution(context): - config = get_lineage_config() - if config.enabled: - context["_datahub_config"] = config - datahub_pre_execution(context) - - # Call original policy - if pre_execution: - pre_execution(context) - - return custom_pre_execution - - -def _wrap_on_failure_callback(on_failure_callback): - def custom_on_failure_callback(context): - config = get_lineage_config() - if config.enabled: - context["_datahub_config"] = config - try: - datahub_task_status_callback(context, status=InstanceRunResult.FAILURE) - except Exception as e: - if not config.graceful_exceptions: - raise e - else: - print(f"Exception: {traceback.format_exc()}") - - # Call original policy - if on_failure_callback: - on_failure_callback(context) - - return custom_on_failure_callback - - -def _wrap_on_success_callback(on_success_callback): - def custom_on_success_callback(context): - config = get_lineage_config() - if config.enabled: - context["_datahub_config"] = config - try: - datahub_task_status_callback(context, status=InstanceRunResult.SUCCESS) - except Exception as e: - if not config.graceful_exceptions: - raise e - else: - print(f"Exception: {traceback.format_exc()}") - - # Call original policy - if on_success_callback: - on_success_callback(context) - - return custom_on_success_callback - - -def task_policy(task: Union[BaseOperator, MappedOperator]) -> None: - task.log.debug(f"Setting task policy for Dag: {task.dag_id} Task: {task.task_id}") - # task.add_inlets(["auto"]) - # task.pre_execute = _wrap_pre_execution(task.pre_execute) - - # MappedOperator's callbacks don't have setters until Airflow 2.X.X - # https://github.com/apache/airflow/issues/24547 - # We can bypass this by going through partial_kwargs for now - if MappedOperator and isinstance(task, MappedOperator): # type: ignore - on_failure_callback_prop: property = getattr( - MappedOperator, TASK_ON_FAILURE_CALLBACK - ) - on_success_callback_prop: property = getattr( - MappedOperator, TASK_ON_SUCCESS_CALLBACK - ) - if not on_failure_callback_prop.fset or not on_success_callback_prop.fset: - task.log.debug( - "Using MappedOperator's partial_kwargs instead of callback properties" - ) - task.partial_kwargs[TASK_ON_FAILURE_CALLBACK] = _wrap_on_failure_callback( - task.on_failure_callback + if _USE_AIRFLOW_LISTENER_INTERFACE: + if HAS_AIRFLOW_DAG_LISTENER_API: + from datahub_airflow_plugin.datahub_listener import ( # type: ignore[misc] + get_airflow_plugin_listener, ) - task.partial_kwargs[TASK_ON_SUCCESS_CALLBACK] = _wrap_on_success_callback( - task.on_success_callback - ) - return - - task.on_failure_callback = _wrap_on_failure_callback(task.on_failure_callback) # type: ignore - task.on_success_callback = _wrap_on_success_callback(task.on_success_callback) # type: ignore - # task.pre_execute = _wrap_pre_execution(task.pre_execute) - - -def _wrap_task_policy(policy): - if policy and hasattr(policy, "_task_policy_patched_by"): - return policy - - def custom_task_policy(task): - policy(task) - task_policy(task) - - # Add a flag to the policy to indicate that we've patched it. - custom_task_policy._task_policy_patched_by = "datahub_plugin" # type: ignore[attr-defined] - return custom_task_policy + listeners: list = list(filter(None, [get_airflow_plugin_listener()])) -def _patch_policy(settings): - if hasattr(settings, "task_policy"): - datahub_task_policy = _wrap_task_policy(settings.task_policy) - settings.task_policy = datahub_task_policy + else: + # On Airflow < 2.5, we need the listener to be a module. + # This is just a quick shim layer to make that work. + # The DAG listener API was added at the same time as this method + # was fixed, so we're reusing the same check variable. + # + # Related Airflow change: https://github.com/apache/airflow/pull/27113. + import datahub_airflow_plugin._datahub_listener_module as _listener_module # type: ignore[misc] + listeners = [_listener_module] -def _patch_datahub_policy(): - with contextlib.suppress(ImportError): - import airflow_local_settings - _patch_policy(airflow_local_settings) - - from airflow.models.dagbag import settings - - _patch_policy(settings) - - -_patch_datahub_policy() - - -class DatahubPlugin(AirflowPlugin): - name = "datahub_plugin" +if not _USE_AIRFLOW_LISTENER_INTERFACE: + # Use the policy patcher mechanism on Airflow 2.2 and below. + import datahub_airflow_plugin.datahub_plugin_v22 # noqa: F401 diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py new file mode 100644 index 0000000000000..046fbb5efaa03 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py @@ -0,0 +1,336 @@ +import contextlib +import logging +import traceback +from typing import Any, Callable, Iterable, List, Optional, Union + +import airflow +from airflow.lineage import PIPELINE_OUTLETS +from airflow.models.baseoperator import BaseOperator +from airflow.utils.module_loading import import_string +from cattr import structure +from datahub.api.entities.dataprocess.dataprocess_instance import InstanceRunResult +from datahub.telemetry import telemetry + +from datahub_airflow_plugin._airflow_shims import ( + MappedOperator, + get_task_inlets, + get_task_outlets, +) +from datahub_airflow_plugin._config import get_lineage_config +from datahub_airflow_plugin.client.airflow_generator import AirflowGenerator +from datahub_airflow_plugin.hooks.datahub import DatahubGenericHook +from datahub_airflow_plugin.lineage.datahub import DatahubLineageConfig + +TASK_ON_FAILURE_CALLBACK = "on_failure_callback" +TASK_ON_SUCCESS_CALLBACK = "on_success_callback" + + +def get_task_inlets_advanced(task: BaseOperator, context: Any) -> Iterable[Any]: + # TODO: Fix for https://github.com/apache/airflow/commit/1b1f3fabc5909a447a6277cafef3a0d4ef1f01ae + # in Airflow 2.4. + # TODO: ignore/handle airflow's dataset type in our lineage + + inlets: List[Any] = [] + task_inlets = get_task_inlets(task) + # From Airflow 2.3 this should be AbstractOperator but due to compatibility reason lets use BaseOperator + if isinstance(task_inlets, (str, BaseOperator)): + inlets = [ + task_inlets, + ] + + if task_inlets and isinstance(task_inlets, list): + inlets = [] + task_ids = ( + {o for o in task_inlets if isinstance(o, str)} + .union(op.task_id for op in task_inlets if isinstance(op, BaseOperator)) + .intersection(task.get_flat_relative_ids(upstream=True)) + ) + + from airflow.lineage import AUTO + + # pick up unique direct upstream task_ids if AUTO is specified + if AUTO.upper() in task_inlets or AUTO.lower() in task_inlets: + print("Picking up unique direct upstream task_ids as AUTO is specified") + task_ids = task_ids.union( + task_ids.symmetric_difference(task.upstream_task_ids) + ) + + inlets = task.xcom_pull( + context, task_ids=list(task_ids), dag_id=task.dag_id, key=PIPELINE_OUTLETS + ) + + # re-instantiate the obtained inlets + inlets = [ + structure(item["data"], import_string(item["type_name"])) + # _get_instance(structure(item, Metadata)) + for sublist in inlets + if sublist + for item in sublist + ] + + for inlet in task_inlets: + if not isinstance(inlet, str): + inlets.append(inlet) + + return inlets + + +def _make_emit_callback( + logger: logging.Logger, +) -> Callable[[Optional[Exception], str], None]: + def emit_callback(err: Optional[Exception], msg: str) -> None: + if err: + logger.error(f"Error sending metadata to datahub: {msg}", exc_info=err) + + return emit_callback + + +def datahub_task_status_callback(context, status): + ti = context["ti"] + task: "BaseOperator" = ti.task + dag = context["dag"] + config: DatahubLineageConfig = context["_datahub_config"] + + # This code is from the original airflow lineage code -> + # https://github.com/apache/airflow/blob/main/airflow/lineage/__init__.py + inlets = get_task_inlets_advanced(task, context) + + emitter = ( + DatahubGenericHook(config.datahub_conn_id).get_underlying_hook().make_emitter() + ) + + dataflow = AirflowGenerator.generate_dataflow( + cluster=config.cluster, + dag=dag, + capture_tags=config.capture_tags_info, + capture_owner=config.capture_ownership_info, + ) + task.log.info(f"Emitting Datahub Dataflow: {dataflow}") + dataflow.emit(emitter, callback=_make_emit_callback(task.log)) + + datajob = AirflowGenerator.generate_datajob( + cluster=config.cluster, + task=task, + dag=dag, + capture_tags=config.capture_tags_info, + capture_owner=config.capture_ownership_info, + ) + + for inlet in inlets: + datajob.inlets.append(inlet.urn) + + task_outlets = get_task_outlets(task) + for outlet in task_outlets: + datajob.outlets.append(outlet.urn) + + task.log.info(f"Emitting Datahub Datajob: {datajob}") + datajob.emit(emitter, callback=_make_emit_callback(task.log)) + + if config.capture_executions: + dpi = AirflowGenerator.run_datajob( + emitter=emitter, + cluster=config.cluster, + ti=ti, + dag=dag, + dag_run=context["dag_run"], + datajob=datajob, + start_timestamp_millis=int(ti.start_date.timestamp() * 1000), + ) + + task.log.info(f"Emitted Start Datahub Dataprocess Instance: {dpi}") + + dpi = AirflowGenerator.complete_datajob( + emitter=emitter, + cluster=config.cluster, + ti=ti, + dag_run=context["dag_run"], + result=status, + dag=dag, + datajob=datajob, + end_timestamp_millis=int(ti.end_date.timestamp() * 1000), + ) + task.log.info(f"Emitted Completed Data Process Instance: {dpi}") + + emitter.flush() + + +def datahub_pre_execution(context): + ti = context["ti"] + task: "BaseOperator" = ti.task + dag = context["dag"] + config: DatahubLineageConfig = context["_datahub_config"] + + task.log.info("Running Datahub pre_execute method") + + emitter = ( + DatahubGenericHook(config.datahub_conn_id).get_underlying_hook().make_emitter() + ) + + # This code is from the original airflow lineage code -> + # https://github.com/apache/airflow/blob/main/airflow/lineage/__init__.py + inlets = get_task_inlets_advanced(task, context) + + datajob = AirflowGenerator.generate_datajob( + cluster=config.cluster, + task=ti.task, + dag=dag, + capture_tags=config.capture_tags_info, + capture_owner=config.capture_ownership_info, + ) + + for inlet in inlets: + datajob.inlets.append(inlet.urn) + + task_outlets = get_task_outlets(task) + + for outlet in task_outlets: + datajob.outlets.append(outlet.urn) + + task.log.info(f"Emitting Datahub dataJob {datajob}") + datajob.emit(emitter, callback=_make_emit_callback(task.log)) + + if config.capture_executions: + dpi = AirflowGenerator.run_datajob( + emitter=emitter, + cluster=config.cluster, + ti=ti, + dag=dag, + dag_run=context["dag_run"], + datajob=datajob, + start_timestamp_millis=int(ti.start_date.timestamp() * 1000), + ) + + task.log.info(f"Emitting Datahub Dataprocess Instance: {dpi}") + + emitter.flush() + + +def _wrap_pre_execution(pre_execution): + def custom_pre_execution(context): + config = get_lineage_config() + if config.enabled: + context["_datahub_config"] = config + datahub_pre_execution(context) + + # Call original policy + if pre_execution: + pre_execution(context) + + return custom_pre_execution + + +def _wrap_on_failure_callback(on_failure_callback): + def custom_on_failure_callback(context): + config = get_lineage_config() + if config.enabled: + context["_datahub_config"] = config + try: + datahub_task_status_callback(context, status=InstanceRunResult.FAILURE) + except Exception as e: + if not config.graceful_exceptions: + raise e + else: + print(f"Exception: {traceback.format_exc()}") + + # Call original policy + if on_failure_callback: + on_failure_callback(context) + + return custom_on_failure_callback + + +def _wrap_on_success_callback(on_success_callback): + def custom_on_success_callback(context): + config = get_lineage_config() + if config.enabled: + context["_datahub_config"] = config + try: + datahub_task_status_callback(context, status=InstanceRunResult.SUCCESS) + except Exception as e: + if not config.graceful_exceptions: + raise e + else: + print(f"Exception: {traceback.format_exc()}") + + # Call original policy + if on_success_callback: + on_success_callback(context) + + return custom_on_success_callback + + +def task_policy(task: Union[BaseOperator, MappedOperator]) -> None: + task.log.debug(f"Setting task policy for Dag: {task.dag_id} Task: {task.task_id}") + # task.add_inlets(["auto"]) + # task.pre_execute = _wrap_pre_execution(task.pre_execute) + + # MappedOperator's callbacks don't have setters until Airflow 2.X.X + # https://github.com/apache/airflow/issues/24547 + # We can bypass this by going through partial_kwargs for now + if MappedOperator and isinstance(task, MappedOperator): # type: ignore + on_failure_callback_prop: property = getattr( + MappedOperator, TASK_ON_FAILURE_CALLBACK + ) + on_success_callback_prop: property = getattr( + MappedOperator, TASK_ON_SUCCESS_CALLBACK + ) + if not on_failure_callback_prop.fset or not on_success_callback_prop.fset: + task.log.debug( + "Using MappedOperator's partial_kwargs instead of callback properties" + ) + task.partial_kwargs[TASK_ON_FAILURE_CALLBACK] = _wrap_on_failure_callback( + task.on_failure_callback + ) + task.partial_kwargs[TASK_ON_SUCCESS_CALLBACK] = _wrap_on_success_callback( + task.on_success_callback + ) + return + + task.on_failure_callback = _wrap_on_failure_callback(task.on_failure_callback) # type: ignore + task.on_success_callback = _wrap_on_success_callback(task.on_success_callback) # type: ignore + # task.pre_execute = _wrap_pre_execution(task.pre_execute) + + +def _wrap_task_policy(policy): + if policy and hasattr(policy, "_task_policy_patched_by"): + return policy + + def custom_task_policy(task): + policy(task) + task_policy(task) + + # Add a flag to the policy to indicate that we've patched it. + custom_task_policy._task_policy_patched_by = "datahub_plugin" # type: ignore[attr-defined] + return custom_task_policy + + +def _patch_policy(settings): + if hasattr(settings, "task_policy"): + datahub_task_policy = _wrap_task_policy(settings.task_policy) + settings.task_policy = datahub_task_policy + + +def _patch_datahub_policy(): + with contextlib.suppress(ImportError): + import airflow_local_settings + + _patch_policy(airflow_local_settings) + + from airflow.models.dagbag import settings + + _patch_policy(settings) + + plugin_config = get_lineage_config() + telemetry.telemetry_instance.ping( + "airflow-plugin-init", + { + "airflow-version": airflow.__version__, + "datahub-airflow-plugin": "v1", + "capture_executions": plugin_config.capture_executions, + "capture_tags": plugin_config.capture_tags_info, + "capture_ownership": plugin_config.capture_ownership_info, + }, + ) + + +_patch_datahub_policy() diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_emission_dag.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_emission_dag.py index f40295c6bb883..0d7cdb6b6e90a 100644 --- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_emission_dag.py +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_emission_dag.py @@ -2,12 +2,11 @@ This example demonstrates how to emit lineage to DataHub within an Airflow DAG. """ - from datetime import timedelta import datahub.emitter.mce_builder as builder from airflow import DAG -from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator +from airflow.operators.bash import BashOperator from airflow.utils.dates import days_ago from datahub_airflow_plugin.operators.datahub import DatahubEmitterOperator @@ -33,23 +32,10 @@ catchup=False, default_view="tree", ) as dag: - # This example shows a SnowflakeOperator followed by a lineage emission. However, the - # same DatahubEmitterOperator can be used to emit lineage in any context. - - sql = """CREATE OR REPLACE TABLE `mydb.schema.tableC` AS - WITH some_table AS ( - SELECT * FROM `mydb.schema.tableA` - ), - some_other_table AS ( - SELECT id, some_column FROM `mydb.schema.tableB` - ) - SELECT * FROM some_table - LEFT JOIN some_other_table ON some_table.unique_id=some_other_table.id""" - transformation_task = SnowflakeOperator( - task_id="snowflake_transformation", + transformation_task = BashOperator( + task_id="transformation_task", dag=dag, - snowflake_conn_id="snowflake_default", - sql=sql, + bash_command="echo 'This is where you might run your data tooling.'", ) emit_lineage_task = DatahubEmitterOperator( diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/hooks/datahub.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/hooks/datahub.py index 8fb7363f8cad1..9604931795ccb 100644 --- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/hooks/datahub.py +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/hooks/datahub.py @@ -1,7 +1,9 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple, Union from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook +from datahub.emitter.generic_emitter import Emitter +from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.metadata.com.linkedin.pegasus2avro.mxe import ( MetadataChangeEvent, MetadataChangeProposal, @@ -11,6 +13,7 @@ from airflow.models.connection import Connection from datahub.emitter.kafka_emitter import DatahubKafkaEmitter from datahub.emitter.rest_emitter import DatahubRestEmitter + from datahub.emitter.synchronized_file_emitter import SynchronizedFileEmitter from datahub.ingestion.sink.datahub_kafka import KafkaSinkConfig @@ -80,17 +83,24 @@ def make_emitter(self) -> "DatahubRestEmitter": return datahub.emitter.rest_emitter.DatahubRestEmitter(*self._get_config()) - def emit_mces(self, mces: List[MetadataChangeEvent]) -> None: + def emit( + self, + items: Sequence[ + Union[ + MetadataChangeEvent, + MetadataChangeProposal, + MetadataChangeProposalWrapper, + ] + ], + ) -> None: emitter = self.make_emitter() - for mce in mces: - emitter.emit_mce(mce) + for item in items: + emitter.emit(item) - def emit_mcps(self, mcps: List[MetadataChangeProposal]) -> None: - emitter = self.make_emitter() - - for mce in mcps: - emitter.emit_mcp(mce) + # Retained for backwards compatibility. + emit_mces = emit + emit_mcps = emit class DatahubKafkaHook(BaseHook): @@ -152,7 +162,16 @@ def make_emitter(self) -> "DatahubKafkaEmitter": sink_config = self._get_config() return datahub.emitter.kafka_emitter.DatahubKafkaEmitter(sink_config) - def emit_mces(self, mces: List[MetadataChangeEvent]) -> None: + def emit( + self, + items: Sequence[ + Union[ + MetadataChangeEvent, + MetadataChangeProposal, + MetadataChangeProposalWrapper, + ] + ], + ) -> None: emitter = self.make_emitter() errors = [] @@ -160,29 +179,50 @@ def callback(exc, msg): if exc: errors.append(exc) - for mce in mces: - emitter.emit_mce_async(mce, callback) + for mce in items: + emitter.emit(mce, callback) emitter.flush() if errors: - raise AirflowException(f"failed to push some MCEs: {errors}") + raise AirflowException(f"failed to push some metadata: {errors}") - def emit_mcps(self, mcps: List[MetadataChangeProposal]) -> None: - emitter = self.make_emitter() - errors = [] + # Retained for backwards compatibility. + emit_mces = emit + emit_mcps = emit - def callback(exc, msg): - if exc: - errors.append(exc) - for mcp in mcps: - emitter.emit_mcp_async(mcp, callback) +class SynchronizedFileHook(BaseHook): + conn_type = "datahub-file" - emitter.flush() + def __init__(self, datahub_conn_id: str) -> None: + super().__init__() + self.datahub_conn_id = datahub_conn_id - if errors: - raise AirflowException(f"failed to push some MCPs: {errors}") + def make_emitter(self) -> "SynchronizedFileEmitter": + from datahub.emitter.synchronized_file_emitter import SynchronizedFileEmitter + + conn = self.get_connection(self.datahub_conn_id) + filename = conn.host + if not filename: + raise AirflowException("filename parameter is required") + + return SynchronizedFileEmitter(filename=filename) + + def emit( + self, + items: Sequence[ + Union[ + MetadataChangeEvent, + MetadataChangeProposal, + MetadataChangeProposalWrapper, + ] + ], + ) -> None: + emitter = self.make_emitter() + + for item in items: + emitter.emit(item) class DatahubGenericHook(BaseHook): @@ -198,7 +238,9 @@ def __init__(self, datahub_conn_id: str) -> None: super().__init__() self.datahub_conn_id = datahub_conn_id - def get_underlying_hook(self) -> Union[DatahubRestHook, DatahubKafkaHook]: + def get_underlying_hook( + self, + ) -> Union[DatahubRestHook, DatahubKafkaHook, SynchronizedFileHook]: conn = self.get_connection(self.datahub_conn_id) # We need to figure out the underlying hook type. First check the @@ -213,6 +255,11 @@ def get_underlying_hook(self) -> Union[DatahubRestHook, DatahubKafkaHook]: or conn.conn_type == DatahubKafkaHook.conn_type.replace("-", "_") ): return DatahubKafkaHook(self.datahub_conn_id) + elif ( + conn.conn_type == SynchronizedFileHook.conn_type + or conn.conn_type == SynchronizedFileHook.conn_type.replace("-", "_") + ): + return SynchronizedFileHook(self.datahub_conn_id) elif "rest" in self.datahub_conn_id: return DatahubRestHook(self.datahub_conn_id) elif "kafka" in self.datahub_conn_id: @@ -222,8 +269,20 @@ def get_underlying_hook(self) -> Union[DatahubRestHook, DatahubKafkaHook]: f"DataHub cannot handle conn_type {conn.conn_type} in {conn}" ) - def make_emitter(self) -> Union["DatahubRestEmitter", "DatahubKafkaEmitter"]: + def make_emitter(self) -> Emitter: return self.get_underlying_hook().make_emitter() - def emit_mces(self, mces: List[MetadataChangeEvent]) -> None: - return self.get_underlying_hook().emit_mces(mces) + def emit( + self, + items: Sequence[ + Union[ + MetadataChangeEvent, + MetadataChangeProposal, + MetadataChangeProposalWrapper, + ] + ], + ) -> None: + return self.get_underlying_hook().emit(items) + + # Retained for backwards compatibility. + emit_mces = emit diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_lineage_core.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/_lineage_core.py similarity index 72% rename from metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_lineage_core.py rename to metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/_lineage_core.py index d91c039ffa718..f5f519fa23b11 100644 --- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_lineage_core.py +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/_lineage_core.py @@ -1,11 +1,10 @@ from datetime import datetime from typing import TYPE_CHECKING, Dict, List -import datahub.emitter.mce_builder as builder from datahub.api.entities.dataprocess.dataprocess_instance import InstanceRunResult -from datahub.configuration.common import ConfigModel from datahub.utilities.urns.dataset_urn import DatasetUrn +from datahub_airflow_plugin._config import DatahubLineageConfig from datahub_airflow_plugin.client.airflow_generator import AirflowGenerator from datahub_airflow_plugin.entities import _Entity @@ -15,39 +14,14 @@ from airflow.models.taskinstance import TaskInstance from datahub_airflow_plugin._airflow_shims import Operator - from datahub_airflow_plugin.hooks.datahub import DatahubGenericHook def _entities_to_urn_list(iolets: List[_Entity]) -> List[DatasetUrn]: return [DatasetUrn.create_from_string(let.urn) for let in iolets] -class DatahubBasicLineageConfig(ConfigModel): - enabled: bool = True - - # DataHub hook connection ID. - datahub_conn_id: str - - # Cluster to associate with the pipelines and tasks. Defaults to "prod". - cluster: str = builder.DEFAULT_FLOW_CLUSTER - - # If true, the owners field of the DAG will be capture as a DataHub corpuser. - capture_ownership_info: bool = True - - # If true, the tags field of the DAG will be captured as DataHub tags. - capture_tags_info: bool = True - - capture_executions: bool = False - - def make_emitter_hook(self) -> "DatahubGenericHook": - # This is necessary to avoid issues with circular imports. - from datahub_airflow_plugin.hooks.datahub import DatahubGenericHook - - return DatahubGenericHook(self.datahub_conn_id) - - def send_lineage_to_datahub( - config: DatahubBasicLineageConfig, + config: DatahubLineageConfig, operator: "Operator", inlets: List[_Entity], outlets: List[_Entity], diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/datahub.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/datahub.py index c41bb2b2a1e37..3ebe7831d08f9 100644 --- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/datahub.py +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/datahub.py @@ -4,8 +4,8 @@ from airflow.configuration import conf from airflow.lineage.backend import LineageBackend -from datahub_airflow_plugin._lineage_core import ( - DatahubBasicLineageConfig, +from datahub_airflow_plugin.lineage._lineage_core import ( + DatahubLineageConfig, send_lineage_to_datahub, ) @@ -13,14 +13,7 @@ from airflow.models.baseoperator import BaseOperator -class DatahubLineageConfig(DatahubBasicLineageConfig): - # If set to true, most runtime errors in the lineage backend will be - # suppressed and will not cause the overall task to fail. Note that - # configuration issues will still throw exceptions. - graceful_exceptions: bool = True - - -def get_lineage_config() -> DatahubLineageConfig: +def get_lineage_backend_config() -> DatahubLineageConfig: """Load the lineage config from airflow.cfg.""" # The kwargs pattern is also used for secret backends. @@ -51,8 +44,7 @@ class DatahubLineageBackend(LineageBackend): datahub_kwargs = { "datahub_conn_id": "datahub_rest_default", "capture_ownership_info": true, - "capture_tags_info": true, - "graceful_exceptions": true } + "capture_tags_info": true } # The above indentation is important! """ @@ -61,7 +53,7 @@ def __init__(self) -> None: # By attempting to get and parse the config, we can detect configuration errors # ahead of time. The init method is only called in Airflow 2.x. - _ = get_lineage_config() + _ = get_lineage_backend_config() # With Airflow 2.0, this can be an instance method. However, with Airflow 1.10.x, this # method is used statically, even though LineageBackend declares it as an instance variable. @@ -72,7 +64,7 @@ def send_lineage( outlets: Optional[List] = None, # unused context: Optional[Dict] = None, ) -> None: - config = get_lineage_config() + config = get_lineage_backend_config() if not config.enabled: return @@ -82,10 +74,4 @@ def send_lineage( config, operator, operator.inlets, operator.outlets, context ) except Exception as e: - if config.graceful_exceptions: - operator.log.error(e) - operator.log.info( - "Suppressing error because graceful_exceptions is set" - ) - else: - raise + operator.log.error(e) diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub.py index 109e7ddfe4dfa..15b50c51a561d 100644 --- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub.py +++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub.py @@ -57,7 +57,7 @@ def __init__( # type: ignore[no-untyped-def] datahub_conn_id=datahub_conn_id, **kwargs, ) - self.mces = mces + self.metadata = mces def execute(self, context): - self.generic_hook.get_underlying_hook().emit_mces(self.mces) + self.generic_hook.get_underlying_hook().emit(self.metadata) diff --git a/metadata-ingestion-modules/airflow-plugin/tests/conftest.py b/metadata-ingestion-modules/airflow-plugin/tests/conftest.py new file mode 100644 index 0000000000000..d2c45e723f1b0 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/conftest.py @@ -0,0 +1,6 @@ +def pytest_addoption(parser): + parser.addoption( + "--update-golden-files", + action="store_true", + default=False, + ) diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/basic_iolets.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/basic_iolets.py new file mode 100644 index 0000000000000..8b0803ab98422 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/basic_iolets.py @@ -0,0 +1,34 @@ +from datetime import datetime + +from airflow import DAG +from airflow.operators.bash import BashOperator + +from datahub_airflow_plugin.entities import Dataset, Urn + +with DAG( + "basic_iolets", + start_date=datetime(2023, 1, 1), + schedule_interval=None, + catchup=False, +) as dag: + task = BashOperator( + task_id="run_data_task", + dag=dag, + bash_command="echo 'This is where you might run your data tooling.'", + inlets=[ + Dataset(platform="snowflake", name="mydb.schema.tableA"), + Dataset(platform="snowflake", name="mydb.schema.tableB", env="DEV"), + Dataset( + platform="snowflake", + name="mydb.schema.tableC", + platform_instance="cloud", + ), + Urn( + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ), + ], + outlets=[ + Dataset("snowflake", "mydb.schema.tableD"), + Dataset("snowflake", "mydb.schema.tableE"), + ], + ) diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py new file mode 100644 index 0000000000000..1dd047f0a6dcc --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py @@ -0,0 +1,34 @@ +from datetime import datetime + +from airflow import DAG +from airflow.operators.bash import BashOperator + +from datahub_airflow_plugin.entities import Dataset, Urn + +with DAG( + "simple_dag", + start_date=datetime(2023, 1, 1), + schedule_interval=None, + catchup=False, + description="A simple DAG that runs a few fake data tasks.", +) as dag: + task1 = BashOperator( + task_id="task_1", + dag=dag, + bash_command="echo 'task 1'", + inlets=[ + Dataset(platform="snowflake", name="mydb.schema.tableA"), + Urn( + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ), + ], + outlets=[Dataset("snowflake", "mydb.schema.tableD")], + ) + + task2 = BashOperator( + task_id="run_another_data_task", + dag=dag, + bash_command="echo 'task 2'", + ) + + task1 >> task2 diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py new file mode 100644 index 0000000000000..347d0f88b0cd0 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py @@ -0,0 +1,32 @@ +from datetime import datetime + +from airflow import DAG +from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator + +SNOWFLAKE_COST_TABLE = "costs" +SNOWFLAKE_PROCESSED_TABLE = "processed_costs" + +with DAG( + "snowflake_operator", + start_date=datetime(2023, 1, 1), + schedule_interval=None, + catchup=False, +) as dag: + transform_cost_table = SnowflakeOperator( + snowflake_conn_id="my_snowflake", + task_id="transform_cost_table", + sql=""" + CREATE OR REPLACE TABLE {{ params.out_table_name }} AS + SELECT + id, + month, + total_cost, + area, + total_cost / area as cost_per_area + FROM {{ params.in_table_name }} + """, + params={ + "in_table_name": SNOWFLAKE_COST_TABLE, + "out_table_name": SNOWFLAKE_PROCESSED_TABLE, + }, + ) diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/sqlite_operator.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/sqlite_operator.py new file mode 100644 index 0000000000000..77faec3c8935a --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/sqlite_operator.py @@ -0,0 +1,75 @@ +from datetime import datetime + +from airflow import DAG +from airflow.providers.sqlite.operators.sqlite import SqliteOperator + +CONN_ID = "my_sqlite" + +COST_TABLE = "costs" +PROCESSED_TABLE = "processed_costs" + +with DAG( + "sqlite_operator", + start_date=datetime(2023, 1, 1), + schedule_interval=None, + catchup=False, +) as dag: + create_cost_table = SqliteOperator( + sqlite_conn_id=CONN_ID, + task_id="create_cost_table", + sql=""" + CREATE TABLE IF NOT EXISTS {{ params.table_name }} ( + id INTEGER PRIMARY KEY, + month TEXT NOT NULL, + total_cost REAL NOT NULL, + area REAL NOT NULL + ) + """, + params={"table_name": COST_TABLE}, + ) + + populate_cost_table = SqliteOperator( + sqlite_conn_id=CONN_ID, + task_id="populate_cost_table", + sql=""" + INSERT INTO {{ params.table_name }} (id, month, total_cost, area) + VALUES + (1, '2021-01', 100, 10), + (2, '2021-02', 200, 20), + (3, '2021-03', 300, 30) + """, + params={"table_name": COST_TABLE}, + ) + + transform_cost_table = SqliteOperator( + sqlite_conn_id=CONN_ID, + task_id="transform_cost_table", + sql=""" + CREATE TABLE IF NOT EXISTS {{ params.out_table_name }} AS + SELECT + id, + month, + total_cost, + area, + total_cost / area as cost_per_area + FROM {{ params.in_table_name }} + """, + params={ + "in_table_name": COST_TABLE, + "out_table_name": PROCESSED_TABLE, + }, + ) + + cleanup_tables = [] + for table_name in [COST_TABLE, PROCESSED_TABLE]: + cleanup_table = SqliteOperator( + sqlite_conn_id=CONN_ID, + task_id=f"cleanup_{table_name}", + sql=""" + DROP TABLE {{ params.table_name }} + """, + params={"table_name": table_name}, + ) + cleanup_tables.append(cleanup_table) + + create_cost_table >> populate_cost_table >> transform_cost_table >> cleanup_tables diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_basic_iolets.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_basic_iolets.json new file mode 100644 index 0000000000000..26aa2afaa831a --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_basic_iolets.json @@ -0,0 +1,533 @@ +[ +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/basic_iolets.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "None", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=basic_iolets", + "name": "basic_iolets" + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_data_task'", + "trigger_rule": "'all_success'", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=basic_iolets&_flt_3_task_id=run_data_task", + "name": "run_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_data_task'", + "trigger_rule": "'all_success'", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=basic_iolets&_flt_3_task_id=run_data_task", + "name": "run_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "0.176536", + "start_date": "2023-09-30 00:49:56.670239+00:00", + "end_date": "2023-09-30 00:49:56.846775+00:00", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "1", + "max_tries": "0", + "external_executor_id": "None", + "state": "success", + "operator": "BashOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets", + "name": "basic_iolets_run_data_task_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696034996670, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696034996670, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 2 + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696034996846, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +} +] \ No newline at end of file diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_simple_dag.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_simple_dag.json new file mode 100644 index 0000000000000..b2e3a1fe47da7 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_simple_dag.json @@ -0,0 +1,718 @@ +[ +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "None", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=simple_dag", + "name": "simple_dag", + "description": "A simple DAG that runs a few fake data tasks." + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'task_1'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'task_1'", + "trigger_rule": "'all_success'", + "wait_for_downstream": "False", + "downstream_task_ids": "['run_another_data_task']", + "inlets": "[]", + "outlets": "[]" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=task_1", + "name": "task_1", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'task_1'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'task_1'", + "trigger_rule": "'all_success'", + "wait_for_downstream": "False", + "downstream_task_ids": "['run_another_data_task']", + "inlets": "[]", + "outlets": "[]" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=task_1", + "name": "task_1", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "0.175983", + "start_date": "2023-09-30 00:48:58.943850+00:00", + "end_date": "2023-09-30 00:48:59.119833+00:00", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "1", + "max_tries": "0", + "external_executor_id": "None", + "state": "success", + "operator": "BashOperator", + "priority_weight": "2", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag", + "name": "simple_dag_task_1_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696034938943, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696034938943, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 2 + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696034939119, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "None", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=simple_dag", + "name": "simple_dag", + "description": "A simple DAG that runs a few fake data tasks." + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_another_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_another_data_task'", + "trigger_rule": "'all_success'", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=run_another_data_task", + "name": "run_another_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_another_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_another_data_task'", + "trigger_rule": "'all_success'", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=run_another_data_task", + "name": "run_another_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "0.129888", + "start_date": "2023-09-30 00:49:02.158752+00:00", + "end_date": "2023-09-30 00:49:02.288640+00:00", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "1", + "max_tries": "0", + "external_executor_id": "None", + "state": "success", + "operator": "BashOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag", + "name": "simple_dag_run_another_data_task_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696034942158, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696034942158, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 2 + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696034942288, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +} +] \ No newline at end of file diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets.json new file mode 100644 index 0000000000000..2e733c2ad40a9 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets.json @@ -0,0 +1,535 @@ +[ +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/basic_iolets.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=basic_iolets", + "name": "basic_iolets" + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_data_task'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[Dataset(platform='snowflake', name='mydb.schema.tableA', env='PROD', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableB', env='DEV', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableC', env='PROD', platform_instance='cloud'), Urn(_urn='urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)')]", + "outlets": "[Dataset(platform='snowflake', name='mydb.schema.tableD', env='PROD', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableE', env='PROD', platform_instance=None)]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'This is where you might run your data tooling.'\", \"dag\": \"<>\", \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"env\": \"DEV\", \"name\": \"mydb.schema.tableB\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableC\", \"platform\": \"snowflake\", \"platform_instance\": \"cloud\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableE\", \"platform\": \"snowflake\"}], \"task_id\": \"run_data_task\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_lock_for_execution\": true, \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'This is where you might run your data tooling.'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"env\": \"DEV\", \"name\": \"mydb.schema.tableB\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableC\", \"platform\": \"snowflake\", \"platform_instance\": \"cloud\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableE\", \"platform\": \"snowflake\"}], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_on_exit_code\": [99], \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"run_data_task\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [], \"wait_for_downstream\": false, \"wait_for_past_depends_before_skipping\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=basic_iolets&_flt_3_task_id=run_data_task", + "name": "run_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 01:13:14.266272+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "BashOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets&map_index=-1", + "name": "basic_iolets_run_data_task_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696036394266, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696036394266, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_data_task'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[Dataset(platform='snowflake', name='mydb.schema.tableA', env='PROD', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableB', env='DEV', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableC', env='PROD', platform_instance='cloud'), Urn(_urn='urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)')]", + "outlets": "[Dataset(platform='snowflake', name='mydb.schema.tableD', env='PROD', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableE', env='PROD', platform_instance=None)]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'This is where you might run your data tooling.'\", \"dag\": \"<>\", \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"env\": \"DEV\", \"name\": \"mydb.schema.tableB\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableC\", \"platform\": \"snowflake\", \"platform_instance\": \"cloud\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableE\", \"platform\": \"snowflake\"}], \"task_id\": \"run_data_task\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_lock_for_execution\": true, \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'This is where you might run your data tooling.'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"env\": \"DEV\", \"name\": \"mydb.schema.tableB\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableC\", \"platform\": \"snowflake\", \"platform_instance\": \"cloud\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableE\", \"platform\": \"snowflake\"}], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_on_exit_code\": [99], \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"run_data_task\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [], \"wait_for_downstream\": false, \"wait_for_past_depends_before_skipping\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=basic_iolets&_flt_3_task_id=run_data_task", + "name": "run_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696036394833, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +} +] \ No newline at end of file diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets_no_dag_listener.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets_no_dag_listener.json new file mode 100644 index 0000000000000..44b288efda954 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets_no_dag_listener.json @@ -0,0 +1,535 @@ +[ +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/basic_iolets.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=basic_iolets", + "name": "basic_iolets" + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_data_task'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[Dataset(platform='snowflake', name='mydb.schema.tableA', env='PROD', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableB', env='DEV', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableC', env='PROD', platform_instance='cloud'), Urn(_urn='urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)')]", + "outlets": "[Dataset(platform='snowflake', name='mydb.schema.tableD', env='PROD', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableE', env='PROD', platform_instance=None)]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'This is where you might run your data tooling.'\", \"dag\": \"<>\", \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"env\": \"DEV\", \"name\": \"mydb.schema.tableB\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableC\", \"platform\": \"snowflake\", \"platform_instance\": \"cloud\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableE\", \"platform\": \"snowflake\"}], \"task_id\": \"run_data_task\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'This is where you might run your data tooling.'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"env\": \"DEV\", \"name\": \"mydb.schema.tableB\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableC\", \"platform\": \"snowflake\", \"platform_instance\": \"cloud\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableE\", \"platform\": \"snowflake\"}], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_exit_code\": 99, \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"run_data_task\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [], \"wait_for_downstream\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=basic_iolets&_flt_3_task_id=run_data_task", + "name": "run_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:59:52.401211+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "BashOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets&map_index=-1", + "name": "basic_iolets_run_data_task_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696057192401, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057192401, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_data_task'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[Dataset(platform='snowflake', name='mydb.schema.tableA', env='PROD', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableB', env='DEV', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableC', env='PROD', platform_instance='cloud'), Urn(_urn='urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)')]", + "outlets": "[Dataset(platform='snowflake', name='mydb.schema.tableD', env='PROD', platform_instance=None), Dataset(platform='snowflake', name='mydb.schema.tableE', env='PROD', platform_instance=None)]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'This is where you might run your data tooling.'\", \"dag\": \"<>\", \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"env\": \"DEV\", \"name\": \"mydb.schema.tableB\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableC\", \"platform\": \"snowflake\", \"platform_instance\": \"cloud\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableE\", \"platform\": \"snowflake\"}], \"task_id\": \"run_data_task\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'This is where you might run your data tooling.'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"env\": \"DEV\", \"name\": \"mydb.schema.tableB\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableC\", \"platform\": \"snowflake\", \"platform_instance\": \"cloud\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}, {\"env\": \"PROD\", \"name\": \"mydb.schema.tableE\", \"platform\": \"snowflake\"}], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_exit_code\": 99, \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"run_data_task\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [], \"wait_for_downstream\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=basic_iolets&_flt_3_task_id=run_data_task", + "name": "run_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,cloud.mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableB,DEV)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableE,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:5d666eaf9015a31b3e305e8bc2dba078", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057192982, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +} +] \ No newline at end of file diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag.json new file mode 100644 index 0000000000000..454c509279e11 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag.json @@ -0,0 +1,666 @@ +[ +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=simple_dag", + "name": "simple_dag", + "description": "A simple DAG that runs a few fake data tasks." + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'task_1'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'task_1'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['run_another_data_task']", + "inlets": "[Dataset(platform='snowflake', name='mydb.schema.tableA', env='PROD', platform_instance=None), Urn(_urn='urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)')]", + "outlets": "[Dataset(platform='snowflake', name='mydb.schema.tableD', env='PROD', platform_instance=None)]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'task 1'\", \"dag\": \"<>\", \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}], \"task_id\": \"task_1\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_lock_for_execution\": true, \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'task 1'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [\"run_another_data_task\"], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_on_exit_code\": [99], \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"task_1\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [], \"wait_for_downstream\": false, \"wait_for_past_depends_before_skipping\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=task_1", + "name": "task_1", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:53:58.219003+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "BashOperator", + "priority_weight": "2", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag&map_index=-1", + "name": "simple_dag_task_1_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696056838219, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056838219, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'task_1'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'task_1'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['run_another_data_task']", + "inlets": "[Dataset(platform='snowflake', name='mydb.schema.tableA', env='PROD', platform_instance=None), Urn(_urn='urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)')]", + "outlets": "[Dataset(platform='snowflake', name='mydb.schema.tableD', env='PROD', platform_instance=None)]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'task 1'\", \"dag\": \"<>\", \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}], \"task_id\": \"task_1\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_lock_for_execution\": true, \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'task 1'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [\"run_another_data_task\"], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_on_exit_code\": [99], \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"task_1\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [], \"wait_for_downstream\": false, \"wait_for_past_depends_before_skipping\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=task_1", + "name": "task_1", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056838648, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_another_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_another_data_task'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'task 2'\", \"dag\": \"<>\", \"task_id\": \"run_another_data_task\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_lock_for_execution\": true, \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'task 2'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [], \"outlets\": [], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_on_exit_code\": [99], \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"run_another_data_task\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [\"task_1\"], \"wait_for_downstream\": false, \"wait_for_past_depends_before_skipping\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=run_another_data_task", + "name": "run_another_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:54:02.407515+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "BashOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag&map_index=-1", + "name": "simple_dag_run_another_data_task_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696056842407, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056842407, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_another_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_another_data_task'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'task 2'\", \"dag\": \"<>\", \"task_id\": \"run_another_data_task\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_lock_for_execution\": true, \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'task 2'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [], \"outlets\": [], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_on_exit_code\": [99], \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"run_another_data_task\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [\"task_1\"], \"wait_for_downstream\": false, \"wait_for_past_depends_before_skipping\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=run_another_data_task", + "name": "run_another_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056842831, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +} +] \ No newline at end of file diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag_no_dag_listener.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag_no_dag_listener.json new file mode 100644 index 0000000000000..73b5765e96b7d --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag_no_dag_listener.json @@ -0,0 +1,722 @@ +[ +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=simple_dag", + "name": "simple_dag", + "description": "A simple DAG that runs a few fake data tasks." + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'task_1'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'task_1'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['run_another_data_task']", + "inlets": "[Dataset(platform='snowflake', name='mydb.schema.tableA', env='PROD', platform_instance=None), Urn(_urn='urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)')]", + "outlets": "[Dataset(platform='snowflake', name='mydb.schema.tableD', env='PROD', platform_instance=None)]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'task 1'\", \"dag\": \"<>\", \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}], \"task_id\": \"task_1\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'task 1'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [\"run_another_data_task\"], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_exit_code\": 99, \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"task_1\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [], \"wait_for_downstream\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=task_1", + "name": "task_1", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:58:56.105026+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "BashOperator", + "priority_weight": "2", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag&map_index=-1", + "name": "simple_dag_task_1_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696057136105, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057136105, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'task_1'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'task_1'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['run_another_data_task']", + "inlets": "[Dataset(platform='snowflake', name='mydb.schema.tableA', env='PROD', platform_instance=None), Urn(_urn='urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)')]", + "outlets": "[Dataset(platform='snowflake', name='mydb.schema.tableD', env='PROD', platform_instance=None)]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'task 1'\", \"dag\": \"<>\", \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}], \"task_id\": \"task_1\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'task 1'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [\"run_another_data_task\"], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableA\", \"platform\": \"snowflake\"}, {\"_urn\": \"urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)\"}], \"outlets\": [{\"env\": \"PROD\", \"name\": \"mydb.schema.tableD\", \"platform\": \"snowflake\"}], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_exit_code\": 99, \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"task_1\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [], \"wait_for_downstream\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=task_1", + "name": "task_1", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableA,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableD,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fdbbbcd638bc0e91bbf8d7775efbecaf", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057136612, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=simple_dag", + "name": "simple_dag", + "description": "A simple DAG that runs a few fake data tasks." + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_another_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_another_data_task'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'task 2'\", \"dag\": \"<>\", \"task_id\": \"run_another_data_task\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'task 2'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [], \"outlets\": [], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_exit_code\": 99, \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"run_another_data_task\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [\"task_1\"], \"wait_for_downstream\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=run_another_data_task", + "name": "run_another_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:58:59.567004+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "BashOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag&map_index=-1", + "name": "simple_dag_run_another_data_task_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696057139567, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057139567, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'run_another_data_task'", + "execution_timeout": "None", + "sla": "None", + "task_id": "'run_another_data_task'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "openlineage_run_facet_unknownSourceAttribute": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"unknownItems\": [{\"name\": \"BashOperator\", \"properties\": {\"_BaseOperator__from_mapped\": false, \"_BaseOperator__init_kwargs\": {\"bash_command\": \"echo 'task 2'\", \"dag\": \"<>\", \"task_id\": \"run_another_data_task\"}, \"_BaseOperator__instantiated\": true, \"_dag\": \"<>\", \"_log\": \"<>\", \"append_env\": false, \"bash_command\": \"echo 'task 2'\", \"depends_on_past\": false, \"do_xcom_push\": true, \"downstream_task_ids\": [], \"email_on_failure\": true, \"email_on_retry\": true, \"executor_config\": {}, \"ignore_first_depends_on_past\": true, \"inlets\": [], \"outlets\": [], \"output_encoding\": \"utf-8\", \"owner\": \"airflow\", \"params\": \"<>\", \"pool\": \"default_pool\", \"pool_slots\": 1, \"priority_weight\": 1, \"queue\": \"default\", \"retries\": 0, \"retry_delay\": \"<>\", \"retry_exponential_backoff\": false, \"skip_exit_code\": 99, \"start_date\": \"<>\", \"task_group\": \"<>\", \"task_id\": \"run_another_data_task\", \"trigger_rule\": \"all_success\", \"upstream_task_ids\": [\"task_1\"], \"wait_for_downstream\": false, \"weight_rule\": \"downstream\"}, \"type\": \"operator\"}]}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=simple_dag&_flt_3_task_id=run_another_data_task", + "name": "run_another_data_task", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:888f71b79d9a0b162fe44acad7b2c2ae", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057140164, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +} +] \ No newline at end of file diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_snowflake_operator.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_snowflake_operator.json new file mode 100644 index 0000000000000..affc395d421da --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_snowflake_operator.json @@ -0,0 +1,507 @@ +[ +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,snowflake_operator,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=snowflake_operator", + "name": "snowflake_operator" + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,snowflake_operator,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,snowflake_operator,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,snowflake_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'transform_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n CREATE OR REPLACE TABLE processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n '", + "task_id": "'transform_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n CREATE OR REPLACE TABLE processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n \"}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=snowflake_operator&_flt_3_task_id=transform_cost_table", + "name": "transform_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,snowflake_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [ + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),id)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD),id)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),month)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD),month)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD),total_cost)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),area)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD),area)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),area)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD),cost_per_area)" + ], + "confidenceScore": 1.0 + } + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,snowflake_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,snowflake_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:3161034cc84e16a7c5e1906225734747", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:55:36.844976+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SnowflakeOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=transform_cost_table&dag_id=snowflake_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=transform_cost_table&dag_id=snowflake_operator&map_index=-1", + "name": "snowflake_operator_transform_cost_table_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696056936844, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:3161034cc84e16a7c5e1906225734747", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,snowflake_operator,prod),transform_cost_table)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:3161034cc84e16a7c5e1906225734747", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:3161034cc84e16a7c5e1906225734747", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:3161034cc84e16a7c5e1906225734747", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056936844, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,snowflake_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'transform_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n CREATE OR REPLACE TABLE processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n '", + "task_id": "'transform_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n CREATE OR REPLACE TABLE processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n \"}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=snowflake_operator&_flt_3_task_id=transform_cost_table", + "name": "transform_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,snowflake_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [ + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),id)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD),id)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),month)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD),month)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD),total_cost)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),area)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD),area)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),area)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD),cost_per_area)" + ], + "confidenceScore": 1.0 + } + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,datahub_test_database.datahub_test_schema.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,snowflake_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,snowflake_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:3161034cc84e16a7c5e1906225734747", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056938096, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "FAILURE", + "nativeResultType": "airflow" + } + } + } +} +] \ No newline at end of file diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json new file mode 100644 index 0000000000000..1a32b38ce055d --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json @@ -0,0 +1,1735 @@ +[ +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/sqlite_operator.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=sqlite_operator", + "name": "sqlite_operator" + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'create_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n CREATE TABLE IF NOT EXISTS costs (\\n id INTEGER PRIMARY KEY,\\n month TEXT NOT NULL,\\n total_cost REAL NOT NULL,\\n area REAL NOT NULL\\n )\\n '", + "task_id": "'create_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['populate_cost_table']", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n CREATE TABLE IF NOT EXISTS costs (\\n id INTEGER PRIMARY KEY,\\n month TEXT NOT NULL,\\n total_cost REAL NOT NULL,\\n area REAL NOT NULL\\n )\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=create_cost_table", + "name": "create_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fbeed1180fa0434e02ac6f75ace87869", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:56:24.632190+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SqliteOperator", + "priority_weight": "5", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=create_cost_table&dag_id=sqlite_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=create_cost_table&dag_id=sqlite_operator&map_index=-1", + "name": "sqlite_operator_create_cost_table_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696056984632, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fbeed1180fa0434e02ac6f75ace87869", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fbeed1180fa0434e02ac6f75ace87869", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fbeed1180fa0434e02ac6f75ace87869", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056984632, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'create_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n CREATE TABLE IF NOT EXISTS costs (\\n id INTEGER PRIMARY KEY,\\n month TEXT NOT NULL,\\n total_cost REAL NOT NULL,\\n area REAL NOT NULL\\n )\\n '", + "task_id": "'create_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['populate_cost_table']", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n CREATE TABLE IF NOT EXISTS costs (\\n id INTEGER PRIMARY KEY,\\n month TEXT NOT NULL,\\n total_cost REAL NOT NULL,\\n area REAL NOT NULL\\n )\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=create_cost_table", + "name": "create_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fbeed1180fa0434e02ac6f75ace87869", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056984947, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'populate_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "\"\\n INSERT INTO costs (id, month, total_cost, area)\\n VALUES\\n (1, '2021-01', 100, 10),\\n (2, '2021-02', 200, 20),\\n (3, '2021-03', 300, 30)\\n \"", + "task_id": "'populate_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['transform_cost_table']", + "inlets": "[]", + "outlets": "[]", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n INSERT INTO costs (id, month, total_cost, area)\\n VALUES\\n (1, '2021-01', 100, 10),\\n (2, '2021-02', 200, 20),\\n (3, '2021-03', 300, 30)\\n \"}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=populate_cost_table", + "name": "populate_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:04e1badac1eacd1c41123d07f579fa92", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:56:28.605901+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SqliteOperator", + "priority_weight": "4", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=populate_cost_table&dag_id=sqlite_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=populate_cost_table&dag_id=sqlite_operator&map_index=-1", + "name": "sqlite_operator_populate_cost_table_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696056988605, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:04e1badac1eacd1c41123d07f579fa92", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:04e1badac1eacd1c41123d07f579fa92", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:04e1badac1eacd1c41123d07f579fa92", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056988605, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'populate_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "\"\\n INSERT INTO costs (id, month, total_cost, area)\\n VALUES\\n (1, '2021-01', 100, 10),\\n (2, '2021-02', 200, 20),\\n (3, '2021-03', 300, 30)\\n \"", + "task_id": "'populate_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['transform_cost_table']", + "inlets": "[]", + "outlets": "[]", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n INSERT INTO costs (id, month, total_cost, area)\\n VALUES\\n (1, '2021-01', 100, 10),\\n (2, '2021-02', 200, 20),\\n (3, '2021-03', 300, 30)\\n \"}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=populate_cost_table", + "name": "populate_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:04e1badac1eacd1c41123d07f579fa92", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056989098, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'transform_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n CREATE TABLE IF NOT EXISTS processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n '", + "task_id": "'transform_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['cleanup_costs', 'cleanup_processed_costs']", + "inlets": "[]", + "outlets": "[]", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n CREATE TABLE IF NOT EXISTS processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n \"}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=transform_cost_table", + "name": "transform_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)" + ], + "fineGrainedLineages": [ + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),id)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),id)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),month)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),month)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),total_cost)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),area)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),cost_per_area)" + ], + "confidenceScore": 1.0 + } + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:56:32.888165+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SqliteOperator", + "priority_weight": "3", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=transform_cost_table&dag_id=sqlite_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=transform_cost_table&dag_id=sqlite_operator&map_index=-1", + "name": "sqlite_operator_transform_cost_table_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696056992888, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056992888, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'transform_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n CREATE TABLE IF NOT EXISTS processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n '", + "task_id": "'transform_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['cleanup_costs', 'cleanup_processed_costs']", + "inlets": "[]", + "outlets": "[]", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n CREATE TABLE IF NOT EXISTS processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n \"}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=transform_cost_table", + "name": "transform_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)" + ], + "fineGrainedLineages": [ + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),id)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),id)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),month)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),month)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),total_cost)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),area)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),cost_per_area)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),id)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),id)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),month)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),month)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),total_cost)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),area)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),cost_per_area)" + ], + "confidenceScore": 1.0 + } + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056993744, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'cleanup_costs'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n DROP TABLE costs\\n '", + "task_id": "'cleanup_costs'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n DROP TABLE costs\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=cleanup_costs", + "name": "cleanup_costs", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:07285de22276959612189d51336cc21a", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:56:37.745717+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SqliteOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=cleanup_costs&dag_id=sqlite_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=cleanup_costs&dag_id=sqlite_operator&map_index=-1", + "name": "sqlite_operator_cleanup_costs_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696056997745, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:07285de22276959612189d51336cc21a", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:07285de22276959612189d51336cc21a", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:07285de22276959612189d51336cc21a", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056997745, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'cleanup_costs'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n DROP TABLE costs\\n '", + "task_id": "'cleanup_costs'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n DROP TABLE costs\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=cleanup_costs", + "name": "cleanup_costs", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:07285de22276959612189d51336cc21a", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696056998672, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'cleanup_processed_costs'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n DROP TABLE processed_costs\\n '", + "task_id": "'cleanup_processed_costs'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n DROP TABLE processed_costs\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=cleanup_processed_costs", + "name": "cleanup_processed_costs", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:bab908abccf3cd6607b50fdaf3003372", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 06:56:42.645806+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SqliteOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=cleanup_processed_costs&dag_id=sqlite_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=cleanup_processed_costs&dag_id=sqlite_operator&map_index=-1", + "name": "sqlite_operator_cleanup_processed_costs_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696057002645, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:bab908abccf3cd6607b50fdaf3003372", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:bab908abccf3cd6607b50fdaf3003372", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:bab908abccf3cd6607b50fdaf3003372", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057002645, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'cleanup_processed_costs'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n DROP TABLE processed_costs\\n '", + "task_id": "'cleanup_processed_costs'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n DROP TABLE processed_costs\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=cleanup_processed_costs", + "name": "cleanup_processed_costs", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:bab908abccf3cd6607b50fdaf3003372", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057003759, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +} +] \ No newline at end of file diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator_no_dag_listener.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator_no_dag_listener.json new file mode 100644 index 0000000000000..c082be693e30c --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator_no_dag_listener.json @@ -0,0 +1,1955 @@ +[ +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/sqlite_operator.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=sqlite_operator", + "name": "sqlite_operator" + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'create_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n CREATE TABLE IF NOT EXISTS costs (\\n id INTEGER PRIMARY KEY,\\n month TEXT NOT NULL,\\n total_cost REAL NOT NULL,\\n area REAL NOT NULL\\n )\\n '", + "task_id": "'create_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['populate_cost_table']", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n CREATE TABLE IF NOT EXISTS costs (\\n id INTEGER PRIMARY KEY,\\n month TEXT NOT NULL,\\n total_cost REAL NOT NULL,\\n area REAL NOT NULL\\n )\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=create_cost_table", + "name": "create_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fbeed1180fa0434e02ac6f75ace87869", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 07:00:45.832554+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SqliteOperator", + "priority_weight": "5", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=create_cost_table&dag_id=sqlite_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=create_cost_table&dag_id=sqlite_operator&map_index=-1", + "name": "sqlite_operator_create_cost_table_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696057245832, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fbeed1180fa0434e02ac6f75ace87869", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fbeed1180fa0434e02ac6f75ace87869", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fbeed1180fa0434e02ac6f75ace87869", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057245832, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'create_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n CREATE TABLE IF NOT EXISTS costs (\\n id INTEGER PRIMARY KEY,\\n month TEXT NOT NULL,\\n total_cost REAL NOT NULL,\\n area REAL NOT NULL\\n )\\n '", + "task_id": "'create_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['populate_cost_table']", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n CREATE TABLE IF NOT EXISTS costs (\\n id INTEGER PRIMARY KEY,\\n month TEXT NOT NULL,\\n total_cost REAL NOT NULL,\\n area REAL NOT NULL\\n )\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=create_cost_table", + "name": "create_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "inputDatajobs": [], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:fbeed1180fa0434e02ac6f75ace87869", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057246734, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/sqlite_operator.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=sqlite_operator", + "name": "sqlite_operator" + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'populate_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "\"\\n INSERT INTO costs (id, month, total_cost, area)\\n VALUES\\n (1, '2021-01', 100, 10),\\n (2, '2021-02', 200, 20),\\n (3, '2021-03', 300, 30)\\n \"", + "task_id": "'populate_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['transform_cost_table']", + "inlets": "[]", + "outlets": "[]", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n INSERT INTO costs (id, month, total_cost, area)\\n VALUES\\n (1, '2021-01', 100, 10),\\n (2, '2021-02', 200, 20),\\n (3, '2021-03', 300, 30)\\n \"}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=populate_cost_table", + "name": "populate_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:04e1badac1eacd1c41123d07f579fa92", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 07:00:49.653938+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SqliteOperator", + "priority_weight": "4", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=populate_cost_table&dag_id=sqlite_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=populate_cost_table&dag_id=sqlite_operator&map_index=-1", + "name": "sqlite_operator_populate_cost_table_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696057249653, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:04e1badac1eacd1c41123d07f579fa92", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:04e1badac1eacd1c41123d07f579fa92", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:04e1badac1eacd1c41123d07f579fa92", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057249653, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'populate_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "\"\\n INSERT INTO costs (id, month, total_cost, area)\\n VALUES\\n (1, '2021-01', 100, 10),\\n (2, '2021-02', 200, 20),\\n (3, '2021-03', 300, 30)\\n \"", + "task_id": "'populate_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['transform_cost_table']", + "inlets": "[]", + "outlets": "[]", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n INSERT INTO costs (id, month, total_cost, area)\\n VALUES\\n (1, '2021-01', 100, 10),\\n (2, '2021-02', 200, 20),\\n (3, '2021-03', 300, 30)\\n \"}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=populate_cost_table", + "name": "populate_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:04e1badac1eacd1c41123d07f579fa92", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057250831, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/sqlite_operator.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=sqlite_operator", + "name": "sqlite_operator" + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'transform_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n CREATE TABLE IF NOT EXISTS processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n '", + "task_id": "'transform_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['cleanup_costs', 'cleanup_processed_costs']", + "inlets": "[]", + "outlets": "[]", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n CREATE TABLE IF NOT EXISTS processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n \"}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=transform_cost_table", + "name": "transform_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)" + ], + "fineGrainedLineages": [ + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),id)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),id)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),month)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),month)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),total_cost)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),area)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),cost_per_area)" + ], + "confidenceScore": 1.0 + } + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 07:00:53.989264+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SqliteOperator", + "priority_weight": "3", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=transform_cost_table&dag_id=sqlite_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=transform_cost_table&dag_id=sqlite_operator&map_index=-1", + "name": "sqlite_operator_transform_cost_table_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696057253989, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceOutput", + "aspect": { + "json": { + "outputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057253989, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'transform_cost_table'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n CREATE TABLE IF NOT EXISTS processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n '", + "task_id": "'transform_cost_table'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "['cleanup_costs', 'cleanup_processed_costs']", + "inlets": "[]", + "outlets": "[]", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n CREATE TABLE IF NOT EXISTS processed_costs AS\\n SELECT\\n id,\\n month,\\n total_cost,\\n area,\\n total_cost / area as cost_per_area\\n FROM costs\\n \"}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=transform_cost_table", + "name": "transform_cost_table", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)" + ], + "fineGrainedLineages": [ + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),id)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),id)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),month)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),month)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),total_cost)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),area)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),cost_per_area)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),id)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),id)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),month)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),month)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),total_cost)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),area)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),area)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD),total_cost)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD),cost_per_area)" + ], + "confidenceScore": 1.0 + } + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:64e5ff8f552e857b607832731e09808b", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057255628, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/sqlite_operator.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=sqlite_operator", + "name": "sqlite_operator" + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'cleanup_costs'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n DROP TABLE costs\\n '", + "task_id": "'cleanup_costs'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n DROP TABLE costs\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=cleanup_costs", + "name": "cleanup_costs", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:07285de22276959612189d51336cc21a", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 07:01:00.421177+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SqliteOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=cleanup_costs&dag_id=sqlite_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=cleanup_costs&dag_id=sqlite_operator&map_index=-1", + "name": "sqlite_operator_cleanup_costs_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696057260421, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:07285de22276959612189d51336cc21a", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:07285de22276959612189d51336cc21a", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:07285de22276959612189d51336cc21a", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057260421, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'cleanup_costs'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n DROP TABLE costs\\n '", + "task_id": "'cleanup_costs'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n DROP TABLE costs\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=cleanup_costs", + "name": "cleanup_costs", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:07285de22276959612189d51336cc21a", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057262258, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "dataFlowInfo", + "aspect": { + "json": { + "customProperties": { + "_access_control": "None", + "catchup": "False", + "fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/sqlite_operator.py'", + "is_paused_upon_creation": "None", + "start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))", + "tags": "[]", + "timezone": "Timezone('UTC')" + }, + "externalUrl": "http://airflow.example.com/tree?dag_id=sqlite_operator", + "name": "sqlite_operator" + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataFlow", + "entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'cleanup_processed_costs'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n DROP TABLE processed_costs\\n '", + "task_id": "'cleanup_processed_costs'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n DROP TABLE processed_costs\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=cleanup_processed_costs", + "name": "cleanup_processed_costs", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:bab908abccf3cd6607b50fdaf3003372", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceProperties", + "aspect": { + "json": { + "customProperties": { + "run_id": "manual_run_test", + "duration": "None", + "start_date": "2023-09-30 07:01:05.540192+00:00", + "end_date": "None", + "execution_date": "2023-09-27 21:34:38+00:00", + "try_number": "0", + "max_tries": "0", + "external_executor_id": "None", + "state": "running", + "operator": "SqliteOperator", + "priority_weight": "1", + "log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=cleanup_processed_costs&dag_id=sqlite_operator&map_index=-1" + }, + "externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=cleanup_processed_costs&dag_id=sqlite_operator&map_index=-1", + "name": "sqlite_operator_cleanup_processed_costs_manual_run_test", + "type": "BATCH_AD_HOC", + "created": { + "time": 1696057265540, + "actor": "urn:li:corpuser:datahub" + } + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:bab908abccf3cd6607b50fdaf3003372", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRelationships", + "aspect": { + "json": { + "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "upstreamInstances": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:bab908abccf3cd6607b50fdaf3003372", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceInput", + "aspect": { + "json": { + "inputs": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:bab908abccf3cd6607b50fdaf3003372", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057265540, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "STARTED", + "attempt": 1 + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInfo", + "aspect": { + "json": { + "customProperties": { + "depends_on_past": "False", + "email": "None", + "label": "'cleanup_processed_costs'", + "execution_timeout": "None", + "sla": "None", + "sql": "'\\n DROP TABLE processed_costs\\n '", + "task_id": "'cleanup_processed_costs'", + "trigger_rule": "", + "wait_for_downstream": "False", + "downstream_task_ids": "[]", + "inlets": "[]", + "outlets": "[]", + "datahub_sql_parser_error": "Can only generate column-level lineage for select-like inner statements, not (outer statement type: )", + "openlineage_job_facet_sql": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet\", \"query\": \"\\n DROP TABLE processed_costs\\n \"}", + "openlineage_run_facet_extractionError": "{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ExtractionErrorRunFacet\", \"errors\": [{\"_producer\": \"https://github.com/OpenLineage/OpenLineage/tree/1.2.0/integration/airflow\", \"_schemaURL\": \"https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet\", \"errorMessage\": \"Can only generate column-level lineage for select-like inner statements, not (outer statement type: )\", \"task\": \"datahub_sql_parser\"}], \"failedTasks\": 1, \"totalTasks\": 1}" + }, + "externalUrl": "http://airflow.example.com/taskinstance/list/?flt1_dag_id_equals=sqlite_operator&_flt_3_task_id=cleanup_processed_costs", + "name": "cleanup_processed_costs", + "type": { + "string": "COMMAND" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "dataJobInputOutput", + "aspect": { + "json": { + "inputDatasets": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)" + ], + "outputDatasets": [], + "inputDatajobs": [ + "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)" + ], + "fineGrainedLineages": [] + } + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.processed_costs,PROD)", + "changeType": "UPSERT", + "aspectName": "status", + "aspect": { + "json": { + "removed": false + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "ownership", + "aspect": { + "json": { + "owners": [ + { + "owner": "urn:li:corpuser:airflow", + "type": "DEVELOPER", + "source": { + "type": "SERVICE" + } + } + ], + "lastModified": { + "time": 0, + "actor": "urn:li:corpuser:airflow" + } + } + } +}, +{ + "entityType": "dataJob", + "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)", + "changeType": "UPSERT", + "aspectName": "globalTags", + "aspect": { + "json": { + "tags": [] + } + } +}, +{ + "entityType": "dataProcessInstance", + "entityUrn": "urn:li:dataProcessInstance:bab908abccf3cd6607b50fdaf3003372", + "changeType": "UPSERT", + "aspectName": "dataProcessInstanceRunEvent", + "aspect": { + "json": { + "timestampMillis": 1696057267631, + "partitionSpec": { + "type": "FULL_TABLE", + "partition": "FULL_TABLE_SNAPSHOT" + }, + "status": "COMPLETE", + "result": { + "type": "SUCCESS", + "nativeResultType": "airflow" + } + } + } +} +] \ No newline at end of file diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/integration_test_dummy.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/integration_test_dummy.py deleted file mode 100644 index 10cf3ad0a608a..0000000000000 --- a/metadata-ingestion-modules/airflow-plugin/tests/integration/integration_test_dummy.py +++ /dev/null @@ -1,2 +0,0 @@ -def test_dummy(): - pass diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py new file mode 100644 index 0000000000000..a2b7fd151a1e4 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py @@ -0,0 +1,392 @@ +import contextlib +import dataclasses +import functools +import logging +import os +import pathlib +import random +import signal +import subprocess +import time +from typing import Iterator, Sequence + +import pytest +import requests +import tenacity +from airflow.models.connection import Connection +from datahub.testing.compare_metadata_json import assert_metadata_files_equal + +from datahub_airflow_plugin._airflow_shims import ( + HAS_AIRFLOW_DAG_LISTENER_API, + HAS_AIRFLOW_LISTENER_API, + HAS_AIRFLOW_STANDALONE_CMD, +) + +pytestmark = pytest.mark.integration + +logger = logging.getLogger(__name__) +IS_LOCAL = os.environ.get("CI", "false") == "false" + +DAGS_FOLDER = pathlib.Path(__file__).parent / "dags" +GOLDENS_FOLDER = pathlib.Path(__file__).parent / "goldens" + + +@dataclasses.dataclass +class AirflowInstance: + airflow_home: pathlib.Path + airflow_port: int + pid: int + env_vars: dict + + username: str + password: str + + metadata_file: pathlib.Path + + @property + def airflow_url(self) -> str: + return f"http://localhost:{self.airflow_port}" + + @functools.cached_property + def session(self) -> requests.Session: + session = requests.Session() + session.auth = (self.username, self.password) + return session + + +@tenacity.retry( + reraise=True, + wait=tenacity.wait_fixed(1), + stop=tenacity.stop_after_delay(60), + retry=tenacity.retry_if_exception_type( + (AssertionError, requests.exceptions.RequestException) + ), +) +def _wait_for_airflow_healthy(airflow_port: int) -> None: + print("Checking if Airflow is ready...") + res = requests.get(f"http://localhost:{airflow_port}/health", timeout=5) + res.raise_for_status() + + airflow_health = res.json() + assert airflow_health["metadatabase"]["status"] == "healthy" + assert airflow_health["scheduler"]["status"] == "healthy" + + +class NotReadyError(Exception): + pass + + +@tenacity.retry( + reraise=True, + wait=tenacity.wait_fixed(1), + stop=tenacity.stop_after_delay(90), + retry=tenacity.retry_if_exception_type(NotReadyError), +) +def _wait_for_dag_finish( + airflow_instance: AirflowInstance, dag_id: str, require_success: bool +) -> None: + print("Checking if DAG is finished") + res = airflow_instance.session.get( + f"{airflow_instance.airflow_url}/api/v1/dags/{dag_id}/dagRuns", timeout=5 + ) + res.raise_for_status() + + dag_runs = res.json()["dag_runs"] + if not dag_runs: + raise NotReadyError("No DAG runs found") + + dag_run = dag_runs[0] + if dag_run["state"] == "failed": + if require_success: + raise ValueError("DAG failed") + # else - success is not required, so we're done. + + elif dag_run["state"] != "success": + raise NotReadyError(f"DAG has not finished yet: {dag_run['state']}") + + +@contextlib.contextmanager +def _run_airflow( + tmp_path: pathlib.Path, dags_folder: pathlib.Path, is_v1: bool +) -> Iterator[AirflowInstance]: + airflow_home = tmp_path / "airflow_home" + print(f"Using airflow home: {airflow_home}") + + if IS_LOCAL: + airflow_port = 11792 + else: + airflow_port = random.randint(10000, 12000) + print(f"Using airflow port: {airflow_port}") + + datahub_connection_name = "datahub_file_default" + meta_file = tmp_path / "datahub_metadata.json" + + environment = { + **os.environ, + "AIRFLOW_HOME": str(airflow_home), + "AIRFLOW__WEBSERVER__WEB_SERVER_PORT": str(airflow_port), + "AIRFLOW__WEBSERVER__BASE_URL": "http://airflow.example.com", + # Point airflow to the DAGs folder. + "AIRFLOW__CORE__LOAD_EXAMPLES": "False", + "AIRFLOW__CORE__DAGS_FOLDER": str(dags_folder), + "AIRFLOW__CORE__DAGS_ARE_PAUSED_AT_CREATION": "False", + # Have the Airflow API use username/password authentication. + "AIRFLOW__API__AUTH_BACKEND": "airflow.api.auth.backend.basic_auth", + # Configure the datahub plugin and have it write the MCPs to a file. + "AIRFLOW__CORE__LAZY_LOAD_PLUGINS": "False" if is_v1 else "True", + "AIRFLOW__DATAHUB__CONN_ID": datahub_connection_name, + f"AIRFLOW_CONN_{datahub_connection_name.upper()}": Connection( + conn_id="datahub_file_default", + conn_type="datahub-file", + host=str(meta_file), + ).get_uri(), + # Configure fake credentials for the Snowflake connection. + "AIRFLOW_CONN_MY_SNOWFLAKE": Connection( + conn_id="my_snowflake", + conn_type="snowflake", + login="fake_username", + password="fake_password", + schema="DATAHUB_TEST_SCHEMA", + extra={ + "account": "fake_account", + "database": "DATAHUB_TEST_DATABASE", + "warehouse": "fake_warehouse", + "role": "fake_role", + "insecure_mode": "true", + }, + ).get_uri(), + "AIRFLOW_CONN_MY_SQLITE": Connection( + conn_id="my_sqlite", + conn_type="sqlite", + host=str(tmp_path / "my_sqlite.db"), + ).get_uri(), + # Convenience settings. + "AIRFLOW__DATAHUB__LOG_LEVEL": "DEBUG", + "AIRFLOW__DATAHUB__DEBUG_EMITTER": "True", + "SQLALCHEMY_SILENCE_UBER_WARNING": "1", + } + + if not HAS_AIRFLOW_STANDALONE_CMD: + raise pytest.skip("Airflow standalone command is not available") + + # Start airflow in a background subprocess. + airflow_process = subprocess.Popen( + ["airflow", "standalone"], + env=environment, + ) + + try: + _wait_for_airflow_healthy(airflow_port) + print("Airflow is ready!") + + # Sleep for a few seconds to make sure the other Airflow processes are ready. + time.sleep(3) + + # Create an extra "airflow" user for easy testing. + if IS_LOCAL: + print("Creating an extra test user...") + subprocess.check_call( + [ + # fmt: off + "airflow", "users", "create", + "--username", "airflow", + "--password", "airflow", + "--firstname", "admin", + "--lastname", "admin", + "--role", "Admin", + "--email", "airflow@example.com", + # fmt: on + ], + env=environment, + ) + + # Sanity check that the plugin got loaded. + if not is_v1: + print("[debug] Listing loaded plugins") + subprocess.check_call( + ["airflow", "plugins", "-v"], + env=environment, + ) + + # Load the admin user's password. This is generated by the + # `airflow standalone` command, and is different from the + # airflow user that we create when running locally. + airflow_username = "admin" + airflow_password = (airflow_home / "standalone_admin_password.txt").read_text() + + airflow_instance = AirflowInstance( + airflow_home=airflow_home, + airflow_port=airflow_port, + pid=airflow_process.pid, + env_vars=environment, + username=airflow_username, + password=airflow_password, + metadata_file=meta_file, + ) + + yield airflow_instance + finally: + try: + # Attempt a graceful shutdown. + print("Shutting down airflow...") + airflow_process.send_signal(signal.SIGINT) + airflow_process.wait(timeout=30) + except subprocess.TimeoutExpired: + # If the graceful shutdown failed, kill the process. + print("Hard shutting down airflow...") + airflow_process.kill() + airflow_process.wait(timeout=3) + + +def check_golden_file( + pytestconfig: pytest.Config, + output_path: pathlib.Path, + golden_path: pathlib.Path, + ignore_paths: Sequence[str] = (), +) -> None: + update_golden = pytestconfig.getoption("--update-golden-files") + + assert_metadata_files_equal( + output_path=output_path, + golden_path=golden_path, + update_golden=update_golden, + copy_output=False, + ignore_paths=ignore_paths, + ignore_order=False, + ) + + +@dataclasses.dataclass +class DagTestCase: + dag_id: str + success: bool = True + + v2_only: bool = False + + +test_cases = [ + DagTestCase("simple_dag"), + DagTestCase("basic_iolets"), + DagTestCase("snowflake_operator", success=False, v2_only=True), + DagTestCase("sqlite_operator", v2_only=True), +] + + +@pytest.mark.parametrize( + ["golden_filename", "test_case", "is_v1"], + [ + # On Airflow <= 2.2, test plugin v1. + *[ + pytest.param( + f"v1_{test_case.dag_id}", + test_case, + True, + id=f"v1_{test_case.dag_id}", + marks=pytest.mark.skipif( + HAS_AIRFLOW_LISTENER_API, + reason="Not testing plugin v1 on newer Airflow versions", + ), + ) + for test_case in test_cases + if not test_case.v2_only + ], + *[ + pytest.param( + # On Airflow 2.3-2.4, test plugin v2 without dataFlows. + f"v2_{test_case.dag_id}" + if HAS_AIRFLOW_DAG_LISTENER_API + else f"v2_{test_case.dag_id}_no_dag_listener", + test_case, + False, + id=f"v2_{test_case.dag_id}" + if HAS_AIRFLOW_DAG_LISTENER_API + else f"v2_{test_case.dag_id}_no_dag_listener", + marks=pytest.mark.skipif( + not HAS_AIRFLOW_LISTENER_API, + reason="Cannot test plugin v2 without the Airflow plugin listener API", + ), + ) + for test_case in test_cases + ], + ], +) +def test_airflow_plugin( + pytestconfig: pytest.Config, + tmp_path: pathlib.Path, + golden_filename: str, + test_case: DagTestCase, + is_v1: bool, +) -> None: + # This test: + # - Configures the plugin. + # - Starts a local airflow instance in a subprocess. + # - Runs a DAG that uses an operator supported by the extractor. + # - Waits for the DAG to complete. + # - Validates the metadata generated against a golden file. + + if not is_v1 and not test_case.success and not HAS_AIRFLOW_DAG_LISTENER_API: + # Saw a number of issues in CI where this would fail to emit the last events + # due to an error in the SQLAlchemy listener. This never happened locally for me. + pytest.skip("Cannot test failure cases without the Airflow DAG listener API") + + golden_path = GOLDENS_FOLDER / f"{golden_filename}.json" + dag_id = test_case.dag_id + + with _run_airflow( + tmp_path, dags_folder=DAGS_FOLDER, is_v1=is_v1 + ) as airflow_instance: + print(f"Running DAG {dag_id}...") + subprocess.check_call( + [ + "airflow", + "dags", + "trigger", + "--exec-date", + "2023-09-27T21:34:38+00:00", + "-r", + "manual_run_test", + dag_id, + ], + env=airflow_instance.env_vars, + ) + + print("Waiting for DAG to finish...") + _wait_for_dag_finish( + airflow_instance, dag_id, require_success=test_case.success + ) + + print("Sleeping for a few seconds to let the plugin finish...") + time.sleep(10) + + check_golden_file( + pytestconfig=pytestconfig, + output_path=airflow_instance.metadata_file, + golden_path=golden_path, + ignore_paths=[ + # Timing-related items. + r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['start_date'\]", + r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['end_date'\]", + r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['duration'\]", + # Host-specific items. + r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['pid'\]", + r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['hostname'\]", + r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['unixname'\]", + # TODO: If we switched to Git urls, maybe we could get this to work consistently. + r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['fileloc'\]", + r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['openlineage_.*'\]", + ], + ) + + +if __name__ == "__main__": + # When run directly, just set up a local airflow instance. + import tempfile + + with _run_airflow( + tmp_path=pathlib.Path(tempfile.mkdtemp("airflow-plugin-test")), + dags_folder=DAGS_FOLDER, + is_v1=not HAS_AIRFLOW_LISTENER_API, + ) as airflow_instance: + # input("Press enter to exit...") + breakpoint() + print("quitting airflow") diff --git a/metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py b/metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py index 9aa901171cfa6..d8620e74d7e30 100644 --- a/metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py +++ b/metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py @@ -14,18 +14,21 @@ import pytest from airflow.lineage import apply_lineage, prepare_lineage from airflow.models import DAG, Connection, DagBag, DagRun, TaskInstance -from datahub_provider import get_provider_info -from datahub_provider._airflow_shims import AIRFLOW_PATCHED, EmptyOperator -from datahub_provider.entities import Dataset, Urn -from datahub_provider.hooks.datahub import DatahubKafkaHook, DatahubRestHook -from datahub_provider.operators.datahub import DatahubEmitterOperator + +from datahub_airflow_plugin import get_provider_info +from datahub_airflow_plugin._airflow_shims import ( + AIRFLOW_PATCHED, + AIRFLOW_VERSION, + EmptyOperator, +) +from datahub_airflow_plugin.entities import Dataset, Urn +from datahub_airflow_plugin.hooks.datahub import DatahubKafkaHook, DatahubRestHook +from datahub_airflow_plugin.operators.datahub import DatahubEmitterOperator assert AIRFLOW_PATCHED # TODO: Remove default_view="tree" arg. Figure out why is default_view being picked as "grid" and how to fix it ? -# Approach suggested by https://stackoverflow.com/a/11887885/5004662. -AIRFLOW_VERSION = packaging.version.parse(airflow.version.version) lineage_mce = builder.make_lineage_mce( [ @@ -105,7 +108,7 @@ def test_datahub_rest_hook(mock_emitter): mock_emitter.assert_called_once_with(config.host, None, None) instance = mock_emitter.return_value - instance.emit_mce.assert_called_with(lineage_mce) + instance.emit.assert_called_with(lineage_mce) @mock.patch("datahub.emitter.rest_emitter.DatahubRestEmitter", autospec=True) @@ -119,7 +122,7 @@ def test_datahub_rest_hook_with_timeout(mock_emitter): mock_emitter.assert_called_once_with(config.host, None, 5) instance = mock_emitter.return_value - instance.emit_mce.assert_called_with(lineage_mce) + instance.emit.assert_called_with(lineage_mce) @mock.patch("datahub.emitter.kafka_emitter.DatahubKafkaEmitter", autospec=True) @@ -131,11 +134,11 @@ def test_datahub_kafka_hook(mock_emitter): mock_emitter.assert_called_once() instance = mock_emitter.return_value - instance.emit_mce_async.assert_called() + instance.emit.assert_called() instance.flush.assert_called_once() -@mock.patch("datahub_provider.hooks.datahub.DatahubRestHook.emit_mces") +@mock.patch("datahub_provider.hooks.datahub.DatahubRestHook.emit") def test_datahub_lineage_operator(mock_emit): with patch_airflow_connection(datahub_rest_connection_config) as config: assert config.conn_id diff --git a/metadata-ingestion-modules/airflow-plugin/tests/unit/test_dummy.py b/metadata-ingestion-modules/airflow-plugin/tests/unit/test_dummy.py deleted file mode 100644 index 10cf3ad0a608a..0000000000000 --- a/metadata-ingestion-modules/airflow-plugin/tests/unit/test_dummy.py +++ /dev/null @@ -1,2 +0,0 @@ -def test_dummy(): - pass diff --git a/metadata-ingestion-modules/airflow-plugin/tests/unit/test_packaging.py b/metadata-ingestion-modules/airflow-plugin/tests/unit/test_packaging.py new file mode 100644 index 0000000000000..1d0ce5835f958 --- /dev/null +++ b/metadata-ingestion-modules/airflow-plugin/tests/unit/test_packaging.py @@ -0,0 +1,8 @@ +import setuptools + + +def test_package_list_match_inits(): + where = "./src" + package_list = set(setuptools.find_packages(where)) + namespace_packages = set(setuptools.find_namespace_packages(where)) + assert package_list == namespace_packages, "are you missing a package init file?" diff --git a/metadata-ingestion-modules/airflow-plugin/tox.ini b/metadata-ingestion-modules/airflow-plugin/tox.ini index 6a1c06aed8cdd..2f05854940d10 100644 --- a/metadata-ingestion-modules/airflow-plugin/tox.ini +++ b/metadata-ingestion-modules/airflow-plugin/tox.ini @@ -4,32 +4,23 @@ # and then run "tox" from this directory. [tox] -envlist = py3-quick,py3-full - -[gh-actions] -python = - 3.6: py3-full - 3.9: py3-full - -# Providing optional features that add dependencies from setup.py as deps here -# allows tox to recreate testenv when new dependencies are added to setup.py. -# Previous approach of using the tox global setting extras is not recommended -# as extras is only called when the testenv is created for the first time! -# see more here -> https://github.com/tox-dev/tox/issues/1105#issuecomment-448596282 +envlist = py38-airflow21, py38-airflow22, py310-airflow24, py310-airflow26, py310-airflow27 [testenv] -deps = - -e ../../metadata-ingestion/[.dev] +use_develop = true +extras = dev,integration-tests,plugin-v1 +deps = + -e ../../metadata-ingestion/ + # Airflow version + airflow21: apache-airflow~=2.1.0 + airflow22: apache-airflow~=2.2.0 + airflow24: apache-airflow~=2.4.0 + airflow26: apache-airflow~=2.6.0 + airflow27: apache-airflow~=2.7.0 commands = - pytest --cov={envsitepackagesdir}/datahub --cov={envsitepackagesdir}/datahub_provider \ - py3-quick: -m 'not integration and not slow_integration' --junit-xml=junit.quick.xml \ - py3-full: --cov-fail-under 65 --junit-xml=junit.full.xml \ - --continue-on-collection-errors \ - -vv + pytest --cov-append {posargs} -setenv = - AIRFLOW_HOME = /tmp/airflow/thisshouldnotexist-{envname} +# For Airflow 2.4+, add the plugin-v2 extra. +[testenv:py310-airflow{24,26,27}] +extras = dev,integration-tests,plugin-v2 -[testenv:py3-full] -deps = - ../../metadata-ingestion/.[dev] diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py index 8fb7b5f29cc22..34afa8cdb39a4 100644 --- a/metadata-ingestion/setup.py +++ b/metadata-ingestion/setup.py @@ -1,4 +1,3 @@ -import os import sys from typing import Dict, Set @@ -9,16 +8,9 @@ exec(fp.read(), package_metadata) -def get_long_description(): - root = os.path.dirname(__file__) - with open(os.path.join(root, "README.md")) as f: - description = f.read() - - return description - - base_requirements = { - "typing_extensions>=3.10.0.2", + # Typing extension should be >=3.10.0.2 ideally but we can't restrict due to a Airflow 2.1 dependency conflict. + "typing_extensions>=3.7.4.3", "mypy_extensions>=0.4.3", # Actual dependencies. "typing-inspect", @@ -270,6 +262,7 @@ def get_long_description(): # Sink plugins. "datahub-kafka": kafka_common, "datahub-rest": rest_common, + "sync-file-emitter": {"filelock"}, "datahub-lite": { "duckdb", "fastapi", @@ -670,7 +663,12 @@ def get_long_description(): }, license="Apache License 2.0", description="A CLI to work with DataHub metadata", - long_description=get_long_description(), + long_description="""\ +The `acryl-datahub` package contains a CLI and SDK for interacting with DataHub, +as well as an integration framework for pulling/pushing metadata from external systems. + +See the [DataHub docs](https://datahubproject.io/docs/metadata-ingestion). +""", long_description_content_type="text/markdown", classifiers=[ "Development Status :: 5 - Production/Stable", diff --git a/metadata-ingestion/src/datahub/api/entities/corpgroup/corpgroup.py b/metadata-ingestion/src/datahub/api/entities/corpgroup/corpgroup.py index 796786beba21b..a898e35bb810e 100644 --- a/metadata-ingestion/src/datahub/api/entities/corpgroup/corpgroup.py +++ b/metadata-ingestion/src/datahub/api/entities/corpgroup/corpgroup.py @@ -2,7 +2,7 @@ import logging from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, Iterable, List, Optional, Union +from typing import Callable, Iterable, List, Optional, Union import pydantic from pydantic import BaseModel @@ -11,9 +11,10 @@ from datahub.api.entities.corpuser.corpuser import CorpUser, CorpUserGenerationConfig from datahub.configuration.common import ConfigurationError from datahub.configuration.validate_field_rename import pydantic_renamed_field +from datahub.emitter.generic_emitter import Emitter from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.rest_emitter import DatahubRestEmitter -from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph +from datahub.ingestion.graph.client import DataHubGraph from datahub.metadata.schema_classes import ( CorpGroupEditableInfoClass, CorpGroupInfoClass, @@ -25,9 +26,6 @@ _Aspect, ) -if TYPE_CHECKING: - from datahub.emitter.kafka_emitter import DatahubKafkaEmitter - logger = logging.getLogger(__name__) @@ -194,30 +192,9 @@ def generate_mcp( entityUrn=urn, aspect=StatusClass(removed=False) ) - @staticmethod - def _datahub_graph_from_datahub_rest_emitter( - rest_emitter: DatahubRestEmitter, - ) -> DataHubGraph: - """ - Create a datahub graph instance from a REST Emitter. - A stop-gap implementation which is expected to be removed after PATCH support is implemented - for membership updates for users <-> groups - """ - graph = DataHubGraph( - config=DatahubClientConfig( - server=rest_emitter._gms_server, - token=rest_emitter._token, - timeout_sec=rest_emitter._connect_timeout_sec, - retry_status_codes=rest_emitter._retry_status_codes, - extra_headers=rest_emitter._session.headers, - disable_ssl_verification=rest_emitter._session.verify is False, - ) - ) - return graph - def emit( self, - emitter: Union[DatahubRestEmitter, "DatahubKafkaEmitter"], + emitter: Emitter, callback: Optional[Callable[[Exception, str], None]] = None, ) -> None: """ @@ -235,7 +212,7 @@ def emit( # who are passing in a DataHubRestEmitter today # we won't need this in the future once PATCH support is implemented as all emitters # will work - datahub_graph = self._datahub_graph_from_datahub_rest_emitter(emitter) + datahub_graph = emitter.to_graph() for mcp in self.generate_mcp( generation_config=CorpGroupGenerationConfig( override_editable=self.overrideEditable, datahub_graph=datahub_graph diff --git a/metadata-ingestion/src/datahub/api/entities/corpuser/corpuser.py b/metadata-ingestion/src/datahub/api/entities/corpuser/corpuser.py index c67eb02a870a5..9fe1ebedafca7 100644 --- a/metadata-ingestion/src/datahub/api/entities/corpuser/corpuser.py +++ b/metadata-ingestion/src/datahub/api/entities/corpuser/corpuser.py @@ -1,14 +1,14 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, Iterable, List, Optional, Union +from typing import Callable, Iterable, List, Optional import pydantic import datahub.emitter.mce_builder as builder from datahub.configuration.common import ConfigModel +from datahub.emitter.generic_emitter import Emitter from datahub.emitter.mcp import MetadataChangeProposalWrapper -from datahub.emitter.rest_emitter import DatahubRestEmitter from datahub.metadata.schema_classes import ( CorpUserEditableInfoClass, CorpUserInfoClass, @@ -16,9 +16,6 @@ StatusClass, ) -if TYPE_CHECKING: - from datahub.emitter.kafka_emitter import DatahubKafkaEmitter - @dataclass class CorpUserGenerationConfig: @@ -144,7 +141,7 @@ def generate_mcp( def emit( self, - emitter: Union[DatahubRestEmitter, "DatahubKafkaEmitter"], + emitter: Emitter, callback: Optional[Callable[[Exception, str], None]] = None, ) -> None: """ diff --git a/metadata-ingestion/src/datahub/api/entities/datajob/dataflow.py b/metadata-ingestion/src/datahub/api/entities/datajob/dataflow.py index 8a04768bc0a72..acd708ee81a5c 100644 --- a/metadata-ingestion/src/datahub/api/entities/datajob/dataflow.py +++ b/metadata-ingestion/src/datahub/api/entities/datajob/dataflow.py @@ -1,18 +1,9 @@ import logging from dataclasses import dataclass, field -from typing import ( - TYPE_CHECKING, - Callable, - Dict, - Iterable, - List, - Optional, - Set, - Union, - cast, -) +from typing import Callable, Dict, Iterable, List, Optional, Set, cast import datahub.emitter.mce_builder as builder +from datahub.emitter.generic_emitter import Emitter from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.metadata.schema_classes import ( AuditStampClass, @@ -29,10 +20,6 @@ ) from datahub.utilities.urns.data_flow_urn import DataFlowUrn -if TYPE_CHECKING: - from datahub.emitter.kafka_emitter import DatahubKafkaEmitter - from datahub.emitter.rest_emitter import DatahubRestEmitter - logger = logging.getLogger(__name__) @@ -170,7 +157,7 @@ def generate_mcp(self) -> Iterable[MetadataChangeProposalWrapper]: def emit( self, - emitter: Union["DatahubRestEmitter", "DatahubKafkaEmitter"], + emitter: Emitter, callback: Optional[Callable[[Exception, str], None]] = None, ) -> None: """ diff --git a/metadata-ingestion/src/datahub/api/entities/datajob/datajob.py b/metadata-ingestion/src/datahub/api/entities/datajob/datajob.py index 7eb6fc8c8d1a9..0face6415bacc 100644 --- a/metadata-ingestion/src/datahub/api/entities/datajob/datajob.py +++ b/metadata-ingestion/src/datahub/api/entities/datajob/datajob.py @@ -1,16 +1,16 @@ from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Optional, Set, Union +from typing import Callable, Dict, Iterable, List, Optional, Set import datahub.emitter.mce_builder as builder +from datahub.emitter.generic_emitter import Emitter from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.metadata.schema_classes import ( AuditStampClass, AzkabanJobTypeClass, DataJobInfoClass, DataJobInputOutputClass, - DataJobSnapshotClass, + FineGrainedLineageClass, GlobalTagsClass, - MetadataChangeEventClass, OwnerClass, OwnershipClass, OwnershipSourceClass, @@ -23,10 +23,6 @@ from datahub.utilities.urns.data_job_urn import DataJobUrn from datahub.utilities.urns.dataset_urn import DatasetUrn -if TYPE_CHECKING: - from datahub.emitter.kafka_emitter import DatahubKafkaEmitter - from datahub.emitter.rest_emitter import DatahubRestEmitter - @dataclass class DataJob: @@ -59,6 +55,7 @@ class DataJob: group_owners: Set[str] = field(default_factory=set) inlets: List[DatasetUrn] = field(default_factory=list) outlets: List[DatasetUrn] = field(default_factory=list) + fine_grained_lineages: List[FineGrainedLineageClass] = field(default_factory=list) upstream_urns: List[DataJobUrn] = field(default_factory=list) def __post_init__(self): @@ -103,31 +100,6 @@ def generate_tags_aspect(self) -> Iterable[GlobalTagsClass]: ) return [tags] - def generate_mce(self) -> MetadataChangeEventClass: - job_mce = MetadataChangeEventClass( - proposedSnapshot=DataJobSnapshotClass( - urn=str(self.urn), - aspects=[ - DataJobInfoClass( - name=self.name if self.name is not None else self.id, - type=AzkabanJobTypeClass.COMMAND, - description=self.description, - customProperties=self.properties, - externalUrl=self.url, - ), - DataJobInputOutputClass( - inputDatasets=[str(urn) for urn in self.inlets], - outputDatasets=[str(urn) for urn in self.outlets], - inputDatajobs=[str(urn) for urn in self.upstream_urns], - ), - *self.generate_ownership_aspect(), - *self.generate_tags_aspect(), - ], - ) - ) - - return job_mce - def generate_mcp(self) -> Iterable[MetadataChangeProposalWrapper]: mcp = MetadataChangeProposalWrapper( entityUrn=str(self.urn), @@ -159,7 +131,7 @@ def generate_mcp(self) -> Iterable[MetadataChangeProposalWrapper]: def emit( self, - emitter: Union["DatahubRestEmitter", "DatahubKafkaEmitter"], + emitter: Emitter, callback: Optional[Callable[[Exception, str], None]] = None, ) -> None: """ @@ -179,6 +151,7 @@ def generate_data_input_output_mcp(self) -> Iterable[MetadataChangeProposalWrapp inputDatasets=[str(urn) for urn in self.inlets], outputDatasets=[str(urn) for urn in self.outlets], inputDatajobs=[str(urn) for urn in self.upstream_urns], + fineGrainedLineages=self.fine_grained_lineages, ), ) yield mcp diff --git a/metadata-ingestion/src/datahub/api/entities/dataprocess/dataprocess_instance.py b/metadata-ingestion/src/datahub/api/entities/dataprocess/dataprocess_instance.py index 9ec389c3a0989..cf6080c7072e6 100644 --- a/metadata-ingestion/src/datahub/api/entities/dataprocess/dataprocess_instance.py +++ b/metadata-ingestion/src/datahub/api/entities/dataprocess/dataprocess_instance.py @@ -1,9 +1,10 @@ import time from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Optional, Union, cast +from typing import Callable, Dict, Iterable, List, Optional, Union, cast from datahub.api.entities.datajob import DataFlow, DataJob +from datahub.emitter.generic_emitter import Emitter from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.mcp_builder import DatahubKey from datahub.metadata.com.linkedin.pegasus2avro.dataprocess import ( @@ -26,10 +27,6 @@ from datahub.utilities.urns.data_process_instance_urn import DataProcessInstanceUrn from datahub.utilities.urns.dataset_urn import DatasetUrn -if TYPE_CHECKING: - from datahub.emitter.kafka_emitter import DatahubKafkaEmitter - from datahub.emitter.rest_emitter import DatahubRestEmitter - class DataProcessInstanceKey(DatahubKey): cluster: str @@ -106,7 +103,7 @@ def start_event_mcp( def emit_process_start( self, - emitter: Union["DatahubRestEmitter", "DatahubKafkaEmitter"], + emitter: Emitter, start_timestamp_millis: int, attempt: Optional[int] = None, emit_template: bool = True, @@ -197,7 +194,7 @@ def end_event_mcp( def emit_process_end( self, - emitter: Union["DatahubRestEmitter", "DatahubKafkaEmitter"], + emitter: Emitter, end_timestamp_millis: int, result: InstanceRunResult, result_type: Optional[str] = None, @@ -207,7 +204,7 @@ def emit_process_end( """ Generate an DataProcessInstance finish event and emits is - :param emitter: (Union[DatahubRestEmitter, DatahubKafkaEmitter]) the datahub emitter to emit generated mcps + :param emitter: (Emitter) the datahub emitter to emit generated mcps :param end_timestamp_millis: (int) the end time of the execution in milliseconds :param result: (InstanceRunResult) The result of the run :param result_type: (string) It identifies the system where the native result comes from like Airflow, Azkaban @@ -261,24 +258,24 @@ def generate_mcp( @staticmethod def _emit_mcp( mcp: MetadataChangeProposalWrapper, - emitter: Union["DatahubRestEmitter", "DatahubKafkaEmitter"], + emitter: Emitter, callback: Optional[Callable[[Exception, str], None]] = None, ) -> None: """ - :param emitter: (Union[DatahubRestEmitter, DatahubKafkaEmitter]) the datahub emitter to emit generated mcps + :param emitter: (Emitter) the datahub emitter to emit generated mcps :param callback: (Optional[Callable[[Exception, str], None]]) the callback method for KafkaEmitter if it is used """ emitter.emit(mcp, callback) def emit( self, - emitter: Union["DatahubRestEmitter", "DatahubKafkaEmitter"], + emitter: Emitter, callback: Optional[Callable[[Exception, str], None]] = None, ) -> None: """ - :param emitter: (Union[DatahubRestEmitter, DatahubKafkaEmitter]) the datahub emitter to emit generated mcps + :param emitter: (Emitter) the datahub emitter to emit generated mcps :param callback: (Optional[Callable[[Exception, str], None]]) the callback method for KafkaEmitter if it is used """ for mcp in self.generate_mcp(): diff --git a/metadata-ingestion/src/datahub/api/entities/dataproduct/dataproduct.py b/metadata-ingestion/src/datahub/api/entities/dataproduct/dataproduct.py index 04f12b4f61d1e..2d9b14ceb2d06 100644 --- a/metadata-ingestion/src/datahub/api/entities/dataproduct/dataproduct.py +++ b/metadata-ingestion/src/datahub/api/entities/dataproduct/dataproduct.py @@ -2,25 +2,15 @@ import time from pathlib import Path -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - Iterable, - List, - Optional, - Tuple, - Union, -) +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union import pydantic from ruamel.yaml import YAML import datahub.emitter.mce_builder as builder from datahub.configuration.common import ConfigModel +from datahub.emitter.generic_emitter import Emitter from datahub.emitter.mcp import MetadataChangeProposalWrapper -from datahub.emitter.rest_emitter import DatahubRestEmitter from datahub.ingestion.graph.client import DataHubGraph from datahub.metadata.schema_classes import ( AuditStampClass, @@ -43,9 +33,6 @@ from datahub.utilities.registries.domain_registry import DomainRegistry from datahub.utilities.urns.urn import Urn -if TYPE_CHECKING: - from datahub.emitter.kafka_emitter import DatahubKafkaEmitter - def patch_list( orig_list: Optional[list], @@ -225,7 +212,6 @@ def _generate_properties_mcp( def generate_mcp( self, upsert: bool ) -> Iterable[Union[MetadataChangeProposalWrapper, MetadataChangeProposalClass]]: - if self._resolved_domain_urn is None: raise Exception( f"Unable to generate MCP-s because we were unable to resolve the domain {self.domain} to an urn." @@ -282,7 +268,7 @@ def generate_mcp( def emit( self, - emitter: Union[DatahubRestEmitter, "DatahubKafkaEmitter"], + emitter: Emitter, upsert: bool, callback: Optional[Callable[[Exception, str], None]] = None, ) -> None: @@ -440,7 +426,6 @@ def patch_yaml( original_dataproduct: DataProduct, output_file: Path, ) -> bool: - update_needed = False if not original_dataproduct._original_yaml_dict: raise Exception("Original Data Product was not loaded from yaml") @@ -523,7 +508,6 @@ def to_yaml( self, file: Path, ) -> None: - with open(file, "w") as fp: yaml = YAML(typ="rt") # default, if not specfied, is 'rt' (round-trip) yaml.indent(mapping=2, sequence=4, offset=2) diff --git a/metadata-ingestion/src/datahub/emitter/generic_emitter.py b/metadata-ingestion/src/datahub/emitter/generic_emitter.py new file mode 100644 index 0000000000000..28138c6182758 --- /dev/null +++ b/metadata-ingestion/src/datahub/emitter/generic_emitter.py @@ -0,0 +1,31 @@ +from typing import Any, Callable, Optional, Union + +from typing_extensions import Protocol + +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.metadata.com.linkedin.pegasus2avro.mxe import ( + MetadataChangeEvent, + MetadataChangeProposal, +) + + +class Emitter(Protocol): + def emit( + self, + item: Union[ + MetadataChangeEvent, + MetadataChangeProposal, + MetadataChangeProposalWrapper, + ], + # NOTE: This signature should have the exception be optional rather than + # required. However, this would be a breaking change that may need + # more careful consideration. + callback: Optional[Callable[[Exception, str], None]] = None, + # TODO: The rest emitter returns timestamps as the return type. For now + # we smooth over that detail using Any, but eventually we should + # standardize on a return type. + ) -> Any: + raise NotImplementedError + + def flush(self) -> None: + pass diff --git a/metadata-ingestion/src/datahub/emitter/kafka_emitter.py b/metadata-ingestion/src/datahub/emitter/kafka_emitter.py index ec0c8f3418a4a..781930011b78f 100644 --- a/metadata-ingestion/src/datahub/emitter/kafka_emitter.py +++ b/metadata-ingestion/src/datahub/emitter/kafka_emitter.py @@ -10,6 +10,7 @@ from datahub.configuration.common import ConfigModel from datahub.configuration.kafka import KafkaProducerConnectionConfig from datahub.configuration.validate_field_rename import pydantic_renamed_field +from datahub.emitter.generic_emitter import Emitter from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.closeable import Closeable from datahub.metadata.schema_classes import ( @@ -55,7 +56,7 @@ def validate_topic_routes(cls, v: Dict[str, str]) -> Dict[str, str]: return v -class DatahubKafkaEmitter(Closeable): +class DatahubKafkaEmitter(Closeable, Emitter): def __init__(self, config: KafkaEmitterConfig): self.config = config schema_registry_conf = { diff --git a/metadata-ingestion/src/datahub/emitter/rest_emitter.py b/metadata-ingestion/src/datahub/emitter/rest_emitter.py index 937e0902d6d8c..afb19df9791af 100644 --- a/metadata-ingestion/src/datahub/emitter/rest_emitter.py +++ b/metadata-ingestion/src/datahub/emitter/rest_emitter.py @@ -4,7 +4,7 @@ import logging import os from json.decoder import JSONDecodeError -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union import requests from deprecated import deprecated @@ -13,6 +13,7 @@ from datahub.cli.cli_utils import get_system_auth from datahub.configuration.common import ConfigurationError, OperationalError +from datahub.emitter.generic_emitter import Emitter from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.request_helper import make_curl_command from datahub.emitter.serialization_helper import pre_json_transform @@ -23,6 +24,9 @@ ) from datahub.metadata.com.linkedin.pegasus2avro.usage import UsageAggregation +if TYPE_CHECKING: + from datahub.ingestion.graph.client import DataHubGraph + logger = logging.getLogger(__name__) _DEFAULT_CONNECT_TIMEOUT_SEC = 30 # 30 seconds should be plenty to connect @@ -42,7 +46,7 @@ ) -class DataHubRestEmitter(Closeable): +class DataHubRestEmitter(Closeable, Emitter): _gms_server: str _token: Optional[str] _session: requests.Session @@ -190,6 +194,11 @@ def test_connection(self) -> dict: message += "\nPlease check your configuration and make sure you are talking to the DataHub GMS (usually :8080) or Frontend GMS API (usually :9002/api/gms)." raise ConfigurationError(message) + def to_graph(self) -> "DataHubGraph": + from datahub.ingestion.graph.client import DataHubGraph + + return DataHubGraph.from_emitter(self) + def emit( self, item: Union[ @@ -198,9 +207,6 @@ def emit( MetadataChangeProposalWrapper, UsageAggregation, ], - # NOTE: This signature should have the exception be optional rather than - # required. However, this would be a breaking change that may need - # more careful consideration. callback: Optional[Callable[[Exception, str], None]] = None, ) -> Tuple[datetime.datetime, datetime.datetime]: start_time = datetime.datetime.now() diff --git a/metadata-ingestion/src/datahub/emitter/synchronized_file_emitter.py b/metadata-ingestion/src/datahub/emitter/synchronized_file_emitter.py new file mode 100644 index 0000000000000..f82882f1a87cc --- /dev/null +++ b/metadata-ingestion/src/datahub/emitter/synchronized_file_emitter.py @@ -0,0 +1,60 @@ +import logging +import pathlib +from typing import Callable, Optional, Union + +import filelock + +from datahub.emitter.generic_emitter import Emitter +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.ingestion.api.closeable import Closeable +from datahub.ingestion.sink.file import write_metadata_file +from datahub.ingestion.source.file import read_metadata_file +from datahub.metadata.com.linkedin.pegasus2avro.mxe import ( + MetadataChangeEvent, + MetadataChangeProposal, +) + +logger = logging.getLogger(__name__) + + +class SynchronizedFileEmitter(Closeable, Emitter): + """ + A multiprocessing-safe emitter that writes to a file. + + This emitter is intended for testing purposes only. It is not performant + because it reads and writes the full file on every emit call to ensure + that the file is always valid JSON. + """ + + def __init__(self, filename: str) -> None: + self._filename = pathlib.Path(filename) + self._lock = filelock.FileLock(self._filename.with_suffix(".lock")) + + def emit( + self, + item: Union[ + MetadataChangeEvent, MetadataChangeProposal, MetadataChangeProposalWrapper + ], + callback: Optional[Callable[[Exception, str], None]] = None, + ) -> None: + with self._lock: + if self._filename.exists(): + metadata = list(read_metadata_file(self._filename)) + else: + metadata = [] + + logger.debug("Emitting metadata: %s", item) + metadata.append(item) + + write_metadata_file(self._filename, metadata) + + def __repr__(self) -> str: + return f"SynchronizedFileEmitter('{self._filename}')" + + def flush(self) -> None: + # No-op. + pass + + def close(self) -> None: + # No-op. + pass diff --git a/metadata-ingestion/src/datahub/ingestion/graph/client.py b/metadata-ingestion/src/datahub/ingestion/graph/client.py index 673ada4f73051..5120d4f643c94 100644 --- a/metadata-ingestion/src/datahub/ingestion/graph/client.py +++ b/metadata-ingestion/src/datahub/ingestion/graph/client.py @@ -138,6 +138,23 @@ def __init__(self, config: DatahubClientConfig) -> None: self.server_id = "missing" logger.debug(f"Failed to get server id due to {e}") + @classmethod + def from_emitter(cls, emitter: DatahubRestEmitter) -> "DataHubGraph": + return cls( + DatahubClientConfig( + server=emitter._gms_server, + token=emitter._token, + timeout_sec=emitter._read_timeout_sec, + retry_status_codes=emitter._retry_status_codes, + retry_max_times=emitter._retry_max_times, + extra_headers=emitter._session.headers, + disable_ssl_verification=emitter._session.verify is False, + # TODO: Support these headers. + # ca_certificate_path=emitter._ca_certificate_path, + # client_certificate_path=emitter._client_certificate_path, + ) + ) + def _send_restli_request(self, method: str, url: str, **kwargs: Any) -> Dict: try: response = self._session.request(method, url, **kwargs) diff --git a/metadata-ingestion/src/datahub/ingestion/source/kafka_connect.py b/metadata-ingestion/src/datahub/ingestion/source/kafka_connect.py index f3344782917ab..5fae0ee5215a3 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/kafka_connect.py +++ b/metadata-ingestion/src/datahub/ingestion/source/kafka_connect.py @@ -28,7 +28,9 @@ ) from datahub.ingestion.api.source import MetadataWorkUnitProcessor, Source from datahub.ingestion.api.workunit import MetadataWorkUnit -from datahub.ingestion.source.sql.sql_common import get_platform_from_sqlalchemy_uri +from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import ( + get_platform_from_sqlalchemy_uri, +) from datahub.ingestion.source.state.stale_entity_removal_handler import ( StaleEntityRemovalHandler, StaleEntityRemovalSourceReport, diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py index 112defe76d957..056be6c2e50ac 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py @@ -1,12 +1,10 @@ import datetime import logging import traceback -from collections import OrderedDict from dataclasses import dataclass, field from typing import ( TYPE_CHECKING, Any, - Callable, Dict, Iterable, List, @@ -103,52 +101,6 @@ MISSING_COLUMN_INFO = "missing column information" -def _platform_alchemy_uri_tester_gen( - platform: str, opt_starts_with: Optional[str] = None -) -> Tuple[str, Callable[[str], bool]]: - return platform, lambda x: x.startswith( - platform if not opt_starts_with else opt_starts_with - ) - - -PLATFORM_TO_SQLALCHEMY_URI_TESTER_MAP: Dict[str, Callable[[str], bool]] = OrderedDict( - [ - _platform_alchemy_uri_tester_gen("athena", "awsathena"), - _platform_alchemy_uri_tester_gen("bigquery"), - _platform_alchemy_uri_tester_gen("clickhouse"), - _platform_alchemy_uri_tester_gen("druid"), - _platform_alchemy_uri_tester_gen("hana"), - _platform_alchemy_uri_tester_gen("hive"), - _platform_alchemy_uri_tester_gen("mongodb"), - _platform_alchemy_uri_tester_gen("mssql"), - _platform_alchemy_uri_tester_gen("mysql"), - _platform_alchemy_uri_tester_gen("oracle"), - _platform_alchemy_uri_tester_gen("pinot"), - _platform_alchemy_uri_tester_gen("presto"), - ( - "redshift", - lambda x: ( - x.startswith(("jdbc:postgres:", "postgresql")) - and x.find("redshift.amazonaws") > 0 - ) - or x.startswith("redshift"), - ), - # Don't move this before redshift. - _platform_alchemy_uri_tester_gen("postgres", "postgresql"), - _platform_alchemy_uri_tester_gen("snowflake"), - _platform_alchemy_uri_tester_gen("trino"), - _platform_alchemy_uri_tester_gen("vertica"), - ] -) - - -def get_platform_from_sqlalchemy_uri(sqlalchemy_uri: str) -> str: - for platform, tester in PLATFORM_TO_SQLALCHEMY_URI_TESTER_MAP.items(): - if tester(sqlalchemy_uri): - return platform - return "external" - - @dataclass class SQLSourceReport(StaleEntityRemovalSourceReport): tables_scanned: int = 0 diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sqlalchemy_uri_mapper.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sqlalchemy_uri_mapper.py new file mode 100644 index 0000000000000..b6a463837228d --- /dev/null +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/sqlalchemy_uri_mapper.py @@ -0,0 +1,47 @@ +from collections import OrderedDict +from typing import Callable, Dict, Optional, Tuple + + +def _platform_alchemy_uri_tester_gen( + platform: str, opt_starts_with: Optional[str] = None +) -> Tuple[str, Callable[[str], bool]]: + return platform, lambda x: x.startswith(opt_starts_with or platform) + + +PLATFORM_TO_SQLALCHEMY_URI_TESTER_MAP: Dict[str, Callable[[str], bool]] = OrderedDict( + [ + _platform_alchemy_uri_tester_gen("athena", "awsathena"), + _platform_alchemy_uri_tester_gen("bigquery"), + _platform_alchemy_uri_tester_gen("clickhouse"), + _platform_alchemy_uri_tester_gen("druid"), + _platform_alchemy_uri_tester_gen("hana"), + _platform_alchemy_uri_tester_gen("hive"), + _platform_alchemy_uri_tester_gen("mongodb"), + _platform_alchemy_uri_tester_gen("mssql"), + _platform_alchemy_uri_tester_gen("mysql"), + _platform_alchemy_uri_tester_gen("oracle"), + _platform_alchemy_uri_tester_gen("pinot"), + _platform_alchemy_uri_tester_gen("presto"), + ( + "redshift", + lambda x: ( + x.startswith(("jdbc:postgres:", "postgresql")) + and x.find("redshift.amazonaws") > 0 + ) + or x.startswith("redshift"), + ), + # Don't move this before redshift. + _platform_alchemy_uri_tester_gen("postgres", "postgresql"), + _platform_alchemy_uri_tester_gen("snowflake"), + _platform_alchemy_uri_tester_gen("sqlite"), + _platform_alchemy_uri_tester_gen("trino"), + _platform_alchemy_uri_tester_gen("vertica"), + ] +) + + +def get_platform_from_sqlalchemy_uri(sqlalchemy_uri: str) -> str: + for platform, tester in PLATFORM_TO_SQLALCHEMY_URI_TESTER_MAP.items(): + if tester(sqlalchemy_uri): + return platform + return "external" diff --git a/metadata-ingestion/src/datahub/ingestion/source/superset.py b/metadata-ingestion/src/datahub/ingestion/source/superset.py index 2a4563439b6ba..14bc4242d2a91 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/superset.py +++ b/metadata-ingestion/src/datahub/ingestion/source/superset.py @@ -21,7 +21,9 @@ ) from datahub.ingestion.api.source import MetadataWorkUnitProcessor, Source from datahub.ingestion.api.workunit import MetadataWorkUnit -from datahub.ingestion.source.sql import sql_common +from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import ( + get_platform_from_sqlalchemy_uri, +) from datahub.ingestion.source.state.stale_entity_removal_handler import ( StaleEntityRemovalHandler, StaleEntityRemovalSourceReport, @@ -202,7 +204,7 @@ def get_platform_from_database_id(self, database_id): sqlalchemy_uri = database_response.get("result", {}).get("sqlalchemy_uri") if sqlalchemy_uri is None: return database_response.get("result", {}).get("backend", "external") - return sql_common.get_platform_from_sqlalchemy_uri(sqlalchemy_uri) + return get_platform_from_sqlalchemy_uri(sqlalchemy_uri) @lru_cache(maxsize=None) def get_datasource_urn_from_id(self, datasource_id): diff --git a/metadata-ingestion/src/datahub/ingestion/source/tableau.py b/metadata-ingestion/src/datahub/ingestion/source/tableau.py index 4cc00a66116e9..6214cba342622 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/tableau.py +++ b/metadata-ingestion/src/datahub/ingestion/source/tableau.py @@ -1179,8 +1179,6 @@ def get_upstream_fields_of_field_in_datasource( def get_upstream_fields_from_custom_sql( self, datasource: dict, datasource_urn: str ) -> List[FineGrainedLineage]: - fine_grained_lineages: List[FineGrainedLineage] = [] - parsed_result = self.parse_custom_sql( datasource=datasource, datasource_urn=datasource_urn, @@ -1194,13 +1192,20 @@ def get_upstream_fields_from_custom_sql( logger.info( f"Failed to extract column level lineage from datasource {datasource_urn}" ) - return fine_grained_lineages + return [] + if parsed_result.debug_info.error: + logger.info( + f"Failed to extract column level lineage from datasource {datasource_urn}: {parsed_result.debug_info.error}" + ) + return [] cll: List[ColumnLineageInfo] = ( parsed_result.column_lineage if parsed_result.column_lineage is not None else [] ) + + fine_grained_lineages: List[FineGrainedLineage] = [] for cll_info in cll: downstream = ( [ diff --git a/metadata-ingestion/src/datahub/integrations/great_expectations/action.py b/metadata-ingestion/src/datahub/integrations/great_expectations/action.py index eabf62a4cda2b..f116550328819 100644 --- a/metadata-ingestion/src/datahub/integrations/great_expectations/action.py +++ b/metadata-ingestion/src/datahub/integrations/great_expectations/action.py @@ -35,7 +35,9 @@ from datahub.cli.cli_utils import get_boolean_env_variable from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.rest_emitter import DatahubRestEmitter -from datahub.ingestion.source.sql.sql_common import get_platform_from_sqlalchemy_uri +from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import ( + get_platform_from_sqlalchemy_uri, +) from datahub.metadata.com.linkedin.pegasus2avro.assertion import ( AssertionInfo, AssertionResult, diff --git a/metadata-ingestion/src/datahub/testing/compare_metadata_json.py b/metadata-ingestion/src/datahub/testing/compare_metadata_json.py index 5c52e1ab4f0b3..54f6a6e984c00 100644 --- a/metadata-ingestion/src/datahub/testing/compare_metadata_json.py +++ b/metadata-ingestion/src/datahub/testing/compare_metadata_json.py @@ -40,6 +40,7 @@ def assert_metadata_files_equal( update_golden: bool, copy_output: bool, ignore_paths: Sequence[str] = (), + ignore_order: bool = True, ) -> None: golden_exists = os.path.isfile(golden_path) @@ -65,7 +66,7 @@ def assert_metadata_files_equal( write_metadata_file(pathlib.Path(temp.name), golden_metadata) golden = load_json_file(temp.name) - diff = diff_metadata_json(output, golden, ignore_paths) + diff = diff_metadata_json(output, golden, ignore_paths, ignore_order=ignore_order) if diff and update_golden: if isinstance(diff, MCPDiff): diff.apply_delta(golden) @@ -91,16 +92,19 @@ def diff_metadata_json( output: MetadataJson, golden: MetadataJson, ignore_paths: Sequence[str] = (), + ignore_order: bool = True, ) -> Union[DeepDiff, MCPDiff]: ignore_paths = (*ignore_paths, *default_exclude_paths, r"root\[\d+].delta_info") try: - golden_map = get_aspects_by_urn(golden) - output_map = get_aspects_by_urn(output) - return MCPDiff.create( - golden=golden_map, - output=output_map, - ignore_paths=ignore_paths, - ) + if ignore_order: + golden_map = get_aspects_by_urn(golden) + output_map = get_aspects_by_urn(output) + return MCPDiff.create( + golden=golden_map, + output=output_map, + ignore_paths=ignore_paths, + ) + # if ignore_order is False, always use DeepDiff except CannotCompareMCPs as e: logger.info(f"{e}, falling back to MCE diff") except AssertionError as e: @@ -111,5 +115,5 @@ def diff_metadata_json( golden, output, exclude_regex_paths=ignore_paths, - ignore_order=True, + ignore_order=ignore_order, ) diff --git a/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py b/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py index f18235af3d1fd..4b3090eaaad31 100644 --- a/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py +++ b/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py @@ -231,6 +231,13 @@ def _table_level_lineage( # In some cases like "MERGE ... then INSERT (col1, col2) VALUES (col1, col2)", # the `this` on the INSERT part isn't a table. if isinstance(expr.this, sqlglot.exp.Table) + } | { + # For CREATE DDL statements, the table name is nested inside + # a Schema object. + _TableName.from_sqlglot_table(expr.this.this) + for expr in statement.find_all(sqlglot.exp.Create) + if isinstance(expr.this, sqlglot.exp.Schema) + and isinstance(expr.this.this, sqlglot.exp.Table) } tables = ( @@ -242,7 +249,7 @@ def _table_level_lineage( - modified # ignore CTEs created in this statement - { - _TableName(database=None, schema=None, table=cte.alias_or_name) + _TableName(database=None, db_schema=None, table=cte.alias_or_name) for cte in statement.find_all(sqlglot.exp.CTE) } ) @@ -906,32 +913,39 @@ def create_lineage_sql_parsed_result( env: str, schema: Optional[str] = None, graph: Optional[DataHubGraph] = None, -) -> Optional["SqlParsingResult"]: - parsed_result: Optional["SqlParsingResult"] = None +) -> SqlParsingResult: + needs_close = False try: - schema_resolver = ( - graph._make_schema_resolver( + if graph: + schema_resolver = graph._make_schema_resolver( platform=platform, platform_instance=platform_instance, env=env, ) - if graph is not None - else SchemaResolver( + else: + needs_close = True + schema_resolver = SchemaResolver( platform=platform, platform_instance=platform_instance, env=env, graph=None, ) - ) - parsed_result = sqlglot_lineage( + return sqlglot_lineage( query, schema_resolver=schema_resolver, default_db=database, default_schema=schema, ) except Exception as e: - logger.debug(f"Fail to prase query {query}", exc_info=e) - logger.warning("Fail to parse custom SQL") - - return parsed_result + return SqlParsingResult( + in_tables=[], + out_tables=[], + column_lineage=None, + debug_info=SqlParsingDebugInfo( + table_error=e, + ), + ) + finally: + if needs_close: + schema_resolver.close() diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_create_table_ddl.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_create_table_ddl.json new file mode 100644 index 0000000000000..4773974545bfa --- /dev/null +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_create_table_ddl.json @@ -0,0 +1,8 @@ +{ + "query_type": "CREATE", + "in_tables": [], + "out_tables": [ + "urn:li:dataset:(urn:li:dataPlatform:sqlite,costs,PROD)" + ], + "column_lineage": null +} \ No newline at end of file diff --git a/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py b/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py index 483c1ac4cc7f9..2a965a9bb1e61 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py +++ b/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py @@ -274,6 +274,21 @@ def test_expand_select_star_basic(): ) +def test_create_table_ddl(): + assert_sql_result( + """ +CREATE TABLE IF NOT EXISTS costs ( + id INTEGER PRIMARY KEY, + month TEXT NOT NULL, + total_cost REAL NOT NULL, + area REAL NOT NULL +) +""", + dialect="sqlite", + expected_file=RESOURCE_DIR / "test_create_table_ddl.json", + ) + + def test_snowflake_column_normalization(): # Technically speaking this is incorrect since the column names are different and both quoted. diff --git a/metadata-ingestion/tests/unit/test_sql_common.py b/metadata-ingestion/tests/unit/test_sql_common.py index 95af0e623e991..808b38192411d 100644 --- a/metadata-ingestion/tests/unit/test_sql_common.py +++ b/metadata-ingestion/tests/unit/test_sql_common.py @@ -4,12 +4,11 @@ import pytest from sqlalchemy.engine.reflection import Inspector -from datahub.ingestion.source.sql.sql_common import ( - PipelineContext, - SQLAlchemySource, +from datahub.ingestion.source.sql.sql_common import PipelineContext, SQLAlchemySource +from datahub.ingestion.source.sql.sql_config import SQLCommonConfig +from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import ( get_platform_from_sqlalchemy_uri, ) -from datahub.ingestion.source.sql.sql_config import SQLCommonConfig class _TestSQLAlchemyConfig(SQLCommonConfig): From e3780c2d75e4dc4dc95e83476d103a4454ee2aae Mon Sep 17 00:00:00 2001 From: Mayuri Nehate <33225191+mayurinehate@users.noreply.github.com> Date: Wed, 4 Oct 2023 16:23:31 +0530 Subject: [PATCH 083/156] =?UTF-8?q?feat(ingest/snowflake):=20initialize=20?= =?UTF-8?q?schema=20resolver=20from=20datahub=20for=20l=E2=80=A6=20(#8903)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/datahub/ingestion/graph/client.py | 8 ++--- .../ingestion/source/bigquery_v2/bigquery.py | 2 +- .../source/snowflake/snowflake_config.py | 4 +-- .../source/snowflake/snowflake_v2.py | 33 ++++++++++++------- .../datahub/ingestion/source/sql_queries.py | 5 ++- .../src/datahub/utilities/sqlglot_lineage.py | 5 +-- 6 files changed, 33 insertions(+), 24 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/graph/client.py b/metadata-ingestion/src/datahub/ingestion/graph/client.py index 5120d4f643c94..ccff677c3a471 100644 --- a/metadata-ingestion/src/datahub/ingestion/graph/client.py +++ b/metadata-ingestion/src/datahub/ingestion/graph/client.py @@ -7,7 +7,7 @@ from dataclasses import dataclass from datetime import datetime from json.decoder import JSONDecodeError -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Tuple, Type +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Type from avro.schema import RecordSchema from deprecated import deprecated @@ -1010,14 +1010,13 @@ def _make_schema_resolver( def initialize_schema_resolver_from_datahub( self, platform: str, platform_instance: Optional[str], env: str - ) -> Tuple["SchemaResolver", Set[str]]: + ) -> "SchemaResolver": logger.info("Initializing schema resolver") schema_resolver = self._make_schema_resolver( platform, platform_instance, env, include_graph=False ) logger.info(f"Fetching schemas for platform {platform}, env {env}") - urns = [] count = 0 with PerfTimer() as timer: for urn, schema_info in self._bulk_fetch_schema_info_by_filter( @@ -1026,7 +1025,6 @@ def initialize_schema_resolver_from_datahub( env=env, ): try: - urns.append(urn) schema_resolver.add_graphql_schema_metadata(urn, schema_info) count += 1 except Exception: @@ -1041,7 +1039,7 @@ def initialize_schema_resolver_from_datahub( ) logger.info("Finished initializing schema resolver") - return schema_resolver, set(urns) + return schema_resolver def parse_sql_lineage( self, diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py index 8a16b1a4a5f6b..f6adbcf033bcc 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py @@ -458,7 +458,7 @@ def _init_schema_resolver(self) -> SchemaResolver: platform=self.platform, platform_instance=self.config.platform_instance, env=self.config.env, - )[0] + ) else: logger.warning( "Failed to load schema info from DataHub as DataHubGraph is missing.", diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_config.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_config.py index 95f6444384408..032bdef178fdf 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_config.py @@ -101,8 +101,8 @@ class SnowflakeV2Config( ) include_view_column_lineage: bool = Field( - default=False, - description="Populates view->view and table->view column lineage.", + default=True, + description="Populates view->view and table->view column lineage using DataHub's sql parser.", ) _check_role_grants_removed = pydantic_removed_field("check_role_grants") diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py index 240e0ffa1a0b6..215116b4c33fb 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py +++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py @@ -301,14 +301,11 @@ def __init__(self, ctx: PipelineContext, config: SnowflakeV2Config): # Caches tables for a single database. Consider moving to disk or S3 when possible. self.db_tables: Dict[str, List[SnowflakeTable]] = {} - self.sql_parser_schema_resolver = SchemaResolver( - platform=self.platform, - platform_instance=self.config.platform_instance, - env=self.config.env, - ) self.view_definitions: FileBackedDict[str] = FileBackedDict() self.add_config_to_report() + self.sql_parser_schema_resolver = self._init_schema_resolver() + @classmethod def create(cls, config_dict: dict, ctx: PipelineContext) -> "Source": config = SnowflakeV2Config.parse_obj(config_dict) @@ -493,6 +490,24 @@ def query(query): return _report + def _init_schema_resolver(self) -> SchemaResolver: + if not self.config.include_technical_schema and self.config.parse_view_ddl: + if self.ctx.graph: + return self.ctx.graph.initialize_schema_resolver_from_datahub( + platform=self.platform, + platform_instance=self.config.platform_instance, + env=self.config.env, + ) + else: + logger.warning( + "Failed to load schema info from DataHub as DataHubGraph is missing.", + ) + return SchemaResolver( + platform=self.platform, + platform_instance=self.config.platform_instance, + env=self.config.env, + ) + def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]: return [ *super().get_workunit_processors(), @@ -764,7 +779,7 @@ def _process_schema( ) self.db_tables[schema_name] = tables - if self.config.include_technical_schema or self.config.parse_view_ddl: + if self.config.include_technical_schema: for table in tables: yield from self._process_table(table, schema_name, db_name) @@ -776,7 +791,7 @@ def _process_schema( if view.view_definition: self.view_definitions[key] = view.view_definition - if self.config.include_technical_schema or self.config.parse_view_ddl: + if self.config.include_technical_schema: for view in views: yield from self._process_view(view, schema_name, db_name) @@ -892,8 +907,6 @@ def _process_table( yield from self._process_tag(tag) yield from self.gen_dataset_workunits(table, schema_name, db_name) - elif self.config.parse_view_ddl: - self.gen_schema_metadata(table, schema_name, db_name) def fetch_sample_data_for_classification( self, table: SnowflakeTable, schema_name: str, db_name: str, dataset_name: str @@ -1004,8 +1017,6 @@ def _process_view( yield from self._process_tag(tag) yield from self.gen_dataset_workunits(view, schema_name, db_name) - elif self.config.parse_view_ddl: - self.gen_schema_metadata(view, schema_name, db_name) def _process_tag(self, tag: SnowflakeTag) -> Iterable[MetadataWorkUnit]: tag_identifier = tag.identifier() diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql_queries.py b/metadata-ingestion/src/datahub/ingestion/source/sql_queries.py index 2fcc93292c2ef..bce4d1ec76e6e 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql_queries.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql_queries.py @@ -103,13 +103,12 @@ def __init__(self, ctx: PipelineContext, config: SqlQueriesSourceConfig): self.builder = SqlParsingBuilder(usage_config=self.config.usage) if self.config.use_schema_resolver: - schema_resolver, urns = self.graph.initialize_schema_resolver_from_datahub( + self.schema_resolver = self.graph.initialize_schema_resolver_from_datahub( platform=self.config.platform, platform_instance=self.config.platform_instance, env=self.config.env, ) - self.schema_resolver = schema_resolver - self.urns = urns + self.urns = self.schema_resolver.get_urns() else: self.schema_resolver = self.graph._make_schema_resolver( platform=self.config.platform, diff --git a/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py b/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py index 4b3090eaaad31..81c43884fdf7d 100644 --- a/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py +++ b/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py @@ -283,6 +283,9 @@ def __init__( shared_connection=shared_conn, ) + def get_urns(self) -> Set[str]: + return set(self._schema_cache.keys()) + def get_urn_for_table(self, table: _TableName, lower: bool = False) -> str: # TODO: Validate that this is the correct 2/3 layer hierarchy for the platform. @@ -397,8 +400,6 @@ def convert_graphql_schema_metadata_to_info( ) } - # TODO add a method to load all from graphql - def close(self) -> None: self._schema_cache.close() From 13508a9d888df519a389b6bd187b5f745772627b Mon Sep 17 00:00:00 2001 From: Upendra Rao Vedullapalli Date: Wed, 4 Oct 2023 15:20:51 +0200 Subject: [PATCH 084/156] feat(bigquery): excluding projects without any datasets from ingestion (#8535) Co-authored-by: Upendra Vedullapalli Co-authored-by: Andrew Sikowitz --- .../ingestion/source/bigquery_v2/bigquery.py | 19 +++++-- .../source/bigquery_v2/bigquery_config.py | 5 ++ .../source/bigquery_v2/bigquery_report.py | 2 + .../tests/unit/test_bigquery_source.py | 53 ++++++++++++++++++- 4 files changed, 72 insertions(+), 7 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py index f6adbcf033bcc..fee181864a2d6 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py @@ -600,9 +600,6 @@ def _process_project( db_views: Dict[str, List[BigqueryView]] = {} project_id = bigquery_project.id - - yield from self.gen_project_id_containers(project_id) - try: bigquery_project.datasets = ( self.bigquery_data_dictionary.get_datasets_for_project_id(project_id) @@ -619,11 +616,23 @@ def _process_project( return None if len(bigquery_project.datasets) == 0: - logger.warning( - f"No dataset found in {project_id}. Either there are no datasets in this project or missing bigquery.datasets.get permission. You can assign predefined roles/bigquery.metadataViewer role to your service account." + more_info = ( + "Either there are no datasets in this project or missing bigquery.datasets.get permission. " + "You can assign predefined roles/bigquery.metadataViewer role to your service account." ) + if self.config.exclude_empty_projects: + self.report.report_dropped(project_id) + warning_message = f"Excluded project '{project_id}' since no were datasets found. {more_info}" + else: + yield from self.gen_project_id_containers(project_id) + warning_message = ( + f"No datasets found in project '{project_id}'. {more_info}" + ) + logger.warning(warning_message) return + yield from self.gen_project_id_containers(project_id) + self.report.num_project_datasets_to_scan[project_id] = len( bigquery_project.datasets ) diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py index 3b06a4699c566..483355a85ac05 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py @@ -265,6 +265,11 @@ def validate_column_lineage(cls, v: bool, values: Dict[str, Any]) -> bool: description="Maximum number of entries for the in-memory caches of FileBacked data structures.", ) + exclude_empty_projects: bool = Field( + default=False, + description="Option to exclude empty projects from being ingested.", + ) + @root_validator(pre=False) def profile_default_settings(cls, values: Dict) -> Dict: # Extra default SQLAlchemy option for better connection pooling and threading. diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_report.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_report.py index 661589a0c58e5..9d92b011ee285 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_report.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_report.py @@ -122,6 +122,8 @@ class BigQueryV2Report(ProfilingSqlReport, IngestionStageReport, BaseTimeWindowR usage_state_size: Optional[str] = None + exclude_empty_projects: Optional[bool] = None + schema_api_perf: BigQuerySchemaApiPerfReport = field( default_factory=BigQuerySchemaApiPerfReport ) diff --git a/metadata-ingestion/tests/unit/test_bigquery_source.py b/metadata-ingestion/tests/unit/test_bigquery_source.py index 4fc6c31626ba8..e9e91361f49f4 100644 --- a/metadata-ingestion/tests/unit/test_bigquery_source.py +++ b/metadata-ingestion/tests/unit/test_bigquery_source.py @@ -3,13 +3,14 @@ import os from datetime import datetime, timedelta, timezone from types import SimpleNamespace -from typing import Any, Dict, Optional, cast +from typing import Any, Dict, List, Optional, cast from unittest.mock import MagicMock, Mock, patch import pytest from google.api_core.exceptions import GoogleAPICallError from google.cloud.bigquery.table import Row, TableListItem +from datahub.configuration.common import AllowDenyPattern from datahub.ingestion.api.common import PipelineContext from datahub.ingestion.source.bigquery_v2.bigquery import BigqueryV2Source from datahub.ingestion.source.bigquery_v2.bigquery_audit import ( @@ -17,9 +18,13 @@ BigqueryTableIdentifier, BigQueryTableRef, ) -from datahub.ingestion.source.bigquery_v2.bigquery_config import BigQueryV2Config +from datahub.ingestion.source.bigquery_v2.bigquery_config import ( + BigQueryConnectionConfig, + BigQueryV2Config, +) from datahub.ingestion.source.bigquery_v2.bigquery_report import BigQueryV2Report from datahub.ingestion.source.bigquery_v2.bigquery_schema import ( + BigqueryDataset, BigqueryProject, BigQuerySchemaApi, BigqueryView, @@ -854,3 +859,47 @@ def test_get_table_name(full_table_name: str, datahub_full_table_name: str) -> N BigqueryTableIdentifier.from_string_name(full_table_name).get_table_name() == datahub_full_table_name ) + + +def test_default_config_for_excluding_projects_and_datasets(): + config = BigQueryV2Config.parse_obj({}) + assert config.exclude_empty_projects is False + config = BigQueryV2Config.parse_obj({"exclude_empty_projects": True}) + assert config.exclude_empty_projects + + +@patch.object(BigQueryConnectionConfig, "get_bigquery_client", new=lambda self: None) +@patch.object(BigQuerySchemaApi, "get_datasets_for_project_id") +def test_excluding_empty_projects_from_ingestion( + get_datasets_for_project_id_mock, +): + project_id_with_datasets = "project-id-with-datasets" + project_id_without_datasets = "project-id-without-datasets" + + def get_datasets_for_project_id_side_effect( + project_id: str, + ) -> List[BigqueryDataset]: + return ( + [] + if project_id == project_id_without_datasets + else [BigqueryDataset("some-dataset")] + ) + + get_datasets_for_project_id_mock.side_effect = ( + get_datasets_for_project_id_side_effect + ) + + base_config = { + "project_ids": [project_id_with_datasets, project_id_without_datasets], + "schema_pattern": AllowDenyPattern(deny=[".*"]), + "include_usage_statistics": False, + "include_table_lineage": False, + } + + config = BigQueryV2Config.parse_obj(base_config) + source = BigqueryV2Source(config=config, ctx=PipelineContext(run_id="test-1")) + assert len({wu.metadata.entityUrn for wu in source.get_workunits()}) == 2 # type: ignore + + config = BigQueryV2Config.parse_obj({**base_config, "exclude_empty_projects": True}) + source = BigqueryV2Source(config=config, ctx=PipelineContext(run_id="test-2")) + assert len({wu.metadata.entityUrn for wu in source.get_workunits()}) == 1 # type: ignore From d3346a04e486fa098129b626e61013cab4f69350 Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Wed, 4 Oct 2023 10:22:45 -0400 Subject: [PATCH 085/156] feat(ingest/unity): Ingest notebooks and their lineage (#8940) --- .../sources/databricks/unity-catalog_pre.md | 1 + metadata-ingestion/setup.py | 2 +- .../src/datahub/emitter/mcp_builder.py | 12 ++ .../ingestion/source/common/subtypes.py | 3 + .../datahub/ingestion/source/unity/config.py | 9 +- .../datahub/ingestion/source/unity/proxy.py | 89 +++++++---- .../ingestion/source/unity/proxy_types.py | 45 +++++- .../datahub/ingestion/source/unity/report.py | 8 +- .../datahub/ingestion/source/unity/source.py | 148 ++++++++++++++---- .../datahub/ingestion/source/unity/usage.py | 12 +- 10 files changed, 257 insertions(+), 72 deletions(-) diff --git a/metadata-ingestion/docs/sources/databricks/unity-catalog_pre.md b/metadata-ingestion/docs/sources/databricks/unity-catalog_pre.md index 2be8846b87bea..ae2883343d7e8 100644 --- a/metadata-ingestion/docs/sources/databricks/unity-catalog_pre.md +++ b/metadata-ingestion/docs/sources/databricks/unity-catalog_pre.md @@ -13,6 +13,7 @@ * Ownership of or `SELECT` privilege on any tables and views you want to ingest * [Ownership documentation](https://docs.databricks.com/data-governance/unity-catalog/manage-privileges/ownership.html) * [Privileges documentation](https://docs.databricks.com/data-governance/unity-catalog/manage-privileges/privileges.html) + + To ingest your workspace's notebooks and respective lineage, your service principal must have `CAN_READ` privileges on the folders containing the notebooks you want to ingest: [guide](https://docs.databricks.com/en/security/auth-authz/access-control/workspace-acl.html#folder-permissions). + To `include_usage_statistics` (enabled by default), your service principal must have `CAN_MANAGE` permissions on any SQL Warehouses you want to ingest: [guide](https://docs.databricks.com/security/auth-authz/access-control/sql-endpoint-acl.html). + To ingest `profiling` information with `call_analyze` (enabled by default), your service principal must have ownership or `MODIFY` privilege on any tables you want to profile. * Alternatively, you can run [ANALYZE TABLE](https://docs.databricks.com/sql/language-manual/sql-ref-syntax-aux-analyze-table.html) yourself on any tables you want to profile, then set `call_analyze` to `false`. diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py index 34afa8cdb39a4..fe8e3be4632c4 100644 --- a/metadata-ingestion/setup.py +++ b/metadata-ingestion/setup.py @@ -250,7 +250,7 @@ databricks = { # 0.1.11 appears to have authentication issues with azure databricks - "databricks-sdk>=0.1.1, != 0.1.11", + "databricks-sdk>=0.9.0", "pyspark", "requests", } diff --git a/metadata-ingestion/src/datahub/emitter/mcp_builder.py b/metadata-ingestion/src/datahub/emitter/mcp_builder.py index 844a29f1c78a3..7419577b367aa 100644 --- a/metadata-ingestion/src/datahub/emitter/mcp_builder.py +++ b/metadata-ingestion/src/datahub/emitter/mcp_builder.py @@ -9,6 +9,7 @@ make_container_urn, make_data_platform_urn, make_dataplatform_instance_urn, + make_dataset_urn_with_platform_instance, ) from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.workunit import MetadataWorkUnit @@ -125,6 +126,17 @@ class BucketKey(ContainerKey): bucket_name: str +class NotebookKey(DatahubKey): + notebook_id: int + platform: str + instance: Optional[str] + + def as_urn(self) -> str: + return make_dataset_urn_with_platform_instance( + platform=self.platform, platform_instance=self.instance, name=self.guid() + ) + + class DatahubKeyJSONEncoder(json.JSONEncoder): # overload method default def default(self, obj: Any) -> Any: diff --git a/metadata-ingestion/src/datahub/ingestion/source/common/subtypes.py b/metadata-ingestion/src/datahub/ingestion/source/common/subtypes.py index a2d89d26112f4..741b4789bef21 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/common/subtypes.py +++ b/metadata-ingestion/src/datahub/ingestion/source/common/subtypes.py @@ -16,6 +16,9 @@ class DatasetSubTypes(str, Enum): SALESFORCE_STANDARD_OBJECT = "Object" POWERBI_DATASET_TABLE = "PowerBI Dataset Table" + # TODO: Create separate entity... + NOTEBOOK = "Notebook" + class DatasetContainerSubTypes(str, Enum): # Generic SubTypes diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py index 94ff755e3b254..a49c789a82f27 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py @@ -127,11 +127,16 @@ class UnityCatalogSourceConfig( description='Attach domains to catalogs, schemas or tables during ingestion using regex patterns. Domain key can be a guid like *urn:li:domain:ec428203-ce86-4db3-985d-5a8ee6df32ba* or a string like "Marketing".) If you provide strings, then datahub will attempt to resolve this name to a guid, and will error out if this fails. There can be multiple domain keys specified.', ) - include_table_lineage: Optional[bool] = pydantic.Field( + include_table_lineage: bool = pydantic.Field( default=True, description="Option to enable/disable lineage generation.", ) + include_notebooks: bool = pydantic.Field( + default=False, + description="Ingest notebooks, represented as DataHub datasets.", + ) + include_ownership: bool = pydantic.Field( default=False, description="Option to enable/disable ownership generation for metastores, catalogs, schemas, and tables.", @@ -141,7 +146,7 @@ class UnityCatalogSourceConfig( "include_table_ownership", "include_ownership" ) - include_column_lineage: Optional[bool] = pydantic.Field( + include_column_lineage: bool = pydantic.Field( default=True, description="Option to enable/disable lineage generation. Currently we have to call a rest call per column to get column level lineage due to the Databrick api which can slow down ingestion. ", ) diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py index e92f4ff07b1ad..2401f1c3d163c 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py @@ -23,6 +23,7 @@ QueryStatementType, QueryStatus, ) +from databricks.sdk.service.workspace import ObjectType import datahub from datahub.ingestion.source.unity.proxy_profiling import ( @@ -33,6 +34,7 @@ Catalog, Column, Metastore, + Notebook, Query, Schema, ServicePrincipal, @@ -137,6 +139,21 @@ def service_principals(self) -> Iterable[ServicePrincipal]: for principal in self._workspace_client.service_principals.list(): yield self._create_service_principal(principal) + def workspace_notebooks(self) -> Iterable[Notebook]: + for obj in self._workspace_client.workspace.list("/", recursive=True): + if obj.object_type == ObjectType.NOTEBOOK: + yield Notebook( + id=obj.object_id, + path=obj.path, + language=obj.language, + created_at=datetime.fromtimestamp( + obj.created_at / 1000, tz=timezone.utc + ), + modified_at=datetime.fromtimestamp( + obj.modified_at / 1000, tz=timezone.utc + ), + ) + def query_history( self, start_time: datetime, @@ -153,7 +170,7 @@ def query_history( "start_time_ms": start_time.timestamp() * 1000, "end_time_ms": end_time.timestamp() * 1000, }, - "statuses": [QueryStatus.FINISHED.value], + "statuses": [QueryStatus.FINISHED], "statement_types": [typ.value for typ in ALLOWED_STATEMENT_TYPES], } ) @@ -196,61 +213,75 @@ def _query_history( method, path, body={**body, "page_token": response["next_page_token"]} ) - def list_lineages_by_table(self, table_name: str) -> dict: + def list_lineages_by_table( + self, table_name: str, include_entity_lineage: bool + ) -> dict: """List table lineage by table name.""" return self._workspace_client.api_client.do( method="GET", - path="/api/2.0/lineage-tracking/table-lineage/get", - body={"table_name": table_name}, + path="/api/2.0/lineage-tracking/table-lineage", + body={ + "table_name": table_name, + "include_entity_lineage": include_entity_lineage, + }, ) def list_lineages_by_column(self, table_name: str, column_name: str) -> dict: """List column lineage by table name and column name.""" return self._workspace_client.api_client.do( "GET", - "/api/2.0/lineage-tracking/column-lineage/get", + "/api/2.0/lineage-tracking/column-lineage", body={"table_name": table_name, "column_name": column_name}, ) - def table_lineage(self, table: Table) -> None: + def table_lineage( + self, table: Table, include_entity_lineage: bool + ) -> Optional[dict]: # Lineage endpoint doesn't exists on 2.1 version try: response: dict = self.list_lineages_by_table( - table_name=f"{table.schema.catalog.name}.{table.schema.name}.{table.name}" + table_name=table.ref.qualified_table_name, + include_entity_lineage=include_entity_lineage, ) - table.upstreams = { - TableReference( - table.schema.catalog.metastore.id, - item["catalog_name"], - item["schema_name"], - item["name"], - ): {} - for item in response.get("upstream_tables", []) - } + + for item in response.get("upstreams") or []: + if "tableInfo" in item: + table_ref = TableReference.create_from_lineage( + item["tableInfo"], table.schema.catalog.metastore.id + ) + if table_ref: + table.upstreams[table_ref] = {} + for notebook in item.get("notebookInfos") or []: + table.upstream_notebooks.add(notebook["notebook_id"]) + + for item in response.get("downstreams") or []: + for notebook in item.get("notebookInfos") or []: + table.downstream_notebooks.add(notebook["notebook_id"]) + + return response except Exception as e: logger.error(f"Error getting lineage: {e}") + return None - def get_column_lineage(self, table: Table) -> None: + def get_column_lineage(self, table: Table, include_entity_lineage: bool) -> None: try: - table_lineage_response: dict = self.list_lineages_by_table( - table_name=f"{table.schema.catalog.name}.{table.schema.name}.{table.name}" + table_lineage = self.table_lineage( + table, include_entity_lineage=include_entity_lineage ) - if table_lineage_response: + if table_lineage: for column in table.columns: response: dict = self.list_lineages_by_column( - table_name=f"{table.schema.catalog.name}.{table.schema.name}.{table.name}", + table_name=table.ref.qualified_table_name, column_name=column.name, ) for item in response.get("upstream_cols", []): - table_ref = TableReference( - table.schema.catalog.metastore.id, - item["catalog_name"], - item["schema_name"], - item["table_name"], + table_ref = TableReference.create_from_lineage( + item, table.schema.catalog.metastore.id ) - table.upstreams.setdefault(table_ref, {}).setdefault( - column.name, [] - ).append(item["name"]) + if table_ref: + table.upstreams.setdefault(table_ref, {}).setdefault( + column.name, [] + ).append(item["name"]) except Exception as e: logger.error(f"Error getting lineage: {e}") diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py index 2b943d8c98e7d..d57f20245913f 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py @@ -1,8 +1,10 @@ # Supported types are available at # https://api-docs.databricks.com/rest/latest/unity-catalog-api-specification-2-1.html?_ga=2.151019001.1795147704.1666247755-2119235717.1666247755 +import dataclasses +import logging from dataclasses import dataclass, field from datetime import datetime -from typing import Dict, List, Optional +from typing import Dict, FrozenSet, List, Optional, Set from databricks.sdk.service.catalog import ( CatalogType, @@ -11,6 +13,7 @@ TableType, ) from databricks.sdk.service.sql import QueryStatementType +from databricks.sdk.service.workspace import Language from datahub.metadata.schema_classes import ( ArrayTypeClass, @@ -26,6 +29,8 @@ TimeTypeClass, ) +logger = logging.getLogger(__name__) + DATA_TYPE_REGISTRY: dict = { ColumnTypeName.BOOLEAN: BooleanTypeClass, ColumnTypeName.BYTE: BytesTypeClass, @@ -66,6 +71,9 @@ ALLOWED_STATEMENT_TYPES = {*OPERATION_STATEMENT_TYPES.keys(), QueryStatementType.SELECT} +NotebookId = int + + @dataclass class CommonProperty: id: str @@ -136,6 +144,19 @@ def create(cls, table: "Table") -> "TableReference": table.name, ) + @classmethod + def create_from_lineage(cls, d: dict, metastore: str) -> Optional["TableReference"]: + try: + return cls( + metastore, + d["catalog_name"], + d["schema_name"], + d.get("table_name", d["name"]), # column vs table query output + ) + except Exception as e: + logger.warning(f"Failed to create TableReference from {d}: {e}") + return None + def __str__(self) -> str: return f"{self.metastore}.{self.catalog}.{self.schema}.{self.table}" @@ -166,6 +187,8 @@ class Table(CommonProperty): view_definition: Optional[str] properties: Dict[str, str] upstreams: Dict[TableReference, Dict[str, List[str]]] = field(default_factory=dict) + upstream_notebooks: Set[NotebookId] = field(default_factory=set) + downstream_notebooks: Set[NotebookId] = field(default_factory=set) ref: TableReference = field(init=False) @@ -228,3 +251,23 @@ def __bool__(self): self.max is not None, ) ) + + +@dataclass +class Notebook: + id: NotebookId + path: str + language: Language + created_at: datetime + modified_at: datetime + + upstreams: FrozenSet[TableReference] = field(default_factory=frozenset) + + @classmethod + def add_upstream(cls, upstream: TableReference, notebook: "Notebook") -> "Notebook": + return cls( + **{ # type: ignore + **dataclasses.asdict(notebook), + "upstreams": frozenset([*notebook.upstreams, upstream]), + } + ) diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/report.py b/metadata-ingestion/src/datahub/ingestion/source/unity/report.py index 8382b31a56add..808172a136bb3 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/report.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/report.py @@ -5,21 +5,23 @@ from datahub.ingestion.source.state.stale_entity_removal_handler import ( StaleEntityRemovalSourceReport, ) +from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport from datahub.utilities.lossy_collections import LossyDict, LossyList @dataclass -class UnityCatalogReport(StaleEntityRemovalSourceReport): +class UnityCatalogReport(IngestionStageReport, StaleEntityRemovalSourceReport): metastores: EntityFilterReport = EntityFilterReport.field(type="metastore") catalogs: EntityFilterReport = EntityFilterReport.field(type="catalog") schemas: EntityFilterReport = EntityFilterReport.field(type="schema") tables: EntityFilterReport = EntityFilterReport.field(type="table/view") table_profiles: EntityFilterReport = EntityFilterReport.field(type="table profile") + notebooks: EntityFilterReport = EntityFilterReport.field(type="notebook") num_queries: int = 0 num_queries_dropped_parse_failure: int = 0 - num_queries_dropped_missing_table: int = 0 # Can be due to pattern filter - num_queries_dropped_duplicate_table: int = 0 + num_queries_missing_table: int = 0 # Can be due to pattern filter + num_queries_duplicate_table: int = 0 num_queries_parsed_by_spark_plan: int = 0 # Distinguish from Operations emitted for created / updated timestamps diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py index 493acb939c3bb..f2da1aece9fd4 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py @@ -2,7 +2,7 @@ import re import time from datetime import timedelta -from typing import Dict, Iterable, List, Optional, Set +from typing import Dict, Iterable, List, Optional, Set, Union from urllib.parse import urljoin from datahub.emitter.mce_builder import ( @@ -18,6 +18,7 @@ CatalogKey, ContainerKey, MetastoreKey, + NotebookKey, UnitySchemaKey, add_dataset_to_container, gen_containers, @@ -56,6 +57,8 @@ Catalog, Column, Metastore, + Notebook, + NotebookId, Schema, ServicePrincipal, Table, @@ -69,6 +72,7 @@ ViewProperties, ) from datahub.metadata.schema_classes import ( + BrowsePathsClass, DataPlatformInstanceClass, DatasetLineageTypeClass, DatasetPropertiesClass, @@ -88,6 +92,7 @@ UpstreamClass, UpstreamLineageClass, ) +from datahub.utilities.file_backed_collections import FileBackedDict from datahub.utilities.hive_schema_to_avro import get_schema_fields_for_hive_column from datahub.utilities.registries.domain_registry import DomainRegistry @@ -157,6 +162,7 @@ def __init__(self, ctx: PipelineContext, config: UnityCatalogSourceConfig): # Global set of table refs self.table_refs: Set[TableReference] = set() self.view_refs: Set[TableReference] = set() + self.notebooks: FileBackedDict[Notebook] = FileBackedDict() @staticmethod def test_connection(config_dict: dict) -> TestConnectionReport: @@ -176,6 +182,7 @@ def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]: ] def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]: + self.report.report_ingestion_stage_start("Start warehouse") wait_on_warehouse = None if self.config.is_profiling_enabled(): # Can take several minutes, so start now and wait later @@ -187,10 +194,23 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]: ) return + self.report.report_ingestion_stage_start("Ingest service principals") self.build_service_principal_map() + if self.config.include_notebooks: + self.report.report_ingestion_stage_start("Ingest notebooks") + yield from self.process_notebooks() + yield from self.process_metastores() + if self.config.include_notebooks: + self.report.report_ingestion_stage_start("Notebook lineage") + for notebook in self.notebooks.values(): + wu = self._gen_notebook_lineage(notebook) + if wu: + yield wu + if self.config.include_usage_statistics: + self.report.report_ingestion_stage_start("Ingest usage") usage_extractor = UnityCatalogUsageExtractor( config=self.config, report=self.report, @@ -203,6 +223,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]: ) if self.config.is_profiling_enabled(): + self.report.report_ingestion_stage_start("Wait on warehouse") assert wait_on_warehouse timeout = timedelta(seconds=self.config.profiling.max_wait_secs) wait_on_warehouse.result(timeout) @@ -212,6 +233,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]: self.unity_catalog_api_proxy, self.gen_dataset_urn, ) + self.report.report_ingestion_stage_start("Profiling") yield from profiling_extractor.get_workunits(self.table_refs) def build_service_principal_map(self) -> None: @@ -223,6 +245,56 @@ def build_service_principal_map(self) -> None: "service-principals", f"Unable to fetch service principals: {e}" ) + def process_notebooks(self) -> Iterable[MetadataWorkUnit]: + for notebook in self.unity_catalog_api_proxy.workspace_notebooks(): + self.notebooks[str(notebook.id)] = notebook + yield from self._gen_notebook_aspects(notebook) + + def _gen_notebook_aspects(self, notebook: Notebook) -> Iterable[MetadataWorkUnit]: + mcps = MetadataChangeProposalWrapper.construct_many( + entityUrn=self.gen_notebook_urn(notebook), + aspects=[ + DatasetPropertiesClass( + name=notebook.path.rsplit("/", 1)[-1], + customProperties={ + "path": notebook.path, + "language": notebook.language.value, + }, + externalUrl=urljoin( + self.config.workspace_url, f"#notebook/{notebook.id}" + ), + created=TimeStampClass(int(notebook.created_at.timestamp() * 1000)), + lastModified=TimeStampClass( + int(notebook.modified_at.timestamp() * 1000) + ), + ), + SubTypesClass(typeNames=[DatasetSubTypes.NOTEBOOK]), + BrowsePathsClass(paths=notebook.path.split("/")), + # TODO: Add DPI aspect + ], + ) + for mcp in mcps: + yield mcp.as_workunit() + + self.report.notebooks.processed(notebook.path) + + def _gen_notebook_lineage(self, notebook: Notebook) -> Optional[MetadataWorkUnit]: + if not notebook.upstreams: + return None + + return MetadataChangeProposalWrapper( + entityUrn=self.gen_notebook_urn(notebook), + aspect=UpstreamLineageClass( + upstreams=[ + UpstreamClass( + dataset=self.gen_dataset_urn(upstream_ref), + type=DatasetLineageTypeClass.COPY, + ) + for upstream_ref in notebook.upstreams + ] + ), + ).as_workunit() + def process_metastores(self) -> Iterable[MetadataWorkUnit]: metastore = self.unity_catalog_api_proxy.assigned_metastore() yield from self.gen_metastore_containers(metastore) @@ -247,6 +319,7 @@ def process_schemas(self, catalog: Catalog) -> Iterable[MetadataWorkUnit]: self.report.schemas.dropped(schema.id) continue + self.report.report_ingestion_stage_start(f"Ingest schema {schema.id}") yield from self.gen_schema_containers(schema) yield from self.process_tables(schema) @@ -282,13 +355,21 @@ def process_table(self, table: Table, schema: Schema) -> Iterable[MetadataWorkUn ownership = self._create_table_ownership_aspect(table) data_platform_instance = self._create_data_platform_instance_aspect(table) - lineage: Optional[UpstreamLineageClass] = None if self.config.include_column_lineage: - self.unity_catalog_api_proxy.get_column_lineage(table) - lineage = self._generate_column_lineage_aspect(dataset_urn, table) + self.unity_catalog_api_proxy.get_column_lineage( + table, include_entity_lineage=self.config.include_notebooks + ) elif self.config.include_table_lineage: - self.unity_catalog_api_proxy.table_lineage(table) - lineage = self._generate_lineage_aspect(dataset_urn, table) + self.unity_catalog_api_proxy.table_lineage( + table, include_entity_lineage=self.config.include_notebooks + ) + lineage = self._generate_lineage_aspect(dataset_urn, table) + + if self.config.include_notebooks: + for notebook_id in table.downstream_notebooks: + self.notebooks[str(notebook_id)] = Notebook.add_upstream( + table.ref, self.notebooks[str(notebook_id)] + ) yield from [ mcp.as_workunit() @@ -308,7 +389,7 @@ def process_table(self, table: Table, schema: Schema) -> Iterable[MetadataWorkUn ) ] - def _generate_column_lineage_aspect( + def _generate_lineage_aspect( self, dataset_urn: str, table: Table ) -> Optional[UpstreamLineageClass]: upstreams: List[UpstreamClass] = [] @@ -318,6 +399,7 @@ def _generate_column_lineage_aspect( ): upstream_urn = self.gen_dataset_urn(upstream_ref) + # Should be empty if config.include_column_lineage is False finegrained_lineages.extend( FineGrainedLineage( upstreamType=FineGrainedLineageUpstreamType.FIELD_SET, @@ -331,38 +413,28 @@ def _generate_column_lineage_aspect( for d_col, u_cols in sorted(downstream_to_upstream_cols.items()) ) - upstream_table = UpstreamClass( - upstream_urn, - DatasetLineageTypeClass.TRANSFORMED, - ) - upstreams.append(upstream_table) - - if upstreams: - return UpstreamLineageClass( - upstreams=upstreams, fineGrainedLineages=finegrained_lineages - ) - else: - return None - - def _generate_lineage_aspect( - self, dataset_urn: str, table: Table - ) -> Optional[UpstreamLineageClass]: - upstreams: List[UpstreamClass] = [] - for upstream in sorted(table.upstreams.keys()): - upstream_urn = make_dataset_urn_with_platform_instance( - self.platform, - f"{table.schema.catalog.metastore.id}.{upstream}", - self.platform_instance_name, + upstreams.append( + UpstreamClass( + dataset=upstream_urn, + type=DatasetLineageTypeClass.TRANSFORMED, + ) ) - upstream_table = UpstreamClass( - upstream_urn, - DatasetLineageTypeClass.TRANSFORMED, + for notebook in table.upstream_notebooks: + upstreams.append( + UpstreamClass( + dataset=self.gen_notebook_urn(notebook), + type=DatasetLineageTypeClass.TRANSFORMED, + ) ) - upstreams.append(upstream_table) if upstreams: - return UpstreamLineageClass(upstreams=upstreams) + return UpstreamLineageClass( + upstreams=upstreams, + fineGrainedLineages=finegrained_lineages + if self.config.include_column_lineage + else None, + ) else: return None @@ -389,6 +461,14 @@ def gen_dataset_urn(self, table_ref: TableReference) -> str: name=str(table_ref), ) + def gen_notebook_urn(self, notebook: Union[Notebook, NotebookId]) -> str: + notebook_id = notebook.id if isinstance(notebook, Notebook) else notebook + return NotebookKey( + notebook_id=notebook_id, + platform=self.platform, + instance=self.config.platform_instance, + ).as_urn() + def gen_schema_containers(self, schema: Schema) -> Iterable[MetadataWorkUnit]: domain_urn = self._gen_domain_urn(f"{schema.catalog.name}.{schema.name}") diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/usage.py b/metadata-ingestion/src/datahub/ingestion/source/unity/usage.py index 49f56b46fb012..ab21c1a318659 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/usage.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/usage.py @@ -214,12 +214,15 @@ def _resolve_tables( self, tables: List[str], table_map: TableMap ) -> List[TableReference]: """Resolve tables to TableReferences, filtering out unrecognized or unresolvable table names.""" + + missing_table = False + duplicate_table = False output = [] for table in tables: table = str(table) if table not in table_map: logger.debug(f"Dropping query with unrecognized table: {table}") - self.report.num_queries_dropped_missing_table += 1 + missing_table = True else: refs = table_map[table] if len(refs) == 1: @@ -228,6 +231,11 @@ def _resolve_tables( logger.warning( f"Could not resolve table ref for {table}: {len(refs)} duplicates." ) - self.report.num_queries_dropped_duplicate_table += 1 + duplicate_table = True + + if missing_table: + self.report.num_queries_missing_table += 1 + if duplicate_table: + self.report.num_queries_duplicate_table += 1 return output From 301d3e6b1ccffaf946f128766578faddbc7ac44e Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Wed, 4 Oct 2023 10:23:13 -0400 Subject: [PATCH 086/156] test(ingest/unity): Add Unity Catalog memory performance testing (#8932) --- .../ingestion/source/unity/proxy_types.py | 1 - .../tests/performance/bigquery/__init__.py | 0 .../bigquery_events.py} | 0 .../{ => bigquery}/test_bigquery_usage.py | 22 +-- .../tests/performance/data_generation.py | 53 ++++- .../tests/performance/data_model.py | 31 ++- .../tests/performance/databricks/__init__.py | 0 .../performance/databricks/test_unity.py | 71 +++++++ .../databricks/unity_proxy_mock.py | 183 ++++++++++++++++++ .../tests/performance/helpers.py | 21 ++ .../tests/unit/test_bigquery_usage.py | 7 +- 11 files changed, 356 insertions(+), 33 deletions(-) create mode 100644 metadata-ingestion/tests/performance/bigquery/__init__.py rename metadata-ingestion/tests/performance/{bigquery.py => bigquery/bigquery_events.py} (100%) rename metadata-ingestion/tests/performance/{ => bigquery}/test_bigquery_usage.py (80%) create mode 100644 metadata-ingestion/tests/performance/databricks/__init__.py create mode 100644 metadata-ingestion/tests/performance/databricks/test_unity.py create mode 100644 metadata-ingestion/tests/performance/databricks/unity_proxy_mock.py create mode 100644 metadata-ingestion/tests/performance/helpers.py diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py index d57f20245913f..54ac2e90d7c7e 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py @@ -175,7 +175,6 @@ class Table(CommonProperty): columns: List[Column] storage_location: Optional[str] data_source_format: Optional[DataSourceFormat] - comment: Optional[str] table_type: TableType owner: Optional[str] generation: Optional[int] diff --git a/metadata-ingestion/tests/performance/bigquery/__init__.py b/metadata-ingestion/tests/performance/bigquery/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/metadata-ingestion/tests/performance/bigquery.py b/metadata-ingestion/tests/performance/bigquery/bigquery_events.py similarity index 100% rename from metadata-ingestion/tests/performance/bigquery.py rename to metadata-ingestion/tests/performance/bigquery/bigquery_events.py diff --git a/metadata-ingestion/tests/performance/test_bigquery_usage.py b/metadata-ingestion/tests/performance/bigquery/test_bigquery_usage.py similarity index 80% rename from metadata-ingestion/tests/performance/test_bigquery_usage.py rename to metadata-ingestion/tests/performance/bigquery/test_bigquery_usage.py index 7e05ef070b45d..bbc3378450bff 100644 --- a/metadata-ingestion/tests/performance/test_bigquery_usage.py +++ b/metadata-ingestion/tests/performance/bigquery/test_bigquery_usage.py @@ -2,13 +2,11 @@ import os import random from datetime import timedelta -from typing import Iterable, Tuple import humanfriendly import psutil from datahub.emitter.mce_builder import make_dataset_urn -from datahub.ingestion.api.workunit import MetadataWorkUnit from datahub.ingestion.source.bigquery_v2.bigquery_config import ( BigQueryUsageConfig, BigQueryV2Config, @@ -16,12 +14,13 @@ from datahub.ingestion.source.bigquery_v2.bigquery_report import BigQueryV2Report from datahub.ingestion.source.bigquery_v2.usage import BigQueryUsageExtractor from datahub.utilities.perf_timer import PerfTimer -from tests.performance.bigquery import generate_events, ref_from_table +from tests.performance.bigquery.bigquery_events import generate_events, ref_from_table from tests.performance.data_generation import ( NormalDistribution, generate_data, generate_queries, ) +from tests.performance.helpers import workunit_sink def run_test(): @@ -33,7 +32,7 @@ def run_test(): num_views=2000, time_range=timedelta(days=7), ) - all_tables = seed_metadata.tables + seed_metadata.views + all_tables = seed_metadata.all_tables config = BigQueryV2Config( start_time=seed_metadata.start_time, @@ -88,21 +87,6 @@ def run_test(): print(f"Hash collisions: {report.num_usage_query_hash_collisions}") -def workunit_sink(workunits: Iterable[MetadataWorkUnit]) -> Tuple[int, int]: - peak_memory_usage = psutil.Process(os.getpid()).memory_info().rss - i: int = 0 - for i, wu in enumerate(workunits): - if i % 10_000 == 0: - peak_memory_usage = max( - peak_memory_usage, psutil.Process(os.getpid()).memory_info().rss - ) - peak_memory_usage = max( - peak_memory_usage, psutil.Process(os.getpid()).memory_info().rss - ) - - return i, peak_memory_usage - - if __name__ == "__main__": root_logger = logging.getLogger() root_logger.setLevel(logging.INFO) diff --git a/metadata-ingestion/tests/performance/data_generation.py b/metadata-ingestion/tests/performance/data_generation.py index c530848f27f5c..67b156896909a 100644 --- a/metadata-ingestion/tests/performance/data_generation.py +++ b/metadata-ingestion/tests/performance/data_generation.py @@ -11,11 +11,14 @@ import uuid from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Iterable, List, TypeVar +from typing import Iterable, List, TypeVar, Union, cast from faker import Faker from tests.performance.data_model import ( + Column, + ColumnMapping, + ColumnType, Container, FieldAccess, Query, @@ -52,15 +55,21 @@ def sample_with_floor(self, floor: int = 1) -> int: @dataclass class SeedMetadata: - containers: List[Container] + # Each list is a layer of containers, e.g. [[databases], [schemas]] + containers: List[List[Container]] + tables: List[Table] views: List[View] start_time: datetime end_time: datetime + @property + def all_tables(self) -> List[Table]: + return self.tables + cast(List[Table], self.views) + def generate_data( - num_containers: int, + num_containers: Union[List[int], int], num_tables: int, num_views: int, columns_per_table: NormalDistribution = NormalDistribution(5, 2), @@ -68,32 +77,52 @@ def generate_data( view_definition_length: NormalDistribution = NormalDistribution(150, 50), time_range: timedelta = timedelta(days=14), ) -> SeedMetadata: - containers = [Container(f"container-{i}") for i in range(num_containers)] + # Assemble containers + if isinstance(num_containers, int): + num_containers = [num_containers] + + containers: List[List[Container]] = [] + for i, num_in_layer in enumerate(num_containers): + layer = [ + Container( + f"{i}-container-{j}", + parent=random.choice(containers[-1]) if containers else None, + ) + for j in range(num_in_layer) + ] + containers.append(layer) + + # Assemble tables tables = [ Table( f"table-{i}", - container=random.choice(containers), + container=random.choice(containers[-1]), columns=[ f"column-{j}-{uuid.uuid4()}" for j in range(columns_per_table.sample_with_floor()) ], + column_mapping=None, ) for i in range(num_tables) ] views = [ View( f"view-{i}", - container=random.choice(containers), + container=random.choice(containers[-1]), columns=[ f"column-{j}-{uuid.uuid4()}" for j in range(columns_per_table.sample_with_floor()) ], + column_mapping=None, definition=f"{uuid.uuid4()}-{'*' * view_definition_length.sample_with_floor(10)}", parents=random.sample(tables, parents_per_view.sample_with_floor()), ) for i in range(num_views) ] + for table in tables + views: + _generate_column_mapping(table) + now = datetime.now(tz=timezone.utc) return SeedMetadata( containers=containers, @@ -162,6 +191,18 @@ def generate_queries( ) +def _generate_column_mapping(table: Table) -> ColumnMapping: + d = {} + for column in table.columns: + d[column] = Column( + name=column, + type=random.choice(list(ColumnType)), + nullable=random.random() < 0.1, # Fixed 10% chance for now + ) + table.column_mapping = d + return d + + def _sample_list(lst: List[T], dist: NormalDistribution, floor: int = 1) -> List[T]: return random.sample(lst, min(dist.sample_with_floor(floor), len(lst))) diff --git a/metadata-ingestion/tests/performance/data_model.py b/metadata-ingestion/tests/performance/data_model.py index c593e69ceb9a7..9425fa827070e 100644 --- a/metadata-ingestion/tests/performance/data_model.py +++ b/metadata-ingestion/tests/performance/data_model.py @@ -1,10 +1,10 @@ from dataclasses import dataclass from datetime import datetime -from typing import List, Optional +from enum import Enum +from typing import Dict, List, Optional from typing_extensions import Literal -Column = str StatementType = Literal[ # SELECT + values from OperationTypeClass "SELECT", "INSERT", @@ -21,13 +21,36 @@ @dataclass class Container: name: str + parent: Optional["Container"] = None + + +class ColumnType(str, Enum): + # Can add types that take parameters in the future + + INTEGER = "INTEGER" + FLOAT = "FLOAT" # Double precision (64 bit) + STRING = "STRING" + BOOLEAN = "BOOLEAN" + DATETIME = "DATETIME" + + +@dataclass +class Column: + name: str + type: ColumnType + nullable: bool + + +ColumnRef = str +ColumnMapping = Dict[ColumnRef, Column] @dataclass class Table: name: str container: Container - columns: List[Column] + columns: List[ColumnRef] + column_mapping: Optional[ColumnMapping] def is_view(self) -> bool: return False @@ -44,7 +67,7 @@ def is_view(self) -> bool: @dataclass class FieldAccess: - column: Column + column: ColumnRef table: Table diff --git a/metadata-ingestion/tests/performance/databricks/__init__.py b/metadata-ingestion/tests/performance/databricks/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/metadata-ingestion/tests/performance/databricks/test_unity.py b/metadata-ingestion/tests/performance/databricks/test_unity.py new file mode 100644 index 0000000000000..cc9558f0692ed --- /dev/null +++ b/metadata-ingestion/tests/performance/databricks/test_unity.py @@ -0,0 +1,71 @@ +import logging +import os +from unittest.mock import patch + +import humanfriendly +import psutil + +from datahub.ingestion.api.common import PipelineContext +from datahub.ingestion.source.unity.config import UnityCatalogSourceConfig +from datahub.ingestion.source.unity.source import UnityCatalogSource +from datahub.utilities.perf_timer import PerfTimer +from tests.performance.data_generation import ( + NormalDistribution, + generate_data, + generate_queries, +) +from tests.performance.databricks.unity_proxy_mock import UnityCatalogApiProxyMock +from tests.performance.helpers import workunit_sink + + +def run_test(): + seed_metadata = generate_data( + num_containers=[1, 100, 5000], + num_tables=50000, + num_views=10000, + columns_per_table=NormalDistribution(100, 50), + parents_per_view=NormalDistribution(5, 5), + view_definition_length=NormalDistribution(1000, 300), + ) + queries = generate_queries( + seed_metadata, + num_selects=100000, + num_operations=100000, + num_unique_queries=10000, + num_users=1000, + ) + proxy_mock = UnityCatalogApiProxyMock( + seed_metadata, queries=queries, num_service_principals=10000 + ) + print("Data generated") + + config = UnityCatalogSourceConfig( + token="", workspace_url="http://localhost:1234", include_usage_statistics=False + ) + ctx = PipelineContext(run_id="test") + with patch( + "datahub.ingestion.source.unity.source.UnityCatalogApiProxy", + lambda *args, **kwargs: proxy_mock, + ): + source: UnityCatalogSource = UnityCatalogSource(ctx, config) + + pre_mem_usage = psutil.Process(os.getpid()).memory_info().rss + print(f"Test data size: {humanfriendly.format_size(pre_mem_usage)}") + + with PerfTimer() as timer: + workunits = source.get_workunits() + num_workunits, peak_memory_usage = workunit_sink(workunits) + print(f"Workunits Generated: {num_workunits}") + print(f"Seconds Elapsed: {timer.elapsed_seconds():.2f} seconds") + + print( + f"Peak Memory Used: {humanfriendly.format_size(peak_memory_usage - pre_mem_usage)}" + ) + print(source.report.aspects) + + +if __name__ == "__main__": + root_logger = logging.getLogger() + root_logger.setLevel(logging.INFO) + root_logger.addHandler(logging.StreamHandler()) + run_test() diff --git a/metadata-ingestion/tests/performance/databricks/unity_proxy_mock.py b/metadata-ingestion/tests/performance/databricks/unity_proxy_mock.py new file mode 100644 index 0000000000000..593163e12bf0a --- /dev/null +++ b/metadata-ingestion/tests/performance/databricks/unity_proxy_mock.py @@ -0,0 +1,183 @@ +import uuid +from collections import defaultdict +from datetime import datetime, timezone +from typing import Dict, Iterable, List + +from databricks.sdk.service.catalog import ColumnTypeName +from databricks.sdk.service.sql import QueryStatementType + +from datahub.ingestion.source.unity.proxy_types import ( + Catalog, + CatalogType, + Column, + Metastore, + Query, + Schema, + ServicePrincipal, + Table, + TableType, +) +from tests.performance import data_model +from tests.performance.data_generation import SeedMetadata +from tests.performance.data_model import ColumnType, StatementType + + +class UnityCatalogApiProxyMock: + """Mimics UnityCatalogApiProxy for performance testing.""" + + def __init__( + self, + seed_metadata: SeedMetadata, + queries: Iterable[data_model.Query] = (), + num_service_principals: int = 0, + ) -> None: + self.seed_metadata = seed_metadata + self.queries = queries + self.num_service_principals = num_service_principals + self.warehouse_id = "invalid-warehouse-id" + + # Cache for performance + self._schema_to_table: Dict[str, List[data_model.Table]] = defaultdict(list) + for table in seed_metadata.all_tables: + self._schema_to_table[table.container.name].append(table) + + def check_basic_connectivity(self) -> bool: + return True + + def assigned_metastore(self) -> Metastore: + container = self.seed_metadata.containers[0][0] + return Metastore( + id=container.name, + name=container.name, + global_metastore_id=container.name, + metastore_id=container.name, + comment=None, + owner=None, + cloud=None, + region=None, + ) + + def catalogs(self, metastore: Metastore) -> Iterable[Catalog]: + for container in self.seed_metadata.containers[1]: + if not container.parent or metastore.name != container.parent.name: + continue + + yield Catalog( + id=f"{metastore.id}.{container.name}", + name=container.name, + metastore=metastore, + comment=None, + owner=None, + type=CatalogType.MANAGED_CATALOG, + ) + + def schemas(self, catalog: Catalog) -> Iterable[Schema]: + for container in self.seed_metadata.containers[2]: + # Assumes all catalog names are unique + if not container.parent or catalog.name != container.parent.name: + continue + + yield Schema( + id=f"{catalog.id}.{container.name}", + name=container.name, + catalog=catalog, + comment=None, + owner=None, + ) + + def tables(self, schema: Schema) -> Iterable[Table]: + for table in self._schema_to_table[schema.name]: + columns = [] + if table.column_mapping: + for i, col_name in enumerate(table.columns): + column = table.column_mapping[col_name] + columns.append( + Column( + id=column.name, + name=column.name, + type_name=self._convert_column_type(column.type), + type_text=column.type.value, + nullable=column.nullable, + position=i, + comment=None, + type_precision=0, + type_scale=0, + ) + ) + + yield Table( + id=f"{schema.id}.{table.name}", + name=table.name, + schema=schema, + table_type=TableType.VIEW if table.is_view() else TableType.MANAGED, + columns=columns, + created_at=datetime.now(tz=timezone.utc), + comment=None, + owner=None, + storage_location=None, + data_source_format=None, + generation=None, + created_by="", + updated_at=None, + updated_by=None, + table_id="", + view_definition=table.definition + if isinstance(table, data_model.View) + else None, + properties={}, + ) + + def service_principals(self) -> Iterable[ServicePrincipal]: + for i in range(self.num_service_principals): + yield ServicePrincipal( + id=str(i), + application_id=str(uuid.uuid4()), + display_name=f"user-{i}", + active=True, + ) + + def query_history( + self, + start_time: datetime, + end_time: datetime, + ) -> Iterable[Query]: + for i, query in enumerate(self.queries): + yield Query( + query_id=str(i), + query_text=query.text, + statement_type=self._convert_statement_type(query.type), + start_time=query.timestamp, + end_time=query.timestamp, + user_id=hash(query.actor), + user_name=query.actor, + executed_as_user_id=hash(query.actor), + executed_as_user_name=None, + ) + + def table_lineage(self, table: Table) -> None: + pass + + def get_column_lineage(self, table: Table) -> None: + pass + + @staticmethod + def _convert_column_type(t: ColumnType) -> ColumnTypeName: + if t == ColumnType.INTEGER: + return ColumnTypeName.INT + elif t == ColumnType.FLOAT: + return ColumnTypeName.DOUBLE + elif t == ColumnType.STRING: + return ColumnTypeName.STRING + elif t == ColumnType.BOOLEAN: + return ColumnTypeName.BOOLEAN + elif t == ColumnType.DATETIME: + return ColumnTypeName.TIMESTAMP + else: + raise ValueError(f"Unknown column type: {t}") + + @staticmethod + def _convert_statement_type(t: StatementType) -> QueryStatementType: + if t == "CUSTOM" or t == "UNKNOWN": + return QueryStatementType.OTHER + else: + return QueryStatementType[t] diff --git a/metadata-ingestion/tests/performance/helpers.py b/metadata-ingestion/tests/performance/helpers.py new file mode 100644 index 0000000000000..eb98e53670c96 --- /dev/null +++ b/metadata-ingestion/tests/performance/helpers.py @@ -0,0 +1,21 @@ +import os +from typing import Iterable, Tuple + +import psutil + +from datahub.ingestion.api.workunit import MetadataWorkUnit + + +def workunit_sink(workunits: Iterable[MetadataWorkUnit]) -> Tuple[int, int]: + peak_memory_usage = psutil.Process(os.getpid()).memory_info().rss + i: int = 0 + for i, wu in enumerate(workunits): + if i % 10_000 == 0: + peak_memory_usage = max( + peak_memory_usage, psutil.Process(os.getpid()).memory_info().rss + ) + peak_memory_usage = max( + peak_memory_usage, psutil.Process(os.getpid()).memory_info().rss + ) + + return i, peak_memory_usage diff --git a/metadata-ingestion/tests/unit/test_bigquery_usage.py b/metadata-ingestion/tests/unit/test_bigquery_usage.py index e06c6fb3fe7e5..1eb5d8b00e27c 100644 --- a/metadata-ingestion/tests/unit/test_bigquery_usage.py +++ b/metadata-ingestion/tests/unit/test_bigquery_usage.py @@ -35,7 +35,7 @@ TimeWindowSizeClass, ) from datahub.testing.compare_metadata_json import diff_metadata_json -from tests.performance.bigquery import generate_events, ref_from_table +from tests.performance.bigquery.bigquery_events import generate_events, ref_from_table from tests.performance.data_generation import generate_data, generate_queries from tests.performance.data_model import Container, FieldAccess, Query, Table, View @@ -45,14 +45,15 @@ ACTOR_2, ACTOR_2_URN = "b@acryl.io", "urn:li:corpuser:b" DATABASE_1 = Container("database_1") DATABASE_2 = Container("database_2") -TABLE_1 = Table("table_1", DATABASE_1, ["id", "name", "age"]) -TABLE_2 = Table("table_2", DATABASE_1, ["id", "table_1_id", "value"]) +TABLE_1 = Table("table_1", DATABASE_1, ["id", "name", "age"], None) +TABLE_2 = Table("table_2", DATABASE_1, ["id", "table_1_id", "value"], None) VIEW_1 = View( name="view_1", container=DATABASE_1, columns=["id", "name", "total"], definition="VIEW DEFINITION 1", parents=[TABLE_1, TABLE_2], + column_mapping=None, ) ALL_TABLES = [TABLE_1, TABLE_2, VIEW_1] From 165aa54d1e6f1a1707f79be3cce39ec06c8a1652 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20G=C3=B3mez=20Villamor?= Date: Wed, 4 Oct 2023 19:24:04 +0200 Subject: [PATCH 087/156] doc: DataHubUpgradeHistory_v1 (#8918) --- docs/deploy/confluent-cloud.md | 5 +++++ docs/how/kafka-config.md | 17 +++++++++++------ 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/docs/deploy/confluent-cloud.md b/docs/deploy/confluent-cloud.md index 794b55d4686bf..096fd9984f474 100644 --- a/docs/deploy/confluent-cloud.md +++ b/docs/deploy/confluent-cloud.md @@ -16,6 +16,11 @@ First, you'll need to create following new topics in the [Confluent Control Cent 6. (Deprecated) **MetadataChangeEvent_v4**: Metadata change proposal messages 7. (Deprecated) **MetadataAuditEvent_v4**: Metadata change log messages 8. (Deprecated) **FailedMetadataChangeEvent_v4**: Failed to process #1 event +9. **MetadataGraphEvent_v4**: +10. **MetadataGraphEvent_v4**: +11. **PlatformEvent_v1** +12. **DataHubUpgradeHistory_v1**: Notifies the end of DataHub Upgrade job so dependants can act accordingly (_eg_, startup). + Note this topic requires special configuration: **Infinite retention**. Also, 1 partition is enough for the occasional traffic. The first five are the most important, and are explained in more depth in [MCP/MCL](../advanced/mcp-mcl.md). The final topics are those which are deprecated but still used under certain circumstances. It is likely that in the future they will be completely diff --git a/docs/how/kafka-config.md b/docs/how/kafka-config.md index f3f81c3d07c01..2f20e8b548f83 100644 --- a/docs/how/kafka-config.md +++ b/docs/how/kafka-config.md @@ -52,16 +52,21 @@ Also see [Kafka Connect Security](https://docs.confluent.io/current/connect/secu By default, DataHub relies on the a set of Kafka topics to operate. By default, they have the following names: -- **MetadataChangeProposal_v1** -- **FailedMetadataChangeProposal_v1** -- **MetadataChangeLog_Versioned_v1** -- **MetadataChangeLog_Timeseries_v1** -- **DataHubUsageEvent_v1**: User behavior tracking event for UI +1. **MetadataChangeProposal_v1** +2. **FailedMetadataChangeProposal_v1** +3. **MetadataChangeLog_Versioned_v1** +4. **MetadataChangeLog_Timeseries_v1** +5. **DataHubUsageEvent_v1**: User behavior tracking event for UI 6. (Deprecated) **MetadataChangeEvent_v4**: Metadata change proposal messages 7. (Deprecated) **MetadataAuditEvent_v4**: Metadata change log messages 8. (Deprecated) **FailedMetadataChangeEvent_v4**: Failed to process #1 event +9. **MetadataGraphEvent_v4**: +10. **MetadataGraphEvent_v4**: +11. **PlatformEvent_v1**: +12. **DataHubUpgradeHistory_v1**: Notifies the end of DataHub Upgrade job so dependants can act accordingly (_eg_, startup). + Note this topic requires special configuration: **Infinite retention**. Also, 1 partition is enough for the occasional traffic. -These topics are discussed at more length in [Metadata Events](../what/mxe.md). +How Metadata Events relate to these topics is discussed at more length in [Metadata Events](../what/mxe.md). We've included environment variables to customize the name each of these topics, for cases where an organization has naming rules for your topics. From 3a9452c2072c95cbd7a4bf1270b4ef07abd1b1eb Mon Sep 17 00:00:00 2001 From: Hyejin Yoon <0327jane@gmail.com> Date: Thu, 5 Oct 2023 03:42:00 +0900 Subject: [PATCH 088/156] fix: fix typo on aws guide (#8944) --- docs/deploy/aws.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/deploy/aws.md b/docs/deploy/aws.md index 228fcb51d1a28..e0f57b4a0b0cb 100644 --- a/docs/deploy/aws.md +++ b/docs/deploy/aws.md @@ -100,7 +100,7 @@ eksctl create iamserviceaccount \ Install the TargetGroupBinding custom resource definition by running the following. ``` -kubectl apply -k "github.com/aws/eks-charts/stable/aws-load-balancer-controller//crds?ref=master" +kubectl apply -k "github.com/aws/eks-charts/stable/aws-load-balancer-controller/crds?ref=master" ``` Add the helm chart repository containing the latest version of the ALB controller. From e2afd44bfeb287e8365b99bc7677d06e4172643b Mon Sep 17 00:00:00 2001 From: ethan-cartwright Date: Wed, 4 Oct 2023 16:38:58 -0400 Subject: [PATCH 089/156] feat(dbt-ingestion): add documentation link from dbt source to institutionalMemory (#8686) Co-authored-by: Ethan Cartwright Co-authored-by: Harshal Sheth --- .../docs/sources/dbt/dbt-cloud_recipe.yml | 8 +-- metadata-ingestion/docs/sources/dbt/dbt.md | 7 ++ .../ingestion/source/dbt/dbt_common.py | 6 ++ .../src/datahub/utilities/mapping.py | 67 ++++++++++++++++++- metadata-ingestion/tests/unit/test_mapping.py | 41 ++++++++++++ 5 files changed, 123 insertions(+), 6 deletions(-) diff --git a/metadata-ingestion/docs/sources/dbt/dbt-cloud_recipe.yml b/metadata-ingestion/docs/sources/dbt/dbt-cloud_recipe.yml index 113303cfc1ad4..ef0776b189ca9 100644 --- a/metadata-ingestion/docs/sources/dbt/dbt-cloud_recipe.yml +++ b/metadata-ingestion/docs/sources/dbt/dbt-cloud_recipe.yml @@ -6,14 +6,14 @@ source: # In the URL https://cloud.getdbt.com/next/deploy/107298/projects/175705/jobs/148094, # 107298 is the account_id, 175705 is the project_id, and 148094 is the job_id - account_id: # set to your dbt cloud account id - project_id: # set to your dbt cloud project id - job_id: # set to your dbt cloud job id + account_id: "${DBT_ACCOUNT_ID}" # set to your dbt cloud account id + project_id: "${DBT_PROJECT_ID}" # set to your dbt cloud project id + job_id: "${DBT_JOB_ID}" # set to your dbt cloud job id run_id: # set to your dbt cloud run id. This is optional, and defaults to the latest run target_platform: postgres # Options - target_platform: "my_target_platform_id" # e.g. bigquery/postgres/etc. + target_platform: "${TARGET_PLATFORM_ID}" # e.g. bigquery/postgres/etc. # sink configs diff --git a/metadata-ingestion/docs/sources/dbt/dbt.md b/metadata-ingestion/docs/sources/dbt/dbt.md index bfc3ebd5bb350..43ced13c3b1f8 100644 --- a/metadata-ingestion/docs/sources/dbt/dbt.md +++ b/metadata-ingestion/docs/sources/dbt/dbt.md @@ -38,6 +38,12 @@ meta_mapping: operation: "add_terms" config: separator: "," + documentation_link: + match: "(?:https?)?\:\/\/\w*[^#]*" + operation: "add_doc_link" + config: + link: {{ $match }} + description: "Documentation Link" column_meta_mapping: terms_list: match: ".*" @@ -57,6 +63,7 @@ We support the following operations: 2. add_term - Requires `term` property in config. 3. add_terms - Accepts an optional `separator` property in config. 4. add_owner - Requires `owner_type` property in config which can be either user or group. Optionally accepts the `owner_category` config property which you can set to one of `['TECHNICAL_OWNER', 'BUSINESS_OWNER', 'DATA_STEWARD', 'DATAOWNER'` (defaults to `DATAOWNER`). +5. add_doc_link - Requires `link` and `description` properties in config. Upon ingestion run, this will overwrite current links in the institutional knowledge section with this new link. The anchor text is defined here in the meta_mappings as `description`. Note: diff --git a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py index 782d94f39e8a5..3edeb695e9f21 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py +++ b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py @@ -1188,9 +1188,15 @@ def _generate_base_aspects( ): aspects.append(meta_aspects.get(Constants.ADD_TERM_OPERATION)) + # add meta links aspect + meta_links_aspect = meta_aspects.get(Constants.ADD_DOC_LINK_OPERATION) + if meta_links_aspect and self.config.enable_meta_mapping: + aspects.append(meta_links_aspect) + # add schema metadata aspect schema_metadata = self.get_schema_metadata(self.report, node, mce_platform) aspects.append(schema_metadata) + return aspects def get_schema_metadata( diff --git a/metadata-ingestion/src/datahub/utilities/mapping.py b/metadata-ingestion/src/datahub/utilities/mapping.py index 793eccfb22c7e..eb2d975ee607f 100644 --- a/metadata-ingestion/src/datahub/utilities/mapping.py +++ b/metadata-ingestion/src/datahub/utilities/mapping.py @@ -2,12 +2,16 @@ import logging import operator import re +import time from functools import reduce -from typing import Any, Dict, List, Match, Optional, Union +from typing import Any, Dict, List, Match, Optional, Union, cast from datahub.emitter import mce_builder from datahub.emitter.mce_builder import OwnerType from datahub.metadata.schema_classes import ( + AuditStampClass, + InstitutionalMemoryClass, + InstitutionalMemoryMetadataClass, OwnerClass, OwnershipClass, OwnershipSourceClass, @@ -39,6 +43,7 @@ def _insert_match_value(original_value: str, match_value: str) -> str: class Constants: + ADD_DOC_LINK_OPERATION = "add_doc_link" ADD_TAG_OPERATION = "add_tag" ADD_TERM_OPERATION = "add_term" ADD_TERMS_OPERATION = "add_terms" @@ -47,6 +52,8 @@ class Constants: OPERATION_CONFIG = "config" TAG = "tag" TERM = "term" + DOC_LINK = "link" + DOC_DESCRIPTION = "description" OWNER_TYPE = "owner_type" OWNER_CATEGORY = "owner_category" MATCH = "match" @@ -163,7 +170,6 @@ def process(self, raw_props: Dict[str, Any]) -> Dict[str, Any]: ) operations_value_list.append(operation) # type: ignore operations_map[operation_type] = operations_value_list - aspect_map = self.convert_to_aspects(operations_map) except Exception as e: logger.error(f"Error while processing operation defs over raw_props: {e}") @@ -173,6 +179,7 @@ def convert_to_aspects( self, operation_map: Dict[str, Union[set, list]] ) -> Dict[str, Any]: aspect_map: Dict[str, Any] = {} + if Constants.ADD_TAG_OPERATION in operation_map: tag_aspect = mce_builder.make_global_tag_aspect_with_tag_list( sorted(operation_map[Constants.ADD_TAG_OPERATION]) @@ -195,11 +202,57 @@ def convert_to_aspects( ] ) aspect_map[Constants.ADD_OWNER_OPERATION] = owner_aspect + if Constants.ADD_TERM_OPERATION in operation_map: term_aspect = mce_builder.make_glossary_terms_aspect_from_urn_list( sorted(operation_map[Constants.ADD_TERM_OPERATION]) ) aspect_map[Constants.ADD_TERM_OPERATION] = term_aspect + + if Constants.ADD_DOC_LINK_OPERATION in operation_map: + try: + if len( + operation_map[Constants.ADD_DOC_LINK_OPERATION] + ) == 1 and isinstance( + operation_map[Constants.ADD_DOC_LINK_OPERATION], list + ): + docs_dict = cast( + List[Dict], operation_map[Constants.ADD_DOC_LINK_OPERATION] + )[0] + if "description" not in docs_dict or "link" not in docs_dict: + raise Exception( + "Documentation_link meta_mapping config needs a description key and a link key" + ) + + now = int(time.time() * 1000) # milliseconds since epoch + institutional_memory_element = InstitutionalMemoryMetadataClass( + url=docs_dict["link"], + description=docs_dict["description"], + createStamp=AuditStampClass( + time=now, actor="urn:li:corpuser:ingestion" + ), + ) + + # create a new institutional memory aspect + institutional_memory_aspect = InstitutionalMemoryClass( + elements=[institutional_memory_element] + ) + + aspect_map[ + Constants.ADD_DOC_LINK_OPERATION + ] = institutional_memory_aspect + else: + raise Exception( + f"Expected 1 item of type list for the documentation_link meta_mapping config," + f" received type of {type(operation_map[Constants.ADD_DOC_LINK_OPERATION])}" + f", and size of {len(operation_map[Constants.ADD_DOC_LINK_OPERATION])}." + ) + + except Exception as e: + logger.error( + f"Error while constructing aspect for documentation link and description : {e}" + ) + return aspect_map def get_operation_value( @@ -248,6 +301,16 @@ def get_operation_value( term = operation_config[Constants.TERM] term = _insert_match_value(term, _get_best_match(match, "term")) return mce_builder.make_term_urn(term) + elif ( + operation_type == Constants.ADD_DOC_LINK_OPERATION + and operation_config[Constants.DOC_LINK] + and operation_config[Constants.DOC_DESCRIPTION] + ): + link = operation_config[Constants.DOC_LINK] + link = _insert_match_value(link, _get_best_match(match, "link")) + description = operation_config[Constants.DOC_DESCRIPTION] + return {"link": link, "description": description} + elif operation_type == Constants.ADD_TERMS_OPERATION: separator = operation_config.get(Constants.SEPARATOR, ",") captured_terms = match.group(0) diff --git a/metadata-ingestion/tests/unit/test_mapping.py b/metadata-ingestion/tests/unit/test_mapping.py index d69dd4a8a96b0..5c258f16535f8 100644 --- a/metadata-ingestion/tests/unit/test_mapping.py +++ b/metadata-ingestion/tests/unit/test_mapping.py @@ -4,6 +4,7 @@ from datahub.metadata.schema_classes import ( GlobalTagsClass, GlossaryTermsClass, + InstitutionalMemoryClass, OwnerClass, OwnershipClass, OwnershipSourceTypeClass, @@ -233,6 +234,46 @@ def test_operation_processor_advanced_matching_tags(): assert tag_aspect.tags[0].tag == "urn:li:tag:case_4567" +def test_operation_processor_institutional_memory(): + raw_props = { + "documentation_link": "https://test.com/documentation#ignore-this", + } + processor = OperationProcessor( + operation_defs={ + "documentation_link": { + "match": r"(?:https?)?\:\/\/\w*[^#]*", + "operation": "add_doc_link", + "config": {"link": "{{ $match }}", "description": "test"}, + }, + }, + ) + aspect_map = processor.process(raw_props) + assert "add_doc_link" in aspect_map + + doc_link_aspect: InstitutionalMemoryClass = aspect_map["add_doc_link"] + + assert doc_link_aspect.elements[0].url == "https://test.com/documentation" + assert doc_link_aspect.elements[0].description == "test" + + +def test_operation_processor_institutional_memory_no_description(): + raw_props = { + "documentation_link": "test.com/documentation#ignore-this", + } + processor = OperationProcessor( + operation_defs={ + "documentation_link": { + "match": r"(?:https?)?\:\/\/\w*[^#]*", + "operation": "add_doc_link", + "config": {"link": "{{ $match }}"}, + }, + }, + ) + # we require a description, so this should stay empty + aspect_map = processor.process(raw_props) + assert aspect_map == {} + + def test_operation_processor_matching_nested_props(): raw_props = { "gdpr": { From 0f8d2757352597ceaed62b93547381255dbc650e Mon Sep 17 00:00:00 2001 From: John Joyce Date: Wed, 4 Oct 2023 20:03:40 -0700 Subject: [PATCH 090/156] refactor(style): Improve search bar input focus + styling (#8955) --- .../src/app/search/SearchBar.tsx | 46 +++++++++++-------- .../src/app/shared/admin/HeaderLinks.tsx | 28 +++++------ .../src/conf/theme/theme_dark.config.json | 4 +- .../src/conf/theme/theme_light.config.json | 4 +- 4 files changed, 46 insertions(+), 36 deletions(-) diff --git a/datahub-web-react/src/app/search/SearchBar.tsx b/datahub-web-react/src/app/search/SearchBar.tsx index fb10e1ca0026e..b4699994bc460 100644 --- a/datahub-web-react/src/app/search/SearchBar.tsx +++ b/datahub-web-react/src/app/search/SearchBar.tsx @@ -6,7 +6,7 @@ import { useHistory } from 'react-router'; import { AutoCompleteResultForEntity, EntityType, FacetFilterInput, ScenarioType } from '../../types.generated'; import EntityRegistry from '../entity/EntityRegistry'; import filterSearchQuery from './utils/filterSearchQuery'; -import { ANTD_GRAY, ANTD_GRAY_V2 } from '../entity/shared/constants'; +import { ANTD_GRAY, ANTD_GRAY_V2, REDESIGN_COLORS } from '../entity/shared/constants'; import { getEntityPath } from '../entity/shared/containers/profile/utils'; import { EXACT_SEARCH_PREFIX } from './utils/constants'; import { useListRecommendationsQuery } from '../../graphql/recommendations.generated'; @@ -20,7 +20,6 @@ import RecommendedOption from './autoComplete/RecommendedOption'; import SectionHeader, { EntityTypeLabel } from './autoComplete/SectionHeader'; import { useUserContext } from '../context/useUserContext'; import { navigateToSearchUrl } from './utils/navigateToSearchUrl'; -import { getQuickFilterDetails } from './autoComplete/quickFilters/utils'; import ViewAllSearchItem from './ViewAllSearchItem'; import { ViewSelect } from '../entity/view/select/ViewSelect'; import { combineSiblingsInAutoComplete } from './utils/combineSiblingsInAutoComplete'; @@ -39,13 +38,14 @@ const StyledSearchBar = styled(Input)` &&& { border-radius: 70px; height: 40px; - font-size: 20px; - color: ${ANTD_GRAY[7]}; - background-color: ${ANTD_GRAY_V2[2]}; - } - > .ant-input { font-size: 14px; + color: ${ANTD_GRAY[7]}; background-color: ${ANTD_GRAY_V2[2]}; + border: 2px solid transparent; + + &:focus-within { + border: 1.5px solid ${REDESIGN_COLORS.BLUE}; + } } > .ant-input::placeholder { color: ${ANTD_GRAY_V2[10]}; @@ -203,23 +203,16 @@ export const SearchBar = ({ const { quickFilters, selectedQuickFilter, setSelectedQuickFilter } = useQuickFiltersContext(); const autoCompleteQueryOptions = useMemo(() => { - const query = suggestions.length ? effectiveQuery : ''; - const selectedQuickFilterLabel = - showQuickFilters && selectedQuickFilter - ? getQuickFilterDetails(selectedQuickFilter, entityRegistry).label - : ''; - const text = query || selectedQuickFilterLabel; - - if (!text) return []; + if (effectiveQuery === '') return []; return [ { - value: `${EXACT_SEARCH_PREFIX}${text}`, - label: , + value: `${EXACT_SEARCH_PREFIX}${effectiveQuery}`, + label: , type: EXACT_AUTOCOMPLETE_OPTION_TYPE, }, ]; - }, [showQuickFilters, suggestions.length, effectiveQuery, selectedQuickFilter, entityRegistry]); + }, [effectiveQuery]); const autoCompleteEntityOptions = useMemo(() => { return suggestions.map((suggestion: AutoCompleteResultForEntity) => { @@ -296,6 +289,22 @@ export const SearchBar = ({ } } + const searchInputRef = useRef(null); + + useEffect(() => { + const handleKeyDown = (event) => { + // Support command-k to select the search bar. + // 75 is the keyCode for 'k' + if ((event.metaKey || event.ctrlKey) && event.keyCode === 75) { + (searchInputRef?.current as any)?.focus(); + } + }; + document.addEventListener('keydown', handleKeyDown); + return () => { + document.removeEventListener('keydown', handleKeyDown); + }; + }, []); + return ( } + ref={searchInputRef} /> diff --git a/datahub-web-react/src/app/shared/admin/HeaderLinks.tsx b/datahub-web-react/src/app/shared/admin/HeaderLinks.tsx index ced7d8642576b..ce1ad93565ba4 100644 --- a/datahub-web-react/src/app/shared/admin/HeaderLinks.tsx +++ b/datahub-web-react/src/app/shared/admin/HeaderLinks.tsx @@ -93,20 +93,6 @@ export function HeaderLinks(props: Props) { )} - {showIngestion && ( - - - - - - )} + {showIngestion && ( + + + + + + )} {showSettings && ( diff --git a/datahub-web-react/src/conf/theme/theme_dark.config.json b/datahub-web-react/src/conf/theme/theme_dark.config.json index 9746c3ddde5f3..54ebebd3b692b 100644 --- a/datahub-web-react/src/conf/theme/theme_dark.config.json +++ b/datahub-web-react/src/conf/theme/theme_dark.config.json @@ -30,7 +30,7 @@ "homepageMessage": "Find data you can count(*) on" }, "search": { - "searchbarMessage": "Search Datasets, People, & more..." + "searchbarMessage": "Search Tables, Dashboards, People, & more..." }, "menu": { "items": [ @@ -52,4 +52,4 @@ ] } } -} +} \ No newline at end of file diff --git a/datahub-web-react/src/conf/theme/theme_light.config.json b/datahub-web-react/src/conf/theme/theme_light.config.json index 906c04e38a1ba..6b9ef3eac52b0 100644 --- a/datahub-web-react/src/conf/theme/theme_light.config.json +++ b/datahub-web-react/src/conf/theme/theme_light.config.json @@ -33,7 +33,7 @@ "homepageMessage": "Find data you can count on" }, "search": { - "searchbarMessage": "Search Datasets, People, & more..." + "searchbarMessage": "Search Tables, Dashboards, People, & more..." }, "menu": { "items": [ @@ -60,4 +60,4 @@ ] } } -} +} \ No newline at end of file From 817c371fbf8f8287480a2150925e9526a28f1f6e Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Wed, 4 Oct 2023 23:11:06 -0400 Subject: [PATCH 091/156] feat: data contracts models + CLI (#8923) Co-authored-by: Shirshanka Das Co-authored-by: John Joyce --- .../linkedin/datahub/graphql/TestUtils.java | 3 + .../test/resources/test-entity-registry.yaml | 8 + .../pet_of_the_week.dhub.dc.yaml | 21 ++ .../api/entities/datacontract/__init__.py | 0 .../datacontract/data_quality_assertion.py | 107 +++++++++ .../api/entities/datacontract/datacontract.py | 213 ++++++++++++++++++ .../datacontract/freshness_assertion.py | 86 +++++++ .../entities/datacontract/schema_assertion.py | 81 +++++++ .../datahub/cli/specific/datacontract_cli.py | 80 +++++++ .../src/datahub/cli/specific/file_loader.py | 26 +-- .../src/datahub/emitter/mce_builder.py | 24 +- .../src/datahub/emitter/mcp_builder.py | 27 +-- metadata-ingestion/src/datahub/entrypoints.py | 2 + .../src/datahub/ingestion/api/closeable.py | 8 +- .../ingestion/source/dbt/dbt_common.py | 28 ++- .../integrations/great_expectations/action.py | 19 +- .../tests/unit/test_mcp_builder.py | 3 +- .../linkedin/assertion/AssertionAction.pdl | 22 ++ .../linkedin/assertion/AssertionActions.pdl | 18 ++ .../com/linkedin/assertion/AssertionInfo.pdl | 49 +++- .../linkedin/assertion/AssertionResult.pdl | 18 +- .../assertion/AssertionResultError.pdl | 45 ++++ .../linkedin/assertion/AssertionRunEvent.pdl | 57 +++-- .../linkedin/assertion/AssertionSource.pdl | 27 +++ .../assertion/AssertionStdAggregation.pdl | 10 +- .../assertion/AssertionValueChangeType.pdl | 16 ++ .../com/linkedin/assertion/AuditLogSpec.pdl | 18 ++ .../assertion/DatasetAssertionInfo.pdl | 19 +- .../assertion/FixedIntervalSchedule.pdl | 10 + .../assertion/FreshnessAssertionInfo.pdl | 53 +++++ .../assertion/FreshnessAssertionSchedule.pdl | 66 ++++++ .../assertion/FreshnessCronSchedule.pdl | 25 ++ .../linkedin/assertion/FreshnessFieldKind.pdl | 17 ++ .../linkedin/assertion/FreshnessFieldSpec.pdl | 14 ++ .../IncrementingSegmentFieldTransformer.pdl | 60 +++++ .../IncrementingSegmentRowCountChange.pdl | 33 +++ .../IncrementingSegmentRowCountTotal.pdl | 27 +++ .../assertion/IncrementingSegmentSpec.pdl | 33 +++ .../com/linkedin/assertion/RowCountChange.pdl | 27 +++ .../com/linkedin/assertion/RowCountTotal.pdl | 22 ++ .../assertion/SchemaAssertionInfo.pdl | 29 +++ .../assertion/VolumeAssertionInfo.pdl | 82 +++++++ .../datacontract/DataContractProperties.pdl | 59 +++++ .../datacontract/DataContractStatus.pdl | 27 +++ .../datacontract/DataQualityContract.pdl | 16 ++ .../datacontract/FreshnessContract.pdl | 13 ++ .../linkedin/datacontract/SchemaContract.pdl | 13 ++ .../com/linkedin/dataset/DatasetFilter.pdl | 30 +++ .../linkedin/metadata/key/DataContractKey.pdl | 14 ++ .../com/linkedin/schema/SchemaFieldSpec.pdl | 21 ++ .../src/main/resources/entity-registry.yml | 9 + 51 files changed, 1641 insertions(+), 94 deletions(-) create mode 100644 metadata-ingestion/examples/data_contract/pet_of_the_week.dhub.dc.yaml create mode 100644 metadata-ingestion/src/datahub/api/entities/datacontract/__init__.py create mode 100644 metadata-ingestion/src/datahub/api/entities/datacontract/data_quality_assertion.py create mode 100644 metadata-ingestion/src/datahub/api/entities/datacontract/datacontract.py create mode 100644 metadata-ingestion/src/datahub/api/entities/datacontract/freshness_assertion.py create mode 100644 metadata-ingestion/src/datahub/api/entities/datacontract/schema_assertion.py create mode 100644 metadata-ingestion/src/datahub/cli/specific/datacontract_cli.py create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionAction.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionActions.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionResultError.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionSource.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionValueChangeType.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/AuditLogSpec.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/FixedIntervalSchedule.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessAssertionInfo.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessAssertionSchedule.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessCronSchedule.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessFieldKind.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessFieldSpec.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentFieldTransformer.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentRowCountChange.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentRowCountTotal.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentSpec.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/RowCountChange.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/RowCountTotal.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/SchemaAssertionInfo.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/VolumeAssertionInfo.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/datacontract/DataContractProperties.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/datacontract/DataContractStatus.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/datacontract/DataQualityContract.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/datacontract/FreshnessContract.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/datacontract/SchemaContract.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/dataset/DatasetFilter.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/metadata/key/DataContractKey.pdl create mode 100644 metadata-models/src/main/pegasus/com/linkedin/schema/SchemaFieldSpec.pdl diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/TestUtils.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/TestUtils.java index 272a93fa1989c..606123cac926d 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/TestUtils.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/TestUtils.java @@ -8,6 +8,7 @@ import com.datahub.plugins.auth.authorization.Authorizer; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.UrnUtils; +import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; @@ -21,6 +22,8 @@ public class TestUtils { public static EntityService getMockEntityService() { + PathSpecBasedSchemaAnnotationVisitor.class.getClassLoader() + .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false); EntityRegistry registry = new ConfigEntityRegistry(TestUtils.class.getResourceAsStream("/test-entity-registry.yaml")); EntityService mockEntityService = Mockito.mock(EntityService.class); Mockito.when(mockEntityService.getEntityRegistry()).thenReturn(registry); diff --git a/datahub-graphql-core/src/test/resources/test-entity-registry.yaml b/datahub-graphql-core/src/test/resources/test-entity-registry.yaml index d694ae53ac42f..efd75a7fb07f5 100644 --- a/datahub-graphql-core/src/test/resources/test-entity-registry.yaml +++ b/datahub-graphql-core/src/test/resources/test-entity-registry.yaml @@ -181,6 +181,7 @@ entities: - assertionInfo - dataPlatformInstance - assertionRunEvent + - assertionActions - status - name: dataHubRetention category: internal @@ -292,4 +293,11 @@ entities: aspects: - ownershipTypeInfo - status +- name: dataContract + category: core + keyAspect: dataContractKey + aspects: + - dataContractProperties + - dataContractStatus + - status events: diff --git a/metadata-ingestion/examples/data_contract/pet_of_the_week.dhub.dc.yaml b/metadata-ingestion/examples/data_contract/pet_of_the_week.dhub.dc.yaml new file mode 100644 index 0000000000000..c73904403f678 --- /dev/null +++ b/metadata-ingestion/examples/data_contract/pet_of_the_week.dhub.dc.yaml @@ -0,0 +1,21 @@ +# id: pet_details_dc # Optional: This is the unique identifier for the data contract +display_name: Data Contract for SampleHiveDataset +entity: urn:li:dataset:(urn:li:dataPlatform:hive,SampleHiveDataset,PROD) +freshness: + time: 0700 + granularity: DAILY +schema: + properties: + field_foo: + type: string + native_type: VARCHAR(100) + field_bar: + type: boolean + required: + - field_bar +data_quality: + - type: column_range + config: + column: field_foo + min: 0 + max: 100 diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/__init__.py b/metadata-ingestion/src/datahub/api/entities/datacontract/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/data_quality_assertion.py b/metadata-ingestion/src/datahub/api/entities/datacontract/data_quality_assertion.py new file mode 100644 index 0000000000000..a665e95e93c43 --- /dev/null +++ b/metadata-ingestion/src/datahub/api/entities/datacontract/data_quality_assertion.py @@ -0,0 +1,107 @@ +from typing import List, Optional, Union + +import pydantic +from typing_extensions import Literal + +import datahub.emitter.mce_builder as builder +from datahub.configuration.common import ConfigModel +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.metadata.schema_classes import ( + AssertionInfoClass, + AssertionStdAggregationClass, + AssertionStdOperatorClass, + AssertionStdParameterClass, + AssertionStdParametersClass, + AssertionStdParameterTypeClass, + AssertionTypeClass, + DatasetAssertionInfoClass, + DatasetAssertionScopeClass, +) + + +class IdConfigMixin(ConfigModel): + id_raw: Optional[str] = pydantic.Field( + default=None, + alias="id", + description="The id of the assertion. If not provided, one will be generated using the type.", + ) + + def generate_default_id(self) -> str: + raise NotImplementedError + + +class CustomSQLAssertion(IdConfigMixin, ConfigModel): + type: Literal["custom_sql"] + + sql: str + + def generate_dataset_assertion_info( + self, entity_urn: str + ) -> DatasetAssertionInfoClass: + return DatasetAssertionInfoClass( + dataset=entity_urn, + scope=DatasetAssertionScopeClass.UNKNOWN, + fields=[], + operator=AssertionStdOperatorClass._NATIVE_, + aggregation=AssertionStdAggregationClass._NATIVE_, + logic=self.sql, + ) + + +class ColumnUniqueAssertion(IdConfigMixin, ConfigModel): + type: Literal["unique"] + + # TODO: support multiple columns? + column: str + + def generate_default_id(self) -> str: + return f"{self.type}-{self.column}" + + def generate_dataset_assertion_info( + self, entity_urn: str + ) -> DatasetAssertionInfoClass: + return DatasetAssertionInfoClass( + dataset=entity_urn, + scope=DatasetAssertionScopeClass.DATASET_COLUMN, + fields=[builder.make_schema_field_urn(entity_urn, self.column)], + operator=AssertionStdOperatorClass.EQUAL_TO, + aggregation=AssertionStdAggregationClass.UNIQUE_PROPOTION, # purposely using the misspelled version to work with gql + parameters=AssertionStdParametersClass( + value=AssertionStdParameterClass( + value="1", type=AssertionStdParameterTypeClass.NUMBER + ) + ), + ) + + +class DataQualityAssertion(ConfigModel): + __root__: Union[ + CustomSQLAssertion, + ColumnUniqueAssertion, + ] = pydantic.Field(discriminator="type") + + @property + def id(self) -> str: + if self.__root__.id_raw: + return self.__root__.id_raw + try: + return self.__root__.generate_default_id() + except NotImplementedError: + return self.__root__.type + + def generate_mcp( + self, assertion_urn: str, entity_urn: str + ) -> List[MetadataChangeProposalWrapper]: + dataset_assertion_info = self.__root__.generate_dataset_assertion_info( + entity_urn + ) + + return [ + MetadataChangeProposalWrapper( + entityUrn=assertion_urn, + aspect=AssertionInfoClass( + type=AssertionTypeClass.DATASET, + datasetAssertion=dataset_assertion_info, + ), + ) + ] diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/datacontract.py b/metadata-ingestion/src/datahub/api/entities/datacontract/datacontract.py new file mode 100644 index 0000000000000..2df446623a9d6 --- /dev/null +++ b/metadata-ingestion/src/datahub/api/entities/datacontract/datacontract.py @@ -0,0 +1,213 @@ +import collections +from typing import Iterable, List, Optional, Tuple + +import pydantic +from ruamel.yaml import YAML +from typing_extensions import Literal + +import datahub.emitter.mce_builder as builder +from datahub.api.entities.datacontract.data_quality_assertion import ( + DataQualityAssertion, +) +from datahub.api.entities.datacontract.freshness_assertion import FreshnessAssertion +from datahub.api.entities.datacontract.schema_assertion import SchemaAssertion +from datahub.configuration.common import ConfigModel +from datahub.emitter.mce_builder import datahub_guid, make_assertion_urn +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.metadata.schema_classes import ( + DataContractPropertiesClass, + DataContractStateClass, + DataContractStatusClass, + DataQualityContractClass, + FreshnessContractClass, + SchemaContractClass, + StatusClass, +) +from datahub.utilities.urns.urn import guess_entity_type + + +class DataContract(ConfigModel): + """A yml representation of a Data Contract. + + This model is used as a simpler, Python-native representation of a DataHub data contract. + It can be easily parsed from a YAML file, and can be easily converted into series of MCPs + that can be emitted to DataHub. + """ + + version: Literal[1] + + id: Optional[str] = pydantic.Field( + default=None, + alias="urn", + description="The data contract urn. If not provided, one will be generated.", + ) + entity: str = pydantic.Field( + description="The entity urn that the Data Contract is associated with" + ) + # TODO: add support for properties + # properties: Optional[Dict[str, str]] = None + + schema_field: Optional[SchemaAssertion] = pydantic.Field( + default=None, alias="schema" + ) + + freshness: Optional[FreshnessAssertion] = pydantic.Field(default=None) + + # TODO: Add a validator to ensure that ids are unique + data_quality: Optional[List[DataQualityAssertion]] = None + + _original_yaml_dict: Optional[dict] = None + + @pydantic.validator("data_quality") + def validate_data_quality( + cls, data_quality: Optional[List[DataQualityAssertion]] + ) -> Optional[List[DataQualityAssertion]]: + if data_quality: + # Raise an error if there are duplicate ids. + id_counts = collections.Counter(dq_check.id for dq_check in data_quality) + duplicates = [id for id, count in id_counts.items() if count > 1] + + if duplicates: + raise ValueError( + f"Got multiple data quality tests with the same type or ID: {duplicates}. Set a unique ID for each data quality test." + ) + + return data_quality + + @property + def urn(self) -> str: + if self.id: + assert guess_entity_type(self.id) == "dataContract" + return self.id + + # Data contract urns are stable + guid_obj = {"entity": self.entity} + urn = f"urn:li:dataContract:{datahub_guid(guid_obj)}" + return urn + + def _generate_freshness_assertion( + self, freshness: FreshnessAssertion + ) -> Tuple[str, List[MetadataChangeProposalWrapper]]: + guid_dict = { + "contract": self.urn, + "entity": self.entity, + "freshness": freshness.id, + } + assertion_urn = builder.make_assertion_urn(builder.datahub_guid(guid_dict)) + + return ( + assertion_urn, + freshness.generate_mcp(assertion_urn, self.entity), + ) + + def _generate_schema_assertion( + self, schema_metadata: SchemaAssertion + ) -> Tuple[str, List[MetadataChangeProposalWrapper]]: + # ingredients for guid -> the contract id, the fact that this is a schema assertion and the entity on which the assertion is made + guid_dict = { + "contract": self.urn, + "entity": self.entity, + "schema": schema_metadata.id, + } + assertion_urn = make_assertion_urn(datahub_guid(guid_dict)) + + return ( + assertion_urn, + schema_metadata.generate_mcp(assertion_urn, self.entity), + ) + + def _generate_data_quality_assertion( + self, data_quality: DataQualityAssertion + ) -> Tuple[str, List[MetadataChangeProposalWrapper]]: + guid_dict = { + "contract": self.urn, + "entity": self.entity, + "data_quality": data_quality.id, + } + assertion_urn = make_assertion_urn(datahub_guid(guid_dict)) + + return ( + assertion_urn, + data_quality.generate_mcp(assertion_urn, self.entity), + ) + + def _generate_dq_assertions( + self, data_quality_spec: List[DataQualityAssertion] + ) -> Tuple[List[str], List[MetadataChangeProposalWrapper]]: + assertion_urns = [] + assertion_mcps = [] + + for dq_check in data_quality_spec: + assertion_urn, assertion_mcp = self._generate_data_quality_assertion( + dq_check + ) + + assertion_urns.append(assertion_urn) + assertion_mcps.extend(assertion_mcp) + + return (assertion_urns, assertion_mcps) + + def generate_mcp( + self, + ) -> Iterable[MetadataChangeProposalWrapper]: + schema_assertion_urn = None + if self.schema_field is not None: + ( + schema_assertion_urn, + schema_assertion_mcps, + ) = self._generate_schema_assertion(self.schema_field) + yield from schema_assertion_mcps + + freshness_assertion_urn = None + if self.freshness: + ( + freshness_assertion_urn, + sla_assertion_mcps, + ) = self._generate_freshness_assertion(self.freshness) + yield from sla_assertion_mcps + + dq_assertions, dq_assertion_mcps = self._generate_dq_assertions( + self.data_quality or [] + ) + yield from dq_assertion_mcps + + # Now that we've generated the assertions, we can generate + # the actual data contract. + yield from MetadataChangeProposalWrapper.construct_many( + entityUrn=self.urn, + aspects=[ + DataContractPropertiesClass( + entity=self.entity, + schema=[SchemaContractClass(assertion=schema_assertion_urn)] + if schema_assertion_urn + else None, + freshness=[ + FreshnessContractClass(assertion=freshness_assertion_urn) + ] + if freshness_assertion_urn + else None, + dataQuality=[ + DataQualityContractClass(assertion=dq_assertion_urn) + for dq_assertion_urn in dq_assertions + ], + ), + # Also emit status. + StatusClass(removed=False), + # Emit the contract state as PENDING. + DataContractStatusClass(state=DataContractStateClass.PENDING) + if True + else None, + ], + ) + + @classmethod + def from_yaml( + cls, + file: str, + ) -> "DataContract": + with open(file) as fp: + yaml = YAML(typ="rt") # default, if not specfied, is 'rt' (round-trip) + orig_dictionary = yaml.load(fp) + parsed_data_contract = DataContract.parse_obj(orig_dictionary) + parsed_data_contract._original_yaml_dict = orig_dictionary + return parsed_data_contract diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/freshness_assertion.py b/metadata-ingestion/src/datahub/api/entities/datacontract/freshness_assertion.py new file mode 100644 index 0000000000000..ee8fa1181e614 --- /dev/null +++ b/metadata-ingestion/src/datahub/api/entities/datacontract/freshness_assertion.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +from datetime import timedelta +from typing import List, Union + +import pydantic +from typing_extensions import Literal + +from datahub.configuration.common import ConfigModel +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.metadata.schema_classes import ( + AssertionInfoClass, + AssertionTypeClass, + CalendarIntervalClass, + FixedIntervalScheduleClass, + FreshnessAssertionInfoClass, + FreshnessAssertionScheduleClass, + FreshnessAssertionScheduleTypeClass, + FreshnessAssertionTypeClass, + FreshnessCronScheduleClass, +) + + +class CronFreshnessAssertion(ConfigModel): + type: Literal["cron"] + + cron: str = pydantic.Field( + description="The cron expression to use. See https://crontab.guru/ for help." + ) + timezone: str = pydantic.Field( + "UTC", + description="The timezone to use for the cron schedule. Defaults to UTC.", + ) + + +class FixedIntervalFreshnessAssertion(ConfigModel): + type: Literal["interval"] + + interval: timedelta + + +class FreshnessAssertion(ConfigModel): + __root__: Union[ + CronFreshnessAssertion, FixedIntervalFreshnessAssertion + ] = pydantic.Field(discriminator="type") + + @property + def id(self): + return self.__root__.type + + def generate_mcp( + self, assertion_urn: str, entity_urn: str + ) -> List[MetadataChangeProposalWrapper]: + freshness = self.__root__ + + if isinstance(freshness, CronFreshnessAssertion): + schedule = FreshnessAssertionScheduleClass( + type=FreshnessAssertionScheduleTypeClass.CRON, + cron=FreshnessCronScheduleClass( + cron=freshness.cron, + timezone=freshness.timezone, + ), + ) + elif isinstance(freshness, FixedIntervalFreshnessAssertion): + schedule = FreshnessAssertionScheduleClass( + type=FreshnessAssertionScheduleTypeClass.FIXED_INTERVAL, + fixedInterval=FixedIntervalScheduleClass( + unit=CalendarIntervalClass.SECOND, + multiple=int(freshness.interval.total_seconds()), + ), + ) + else: + raise ValueError(f"Unknown freshness type {freshness}") + + assertionInfo = AssertionInfoClass( + type=AssertionTypeClass.FRESHNESS, + freshnessAssertion=FreshnessAssertionInfoClass( + entity=entity_urn, + type=FreshnessAssertionTypeClass.DATASET_CHANGE, + schedule=schedule, + ), + ) + + return [ + MetadataChangeProposalWrapper(entityUrn=assertion_urn, aspect=assertionInfo) + ] diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/schema_assertion.py b/metadata-ingestion/src/datahub/api/entities/datacontract/schema_assertion.py new file mode 100644 index 0000000000000..b5b592e01f58f --- /dev/null +++ b/metadata-ingestion/src/datahub/api/entities/datacontract/schema_assertion.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import json +from typing import List, Union + +import pydantic +from typing_extensions import Literal + +from datahub.configuration.common import ConfigModel +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.ingestion.extractor.json_schema_util import get_schema_metadata +from datahub.metadata.schema_classes import ( + AssertionInfoClass, + AssertionTypeClass, + SchemaAssertionInfoClass, + SchemaFieldClass, + SchemalessClass, + SchemaMetadataClass, +) + + +class JsonSchemaContract(ConfigModel): + type: Literal["json-schema"] + + json_schema: dict = pydantic.Field(alias="json-schema") + + _schema_metadata: SchemaMetadataClass + + def _init_private_attributes(self) -> None: + super()._init_private_attributes() + self._schema_metadata = get_schema_metadata( + platform="urn:li:dataPlatform:datahub", + name="", + json_schema=self.json_schema, + raw_schema_string=json.dumps(self.json_schema), + ) + + +class FieldListSchemaContract(ConfigModel, arbitrary_types_allowed=True): + type: Literal["field-list"] + + fields: List[SchemaFieldClass] + + _schema_metadata: SchemaMetadataClass + + def _init_private_attributes(self) -> None: + super()._init_private_attributes() + self._schema_metadata = SchemaMetadataClass( + schemaName="", + platform="urn:li:dataPlatform:datahub", + version=0, + hash="", + platformSchema=SchemalessClass(), + fields=self.fields, + ) + + +class SchemaAssertion(ConfigModel): + __root__: Union[JsonSchemaContract, FieldListSchemaContract] = pydantic.Field( + discriminator="type" + ) + + @property + def id(self): + return self.__root__.type + + def generate_mcp( + self, assertion_urn: str, entity_urn: str + ) -> List[MetadataChangeProposalWrapper]: + schema_metadata = self.__root__._schema_metadata + + assertionInfo = AssertionInfoClass( + type=AssertionTypeClass.DATA_SCHEMA, + schemaAssertion=SchemaAssertionInfoClass( + entity=entity_urn, schema=schema_metadata + ), + ) + + return [ + MetadataChangeProposalWrapper(entityUrn=assertion_urn, aspect=assertionInfo) + ] diff --git a/metadata-ingestion/src/datahub/cli/specific/datacontract_cli.py b/metadata-ingestion/src/datahub/cli/specific/datacontract_cli.py new file mode 100644 index 0000000000000..3745943c8c96a --- /dev/null +++ b/metadata-ingestion/src/datahub/cli/specific/datacontract_cli.py @@ -0,0 +1,80 @@ +import logging +from typing import Optional + +import click +from click_default_group import DefaultGroup + +from datahub.api.entities.datacontract.datacontract import DataContract +from datahub.ingestion.graph.client import get_default_graph +from datahub.telemetry import telemetry +from datahub.upgrade import upgrade + +logger = logging.getLogger(__name__) + + +@click.group(cls=DefaultGroup, default="upsert") +def datacontract() -> None: + """A group of commands to interact with the DataContract entity in DataHub.""" + pass + + +@datacontract.command() +@click.option("-f", "--file", required=True, type=click.Path(exists=True)) +@upgrade.check_upgrade +@telemetry.with_telemetry() +def upsert(file: str) -> None: + """Upsert (create or update) a Data Contract in DataHub.""" + + data_contract: DataContract = DataContract.from_yaml(file) + urn = data_contract.urn + + with get_default_graph() as graph: + if not graph.exists(data_contract.entity): + raise ValueError( + f"Cannot define a data contract for non-existent entity {data_contract.entity}" + ) + + try: + for mcp in data_contract.generate_mcp(): + graph.emit(mcp) + click.secho(f"Update succeeded for urn {urn}.", fg="green") + except Exception as e: + logger.exception(e) + click.secho( + f"Update failed for {urn}: {e}", + fg="red", + ) + + +@datacontract.command() +@click.option( + "--urn", required=False, type=str, help="The urn for the data contract to delete" +) +@click.option( + "-f", + "--file", + required=False, + type=click.Path(exists=True), + help="The file containing the data contract definition", +) +@click.option("--hard/--soft", required=False, is_flag=True, default=False) +@upgrade.check_upgrade +@telemetry.with_telemetry() +def delete(urn: Optional[str], file: Optional[str], hard: bool) -> None: + """Delete a Data Contract in DataHub. Defaults to a soft-delete. Use --hard to completely erase metadata.""" + + if not urn: + if not file: + raise click.UsageError( + "Must provide either an urn or a file to delete a data contract" + ) + + data_contract = DataContract.from_yaml(file) + urn = data_contract.urn + + with get_default_graph() as graph: + if not graph.exists(urn): + raise ValueError(f"Data Contract {urn} does not exist") + + graph.delete_entity(urn, hard=hard) + click.secho(f"Data Contract {urn} deleted") diff --git a/metadata-ingestion/src/datahub/cli/specific/file_loader.py b/metadata-ingestion/src/datahub/cli/specific/file_loader.py index 54f12e024d294..a9787343fdb91 100644 --- a/metadata-ingestion/src/datahub/cli/specific/file_loader.py +++ b/metadata-ingestion/src/datahub/cli/specific/file_loader.py @@ -1,9 +1,7 @@ -import io from pathlib import Path from typing import Union -from datahub.configuration.common import ConfigurationError -from datahub.configuration.yaml import YamlConfigurationMechanism +from datahub.configuration.config_loader import load_config_file def load_file(config_file: Path) -> Union[dict, list]: @@ -17,19 +15,11 @@ def load_file(config_file: Path) -> Union[dict, list]: evolve to becoming a standard function that all the specific. cli variants will use to load up the models from external files """ - if not isinstance(config_file, Path): - config_file = Path(config_file) - if not config_file.is_file(): - raise ConfigurationError(f"Cannot open config file {config_file}") - if config_file.suffix in {".yaml", ".yml"}: - config_mech: YamlConfigurationMechanism = YamlConfigurationMechanism() - else: - raise ConfigurationError( - f"Only .yaml and .yml are supported. Cannot process file type {config_file.suffix}" - ) - - raw_config_file = config_file.read_text() - config_fp = io.StringIO(raw_config_file) - raw_config = config_mech.load_config(config_fp) - return raw_config + res = load_config_file( + config_file, + squirrel_original_config=False, + resolve_env_vars=False, + allow_stdin=False, + ) + return res diff --git a/metadata-ingestion/src/datahub/emitter/mce_builder.py b/metadata-ingestion/src/datahub/emitter/mce_builder.py index 0928818c7005c..64c9ec1bb5704 100644 --- a/metadata-ingestion/src/datahub/emitter/mce_builder.py +++ b/metadata-ingestion/src/datahub/emitter/mce_builder.py @@ -1,11 +1,11 @@ """Convenience functions for creating MCEs""" +import hashlib import json import logging import os import re import time from enum import Enum -from hashlib import md5 from typing import ( TYPE_CHECKING, Any, @@ -21,7 +21,6 @@ import typing_inspect from datahub.configuration.source_common import DEFAULT_ENV as DEFAULT_ENV_CONFIGURATION -from datahub.emitter.serialization_helper import pre_json_transform from datahub.metadata.schema_classes import ( AssertionKeyClass, AuditStampClass, @@ -159,11 +158,24 @@ def container_urn_to_key(guid: str) -> Optional[ContainerKeyClass]: return None +class _DatahubKeyJSONEncoder(json.JSONEncoder): + # overload method default + def default(self, obj: Any) -> Any: + if hasattr(obj, "guid"): + return obj.guid() + # Call the default method for other types + return json.JSONEncoder.default(self, obj) + + def datahub_guid(obj: dict) -> str: - obj_str = json.dumps( - pre_json_transform(obj), separators=(",", ":"), sort_keys=True - ).encode("utf-8") - return md5(obj_str).hexdigest() + json_key = json.dumps( + obj, + separators=(",", ":"), + sort_keys=True, + cls=_DatahubKeyJSONEncoder, + ) + md5_hash = hashlib.md5(json_key.encode("utf-8")) + return str(md5_hash.hexdigest()) def make_assertion_urn(assertion_id: str) -> str: diff --git a/metadata-ingestion/src/datahub/emitter/mcp_builder.py b/metadata-ingestion/src/datahub/emitter/mcp_builder.py index 7419577b367aa..06f689dfd317b 100644 --- a/metadata-ingestion/src/datahub/emitter/mcp_builder.py +++ b/metadata-ingestion/src/datahub/emitter/mcp_builder.py @@ -1,11 +1,10 @@ -import hashlib -import json -from typing import Any, Dict, Iterable, List, Optional, TypeVar +from typing import Dict, Iterable, List, Optional, TypeVar from pydantic.fields import Field from pydantic.main import BaseModel from datahub.emitter.mce_builder import ( + datahub_guid, make_container_urn, make_data_platform_urn, make_dataplatform_instance_urn, @@ -33,24 +32,13 @@ ) -def _stable_guid_from_dict(d: dict) -> str: - json_key = json.dumps( - d, - separators=(",", ":"), - sort_keys=True, - cls=DatahubKeyJSONEncoder, - ) - md5_hash = hashlib.md5(json_key.encode("utf-8")) - return str(md5_hash.hexdigest()) - - class DatahubKey(BaseModel): def guid_dict(self) -> Dict[str, str]: return self.dict(by_alias=True, exclude_none=True) def guid(self) -> str: bag = self.guid_dict() - return _stable_guid_from_dict(bag) + return datahub_guid(bag) class ContainerKey(DatahubKey): @@ -137,15 +125,6 @@ def as_urn(self) -> str: ) -class DatahubKeyJSONEncoder(json.JSONEncoder): - # overload method default - def default(self, obj: Any) -> Any: - if hasattr(obj, "guid"): - return obj.guid() - # Call the default method for other types - return json.JSONEncoder.default(self, obj) - - KeyType = TypeVar("KeyType", bound=ContainerKey) diff --git a/metadata-ingestion/src/datahub/entrypoints.py b/metadata-ingestion/src/datahub/entrypoints.py index 84615fd9a6148..5bfab3b841fa3 100644 --- a/metadata-ingestion/src/datahub/entrypoints.py +++ b/metadata-ingestion/src/datahub/entrypoints.py @@ -21,6 +21,7 @@ from datahub.cli.ingest_cli import ingest from datahub.cli.migrate import migrate from datahub.cli.put_cli import put +from datahub.cli.specific.datacontract_cli import datacontract from datahub.cli.specific.dataproduct_cli import dataproduct from datahub.cli.specific.group_cli import group from datahub.cli.specific.user_cli import user @@ -158,6 +159,7 @@ def init() -> None: datahub.add_command(user) datahub.add_command(group) datahub.add_command(dataproduct) +datahub.add_command(datacontract) try: from datahub.cli.lite_cli import lite diff --git a/metadata-ingestion/src/datahub/ingestion/api/closeable.py b/metadata-ingestion/src/datahub/ingestion/api/closeable.py index 523174b9978b3..80a5008ed6368 100644 --- a/metadata-ingestion/src/datahub/ingestion/api/closeable.py +++ b/metadata-ingestion/src/datahub/ingestion/api/closeable.py @@ -1,7 +1,9 @@ from abc import abstractmethod from contextlib import AbstractContextManager from types import TracebackType -from typing import Optional, Type +from typing import Optional, Type, TypeVar + +_Self = TypeVar("_Self", bound="Closeable") class Closeable(AbstractContextManager): @@ -9,6 +11,10 @@ class Closeable(AbstractContextManager): def close(self) -> None: pass + def __enter__(self: _Self) -> _Self: + # This method is mainly required for type checking. + return self + def __exit__( self, exc_type: Optional[Type[BaseException]], diff --git a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py index 3edeb695e9f21..f9b71892975b4 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py +++ b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py @@ -701,18 +701,22 @@ def create_test_entity_mcps( assertion_urn = mce_builder.make_assertion_urn( mce_builder.datahub_guid( { - "platform": DBT_PLATFORM, - "name": node.dbt_name, - "instance": self.config.platform_instance, - **( - # Ideally we'd include the env unconditionally. However, we started out - # not including env in the guid, so we need to maintain backwards compatibility - # with existing PROD assertions. - {"env": self.config.env} - if self.config.env != mce_builder.DEFAULT_ENV - and self.config.include_env_in_assertion_guid - else {} - ), + k: v + for k, v in { + "platform": DBT_PLATFORM, + "name": node.dbt_name, + "instance": self.config.platform_instance, + **( + # Ideally we'd include the env unconditionally. However, we started out + # not including env in the guid, so we need to maintain backwards compatibility + # with existing PROD assertions. + {"env": self.config.env} + if self.config.env != mce_builder.DEFAULT_ENV + and self.config.include_env_in_assertion_guid + else {} + ), + }.items() + if v is not None } ) ) diff --git a/metadata-ingestion/src/datahub/integrations/great_expectations/action.py b/metadata-ingestion/src/datahub/integrations/great_expectations/action.py index f116550328819..8b393a8f6f1c6 100644 --- a/metadata-ingestion/src/datahub/integrations/great_expectations/action.py +++ b/metadata-ingestion/src/datahub/integrations/great_expectations/action.py @@ -35,6 +35,7 @@ from datahub.cli.cli_utils import get_boolean_env_variable from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.rest_emitter import DatahubRestEmitter +from datahub.emitter.serialization_helper import pre_json_transform from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import ( get_platform_from_sqlalchemy_uri, ) @@ -253,13 +254,15 @@ def get_assertions_with_results( # possibly for each validation run assertionUrn = builder.make_assertion_urn( builder.datahub_guid( - { - "platform": GE_PLATFORM_NAME, - "nativeType": expectation_type, - "nativeParameters": kwargs, - "dataset": assertion_datasets[0], - "fields": assertion_fields, - } + pre_json_transform( + { + "platform": GE_PLATFORM_NAME, + "nativeType": expectation_type, + "nativeParameters": kwargs, + "dataset": assertion_datasets[0], + "fields": assertion_fields, + } + ) ) ) logger.debug( @@ -638,7 +641,7 @@ def get_dataset_partitions(self, batch_identifier, data_asset): ].batch_request.runtime_parameters["query"] partitionSpec = PartitionSpecClass( type=PartitionTypeClass.QUERY, - partition=f"Query_{builder.datahub_guid(query)}", + partition=f"Query_{builder.datahub_guid(pre_json_transform(query))}", ) batchSpec = BatchSpec( diff --git a/metadata-ingestion/tests/unit/test_mcp_builder.py b/metadata-ingestion/tests/unit/test_mcp_builder.py index 23f2bddc2084e..561b782ef9e46 100644 --- a/metadata-ingestion/tests/unit/test_mcp_builder.py +++ b/metadata-ingestion/tests/unit/test_mcp_builder.py @@ -1,5 +1,4 @@ import datahub.emitter.mcp_builder as builder -from datahub.emitter.mce_builder import datahub_guid def test_guid_generator(): @@ -80,7 +79,7 @@ def test_guid_generators(): key = builder.SchemaKey( database="test", schema="Test", platform="mysql", instance="TestInstance" ) - guid_datahub = datahub_guid(key.dict(by_alias=True)) + guid_datahub = key.guid() guid = key.guid() assert guid == guid_datahub diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionAction.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionAction.pdl new file mode 100644 index 0000000000000..df6620b66bfd8 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionAction.pdl @@ -0,0 +1,22 @@ +namespace com.linkedin.assertion + +/** + * The Actions about an Assertion. + * In the future, we'll likely extend this model to support additional + * parameters or options related to the assertion actions. + */ +record AssertionAction { + /** + * The type of the Action + */ + type: enum AssertionActionType { + /** + * Raise an incident. + */ + RAISE_INCIDENT + /** + * Resolve open incidents related to the assertion. + */ + RESOLVE_INCIDENT + } +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionActions.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionActions.pdl new file mode 100644 index 0000000000000..61846c1ba9c12 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionActions.pdl @@ -0,0 +1,18 @@ +namespace com.linkedin.assertion + +/** + * The Actions about an Assertion + */ +@Aspect = { + "name": "assertionActions" +} +record AssertionActions { + /** + * Actions to be executed on successful assertion run. + */ + onSuccess: array[AssertionAction] = [] + /** + * Actions to be executed on failed assertion run. + */ + onFailure: array[AssertionAction] = [] +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionInfo.pdl index 77ee147a781e2..ae2a58028057b 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionInfo.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionInfo.pdl @@ -13,13 +13,58 @@ record AssertionInfo includes CustomProperties, ExternalReference { /** * Type of assertion. Assertion types can evolve to span Datasets, Flows (Pipelines), Models, Features etc. */ + @Searchable = { } type: enum AssertionType { - // A single-dataset assertion. When this is the value, the datasetAssertion field will be populated. + /** + * A single-dataset assertion. When this is the value, the datasetAssertion field will be populated. + */ DATASET + + /** + * A freshness assertion, or an assertion which indicates when a particular operation should occur + * to an asset. + */ + FRESHNESS + + /** + * A volume assertion, or an assertion which indicates how much data should be available for a + * particular asset. + */ + VOLUME + + /** + * A schema or structural assertion. + * + * Would have named this SCHEMA but the codegen for PDL does not allow this (reserved word). + */ + DATA_SCHEMA } /** - * Dataset Assertion information when type is DATASET + * A Dataset Assertion definition. This field is populated when the type is DATASET. */ datasetAssertion: optional DatasetAssertionInfo + + /** + * An Freshness Assertion definition. This field is populated when the type is FRESHNESS. + */ + freshnessAssertion: optional FreshnessAssertionInfo + + /** + * An Volume Assertion definition. This field is populated when the type is VOLUME. + */ + volumeAssertion: optional VolumeAssertionInfo + + /** + * An schema Assertion definition. This field is populated when the type is DATASET_SCHEMA + */ + schemaAssertion: optional SchemaAssertionInfo + + /** + * The source or origin of the Assertion definition. + * + * If the source type of the Assertion is EXTERNAL, it is expected to have a corresponding dataPlatformInstance aspect detailing + * the platform where it was ingested from. + */ + source: optional AssertionSource } \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionResult.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionResult.pdl index decbfc08263de..ded84e1969153 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionResult.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionResult.pdl @@ -5,10 +5,15 @@ namespace com.linkedin.assertion */ record AssertionResult { /** - * The final result, e.g. either SUCCESS or FAILURE. + * The final result, e.g. either SUCCESS, FAILURE, or ERROR. */ @TimeseriesField = {} + @Searchable = {} type: enum AssertionResultType { + /** + * The Assertion has not yet been fully evaluated + */ + INIT /** * The Assertion Succeeded */ @@ -17,6 +22,10 @@ record AssertionResult { * The Assertion Failed */ FAILURE + /** + * The Assertion encountered an Error + */ + ERROR } /** @@ -45,8 +54,13 @@ record AssertionResult { nativeResults: optional map[string, string] /** - * URL where full results are available + * External URL where full results are available. Only present when assertion source is not native. */ externalUrl: optional string + /** + * The error object if AssertionResultType is an Error + */ + error: optional AssertionResultError + } \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionResultError.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionResultError.pdl new file mode 100644 index 0000000000000..e768fe8521942 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionResultError.pdl @@ -0,0 +1,45 @@ +namespace com.linkedin.assertion + +/** + * An error encountered when evaluating an AssertionResult + */ +record AssertionResultError { + /** + * The type of error encountered + */ + type: enum AssertionResultErrorType { + /** + * Source is unreachable + */ + SOURCE_CONNECTION_ERROR + /** + * Source query failed to execute + */ + SOURCE_QUERY_FAILED + /** + * Insufficient data to evaluate the assertion + */ + INSUFFICIENT_DATA + /** + * Invalid parameters were detected + */ + INVALID_PARAMETERS + /** + * Event type not supported by the specified source + */ + INVALID_SOURCE_TYPE + /** + * Unsupported platform + */ + UNSUPPORTED_PLATFORM + /** + * Unknown error + */ + UNKNOWN_ERROR + } + + /** + * Additional metadata depending on the type of error + */ + properties: optional map[string, string] +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionRunEvent.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionRunEvent.pdl index 9e75f96fafd06..14f1204232740 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionRunEvent.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionRunEvent.pdl @@ -1,6 +1,7 @@ namespace com.linkedin.assertion -import com.linkedin.timeseries.TimeseriesAspectBase +import com.linkedin.timeseries.PartitionSpec +import com.linkedin.timeseries.TimeWindowSize import com.linkedin.common.ExternalReference import com.linkedin.common.Urn @@ -12,36 +13,31 @@ import com.linkedin.common.Urn "name": "assertionRunEvent", "type": "timeseries", } -record AssertionRunEvent includes TimeseriesAspectBase { +record AssertionRunEvent { + + /** + * The event timestamp field as epoch at UTC in milli seconds. + */ + @Searchable = { + "fieldName": "lastCompletedTime", + "fieldType": "DATETIME" + } + timestampMillis: long /** * Native (platform-specific) identifier for this run */ - //Multiple assertions could occur in same evaluator run runId: string - /* - * Urn of assertion which is evaluated - */ - @TimeseriesField = {} - assertionUrn: Urn - /* * Urn of entity on which the assertion is applicable */ - //example - dataset urn, if dataset is being asserted @TimeseriesField = {} asserteeUrn: Urn - - /** - * Specification of the batch which this run is evaluating - */ - batchSpec: optional BatchSpec /** * The status of the assertion run as per this timeseries event. */ - // Currently just supports COMPLETE, but should evolve to support other statuses like STARTED, RUNNING, etc. @TimeseriesField = {} status: enum AssertionRunStatus { /** @@ -59,4 +55,33 @@ record AssertionRunEvent includes TimeseriesAspectBase { * Runtime parameters of evaluation */ runtimeContext: optional map[string, string] + + /** + * Specification of the batch which this run is evaluating + */ + batchSpec: optional BatchSpec + + /* + * Urn of assertion which is evaluated + */ + @TimeseriesField = {} + assertionUrn: Urn + + /** + * Granularity of the event if applicable + */ + eventGranularity: optional TimeWindowSize + + /** + * The optional partition specification. + */ + partitionSpec: optional PartitionSpec = { + "type":"FULL_TABLE", + "partition":"FULL_TABLE_SNAPSHOT" + } + + /** + * The optional messageId, if provided serves as a custom user-defined unique identifier for an aspect value. + */ + messageId: optional string } \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionSource.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionSource.pdl new file mode 100644 index 0000000000000..d8892c0c71c6f --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionSource.pdl @@ -0,0 +1,27 @@ +namespace com.linkedin.assertion + +/** + * The source of an assertion + */ +record AssertionSource { + /** + * The type of the Assertion Source + */ + @Searchable = { + "fieldName": "sourceType" + } + type: enum AssertionSourceType { + /** + * The assertion was defined natively on DataHub by a user. + */ + NATIVE + /** + * The assertion was defined and managed externally of DataHub. + */ + EXTERNAL + /** + * The assertion was inferred, e.g. from offline AI / ML models. + */ + INFERRED + } +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionStdAggregation.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionStdAggregation.pdl index b79b96f9379b0..968944165a1c8 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionStdAggregation.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionStdAggregation.pdl @@ -4,6 +4,7 @@ namespace com.linkedin.assertion * The function that is applied to the aggregation input (schema, rows, column values) before evaluating an operator. */ enum AssertionStdAggregation { + /** * Assertion is applied on number of rows. */ @@ -20,7 +21,7 @@ enum AssertionStdAggregation { COLUMN_COUNT /** - * Assertion is applied on individual column value. + * Assertion is applied on individual column value. (No aggregation) */ IDENTITY @@ -42,6 +43,13 @@ enum AssertionStdAggregation { /** * Assertion is applied on proportion of distinct values in column */ + UNIQUE_PROPORTION + + /** + * Assertion is applied on proportion of distinct values in column + * + * Deprecated! Use UNIQUE_PROPORTION instead. + */ UNIQUE_PROPOTION /** diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionValueChangeType.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionValueChangeType.pdl new file mode 100644 index 0000000000000..5a1ff4fa73ffb --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionValueChangeType.pdl @@ -0,0 +1,16 @@ +namespace com.linkedin.assertion + +/** +* An enum to represent a type of change in an assertion value, metric, or measurement. +*/ +enum AssertionValueChangeType { + /** + * A change that is defined in absolute terms. + */ + ABSOLUTE + /** + * A change that is defined in relative terms using percentage change + * from the original value. + */ + PERCENTAGE +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AuditLogSpec.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AuditLogSpec.pdl new file mode 100644 index 0000000000000..4d5bf261cbf89 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AuditLogSpec.pdl @@ -0,0 +1,18 @@ +namespace com.linkedin.assertion + +import com.linkedin.schema.SchemaFieldDataType + +/** +* Information about the Audit Log operation to use in evaluating an assertion. +**/ +record AuditLogSpec { + /** + * The list of operation types that should be monitored. If not provided, a default set will be used. + */ + operationTypes: optional array [string] + + /** + * Optional: The user name associated with the operation. + */ + userName: optional string +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/DatasetAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/DatasetAssertionInfo.pdl index c411c7ff8a572..2a8bf28f1ff11 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/assertion/DatasetAssertionInfo.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/DatasetAssertionInfo.pdl @@ -18,9 +18,10 @@ record DatasetAssertionInfo { /** * Scope of the Assertion. What part of the dataset does this assertion apply to? **/ + @Searchable = {} scope: enum DatasetAssertionScope { /** - * This assertion applies to dataset columns + * This assertion applies to dataset column(s) */ DATASET_COLUMN @@ -29,6 +30,11 @@ record DatasetAssertionInfo { */ DATASET_ROWS + /** + * This assertion applies to the storage size of the dataset + */ + DATASET_STORAGE_SIZE + /** * This assertion applies to the schema of the dataset */ @@ -41,7 +47,9 @@ record DatasetAssertionInfo { } /** - * One or more dataset schema fields that are targeted by this assertion + * One or more dataset schema fields that are targeted by this assertion. + * + * This field is expected to be provided if the assertion scope is DATASET_COLUMN. */ @Relationship = { "/*": { @@ -49,11 +57,18 @@ record DatasetAssertionInfo { "entityTypes": [ "schemaField" ] } } + @Searchable = { + "/*": { + "fieldType": "URN" + } + } fields: optional array[Urn] /** * Standardized assertion operator + * This field is left blank if there is no selected aggregation or metric for a particular column. */ + @Searchable = {} aggregation: optional AssertionStdAggregation /** diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/FixedIntervalSchedule.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/FixedIntervalSchedule.pdl new file mode 100644 index 0000000000000..c08c33ffb92d3 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/FixedIntervalSchedule.pdl @@ -0,0 +1,10 @@ +namespace com.linkedin.assertion + +import com.linkedin.common.Urn +import com.linkedin.timeseries.TimeWindowSize + +/** +* Attributes defining a relative fixed interval SLA schedule. +*/ +record FixedIntervalSchedule includes TimeWindowSize { +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessAssertionInfo.pdl new file mode 100644 index 0000000000000..4445a11ff40a7 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessAssertionInfo.pdl @@ -0,0 +1,53 @@ +namespace com.linkedin.assertion + +import com.linkedin.common.Urn +import com.linkedin.dataset.DatasetFilter + +/** +* Attributes defining a Freshness Assertion. +**/ +record FreshnessAssertionInfo { + /** + * The type of the freshness assertion being monitored. + */ + @Searchable = {} + type: enum FreshnessAssertionType { + /** + * An Freshness based on Operations performed on a particular Dataset (insert, update, delete, etc) and sourced from an audit log, as + * opposed to based on the highest watermark in a timestamp column (e.g. a query). Only valid when entity is of type "dataset". + */ + DATASET_CHANGE + /** + * An Freshness based on a successful execution of a Data Job. + */ + DATA_JOB_RUN + } + + /** + * The entity targeted by this Freshness check. + */ + @Searchable = { + "fieldType": "URN" + } + @Relationship = { + "name": "Asserts", + "entityTypes": [ "dataset", "dataJob" ] + } + entity: Urn + + /** + * Produce FAILURE Assertion Result if the asset is not updated on the cadence and within the time range described by the schedule. + */ + @Searchable = { + "/type": { + "fieldName": "scheduleType" + } + } + schedule: FreshnessAssertionSchedule + + /** + * A definition of the specific filters that should be applied, when performing monitoring. + * If not provided, there is no filter, and the full table is under consideration. + */ + filter: optional DatasetFilter +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessAssertionSchedule.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessAssertionSchedule.pdl new file mode 100644 index 0000000000000..a87342ad4f5ed --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessAssertionSchedule.pdl @@ -0,0 +1,66 @@ +namespace com.linkedin.assertion + +import com.linkedin.common.Urn + +/** +* Attributes defining a single Freshness schedule. +*/ +record FreshnessAssertionSchedule { + + /** + * The type of a Freshness Assertion Schedule. + * + * Once we support data-time-relative schedules (e.g. schedules relative to time partitions), + * we will add those schedule types here. + */ + type: enum FreshnessAssertionScheduleType { + /** + * An highly configurable recurring schedule which describes the times of events described + * by a CRON schedule, with the evaluation schedule assuming to be matching the cron schedule. + * + * In a CRON schedule type, we compute the look-back window to be the time between the last scheduled event + * and the current event (evaluation time). This means that the evaluation schedule must match exactly + * the schedule defined inside the cron schedule. + * + * For example, a CRON schedule defined as "0 8 * * *" would represent a schedule of "every day by 8am". Assuming + * that the assertion evaluation schedule is defined to match this, the freshness assertion would be evaluated in the following way: + * + * 1. Compute the "last scheduled occurrence" of the event using the CRON schedule. For example, yesterday at 8am. + * 2. Compute the bounds of a time window between the "last scheduled occurrence" (yesterday at 8am) until the "current occurrence" (today at 8am) + * 3. Verify that the target event has occurred within the CRON-interval window. + * 4. If the target event has occurred within the time window, then assertion passes. + * 5. If the target event has not occurred within the time window, then the assertion fails. + * + */ + CRON + /** + * A fixed interval which is used to compute a look-back window for use when evaluating the assertion relative + * to the Evaluation Time of the Assertion. + * + * To compute the valid look-back window, we subtract the fixed interval from the evaluation time. Then, we verify + * that the target event has occurred within that window. + * + * For example, a fixed interval of "24h" would represent a schedule of "in the last 24 hours". + * The 24 hour interval is relative to the evaluation time of the assertion. For example if we schedule the assertion + * to be evaluated each hour, we'd compute the result as follows: + * + * 1. Subtract the fixed interval from the current time (Evaluation time) to compute the bounds of a fixed look-back window. + * 2. Verify that the target event has occurred within the CRON-interval window. + * 3. If the target event has occurred within the time window, then assertion passes. + * 4. If the target event has not occurred within the time window, then the assertion fails. + * + */ + FIXED_INTERVAL + } + + /** + * A cron schedule. This field is required when type is CRON. + */ + cron: optional FreshnessCronSchedule + + /** + * A fixed interval schedule. This field is required when type is FIXED_INTERVAL. + */ + fixedInterval: optional FixedIntervalSchedule + +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessCronSchedule.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessCronSchedule.pdl new file mode 100644 index 0000000000000..d48900690c51d --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessCronSchedule.pdl @@ -0,0 +1,25 @@ +namespace com.linkedin.assertion + +/** +* Attributes defining a CRON-formatted schedule used for defining a freshness assertion. +*/ +record FreshnessCronSchedule { + /** + * A cron-formatted execution interval, as a cron string, e.g. 1 * * * * + */ + cron: string + + /** + * Timezone in which the cron interval applies, e.g. America/Los Angeles + */ + timezone: string + + /** + * An optional offset in milliseconds to SUBTRACT from the timestamp generated by the cron schedule + * to generate the lower bounds of the "freshness window", or the window of time in which an event must have occurred in order for the Freshness check + * to be considering passing. + * + * If left empty, the start of the SLA window will be the _end_ of the previously evaluated Freshness window. + */ + windowStartOffsetMs: optional long +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessFieldKind.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessFieldKind.pdl new file mode 100644 index 0000000000000..7b25589e500da --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessFieldKind.pdl @@ -0,0 +1,17 @@ +namespace com.linkedin.assertion + +enum FreshnessFieldKind { + /** + * Determine that a change has occurred by inspecting an last modified field which + * represents the last time at which a row was changed. + */ + LAST_MODIFIED, + /** + * Determine that a change has occurred by inspecting a field which should be tracked as the + * "high watermark" for the table. This should be an ascending number or date field. + * + * If rows with this column have not been added since the previous check + * then the Freshness Assertion will fail. + */ + HIGH_WATERMARK +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessFieldSpec.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessFieldSpec.pdl new file mode 100644 index 0000000000000..04acd1c71352d --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessFieldSpec.pdl @@ -0,0 +1,14 @@ +namespace com.linkedin.assertion + +import com.linkedin.schema.SchemaFieldSpec + + +/** +* Lightweight spec used for referencing a particular schema field. +**/ +record FreshnessFieldSpec includes SchemaFieldSpec { + /** + * The type of the field being used to verify the Freshness Assertion. + */ + kind: optional FreshnessFieldKind +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentFieldTransformer.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentFieldTransformer.pdl new file mode 100644 index 0000000000000..d1d3e7b23b666 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentFieldTransformer.pdl @@ -0,0 +1,60 @@ +namespace com.linkedin.assertion + +/** +* The definition of the transformer function that should be applied to a given field / column value in a dataset +* in order to determine the segment or bucket that it belongs to, which in turn is used to evaluate +* volume assertions. +*/ +record IncrementingSegmentFieldTransformer { + /** + * A 'standard' transformer type. Note that not all source systems will support all operators. + */ + type: enum IncrementingSegmentFieldTransformerType { + /** + * Rounds a timestamp (in seconds) down to the start of the month. + */ + TIMESTAMP_MS_TO_MINUTE + + /** + * Rounds a timestamp (in milliseconds) down to the nearest hour. + */ + TIMESTAMP_MS_TO_HOUR + + /** + * Rounds a timestamp (in milliseconds) down to the start of the day. + */ + TIMESTAMP_MS_TO_DATE + + /** + * Rounds a timestamp (in milliseconds) down to the start of the month + */ + TIMESTAMP_MS_TO_MONTH + + /** + * Rounds a timestamp (in milliseconds) down to the start of the year + */ + TIMESTAMP_MS_TO_YEAR + + /** + * Rounds a numeric value down to the nearest integer. + */ + FLOOR + + /** + * Rounds a numeric value up to the nearest integer. + */ + CEILING + + /** + * A backdoor to provide a native operator type specific to a given source system like + * Snowflake, Redshift, BQ, etc. + */ + NATIVE + } + + /** + * The 'native' transformer type, useful as a back door if a custom operator is required. + * This field is required if the type is NATIVE. + */ + nativeType: optional string +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentRowCountChange.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentRowCountChange.pdl new file mode 100644 index 0000000000000..7c4c73f2ea887 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentRowCountChange.pdl @@ -0,0 +1,33 @@ +namespace com.linkedin.assertion + + +/** +* Attributes defining an INCREMENTING_SEGMENT_ROW_COUNT_CHANGE volume assertion. +*/ +record IncrementingSegmentRowCountChange { + /** + * A specification of how the 'segment' can be derived using a column and an optional transformer function. + */ + segment: IncrementingSegmentSpec + + /** + * The type of the value used to evaluate the assertion: a fixed absolute value or a relative percentage. + */ + type: AssertionValueChangeType + + /** + * The operator you'd like to apply to the row count value + * + * Note that only numeric operators are valid inputs: + * GREATER_THAN, GREATER_THAN_OR_EQUAL_TO, EQUAL_TO, LESS_THAN, LESS_THAN_OR_EQUAL_TO, + * BETWEEN. + */ + operator: AssertionStdOperator + + /** + * The parameters you'd like to provide as input to the operator. + * + * Note that only numeric parameter types are valid inputs: NUMBER. + */ + parameters: AssertionStdParameters +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentRowCountTotal.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentRowCountTotal.pdl new file mode 100644 index 0000000000000..6b035107aae09 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentRowCountTotal.pdl @@ -0,0 +1,27 @@ +namespace com.linkedin.assertion + +/** +* Attributes defining an INCREMENTING_SEGMENT_ROW_COUNT_TOTAL volume assertion. +*/ +record IncrementingSegmentRowCountTotal { + /** + * A specification of how the 'segment' can be derived using a column and an optional transformer function. + */ + segment: IncrementingSegmentSpec + + /** + * The operator you'd like to apply. + * + * Note that only numeric operators are valid inputs: + * GREATER_THAN, GREATER_THAN_OR_EQUAL_TO, EQUAL_TO, LESS_THAN, LESS_THAN_OR_EQUAL_TO, + * BETWEEN. + */ + operator: AssertionStdOperator + + /** + * The parameters you'd like to provide as input to the operator. + * + * Note that only numeric parameter types are valid inputs: NUMBER. + */ + parameters: AssertionStdParameters +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentSpec.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentSpec.pdl new file mode 100644 index 0000000000000..eddd0c3da3df7 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/IncrementingSegmentSpec.pdl @@ -0,0 +1,33 @@ +namespace com.linkedin.assertion + +import com.linkedin.schema.SchemaFieldSpec + +/** +* Core attributes required to identify an incrementing segment in a table. This type is mainly useful +* for tables that constantly increase with new rows being added on a particular cadence (e.g. fact or event tables) +* +* An incrementing segment represents a logical chunk of data which is INSERTED +* into a dataset on a regular interval, along with the presence of a constantly-incrementing column +* value such as an event time, date partition, or last modified column. +* +* An incrementing segment is principally identified by 2 key attributes combined: +* +* 1. A field or column that represents the incrementing value. New rows that are inserted will be identified using this column. +* Note that the value of this column may not by itself represent the "bucket" or the "segment" in which the row falls. +* +* 2. [Optional] An transformer function that may be applied to the selected column value in order +* to obtain the final "segment identifier" or "bucket identifier". Rows that have the same value after applying the transformation +* will be grouped into the same segment, using which the final value (e.g. row count) will be determined. +*/ +record IncrementingSegmentSpec { + /** + * The field to use to generate segments. It must be constantly incrementing as new rows are inserted. + */ + field: SchemaFieldSpec + + /** + * Optional transformer function to apply to the field in order to obtain the final segment or bucket identifier. + * If not provided, then no operator will be applied to the field. (identity function) + */ + transformer: optional IncrementingSegmentFieldTransformer +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/RowCountChange.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/RowCountChange.pdl new file mode 100644 index 0000000000000..85a915066f584 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/RowCountChange.pdl @@ -0,0 +1,27 @@ +namespace com.linkedin.assertion + +/** +* Attributes defining a ROW_COUNT_CHANGE volume assertion. +*/ +record RowCountChange { + /** + * The type of the value used to evaluate the assertion: a fixed absolute value or a relative percentage. + */ + type: AssertionValueChangeType + + /** + * The operator you'd like to apply. + * + * Note that only numeric operators are valid inputs: + * GREATER_THAN, GREATER_THAN_OR_EQUAL_TO, EQUAL_TO, LESS_THAN, LESS_THAN_OR_EQUAL_TO, + * BETWEEN. + */ + operator: AssertionStdOperator + + /** + * The parameters you'd like to provide as input to the operator. + * + * Note that only numeric parameter types are valid inputs: NUMBER. + */ + parameters: AssertionStdParameters +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/RowCountTotal.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/RowCountTotal.pdl new file mode 100644 index 0000000000000..f691f15f62e04 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/RowCountTotal.pdl @@ -0,0 +1,22 @@ +namespace com.linkedin.assertion + +/** +* Attributes defining a ROW_COUNT_TOTAL volume assertion. +*/ +record RowCountTotal { + /** + * The operator you'd like to apply. + * + * Note that only numeric operators are valid inputs: + * GREATER_THAN, GREATER_THAN_OR_EQUAL_TO, EQUAL_TO, LESS_THAN, LESS_THAN_OR_EQUAL_TO, + * BETWEEN. + */ + operator: AssertionStdOperator + + /** + * The parameters you'd like to provide as input to the operator. + * + * Note that only numeric parameter types are valid inputs: NUMBER. + */ + parameters: AssertionStdParameters +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/SchemaAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/SchemaAssertionInfo.pdl new file mode 100644 index 0000000000000..fd246e0c7cfc4 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/SchemaAssertionInfo.pdl @@ -0,0 +1,29 @@ +namespace com.linkedin.assertion + +import com.linkedin.common.Urn +import com.linkedin.schema.SchemaMetadata + +/** +* Attributes that are applicable to schema assertions +**/ +record SchemaAssertionInfo { + /** + * The entity targeted by the assertion + */ + @Searchable = { + "fieldType": "URN" + } + @Relationship = { + "name": "Asserts", + "entityTypes": [ "dataset", "dataJob" ] + } + entity: Urn + + /** + * A definition of the expected structure for the asset + * + * Note that many of the fields of this model, especially those related to metadata (tags, terms) + * will go unused in this context. + */ + schema: SchemaMetadata +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/VolumeAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/VolumeAssertionInfo.pdl new file mode 100644 index 0000000000000..327b76f95762e --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/VolumeAssertionInfo.pdl @@ -0,0 +1,82 @@ +namespace com.linkedin.assertion + +import com.linkedin.common.Urn +import com.linkedin.dataset.DatasetFilter + +/** +* Attributes defining a dataset Volume Assertion +*/ +record VolumeAssertionInfo { + /** + * The type of the freshness assertion being monitored. + */ + @Searchable = {} + type: enum VolumeAssertionType { + /** + * A volume assertion that is evaluated against the total row count of a dataset. + */ + ROW_COUNT_TOTAL + /** + * A volume assertion that is evaluated against an incremental row count of a dataset, + * or a row count change. + */ + ROW_COUNT_CHANGE + /** + * A volume assertion that checks the latest "segment" in a table based on an incrementing + * column to check whether it's row count falls into a particular range. + * + * This can be used to monitor the row count of an incrementing date-partition column segment. + */ + INCREMENTING_SEGMENT_ROW_COUNT_TOTAL + /** + * A volume assertion that compares the row counts in neighboring "segments" or "partitions" + * of an incrementing column. + * This can be used to track changes between subsequent date partition + * in a table, for example. + */ + INCREMENTING_SEGMENT_ROW_COUNT_CHANGE + } + + /** + * The entity targeted by this Volume check. + */ + @Searchable = { + "fieldType": "URN" + } + @Relationship = { + "name": "Asserts", + "entityTypes": [ "dataset" ] + } + entity: Urn + + /** + * Produce FAILURE Assertion Result if the row count of the asset does not meet specific requirements. + * Required if type is 'ROW_COUNT_TOTAL' + */ + rowCountTotal: optional RowCountTotal + + /** + * Produce FAILURE Assertion Result if the delta row count of the asset does not meet specific requirements + * within a given period of time. + * Required if type is 'ROW_COUNT_CHANGE' + */ + rowCountChange: optional RowCountChange + + /** + * Produce FAILURE Assertion Result if the asset's latest incrementing segment row count total + * does not meet specific requirements. Required if type is 'INCREMENTING_SEGMENT_ROW_COUNT_TOTAL' + */ + incrementingSegmentRowCountTotal: optional IncrementingSegmentRowCountTotal + + /** + * Produce FAILURE Assertion Result if the asset's incrementing segment row count delta + * does not meet specific requirements. Required if type is 'INCREMENTING_SEGMENT_ROW_COUNT_CHANGE' + */ + incrementingSegmentRowCountChange: optional IncrementingSegmentRowCountChange + + /** + * A definition of the specific filters that should be applied, when performing monitoring. + * If not provided, there is no filter, and the full table is under consideration. + */ + filter: optional DatasetFilter +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/datacontract/DataContractProperties.pdl b/metadata-models/src/main/pegasus/com/linkedin/datacontract/DataContractProperties.pdl new file mode 100644 index 0000000000000..a623f585df30c --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/datacontract/DataContractProperties.pdl @@ -0,0 +1,59 @@ +namespace com.linkedin.datacontract + +import com.linkedin.common.Urn + +/** + * Information about a data contract + */ +@Aspect = { + "name": "dataContractProperties" +} +record DataContractProperties { + /** + * The entity that this contract is associated with. Currently, we only support Dataset contracts, but + * in the future we may also support Data Product level contracts. + */ + @Relationship = { + "name": "ContractFor", + "entityTypes": [ "dataset" ] + } + entity: Urn + + /** + * An optional set of schema contracts. If this is a dataset contract, there will only be one. + */ + @Relationship = { + "/*/assertion": { + "name": "IncludesSchemaAssertion", + "entityTypes": [ "assertion" ] + } + } + schema: optional array[SchemaContract] + + /** + * An optional set of FRESHNESS contracts. If this is a dataset contract, there will only be one. + */ + @Relationship = { + "/*/assertion": { + "name": "IncludesFreshnessAssertion", + "entityTypes": [ "assertion" ] + } + } + freshness: optional array[FreshnessContract] + + /** + * An optional set of Data Quality contracts, e.g. table and column level contract constraints. + */ + @Relationship = { + "/*/assertion": { + "name": "IncludesDataQualityAssertion", + "entityTypes": [ "assertion" ] + } + } + dataQuality: optional array[DataQualityContract] + + /** + * YAML-formatted contract definition + */ + rawContract: optional string +} diff --git a/metadata-models/src/main/pegasus/com/linkedin/datacontract/DataContractStatus.pdl b/metadata-models/src/main/pegasus/com/linkedin/datacontract/DataContractStatus.pdl new file mode 100644 index 0000000000000..d61fb191ae53d --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/datacontract/DataContractStatus.pdl @@ -0,0 +1,27 @@ +namespace com.linkedin.datacontract + +import com.linkedin.common.Urn +import com.linkedin.common.CustomProperties + +/** + * Information about the status of a data contract + */ +@Aspect = { + "name": "dataContractStatus" +} +record DataContractStatus includes CustomProperties { + /** + * The latest state of the data contract + */ + @Searchable = {} + state: enum DataContractState { + /** + * The data contract is active. + */ + ACTIVE + /** + * The data contract is pending implementation. + */ + PENDING + } +} diff --git a/metadata-models/src/main/pegasus/com/linkedin/datacontract/DataQualityContract.pdl b/metadata-models/src/main/pegasus/com/linkedin/datacontract/DataQualityContract.pdl new file mode 100644 index 0000000000000..273d2c2a56f95 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/datacontract/DataQualityContract.pdl @@ -0,0 +1,16 @@ +namespace com.linkedin.datacontract + +import com.linkedin.common.Urn + + +/** + * A data quality contract pertaining to a physical data asset + * Data Quality contracts are used to make assertions about data quality metrics for a physical data asset + */ +record DataQualityContract { + /** + * The assertion representing the Data Quality contract. + * E.g. a table or column-level assertion. + */ + assertion: Urn +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/datacontract/FreshnessContract.pdl b/metadata-models/src/main/pegasus/com/linkedin/datacontract/FreshnessContract.pdl new file mode 100644 index 0000000000000..8cfa66846d505 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/datacontract/FreshnessContract.pdl @@ -0,0 +1,13 @@ +namespace com.linkedin.datacontract + +import com.linkedin.common.Urn + +/** + * A contract pertaining to the operational SLAs of a physical data asset + */ +record FreshnessContract { + /** + * The assertion representing the SLA contract. + */ + assertion: Urn +} diff --git a/metadata-models/src/main/pegasus/com/linkedin/datacontract/SchemaContract.pdl b/metadata-models/src/main/pegasus/com/linkedin/datacontract/SchemaContract.pdl new file mode 100644 index 0000000000000..6c11e0da5b128 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/datacontract/SchemaContract.pdl @@ -0,0 +1,13 @@ +namespace com.linkedin.datacontract + +import com.linkedin.common.Urn + +/** + * Expectations for a logical schema + */ +record SchemaContract { + /** + * The assertion representing the schema contract. + */ + assertion: Urn +} diff --git a/metadata-models/src/main/pegasus/com/linkedin/dataset/DatasetFilter.pdl b/metadata-models/src/main/pegasus/com/linkedin/dataset/DatasetFilter.pdl new file mode 100644 index 0000000000000..6823398f79f3d --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/dataset/DatasetFilter.pdl @@ -0,0 +1,30 @@ +namespace com.linkedin.dataset + +/** + * A definition of filters that should be used when + * querying an external Dataset or Table. + * + * Note that this models should NOT be used for working with + * search / filter on DataHub Platform itself. + */ +record DatasetFilter { + /** + * How the partition will be represented in this model. + * + * In the future, we'll likely add support for more structured + * predicates. + */ + type: enum DatasetFilterType { + /** + * The partition is represented as a an opaque, raw SQL + * clause. + */ + SQL + } + + /** + * The raw where clause string which will be used for monitoring. + * Required if the type is SQL. + */ + sql: optional string +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/metadata/key/DataContractKey.pdl b/metadata-models/src/main/pegasus/com/linkedin/metadata/key/DataContractKey.pdl new file mode 100644 index 0000000000000..f1d4a709cd6bf --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/metadata/key/DataContractKey.pdl @@ -0,0 +1,14 @@ +namespace com.linkedin.metadata.key + +/** + * Key for a Data Contract + */ +@Aspect = { + "name": "dataContractKey" +} +record DataContractKey { + /** + * Unique id for the contract + */ + id: string +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/schema/SchemaFieldSpec.pdl b/metadata-models/src/main/pegasus/com/linkedin/schema/SchemaFieldSpec.pdl new file mode 100644 index 0000000000000..e875ff7a84403 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/schema/SchemaFieldSpec.pdl @@ -0,0 +1,21 @@ +namespace com.linkedin.schema + +/** +* Lightweight spec used for referencing a particular schema field. +**/ +record SchemaFieldSpec { + /** + * The field path + */ + path: string + + /** + * The DataHub standard schema field type. + */ + type: string + + /** + * The native field type + */ + nativeType: string +} \ No newline at end of file diff --git a/metadata-models/src/main/resources/entity-registry.yml b/metadata-models/src/main/resources/entity-registry.yml index 56fc5f6568eb7..11d0f74305d7b 100644 --- a/metadata-models/src/main/resources/entity-registry.yml +++ b/metadata-models/src/main/resources/entity-registry.yml @@ -262,6 +262,7 @@ entities: - assertionInfo - dataPlatformInstance - assertionRunEvent + - assertionActions - status - name: dataHubRetention category: internal @@ -457,4 +458,12 @@ entities: aspects: - ownershipTypeInfo - status + - name: dataContract + category: core + keyAspect: dataContractKey + aspects: + - dataContractProperties + - dataContractStatus + - status + events: From 2bc685d3b98f879d1c3051a8484a78489359d910 Mon Sep 17 00:00:00 2001 From: Aseem Bansal Date: Thu, 5 Oct 2023 09:31:32 +0530 Subject: [PATCH 092/156] ci: tweak ci to decrease wait time of devs (#8945) --- .github/workflows/build-and-test.yml | 14 ++++++++++---- .github/workflows/metadata-ingestion.yml | 7 ++++--- .../integration/powerbi/test_admin_only_api.py | 3 +++ .../tests/integration/powerbi/test_m_parser.py | 2 +- .../tests/integration/powerbi/test_powerbi.py | 2 +- .../tests/integration/snowflake/test_snowflake.py | 4 ++-- .../integration/tableau/test_tableau_ingest.py | 2 +- .../tests/integration/trino/test_trino.py | 5 ++--- 8 files changed, 24 insertions(+), 15 deletions(-) diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml index 3f409878b191f..96b9bb2a14933 100644 --- a/.github/workflows/build-and-test.yml +++ b/.github/workflows/build-and-test.yml @@ -27,8 +27,8 @@ jobs: command: [ # metadata-ingestion and airflow-plugin each have dedicated build jobs - "./gradlew build -x :metadata-ingestion:build -x :metadata-ingestion:check -x docs-website:build -x :metadata-integration:java:spark-lineage:test -x :metadata-io:test -x :metadata-ingestion-modules:airflow-plugin:build -x :metadata-ingestion-modules:airflow-plugin:check -x :datahub-frontend:build -x :datahub-web-react:build --parallel", - "./gradlew :datahub-frontend:build :datahub-web-react:build --parallel", + "except_metadata_ingestion", + "frontend" ] timezone: [ @@ -53,9 +53,15 @@ jobs: with: python-version: "3.10" cache: pip - - name: Gradle build (and test) + - name: Gradle build (and test) for metadata ingestion + # we only need the timezone runs for frontend tests + if: ${{ matrix.command == 'except_metadata_ingestion' && matrix.timezone == 'America/New_York' }} run: | - ${{ matrix.command }} + ./gradlew build -x :metadata-ingestion:build -x :metadata-ingestion:check -x docs-website:build -x :metadata-integration:java:spark-lineage:test -x :metadata-io:test -x :metadata-ingestion-modules:airflow-plugin:build -x :metadata-ingestion-modules:airflow-plugin:check -x :datahub-frontend:build -x :datahub-web-react:build --parallel + - name: Gradle build (and test) for frontend + if: ${{ matrix.command == 'frontend' }} + run: | + ./gradlew :datahub-frontend:build :datahub-web-react:build --parallel env: NODE_OPTIONS: "--max-old-space-size=3072" - uses: actions/upload-artifact@v3 diff --git a/.github/workflows/metadata-ingestion.yml b/.github/workflows/metadata-ingestion.yml index 8d56a0adf5bd5..dea4603868f8e 100644 --- a/.github/workflows/metadata-ingestion.yml +++ b/.github/workflows/metadata-ingestion.yml @@ -34,7 +34,6 @@ jobs: python-version: ["3.7", "3.10"] command: [ - "lint", "testQuick", "testIntegrationBatch0", "testIntegrationBatch1", @@ -54,6 +53,9 @@ jobs: run: ./metadata-ingestion/scripts/install_deps.sh - name: Install package run: ./gradlew :metadata-ingestion:installPackageOnly + - name: Run lint alongwith testQuick + if: ${{ matrix.command == 'testQuick' }} + run: ./gradlew :metadata-ingestion:lint - name: Run metadata-ingestion tests run: ./gradlew :metadata-ingestion:${{ matrix.command }} - name: Debug info @@ -65,7 +67,6 @@ jobs: docker image ls docker system df - uses: actions/upload-artifact@v3 - if: ${{ always() && matrix.command != 'lint' }} with: name: Test Results (metadata ingestion ${{ matrix.python-version }}) path: | @@ -73,7 +74,7 @@ jobs: **/build/test-results/test/** **/junit.*.xml - name: Upload coverage to Codecov - if: ${{ always() && matrix.python-version == '3.10' && matrix.command != 'lint' }} + if: ${{ always() && matrix.python-version == '3.10' }} uses: codecov/codecov-action@v3 with: token: ${{ secrets.CODECOV_TOKEN }} diff --git a/metadata-ingestion/tests/integration/powerbi/test_admin_only_api.py b/metadata-ingestion/tests/integration/powerbi/test_admin_only_api.py index f95fd81681a9a..6f45dcf97f1dd 100644 --- a/metadata-ingestion/tests/integration/powerbi/test_admin_only_api.py +++ b/metadata-ingestion/tests/integration/powerbi/test_admin_only_api.py @@ -3,11 +3,14 @@ from typing import Any, Dict from unittest import mock +import pytest from freezegun import freeze_time from datahub.ingestion.run.pipeline import Pipeline from tests.test_helpers import mce_helpers +pytestmark = pytest.mark.integration_batch_2 + FROZEN_TIME = "2022-02-03 07:00:00" diff --git a/metadata-ingestion/tests/integration/powerbi/test_m_parser.py b/metadata-ingestion/tests/integration/powerbi/test_m_parser.py index 2e9c02ef759a5..e3cc6c8101650 100644 --- a/metadata-ingestion/tests/integration/powerbi/test_m_parser.py +++ b/metadata-ingestion/tests/integration/powerbi/test_m_parser.py @@ -19,7 +19,7 @@ from datahub.ingestion.source.powerbi.m_query.resolver import DataPlatformTable, Lineage from datahub.utilities.sqlglot_lineage import ColumnLineageInfo, DownstreamColumnRef -pytestmark = pytest.mark.slow +pytestmark = pytest.mark.integration_batch_2 M_QUERIES = [ 'let\n Source = Snowflake.Databases("bu10758.ap-unknown-2.fakecomputing.com","PBI_TEST_WAREHOUSE_PROD",[Role="PBI_TEST_MEMBER"]),\n PBI_TEST_Database = Source{[Name="PBI_TEST",Kind="Database"]}[Data],\n TEST_Schema = PBI_TEST_Database{[Name="TEST",Kind="Schema"]}[Data],\n TESTTABLE_Table = TEST_Schema{[Name="TESTTABLE",Kind="Table"]}[Data]\nin\n TESTTABLE_Table', diff --git a/metadata-ingestion/tests/integration/powerbi/test_powerbi.py b/metadata-ingestion/tests/integration/powerbi/test_powerbi.py index b0695e3ea9954..7232d2a38da1d 100644 --- a/metadata-ingestion/tests/integration/powerbi/test_powerbi.py +++ b/metadata-ingestion/tests/integration/powerbi/test_powerbi.py @@ -21,7 +21,7 @@ ) from tests.test_helpers import mce_helpers -pytestmark = pytest.mark.slow +pytestmark = pytest.mark.integration_batch_2 FROZEN_TIME = "2022-02-03 07:00:00" diff --git a/metadata-ingestion/tests/integration/snowflake/test_snowflake.py b/metadata-ingestion/tests/integration/snowflake/test_snowflake.py index dec50aefd19f0..2c77ace8b53e5 100644 --- a/metadata-ingestion/tests/integration/snowflake/test_snowflake.py +++ b/metadata-ingestion/tests/integration/snowflake/test_snowflake.py @@ -30,6 +30,8 @@ from tests.integration.snowflake.common import FROZEN_TIME, default_query_results from tests.test_helpers import mce_helpers +pytestmark = pytest.mark.integration_batch_2 + def random_email(): return ( @@ -55,7 +57,6 @@ def random_cloud_region(): ) -@pytest.mark.integration def test_snowflake_basic(pytestconfig, tmp_path, mock_time, mock_datahub_graph): test_resources_dir = pytestconfig.rootpath / "tests/integration/snowflake" @@ -183,7 +184,6 @@ def test_snowflake_basic(pytestconfig, tmp_path, mock_time, mock_datahub_graph): @freeze_time(FROZEN_TIME) -@pytest.mark.integration def test_snowflake_private_link(pytestconfig, tmp_path, mock_time, mock_datahub_graph): test_resources_dir = pytestconfig.rootpath / "tests/integration/snowflake" diff --git a/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py b/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py index 53b8519a886d3..c31867f5aa904 100644 --- a/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py +++ b/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py @@ -757,7 +757,7 @@ def test_tableau_no_verify(): @freeze_time(FROZEN_TIME) -@pytest.mark.slow +@pytest.mark.integration_batch_2 def test_tableau_signout_timeout(pytestconfig, tmp_path, mock_datahub_graph): enable_logging() output_file_name: str = "tableau_signout_timeout_mces.json" diff --git a/metadata-ingestion/tests/integration/trino/test_trino.py b/metadata-ingestion/tests/integration/trino/test_trino.py index 22e5f6f91a06e..177c273c0d242 100644 --- a/metadata-ingestion/tests/integration/trino/test_trino.py +++ b/metadata-ingestion/tests/integration/trino/test_trino.py @@ -13,6 +13,8 @@ from tests.test_helpers import fs_helpers, mce_helpers from tests.test_helpers.docker_helpers import wait_for_port +pytestmark = pytest.mark.integration_batch_1 + FROZEN_TIME = "2021-09-23 12:00:00" data_platform = "trino" @@ -51,7 +53,6 @@ def loaded_trino(trino_runner): @freeze_time(FROZEN_TIME) -@pytest.mark.integration @pytest.mark.xfail def test_trino_ingest( loaded_trino, test_resources_dir, pytestconfig, tmp_path, mock_time @@ -111,7 +112,6 @@ def test_trino_ingest( @freeze_time(FROZEN_TIME) -@pytest.mark.integration def test_trino_hive_ingest( loaded_trino, test_resources_dir, pytestconfig, tmp_path, mock_time ): @@ -167,7 +167,6 @@ def test_trino_hive_ingest( @freeze_time(FROZEN_TIME) -@pytest.mark.integration def test_trino_instance_ingest( loaded_trino, test_resources_dir, pytestconfig, tmp_path, mock_time ): From 2fcced6db9d30228c421d0773c8249c889cd0d9f Mon Sep 17 00:00:00 2001 From: Mayuri Nehate <33225191+mayurinehate@users.noreply.github.com> Date: Thu, 5 Oct 2023 09:31:57 +0530 Subject: [PATCH 093/156] docs(ingest): add permissions required for athena ingestion (#8948) --- .../docs/sources/athena/athena_pre.md | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 metadata-ingestion/docs/sources/athena/athena_pre.md diff --git a/metadata-ingestion/docs/sources/athena/athena_pre.md b/metadata-ingestion/docs/sources/athena/athena_pre.md new file mode 100644 index 0000000000000..a56457d3f84fc --- /dev/null +++ b/metadata-ingestion/docs/sources/athena/athena_pre.md @@ -0,0 +1,72 @@ +### Prerequisities + +In order to execute this source, you will need to create a policy with below permissions and attach it to the the aws role or credentials used in ingestion recipe. + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "VisualEditor0", + "Effect": "Allow", + "Action": [ + "athena:GetTableMetadata", + "athena:StartQueryExecution", + "athena:GetQueryResults", + "athena:GetDatabase", + "athena:ListDataCatalogs", + "athena:GetDataCatalog", + "athena:ListQueryExecutions", + "athena:GetWorkGroup", + "athena:StopQueryExecution", + "athena:GetQueryResultsStream", + "athena:ListDatabases", + "athena:GetQueryExecution", + "athena:ListTableMetadata", + "athena:BatchGetQueryExecution", + "glue:GetTables", + "glue:GetDatabases", + "glue:GetTable", + "glue:GetDatabase", + "glue:SearchTables", + "glue:GetTableVersions", + "glue:GetTableVersion", + "glue:GetPartition", + "glue:GetPartitions", + "s3:GetObject", + "s3:ListBucket", + "s3:GetBucketLocation", + ], + "Resource": [ + "arn:aws:athena:${region-id}:${account-id}:datacatalog/*", + "arn:aws:athena:${region-id}:${account-id}:workgroup/*", + "arn:aws:glue:${region-id}:${account-id}:tableVersion/*/*/*", + "arn:aws:glue:${region-id}:${account-id}:table/*/*", + "arn:aws:glue:${region-id}:${account-id}:catalog", + "arn:aws:glue:${region-id}:${account-id}:database/*", + "arn:aws:s3:::${datasets-bucket}", + "arn:aws:s3:::${datasets-bucket}/*" + ] + }, + { + "Sid": "VisualEditor1", + "Effect": "Allow", + "Action": [ + "s3:PutObject", + "s3:GetObject", + "s3:ListBucketMultipartUploads", + "s3:AbortMultipartUpload", + "s3:ListBucket", + "s3:GetBucketLocation", + "s3:ListMultipartUploadParts" + ], + "Resource": [ + "arn:aws:s3:::${athena-query-result-bucket}/*", + "arn:aws:s3:::${athena-query-result-bucket}" + ] + }, + ] +} +``` + +Replace `${var}` with appropriate values as per your athena setup. \ No newline at end of file From 6310e51eb09711e98d86625578127349c5144c66 Mon Sep 17 00:00:00 2001 From: Jinlin Yang <86577891+jinlintt@users.noreply.github.com> Date: Wed, 4 Oct 2023 21:03:31 -0700 Subject: [PATCH 094/156] feat(ingestion/dynamodb): implement pagination for list_tables (#8910) --- .../app/ingest/source/builder/sources.json | 4 +- .../docs/sources/dynamodb/dynamodb_post.md | 13 ++- .../docs/sources/dynamodb/dynamodb_pre.md | 6 +- .../docs/sources/dynamodb/dynamodb_recipe.yml | 16 ++-- .../ingestion/source/dynamodb/dynamodb.py | 85 +++++++++++-------- 5 files changed, 65 insertions(+), 59 deletions(-) diff --git a/datahub-web-react/src/app/ingest/source/builder/sources.json b/datahub-web-react/src/app/ingest/source/builder/sources.json index 1bd5b6f1f768b..b18384909c33f 100644 --- a/datahub-web-react/src/app/ingest/source/builder/sources.json +++ b/datahub-web-react/src/app/ingest/source/builder/sources.json @@ -130,7 +130,7 @@ "name": "dynamodb", "displayName": "DynamoDB", "docsUrl": "https://datahubproject.io/docs/metadata-ingestion/", - "recipe": "source:\n type: dynamodb\n config:\n platform_instance: \"AWS_ACCOUNT_ID\"\n aws_access_key_id : '${AWS_ACCESS_KEY_ID}'\n aws_secret_access_key : '${AWS_SECRET_ACCESS_KEY}'\n # User could use the below option to provide a list of primary keys of a table in dynamodb format,\n # those items from given primary keys will be included when we scan the table.\n # For each table we can retrieve up to 16 MB of data, which can contain as many as 100 items.\n # We'll enforce the the primary keys list size not to exceed 100\n # The total items we'll try to retrieve in these two scenarios:\n # 1. If user don't specify include_table_item: we'll retrieve up to 100 items\n # 2. If user specifies include_table_item: we'll retrieve up to 100 items plus user specified items in\n # the table, with a total not more than 200 items\n # include_table_item:\n # table_name:\n # [\n # {\n # 'partition_key_name': { 'attribute_type': 'attribute_value' },\n # 'sort_key_name': { 'attribute_type': 'attribute_value' },\n # },\n # ]" + "recipe": "source:\n type: dynamodb\n config:\n platform_instance: \"AWS_ACCOUNT_ID\"\n aws_access_key_id : '${AWS_ACCESS_KEY_ID}'\n aws_secret_access_key : '${AWS_SECRET_ACCESS_KEY}'\n # If there are items that have most representative fields of the table, users could use the\n # `include_table_item` option to provide a list of primary keys of the table in dynamodb format.\n # For each `region.table`, the list of primary keys can be at most 100.\n # We include these items in addition to the first 100 items in the table when we scan it.\n # include_table_item:\n # region.table_name:\n # [\n # {\n # 'partition_key_name': { 'attribute_type': 'attribute_value' },\n # 'sort_key_name': { 'attribute_type': 'attribute_value' },\n # },\n # ]" }, { "urn": "urn:li:dataPlatform:glue", @@ -223,4 +223,4 @@ "docsUrl": "https://datahubproject.io/docs/metadata-ingestion/", "recipe": "source:\n type: \n config:\n # Source-type specifics config\n " } -] \ No newline at end of file +] diff --git a/metadata-ingestion/docs/sources/dynamodb/dynamodb_post.md b/metadata-ingestion/docs/sources/dynamodb/dynamodb_post.md index 7f9a0324c7bc6..a1c0a6e2d4d21 100644 --- a/metadata-ingestion/docs/sources/dynamodb/dynamodb_post.md +++ b/metadata-ingestion/docs/sources/dynamodb/dynamodb_post.md @@ -1,21 +1,18 @@ -## Limitations - -For each region, the list table operation returns maximum number 100 tables, we need to further improve it by implementing pagination for listing tables - ## Advanced Configurations ### Using `include_table_item` config -If there are items that have most representative fields of the table, user could use the `include_table_item` option to provide a list of primary keys of a table in dynamodb format, those items from given primary keys will be included when we scan the table. +If there are items that have most representative fields of the table, users could use the `include_table_item` option to provide a list of primary keys of the table in dynamodb format. We include these items in addition to the first 100 items in the table when we scan it. -Take [AWS DynamoDB Developer Guide Example tables and data](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/AppendixSampleTables.html) as an example, if user has a table `Reply` with composite primary key `Id` and `ReplyDateTime`, user can use `include_table_item` to include 2 items as following: +Take [AWS DynamoDB Developer Guide Example tables and data](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/AppendixSampleTables.html) as an example, if a account has a table `Reply` in the `us-west-2` region with composite primary key `Id` and `ReplyDateTime`, users can use `include_table_item` to include 2 items as following: Example: ```yml -# put the table name and composite key in DynamoDB format +# The table name should be in the format of region.table_name +# The primary keys should be in the DynamoDB format include_table_item: - Reply: + us-west-2.Reply: [ { "ReplyDateTime": { "S": "2015-09-22T19:58:22.947Z" }, diff --git a/metadata-ingestion/docs/sources/dynamodb/dynamodb_pre.md b/metadata-ingestion/docs/sources/dynamodb/dynamodb_pre.md index a48e8d5be04aa..598d0ecdb3786 100644 --- a/metadata-ingestion/docs/sources/dynamodb/dynamodb_pre.md +++ b/metadata-ingestion/docs/sources/dynamodb/dynamodb_pre.md @@ -1,8 +1,8 @@ ### Prerequisities -In order to execute this source, you will need to create access key and secret keys that have DynamoDB read access. You can create these policies and attach to your account or can ask your account admin to attach these policies to your account. +In order to execute this source, you need to attach the `AmazonDynamoDBReadOnlyAccess` policy to a user in your AWS account. Then create an API access key and secret for the user. -For access key permissions, you can create a policy with permissions below and attach to your account, you can find more details in [Managing access keys for IAM users](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_access-keys.html) +For a user to be able to create API access key, it needs the following access key permissions. Your AWS account admin can create a policy with these permissions and attach to the user, you can find more details in [Managing access keys for IAM users](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_access-keys.html) ```json { @@ -22,5 +22,3 @@ For access key permissions, you can create a policy with permissions below and a ] } ``` - -For DynamoDB read access, you can simply attach AWS managed policy `AmazonDynamoDBReadOnlyAccess` to your account, you can find more details in [Attaching a policy to an IAM user group](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_groups_manage_attach-policy.html) diff --git a/metadata-ingestion/docs/sources/dynamodb/dynamodb_recipe.yml b/metadata-ingestion/docs/sources/dynamodb/dynamodb_recipe.yml index bd41637907b5c..4f4edc9a7d496 100644 --- a/metadata-ingestion/docs/sources/dynamodb/dynamodb_recipe.yml +++ b/metadata-ingestion/docs/sources/dynamodb/dynamodb_recipe.yml @@ -4,16 +4,14 @@ source: platform_instance: "AWS_ACCOUNT_ID" aws_access_key_id: "${AWS_ACCESS_KEY_ID}" aws_secret_access_key: "${AWS_SECRET_ACCESS_KEY}" - # User could use the below option to provide a list of primary keys of a table in dynamodb format, - # those items from given primary keys will be included when we scan the table. - # For each table we can retrieve up to 16 MB of data, which can contain as many as 100 items. - # We'll enforce the the primary keys list size not to exceed 100 - # The total items we'll try to retrieve in these two scenarios: - # 1. If user don't specify include_table_item: we'll retrieve up to 100 items - # 2. If user specifies include_table_item: we'll retrieve up to 100 items plus user specified items in - # the table, with a total not more than 200 items + # + # If there are items that have most representative fields of the table, users could use the + # `include_table_item` option to provide a list of primary keys of the table in dynamodb format. + # For each `region.table`, the list of primary keys can be at most 100. + # We include these items in addition to the first 100 items in the table when we scan it. + # # include_table_item: - # table_name: + # region.table_name: # [ # { # "partition_key_name": { "attribute_type": "attribute_value" }, diff --git a/metadata-ingestion/src/datahub/ingestion/source/dynamodb/dynamodb.py b/metadata-ingestion/src/datahub/ingestion/source/dynamodb/dynamodb.py index 6b7c118373673..d7f3dfb9279fb 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/dynamodb/dynamodb.py +++ b/metadata-ingestion/src/datahub/ingestion/source/dynamodb/dynamodb.py @@ -1,5 +1,5 @@ import logging -from dataclasses import field +from dataclasses import dataclass, field from typing import Any, Counter, Dict, Iterable, List, Optional, Type, Union import boto3 @@ -79,12 +79,13 @@ class DynamoDBConfig(DatasetSourceConfigMixin, StatefulIngestionConfigBase): table_pattern: AllowDenyPattern = Field( default=AllowDenyPattern.allow_all(), - description="regex patterns for tables to filter in ingestion.", + description="Regex patterns for tables to filter in ingestion. The table name format is 'region.table'", ) # Custom Stateful Ingestion settings stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None +@dataclass class DynamoDBSourceReport(StaleEntityRemovalSourceReport): filtered: List[str] = field(default_factory=list) @@ -175,39 +176,30 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]: # traverse databases in sorted order so output is consistent for region in dynamodb_regions: - try: - # create a new dynamodb client for each region, - # it seems for one client we could only list the table of one specific region, - # the list_tables() method don't take any config that related to region - # TODO: list table returns maximum number 100, need to implement pagination here - dynamodb_client = boto3.client( - "dynamodb", - region_name=region, - aws_access_key_id=self.config.aws_access_key_id - if self.config.aws_access_key_id - else None, - aws_secret_access_key=self.config.aws_secret_access_key.get_secret_value() - if self.config.aws_secret_access_key - else None, - ) - table_names: List[str] = dynamodb_client.list_tables()["TableNames"] - except Exception as ex: - # TODO: If regions is config input then this would be self.report.report_warning, - # we can create dynamodb client to take aws region or regions as user input - logger.info(f"exception happen in region {region}, skipping: {ex}") - continue - for table_name in sorted(table_names): - if not self.config.table_pattern.allowed(table_name): + logger.info(f"Processing region {region}") + # create a new dynamodb client for each region, + # it seems for one client we could only list the table of one specific region, + # the list_tables() method don't take any config that related to region + dynamodb_client = boto3.client( + "dynamodb", + region_name=region, + aws_access_key_id=self.config.aws_access_key_id, + aws_secret_access_key=self.config.aws_secret_access_key.get_secret_value(), + ) + + for table_name in self._list_tables(dynamodb_client): + dataset_name = f"{region}.{table_name}" + if not self.config.table_pattern.allowed(dataset_name): + logger.debug(f"skipping table: {dataset_name}") + self.report.report_dropped(dataset_name) continue + + logger.debug(f"Processing table: {dataset_name}") table_info = dynamodb_client.describe_table(TableName=table_name)[ "Table" ] account_id = table_info["TableArn"].split(":")[4] - if not self.config.table_pattern.allowed(table_name): - self.report.report_dropped(table_name) - continue platform_instance = self.config.platform_instance or account_id - dataset_name = f"{region}.{table_name}" dataset_urn = make_dataset_urn_with_platform_instance( platform=self.platform, platform_instance=platform_instance, @@ -222,7 +214,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]: ) primary_key_dict = self.extract_primary_key_from_key_schema(table_info) table_schema = self.construct_schema_from_dynamodb( - dynamodb_client, table_name + dynamodb_client, region, table_name ) schema_metadata = self.construct_schema_metadata( table_name, @@ -254,9 +246,25 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]: aspect=platform_instance_aspect, ).as_workunit() + def _list_tables( + self, + dynamodb_client: BaseClient, + ) -> Iterable[str]: + # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/paginator/ListTables.html + try: + for page in dynamodb_client.get_paginator("list_tables").paginate(): + table_names = page.get("TableNames") + if table_names: + yield from table_names + except Exception as ex: + # TODO: If regions is config input then this would be self.report.report_warning, + # we can create dynamodb client to take aws region or regions as user input + logger.info(f"Exception happened while listing tables, skipping: {ex}") + def construct_schema_from_dynamodb( self, dynamodb_client: BaseClient, + region: str, table_name: str, ) -> Dict[str, SchemaDescription]: """ @@ -275,7 +283,7 @@ def construct_schema_from_dynamodb( The MaxItems is the total number of items to return, and PageSize is the size of each page, we are assigning same value to these two config. If MaxItems is more than PageSize then we expect MaxItems / PageSize pages in response_iterator will return """ - self.include_table_item_to_schema(dynamodb_client, table_name, schema) + self.include_table_item_to_schema(dynamodb_client, region, table_name, schema) response_iterator = paginator.paginate( TableName=table_name, PaginationConfig={ @@ -294,33 +302,38 @@ def construct_schema_from_dynamodb( def include_table_item_to_schema( self, dynamodb_client: Any, + region: str, table_name: str, schema: Dict[str, SchemaDescription], ) -> None: """ - It will look up in the config include_table_item dict to see if the current table name exists as key, + It will look up in the config include_table_item dict to see if "region.table_name" exists as key, if it exists then get the items by primary key from the table and put it to schema """ if self.config.include_table_item is None: return - if table_name not in self.config.include_table_item.keys(): + dataset_name = f"{region}.{table_name}" + if dataset_name not in self.config.include_table_item.keys(): return - primary_key_list = self.config.include_table_item.get(table_name) + primary_key_list = self.config.include_table_item.get(dataset_name) assert isinstance(primary_key_list, List) if len(primary_key_list) > MAX_PRIMARY_KEYS_SIZE: logger.info( - f"the provided primary keys list size exceeded the max size for table {table_name}, we'll only process the first {MAX_PRIMARY_KEYS_SIZE} items" + f"the provided primary keys list size exceeded the max size for table {dataset_name}, we'll only process the first {MAX_PRIMARY_KEYS_SIZE} items" ) primary_key_list = primary_key_list[0:MAX_PRIMARY_KEYS_SIZE] items = [] response = dynamodb_client.batch_get_item( RequestItems={table_name: {"Keys": primary_key_list}} - ).get("Responses", None) + ).get("Responses") if response is None: logger.error( f"failed to retrieve item from table {table_name} by the given key {primary_key_list}" ) return + logger.debug( + f"successfully retrieved {len(primary_key_list)} items based on supplied primary key list" + ) items = response.get(table_name) self.construct_schema_from_items(items, schema) From c9309ff1579e31c79d2d8e764a89f7c5e3ff483c Mon Sep 17 00:00:00 2001 From: Shirshanka Das Date: Thu, 5 Oct 2023 09:07:12 -0700 Subject: [PATCH 095/156] feat(ci): enable ci to run on PR-s targeting all branches (#8933) --- .github/workflows/airflow-plugin.yml | 2 +- .github/workflows/build-and-test.yml | 11 +++-------- .github/workflows/check-datahub-jars.yml | 9 ++------- .github/workflows/close-stale-issues.yml | 4 +++- .github/workflows/code-checks.yml | 13 ++++--------- .github/workflows/docker-postgres-setup.yml | 3 +-- .github/workflows/docker-unified.yml | 7 +++---- .github/workflows/documentation.yml | 2 +- .github/workflows/lint-actions.yml | 4 +++- .github/workflows/metadata-ingestion.yml | 2 +- .github/workflows/metadata-io.yml | 2 +- .github/workflows/spark-smoke-test.yml | 2 +- 12 files changed, 24 insertions(+), 37 deletions(-) diff --git a/.github/workflows/airflow-plugin.yml b/.github/workflows/airflow-plugin.yml index a250bddcc16d1..54042d104d906 100644 --- a/.github/workflows/airflow-plugin.yml +++ b/.github/workflows/airflow-plugin.yml @@ -10,7 +10,7 @@ on: - "metadata-models/**" pull_request: branches: - - master + - "**" paths: - ".github/**" - "metadata-ingestion-modules/airflow-plugin/**" diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml index 96b9bb2a14933..25f3957e8f086 100644 --- a/.github/workflows/build-and-test.yml +++ b/.github/workflows/build-and-test.yml @@ -8,7 +8,7 @@ on: - "**.md" pull_request: branches: - - master + - "**" paths-ignore: - "docs/**" - "**.md" @@ -24,17 +24,12 @@ jobs: strategy: fail-fast: false matrix: - command: - [ + command: [ # metadata-ingestion and airflow-plugin each have dedicated build jobs "except_metadata_ingestion", "frontend" ] - timezone: - [ - "UTC", - "America/New_York", - ] + timezone: ["UTC", "America/New_York"] runs-on: ubuntu-latest timeout-minutes: 60 steps: diff --git a/.github/workflows/check-datahub-jars.yml b/.github/workflows/check-datahub-jars.yml index 841a9ed5f9bc7..9a17a70e7f8d4 100644 --- a/.github/workflows/check-datahub-jars.yml +++ b/.github/workflows/check-datahub-jars.yml @@ -10,7 +10,7 @@ on: - "**.md" pull_request: branches: - - master + - "**" paths-ignore: - "docker/**" - "docs/**" @@ -28,12 +28,7 @@ jobs: max-parallel: 1 fail-fast: false matrix: - command: - [ - "datahub-client", - "datahub-protobuf", - "spark-lineage" - ] + command: ["datahub-client", "datahub-protobuf", "spark-lineage"] runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 diff --git a/.github/workflows/close-stale-issues.yml b/.github/workflows/close-stale-issues.yml index a7809087702ac..98e3041f28804 100644 --- a/.github/workflows/close-stale-issues.yml +++ b/.github/workflows/close-stale-issues.yml @@ -18,7 +18,9 @@ jobs: days-before-issue-stale: 30 days-before-issue-close: 30 stale-issue-label: "stale" - stale-issue-message: "This issue is stale because it has been open for 30 days with no activity. If you believe this is still an issue on the latest DataHub release please leave a comment with the version that you tested it with. If this is a question/discussion please head to https://slack.datahubproject.io. For feature requests please use https://feature-requests.datahubproject.io" + stale-issue-message: + "This issue is stale because it has been open for 30 days with no activity. If you believe this is still an issue on the latest DataHub release please leave a comment with the version that you tested it with. If this is a question/discussion please head to https://slack.datahubproject.io.\ + \ For feature requests please use https://feature-requests.datahubproject.io" close-issue-message: "This issue was closed because it has been inactive for 30 days since being marked as stale." days-before-pr-stale: -1 days-before-pr-close: -1 diff --git a/.github/workflows/code-checks.yml b/.github/workflows/code-checks.yml index 6ce19a5b4616e..e12971b8a6208 100644 --- a/.github/workflows/code-checks.yml +++ b/.github/workflows/code-checks.yml @@ -10,7 +10,7 @@ on: - ".github/workflows/code-checks.yml" pull_request: branches: - - master + - "**" paths: - "metadata-io/**" - "datahub-web-react/**" @@ -21,17 +21,12 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true - jobs: code_check: strategy: fail-fast: false matrix: - command: - [ - "check_event_type.py", - "check_policies.py" - ] + command: ["check_event_type.py", "check_policies.py"] name: run code checks runs-on: ubuntu-latest steps: @@ -43,5 +38,5 @@ jobs: with: python-version: "3.10" - name: run check ${{ matrix.command }} - run: | - python .github/scripts/${{ matrix.command }} \ No newline at end of file + run: |- + python .github/scripts/${{ matrix.command }} diff --git a/.github/workflows/docker-postgres-setup.yml b/.github/workflows/docker-postgres-setup.yml index a5d421d4b7ff5..fda4349f90bf7 100644 --- a/.github/workflows/docker-postgres-setup.yml +++ b/.github/workflows/docker-postgres-setup.yml @@ -8,7 +8,7 @@ on: - ".github/workflows/docker-postgres-setup.yml" pull_request: branches: - - master + - "**" paths: - "docker/postgres-setup/**" - ".github/workflows/docker-postgres-setup.yml" @@ -61,4 +61,3 @@ jobs: context: . file: ./docker/postgres-setup/Dockerfile platforms: linux/amd64,linux/arm64 - diff --git a/.github/workflows/docker-unified.yml b/.github/workflows/docker-unified.yml index 2aae6bf51529d..8666a5e2e2171 100644 --- a/.github/workflows/docker-unified.yml +++ b/.github/workflows/docker-unified.yml @@ -8,7 +8,7 @@ on: - "**.md" pull_request: branches: - - master + - "**" paths-ignore: - "docs/**" - "**.md" @@ -545,7 +545,6 @@ jobs: id: tag run: echo "tag=${{ steps.filter.outputs.datahub-ingestion-base == 'true' && needs.setup.outputs.unique_full_tag || 'head' }}" >> $GITHUB_OUTPUT - datahub_ingestion_slim_build: name: Build and Push DataHub Ingestion Docker Images runs-on: ubuntu-latest @@ -809,8 +808,8 @@ jobs: DATAHUB_VERSION: ${{ needs.setup.outputs.unique_tag }} DATAHUB_ACTIONS_IMAGE: ${{ env.DATAHUB_INGESTION_IMAGE }} ACTIONS_VERSION: ${{ needs.datahub_ingestion_slim_build.outputs.tag }} - ACTIONS_EXTRA_PACKAGES: 'acryl-datahub-actions[executor]==0.0.13 acryl-datahub-actions==0.0.13 acryl-datahub==0.10.5' - ACTIONS_CONFIG: 'https://raw.githubusercontent.com/acryldata/datahub-actions/main/docker/config/executor.yaml' + ACTIONS_EXTRA_PACKAGES: "acryl-datahub-actions[executor]==0.0.13 acryl-datahub-actions==0.0.13 acryl-datahub==0.10.5" + ACTIONS_CONFIG: "https://raw.githubusercontent.com/acryldata/datahub-actions/main/docker/config/executor.yaml" run: | ./smoke-test/run-quickstart.sh - name: sleep 60s diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml index 68432a4feb13d..ebe2990f3a3cd 100644 --- a/.github/workflows/documentation.yml +++ b/.github/workflows/documentation.yml @@ -3,7 +3,7 @@ name: documentation on: pull_request: branches: - - master + - "**" push: branches: - master diff --git a/.github/workflows/lint-actions.yml b/.github/workflows/lint-actions.yml index b285e46da4857..6f34bf292bf51 100644 --- a/.github/workflows/lint-actions.yml +++ b/.github/workflows/lint-actions.yml @@ -2,8 +2,10 @@ name: Lint actions on: pull_request: paths: - - '.github/workflows/**' + - ".github/workflows/**" + branches: + - "**" jobs: actionlint: runs-on: ubuntu-latest diff --git a/.github/workflows/metadata-ingestion.yml b/.github/workflows/metadata-ingestion.yml index dea4603868f8e..699ca330ce0ac 100644 --- a/.github/workflows/metadata-ingestion.yml +++ b/.github/workflows/metadata-ingestion.yml @@ -9,7 +9,7 @@ on: - "metadata-models/**" pull_request: branches: - - master + - "**" paths: - ".github/**" - "metadata-ingestion/**" diff --git a/.github/workflows/metadata-io.yml b/.github/workflows/metadata-io.yml index e37ddd0ce4e86..48f230ce14c8d 100644 --- a/.github/workflows/metadata-io.yml +++ b/.github/workflows/metadata-io.yml @@ -10,7 +10,7 @@ on: - "metadata-io/**" pull_request: branches: - - master + - "**" paths: - "**/*.gradle" - "li-utils/**" diff --git a/.github/workflows/spark-smoke-test.yml b/.github/workflows/spark-smoke-test.yml index b2482602e7548..541b2019b93ef 100644 --- a/.github/workflows/spark-smoke-test.yml +++ b/.github/workflows/spark-smoke-test.yml @@ -12,7 +12,7 @@ on: - ".github/workflows/spark-smoke-test.yml" pull_request: branches: - - master + - "**" paths: - "metadata_models/**" - "metadata-integration/java/datahub-client/**" From 3cede10ab30e22dcad286bd42bcd154732e40942 Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Thu, 5 Oct 2023 13:29:47 -0400 Subject: [PATCH 096/156] feat(ingest/dbt): support `use_compiled_code` and `test_warnings_are_errors` (#8956) --- .../datahub/configuration/source_common.py | 2 +- ...ation.py => validate_field_deprecation.py} | 14 +++++-- .../ingestion/source/dbt/dbt_common.py | 41 ++++++++++++++----- .../src/datahub/ingestion/source/file.py | 2 +- .../ingestion/source/powerbi/config.py | 2 +- .../ingestion/source/redshift/config.py | 2 +- .../src/datahub/ingestion/source/s3/config.py | 2 +- .../ingestion/source/sql/clickhouse.py | 2 +- .../ingestion/source/sql/sql_config.py | 2 +- .../src/datahub/ingestion/source/tableau.py | 2 +- .../tests/unit/test_pydantic_validators.py | 2 +- 11 files changed, 51 insertions(+), 22 deletions(-) rename metadata-ingestion/src/datahub/configuration/{pydantic_field_deprecation.py => validate_field_deprecation.py} (74%) diff --git a/metadata-ingestion/src/datahub/configuration/source_common.py b/metadata-ingestion/src/datahub/configuration/source_common.py index 37b93f3e598e1..a9f891ddb7b1e 100644 --- a/metadata-ingestion/src/datahub/configuration/source_common.py +++ b/metadata-ingestion/src/datahub/configuration/source_common.py @@ -4,7 +4,7 @@ from pydantic.fields import Field from datahub.configuration.common import ConfigModel, ConfigurationError -from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated +from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.metadata.schema_classes import FabricTypeClass DEFAULT_ENV = FabricTypeClass.PROD diff --git a/metadata-ingestion/src/datahub/configuration/pydantic_field_deprecation.py b/metadata-ingestion/src/datahub/configuration/validate_field_deprecation.py similarity index 74% rename from metadata-ingestion/src/datahub/configuration/pydantic_field_deprecation.py rename to metadata-ingestion/src/datahub/configuration/validate_field_deprecation.py index ed82acb594ed7..6134c4dab4817 100644 --- a/metadata-ingestion/src/datahub/configuration/pydantic_field_deprecation.py +++ b/metadata-ingestion/src/datahub/configuration/validate_field_deprecation.py @@ -1,20 +1,28 @@ import warnings -from typing import Optional, Type +from typing import Any, Optional, Type import pydantic from datahub.configuration.common import ConfigurationWarning from datahub.utilities.global_warning_util import add_global_warning +_unset = object() -def pydantic_field_deprecated(field: str, message: Optional[str] = None) -> classmethod: + +def pydantic_field_deprecated( + field: str, + warn_if_value_is_not: Any = _unset, + message: Optional[str] = None, +) -> classmethod: if message: output = message else: output = f"{field} is deprecated and will be removed in a future release. Please remove it from your config." def _validate_deprecated(cls: Type, values: dict) -> dict: - if field in values: + if field in values and ( + warn_if_value_is_not is _unset or values[field] != warn_if_value_is_not + ): add_global_warning(output) warnings.warn(output, ConfigurationWarning, stacklevel=2) return values diff --git a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py index f9b71892975b4..0f5c08eb6ac54 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py +++ b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py @@ -18,8 +18,8 @@ ConfigurationError, LineageConfig, ) -from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated from datahub.configuration.source_common import DatasetSourceConfigMixin +from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.emitter import mce_builder from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.common import PipelineContext @@ -214,7 +214,9 @@ class DBTCommonConfig( default=False, description="Use model identifier instead of model name if defined (if not, default to model name).", ) - _deprecate_use_identifiers = pydantic_field_deprecated("use_identifiers") + _deprecate_use_identifiers = pydantic_field_deprecated( + "use_identifiers", warn_if_value_is_not=False + ) entities_enabled: DBTEntitiesEnabled = Field( DBTEntitiesEnabled(), @@ -278,6 +280,14 @@ class DBTCommonConfig( description="When enabled, converts column URNs to lowercase to ensure cross-platform compatibility. " "If `target_platform` is Snowflake, the default is True.", ) + use_compiled_code: bool = Field( + default=False, + description="When enabled, uses the compiled dbt code instead of the raw dbt node definition.", + ) + test_warnings_are_errors: bool = Field( + default=False, + description="When enabled, dbt test warnings will be treated as failures.", + ) @validator("target_platform") def validate_target_platform_value(cls, target_platform: str) -> str: @@ -811,7 +821,7 @@ def _make_assertion_from_test( mce_builder.make_schema_field_urn(upstream_urn, column_name) ], nativeType=node.name, - logic=node.compiled_code if node.compiled_code else node.raw_code, + logic=node.compiled_code or node.raw_code, aggregation=AssertionStdAggregationClass._NATIVE_, nativeParameters=string_map(kw_args), ), @@ -825,7 +835,7 @@ def _make_assertion_from_test( dataset=upstream_urn, scope=DatasetAssertionScopeClass.DATASET_ROWS, operator=AssertionStdOperatorClass._NATIVE_, - logic=node.compiled_code if node.compiled_code else node.raw_code, + logic=node.compiled_code or node.raw_code, nativeType=node.name, aggregation=AssertionStdAggregationClass._NATIVE_, nativeParameters=string_map(kw_args), @@ -856,6 +866,10 @@ def _make_assertion_result_from_test( result=AssertionResultClass( type=AssertionResultTypeClass.SUCCESS if test_result.status == "pass" + or ( + not self.config.test_warnings_are_errors + and test_result.status == "warn" + ) else AssertionResultTypeClass.FAILURE, nativeResults=test_result.native_results, ), @@ -1007,8 +1021,8 @@ def create_platform_mces( aspects.append(upstream_lineage_class) # add view properties aspect - if node.raw_code and node.language == "sql": - view_prop_aspect = self._create_view_properties_aspect(node) + view_prop_aspect = self._create_view_properties_aspect(node) + if view_prop_aspect: aspects.append(view_prop_aspect) # emit subtype mcp @@ -1133,14 +1147,21 @@ def _create_dataset_properties_aspect( def get_external_url(self, node: DBTNode) -> Optional[str]: pass - def _create_view_properties_aspect(self, node: DBTNode) -> ViewPropertiesClass: + def _create_view_properties_aspect( + self, node: DBTNode + ) -> Optional[ViewPropertiesClass]: + view_logic = ( + node.compiled_code if self.config.use_compiled_code else node.raw_code + ) + + if node.language != "sql" or not view_logic: + return None + materialized = node.materialization in {"table", "incremental", "snapshot"} - # this function is only called when raw sql is present. assert is added to satisfy lint checks - assert node.raw_code is not None view_properties = ViewPropertiesClass( materialized=materialized, viewLanguage="SQL", - viewLogic=node.raw_code, + viewLogic=view_logic, ) return view_properties diff --git a/metadata-ingestion/src/datahub/ingestion/source/file.py b/metadata-ingestion/src/datahub/ingestion/source/file.py index de61fa8481c58..590aa59f7b5b6 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/file.py +++ b/metadata-ingestion/src/datahub/ingestion/source/file.py @@ -16,7 +16,7 @@ from pydantic.fields import Field from datahub.configuration.common import ConfigEnum, ConfigModel, ConfigurationError -from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated +from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.configuration.validate_field_rename import pydantic_renamed_field from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.common import PipelineContext diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py index a8c7e48f3785c..96729f4c60c6c 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py @@ -9,8 +9,8 @@ import datahub.emitter.mce_builder as builder from datahub.configuration.common import AllowDenyPattern, ConfigModel -from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated from datahub.configuration.source_common import DEFAULT_ENV, DatasetSourceConfigMixin +from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.ingestion.source.common.subtypes import BIAssetSubTypes from datahub.ingestion.source.state.stale_entity_removal_handler import ( StaleEntityRemovalSourceReport, diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py index 93850607e551e..804a14b0fe1cf 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py @@ -7,8 +7,8 @@ from datahub.configuration import ConfigModel from datahub.configuration.common import AllowDenyPattern -from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated from datahub.configuration.source_common import DatasetLineageProviderConfigBase +from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.ingestion.source.data_lake_common.path_spec import PathSpec from datahub.ingestion.source.sql.postgres import BasePostgresConfig from datahub.ingestion.source.state.stateful_ingestion_base import ( diff --git a/metadata-ingestion/src/datahub/ingestion/source/s3/config.py b/metadata-ingestion/src/datahub/ingestion/source/s3/config.py index f1dd622efb746..9b5296f0b9dd5 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/s3/config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/s3/config.py @@ -5,8 +5,8 @@ from pydantic.fields import Field from datahub.configuration.common import AllowDenyPattern -from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated from datahub.configuration.source_common import DatasetSourceConfigMixin +from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.configuration.validate_field_rename import pydantic_renamed_field from datahub.ingestion.source.aws.aws_common import AwsConnectionConfig from datahub.ingestion.source.data_lake_common.config import PathSpecsConfigMixin diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py b/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py index 1626f86b92545..8873038079bad 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py @@ -19,9 +19,9 @@ from sqlalchemy.types import BOOLEAN, DATE, DATETIME, INTEGER import datahub.emitter.mce_builder as builder -from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated from datahub.configuration.source_common import DatasetLineageProviderConfigBase from datahub.configuration.time_window_config import BaseTimeWindowConfig +from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.emitter import mce_builder from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.decorators import ( diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_config.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_config.py index 8f1e04b915f3b..677d32c8bac08 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_config.py @@ -7,8 +7,8 @@ from pydantic import Field from datahub.configuration.common import AllowDenyPattern, ConfigModel -from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated from datahub.configuration.source_common import DatasetSourceConfigMixin +from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig from datahub.ingestion.source.state.stale_entity_removal_handler import ( StatefulStaleMetadataRemovalConfig, diff --git a/metadata-ingestion/src/datahub/ingestion/source/tableau.py b/metadata-ingestion/src/datahub/ingestion/source/tableau.py index 6214cba342622..e347cd26d245a 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/tableau.py +++ b/metadata-ingestion/src/datahub/ingestion/source/tableau.py @@ -37,11 +37,11 @@ ConfigModel, ConfigurationError, ) -from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated from datahub.configuration.source_common import ( DatasetLineageProviderConfigBase, DatasetSourceConfigMixin, ) +from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.mcp_builder import ( ContainerKey, diff --git a/metadata-ingestion/tests/unit/test_pydantic_validators.py b/metadata-ingestion/tests/unit/test_pydantic_validators.py index 07d86043a35bf..3e9ec6cbaf357 100644 --- a/metadata-ingestion/tests/unit/test_pydantic_validators.py +++ b/metadata-ingestion/tests/unit/test_pydantic_validators.py @@ -4,7 +4,7 @@ from pydantic import ValidationError from datahub.configuration.common import ConfigModel -from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated +from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.configuration.validate_field_removal import pydantic_removed_field from datahub.configuration.validate_field_rename import pydantic_renamed_field from datahub.utilities.global_warning_util import get_global_warnings From debac3cf5c31b471a5a82da8d18fb8303cc8b9d0 Mon Sep 17 00:00:00 2001 From: Patrick Franco Braz Date: Thu, 5 Oct 2023 17:47:10 -0300 Subject: [PATCH 097/156] refactor(boot): increases wait timeout for servlets initialization (#8947) Co-authored-by: RyanHolstien --- .../configuration/src/main/resources/application.yml | 3 +++ .../metadata/boot/OnBootApplicationListener.java | 9 +++++++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/metadata-service/configuration/src/main/resources/application.yml b/metadata-service/configuration/src/main/resources/application.yml index 4be31b2b6bb15..4dfd96ac75c6c 100644 --- a/metadata-service/configuration/src/main/resources/application.yml +++ b/metadata-service/configuration/src/main/resources/application.yml @@ -276,6 +276,9 @@ bootstrap: enabled: ${UPGRADE_DEFAULT_BROWSE_PATHS_ENABLED:false} # enable to run the upgrade to migrate legacy default browse paths to new ones backfillBrowsePathsV2: enabled: ${BACKFILL_BROWSE_PATHS_V2:false} # Enables running the backfill of browsePathsV2 upgrade step. There are concerns about the load of this step so hiding it behind a flag. Deprecating in favor of running through SystemUpdate + servlets: + waitTimeout: ${BOOTSTRAP_SERVLETS_WAITTIMEOUT:60} # Total waiting time in seconds for servlets to initialize + systemUpdate: initialBackOffMs: ${BOOTSTRAP_SYSTEM_UPDATE_INITIAL_BACK_OFF_MILLIS:5000} diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/OnBootApplicationListener.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/OnBootApplicationListener.java index 980cafaceae27..032b934a7ba87 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/OnBootApplicationListener.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/OnBootApplicationListener.java @@ -15,15 +15,18 @@ import org.apache.http.impl.client.HttpClients; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; +import org.springframework.beans.factory.annotation.Value; import org.springframework.context.event.ContextRefreshedEvent; import org.springframework.context.event.EventListener; import org.springframework.stereotype.Component; import org.springframework.web.context.WebApplicationContext; +import org.springframework.context.annotation.Configuration; /** * Responsible for coordinating starting steps that happen before the application starts up. */ +@Configuration @Slf4j @Component public class OnBootApplicationListener { @@ -44,6 +47,8 @@ public class OnBootApplicationListener { @Qualifier("configurationProvider") private ConfigurationProvider provider; + @Value("${bootstrap.servlets.waitTimeout}") + private int _servletsWaitTimeout; @EventListener(ContextRefreshedEvent.class) public void onApplicationEvent(@Nonnull ContextRefreshedEvent event) { @@ -62,7 +67,7 @@ public void onApplicationEvent(@Nonnull ContextRefreshedEvent event) { public Runnable isSchemaRegistryAPIServletReady() { return () -> { final HttpGet request = new HttpGet(provider.getKafka().getSchemaRegistry().getUrl()); - int timeouts = 30; + int timeouts = _servletsWaitTimeout; boolean openAPIServeletReady = false; while (!openAPIServeletReady && timeouts > 0) { try { @@ -79,7 +84,7 @@ public Runnable isSchemaRegistryAPIServletReady() { timeouts--; } if (!openAPIServeletReady) { - log.error("Failed to bootstrap DataHub, OpenAPI servlet was not ready after 30 seconds"); + log.error("Failed to bootstrap DataHub, OpenAPI servlet was not ready after {} seconds", timeouts); System.exit(1); } else { _bootstrapManager.start(); From 26bc039b967d3a62a7079522b702e97ed8ad8d27 Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Thu, 5 Oct 2023 23:23:15 -0400 Subject: [PATCH 098/156] fix(ingest/unity): Remove metastore from ingestion and urns; standardize platform instance; add notebook filter (#8943) --- docs/how/updating-datahub.md | 5 + .../src/datahub/emitter/mcp_builder.py | 10 +- .../datahub/ingestion/source/unity/config.py | 45 ++++++++- .../datahub/ingestion/source/unity/proxy.py | 16 +-- .../ingestion/source/unity/proxy_types.py | 19 ++-- .../datahub/ingestion/source/unity/source.py | 99 ++++++++++++------- 6 files changed, 145 insertions(+), 49 deletions(-) diff --git a/docs/how/updating-datahub.md b/docs/how/updating-datahub.md index 4df8d435cf1c4..5d0ad5eaf8f7e 100644 --- a/docs/how/updating-datahub.md +++ b/docs/how/updating-datahub.md @@ -9,6 +9,11 @@ This file documents any backwards-incompatible changes in DataHub and assists pe - #8810 - Removed support for SQLAlchemy 1.3.x. Only SQLAlchemy 1.4.x is supported now. - #8853 - The Airflow plugin no longer supports Airflow 2.0.x or Python 3.7. See the docs for more details. - #8853 - Introduced the Airflow plugin v2. If you're using Airflow 2.3+, the v2 plugin will be enabled by default, and so you'll need to switch your requirements to include `pip install 'acryl-datahub-airflow-plugin[plugin-v2]'`. To continue using the v1 plugin, set the `DATAHUB_AIRFLOW_PLUGIN_USE_V1_PLUGIN` environment variable to `true`. +- #8943 The Unity Catalog ingestion source has a new option `include_metastore`, which will cause all urns to be changed when disabled. +This is currently enabled by default to preserve compatibility, but will be disabled by default and then removed in the future. +If stateful ingestion is enabled, simply setting `include_metastore: false` will perform all required cleanup. +Otherwise, we recommend soft deleting all databricks data via the DataHub CLI: +`datahub delete --platform databricks --soft` and then reingesting with `include_metastore: false`. ### Potential Downtime diff --git a/metadata-ingestion/src/datahub/emitter/mcp_builder.py b/metadata-ingestion/src/datahub/emitter/mcp_builder.py index 06f689dfd317b..65e0c0d6ba60d 100644 --- a/metadata-ingestion/src/datahub/emitter/mcp_builder.py +++ b/metadata-ingestion/src/datahub/emitter/mcp_builder.py @@ -94,7 +94,15 @@ class MetastoreKey(ContainerKey): metastore: str -class CatalogKey(MetastoreKey): +class CatalogKeyWithMetastore(MetastoreKey): + catalog: str + + +class UnitySchemaKeyWithMetastore(CatalogKeyWithMetastore): + unity_schema: str + + +class CatalogKey(ContainerKey): catalog: str diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py index a49c789a82f27..f259fa260f653 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py @@ -1,3 +1,4 @@ +import logging import os from datetime import datetime, timedelta, timezone from typing import Any, Dict, Optional @@ -21,6 +22,9 @@ OperationConfig, is_profiling_enabled, ) +from datahub.utilities.global_warning_util import add_global_warning + +logger = logging.getLogger(__name__) class UnityCatalogProfilerConfig(ConfigModel): @@ -97,9 +101,25 @@ class UnityCatalogSourceConfig( description="Name of the workspace. Default to deployment name present in workspace_url", ) + include_metastore: bool = pydantic.Field( + default=True, + description=( + "Whether to ingest the workspace's metastore as a container and include it in all urns." + " Changing this will affect the urns of all entities in the workspace." + " This will be disabled by default in the future," + " so it is recommended to set this to `False` for new ingestions." + " If you have an existing unity catalog ingestion, you'll want to avoid duplicates by soft deleting existing data." + " If stateful ingestion is enabled, running with `include_metastore: false` should be sufficient." + " Otherwise, we recommend deleting via the cli: `datahub delete --platform databricks` and re-ingesting with `include_metastore: false`." + ), + ) + ingest_data_platform_instance_aspect: Optional[bool] = pydantic.Field( default=False, - description="Option to enable/disable ingestion of the data platform instance aspect. The default data platform instance id for a dataset is workspace_name", + description=( + "Option to enable/disable ingestion of the data platform instance aspect." + " The default data platform instance id for a dataset is workspace_name" + ), ) _only_ingest_assigned_metastore_removed = pydantic_removed_field( @@ -122,6 +142,16 @@ class UnityCatalogSourceConfig( default=AllowDenyPattern.allow_all(), description="Regex patterns for tables to filter in ingestion. Specify regex to match the entire table name in `catalog.schema.table` format. e.g. to match all tables starting with customer in Customer catalog and public schema, use the regex `Customer\\.public\\.customer.*`.", ) + + notebook_pattern: AllowDenyPattern = Field( + default=AllowDenyPattern.allow_all(), + description=( + "Regex patterns for notebooks to filter in ingestion, based on notebook *path*." + " Specify regex to match the entire notebook path in `//.../` format." + " e.g. to match all notebooks in the root Shared directory, use the regex `/Shared/.*`." + ), + ) + domain: Dict[str, AllowDenyPattern] = Field( default=dict(), description='Attach domains to catalogs, schemas or tables during ingestion using regex patterns. Domain key can be a guid like *urn:li:domain:ec428203-ce86-4db3-985d-5a8ee6df32ba* or a string like "Marketing".) If you provide strings, then datahub will attempt to resolve this name to a guid, and will error out if this fails. There can be multiple domain keys specified.', @@ -182,3 +212,16 @@ def workspace_url_should_start_with_http_scheme(cls, workspace_url: str) -> str: "Workspace URL must start with http scheme. e.g. https://my-workspace.cloud.databricks.com" ) return workspace_url + + @pydantic.validator("include_metastore") + def include_metastore_warning(cls, v: bool) -> bool: + if v: + msg = ( + "`include_metastore` is enabled." + " This is not recommended and will be disabled by default in the future, which is a breaking change." + " All databricks urns will change if you re-ingest with this disabled." + " We recommend soft deleting all databricks data and re-ingesting with `include_metastore` set to `False`." + ) + logger.warning(msg) + add_global_warning(msg) + return v diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py index 2401f1c3d163c..529d9e7b563a5 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py @@ -97,14 +97,13 @@ def __init__( self.report = report def check_basic_connectivity(self) -> bool: - self._workspace_client.metastores.summary() - return True + return bool(self._workspace_client.catalogs.list()) def assigned_metastore(self) -> Metastore: response = self._workspace_client.metastores.summary() return self._create_metastore(response) - def catalogs(self, metastore: Metastore) -> Iterable[Catalog]: + def catalogs(self, metastore: Optional[Metastore]) -> Iterable[Catalog]: response = self._workspace_client.catalogs.list() if not response: logger.info("Catalogs not found") @@ -247,7 +246,7 @@ def table_lineage( for item in response.get("upstreams") or []: if "tableInfo" in item: table_ref = TableReference.create_from_lineage( - item["tableInfo"], table.schema.catalog.metastore.id + item["tableInfo"], table.schema.catalog.metastore ) if table_ref: table.upstreams[table_ref] = {} @@ -276,7 +275,7 @@ def get_column_lineage(self, table: Table, include_entity_lineage: bool) -> None ) for item in response.get("upstream_cols", []): table_ref = TableReference.create_from_lineage( - item, table.schema.catalog.metastore.id + item, table.schema.catalog.metastore ) if table_ref: table.upstreams.setdefault(table_ref, {}).setdefault( @@ -305,10 +304,13 @@ def _create_metastore( comment=None, ) - def _create_catalog(self, metastore: Metastore, obj: CatalogInfo) -> Catalog: + def _create_catalog( + self, metastore: Optional[Metastore], obj: CatalogInfo + ) -> Catalog: + catalog_name = self._escape_sequence(obj.name) return Catalog( name=obj.name, - id=f"{metastore.id}.{self._escape_sequence(obj.name)}", + id=f"{metastore.id}.{catalog_name}" if metastore else catalog_name, metastore=metastore, comment=obj.comment, owner=obj.owner, diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py index 54ac2e90d7c7e..18ac2475b51e0 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py @@ -92,7 +92,7 @@ class Metastore(CommonProperty): @dataclass class Catalog(CommonProperty): - metastore: Metastore + metastore: Optional[Metastore] owner: Optional[str] type: CatalogType @@ -130,7 +130,7 @@ class ServicePrincipal: @dataclass(frozen=True, order=True) class TableReference: - metastore: str + metastore: Optional[str] catalog: str schema: str table: str @@ -138,17 +138,21 @@ class TableReference: @classmethod def create(cls, table: "Table") -> "TableReference": return cls( - table.schema.catalog.metastore.id, + table.schema.catalog.metastore.id + if table.schema.catalog.metastore + else None, table.schema.catalog.name, table.schema.name, table.name, ) @classmethod - def create_from_lineage(cls, d: dict, metastore: str) -> Optional["TableReference"]: + def create_from_lineage( + cls, d: dict, metastore: Optional[Metastore] + ) -> Optional["TableReference"]: try: return cls( - metastore, + metastore.id if metastore else None, d["catalog_name"], d["schema_name"], d.get("table_name", d["name"]), # column vs table query output @@ -158,7 +162,10 @@ def create_from_lineage(cls, d: dict, metastore: str) -> Optional["TableReferenc return None def __str__(self) -> str: - return f"{self.metastore}.{self.catalog}.{self.schema}.{self.table}" + if self.metastore: + return f"{self.metastore}.{self.catalog}.{self.schema}.{self.table}" + else: + return self.qualified_table_name @property def qualified_table_name(self) -> str: diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py index f2da1aece9fd4..4f7866aee7681 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py @@ -16,10 +16,12 @@ from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.mcp_builder import ( CatalogKey, + CatalogKeyWithMetastore, ContainerKey, MetastoreKey, NotebookKey, UnitySchemaKey, + UnitySchemaKeyWithMetastore, add_dataset_to_container, gen_containers, ) @@ -127,7 +129,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource): config: UnityCatalogSourceConfig unity_catalog_api_proxy: UnityCatalogApiProxy platform: str = "databricks" - platform_instance_name: str + platform_instance_name: Optional[str] def get_report(self) -> UnityCatalogReport: return self.report @@ -146,11 +148,13 @@ def __init__(self, ctx: PipelineContext, config: UnityCatalogSourceConfig): self.external_url_base = urljoin(self.config.workspace_url, "/explore/data") # Determine the platform_instance_name - self.platform_instance_name = ( - config.workspace_name - if config.workspace_name is not None - else config.workspace_url.split("//")[1].split(".")[0] - ) + self.platform_instance_name = self.config.platform_instance + if self.config.include_metastore: + self.platform_instance_name = ( + config.workspace_name + if config.workspace_name is not None + else config.workspace_url.split("//")[1].split(".")[0] + ) if self.config.domain: self.domain_registry = DomainRegistry( @@ -247,10 +251,14 @@ def build_service_principal_map(self) -> None: def process_notebooks(self) -> Iterable[MetadataWorkUnit]: for notebook in self.unity_catalog_api_proxy.workspace_notebooks(): + if not self.config.notebook_pattern.allowed(notebook.path): + self.report.notebooks.dropped(notebook.path) + continue + self.notebooks[str(notebook.id)] = notebook - yield from self._gen_notebook_aspects(notebook) + yield from self._gen_notebook_workunits(notebook) - def _gen_notebook_aspects(self, notebook: Notebook) -> Iterable[MetadataWorkUnit]: + def _gen_notebook_workunits(self, notebook: Notebook) -> Iterable[MetadataWorkUnit]: mcps = MetadataChangeProposalWrapper.construct_many( entityUrn=self.gen_notebook_urn(notebook), aspects=[ @@ -270,7 +278,7 @@ def _gen_notebook_aspects(self, notebook: Notebook) -> Iterable[MetadataWorkUnit ), SubTypesClass(typeNames=[DatasetSubTypes.NOTEBOOK]), BrowsePathsClass(paths=notebook.path.split("/")), - # TODO: Add DPI aspect + self._create_data_platform_instance_aspect(), ], ) for mcp in mcps: @@ -296,13 +304,17 @@ def _gen_notebook_lineage(self, notebook: Notebook) -> Optional[MetadataWorkUnit ).as_workunit() def process_metastores(self) -> Iterable[MetadataWorkUnit]: - metastore = self.unity_catalog_api_proxy.assigned_metastore() - yield from self.gen_metastore_containers(metastore) + metastore: Optional[Metastore] = None + if self.config.include_metastore: + metastore = self.unity_catalog_api_proxy.assigned_metastore() + yield from self.gen_metastore_containers(metastore) yield from self.process_catalogs(metastore) + if metastore and self.config.include_metastore: + self.report.metastores.processed(metastore.id) - self.report.metastores.processed(metastore.id) - - def process_catalogs(self, metastore: Metastore) -> Iterable[MetadataWorkUnit]: + def process_catalogs( + self, metastore: Optional[Metastore] + ) -> Iterable[MetadataWorkUnit]: for catalog in self.unity_catalog_api_proxy.catalogs(metastore=metastore): if not self.config.catalog_pattern.allowed(catalog.id): self.report.catalogs.dropped(catalog.id) @@ -353,7 +365,7 @@ def process_table(self, table: Table, schema: Schema) -> Iterable[MetadataWorkUn operation = self._create_table_operation_aspect(table) domain = self._get_domain_aspect(dataset_name=table.ref.qualified_table_name) ownership = self._create_table_ownership_aspect(table) - data_platform_instance = self._create_data_platform_instance_aspect(table) + data_platform_instance = self._create_data_platform_instance_aspect() if self.config.include_column_lineage: self.unity_catalog_api_proxy.get_column_lineage( @@ -503,27 +515,37 @@ def gen_metastore_containers( def gen_catalog_containers(self, catalog: Catalog) -> Iterable[MetadataWorkUnit]: domain_urn = self._gen_domain_urn(catalog.name) - metastore_container_key = self.gen_metastore_key(catalog.metastore) catalog_container_key = self.gen_catalog_key(catalog) yield from gen_containers( container_key=catalog_container_key, name=catalog.name, sub_types=[DatasetContainerSubTypes.CATALOG], domain_urn=domain_urn, - parent_container_key=metastore_container_key, + parent_container_key=self.gen_metastore_key(catalog.metastore) + if self.config.include_metastore and catalog.metastore + else None, description=catalog.comment, owner_urn=self.get_owner_urn(catalog.owner), external_url=f"{self.external_url_base}/{catalog.name}", ) def gen_schema_key(self, schema: Schema) -> ContainerKey: - return UnitySchemaKey( - unity_schema=schema.name, - platform=self.platform, - instance=self.config.platform_instance, - catalog=schema.catalog.name, - metastore=schema.catalog.metastore.name, - ) + if self.config.include_metastore: + assert schema.catalog.metastore + return UnitySchemaKeyWithMetastore( + unity_schema=schema.name, + platform=self.platform, + instance=self.config.platform_instance, + catalog=schema.catalog.name, + metastore=schema.catalog.metastore.name, + ) + else: + return UnitySchemaKey( + unity_schema=schema.name, + platform=self.platform, + instance=self.config.platform_instance, + catalog=schema.catalog.name, + ) def gen_metastore_key(self, metastore: Metastore) -> MetastoreKey: return MetastoreKey( @@ -532,13 +554,21 @@ def gen_metastore_key(self, metastore: Metastore) -> MetastoreKey: instance=self.config.platform_instance, ) - def gen_catalog_key(self, catalog: Catalog) -> CatalogKey: - return CatalogKey( - catalog=catalog.name, - metastore=catalog.metastore.name, - platform=self.platform, - instance=self.config.platform_instance, - ) + def gen_catalog_key(self, catalog: Catalog) -> ContainerKey: + if self.config.include_metastore: + assert catalog.metastore + return CatalogKeyWithMetastore( + catalog=catalog.name, + metastore=catalog.metastore.name, + platform=self.platform, + instance=self.config.platform_instance, + ) + else: + return CatalogKey( + catalog=catalog.name, + platform=self.platform, + instance=self.config.platform_instance, + ) def _gen_domain_urn(self, dataset_name: str) -> Optional[str]: domain_urn: Optional[str] = None @@ -643,15 +673,16 @@ def _create_table_ownership_aspect(self, table: Table) -> Optional[OwnershipClas return None def _create_data_platform_instance_aspect( - self, table: Table + self, ) -> Optional[DataPlatformInstanceClass]: - # Only ingest the DPI aspect if the flag is true if self.config.ingest_data_platform_instance_aspect: return DataPlatformInstanceClass( platform=make_data_platform_urn(self.platform), instance=make_dataplatform_instance_urn( self.platform, self.platform_instance_name - ), + ) + if self.platform_instance_name + else None, ) return None From ea87febd2bdf0aebf603532be9448e6435f1fea9 Mon Sep 17 00:00:00 2001 From: Hyejin Yoon <0327jane@gmail.com> Date: Fri, 6 Oct 2023 14:36:32 +0900 Subject: [PATCH 099/156] fix: add retry for fetch_url (#8958) --- docs-website/download_historical_versions.py | 34 ++++++++++++++------ 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/docs-website/download_historical_versions.py b/docs-website/download_historical_versions.py index 83157edc1972c..53ee9cf1e63ef 100644 --- a/docs-website/download_historical_versions.py +++ b/docs-website/download_historical_versions.py @@ -1,6 +1,7 @@ import json import os import tarfile +import time import urllib.request repo_url = "https://api.github.com/repos/datahub-project/static-assets" @@ -16,17 +17,30 @@ def download_file(url, destination): f.write(chunk) -def fetch_urls(repo_url: str, folder_path: str, file_format: str): +def fetch_urls( + repo_url: str, folder_path: str, file_format: str, max_retries=3, retry_delay=5 +): api_url = f"{repo_url}/contents/{folder_path}" - response = urllib.request.urlopen(api_url) - data = response.read().decode("utf-8") - urls = [ - file["download_url"] - for file in json.loads(data) - if file["name"].endswith(file_format) - ] - print(urls) - return urls + for attempt in range(max_retries + 1): + try: + response = urllib.request.urlopen(api_url) + if response.status == 403 or (500 <= response.status < 600): + raise Exception(f"HTTP Error {response.status}: {response.reason}") + data = response.read().decode("utf-8") + urls = [ + file["download_url"] + for file in json.loads(data) + if file["name"].endswith(file_format) + ] + print(urls) + return urls + except Exception as e: + if attempt < max_retries: + print(f"Attempt {attempt + 1}/{max_retries}: {e}") + time.sleep(retry_delay) + else: + print(f"Max retries reached. Unable to fetch data.") + raise def extract_tar_file(destination_path): From c80da8f949aea340af73c992ff6d2bd129eb55fe Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Fri, 6 Oct 2023 10:06:36 -0400 Subject: [PATCH 100/156] feat(ingest/unity): Use ThreadPoolExecutor for CLL (#8952) --- .../datahub/ingestion/source/unity/config.py | 11 +++++ .../datahub/ingestion/source/unity/proxy.py | 46 ++++++++----------- .../datahub/ingestion/source/unity/report.py | 2 + .../datahub/ingestion/source/unity/source.py | 33 +++++++++---- 4 files changed, 57 insertions(+), 35 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py index f259fa260f653..51390873712d3 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py @@ -181,6 +181,17 @@ class UnityCatalogSourceConfig( description="Option to enable/disable lineage generation. Currently we have to call a rest call per column to get column level lineage due to the Databrick api which can slow down ingestion. ", ) + column_lineage_column_limit: int = pydantic.Field( + default=300, + description="Limit the number of columns to get column level lineage. ", + ) + + lineage_max_workers: int = pydantic.Field( + default=5 * (os.cpu_count() or 4), + description="Number of worker threads to use for column lineage thread pool executor. Set to 1 to disable.", + hidden_from_docs=True, + ) + include_usage_statistics: bool = Field( default=True, description="Generate usage statistics.", diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py index 529d9e7b563a5..9bcdb200f180e 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py @@ -233,9 +233,7 @@ def list_lineages_by_column(self, table_name: str, column_name: str) -> dict: body={"table_name": table_name, "column_name": column_name}, ) - def table_lineage( - self, table: Table, include_entity_lineage: bool - ) -> Optional[dict]: + def table_lineage(self, table: Table, include_entity_lineage: bool) -> None: # Lineage endpoint doesn't exists on 2.1 version try: response: dict = self.list_lineages_by_table( @@ -256,34 +254,30 @@ def table_lineage( for item in response.get("downstreams") or []: for notebook in item.get("notebookInfos") or []: table.downstream_notebooks.add(notebook["notebook_id"]) - - return response except Exception as e: - logger.error(f"Error getting lineage: {e}") - return None + logger.warning( + f"Error getting lineage on table {table.ref}: {e}", exc_info=True + ) - def get_column_lineage(self, table: Table, include_entity_lineage: bool) -> None: + def get_column_lineage(self, table: Table, column_name: str) -> None: try: - table_lineage = self.table_lineage( - table, include_entity_lineage=include_entity_lineage + response: dict = self.list_lineages_by_column( + table_name=table.ref.qualified_table_name, + column_name=column_name, ) - if table_lineage: - for column in table.columns: - response: dict = self.list_lineages_by_column( - table_name=table.ref.qualified_table_name, - column_name=column.name, - ) - for item in response.get("upstream_cols", []): - table_ref = TableReference.create_from_lineage( - item, table.schema.catalog.metastore - ) - if table_ref: - table.upstreams.setdefault(table_ref, {}).setdefault( - column.name, [] - ).append(item["name"]) - + for item in response.get("upstream_cols") or []: + table_ref = TableReference.create_from_lineage( + item, table.schema.catalog.metastore + ) + if table_ref: + table.upstreams.setdefault(table_ref, {}).setdefault( + column_name, [] + ).append(item["name"]) except Exception as e: - logger.error(f"Error getting lineage: {e}") + logger.warning( + f"Error getting column lineage on table {table.ref}, column {column_name}: {e}", + exc_info=True, + ) @staticmethod def _escape_sequence(value: str) -> str: diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/report.py b/metadata-ingestion/src/datahub/ingestion/source/unity/report.py index 808172a136bb3..fa61571fa92cb 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/report.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/report.py @@ -18,6 +18,8 @@ class UnityCatalogReport(IngestionStageReport, StaleEntityRemovalSourceReport): table_profiles: EntityFilterReport = EntityFilterReport.field(type="table profile") notebooks: EntityFilterReport = EntityFilterReport.field(type="notebook") + num_column_lineage_skipped_column_count: int = 0 + num_queries: int = 0 num_queries_dropped_parse_failure: int = 0 num_queries_missing_table: int = 0 # Can be due to pattern filter diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py index 4f7866aee7681..27c1f341aa84d 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py @@ -1,6 +1,7 @@ import logging import re import time +from concurrent.futures import ThreadPoolExecutor from datetime import timedelta from typing import Dict, Iterable, List, Optional, Set, Union from urllib.parse import urljoin @@ -367,15 +368,7 @@ def process_table(self, table: Table, schema: Schema) -> Iterable[MetadataWorkUn ownership = self._create_table_ownership_aspect(table) data_platform_instance = self._create_data_platform_instance_aspect() - if self.config.include_column_lineage: - self.unity_catalog_api_proxy.get_column_lineage( - table, include_entity_lineage=self.config.include_notebooks - ) - elif self.config.include_table_lineage: - self.unity_catalog_api_proxy.table_lineage( - table, include_entity_lineage=self.config.include_notebooks - ) - lineage = self._generate_lineage_aspect(dataset_urn, table) + lineage = self.ingest_lineage(table) if self.config.include_notebooks: for notebook_id in table.downstream_notebooks: @@ -401,6 +394,28 @@ def process_table(self, table: Table, schema: Schema) -> Iterable[MetadataWorkUn ) ] + def ingest_lineage(self, table: Table) -> Optional[UpstreamLineageClass]: + if self.config.include_table_lineage: + self.unity_catalog_api_proxy.table_lineage( + table, include_entity_lineage=self.config.include_notebooks + ) + + if self.config.include_column_lineage and table.upstreams: + if len(table.columns) > self.config.column_lineage_column_limit: + self.report.num_column_lineage_skipped_column_count += 1 + + with ThreadPoolExecutor( + max_workers=self.config.lineage_max_workers + ) as executor: + for column in table.columns[: self.config.column_lineage_column_limit]: + executor.submit( + self.unity_catalog_api_proxy.get_column_lineage, + table, + column.name, + ) + + return self._generate_lineage_aspect(self.gen_dataset_urn(table.ref), table) + def _generate_lineage_aspect( self, dataset_urn: str, table: Table ) -> Optional[UpstreamLineageClass]: From 8e7f286e71b36a07b4fedc0de1807354064a4fa5 Mon Sep 17 00:00:00 2001 From: Mayuri Nehate <33225191+mayurinehate@users.noreply.github.com> Date: Fri, 6 Oct 2023 20:12:39 +0530 Subject: [PATCH 101/156] feat(ingest/snowflake): support profiling with sampling (#8902) Co-authored-by: Andrew Sikowitz --- .../ingestion/source/bigquery_v2/profiler.py | 127 ++++++---------- .../ingestion/source/ge_data_profiler.py | 32 +++-- .../ingestion/source/ge_profiling_config.py | 4 +- .../ingestion/source/redshift/profile.py | 93 ++---------- .../source/snowflake/snowflake_profiler.py | 135 +++++------------- .../source/sql/sql_generic_profiler.py | 105 +++++++++++++- 6 files changed, 209 insertions(+), 287 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/profiler.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/profiler.py index b3e88459917b3..8ae17600e0eea 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/profiler.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/profiler.py @@ -1,12 +1,9 @@ -import dataclasses import logging from datetime import datetime from typing import Dict, Iterable, List, Optional, Tuple, cast from dateutil.relativedelta import relativedelta -from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance -from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.workunit import MetadataWorkUnit from datahub.ingestion.source.bigquery_v2.bigquery_audit import BigqueryTableIdentifier from datahub.ingestion.source.bigquery_v2.bigquery_config import BigQueryV2Config @@ -15,7 +12,7 @@ RANGE_PARTITION_NAME, BigqueryTable, ) -from datahub.ingestion.source.ge_data_profiler import GEProfilerRequest +from datahub.ingestion.source.sql.sql_generic import BaseTable from datahub.ingestion.source.sql.sql_generic_profiler import ( GenericProfiler, TableProfilerRequest, @@ -25,12 +22,6 @@ logger = logging.getLogger(__name__) -@dataclasses.dataclass -class BigqueryProfilerRequest(GEProfilerRequest): - table: BigqueryTable - profile_table_level_only: bool = False - - class BigqueryProfiler(GenericProfiler): config: BigQueryV2Config report: BigQueryV2Report @@ -183,84 +174,54 @@ def get_workunits( ) # Emit the profile work unit - profile_request = self.get_bigquery_profile_request( - project=project_id, dataset=dataset, table=table - ) + profile_request = self.get_profile_request(table, dataset, project_id) if profile_request is not None: + self.report.report_entity_profiled(profile_request.pretty_name) profile_requests.append(profile_request) if len(profile_requests) == 0: return - yield from self.generate_wu_from_profile_requests(profile_requests) - - def generate_wu_from_profile_requests( - self, profile_requests: List[BigqueryProfilerRequest] - ) -> Iterable[MetadataWorkUnit]: - table_profile_requests = cast(List[TableProfilerRequest], profile_requests) - for request, profile in self.generate_profiles( - table_profile_requests, + yield from self.generate_profile_workunits( + profile_requests, self.config.profiling.max_workers, platform=self.platform, profiler_args=self.get_profile_args(), - ): - if request is None or profile is None: - continue - - request = cast(BigqueryProfilerRequest, request) - profile.sizeInBytes = request.table.size_in_bytes - # If table is partitioned we profile only one partition (if nothing set then the last one) - # but for table level we can use the rows_count from the table metadata - # This way even though column statistics only reflects one partition data but the rows count - # shows the proper count. - if profile.partitionSpec and profile.partitionSpec.partition: - profile.rowCount = request.table.rows_count - - dataset_name = request.pretty_name - dataset_urn = make_dataset_urn_with_platform_instance( - self.platform, - dataset_name, - self.config.platform_instance, - self.config.env, - ) - # We don't add to the profiler state if we only do table level profiling as it always happens - if self.state_handler and not request.profile_table_level_only: - self.state_handler.add_to_state( - dataset_urn, int(datetime.now().timestamp() * 1000) - ) - - yield MetadataChangeProposalWrapper( - entityUrn=dataset_urn, aspect=profile - ).as_workunit() + ) - def get_bigquery_profile_request( - self, project: str, dataset: str, table: BigqueryTable - ) -> Optional[BigqueryProfilerRequest]: - skip_profiling = False - profile_table_level_only = self.config.profiling.profile_table_level_only - dataset_name = BigqueryTableIdentifier( - project_id=project, dataset=dataset, table=table.name + def get_dataset_name(self, table_name: str, schema_name: str, db_name: str) -> str: + return BigqueryTableIdentifier( + project_id=db_name, dataset=schema_name, table=table_name ).get_table_name() - if not self.is_dataset_eligible_for_profiling( - dataset_name, table.last_altered, table.size_in_bytes, table.rows_count - ): - profile_table_level_only = True - self.report.num_tables_not_eligible_profiling[f"{project}.{dataset}"] += 1 - if not table.column_count: - skip_profiling = True + def get_batch_kwargs( + self, table: BaseTable, schema_name: str, db_name: str + ) -> dict: + return dict( + schema=db_name, # + table=f"{schema_name}.{table.name}", # . + ) - if skip_profiling: - if self.config.profiling.report_dropped_profiles: - self.report.report_dropped(f"profile of {dataset_name}") + def get_profile_request( + self, table: BaseTable, schema_name: str, db_name: str + ) -> Optional[TableProfilerRequest]: + profile_request = super().get_profile_request(table, schema_name, db_name) + + if not profile_request: return None + # Below code handles profiling changes required for partitioned or sharded tables + # 1. Skip profile if partition profiling is disabled. + # 2. Else update `profile_request.batch_kwargs` with partition and custom_sql + + bq_table = cast(BigqueryTable, table) (partition, custom_sql) = self.generate_partition_profiler_query( - project, dataset, table, self.config.profiling.partition_datetime + db_name, schema_name, bq_table, self.config.profiling.partition_datetime ) - if partition is None and table.partition_info: + + if partition is None and bq_table.partition_info: self.report.report_warning( "profile skipped as partitioned table is empty or partition id or type was invalid", - dataset_name, + profile_request.pretty_name, ) return None if ( @@ -268,24 +229,20 @@ def get_bigquery_profile_request( and not self.config.profiling.partition_profiling_enabled ): logger.debug( - f"{dataset_name} and partition {partition} is skipped because profiling.partition_profiling_enabled property is disabled" + f"{profile_request.pretty_name} and partition {partition} is skipped because profiling.partition_profiling_enabled property is disabled" ) self.report.profiling_skipped_partition_profiling_disabled.append( - dataset_name + profile_request.pretty_name ) return None - self.report.report_entity_profiled(dataset_name) - logger.debug(f"Preparing profiling request for {dataset_name}") - profile_request = BigqueryProfilerRequest( - pretty_name=dataset_name, - batch_kwargs=dict( - schema=project, - table=f"{dataset}.{table.name}", - custom_sql=custom_sql, - partition=partition, - ), - table=table, - profile_table_level_only=profile_table_level_only, - ) + if partition: + logger.debug("Updating profiling request for partitioned/sharded tables") + profile_request.batch_kwargs.update( + dict( + custom_sql=custom_sql, + partition=partition, + ) + ) + return profile_request diff --git a/metadata-ingestion/src/datahub/ingestion/source/ge_data_profiler.py b/metadata-ingestion/src/datahub/ingestion/source/ge_data_profiler.py index 01e083d566168..9f6ac9dd21164 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/ge_data_profiler.py +++ b/metadata-ingestion/src/datahub/ingestion/source/ge_data_profiler.py @@ -273,6 +273,7 @@ class _SingleDatasetProfiler(BasicDatasetProfilerBase): partition: Optional[str] config: GEProfilingConfig report: SQLSourceReport + custom_sql: Optional[str] query_combiner: SQLAlchemyQueryCombiner @@ -596,16 +597,8 @@ def generate_dataset_profile( # noqa: C901 (complexity) "catch_exceptions", self.config.catch_exceptions ) - profile = DatasetProfileClass(timestampMillis=get_sys_time()) - if self.partition: - profile.partitionSpec = PartitionSpecClass(partition=self.partition) - elif self.config.limit and self.config.offset: - profile.partitionSpec = PartitionSpecClass( - type=PartitionTypeClass.QUERY, - partition=json.dumps( - dict(limit=self.config.limit, offset=self.config.offset) - ), - ) + profile = self.init_profile() + profile.fieldProfiles = [] self._get_dataset_rows(profile) @@ -740,6 +733,24 @@ def generate_dataset_profile( # noqa: C901 (complexity) self.query_combiner.flush() return profile + def init_profile(self): + profile = DatasetProfileClass(timestampMillis=get_sys_time()) + if self.partition: + profile.partitionSpec = PartitionSpecClass(partition=self.partition) + elif self.config.limit: + profile.partitionSpec = PartitionSpecClass( + type=PartitionTypeClass.QUERY, + partition=json.dumps( + dict(limit=self.config.limit, offset=self.config.offset) + ), + ) + elif self.custom_sql: + profile.partitionSpec = PartitionSpecClass( + type=PartitionTypeClass.QUERY, partition="SAMPLE" + ) + + return profile + def update_dataset_batch_use_sampling(self, profile: DatasetProfileClass) -> None: if ( self.dataset.engine.dialect.name.lower() == BIGQUERY @@ -1064,6 +1075,7 @@ def _generate_single_profile( partition, self.config, self.report, + custom_sql, query_combiner, ).generate_dataset_profile() diff --git a/metadata-ingestion/src/datahub/ingestion/source/ge_profiling_config.py b/metadata-ingestion/src/datahub/ingestion/source/ge_profiling_config.py index 77761c529ba0b..24a3e520d8caf 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/ge_profiling_config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/ge_profiling_config.py @@ -157,12 +157,12 @@ class GEProfilingConfig(ConfigModel): ) use_sampling: bool = Field( default=True, - description="Whether to profile column level stats on sample of table. Only BigQuery supports this. " + description="Whether to profile column level stats on sample of table. Only BigQuery and Snowflake support this. " "If enabled, profiling is done on rows sampled from table. Sampling is not done for smaller tables. ", ) sample_size: int = Field( - default=1000, + default=10000, description="Number of rows to be sampled from table for column level profiling." "Applicable only if `use_sampling` is set to True.", ) diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/profile.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/profile.py index e983734082b1d..771636e8498a3 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/redshift/profile.py +++ b/metadata-ingestion/src/datahub/ingestion/source/redshift/profile.py @@ -1,33 +1,19 @@ -import dataclasses import logging -from datetime import datetime -from typing import Dict, Iterable, List, Optional, Union, cast +from typing import Dict, Iterable, List, Optional, Union -from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance -from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.workunit import MetadataWorkUnit -from datahub.ingestion.source.ge_data_profiler import GEProfilerRequest from datahub.ingestion.source.redshift.config import RedshiftConfig from datahub.ingestion.source.redshift.redshift_schema import ( RedshiftTable, RedshiftView, ) from datahub.ingestion.source.redshift.report import RedshiftReport -from datahub.ingestion.source.sql.sql_generic_profiler import ( - GenericProfiler, - TableProfilerRequest, -) +from datahub.ingestion.source.sql.sql_generic_profiler import GenericProfiler from datahub.ingestion.source.state.profiling_state_handler import ProfilingHandler logger = logging.getLogger(__name__) -@dataclasses.dataclass -class RedshiftProfilerRequest(GEProfilerRequest): - table: Union[RedshiftTable, RedshiftView] - profile_table_level_only: bool = False - - class RedshiftProfiler(GenericProfiler): config: RedshiftConfig report: RedshiftReport @@ -63,80 +49,21 @@ def get_workunits( continue for table in tables[db].get(schema, {}): # Emit the profile work unit - profile_request = self.get_redshift_profile_request( - table, schema, db - ) + profile_request = self.get_profile_request(table, schema, db) if profile_request is not None: + self.report.report_entity_profiled(profile_request.pretty_name) profile_requests.append(profile_request) if len(profile_requests) == 0: continue - table_profile_requests = cast(List[TableProfilerRequest], profile_requests) - for request, profile in self.generate_profiles( - table_profile_requests, + + yield from self.generate_profile_workunits( + profile_requests, self.config.profiling.max_workers, db, platform=self.platform, profiler_args=self.get_profile_args(), - ): - if profile is None: - continue - request = cast(RedshiftProfilerRequest, request) - - profile.sizeInBytes = request.table.size_in_bytes - dataset_name = request.pretty_name - dataset_urn = make_dataset_urn_with_platform_instance( - self.platform, - dataset_name, - self.config.platform_instance, - self.config.env, - ) - - # We don't add to the profiler state if we only do table level profiling as it always happens - if self.state_handler and not request.profile_table_level_only: - self.state_handler.add_to_state( - dataset_urn, int(datetime.now().timestamp() * 1000) - ) - - yield MetadataChangeProposalWrapper( - entityUrn=dataset_urn, aspect=profile - ).as_workunit() - - def get_redshift_profile_request( - self, - table: Union[RedshiftTable, RedshiftView], - schema_name: str, - db_name: str, - ) -> Optional[RedshiftProfilerRequest]: - skip_profiling = False - profile_table_level_only = self.config.profiling.profile_table_level_only - dataset_name = f"{db_name}.{schema_name}.{table.name}".lower() - if not self.is_dataset_eligible_for_profiling( - dataset_name, table.last_altered, table.size_in_bytes, table.rows_count - ): - # Profile only table level if dataset is filtered from profiling - # due to size limits alone - if self.is_dataset_eligible_for_profiling( - dataset_name, table.last_altered, 0, 0 - ): - profile_table_level_only = True - else: - skip_profiling = True - - if len(table.columns) == 0: - skip_profiling = True - - if skip_profiling: - if self.config.profiling.report_dropped_profiles: - self.report.report_dropped(f"profile of {dataset_name}") - return None + ) - self.report.report_entity_profiled(dataset_name) - logger.debug(f"Preparing profiling request for {dataset_name}") - profile_request = RedshiftProfilerRequest( - pretty_name=dataset_name, - batch_kwargs=dict(schema=schema_name, table=table.name), - table=table, - profile_table_level_only=profile_table_level_only, - ) - return profile_request + def get_dataset_name(self, table_name: str, schema_name: str, db_name: str) -> str: + return f"{db_name}.{schema_name}.{table_name}".lower() diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_profiler.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_profiler.py index 5f5e8e4bcdea3..24275dcdff34d 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_profiler.py +++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_profiler.py @@ -1,20 +1,12 @@ -import dataclasses import logging -from datetime import datetime -from typing import Callable, Dict, Iterable, List, Optional, cast +from typing import Callable, Dict, Iterable, List, Optional from snowflake.sqlalchemy import snowdialect from sqlalchemy import create_engine, inspect from sqlalchemy.sql import sqltypes -from datahub.configuration.pattern_utils import is_schema_allowed -from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance -from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.workunit import MetadataWorkUnit -from datahub.ingestion.source.ge_data_profiler import ( - DatahubGEProfiler, - GEProfilerRequest, -) +from datahub.ingestion.source.ge_data_profiler import DatahubGEProfiler from datahub.ingestion.source.snowflake.snowflake_config import SnowflakeV2Config from datahub.ingestion.source.snowflake.snowflake_query import SnowflakeQuery from datahub.ingestion.source.snowflake.snowflake_report import SnowflakeV2Report @@ -23,10 +15,8 @@ SnowflakeTable, ) from datahub.ingestion.source.snowflake.snowflake_utils import SnowflakeCommonMixin -from datahub.ingestion.source.sql.sql_generic_profiler import ( - GenericProfiler, - TableProfilerRequest, -) +from datahub.ingestion.source.sql.sql_generic import BaseTable +from datahub.ingestion.source.sql.sql_generic_profiler import GenericProfiler from datahub.ingestion.source.state.profiling_state_handler import ProfilingHandler snowdialect.ischema_names["GEOGRAPHY"] = sqltypes.NullType @@ -35,12 +25,6 @@ logger = logging.getLogger(__name__) -@dataclasses.dataclass -class SnowflakeProfilerRequest(GEProfilerRequest): - table: SnowflakeTable - profile_table_level_only: bool = False - - class SnowflakeProfiler(GenericProfiler, SnowflakeCommonMixin): def __init__( self, @@ -65,101 +49,52 @@ def get_workunits( profile_requests = [] for schema in database.schemas: - if not is_schema_allowed( - self.config.schema_pattern, - schema.name, - database.name, - self.config.match_fully_qualified_names, - ): - continue - for table in db_tables[schema.name]: - profile_request = self.get_snowflake_profile_request( + profile_request = self.get_profile_request( table, schema.name, database.name ) if profile_request is not None: + self.report.report_entity_profiled(profile_request.pretty_name) profile_requests.append(profile_request) if len(profile_requests) == 0: return - table_profile_requests = cast(List[TableProfilerRequest], profile_requests) - - for request, profile in self.generate_profiles( - table_profile_requests, + yield from self.generate_profile_workunits( + profile_requests, self.config.profiling.max_workers, database.name, platform=self.platform, profiler_args=self.get_profile_args(), - ): - if profile is None: - continue - profile.sizeInBytes = cast( - SnowflakeProfilerRequest, request - ).table.size_in_bytes - dataset_name = request.pretty_name - dataset_urn = make_dataset_urn_with_platform_instance( - self.platform, - dataset_name, - self.config.platform_instance, - self.config.env, - ) - - # We don't add to the profiler state if we only do table level profiling as it always happens - if self.state_handler: - self.state_handler.add_to_state( - dataset_urn, int(datetime.now().timestamp() * 1000) - ) - - yield MetadataChangeProposalWrapper( - entityUrn=dataset_urn, aspect=profile - ).as_workunit() + ) - def get_snowflake_profile_request( - self, - table: SnowflakeTable, - schema_name: str, - db_name: str, - ) -> Optional[SnowflakeProfilerRequest]: - skip_profiling = False - profile_table_level_only = self.config.profiling.profile_table_level_only - dataset_name = self.get_dataset_identifier(table.name, schema_name, db_name) - if not self.is_dataset_eligible_for_profiling( - dataset_name, table.last_altered, table.size_in_bytes, table.rows_count + def get_dataset_name(self, table_name: str, schema_name: str, db_name: str) -> str: + return self.get_dataset_identifier(table_name, schema_name, db_name) + + def get_batch_kwargs( + self, table: BaseTable, schema_name: str, db_name: str + ) -> dict: + custom_sql = None + if ( + not self.config.profiling.limit + and self.config.profiling.use_sampling + and table.rows_count + and table.rows_count > self.config.profiling.sample_size ): - # Profile only table level if dataset is filtered from profiling - # due to size limits alone - if self.is_dataset_eligible_for_profiling( - dataset_name, table.last_altered, 0, 0 - ): - profile_table_level_only = True - else: - skip_profiling = True - - if len(table.columns) == 0: - skip_profiling = True - - if skip_profiling: - if self.config.profiling.report_dropped_profiles: - self.report.report_dropped(f"profile of {dataset_name}") - return None - - self.report.report_entity_profiled(dataset_name) - logger.debug(f"Preparing profiling request for {dataset_name}") - profile_request = SnowflakeProfilerRequest( - pretty_name=dataset_name, - batch_kwargs=dict( - schema=schema_name, - table=table.name, - # Lowercase/Mixedcase table names in Snowflake do not work by default. - # We need to pass `use_quoted_name=True` for such tables as mentioned here - - # https://github.com/great-expectations/great_expectations/pull/2023 - use_quoted_name=(table.name != table.name.upper()), - ), - table=table, - profile_table_level_only=profile_table_level_only, - ) - return profile_request + # GX creates a temporary table from query if query is passed as batch kwargs. + # We are using fraction-based sampling here, instead of fixed-size sampling because + # Fixed-size sampling can be slower than equivalent fraction-based sampling + # as per https://docs.snowflake.com/en/sql-reference/constructs/sample#performance-considerations + sample_pc = 100 * self.config.profiling.sample_size / table.rows_count + custom_sql = f'select * from "{db_name}"."{schema_name}"."{table.name}" TABLESAMPLE ({sample_pc:.3f})' + return { + **super().get_batch_kwargs(table, schema_name, db_name), + # Lowercase/Mixedcase table names in Snowflake do not work by default. + # We need to pass `use_quoted_name=True` for such tables as mentioned here - + # https://github.com/great-expectations/great_expectations/pull/2023 + "use_quoted_name": (table.name != table.name.upper()), + "custom_sql": custom_sql, + } def get_profiler_instance( self, db_name: Optional[str] = None diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_generic_profiler.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_generic_profiler.py index 344c114d464a9..aaeee5717a867 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_generic_profiler.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_generic_profiler.py @@ -1,12 +1,15 @@ import logging +from abc import abstractmethod from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone -from typing import Dict, Iterable, List, Optional, Tuple, Union, cast +from typing import Dict, Iterable, List, Optional, Union, cast from sqlalchemy import create_engine, inspect from sqlalchemy.engine.reflection import Inspector from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.ingestion.api.workunit import MetadataWorkUnit from datahub.ingestion.source.ge_data_profiler import ( DatahubGEProfiler, GEProfilerRequest, @@ -16,7 +19,7 @@ from datahub.ingestion.source.sql.sql_generic import BaseTable, BaseView from datahub.ingestion.source.state.profiling_state_handler import ProfilingHandler from datahub.metadata.com.linkedin.pegasus2avro.dataset import DatasetProfile -from datahub.metadata.schema_classes import DatasetProfileClass +from datahub.metadata.com.linkedin.pegasus2avro.timeseries import PartitionType from datahub.utilities.stats_collections import TopKDict, int_top_k_dict @@ -63,14 +66,14 @@ def __init__( self.platform = platform self.state_handler = state_handler - def generate_profiles( + def generate_profile_workunits( self, requests: List[TableProfilerRequest], max_workers: int, db_name: Optional[str] = None, platform: Optional[str] = None, profiler_args: Optional[Dict] = None, - ) -> Iterable[Tuple[GEProfilerRequest, Optional[DatasetProfileClass]]]: + ) -> Iterable[MetadataWorkUnit]: ge_profile_requests: List[GEProfilerRequest] = [ cast(GEProfilerRequest, request) for request in requests @@ -80,21 +83,109 @@ def generate_profiles( request for request in requests if request.profile_table_level_only ] for request in table_level_profile_requests: - profile = DatasetProfile( + table_level_profile = DatasetProfile( timestampMillis=int(datetime.now().timestamp() * 1000), columnCount=request.table.column_count, rowCount=request.table.rows_count, sizeInBytes=request.table.size_in_bytes, ) - yield (request, profile) + dataset_urn = self.dataset_urn_builder(request.pretty_name) + yield MetadataChangeProposalWrapper( + entityUrn=dataset_urn, aspect=table_level_profile + ).as_workunit() if not ge_profile_requests: return # Otherwise, if column level profiling is enabled, use GE profiler. ge_profiler = self.get_profiler_instance(db_name) - yield from ge_profiler.generate_profiles( + + for ge_profiler_request, profile in ge_profiler.generate_profiles( ge_profile_requests, max_workers, platform, profiler_args + ): + if profile is None: + continue + + request = cast(TableProfilerRequest, ge_profiler_request) + profile.sizeInBytes = request.table.size_in_bytes + + # If table is partitioned we profile only one partition (if nothing set then the last one) + # but for table level we can use the rows_count from the table metadata + # This way even though column statistics only reflects one partition data but the rows count + # shows the proper count. + if ( + profile.partitionSpec + and profile.partitionSpec.type != PartitionType.FULL_TABLE + ): + profile.rowCount = request.table.rows_count + + dataset_urn = self.dataset_urn_builder(request.pretty_name) + + # We don't add to the profiler state if we only do table level profiling as it always happens + if self.state_handler: + self.state_handler.add_to_state( + dataset_urn, int(datetime.now().timestamp() * 1000) + ) + yield MetadataChangeProposalWrapper( + entityUrn=dataset_urn, aspect=profile + ).as_workunit() + + def dataset_urn_builder(self, dataset_name: str) -> str: + return make_dataset_urn_with_platform_instance( + self.platform, + dataset_name, + self.config.platform_instance, + self.config.env, + ) + + @abstractmethod + def get_dataset_name(self, table_name: str, schema_name: str, db_name: str) -> str: + pass + + def get_profile_request( + self, table: BaseTable, schema_name: str, db_name: str + ) -> Optional[TableProfilerRequest]: + skip_profiling = False + profile_table_level_only = self.config.profiling.profile_table_level_only + dataset_name = self.get_dataset_name(table.name, schema_name, db_name) + if not self.is_dataset_eligible_for_profiling( + dataset_name, table.last_altered, table.size_in_bytes, table.rows_count + ): + # Profile only table level if dataset is filtered from profiling + # due to size limits alone + if self.is_dataset_eligible_for_profiling( + dataset_name, table.last_altered, 0, 0 + ): + profile_table_level_only = True + else: + skip_profiling = True + self.report.num_tables_not_eligible_profiling[ + f"{db_name}.{schema_name}" + ] += 1 + + if table.column_count == 0: + skip_profiling = True + + if skip_profiling: + if self.config.profiling.report_dropped_profiles: + self.report.report_dropped(f"profile of {dataset_name}") + return None + + logger.debug(f"Preparing profiling request for {dataset_name}") + profile_request = TableProfilerRequest( + pretty_name=dataset_name, + batch_kwargs=self.get_batch_kwargs(table, schema_name, db_name), + table=table, + profile_table_level_only=profile_table_level_only, + ) + return profile_request + + def get_batch_kwargs( + self, table: BaseTable, schema_name: str, db_name: str + ) -> dict: + return dict( + schema=schema_name, + table=table.name, ) def get_inspectors(self) -> Iterable[Inspector]: From c0feceb76fbf607e2883b7f2960eaf6c757629e4 Mon Sep 17 00:00:00 2001 From: Kos Korchak <97058061+kkorchak@users.noreply.github.com> Date: Fri, 6 Oct 2023 17:10:24 -0400 Subject: [PATCH 102/156] test(): Manage Access Tokens Cypress test (#8936) --- .../src/app/settings/AccessTokenModal.tsx | 4 +- .../src/app/settings/AccessTokens.tsx | 7 ++- .../src/app/settings/CreateTokenModal.tsx | 18 +++++--- .../e2e/settings/manage_access_tokens.js | 43 +++++++++++++++++++ .../tests/cypress/cypress/support/commands.js | 6 +++ 5 files changed, 70 insertions(+), 8 deletions(-) create mode 100644 smoke-test/tests/cypress/cypress/e2e/settings/manage_access_tokens.js diff --git a/datahub-web-react/src/app/settings/AccessTokenModal.tsx b/datahub-web-react/src/app/settings/AccessTokenModal.tsx index 0303db656c2a8..10427210d0692 100644 --- a/datahub-web-react/src/app/settings/AccessTokenModal.tsx +++ b/datahub-web-react/src/app/settings/AccessTokenModal.tsx @@ -60,7 +60,7 @@ export const AccessTokenModal = ({ visible, onClose, accessToken, expiresInText onCancel={onClose} footer={ <> - @@ -81,7 +81,7 @@ export const AccessTokenModal = ({ visible, onClose, accessToken, expiresInText Token{expiresInText} - 
{accessToken}
+ 
{accessToken}
 diff --git a/datahub-web-react/src/app/settings/AccessTokens.tsx b/datahub-web-react/src/app/settings/AccessTokens.tsx index 02ff3f1cd304c..c7a015de392da 100644 --- a/datahub-web-react/src/app/settings/AccessTokens.tsx +++ b/datahub-web-react/src/app/settings/AccessTokens.tsx @@ -199,7 +199,12 @@ export const AccessTokens = () => { key: 'x', render: (_, record: any) => ( - diff --git a/datahub-web-react/src/app/settings/CreateTokenModal.tsx b/datahub-web-react/src/app/settings/CreateTokenModal.tsx index 6038a86e23303..3cc446651efcb 100644 --- a/datahub-web-react/src/app/settings/CreateTokenModal.tsx +++ b/datahub-web-react/src/app/settings/CreateTokenModal.tsx @@ -117,10 +117,15 @@ export default function CreateTokenModal({ currentUserUrn, visible, onClose, onC onCancel={onModalClose} footer={ <> - - @@ -148,18 +153,21 @@ export default function CreateTokenModal({ currentUserUrn, visible, onClose, onC ]} hasFeedback > - + Description}> An optional description for your new token. - + Expires in - + {ACCESS_TOKEN_DURATIONS.map((duration) => ( diff --git a/smoke-test/tests/cypress/cypress/e2e/settings/manage_access_tokens.js b/smoke-test/tests/cypress/cypress/e2e/settings/manage_access_tokens.js new file mode 100644 index 0000000000000..7a77c2b77df5b --- /dev/null +++ b/smoke-test/tests/cypress/cypress/e2e/settings/manage_access_tokens.js @@ -0,0 +1,43 @@ +import { aliasQuery, hasOperationName } from "../utils"; +const test_id = Math.floor(Math.random() * 100000); + +describe("manage access tokens", () => { + before(() => { + cy.intercept("POST", "/api/v2/graphql", (req) => { + aliasQuery(req, "appConfig"); + }); + }); + + const setTokenAuthEnabledFlag = (isOn) => { + cy.intercept("POST", "/api/v2/graphql", (req) => { + if (hasOperationName(req, "appConfig")) { + req.reply((res) => { + res.body.data.appConfig.authConfig.tokenAuthEnabled = isOn; + }); + } + }); + }; + + it("create and revoke access token", () => { + //create access token, verify token on ui + setTokenAuthEnabledFlag(true); + cy.loginWithCredentials(); + cy.goToAccessTokenSettings(); + cy.clickOptionWithTestId("add-token-button"); + cy.enterTextInTestId("create-access-token-name", "Token Name" + test_id); + cy.enterTextInTestId("create-access-token-description", "Token Description" + test_id); + cy.clickOptionWithTestId("create-access-token-button"); + cy.waitTextVisible("New Personal Access Token"); + cy.get('[data-testid="access-token-value"]').should("be.visible"); + cy.get('[data-testid="access-token-value"]').invoke('text').should('match', /^[a-zA-Z0-9-_]+\.[a-zA-Z0-9-_]+\.[a-zA-Z0-9-_]+$/); + cy.clickOptionWithTestId("access-token-modal-close-button"); + //revoke access token, verify token removed from ui + cy.waitTextVisible("Token Name" + test_id); + cy.waitTextVisible("Token Description" + test_id); + cy.clickOptionWithTestId("revoke-token-button"); + cy.waitTextVisible("Are you sure you want to revoke this token?"); + cy.clickOptionWithText("Yes"); + cy.ensureTextNotPresent("Token Name" + test_id); + cy.ensureTextNotPresent("Token Description" + test_id); + }); +}); \ No newline at end of file diff --git a/smoke-test/tests/cypress/cypress/support/commands.js b/smoke-test/tests/cypress/cypress/support/commands.js index 8bfe7305c001f..64bc1253fc383 100644 --- a/smoke-test/tests/cypress/cypress/support/commands.js +++ b/smoke-test/tests/cypress/cypress/support/commands.js @@ -84,6 +84,12 @@ Cypress.Commands.add("goToOwnershipTypesSettings", () => { cy.waitTextVisible("Manage Ownership"); }); +Cypress.Commands.add("goToAccessTokenSettings", () => { + cy.visit("/settings/tokens"); + cy.waitTextVisible("Manage Access Tokens"); + cy.wait(3000); +}); + Cypress.Commands.add("goToIngestionPage", () => { cy.visit("/ingestion"); cy.waitTextVisible("Manage Ingestion"); From b191abbc5bb32a0a3c895facdff14d146da9fb74 Mon Sep 17 00:00:00 2001 From: Kos Korchak <97058061+kkorchak@users.noreply.github.com> Date: Fri, 6 Oct 2023 17:11:57 -0400 Subject: [PATCH 103/156] test(): Nested domains cypress test (#8879) --- .../src/app/domain/CreateDomainModal.tsx | 5 +- .../nestedDomains/ManageDomainsPageV2.tsx | 7 ++- .../domainNavigator/DomainNode.tsx | 2 +- .../shared/EntityDropdown/EntityDropdown.tsx | 4 +- .../shared/EntityDropdown/MoveDomainModal.tsx | 5 +- .../cypress/e2e/domains/nested_domains.js | 53 +++++++++++++++++++ 6 files changed, 70 insertions(+), 6 deletions(-) create mode 100644 smoke-test/tests/cypress/cypress/e2e/domains/nested_domains.js diff --git a/datahub-web-react/src/app/domain/CreateDomainModal.tsx b/datahub-web-react/src/app/domain/CreateDomainModal.tsx index ca1bc30596003..606444d34bdc9 100644 --- a/datahub-web-react/src/app/domain/CreateDomainModal.tsx +++ b/datahub-web-react/src/app/domain/CreateDomainModal.tsx @@ -191,7 +191,10 @@ export default function CreateDomainModal({ onClose, onCreate }: Props) { rules={[{ whitespace: true }, { min: 1, max: 500 }]} hasFeedback > - + diff --git a/datahub-web-react/src/app/domain/nestedDomains/ManageDomainsPageV2.tsx b/datahub-web-react/src/app/domain/nestedDomains/ManageDomainsPageV2.tsx index 0e5c035df00c1..b69f0c5458b5d 100644 --- a/datahub-web-react/src/app/domain/nestedDomains/ManageDomainsPageV2.tsx +++ b/datahub-web-react/src/app/domain/nestedDomains/ManageDomainsPageV2.tsx @@ -42,7 +42,12 @@ export default function ManageDomainsPageV2() {
-
diff --git a/datahub-web-react/src/app/domain/nestedDomains/domainNavigator/DomainNode.tsx b/datahub-web-react/src/app/domain/nestedDomains/domainNavigator/DomainNode.tsx index 09c8e13853bb7..bf70bd043fd4a 100644 --- a/datahub-web-react/src/app/domain/nestedDomains/domainNavigator/DomainNode.tsx +++ b/datahub-web-react/src/app/domain/nestedDomains/domainNavigator/DomainNode.tsx @@ -103,7 +103,7 @@ export default function DomainNode({ domain, numDomainChildren, domainUrnToHide, return ( <> - + {hasDomainChildren && ( diff --git a/datahub-web-react/src/app/entity/shared/EntityDropdown/EntityDropdown.tsx b/datahub-web-react/src/app/entity/shared/EntityDropdown/EntityDropdown.tsx index be975249b2670..bfb7ff7e540c4 100644 --- a/datahub-web-react/src/app/entity/shared/EntityDropdown/EntityDropdown.tsx +++ b/datahub-web-react/src/app/entity/shared/EntityDropdown/EntityDropdown.tsx @@ -203,7 +203,7 @@ function EntityDropdown(props: Props) { disabled={isMoveDisabled(entityType, entityData, me.platformPrivileges)} onClick={() => setIsMoveModalVisible(true)} > - +  Move @@ -223,7 +223,7 @@ function EntityDropdown(props: Props) { : undefined } > - +  Delete diff --git a/datahub-web-react/src/app/entity/shared/EntityDropdown/MoveDomainModal.tsx b/datahub-web-react/src/app/entity/shared/EntityDropdown/MoveDomainModal.tsx index cdbf6fdabf3c9..3826f934c1c25 100644 --- a/datahub-web-react/src/app/entity/shared/EntityDropdown/MoveDomainModal.tsx +++ b/datahub-web-react/src/app/entity/shared/EntityDropdown/MoveDomainModal.tsx @@ -67,6 +67,7 @@ function MoveDomainModal(props: Props) { return ( Cancel - + } > diff --git a/smoke-test/tests/cypress/cypress/e2e/domains/nested_domains.js b/smoke-test/tests/cypress/cypress/e2e/domains/nested_domains.js new file mode 100644 index 0000000000000..a2d4de0f51659 --- /dev/null +++ b/smoke-test/tests/cypress/cypress/e2e/domains/nested_domains.js @@ -0,0 +1,53 @@ +const domainName = "CypressNestedDomain"; +const domainDescription = "CypressNestedDomainDescription"; + +describe("nested domains test", () => { + + it("create a domain, move under parent, remove domain", () => { + // Create a new domain without a parent + cy.loginWithCredentials(); + cy.goToDomainList(); + cy.clickOptionWithTestId("domains-new-domain-button"); + cy.get('[data-testid="create-domain-name"]').click().type(domainName); + cy.get('[data-testid="create-domain-description"]').click().type(domainDescription); + cy.clickOptionWithTestId("create-domain-button"); + cy.waitTextVisible(domainName); + + // Ensure the new domain has no parent in the navigation sidebar + cy.waitTextVisible(domainDescription); + + // Move a domain from the root level to be under a parent domain + cy.clickOptionWithText(domainName); + cy.openThreeDotDropdown(); + cy.clickOptionWithTestId("entity-menu-move-button"); + cy.get('[data-testid="move-domain-modal"]').contains("Marketing").click({force: true}); + cy.get('[data-testid="move-domain-modal"]').contains("Marketing").should("be.visible"); + cy.clickOptionWithTestId("move-domain-modal-move-button").wait(5000); + + // Wnsure domain is no longer on the sidebar navigator at the top level but shows up under the parent + cy.goToDomainList(); + cy.ensureTextNotPresent(domainName); + cy.ensureTextNotPresent(domainDescription); + cy.waitTextVisible("1 sub-domain"); + + // Move a domain from under a parent domain to the root level + cy.get('[data-testid="domain-list-item"]').contains("Marketing").prev().click(); + cy.clickOptionWithText(domainName); + cy.openThreeDotDropdown(); + cy.clickOptionWithTestId("entity-menu-move-button"); + cy.clickOptionWithTestId("move-domain-modal-move-button").wait(5000); + cy.goToDomainList(); + cy.waitTextVisible(domainName); + cy.waitTextVisible(domainDescription); + + // Delete a domain + cy.clickOptionWithText(domainName).wait(3000); + cy.openThreeDotDropdown(); + cy.clickOptionWithTestId("entity-menu-delete-button"); + cy.waitTextVisible("Are you sure you want to remove this Domain?"); + cy.clickOptionWithText("Yes"); + cy.waitTextVisible("Deleted Domain!"); + cy.ensureTextNotPresent(domainName); + cy.ensureTextNotPresent(domainDescription); + }); +}); \ No newline at end of file From 93958302d529a65021c78f880347930297854692 Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Sun, 8 Oct 2023 13:26:48 -0400 Subject: [PATCH 104/156] feat(models/assertion): Add SQL Assertions (#8969) --- .../com/linkedin/assertion/AssertionInfo.pdl | 17 ++++- .../linkedin/assertion/SqlAssertionInfo.pdl | 67 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 metadata-models/src/main/pegasus/com/linkedin/assertion/SqlAssertionInfo.pdl diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionInfo.pdl index ae2a58028057b..e161270145a88 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionInfo.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionInfo.pdl @@ -32,6 +32,11 @@ record AssertionInfo includes CustomProperties, ExternalReference { */ VOLUME + /** + * A raw SQL-statement based assertion + */ + SQL + /** * A schema or structural assertion. * @@ -56,7 +61,12 @@ record AssertionInfo includes CustomProperties, ExternalReference { volumeAssertion: optional VolumeAssertionInfo /** - * An schema Assertion definition. This field is populated when the type is DATASET_SCHEMA + * A SQL Assertion definition. This field is populated when the type is SQL. + */ + sqlAssertion: optional SqlAssertionInfo + + /** + * An schema Assertion definition. This field is populated when the type is DATA_SCHEMA */ schemaAssertion: optional SchemaAssertionInfo @@ -67,4 +77,9 @@ record AssertionInfo includes CustomProperties, ExternalReference { * the platform where it was ingested from. */ source: optional AssertionSource + + /** + * An optional human-readable description of the assertion + */ + description: optional string } \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/SqlAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/SqlAssertionInfo.pdl new file mode 100644 index 0000000000000..f6ce738252f35 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/assertion/SqlAssertionInfo.pdl @@ -0,0 +1,67 @@ +namespace com.linkedin.assertion + +import com.linkedin.common.Urn +import com.linkedin.dataset.DatasetFilter + +/** +* Attributes defining a SQL Assertion +*/ +record SqlAssertionInfo { + /** + * The type of the SQL assertion being monitored. + */ + @Searchable = {} + type: enum SqlAssertionType { + /** + * A SQL Metric Assertion, e.g. one based on a numeric value returned by an arbitrary SQL query. + */ + METRIC + /** + * A SQL assertion that is evaluated against the CHANGE in a metric assertion + * over time. + */ + METRIC_CHANGE + } + + /** + * The entity targeted by this SQL check. + */ + @Searchable = { + "fieldType": "URN" + } + @Relationship = { + "name": "Asserts", + "entityTypes": [ "dataset" ] + } + entity: Urn + + /** + * The SQL statement to be executed when evaluating the assertion (or computing the metric). + * This should be a valid and complete statement, executable by itself. + * + * Usually this should be a SELECT query statement. + */ + statement: string + + /** + * The type of the value used to evaluate the assertion: a fixed absolute value or a relative percentage. + * This value is required if the type is METRIC_CHANGE. + */ + changeType: optional AssertionValueChangeType + + /** + * The operator you'd like to apply to the result of the SQL query. + * + * Note that at this time, only numeric operators are valid inputs: + * GREATER_THAN, GREATER_THAN_OR_EQUAL_TO, EQUAL_TO, LESS_THAN, LESS_THAN_OR_EQUAL_TO, + * BETWEEN. + */ + operator: AssertionStdOperator + + /** + * The parameters you'd like to provide as input to the operator. + * + * Note that only numeric parameter types are valid inputs: NUMBER. + */ + parameters: AssertionStdParameters +} \ No newline at end of file From 8d175ef7ef1ae8ffada7b2df2fb711ac02a6785d Mon Sep 17 00:00:00 2001 From: Mayuri Nehate <33225191+mayurinehate@users.noreply.github.com> Date: Tue, 10 Oct 2023 02:04:25 +0530 Subject: [PATCH 105/156] feat(ingest): incremental lineage source helper (#8941) Co-authored-by: Harshal Sheth --- .../datahub/ingestion/api/source_helpers.py | 138 +++++++++- .../ingestion/source/bigquery_v2/bigquery.py | 3 +- .../source/snowflake/snowflake_v2.py | 9 + .../snowflake_privatelink_golden.json | 243 +++++++++++------ .../integration/snowflake/test_snowflake.py | 2 + .../snowflake/test_snowflake_failures.py | 6 +- .../snowflake/test_snowflake_stateful.py | 3 +- ...l_less_upstreams_in_gms_aspect_golden.json | 106 ++++++++ ...l_more_upstreams_in_gms_aspect_golden.json | 120 +++++++++ .../incremental_table_lineage_golden.json | 41 +++ .../test_incremental_lineage_helper.py | 244 ++++++++++++++++++ .../source_helpers}/test_source_helpers.py | 0 12 files changed, 829 insertions(+), 86 deletions(-) create mode 100644 metadata-ingestion/tests/unit/api/source_helpers/incremental_cll_less_upstreams_in_gms_aspect_golden.json create mode 100644 metadata-ingestion/tests/unit/api/source_helpers/incremental_cll_more_upstreams_in_gms_aspect_golden.json create mode 100644 metadata-ingestion/tests/unit/api/source_helpers/incremental_table_lineage_golden.json create mode 100644 metadata-ingestion/tests/unit/api/source_helpers/test_incremental_lineage_helper.py rename metadata-ingestion/tests/unit/{ => api/source_helpers}/test_source_helpers.py (100%) diff --git a/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py b/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py index 7fc15cf829678..42f970e97c95f 100644 --- a/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py +++ b/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py @@ -1,3 +1,4 @@ +import copy import logging from datetime import datetime, timezone from typing import ( @@ -15,9 +16,14 @@ ) from datahub.configuration.time_window_config import BaseTimeWindowConfig -from datahub.emitter.mce_builder import make_dataplatform_instance_urn +from datahub.emitter.mce_builder import ( + datahub_guid, + make_dataplatform_instance_urn, + set_aspect, +) from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.workunit import MetadataWorkUnit +from datahub.ingestion.graph.client import DataHubGraph from datahub.metadata.schema_classes import ( BrowsePathEntryClass, BrowsePathsClass, @@ -25,12 +31,17 @@ ChangeTypeClass, ContainerClass, DatasetUsageStatisticsClass, + FineGrainedLineageClass, MetadataChangeEventClass, MetadataChangeProposalClass, StatusClass, + SystemMetadataClass, TagKeyClass, TimeWindowSizeClass, + UpstreamClass, + UpstreamLineageClass, ) +from datahub.specific.dataset import DatasetPatchBuilder from datahub.telemetry import telemetry from datahub.utilities.urns.dataset_urn import DatasetUrn from datahub.utilities.urns.tag_urn import TagUrn @@ -366,3 +377,128 @@ def _prepend_platform_instance( return [BrowsePathEntryClass(id=urn, urn=urn)] + entries return entries + + +def auto_incremental_lineage( + graph: Optional[DataHubGraph], + incremental_lineage: bool, + include_column_level_lineage: bool, + stream: Iterable[MetadataWorkUnit], +) -> Iterable[MetadataWorkUnit]: + if not incremental_lineage: + yield from stream + return # early exit + + for wu in stream: + lineage_aspect: Optional[UpstreamLineageClass] = wu.get_aspect_of_type( + UpstreamLineageClass + ) + urn = wu.get_urn() + + if lineage_aspect: + if isinstance(wu.metadata, MetadataChangeEventClass): + set_aspect( + wu.metadata, None, UpstreamLineageClass + ) # we'll emit upstreamLineage separately below + if len(wu.metadata.proposedSnapshot.aspects) > 0: + yield wu + + yield _lineage_wu_via_read_modify_write( + graph, urn, lineage_aspect, wu.metadata.systemMetadata + ) if include_column_level_lineage else _convert_upstream_lineage_to_patch( + urn, lineage_aspect, wu.metadata.systemMetadata + ) + else: + yield wu + + +def _convert_upstream_lineage_to_patch( + urn: str, + aspect: UpstreamLineageClass, + system_metadata: Optional[SystemMetadataClass], +) -> MetadataWorkUnit: + patch_builder = DatasetPatchBuilder(urn, system_metadata) + for upstream in aspect.upstreams: + patch_builder.add_upstream_lineage(upstream) + mcp = next(iter(patch_builder.build())) + return MetadataWorkUnit(id=f"{urn}-upstreamLineage", mcp_raw=mcp) + + +def _lineage_wu_via_read_modify_write( + graph: Optional[DataHubGraph], + urn: str, + aspect: UpstreamLineageClass, + system_metadata: Optional[SystemMetadataClass], +) -> MetadataWorkUnit: + if graph is None: + raise ValueError( + "Failed to handle incremental lineage, DataHubGraph is missing. " + "Use `datahub-rest` sink OR provide `datahub-api` config in recipe. " + ) + gms_aspect = graph.get_aspect(urn, UpstreamLineageClass) + if gms_aspect: + new_aspect = _merge_upstream_lineage(aspect, gms_aspect) + else: + new_aspect = aspect + + return MetadataChangeProposalWrapper( + entityUrn=urn, aspect=new_aspect, systemMetadata=system_metadata + ).as_workunit() + + +def _merge_upstream_lineage( + new_aspect: UpstreamLineageClass, gms_aspect: UpstreamLineageClass +) -> UpstreamLineageClass: + merged_aspect = copy.deepcopy(gms_aspect) + + upstreams_map: Dict[str, UpstreamClass] = { + upstream.dataset: upstream for upstream in merged_aspect.upstreams + } + + upstreams_updated = False + fine_upstreams_updated = False + + for table_upstream in new_aspect.upstreams: + if table_upstream.dataset not in upstreams_map or ( + table_upstream.auditStamp.time + > upstreams_map[table_upstream.dataset].auditStamp.time + ): + upstreams_map[table_upstream.dataset] = table_upstream + upstreams_updated = True + + if upstreams_updated: + merged_aspect.upstreams = list(upstreams_map.values()) + + if new_aspect.fineGrainedLineages and merged_aspect.fineGrainedLineages: + fine_upstreams_map: Dict[str, FineGrainedLineageClass] = { + get_fine_grained_lineage_key(fine_upstream): fine_upstream + for fine_upstream in merged_aspect.fineGrainedLineages + } + for column_upstream in new_aspect.fineGrainedLineages: + column_upstream_key = get_fine_grained_lineage_key(column_upstream) + + if column_upstream_key not in fine_upstreams_map or ( + column_upstream.confidenceScore + > fine_upstreams_map[column_upstream_key].confidenceScore + ): + fine_upstreams_map[column_upstream_key] = column_upstream + fine_upstreams_updated = True + + if fine_upstreams_updated: + merged_aspect.fineGrainedLineages = list(fine_upstreams_map.values()) + else: + merged_aspect.fineGrainedLineages = ( + new_aspect.fineGrainedLineages or gms_aspect.fineGrainedLineages + ) + + return merged_aspect + + +def get_fine_grained_lineage_key(fine_upstream: FineGrainedLineageClass) -> str: + return datahub_guid( + { + "upstreams": sorted(fine_upstream.upstreams or []), + "downstreams": sorted(fine_upstream.downstreams or []), + "transformOperation": fine_upstream.transformOperation, + } + ) diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py index fee181864a2d6..b4a04d96b532b 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py @@ -461,7 +461,8 @@ def _init_schema_resolver(self) -> SchemaResolver: ) else: logger.warning( - "Failed to load schema info from DataHub as DataHubGraph is missing.", + "Failed to load schema info from DataHub as DataHubGraph is missing. " + "Use `datahub-rest` sink OR provide `datahub-api` config in recipe. ", ) return SchemaResolver(platform=self.platform, env=self.config.env) diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py index 215116b4c33fb..e0848b5f9ab34 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py +++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py @@ -4,6 +4,7 @@ import os.path import platform from dataclasses import dataclass +from functools import partial from typing import Callable, Dict, Iterable, List, Optional, Union import pandas as pd @@ -35,6 +36,7 @@ TestableSource, TestConnectionReport, ) +from datahub.ingestion.api.source_helpers import auto_incremental_lineage from datahub.ingestion.api.workunit import MetadataWorkUnit from datahub.ingestion.glossary.classification_mixin import ClassificationHandler from datahub.ingestion.source.common.subtypes import ( @@ -511,6 +513,13 @@ def _init_schema_resolver(self) -> SchemaResolver: def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]: return [ *super().get_workunit_processors(), + partial( + auto_incremental_lineage, + self.ctx.graph, + self.config.incremental_lineage, + self.config.include_column_lineage + or self.config.include_view_column_lineage, + ), StaleEntityRemovalHandler.create( self, self.config, self.ctx ).workunit_processor, diff --git a/metadata-ingestion/tests/integration/snowflake/snowflake_privatelink_golden.json b/metadata-ingestion/tests/integration/snowflake/snowflake_privatelink_golden.json index 7687b99ac8d6d..5057dacd5b0c8 100644 --- a/metadata-ingestion/tests/integration/snowflake/snowflake_privatelink_golden.json +++ b/metadata-ingestion/tests/integration/snowflake/snowflake_privatelink_golden.json @@ -24,7 +24,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -39,7 +40,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -54,7 +56,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -71,7 +74,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -86,7 +90,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -115,7 +120,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -130,7 +136,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -145,7 +152,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -162,7 +170,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -177,7 +186,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -197,7 +207,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -212,7 +223,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -375,7 +387,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -401,7 +414,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -416,7 +430,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -433,7 +448,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -457,7 +473,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -472,7 +489,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -635,7 +653,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -661,7 +680,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -676,7 +696,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -693,7 +714,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -717,7 +739,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -732,7 +755,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -895,7 +919,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -921,7 +946,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -936,7 +962,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -953,7 +980,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -977,7 +1005,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -992,7 +1021,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1155,7 +1185,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1181,7 +1212,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1196,7 +1228,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1213,7 +1246,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1237,7 +1271,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1252,7 +1287,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1415,7 +1451,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1441,7 +1478,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1456,7 +1494,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1473,7 +1512,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1497,7 +1537,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1512,7 +1553,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1675,7 +1717,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1701,7 +1744,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1716,7 +1760,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1733,7 +1778,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1757,7 +1803,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1772,7 +1819,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1935,7 +1983,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1961,7 +2010,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1976,7 +2026,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -1993,7 +2044,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2017,7 +2069,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2032,7 +2085,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2195,7 +2249,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2221,7 +2276,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2236,7 +2292,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2253,7 +2310,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2277,7 +2335,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2292,7 +2351,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2455,7 +2515,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2481,7 +2542,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2496,7 +2558,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2513,7 +2576,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2537,7 +2601,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2552,7 +2617,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2715,7 +2781,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2741,7 +2808,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2756,7 +2824,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2773,7 +2842,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2797,7 +2867,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2821,7 +2892,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2845,7 +2917,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2869,7 +2942,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2893,7 +2967,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2917,7 +2992,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2941,7 +3017,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2965,7 +3042,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -2989,7 +3067,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -3013,7 +3092,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -3037,7 +3117,8 @@ }, "systemMetadata": { "lastObserved": 1654621200000, - "runId": "snowflake-2022_06_07-17_00_00" + "runId": "snowflake-2022_06_07-17_00_00", + "lastRunId": "no-run-id-provided" } } ] \ No newline at end of file diff --git a/metadata-ingestion/tests/integration/snowflake/test_snowflake.py b/metadata-ingestion/tests/integration/snowflake/test_snowflake.py index 2c77ace8b53e5..3dafe85ef950a 100644 --- a/metadata-ingestion/tests/integration/snowflake/test_snowflake.py +++ b/metadata-ingestion/tests/integration/snowflake/test_snowflake.py @@ -125,6 +125,7 @@ def test_snowflake_basic(pytestconfig, tmp_path, mock_time, mock_datahub_graph): validate_upstreams_against_patterns=False, include_operational_stats=True, email_as_user_identifier=True, + incremental_lineage=False, start_time=datetime(2022, 6, 6, 0, 0, 0, 0).replace( tzinfo=timezone.utc ), @@ -213,6 +214,7 @@ def test_snowflake_private_link(pytestconfig, tmp_path, mock_time, mock_datahub_ include_views=False, include_view_lineage=False, include_usage_stats=False, + incremental_lineage=False, include_operational_stats=False, start_time=datetime(2022, 6, 6, 0, 0, 0, 0).replace( tzinfo=timezone.utc diff --git a/metadata-ingestion/tests/integration/snowflake/test_snowflake_failures.py b/metadata-ingestion/tests/integration/snowflake/test_snowflake_failures.py index bba53c1e97a47..cd53b8f7db4f6 100644 --- a/metadata-ingestion/tests/integration/snowflake/test_snowflake_failures.py +++ b/metadata-ingestion/tests/integration/snowflake/test_snowflake_failures.py @@ -283,10 +283,12 @@ def test_snowflake_unexpected_snowflake_view_lineage_error_causes_pipeline_warni ) snowflake_pipeline_config1 = snowflake_pipeline_config.copy() - cast( + config = cast( SnowflakeV2Config, cast(PipelineConfig, snowflake_pipeline_config1).source.config, - ).include_view_lineage = True + ) + config.include_view_lineage = True + config.incremental_lineage = False pipeline = Pipeline(snowflake_pipeline_config1) pipeline.run() pipeline.raise_from_status() # pipeline should not fail diff --git a/metadata-ingestion/tests/integration/snowflake/test_snowflake_stateful.py b/metadata-ingestion/tests/integration/snowflake/test_snowflake_stateful.py index f72bd5b72d2cd..7e2ac94fa4e35 100644 --- a/metadata-ingestion/tests/integration/snowflake/test_snowflake_stateful.py +++ b/metadata-ingestion/tests/integration/snowflake/test_snowflake_stateful.py @@ -31,6 +31,7 @@ def stateful_pipeline_config(include_tables: bool) -> PipelineConfig: match_fully_qualified_names=True, schema_pattern=AllowDenyPattern(allow=["test_db.test_schema"]), include_tables=include_tables, + incremental_lineage=False, stateful_ingestion=StatefulStaleMetadataRemovalConfig.parse_obj( { "enabled": True, @@ -49,7 +50,7 @@ def stateful_pipeline_config(include_tables: bool) -> PipelineConfig: @freeze_time(FROZEN_TIME) -def test_tableau_stateful(mock_datahub_graph): +def test_stale_metadata_removal(mock_datahub_graph): with mock.patch( "datahub.ingestion.source.state_provider.datahub_ingestion_checkpointing_provider.DataHubGraph", mock_datahub_graph, diff --git a/metadata-ingestion/tests/unit/api/source_helpers/incremental_cll_less_upstreams_in_gms_aspect_golden.json b/metadata-ingestion/tests/unit/api/source_helpers/incremental_cll_less_upstreams_in_gms_aspect_golden.json new file mode 100644 index 0000000000000..812566143014b --- /dev/null +++ b/metadata-ingestion/tests/unit/api/source_helpers/incremental_cll_less_upstreams_in_gms_aspect_golden.json @@ -0,0 +1,106 @@ +[ +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD)", + "changeType": "UPSERT", + "aspectName": "upstreamLineage", + "aspect": { + "json": { + "upstreams": [ + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD)", + "type": "TRANSFORMED" + }, + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD)", + "type": "TRANSFORMED" + } + ], + "fineGrainedLineages": [ + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_a)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_a)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_b)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_b)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_c)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_c)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_a)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD),col_a)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_a)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_b)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD),col_b)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_b)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_c)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD),col_c)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_c)" + ], + "confidenceScore": 1.0 + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "run-id", + "lastRunId": "no-run-id-provided" + } +} +] \ No newline at end of file diff --git a/metadata-ingestion/tests/unit/api/source_helpers/incremental_cll_more_upstreams_in_gms_aspect_golden.json b/metadata-ingestion/tests/unit/api/source_helpers/incremental_cll_more_upstreams_in_gms_aspect_golden.json new file mode 100644 index 0000000000000..17f4d10728268 --- /dev/null +++ b/metadata-ingestion/tests/unit/api/source_helpers/incremental_cll_more_upstreams_in_gms_aspect_golden.json @@ -0,0 +1,120 @@ +[ +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD)", + "changeType": "UPSERT", + "aspectName": "upstreamLineage", + "aspect": { + "json": { + "upstreams": [ + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD)", + "type": "TRANSFORMED" + }, + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD)", + "type": "TRANSFORMED" + }, + { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:platform,upstream3,PROD)", + "type": "TRANSFORMED" + } + ], + "fineGrainedLineages": [ + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_a)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD),col_a)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream3,PROD),col_a)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_a)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_b)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD),col_b)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream3,PROD),col_b)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_b)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_c)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD),col_c)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream3,PROD),col_c)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_c)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_a)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD),col_a)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_a)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_b)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD),col_b)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_b)" + ], + "confidenceScore": 1.0 + }, + { + "upstreamType": "FIELD_SET", + "upstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD),col_c)", + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD),col_c)" + ], + "downstreamType": "FIELD", + "downstreams": [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD),col_c)" + ], + "confidenceScore": 1.0 + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "run-id", + "lastRunId": "no-run-id-provided" + } +} +] \ No newline at end of file diff --git a/metadata-ingestion/tests/unit/api/source_helpers/incremental_table_lineage_golden.json b/metadata-ingestion/tests/unit/api/source_helpers/incremental_table_lineage_golden.json new file mode 100644 index 0000000000000..c828373c73080 --- /dev/null +++ b/metadata-ingestion/tests/unit/api/source_helpers/incremental_table_lineage_golden.json @@ -0,0 +1,41 @@ +[ +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:platform,dataset1,PROD)", + "changeType": "PATCH", + "aspectName": "upstreamLineage", + "aspect": { + "json": [ + { + "op": "add", + "path": "/upstreams/urn%3Ali%3Adataset%3A%28urn%3Ali%3AdataPlatform%3Aplatform%2Cupstream1%2CPROD%29", + "value": { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:platform,upstream1,PROD)", + "type": "TRANSFORMED" + } + }, + { + "op": "add", + "path": "/upstreams/urn%3Ali%3Adataset%3A%28urn%3Ali%3AdataPlatform%3Aplatform%2Cupstream2%2CPROD%29", + "value": { + "auditStamp": { + "time": 0, + "actor": "urn:li:corpuser:unknown" + }, + "dataset": "urn:li:dataset:(urn:li:dataPlatform:platform,upstream2,PROD)", + "type": "TRANSFORMED" + } + } + ] + }, + "systemMetadata": { + "lastObserved": 1643871600000, + "runId": "run-id", + "lastRunId": "no-run-id-provided" + } +} +] \ No newline at end of file diff --git a/metadata-ingestion/tests/unit/api/source_helpers/test_incremental_lineage_helper.py b/metadata-ingestion/tests/unit/api/source_helpers/test_incremental_lineage_helper.py new file mode 100644 index 0000000000000..4078bda26c743 --- /dev/null +++ b/metadata-ingestion/tests/unit/api/source_helpers/test_incremental_lineage_helper.py @@ -0,0 +1,244 @@ +from typing import List, Optional +from unittest.mock import MagicMock + +import pytest + +import datahub.metadata.schema_classes as models +from datahub.emitter.mce_builder import make_dataset_urn, make_schema_field_urn +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.ingestion.api.source_helpers import auto_incremental_lineage +from datahub.ingestion.api.workunit import MetadataWorkUnit +from datahub.ingestion.sink.file import write_metadata_file +from tests.test_helpers import mce_helpers + +platform = "platform" +system_metadata = models.SystemMetadataClass(lastObserved=1643871600000, runId="run-id") + + +def make_lineage_aspect( + dataset_name: str, + upstreams: List[str], + timestamp: int = 0, + columns: List[str] = [], + include_cll: bool = False, +) -> models.UpstreamLineageClass: + """ + Generates dataset properties and upstream lineage aspects + with simple column to column lineage between current dataset and all upstreams + """ + + dataset_urn = make_dataset_urn(platform, dataset_name) + return models.UpstreamLineageClass( + upstreams=[ + models.UpstreamClass( + dataset=upstream_urn, + type=models.DatasetLineageTypeClass.TRANSFORMED, + auditStamp=models.AuditStampClass( + time=timestamp, actor="urn:li:corpuser:unknown" + ), + ) + for upstream_urn in upstreams + ], + fineGrainedLineages=[ + models.FineGrainedLineageClass( + upstreamType=models.FineGrainedLineageUpstreamTypeClass.FIELD_SET, + downstreamType=models.FineGrainedLineageDownstreamTypeClass.FIELD, + upstreams=[ + make_schema_field_urn(upstream_urn, col) + for upstream_urn in upstreams + ], + downstreams=[make_schema_field_urn(dataset_urn, col)], + ) + for col in columns + ] + if include_cll + else None, + ) + + +def base_table_lineage_aspect() -> models.UpstreamLineageClass: + return make_lineage_aspect( + "dataset1", + upstreams=[ + make_dataset_urn(platform, name) for name in ["upstream1", "upstream2"] + ], + ) + + +def base_cll_aspect(timestamp: int = 0) -> models.UpstreamLineageClass: + return make_lineage_aspect( + "dataset1", + upstreams=[ + make_dataset_urn(platform, name) for name in ["upstream1", "upstream2"] + ], + timestamp=timestamp, + columns=["col_a", "col_b", "col_c"], + include_cll=True, + ) + + +def test_incremental_table_lineage(tmp_path, pytestconfig): + test_resources_dir = pytestconfig.rootpath / "tests/unit/api/source_helpers" + test_file = tmp_path / "incremental_table_lineage.json" + golden_file = test_resources_dir / "incremental_table_lineage_golden.json" + + urn = make_dataset_urn(platform, "dataset1") + aspect = base_table_lineage_aspect() + + processed_wus = auto_incremental_lineage( + graph=None, + incremental_lineage=True, + include_column_level_lineage=False, + stream=[ + MetadataChangeProposalWrapper( + entityUrn=urn, aspect=aspect, systemMetadata=system_metadata + ).as_workunit() + ], + ) + + write_metadata_file( + test_file, + [wu.metadata for wu in processed_wus], + ) + mce_helpers.check_golden_file( + pytestconfig=pytestconfig, output_path=test_file, golden_path=golden_file + ) + + +@pytest.mark.parametrize( + "gms_aspect,current_aspect,output_aspect", + [ + # emitting CLL upstreamLineage over table level upstreamLineage + [ + base_table_lineage_aspect(), + base_cll_aspect(), + base_cll_aspect(), + ], + # emitting upstreamLineage for the first time + [ + None, + base_cll_aspect(), + base_cll_aspect(), + ], + # emitting CLL upstreamLineage over same CLL upstreamLineage + [ + base_cll_aspect(), + base_cll_aspect(), + base_cll_aspect(), + ], + # emitting CLL upstreamLineage over same CLL upstreamLineage but with earlier timestamp + [ + base_cll_aspect(), # default timestamp is 0 + base_cll_aspect(timestamp=1643871600000), + base_cll_aspect(timestamp=1643871600000), + ], + ], +) +def test_incremental_column_level_lineage( + gms_aspect: Optional[models.UpstreamLineageClass], + current_aspect: models.UpstreamLineageClass, + output_aspect: models.UpstreamLineageClass, +) -> None: + mock_graph = MagicMock() + mock_graph.get_aspect.return_value = gms_aspect + dataset_urn = make_dataset_urn(platform, "dataset1") + + processed_wus = auto_incremental_lineage( + graph=mock_graph, + incremental_lineage=True, + include_column_level_lineage=True, + stream=[ + MetadataChangeProposalWrapper( + entityUrn=dataset_urn, + aspect=current_aspect, + systemMetadata=system_metadata, + ).as_workunit() + ], + ) + + wu: MetadataWorkUnit = next(iter(processed_wus)) + aspect = wu.get_aspect_of_type(models.UpstreamLineageClass) + assert aspect == output_aspect + + +def test_incremental_column_lineage_less_upstreams_in_gms_aspect( + tmp_path, pytestconfig +): + test_resources_dir = pytestconfig.rootpath / "tests/unit/api/source_helpers" + test_file = tmp_path / "incremental_cll_less_upstreams_in_gms_aspect.json" + golden_file = ( + test_resources_dir / "incremental_cll_less_upstreams_in_gms_aspect_golden.json" + ) + + urn = make_dataset_urn(platform, "dataset1") + aspect = base_cll_aspect() + + mock_graph = MagicMock() + mock_graph.get_aspect.return_value = make_lineage_aspect( + "dataset1", + upstreams=[make_dataset_urn(platform, name) for name in ["upstream1"]], + columns=["col_a", "col_b", "col_c"], + include_cll=True, + ) + + processed_wus = auto_incremental_lineage( + graph=mock_graph, + incremental_lineage=True, + include_column_level_lineage=True, + stream=[ + MetadataChangeProposalWrapper( + entityUrn=urn, aspect=aspect, systemMetadata=system_metadata + ).as_workunit() + ], + ) + + write_metadata_file( + test_file, + [wu.metadata for wu in processed_wus], + ) + mce_helpers.check_golden_file( + pytestconfig=pytestconfig, output_path=test_file, golden_path=golden_file + ) + + +def test_incremental_column_lineage_more_upstreams_in_gms_aspect( + tmp_path, pytestconfig +): + test_resources_dir = pytestconfig.rootpath / "tests/unit/api/source_helpers" + test_file = tmp_path / "incremental_cll_more_upstreams_in_gms_aspect.json" + golden_file = ( + test_resources_dir / "incremental_cll_more_upstreams_in_gms_aspect_golden.json" + ) + + urn = make_dataset_urn(platform, "dataset1") + aspect = base_cll_aspect() + + mock_graph = MagicMock() + mock_graph.get_aspect.return_value = make_lineage_aspect( + "dataset1", + upstreams=[ + make_dataset_urn(platform, name) + for name in ["upstream1", "upstream2", "upstream3"] + ], + columns=["col_a", "col_b", "col_c"], + include_cll=True, + ) + + processed_wus = auto_incremental_lineage( + graph=mock_graph, + incremental_lineage=True, + include_column_level_lineage=True, + stream=[ + MetadataChangeProposalWrapper( + entityUrn=urn, aspect=aspect, systemMetadata=system_metadata + ).as_workunit() + ], + ) + + write_metadata_file( + test_file, + [wu.metadata for wu in processed_wus], + ) + mce_helpers.check_golden_file( + pytestconfig=pytestconfig, output_path=test_file, golden_path=golden_file + ) diff --git a/metadata-ingestion/tests/unit/test_source_helpers.py b/metadata-ingestion/tests/unit/api/source_helpers/test_source_helpers.py similarity index 100% rename from metadata-ingestion/tests/unit/test_source_helpers.py rename to metadata-ingestion/tests/unit/api/source_helpers/test_source_helpers.py From 57f855ecd11632e884b12fda0fc57e2694ee26a5 Mon Sep 17 00:00:00 2001 From: Mayuri Nehate <33225191+mayurinehate@users.noreply.github.com> Date: Tue, 10 Oct 2023 12:18:21 +0530 Subject: [PATCH 106/156] feat(ingest): refactor + simplify incremental lineage helper (#8976) --- .../api/incremental_lineage_helper.py | 139 ++++++++++++++++++ .../datahub/ingestion/api/source_helpers.py | 138 +---------------- .../source/snowflake/snowflake_v2.py | 4 +- .../test_incremental_lineage_helper.py | 6 +- 4 files changed, 142 insertions(+), 145 deletions(-) create mode 100644 metadata-ingestion/src/datahub/ingestion/api/incremental_lineage_helper.py diff --git a/metadata-ingestion/src/datahub/ingestion/api/incremental_lineage_helper.py b/metadata-ingestion/src/datahub/ingestion/api/incremental_lineage_helper.py new file mode 100644 index 0000000000000..9478c5cf7efa2 --- /dev/null +++ b/metadata-ingestion/src/datahub/ingestion/api/incremental_lineage_helper.py @@ -0,0 +1,139 @@ +import copy +from typing import Dict, Iterable, Optional + +from datahub.emitter.mce_builder import datahub_guid, set_aspect +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.ingestion.api.workunit import MetadataWorkUnit +from datahub.ingestion.graph.client import DataHubGraph +from datahub.metadata.schema_classes import ( + FineGrainedLineageClass, + MetadataChangeEventClass, + SystemMetadataClass, + UpstreamClass, + UpstreamLineageClass, +) +from datahub.specific.dataset import DatasetPatchBuilder + + +def _convert_upstream_lineage_to_patch( + urn: str, + aspect: UpstreamLineageClass, + system_metadata: Optional[SystemMetadataClass], +) -> MetadataWorkUnit: + patch_builder = DatasetPatchBuilder(urn, system_metadata) + for upstream in aspect.upstreams: + patch_builder.add_upstream_lineage(upstream) + mcp = next(iter(patch_builder.build())) + return MetadataWorkUnit(id=f"{urn}-upstreamLineage", mcp_raw=mcp) + + +def get_fine_grained_lineage_key(fine_upstream: FineGrainedLineageClass) -> str: + return datahub_guid( + { + "upstreams": sorted(fine_upstream.upstreams or []), + "downstreams": sorted(fine_upstream.downstreams or []), + "transformOperation": fine_upstream.transformOperation, + } + ) + + +def _merge_upstream_lineage( + new_aspect: UpstreamLineageClass, gms_aspect: UpstreamLineageClass +) -> UpstreamLineageClass: + merged_aspect = copy.deepcopy(gms_aspect) + + upstreams_map: Dict[str, UpstreamClass] = { + upstream.dataset: upstream for upstream in merged_aspect.upstreams + } + + upstreams_updated = False + fine_upstreams_updated = False + + for table_upstream in new_aspect.upstreams: + if table_upstream.dataset not in upstreams_map or ( + table_upstream.auditStamp.time + > upstreams_map[table_upstream.dataset].auditStamp.time + ): + upstreams_map[table_upstream.dataset] = table_upstream + upstreams_updated = True + + if upstreams_updated: + merged_aspect.upstreams = list(upstreams_map.values()) + + if new_aspect.fineGrainedLineages and merged_aspect.fineGrainedLineages: + fine_upstreams_map: Dict[str, FineGrainedLineageClass] = { + get_fine_grained_lineage_key(fine_upstream): fine_upstream + for fine_upstream in merged_aspect.fineGrainedLineages + } + for column_upstream in new_aspect.fineGrainedLineages: + column_upstream_key = get_fine_grained_lineage_key(column_upstream) + + if column_upstream_key not in fine_upstreams_map or ( + column_upstream.confidenceScore + > fine_upstreams_map[column_upstream_key].confidenceScore + ): + fine_upstreams_map[column_upstream_key] = column_upstream + fine_upstreams_updated = True + + if fine_upstreams_updated: + merged_aspect.fineGrainedLineages = list(fine_upstreams_map.values()) + else: + merged_aspect.fineGrainedLineages = ( + new_aspect.fineGrainedLineages or gms_aspect.fineGrainedLineages + ) + + return merged_aspect + + +def _lineage_wu_via_read_modify_write( + graph: Optional[DataHubGraph], + urn: str, + aspect: UpstreamLineageClass, + system_metadata: Optional[SystemMetadataClass], +) -> MetadataWorkUnit: + if graph is None: + raise ValueError( + "Failed to handle incremental lineage, DataHubGraph is missing. " + "Use `datahub-rest` sink OR provide `datahub-api` config in recipe. " + ) + gms_aspect = graph.get_aspect(urn, UpstreamLineageClass) + if gms_aspect: + new_aspect = _merge_upstream_lineage(aspect, gms_aspect) + else: + new_aspect = aspect + + return MetadataChangeProposalWrapper( + entityUrn=urn, aspect=new_aspect, systemMetadata=system_metadata + ).as_workunit() + + +def auto_incremental_lineage( + graph: Optional[DataHubGraph], + incremental_lineage: bool, + stream: Iterable[MetadataWorkUnit], +) -> Iterable[MetadataWorkUnit]: + if not incremental_lineage: + yield from stream + return # early exit + + for wu in stream: + lineage_aspect: Optional[UpstreamLineageClass] = wu.get_aspect_of_type( + UpstreamLineageClass + ) + urn = wu.get_urn() + + if lineage_aspect: + if isinstance(wu.metadata, MetadataChangeEventClass): + set_aspect( + wu.metadata, None, UpstreamLineageClass + ) # we'll emit upstreamLineage separately below + if len(wu.metadata.proposedSnapshot.aspects) > 0: + yield wu + + yield _lineage_wu_via_read_modify_write( + graph, urn, lineage_aspect, wu.metadata.systemMetadata + ) if lineage_aspect.fineGrainedLineages else _convert_upstream_lineage_to_patch( + urn, lineage_aspect, wu.metadata.systemMetadata + ) + else: + yield wu diff --git a/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py b/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py index 42f970e97c95f..7fc15cf829678 100644 --- a/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py +++ b/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py @@ -1,4 +1,3 @@ -import copy import logging from datetime import datetime, timezone from typing import ( @@ -16,14 +15,9 @@ ) from datahub.configuration.time_window_config import BaseTimeWindowConfig -from datahub.emitter.mce_builder import ( - datahub_guid, - make_dataplatform_instance_urn, - set_aspect, -) +from datahub.emitter.mce_builder import make_dataplatform_instance_urn from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.workunit import MetadataWorkUnit -from datahub.ingestion.graph.client import DataHubGraph from datahub.metadata.schema_classes import ( BrowsePathEntryClass, BrowsePathsClass, @@ -31,17 +25,12 @@ ChangeTypeClass, ContainerClass, DatasetUsageStatisticsClass, - FineGrainedLineageClass, MetadataChangeEventClass, MetadataChangeProposalClass, StatusClass, - SystemMetadataClass, TagKeyClass, TimeWindowSizeClass, - UpstreamClass, - UpstreamLineageClass, ) -from datahub.specific.dataset import DatasetPatchBuilder from datahub.telemetry import telemetry from datahub.utilities.urns.dataset_urn import DatasetUrn from datahub.utilities.urns.tag_urn import TagUrn @@ -377,128 +366,3 @@ def _prepend_platform_instance( return [BrowsePathEntryClass(id=urn, urn=urn)] + entries return entries - - -def auto_incremental_lineage( - graph: Optional[DataHubGraph], - incremental_lineage: bool, - include_column_level_lineage: bool, - stream: Iterable[MetadataWorkUnit], -) -> Iterable[MetadataWorkUnit]: - if not incremental_lineage: - yield from stream - return # early exit - - for wu in stream: - lineage_aspect: Optional[UpstreamLineageClass] = wu.get_aspect_of_type( - UpstreamLineageClass - ) - urn = wu.get_urn() - - if lineage_aspect: - if isinstance(wu.metadata, MetadataChangeEventClass): - set_aspect( - wu.metadata, None, UpstreamLineageClass - ) # we'll emit upstreamLineage separately below - if len(wu.metadata.proposedSnapshot.aspects) > 0: - yield wu - - yield _lineage_wu_via_read_modify_write( - graph, urn, lineage_aspect, wu.metadata.systemMetadata - ) if include_column_level_lineage else _convert_upstream_lineage_to_patch( - urn, lineage_aspect, wu.metadata.systemMetadata - ) - else: - yield wu - - -def _convert_upstream_lineage_to_patch( - urn: str, - aspect: UpstreamLineageClass, - system_metadata: Optional[SystemMetadataClass], -) -> MetadataWorkUnit: - patch_builder = DatasetPatchBuilder(urn, system_metadata) - for upstream in aspect.upstreams: - patch_builder.add_upstream_lineage(upstream) - mcp = next(iter(patch_builder.build())) - return MetadataWorkUnit(id=f"{urn}-upstreamLineage", mcp_raw=mcp) - - -def _lineage_wu_via_read_modify_write( - graph: Optional[DataHubGraph], - urn: str, - aspect: UpstreamLineageClass, - system_metadata: Optional[SystemMetadataClass], -) -> MetadataWorkUnit: - if graph is None: - raise ValueError( - "Failed to handle incremental lineage, DataHubGraph is missing. " - "Use `datahub-rest` sink OR provide `datahub-api` config in recipe. " - ) - gms_aspect = graph.get_aspect(urn, UpstreamLineageClass) - if gms_aspect: - new_aspect = _merge_upstream_lineage(aspect, gms_aspect) - else: - new_aspect = aspect - - return MetadataChangeProposalWrapper( - entityUrn=urn, aspect=new_aspect, systemMetadata=system_metadata - ).as_workunit() - - -def _merge_upstream_lineage( - new_aspect: UpstreamLineageClass, gms_aspect: UpstreamLineageClass -) -> UpstreamLineageClass: - merged_aspect = copy.deepcopy(gms_aspect) - - upstreams_map: Dict[str, UpstreamClass] = { - upstream.dataset: upstream for upstream in merged_aspect.upstreams - } - - upstreams_updated = False - fine_upstreams_updated = False - - for table_upstream in new_aspect.upstreams: - if table_upstream.dataset not in upstreams_map or ( - table_upstream.auditStamp.time - > upstreams_map[table_upstream.dataset].auditStamp.time - ): - upstreams_map[table_upstream.dataset] = table_upstream - upstreams_updated = True - - if upstreams_updated: - merged_aspect.upstreams = list(upstreams_map.values()) - - if new_aspect.fineGrainedLineages and merged_aspect.fineGrainedLineages: - fine_upstreams_map: Dict[str, FineGrainedLineageClass] = { - get_fine_grained_lineage_key(fine_upstream): fine_upstream - for fine_upstream in merged_aspect.fineGrainedLineages - } - for column_upstream in new_aspect.fineGrainedLineages: - column_upstream_key = get_fine_grained_lineage_key(column_upstream) - - if column_upstream_key not in fine_upstreams_map or ( - column_upstream.confidenceScore - > fine_upstreams_map[column_upstream_key].confidenceScore - ): - fine_upstreams_map[column_upstream_key] = column_upstream - fine_upstreams_updated = True - - if fine_upstreams_updated: - merged_aspect.fineGrainedLineages = list(fine_upstreams_map.values()) - else: - merged_aspect.fineGrainedLineages = ( - new_aspect.fineGrainedLineages or gms_aspect.fineGrainedLineages - ) - - return merged_aspect - - -def get_fine_grained_lineage_key(fine_upstream: FineGrainedLineageClass) -> str: - return datahub_guid( - { - "upstreams": sorted(fine_upstream.upstreams or []), - "downstreams": sorted(fine_upstream.downstreams or []), - "transformOperation": fine_upstream.transformOperation, - } - ) diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py index e0848b5f9ab34..a5c07d9a3870c 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py +++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py @@ -27,6 +27,7 @@ platform_name, support_status, ) +from datahub.ingestion.api.incremental_lineage_helper import auto_incremental_lineage from datahub.ingestion.api.source import ( CapabilityReport, MetadataWorkUnitProcessor, @@ -36,7 +37,6 @@ TestableSource, TestConnectionReport, ) -from datahub.ingestion.api.source_helpers import auto_incremental_lineage from datahub.ingestion.api.workunit import MetadataWorkUnit from datahub.ingestion.glossary.classification_mixin import ClassificationHandler from datahub.ingestion.source.common.subtypes import ( @@ -517,8 +517,6 @@ def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]: auto_incremental_lineage, self.ctx.graph, self.config.incremental_lineage, - self.config.include_column_lineage - or self.config.include_view_column_lineage, ), StaleEntityRemovalHandler.create( self, self.config, self.ctx diff --git a/metadata-ingestion/tests/unit/api/source_helpers/test_incremental_lineage_helper.py b/metadata-ingestion/tests/unit/api/source_helpers/test_incremental_lineage_helper.py index 4078bda26c743..54a22d860285c 100644 --- a/metadata-ingestion/tests/unit/api/source_helpers/test_incremental_lineage_helper.py +++ b/metadata-ingestion/tests/unit/api/source_helpers/test_incremental_lineage_helper.py @@ -6,7 +6,7 @@ import datahub.metadata.schema_classes as models from datahub.emitter.mce_builder import make_dataset_urn, make_schema_field_urn from datahub.emitter.mcp import MetadataChangeProposalWrapper -from datahub.ingestion.api.source_helpers import auto_incremental_lineage +from datahub.ingestion.api.incremental_lineage_helper import auto_incremental_lineage from datahub.ingestion.api.workunit import MetadataWorkUnit from datahub.ingestion.sink.file import write_metadata_file from tests.test_helpers import mce_helpers @@ -88,7 +88,6 @@ def test_incremental_table_lineage(tmp_path, pytestconfig): processed_wus = auto_incremental_lineage( graph=None, incremental_lineage=True, - include_column_level_lineage=False, stream=[ MetadataChangeProposalWrapper( entityUrn=urn, aspect=aspect, systemMetadata=system_metadata @@ -146,7 +145,6 @@ def test_incremental_column_level_lineage( processed_wus = auto_incremental_lineage( graph=mock_graph, incremental_lineage=True, - include_column_level_lineage=True, stream=[ MetadataChangeProposalWrapper( entityUrn=dataset_urn, @@ -184,7 +182,6 @@ def test_incremental_column_lineage_less_upstreams_in_gms_aspect( processed_wus = auto_incremental_lineage( graph=mock_graph, incremental_lineage=True, - include_column_level_lineage=True, stream=[ MetadataChangeProposalWrapper( entityUrn=urn, aspect=aspect, systemMetadata=system_metadata @@ -227,7 +224,6 @@ def test_incremental_column_lineage_more_upstreams_in_gms_aspect( processed_wus = auto_incremental_lineage( graph=mock_graph, incremental_lineage=True, - include_column_level_lineage=True, stream=[ MetadataChangeProposalWrapper( entityUrn=urn, aspect=aspect, systemMetadata=system_metadata From bb39d5418fcbf8bebbae1b510c63a1170865a072 Mon Sep 17 00:00:00 2001 From: Aseem Bansal Date: Tue, 10 Oct 2023 16:08:34 +0530 Subject: [PATCH 107/156] fix(lint): run black, isort (#8978) --- .../tests/assertions/assertions_test.py | 33 ++-- smoke-test/tests/browse/browse_test.py | 51 +++++-- smoke-test/tests/cli/datahub-cli.py | 76 +++++++--- smoke-test/tests/cli/datahub_graph_test.py | 12 +- .../cli/delete_cmd/test_timeseries_delete.py | 12 +- .../ingest_cmd/test_timeseries_rollback.py | 6 +- .../cli/user_groups_cmd/test_group_cmd.py | 3 +- smoke-test/tests/conftest.py | 4 +- smoke-test/tests/consistency_utils.py | 16 +- .../tests/containers/containers_test.py | 4 +- smoke-test/tests/cypress/integration_test.py | 23 ++- .../tests/dataproduct/test_dataproduct.py | 4 +- smoke-test/tests/delete/delete_test.py | 18 +-- .../tests/deprecation/deprecation_test.py | 9 +- smoke-test/tests/domains/domains_test.py | 15 +- .../managed_ingestion_test.py | 3 +- smoke-test/tests/patch/common_patch_tests.py | 52 ++----- .../tests/patch/test_datajob_patches.py | 23 +-- .../tests/patch/test_dataset_patches.py | 18 ++- smoke-test/tests/policies/test_policies.py | 10 +- .../tests/setup/lineage/helper_classes.py | 5 +- .../setup/lineage/ingest_data_job_change.py | 42 ++---- .../lineage/ingest_dataset_join_change.py | 36 ++--- .../lineage/ingest_input_datasets_change.py | 42 ++---- .../setup/lineage/ingest_time_lineage.py | 18 ++- smoke-test/tests/setup/lineage/utils.py | 85 +++++------ .../tags-and-terms/tags_and_terms_test.py | 4 +- smoke-test/tests/telemetry/telemetry_test.py | 4 +- smoke-test/tests/test_result_msg.py | 23 ++- smoke-test/tests/test_stateful_ingestion.py | 14 +- smoke-test/tests/tests/tests_test.py | 7 +- smoke-test/tests/timeline/timeline_test.py | 67 +++++---- .../tokens/revokable_access_token_test.py | 12 +- smoke-test/tests/utils.py | 17 +-- smoke-test/tests/views/views_test.py | 142 +++++++++--------- 35 files changed, 457 insertions(+), 453 deletions(-) diff --git a/smoke-test/tests/assertions/assertions_test.py b/smoke-test/tests/assertions/assertions_test.py index 4aa64c512f684..48f3564e6cd97 100644 --- a/smoke-test/tests/assertions/assertions_test.py +++ b/smoke-test/tests/assertions/assertions_test.py @@ -2,28 +2,29 @@ import urllib import pytest -import requests_wrapper as requests import tenacity from datahub.emitter.mce_builder import make_dataset_urn, make_schema_field_urn from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.common import PipelineContext, RecordEnvelope from datahub.ingestion.api.sink import NoopWriteCallback from datahub.ingestion.sink.file import FileSink, FileSinkConfig -from datahub.metadata.com.linkedin.pegasus2avro.assertion import AssertionStdAggregation -from datahub.metadata.schema_classes import ( - AssertionInfoClass, - AssertionResultClass, - AssertionResultTypeClass, - AssertionRunEventClass, - AssertionRunStatusClass, - AssertionStdOperatorClass, - AssertionTypeClass, - DatasetAssertionInfoClass, - DatasetAssertionScopeClass, - PartitionSpecClass, - PartitionTypeClass, -) -from tests.utils import delete_urns_from_file, get_gms_url, ingest_file_via_rest, wait_for_healthcheck_util, get_sleep_info +from datahub.metadata.com.linkedin.pegasus2avro.assertion import \ + AssertionStdAggregation +from datahub.metadata.schema_classes import (AssertionInfoClass, + AssertionResultClass, + AssertionResultTypeClass, + AssertionRunEventClass, + AssertionRunStatusClass, + AssertionStdOperatorClass, + AssertionTypeClass, + DatasetAssertionInfoClass, + DatasetAssertionScopeClass, + PartitionSpecClass, + PartitionTypeClass) + +import requests_wrapper as requests +from tests.utils import (delete_urns_from_file, get_gms_url, get_sleep_info, + ingest_file_via_rest, wait_for_healthcheck_util) restli_default_headers = { "X-RestLi-Protocol-Version": "2.0.0", diff --git a/smoke-test/tests/browse/browse_test.py b/smoke-test/tests/browse/browse_test.py index b9d2143d13ec7..550f0062d5a39 100644 --- a/smoke-test/tests/browse/browse_test.py +++ b/smoke-test/tests/browse/browse_test.py @@ -1,9 +1,10 @@ import time import pytest -import requests_wrapper as requests -from tests.utils import delete_urns_from_file, get_frontend_url, ingest_file_via_rest +import requests_wrapper as requests +from tests.utils import (delete_urns_from_file, get_frontend_url, + ingest_file_via_rest) TEST_DATASET_1_URN = "urn:li:dataset:(urn:li:dataPlatform:kafka,test-browse-1,PROD)" TEST_DATASET_2_URN = "urn:li:dataset:(urn:li:dataPlatform:kafka,test-browse-2,PROD)" @@ -51,7 +52,9 @@ def test_get_browse_paths(frontend_session, ingest_cleanup_data): # /prod -- There should be one entity get_browse_paths_json = { "query": get_browse_paths_query, - "variables": {"input": { "type": "DATASET", "path": ["prod"], "start": 0, "count": 100 } }, + "variables": { + "input": {"type": "DATASET", "path": ["prod"], "start": 0, "count": 100} + }, } response = frontend_session.post( @@ -67,12 +70,19 @@ def test_get_browse_paths(frontend_session, ingest_cleanup_data): browse = res_data["data"]["browse"] print(browse) - assert browse["entities"] == [{ "urn": TEST_DATASET_3_URN }] + assert browse["entities"] == [{"urn": TEST_DATASET_3_URN}] # /prod/kafka1 get_browse_paths_json = { "query": get_browse_paths_query, - "variables": {"input": { "type": "DATASET", "path": ["prod", "kafka1"], "start": 0, "count": 10 } }, + "variables": { + "input": { + "type": "DATASET", + "path": ["prod", "kafka1"], + "start": 0, + "count": 10, + } + }, } response = frontend_session.post( @@ -88,16 +98,27 @@ def test_get_browse_paths(frontend_session, ingest_cleanup_data): browse = res_data["data"]["browse"] assert browse == { - "total": 3, - "entities": [{ "urn": TEST_DATASET_1_URN }, { "urn": TEST_DATASET_2_URN }, { "urn": TEST_DATASET_3_URN }], - "groups": [], - "metadata": { "path": ["prod", "kafka1"], "totalNumEntities": 0 } + "total": 3, + "entities": [ + {"urn": TEST_DATASET_1_URN}, + {"urn": TEST_DATASET_2_URN}, + {"urn": TEST_DATASET_3_URN}, + ], + "groups": [], + "metadata": {"path": ["prod", "kafka1"], "totalNumEntities": 0}, } # /prod/kafka2 get_browse_paths_json = { "query": get_browse_paths_query, - "variables": {"input": { "type": "DATASET", "path": ["prod", "kafka2"], "start": 0, "count": 10 } }, + "variables": { + "input": { + "type": "DATASET", + "path": ["prod", "kafka2"], + "start": 0, + "count": 10, + } + }, } response = frontend_session.post( @@ -113,10 +134,8 @@ def test_get_browse_paths(frontend_session, ingest_cleanup_data): browse = res_data["data"]["browse"] assert browse == { - "total": 2, - "entities": [{ "urn": TEST_DATASET_1_URN }, { "urn": TEST_DATASET_2_URN }], - "groups": [], - "metadata": { "path": ["prod", "kafka2"], "totalNumEntities": 0 } + "total": 2, + "entities": [{"urn": TEST_DATASET_1_URN}, {"urn": TEST_DATASET_2_URN}], + "groups": [], + "metadata": {"path": ["prod", "kafka2"], "totalNumEntities": 0}, } - - diff --git a/smoke-test/tests/cli/datahub-cli.py b/smoke-test/tests/cli/datahub-cli.py index 1d0080bdd9d48..c3db6028efceb 100644 --- a/smoke-test/tests/cli/datahub-cli.py +++ b/smoke-test/tests/cli/datahub-cli.py @@ -1,8 +1,11 @@ import json -import pytest from time import sleep -from datahub.cli.cli_utils import guess_entity_type, post_entity, get_aspects_for_entity + +import pytest +from datahub.cli.cli_utils import (get_aspects_for_entity, guess_entity_type, + post_entity) from datahub.cli.ingest_cli import get_session_and_host, rollback + from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync ingested_dataset_run_id = "" @@ -24,24 +27,46 @@ def test_setup(): session, gms_host = get_session_and_host() - assert "browsePaths" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False) - assert "editableDatasetProperties" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False) + assert "browsePaths" not in get_aspects_for_entity( + entity_urn=dataset_urn, aspects=["browsePaths"], typed=False + ) + assert "editableDatasetProperties" not in get_aspects_for_entity( + entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False + ) - ingested_dataset_run_id = ingest_file_via_rest("tests/cli/cli_test_data.json").config.run_id + ingested_dataset_run_id = ingest_file_via_rest( + "tests/cli/cli_test_data.json" + ).config.run_id print("Setup ingestion id: " + ingested_dataset_run_id) - assert "browsePaths" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False) + assert "browsePaths" in get_aspects_for_entity( + entity_urn=dataset_urn, aspects=["browsePaths"], typed=False + ) yield # Clean up rollback_url = f"{gms_host}/runs?action=rollback" - session.post(rollback_url, data=json.dumps({"runId": ingested_editable_run_id, "dryRun": False, "hardDelete": True})) - session.post(rollback_url, data=json.dumps({"runId": ingested_dataset_run_id, "dryRun": False, "hardDelete": True})) + session.post( + rollback_url, + data=json.dumps( + {"runId": ingested_editable_run_id, "dryRun": False, "hardDelete": True} + ), + ) + session.post( + rollback_url, + data=json.dumps( + {"runId": ingested_dataset_run_id, "dryRun": False, "hardDelete": True} + ), + ) - assert "browsePaths" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False) - assert "editableDatasetProperties" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False) + assert "browsePaths" not in get_aspects_for_entity( + entity_urn=dataset_urn, aspects=["browsePaths"], typed=False + ) + assert "editableDatasetProperties" not in get_aspects_for_entity( + entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False + ) @pytest.mark.dependency() @@ -49,9 +74,7 @@ def test_rollback_editable(): global ingested_dataset_run_id global ingested_editable_run_id platform = "urn:li:dataPlatform:kafka" - dataset_name = ( - "test-rollback" - ) + dataset_name = "test-rollback" env = "PROD" dataset_urn = f"urn:li:dataset:({platform},{dataset_name},{env})" @@ -59,23 +82,38 @@ def test_rollback_editable(): print("Ingested dataset id:", ingested_dataset_run_id) # Assert that second data ingestion worked - assert "browsePaths" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False) + assert "browsePaths" in get_aspects_for_entity( + entity_urn=dataset_urn, aspects=["browsePaths"], typed=False + ) # Make editable change - ingested_editable_run_id = ingest_file_via_rest("tests/cli/cli_editable_test_data.json").config.run_id + ingested_editable_run_id = ingest_file_via_rest( + "tests/cli/cli_editable_test_data.json" + ).config.run_id print("ingested editable id:", ingested_editable_run_id) # Assert that second data ingestion worked - assert "editableDatasetProperties" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False) + assert "editableDatasetProperties" in get_aspects_for_entity( + entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False + ) # rollback ingestion 1 rollback_url = f"{gms_host}/runs?action=rollback" - session.post(rollback_url, data=json.dumps({"runId": ingested_dataset_run_id, "dryRun": False, "hardDelete": False})) + session.post( + rollback_url, + data=json.dumps( + {"runId": ingested_dataset_run_id, "dryRun": False, "hardDelete": False} + ), + ) # Allow async MCP processor to handle ingestions & rollbacks wait_for_writes_to_sync() # EditableDatasetProperties should still be part of the entity that was soft deleted. - assert "editableDatasetProperties" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False) + assert "editableDatasetProperties" in get_aspects_for_entity( + entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False + ) # But first ingestion aspects should not be present - assert "browsePaths" not in get_aspects_for_entity(entity_urn=dataset_urn, typed=False) + assert "browsePaths" not in get_aspects_for_entity( + entity_urn=dataset_urn, typed=False + ) diff --git a/smoke-test/tests/cli/datahub_graph_test.py b/smoke-test/tests/cli/datahub_graph_test.py index 16925d26f6983..17c8924fb0998 100644 --- a/smoke-test/tests/cli/datahub_graph_test.py +++ b/smoke-test/tests/cli/datahub_graph_test.py @@ -1,13 +1,11 @@ import pytest import tenacity from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph -from datahub.metadata.schema_classes import KafkaSchemaClass, SchemaMetadataClass -from tests.utils import ( - delete_urns_from_file, - get_gms_url, - get_sleep_info, - ingest_file_via_rest, -) +from datahub.metadata.schema_classes import (KafkaSchemaClass, + SchemaMetadataClass) + +from tests.utils import (delete_urns_from_file, get_gms_url, get_sleep_info, + ingest_file_via_rest) sleep_sec, sleep_times = get_sleep_info() diff --git a/smoke-test/tests/cli/delete_cmd/test_timeseries_delete.py b/smoke-test/tests/cli/delete_cmd/test_timeseries_delete.py index 4288a61b7a0c1..106da7cd8d71e 100644 --- a/smoke-test/tests/cli/delete_cmd/test_timeseries_delete.py +++ b/smoke-test/tests/cli/delete_cmd/test_timeseries_delete.py @@ -1,21 +1,22 @@ import json import logging +import sys import tempfile import time -import sys from json import JSONDecodeError from typing import Any, Dict, List, Optional -from click.testing import CliRunner, Result - import datahub.emitter.mce_builder as builder +from click.testing import CliRunner, Result from datahub.emitter.serialization_helper import pre_json_transform from datahub.entrypoints import datahub from datahub.metadata.schema_classes import DatasetProfileClass + +import requests_wrapper as requests from tests.aspect_generators.timeseries.dataset_profile_gen import \ gen_dataset_profiles -from tests.utils import get_strftime_from_timestamp_millis, wait_for_writes_to_sync -import requests_wrapper as requests +from tests.utils import (get_strftime_from_timestamp_millis, + wait_for_writes_to_sync) logger = logging.getLogger(__name__) @@ -33,6 +34,7 @@ def sync_elastic() -> None: wait_for_writes_to_sync() + def datahub_put_profile(dataset_profile: DatasetProfileClass) -> None: with tempfile.NamedTemporaryFile("w+t", suffix=".json") as aspect_file: aspect_text: str = json.dumps(pre_json_transform(dataset_profile.to_obj())) diff --git a/smoke-test/tests/cli/ingest_cmd/test_timeseries_rollback.py b/smoke-test/tests/cli/ingest_cmd/test_timeseries_rollback.py index 61e7a5a65b494..e962b1a5cafd6 100644 --- a/smoke-test/tests/cli/ingest_cmd/test_timeseries_rollback.py +++ b/smoke-test/tests/cli/ingest_cmd/test_timeseries_rollback.py @@ -2,14 +2,14 @@ import time from typing import Any, Dict, List, Optional -from click.testing import CliRunner, Result - import datahub.emitter.mce_builder as builder +from click.testing import CliRunner, Result from datahub.emitter.serialization_helper import post_json_transform from datahub.entrypoints import datahub from datahub.metadata.schema_classes import DatasetProfileClass -from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync + import requests_wrapper as requests +from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync runner = CliRunner(mix_stderr=False) diff --git a/smoke-test/tests/cli/user_groups_cmd/test_group_cmd.py b/smoke-test/tests/cli/user_groups_cmd/test_group_cmd.py index 405e061c016f9..7b986d3be0444 100644 --- a/smoke-test/tests/cli/user_groups_cmd/test_group_cmd.py +++ b/smoke-test/tests/cli/user_groups_cmd/test_group_cmd.py @@ -1,6 +1,7 @@ import json import sys import tempfile +import time from typing import Any, Dict, Iterable, List import yaml @@ -8,7 +9,7 @@ from datahub.api.entities.corpgroup.corpgroup import CorpGroup from datahub.entrypoints import datahub from datahub.ingestion.graph.client import DataHubGraph, get_default_graph -import time + import requests_wrapper as requests from tests.utils import wait_for_writes_to_sync diff --git a/smoke-test/tests/conftest.py b/smoke-test/tests/conftest.py index eed7a983197ef..57b92a2db1c19 100644 --- a/smoke-test/tests/conftest.py +++ b/smoke-test/tests/conftest.py @@ -2,8 +2,8 @@ import pytest -from tests.utils import wait_for_healthcheck_util, get_frontend_session from tests.test_result_msg import send_message +from tests.utils import get_frontend_session, wait_for_healthcheck_util # Disable telemetry os.environ["DATAHUB_TELEMETRY_ENABLED"] = "false" @@ -28,5 +28,5 @@ def test_healthchecks(wait_for_healthchecks): def pytest_sessionfinish(session, exitstatus): - """ whole test run finishes. """ + """whole test run finishes.""" send_message(exitstatus) diff --git a/smoke-test/tests/consistency_utils.py b/smoke-test/tests/consistency_utils.py index 15993733c592b..607835bf3649c 100644 --- a/smoke-test/tests/consistency_utils.py +++ b/smoke-test/tests/consistency_utils.py @@ -1,10 +1,16 @@ -import time +import logging import os import subprocess +import time _ELASTIC_BUFFER_WRITES_TIME_IN_SEC: int = 1 USE_STATIC_SLEEP: bool = bool(os.getenv("USE_STATIC_SLEEP", False)) -ELASTICSEARCH_REFRESH_INTERVAL_SECONDS: int = int(os.getenv("ELASTICSEARCH_REFRESH_INTERVAL_SECONDS", 5)) +ELASTICSEARCH_REFRESH_INTERVAL_SECONDS: int = int( + os.getenv("ELASTICSEARCH_REFRESH_INTERVAL_SECONDS", 5) +) + +logger = logging.getLogger(__name__) + def wait_for_writes_to_sync(max_timeout_in_sec: int = 120) -> None: if USE_STATIC_SLEEP: @@ -30,7 +36,9 @@ def wait_for_writes_to_sync(max_timeout_in_sec: int = 120) -> None: lag_zero = True if not lag_zero: - logger.warning(f"Exiting early from waiting for elastic to catch up due to a timeout. Current lag is {lag_values}") + logger.warning( + f"Exiting early from waiting for elastic to catch up due to a timeout. Current lag is {lag_values}" + ) else: # we want to sleep for an additional period of time for Elastic writes buffer to clear - time.sleep(_ELASTIC_BUFFER_WRITES_TIME_IN_SEC) \ No newline at end of file + time.sleep(_ELASTIC_BUFFER_WRITES_TIME_IN_SEC) diff --git a/smoke-test/tests/containers/containers_test.py b/smoke-test/tests/containers/containers_test.py index 575e3def6cf23..05a45239dabf8 100644 --- a/smoke-test/tests/containers/containers_test.py +++ b/smoke-test/tests/containers/containers_test.py @@ -1,5 +1,7 @@ import pytest -from tests.utils import delete_urns_from_file, get_frontend_url, ingest_file_via_rest + +from tests.utils import (delete_urns_from_file, get_frontend_url, + ingest_file_via_rest) @pytest.fixture(scope="module", autouse=False) diff --git a/smoke-test/tests/cypress/integration_test.py b/smoke-test/tests/cypress/integration_test.py index b3bacf39ac7ae..4ad2bc53fa87d 100644 --- a/smoke-test/tests/cypress/integration_test.py +++ b/smoke-test/tests/cypress/integration_test.py @@ -1,18 +1,16 @@ -from typing import Set, List - import datetime -import pytest -import subprocess import os +import subprocess +from typing import List, Set + +import pytest + +from tests.setup.lineage.ingest_time_lineage import (get_time_lineage_urns, + ingest_time_lineage) +from tests.utils import (create_datahub_step_state_aspects, delete_urns, + delete_urns_from_file, get_admin_username, + ingest_file_via_rest) -from tests.utils import ( - create_datahub_step_state_aspects, - get_admin_username, - ingest_file_via_rest, - delete_urns_from_file, - delete_urns, -) -from tests.setup.lineage.ingest_time_lineage import ingest_time_lineage, get_time_lineage_urns CYPRESS_TEST_DATA_DIR = "tests/cypress" TEST_DATA_FILENAME = "data.json" @@ -145,7 +143,6 @@ def ingest_cleanup_data(): delete_urns_from_file(f"{CYPRESS_TEST_DATA_DIR}/{TEST_ONBOARDING_DATA_FILENAME}") delete_urns(get_time_lineage_urns()) - print_now() print("deleting onboarding data file") if os.path.exists(f"{CYPRESS_TEST_DATA_DIR}/{TEST_ONBOARDING_DATA_FILENAME}"): diff --git a/smoke-test/tests/dataproduct/test_dataproduct.py b/smoke-test/tests/dataproduct/test_dataproduct.py index db198098f21fa..baef1cb1cb3ba 100644 --- a/smoke-test/tests/dataproduct/test_dataproduct.py +++ b/smoke-test/tests/dataproduct/test_dataproduct.py @@ -1,4 +1,6 @@ +import logging import os +import subprocess import tempfile import time from random import randint @@ -17,8 +19,6 @@ DomainPropertiesClass, DomainsClass) from datahub.utilities.urns.urn import Urn -import subprocess -import logging logger = logging.getLogger(__name__) diff --git a/smoke-test/tests/delete/delete_test.py b/smoke-test/tests/delete/delete_test.py index 68e001f983fbf..d920faaf3a89a 100644 --- a/smoke-test/tests/delete/delete_test.py +++ b/smoke-test/tests/delete/delete_test.py @@ -1,16 +1,14 @@ -import os import json -import pytest +import os from time import sleep + +import pytest from datahub.cli.cli_utils import get_aspects_for_entity from datahub.cli.ingest_cli import get_session_and_host -from tests.utils import ( - ingest_file_via_rest, - wait_for_healthcheck_util, - delete_urns_from_file, - wait_for_writes_to_sync, - get_datahub_graph, -) + +from tests.utils import (delete_urns_from_file, get_datahub_graph, + ingest_file_via_rest, wait_for_healthcheck_util, + wait_for_writes_to_sync) # Disable telemetry os.environ["DATAHUB_TELEMETRY_ENABLED"] = "false" @@ -102,7 +100,7 @@ def test_delete_reference(test_setup, depends=["test_healthchecks"]): graph.delete_references_to_urn(tag_urn, dry_run=False) wait_for_writes_to_sync() - + # Validate that references no longer exist references_count, related_aspects = graph.delete_references_to_urn( tag_urn, dry_run=True diff --git a/smoke-test/tests/deprecation/deprecation_test.py b/smoke-test/tests/deprecation/deprecation_test.py index 1149a970aa8e5..a8969804d03d7 100644 --- a/smoke-test/tests/deprecation/deprecation_test.py +++ b/smoke-test/tests/deprecation/deprecation_test.py @@ -1,10 +1,7 @@ import pytest -from tests.utils import ( - delete_urns_from_file, - get_frontend_url, - ingest_file_via_rest, - get_root_urn, -) + +from tests.utils import (delete_urns_from_file, get_frontend_url, get_root_urn, + ingest_file_via_rest) @pytest.fixture(scope="module", autouse=True) diff --git a/smoke-test/tests/domains/domains_test.py b/smoke-test/tests/domains/domains_test.py index 7ffe1682cafd8..fa8c918e3cbe1 100644 --- a/smoke-test/tests/domains/domains_test.py +++ b/smoke-test/tests/domains/domains_test.py @@ -1,12 +1,8 @@ import pytest import tenacity -from tests.utils import ( - delete_urns_from_file, - get_frontend_url, - get_gms_url, - ingest_file_via_rest, - get_sleep_info, -) + +from tests.utils import (delete_urns_from_file, get_frontend_url, get_gms_url, + get_sleep_info, ingest_file_via_rest) sleep_sec, sleep_times = get_sleep_info() @@ -240,4 +236,7 @@ def test_set_unset_domain(frontend_session, ingest_cleanup_data): assert res_data assert res_data["data"]["dataset"]["domain"]["domain"]["urn"] == domain_urn - assert res_data["data"]["dataset"]["domain"]["domain"]["properties"]["name"] == "Engineering" + assert ( + res_data["data"]["dataset"]["domain"]["domain"]["properties"]["name"] + == "Engineering" + ) diff --git a/smoke-test/tests/managed-ingestion/managed_ingestion_test.py b/smoke-test/tests/managed-ingestion/managed_ingestion_test.py index 1238a1dd5730a..b5e408731334e 100644 --- a/smoke-test/tests/managed-ingestion/managed_ingestion_test.py +++ b/smoke-test/tests/managed-ingestion/managed_ingestion_test.py @@ -3,7 +3,8 @@ import pytest import tenacity -from tests.utils import get_frontend_url, get_sleep_info, wait_for_healthcheck_util +from tests.utils import (get_frontend_url, get_sleep_info, + wait_for_healthcheck_util) sleep_sec, sleep_times = get_sleep_info() diff --git a/smoke-test/tests/patch/common_patch_tests.py b/smoke-test/tests/patch/common_patch_tests.py index 574e4fd4e4c88..f1d6abf5da794 100644 --- a/smoke-test/tests/patch/common_patch_tests.py +++ b/smoke-test/tests/patch/common_patch_tests.py @@ -2,25 +2,17 @@ import uuid from typing import Dict, Optional, Type -from datahub.emitter.mce_builder import ( - make_tag_urn, - make_term_urn, - make_user_urn, -) +from datahub.emitter.mce_builder import (make_tag_urn, make_term_urn, + make_user_urn) from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.mcp_patch_builder import MetadataPatchProposal from datahub.ingestion.graph.client import DataHubGraph, DataHubGraphConfig -from datahub.metadata.schema_classes import ( - AuditStampClass, - GlobalTagsClass, - GlossaryTermAssociationClass, - GlossaryTermsClass, - OwnerClass, - OwnershipClass, - OwnershipTypeClass, - TagAssociationClass, - _Aspect, -) +from datahub.metadata.schema_classes import (AuditStampClass, GlobalTagsClass, + GlossaryTermAssociationClass, + GlossaryTermsClass, OwnerClass, + OwnershipClass, + OwnershipTypeClass, + TagAssociationClass, _Aspect) def helper_test_entity_terms_patch( @@ -34,18 +26,14 @@ def get_terms(graph, entity_urn): term_urn = make_term_urn(term=f"testTerm-{uuid.uuid4()}") - term_association = GlossaryTermAssociationClass( - urn=term_urn, context="test" - ) + term_association = GlossaryTermAssociationClass(urn=term_urn, context="test") global_terms = GlossaryTermsClass( terms=[term_association], auditStamp=AuditStampClass( time=int(time.time() * 1000.0), actor=make_user_urn("tester") ), ) - mcpw = MetadataChangeProposalWrapper( - entityUrn=test_entity_urn, aspect=global_terms - ) + mcpw = MetadataChangeProposalWrapper(entityUrn=test_entity_urn, aspect=global_terms) with DataHubGraph(DataHubGraphConfig()) as graph: graph.emit_mcp(mcpw) @@ -88,9 +76,7 @@ def helper_test_dataset_tags_patch( tag_association = TagAssociationClass(tag=tag_urn, context="test") global_tags = GlobalTagsClass(tags=[tag_association]) - mcpw = MetadataChangeProposalWrapper( - entityUrn=test_entity_urn, aspect=global_tags - ) + mcpw = MetadataChangeProposalWrapper(entityUrn=test_entity_urn, aspect=global_tags) with DataHubGraph(DataHubGraphConfig()) as graph: graph.emit_mcp(mcpw) @@ -153,15 +139,11 @@ def helper_test_ownership_patch( assert owner.owners[0].owner == make_user_urn("jdoe") for patch_mcp in ( - patch_builder_class(test_entity_urn) - .add_owner(owner_to_add) - .build() + patch_builder_class(test_entity_urn).add_owner(owner_to_add).build() ): graph.emit_mcp(patch_mcp) - owner = graph.get_aspect( - entity_urn=test_entity_urn, aspect_type=OwnershipClass - ) + owner = graph.get_aspect(entity_urn=test_entity_urn, aspect_type=OwnershipClass) assert len(owner.owners) == 2 for patch_mcp in ( @@ -171,9 +153,7 @@ def helper_test_ownership_patch( ): graph.emit_mcp(patch_mcp) - owner = graph.get_aspect( - entity_urn=test_entity_urn, aspect_type=OwnershipClass - ) + owner = graph.get_aspect(entity_urn=test_entity_urn, aspect_type=OwnershipClass) assert len(owner.owners) == 1 assert owner.owners[0].owner == make_user_urn("jdoe") @@ -199,9 +179,7 @@ def get_custom_properties( orig_aspect = base_aspect assert hasattr(orig_aspect, "customProperties") orig_aspect.customProperties = base_property_map - mcpw = MetadataChangeProposalWrapper( - entityUrn=test_entity_urn, aspect=orig_aspect - ) + mcpw = MetadataChangeProposalWrapper(entityUrn=test_entity_urn, aspect=orig_aspect) with DataHubGraph(DataHubGraphConfig()) as graph: graph.emit(mcpw) diff --git a/smoke-test/tests/patch/test_datajob_patches.py b/smoke-test/tests/patch/test_datajob_patches.py index 407410ee89914..342d5d683228a 100644 --- a/smoke-test/tests/patch/test_datajob_patches.py +++ b/smoke-test/tests/patch/test_datajob_patches.py @@ -3,19 +3,14 @@ from datahub.emitter.mce_builder import make_data_job_urn, make_dataset_urn from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.graph.client import DataHubGraph, DataHubGraphConfig -from datahub.metadata.schema_classes import ( - DataJobInfoClass, - DataJobInputOutputClass, - EdgeClass, -) +from datahub.metadata.schema_classes import (DataJobInfoClass, + DataJobInputOutputClass, + EdgeClass) from datahub.specific.datajob import DataJobPatchBuilder from tests.patch.common_patch_tests import ( - helper_test_custom_properties_patch, - helper_test_dataset_tags_patch, - helper_test_entity_terms_patch, - helper_test_ownership_patch, -) + helper_test_custom_properties_patch, helper_test_dataset_tags_patch, + helper_test_entity_terms_patch, helper_test_ownership_patch) def _make_test_datajob_urn( @@ -37,16 +32,12 @@ def test_datajob_ownership_patch(wait_for_healthchecks): # Tags def test_datajob_tags_patch(wait_for_healthchecks): - helper_test_dataset_tags_patch( - _make_test_datajob_urn(), DataJobPatchBuilder - ) + helper_test_dataset_tags_patch(_make_test_datajob_urn(), DataJobPatchBuilder) # Terms def test_dataset_terms_patch(wait_for_healthchecks): - helper_test_entity_terms_patch( - _make_test_datajob_urn(), DataJobPatchBuilder - ) + helper_test_entity_terms_patch(_make_test_datajob_urn(), DataJobPatchBuilder) # Custom Properties diff --git a/smoke-test/tests/patch/test_dataset_patches.py b/smoke-test/tests/patch/test_dataset_patches.py index 239aab64675d8..6704d19760fb9 100644 --- a/smoke-test/tests/patch/test_dataset_patches.py +++ b/smoke-test/tests/patch/test_dataset_patches.py @@ -20,7 +20,10 @@ UpstreamClass, UpstreamLineageClass) from datahub.specific.dataset import DatasetPatchBuilder -from tests.patch.common_patch_tests import helper_test_entity_terms_patch, helper_test_dataset_tags_patch, helper_test_ownership_patch, helper_test_custom_properties_patch + +from tests.patch.common_patch_tests import ( + helper_test_custom_properties_patch, helper_test_dataset_tags_patch, + helper_test_entity_terms_patch, helper_test_ownership_patch) # Common Aspect Patch Tests @@ -31,6 +34,7 @@ def test_dataset_ownership_patch(wait_for_healthchecks): ) helper_test_ownership_patch(dataset_urn, DatasetPatchBuilder) + # Tags def test_dataset_tags_patch(wait_for_healthchecks): dataset_urn = make_dataset_urn( @@ -38,6 +42,7 @@ def test_dataset_tags_patch(wait_for_healthchecks): ) helper_test_dataset_tags_patch(dataset_urn, DatasetPatchBuilder) + # Terms def test_dataset_terms_patch(wait_for_healthchecks): dataset_urn = make_dataset_urn( @@ -284,8 +289,15 @@ def test_custom_properties_patch(wait_for_healthchecks): dataset_urn = make_dataset_urn( platform="hive", name=f"SampleHiveDataset-{uuid.uuid4()}", env="PROD" ) - orig_dataset_properties = DatasetPropertiesClass(name="test_name", description="test_description") - helper_test_custom_properties_patch(test_entity_urn=dataset_urn, patch_builder_class=DatasetPatchBuilder, custom_properties_aspect_class=DatasetPropertiesClass, base_aspect=orig_dataset_properties) + orig_dataset_properties = DatasetPropertiesClass( + name="test_name", description="test_description" + ) + helper_test_custom_properties_patch( + test_entity_urn=dataset_urn, + patch_builder_class=DatasetPatchBuilder, + custom_properties_aspect_class=DatasetPropertiesClass, + base_aspect=orig_dataset_properties, + ) with DataHubGraph(DataHubGraphConfig()) as graph: # Patch custom properties along with name diff --git a/smoke-test/tests/policies/test_policies.py b/smoke-test/tests/policies/test_policies.py index b7091541894dd..67142181d2b96 100644 --- a/smoke-test/tests/policies/test_policies.py +++ b/smoke-test/tests/policies/test_policies.py @@ -1,12 +1,8 @@ import pytest import tenacity -from tests.utils import ( - get_frontend_url, - wait_for_healthcheck_util, - get_frontend_session, - get_sleep_info, - get_root_urn, -) + +from tests.utils import (get_frontend_session, get_frontend_url, get_root_urn, + get_sleep_info, wait_for_healthcheck_util) TEST_POLICY_NAME = "Updated Platform Policy" diff --git a/smoke-test/tests/setup/lineage/helper_classes.py b/smoke-test/tests/setup/lineage/helper_classes.py index 53f77b08d15ed..d550f3093be85 100644 --- a/smoke-test/tests/setup/lineage/helper_classes.py +++ b/smoke-test/tests/setup/lineage/helper_classes.py @@ -1,10 +1,7 @@ from dataclasses import dataclass from typing import Any, Dict, List, Optional -from datahub.metadata.schema_classes import ( - EdgeClass, - SchemaFieldDataTypeClass, -) +from datahub.metadata.schema_classes import EdgeClass, SchemaFieldDataTypeClass @dataclass diff --git a/smoke-test/tests/setup/lineage/ingest_data_job_change.py b/smoke-test/tests/setup/lineage/ingest_data_job_change.py index 8e3e9c5352922..588a1625419bc 100644 --- a/smoke-test/tests/setup/lineage/ingest_data_job_change.py +++ b/smoke-test/tests/setup/lineage/ingest_data_job_change.py @@ -1,36 +1,20 @@ from typing import List -from datahub.emitter.mce_builder import ( - make_dataset_urn, - make_data_flow_urn, - make_data_job_urn_with_flow, -) +from datahub.emitter.mce_builder import (make_data_flow_urn, + make_data_job_urn_with_flow, + make_dataset_urn) from datahub.emitter.rest_emitter import DatahubRestEmitter -from datahub.metadata.schema_classes import ( - DateTypeClass, - NumberTypeClass, - SchemaFieldDataTypeClass, - StringTypeClass, -) +from datahub.metadata.schema_classes import (DateTypeClass, NumberTypeClass, + SchemaFieldDataTypeClass, + StringTypeClass) -from tests.setup.lineage.constants import ( - AIRFLOW_DATA_PLATFORM, - SNOWFLAKE_DATA_PLATFORM, - TIMESTAMP_MILLIS_EIGHT_DAYS_AGO, - TIMESTAMP_MILLIS_ONE_DAY_AGO, -) -from tests.setup.lineage.helper_classes import ( - Field, - Dataset, - Task, - Pipeline, -) -from tests.setup.lineage.utils import ( - create_edge, - create_node, - create_nodes_and_edges, - emit_mcps, -) +from tests.setup.lineage.constants import (AIRFLOW_DATA_PLATFORM, + SNOWFLAKE_DATA_PLATFORM, + TIMESTAMP_MILLIS_EIGHT_DAYS_AGO, + TIMESTAMP_MILLIS_ONE_DAY_AGO) +from tests.setup.lineage.helper_classes import Dataset, Field, Pipeline, Task +from tests.setup.lineage.utils import (create_edge, create_node, + create_nodes_and_edges, emit_mcps) # Constants for Case 2 DAILY_TEMPERATURE_DATASET_ID = "climate.daily_temperature" diff --git a/smoke-test/tests/setup/lineage/ingest_dataset_join_change.py b/smoke-test/tests/setup/lineage/ingest_dataset_join_change.py index 35a8e6d5cf02e..bb9f51b6b5e9b 100644 --- a/smoke-test/tests/setup/lineage/ingest_dataset_join_change.py +++ b/smoke-test/tests/setup/lineage/ingest_dataset_join_change.py @@ -1,32 +1,18 @@ from typing import List -from datahub.emitter.mce_builder import ( - make_dataset_urn, -) +from datahub.emitter.mce_builder import make_dataset_urn from datahub.emitter.rest_emitter import DatahubRestEmitter -from datahub.metadata.schema_classes import ( - NumberTypeClass, - SchemaFieldDataTypeClass, - StringTypeClass, - UpstreamClass, -) +from datahub.metadata.schema_classes import (NumberTypeClass, + SchemaFieldDataTypeClass, + StringTypeClass, UpstreamClass) -from tests.setup.lineage.constants import ( - DATASET_ENTITY_TYPE, - SNOWFLAKE_DATA_PLATFORM, - TIMESTAMP_MILLIS_EIGHT_DAYS_AGO, - TIMESTAMP_MILLIS_ONE_DAY_AGO, -) -from tests.setup.lineage.helper_classes import ( - Field, - Dataset, -) -from tests.setup.lineage.utils import ( - create_node, - create_upstream_edge, - create_upstream_mcp, - emit_mcps, -) +from tests.setup.lineage.constants import (DATASET_ENTITY_TYPE, + SNOWFLAKE_DATA_PLATFORM, + TIMESTAMP_MILLIS_EIGHT_DAYS_AGO, + TIMESTAMP_MILLIS_ONE_DAY_AGO) +from tests.setup.lineage.helper_classes import Dataset, Field +from tests.setup.lineage.utils import (create_node, create_upstream_edge, + create_upstream_mcp, emit_mcps) # Constants for Case 3 GDP_DATASET_ID = "economic_data.gdp" diff --git a/smoke-test/tests/setup/lineage/ingest_input_datasets_change.py b/smoke-test/tests/setup/lineage/ingest_input_datasets_change.py index f4fb795147478..6079d7a3d2b63 100644 --- a/smoke-test/tests/setup/lineage/ingest_input_datasets_change.py +++ b/smoke-test/tests/setup/lineage/ingest_input_datasets_change.py @@ -1,36 +1,20 @@ from typing import List -from datahub.emitter.mce_builder import ( - make_dataset_urn, - make_data_flow_urn, - make_data_job_urn_with_flow, -) +from datahub.emitter.mce_builder import (make_data_flow_urn, + make_data_job_urn_with_flow, + make_dataset_urn) from datahub.emitter.rest_emitter import DatahubRestEmitter -from datahub.metadata.schema_classes import ( - NumberTypeClass, - SchemaFieldDataTypeClass, - StringTypeClass, -) - -from tests.setup.lineage.constants import ( - AIRFLOW_DATA_PLATFORM, - BQ_DATA_PLATFORM, - TIMESTAMP_MILLIS_EIGHT_DAYS_AGO, - TIMESTAMP_MILLIS_ONE_DAY_AGO, -) -from tests.setup.lineage.helper_classes import ( - Field, - Dataset, - Task, - Pipeline, -) -from tests.setup.lineage.utils import ( - create_edge, - create_node, - create_nodes_and_edges, - emit_mcps, -) +from datahub.metadata.schema_classes import (NumberTypeClass, + SchemaFieldDataTypeClass, + StringTypeClass) +from tests.setup.lineage.constants import (AIRFLOW_DATA_PLATFORM, + BQ_DATA_PLATFORM, + TIMESTAMP_MILLIS_EIGHT_DAYS_AGO, + TIMESTAMP_MILLIS_ONE_DAY_AGO) +from tests.setup.lineage.helper_classes import Dataset, Field, Pipeline, Task +from tests.setup.lineage.utils import (create_edge, create_node, + create_nodes_and_edges, emit_mcps) # Constants for Case 1 TRANSACTIONS_DATASET_ID = "transactions.transactions" diff --git a/smoke-test/tests/setup/lineage/ingest_time_lineage.py b/smoke-test/tests/setup/lineage/ingest_time_lineage.py index cae8e0124d501..3aec979707290 100644 --- a/smoke-test/tests/setup/lineage/ingest_time_lineage.py +++ b/smoke-test/tests/setup/lineage/ingest_time_lineage.py @@ -1,12 +1,14 @@ +import os from typing import List from datahub.emitter.rest_emitter import DatahubRestEmitter -from tests.setup.lineage.ingest_input_datasets_change import ingest_input_datasets_change, get_input_datasets_change_urns -from tests.setup.lineage.ingest_data_job_change import ingest_data_job_change, get_data_job_change_urns -from tests.setup.lineage.ingest_dataset_join_change import ingest_dataset_join_change, get_dataset_join_change_urns - -import os +from tests.setup.lineage.ingest_data_job_change import ( + get_data_job_change_urns, ingest_data_job_change) +from tests.setup.lineage.ingest_dataset_join_change import ( + get_dataset_join_change_urns, ingest_dataset_join_change) +from tests.setup.lineage.ingest_input_datasets_change import ( + get_input_datasets_change_urns, ingest_input_datasets_change) SERVER = os.getenv("DATAHUB_SERVER") or "http://localhost:8080" TOKEN = os.getenv("DATAHUB_TOKEN") or "" @@ -20,4 +22,8 @@ def ingest_time_lineage() -> None: def get_time_lineage_urns() -> List[str]: - return get_input_datasets_change_urns() + get_data_job_change_urns() + get_dataset_join_change_urns() + return ( + get_input_datasets_change_urns() + + get_data_job_change_urns() + + get_dataset_join_change_urns() + ) diff --git a/smoke-test/tests/setup/lineage/utils.py b/smoke-test/tests/setup/lineage/utils.py index 672f7a945a6af..c72f6ccb89b7a 100644 --- a/smoke-test/tests/setup/lineage/utils.py +++ b/smoke-test/tests/setup/lineage/utils.py @@ -1,41 +1,30 @@ import datetime -from datahub.emitter.mce_builder import ( - make_data_platform_urn, - make_dataset_urn, - make_data_job_urn_with_flow, - make_data_flow_urn, -) +from typing import List + +from datahub.emitter.mce_builder import (make_data_flow_urn, + make_data_job_urn_with_flow, + make_data_platform_urn, + make_dataset_urn) from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.rest_emitter import DatahubRestEmitter from datahub.metadata.com.linkedin.pegasus2avro.dataset import UpstreamLineage -from datahub.metadata.schema_classes import ( - AuditStampClass, - ChangeTypeClass, - DatasetLineageTypeClass, - DatasetPropertiesClass, - DataFlowInfoClass, - DataJobInputOutputClass, - DataJobInfoClass, - EdgeClass, - MySqlDDLClass, - SchemaFieldClass, - SchemaMetadataClass, - UpstreamClass, -) -from typing import List - -from tests.setup.lineage.constants import ( - DATASET_ENTITY_TYPE, - DATA_JOB_ENTITY_TYPE, - DATA_FLOW_ENTITY_TYPE, - DATA_FLOW_INFO_ASPECT_NAME, - DATA_JOB_INFO_ASPECT_NAME, - DATA_JOB_INPUT_OUTPUT_ASPECT_NAME, -) -from tests.setup.lineage.helper_classes import ( - Dataset, - Pipeline, -) +from datahub.metadata.schema_classes import (AuditStampClass, ChangeTypeClass, + DataFlowInfoClass, + DataJobInfoClass, + DataJobInputOutputClass, + DatasetLineageTypeClass, + DatasetPropertiesClass, EdgeClass, + MySqlDDLClass, SchemaFieldClass, + SchemaMetadataClass, + UpstreamClass) + +from tests.setup.lineage.constants import (DATA_FLOW_ENTITY_TYPE, + DATA_FLOW_INFO_ASPECT_NAME, + DATA_JOB_ENTITY_TYPE, + DATA_JOB_INFO_ASPECT_NAME, + DATA_JOB_INPUT_OUTPUT_ASPECT_NAME, + DATASET_ENTITY_TYPE) +from tests.setup.lineage.helper_classes import Dataset, Pipeline def create_node(dataset: Dataset) -> List[MetadataChangeProposalWrapper]: @@ -85,10 +74,10 @@ def create_node(dataset: Dataset) -> List[MetadataChangeProposalWrapper]: def create_edge( - source_urn: str, - destination_urn: str, - created_timestamp_millis: int, - updated_timestamp_millis: int, + source_urn: str, + destination_urn: str, + created_timestamp_millis: int, + updated_timestamp_millis: int, ) -> EdgeClass: created_audit_stamp: AuditStampClass = AuditStampClass( time=created_timestamp_millis, actor="urn:li:corpuser:unknown" @@ -105,7 +94,7 @@ def create_edge( def create_nodes_and_edges( - airflow_dag: Pipeline, + airflow_dag: Pipeline, ) -> List[MetadataChangeProposalWrapper]: mcps = [] data_flow_urn = make_data_flow_urn( @@ -160,9 +149,9 @@ def create_nodes_and_edges( def create_upstream_edge( - upstream_entity_urn: str, - created_timestamp_millis: int, - updated_timestamp_millis: int, + upstream_entity_urn: str, + created_timestamp_millis: int, + updated_timestamp_millis: int, ): created_audit_stamp: AuditStampClass = AuditStampClass( time=created_timestamp_millis, actor="urn:li:corpuser:unknown" @@ -180,11 +169,11 @@ def create_upstream_edge( def create_upstream_mcp( - entity_type: str, - entity_urn: str, - upstreams: List[UpstreamClass], - timestamp_millis: int, - run_id: str = "", + entity_type: str, + entity_urn: str, + upstreams: List[UpstreamClass], + timestamp_millis: int, + run_id: str = "", ) -> MetadataChangeProposalWrapper: print(f"Creating upstreamLineage aspect for {entity_urn}") timestamp_millis: int = int(datetime.datetime.now().timestamp() * 1000) @@ -203,7 +192,7 @@ def create_upstream_mcp( def emit_mcps( - emitter: DatahubRestEmitter, mcps: List[MetadataChangeProposalWrapper] + emitter: DatahubRestEmitter, mcps: List[MetadataChangeProposalWrapper] ) -> None: for mcp in mcps: emitter.emit_mcp(mcp) diff --git a/smoke-test/tests/tags-and-terms/tags_and_terms_test.py b/smoke-test/tests/tags-and-terms/tags_and_terms_test.py index b0ca29b544cfe..6ac75765286f0 100644 --- a/smoke-test/tests/tags-and-terms/tags_and_terms_test.py +++ b/smoke-test/tests/tags-and-terms/tags_and_terms_test.py @@ -1,5 +1,7 @@ import pytest -from tests.utils import delete_urns_from_file, get_frontend_url, ingest_file_via_rest, wait_for_healthcheck_util + +from tests.utils import (delete_urns_from_file, get_frontend_url, + ingest_file_via_rest, wait_for_healthcheck_util) @pytest.fixture(scope="module", autouse=True) diff --git a/smoke-test/tests/telemetry/telemetry_test.py b/smoke-test/tests/telemetry/telemetry_test.py index 3672abcda948d..3127061c9f506 100644 --- a/smoke-test/tests/telemetry/telemetry_test.py +++ b/smoke-test/tests/telemetry/telemetry_test.py @@ -7,5 +7,7 @@ def test_no_clientID(): client_id_urn = "urn:li:telemetry:clientId" aspect = ["telemetryClientId"] - res_data = json.dumps(get_aspects_for_entity(entity_urn=client_id_urn, aspects=aspect, typed=False)) + res_data = json.dumps( + get_aspects_for_entity(entity_urn=client_id_urn, aspects=aspect, typed=False) + ) assert res_data == "{}" diff --git a/smoke-test/tests/test_result_msg.py b/smoke-test/tests/test_result_msg.py index e3b336db9d66c..b9775e8ee4acd 100644 --- a/smoke-test/tests/test_result_msg.py +++ b/smoke-test/tests/test_result_msg.py @@ -1,6 +1,6 @@ -from slack_sdk import WebClient import os +from slack_sdk import WebClient datahub_stats = {} @@ -10,10 +10,10 @@ def add_datahub_stats(stat_name, stat_val): def send_to_slack(passed: str): - slack_api_token = os.getenv('SLACK_API_TOKEN') - slack_channel = os.getenv('SLACK_CHANNEL') - slack_thread_ts = os.getenv('SLACK_THREAD_TS') - test_identifier = os.getenv('TEST_IDENTIFIER', 'LOCAL_TEST') + slack_api_token = os.getenv("SLACK_API_TOKEN") + slack_channel = os.getenv("SLACK_CHANNEL") + slack_thread_ts = os.getenv("SLACK_THREAD_TS") + test_identifier = os.getenv("TEST_IDENTIFIER", "LOCAL_TEST") if slack_api_token is None or slack_channel is None: return client = WebClient(token=slack_api_token) @@ -26,14 +26,21 @@ def send_to_slack(passed: str): message += f"Num {entity_type} is {val}\n" if slack_thread_ts is None: - client.chat_postMessage(channel=slack_channel, text=f'{test_identifier} Status - {passed}\n{message}') + client.chat_postMessage( + channel=slack_channel, + text=f"{test_identifier} Status - {passed}\n{message}", + ) else: - client.chat_postMessage(channel=slack_channel, text=f'{test_identifier} Status - {passed}\n{message}', thread_ts=slack_thread_ts) + client.chat_postMessage( + channel=slack_channel, + text=f"{test_identifier} Status - {passed}\n{message}", + thread_ts=slack_thread_ts, + ) def send_message(exitstatus): try: - send_to_slack('PASSED' if exitstatus == 0 else 'FAILED') + send_to_slack("PASSED" if exitstatus == 0 else "FAILED") except Exception as e: # We don't want to fail pytest at all print(f"Exception happened for sending msg to slack {e}") diff --git a/smoke-test/tests/test_stateful_ingestion.py b/smoke-test/tests/test_stateful_ingestion.py index a10cf13a08029..c6adb402e5d51 100644 --- a/smoke-test/tests/test_stateful_ingestion.py +++ b/smoke-test/tests/test_stateful_ingestion.py @@ -4,17 +4,15 @@ from datahub.ingestion.run.pipeline import Pipeline from datahub.ingestion.source.sql.mysql import MySQLConfig, MySQLSource from datahub.ingestion.source.state.checkpoint import Checkpoint -from datahub.ingestion.source.state.entity_removal_state import GenericCheckpointState -from datahub.ingestion.source.state.stale_entity_removal_handler import StaleEntityRemovalHandler +from datahub.ingestion.source.state.entity_removal_state import \ + GenericCheckpointState +from datahub.ingestion.source.state.stale_entity_removal_handler import \ + StaleEntityRemovalHandler from sqlalchemy import create_engine from sqlalchemy.sql import text -from tests.utils import ( - get_gms_url, - get_mysql_password, - get_mysql_url, - get_mysql_username, -) +from tests.utils import (get_gms_url, get_mysql_password, get_mysql_url, + get_mysql_username) def test_stateful_ingestion(wait_for_healthchecks): diff --git a/smoke-test/tests/tests/tests_test.py b/smoke-test/tests/tests/tests_test.py index 0b87f90a92c58..213a2ea087b7a 100644 --- a/smoke-test/tests/tests/tests_test.py +++ b/smoke-test/tests/tests/tests_test.py @@ -1,9 +1,13 @@ import pytest import tenacity -from tests.utils import delete_urns_from_file, get_frontend_url, ingest_file_via_rest, wait_for_healthcheck_util, get_sleep_info + +from tests.utils import (delete_urns_from_file, get_frontend_url, + get_sleep_info, ingest_file_via_rest, + wait_for_healthcheck_util) sleep_sec, sleep_times = get_sleep_info() + @pytest.fixture(scope="module", autouse=True) def ingest_cleanup_data(request): print("ingesting test data") @@ -18,6 +22,7 @@ def wait_for_healthchecks(): wait_for_healthcheck_util() yield + @pytest.mark.dependency() def test_healthchecks(wait_for_healthchecks): # Call to wait_for_healthchecks fixture will do the actual functionality. diff --git a/smoke-test/tests/timeline/timeline_test.py b/smoke-test/tests/timeline/timeline_test.py index a73d585c6c72d..4705343c1a2ba 100644 --- a/smoke-test/tests/timeline/timeline_test.py +++ b/smoke-test/tests/timeline/timeline_test.py @@ -3,14 +3,14 @@ from datahub.cli import timeline_cli from datahub.cli.cli_utils import guess_entity_type, post_entity -from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync, get_datahub_graph + +from tests.utils import (get_datahub_graph, ingest_file_via_rest, + wait_for_writes_to_sync) def test_all(): platform = "urn:li:dataPlatform:kafka" - dataset_name = ( - "test-timeline-sample-kafka" - ) + dataset_name = "test-timeline-sample-kafka" env = "PROD" dataset_urn = f"urn:li:dataset:({platform},{dataset_name},{env})" @@ -18,8 +18,13 @@ def test_all(): ingest_file_via_rest("tests/timeline/timeline_test_datav2.json") ingest_file_via_rest("tests/timeline/timeline_test_datav3.json") - res_data = timeline_cli.get_timeline(dataset_urn, ["TAG", "DOCUMENTATION", "TECHNICAL_SCHEMA", "GLOSSARY_TERM", - "OWNER"], None, None, False) + res_data = timeline_cli.get_timeline( + dataset_urn, + ["TAG", "DOCUMENTATION", "TECHNICAL_SCHEMA", "GLOSSARY_TERM", "OWNER"], + None, + None, + False, + ) get_datahub_graph().hard_delete_entity(urn=dataset_urn) assert res_data @@ -35,9 +40,7 @@ def test_all(): def test_schema(): platform = "urn:li:dataPlatform:kafka" - dataset_name = ( - "test-timeline-sample-kafka" - ) + dataset_name = "test-timeline-sample-kafka" env = "PROD" dataset_urn = f"urn:li:dataset:({platform},{dataset_name},{env})" @@ -45,7 +48,9 @@ def test_schema(): put(dataset_urn, "schemaMetadata", "test_resources/timeline/newschemav2.json") put(dataset_urn, "schemaMetadata", "test_resources/timeline/newschemav3.json") - res_data = timeline_cli.get_timeline(dataset_urn, ["TECHNICAL_SCHEMA"], None, None, False) + res_data = timeline_cli.get_timeline( + dataset_urn, ["TECHNICAL_SCHEMA"], None, None, False + ) get_datahub_graph().hard_delete_entity(urn=dataset_urn) assert res_data @@ -61,9 +66,7 @@ def test_schema(): def test_glossary(): platform = "urn:li:dataPlatform:kafka" - dataset_name = ( - "test-timeline-sample-kafka" - ) + dataset_name = "test-timeline-sample-kafka" env = "PROD" dataset_urn = f"urn:li:dataset:({platform},{dataset_name},{env})" @@ -71,7 +74,9 @@ def test_glossary(): put(dataset_urn, "glossaryTerms", "test_resources/timeline/newglossaryv2.json") put(dataset_urn, "glossaryTerms", "test_resources/timeline/newglossaryv3.json") - res_data = timeline_cli.get_timeline(dataset_urn, ["GLOSSARY_TERM"], None, None, False) + res_data = timeline_cli.get_timeline( + dataset_urn, ["GLOSSARY_TERM"], None, None, False + ) get_datahub_graph().hard_delete_entity(urn=dataset_urn) assert res_data @@ -87,17 +92,29 @@ def test_glossary(): def test_documentation(): platform = "urn:li:dataPlatform:kafka" - dataset_name = ( - "test-timeline-sample-kafka" - ) + dataset_name = "test-timeline-sample-kafka" env = "PROD" dataset_urn = f"urn:li:dataset:({platform},{dataset_name},{env})" - put(dataset_urn, "institutionalMemory", "test_resources/timeline/newdocumentation.json") - put(dataset_urn, "institutionalMemory", "test_resources/timeline/newdocumentationv2.json") - put(dataset_urn, "institutionalMemory", "test_resources/timeline/newdocumentationv3.json") + put( + dataset_urn, + "institutionalMemory", + "test_resources/timeline/newdocumentation.json", + ) + put( + dataset_urn, + "institutionalMemory", + "test_resources/timeline/newdocumentationv2.json", + ) + put( + dataset_urn, + "institutionalMemory", + "test_resources/timeline/newdocumentationv3.json", + ) - res_data = timeline_cli.get_timeline(dataset_urn, ["DOCUMENTATION"], None, None, False) + res_data = timeline_cli.get_timeline( + dataset_urn, ["DOCUMENTATION"], None, None, False + ) get_datahub_graph().hard_delete_entity(urn=dataset_urn) assert res_data @@ -113,9 +130,7 @@ def test_documentation(): def test_tags(): platform = "urn:li:dataPlatform:kafka" - dataset_name = ( - "test-timeline-sample-kafka" - ) + dataset_name = "test-timeline-sample-kafka" env = "PROD" dataset_urn = f"urn:li:dataset:({platform},{dataset_name},{env})" @@ -139,9 +154,7 @@ def test_tags(): def test_ownership(): platform = "urn:li:dataPlatform:kafka" - dataset_name = ( - "test-timeline-sample-kafka" - ) + dataset_name = "test-timeline-sample-kafka" env = "PROD" dataset_urn = f"urn:li:dataset:({platform},{dataset_name},{env})" diff --git a/smoke-test/tests/tokens/revokable_access_token_test.py b/smoke-test/tests/tokens/revokable_access_token_test.py index b10ad3aa3fc2a..55f3de594af4e 100644 --- a/smoke-test/tests/tokens/revokable_access_token_test.py +++ b/smoke-test/tests/tokens/revokable_access_token_test.py @@ -1,15 +1,11 @@ import os -import pytest -import requests from time import sleep -from tests.utils import ( - get_frontend_url, - wait_for_healthcheck_util, - get_admin_credentials, - wait_for_writes_to_sync, -) +import pytest +import requests +from tests.utils import (get_admin_credentials, get_frontend_url, + wait_for_healthcheck_util, wait_for_writes_to_sync) # Disable telemetry os.environ["DATAHUB_TELEMETRY_ENABLED"] = "false" diff --git a/smoke-test/tests/utils.py b/smoke-test/tests/utils.py index af03efd4f71f8..bd75b13d1910f 100644 --- a/smoke-test/tests/utils.py +++ b/smoke-test/tests/utils.py @@ -1,19 +1,20 @@ import functools import json +import logging import os -from datetime import datetime, timedelta, timezone import subprocess import time -from typing import Any, Dict, List, Tuple +from datetime import datetime, timedelta, timezone from time import sleep -from joblib import Parallel, delayed +from typing import Any, Dict, List, Tuple -import requests_wrapper as requests -import logging from datahub.cli import cli_utils from datahub.cli.cli_utils import get_system_auth -from datahub.ingestion.graph.client import DataHubGraph, DatahubClientConfig +from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph from datahub.ingestion.run.pipeline import Pipeline +from joblib import Parallel, delayed + +import requests_wrapper as requests from tests.consistency_utils import wait_for_writes_to_sync TIME: int = 1581407189000 @@ -174,6 +175,7 @@ def delete(entry): wait_for_writes_to_sync() + # Fixed now value NOW: datetime = datetime.now() @@ -232,6 +234,3 @@ def create_datahub_step_state_aspects( ] with open(onboarding_filename, "w") as f: json.dump(aspects_dict, f, indent=2) - - - diff --git a/smoke-test/tests/views/views_test.py b/smoke-test/tests/views/views_test.py index 4da69750a167b..685c3bd80b04d 100644 --- a/smoke-test/tests/views/views_test.py +++ b/smoke-test/tests/views/views_test.py @@ -1,16 +1,14 @@ -import pytest import time + +import pytest import tenacity -from tests.utils import ( - delete_urns_from_file, - get_frontend_url, - get_gms_url, - ingest_file_via_rest, - get_sleep_info, -) + +from tests.utils import (delete_urns_from_file, get_frontend_url, get_gms_url, + get_sleep_info, ingest_file_via_rest) sleep_sec, sleep_times = get_sleep_info() + @pytest.mark.dependency() def test_healthchecks(wait_for_healthchecks): # Call to wait_for_healthchecks fixture will do the actual functionality. @@ -40,6 +38,7 @@ def _ensure_more_views(frontend_session, list_views_json, query_name, before_cou assert after_count == before_count + 1 return after_count + @tenacity.retry( stop=tenacity.stop_after_attempt(sleep_times), wait=tenacity.wait_fixed(sleep_sec) ) @@ -111,18 +110,18 @@ def test_create_list_delete_global_view(frontend_session): new_view_name = "Test View" new_view_description = "Test Description" new_view_definition = { - "entityTypes": ["DATASET", "DASHBOARD"], - "filter": { - "operator": "AND", - "filters": [ - { - "field": "tags", - "values": ["urn:li:tag:test"], - "negated": False, - "condition": "EQUAL" - } - ] - } + "entityTypes": ["DATASET", "DASHBOARD"], + "filter": { + "operator": "AND", + "filters": [ + { + "field": "tags", + "values": ["urn:li:tag:test"], + "negated": False, + "condition": "EQUAL", + } + ], + }, } # Create new View @@ -137,7 +136,7 @@ def test_create_list_delete_global_view(frontend_session): "viewType": "GLOBAL", "name": new_view_name, "description": new_view_description, - "definition": new_view_definition + "definition": new_view_definition, } }, } @@ -169,9 +168,7 @@ def test_create_list_delete_global_view(frontend_session): "query": """mutation deleteView($urn: String!) {\n deleteView(urn: $urn) }""", - "variables": { - "urn": view_urn - }, + "variables": {"urn": view_urn}, } response = frontend_session.post( @@ -189,7 +186,9 @@ def test_create_list_delete_global_view(frontend_session): ) -@pytest.mark.dependency(depends=["test_healthchecks", "test_create_list_delete_global_view"]) +@pytest.mark.dependency( + depends=["test_healthchecks", "test_create_list_delete_global_view"] +) def test_create_list_delete_personal_view(frontend_session): # Get count of existing views @@ -237,18 +236,18 @@ def test_create_list_delete_personal_view(frontend_session): new_view_name = "Test View" new_view_description = "Test Description" new_view_definition = { - "entityTypes": ["DATASET", "DASHBOARD"], - "filter": { - "operator": "AND", - "filters": [ - { - "field": "tags", - "values": ["urn:li:tag:test"], - "negated": False, - "condition": "EQUAL" - } - ] - } + "entityTypes": ["DATASET", "DASHBOARD"], + "filter": { + "operator": "AND", + "filters": [ + { + "field": "tags", + "values": ["urn:li:tag:test"], + "negated": False, + "condition": "EQUAL", + } + ], + }, } # Create new View @@ -263,7 +262,7 @@ def test_create_list_delete_personal_view(frontend_session): "viewType": "PERSONAL", "name": new_view_name, "description": new_view_description, - "definition": new_view_definition + "definition": new_view_definition, } }, } @@ -293,9 +292,7 @@ def test_create_list_delete_personal_view(frontend_session): "query": """mutation deleteView($urn: String!) {\n deleteView(urn: $urn) }""", - "variables": { - "urn": view_urn - }, + "variables": {"urn": view_urn}, } response = frontend_session.post( @@ -312,25 +309,28 @@ def test_create_list_delete_personal_view(frontend_session): before_count=new_count, ) -@pytest.mark.dependency(depends=["test_healthchecks", "test_create_list_delete_personal_view"]) + +@pytest.mark.dependency( + depends=["test_healthchecks", "test_create_list_delete_personal_view"] +) def test_update_global_view(frontend_session): # First create a view new_view_name = "Test View" new_view_description = "Test Description" new_view_definition = { - "entityTypes": ["DATASET", "DASHBOARD"], - "filter": { - "operator": "AND", - "filters": [ - { - "field": "tags", - "values": ["urn:li:tag:test"], - "negated": False, - "condition": "EQUAL" - } - ] - } + "entityTypes": ["DATASET", "DASHBOARD"], + "filter": { + "operator": "AND", + "filters": [ + { + "field": "tags", + "values": ["urn:li:tag:test"], + "negated": False, + "condition": "EQUAL", + } + ], + }, } # Create new View @@ -345,7 +345,7 @@ def test_update_global_view(frontend_session): "viewType": "PERSONAL", "name": new_view_name, "description": new_view_description, - "definition": new_view_definition + "definition": new_view_definition, } }, } @@ -366,18 +366,18 @@ def test_update_global_view(frontend_session): new_view_name = "New Test View" new_view_description = "New Test Description" new_view_definition = { - "entityTypes": ["DATASET", "DASHBOARD", "CHART", "DATA_FLOW"], - "filter": { - "operator": "OR", - "filters": [ - { - "field": "glossaryTerms", - "values": ["urn:li:glossaryTerm:test"], - "negated": True, - "condition": "CONTAIN" - } - ] - } + "entityTypes": ["DATASET", "DASHBOARD", "CHART", "DATA_FLOW"], + "filter": { + "operator": "OR", + "filters": [ + { + "field": "glossaryTerms", + "values": ["urn:li:glossaryTerm:test"], + "negated": True, + "condition": "CONTAIN", + } + ], + }, } update_view_json = { @@ -391,8 +391,8 @@ def test_update_global_view(frontend_session): "input": { "name": new_view_name, "description": new_view_description, - "definition": new_view_definition - } + "definition": new_view_definition, + }, }, } @@ -411,9 +411,7 @@ def test_update_global_view(frontend_session): "query": """mutation deleteView($urn: String!) {\n deleteView(urn: $urn) }""", - "variables": { - "urn": view_urn - }, + "variables": {"urn": view_urn}, } response = frontend_session.post( From 6ecdeda5ff590456c6bfadfa5c37821f7281169e Mon Sep 17 00:00:00 2001 From: Aseem Bansal Date: Tue, 10 Oct 2023 16:28:40 +0530 Subject: [PATCH 108/156] fix(setup): drop older table if exists (#8979) --- docker/mariadb/init.sql | 2 ++ docker/mysql-setup/init.sql | 2 ++ docker/mysql/init.sql | 2 ++ docker/postgres-setup/init.sql | 2 ++ docker/postgres/init.sql | 2 ++ 5 files changed, 10 insertions(+) diff --git a/docker/mariadb/init.sql b/docker/mariadb/init.sql index c4132575cf442..95c8cabbc5ca4 100644 --- a/docker/mariadb/init.sql +++ b/docker/mariadb/init.sql @@ -28,3 +28,5 @@ insert into metadata_aspect_v2 (urn, aspect, version, metadata, createdon, creat now(), 'urn:li:corpuser:__datahub_system' ); + +DROP TABLE IF EXISTS metadata_index; diff --git a/docker/mysql-setup/init.sql b/docker/mysql-setup/init.sql index 2370a971941d2..b789329ddfd17 100644 --- a/docker/mysql-setup/init.sql +++ b/docker/mysql-setup/init.sql @@ -39,3 +39,5 @@ INSERT INTO metadata_aspect_v2 SELECT * FROM temp_metadata_aspect_v2 WHERE NOT EXISTS (SELECT * from metadata_aspect_v2); DROP TABLE temp_metadata_aspect_v2; + +DROP TABLE IF EXISTS metadata_index; diff --git a/docker/mysql/init.sql b/docker/mysql/init.sql index b4b4e4617806c..aca57d7cd444c 100644 --- a/docker/mysql/init.sql +++ b/docker/mysql/init.sql @@ -27,3 +27,5 @@ INSERT INTO metadata_aspect_v2 (urn, aspect, version, metadata, createdon, creat now(), 'urn:li:corpuser:__datahub_system' ); + +DROP TABLE IF EXISTS metadata_index; diff --git a/docker/postgres-setup/init.sql b/docker/postgres-setup/init.sql index 12fff7aec7fe6..72b2f73192e00 100644 --- a/docker/postgres-setup/init.sql +++ b/docker/postgres-setup/init.sql @@ -35,3 +35,5 @@ INSERT INTO metadata_aspect_v2 SELECT * FROM temp_metadata_aspect_v2 WHERE NOT EXISTS (SELECT * from metadata_aspect_v2); DROP TABLE temp_metadata_aspect_v2; + +DROP TABLE IF EXISTS metadata_index; diff --git a/docker/postgres/init.sql b/docker/postgres/init.sql index cf477c135422e..87c8dd3337fac 100644 --- a/docker/postgres/init.sql +++ b/docker/postgres/init.sql @@ -28,3 +28,5 @@ insert into metadata_aspect_v2 (urn, aspect, version, metadata, createdon, creat now(), 'urn:li:corpuser:__datahub_system' ); + +DROP TABLE IF EXISTS metadata_index; From 1a72fa499c3404c6c3d2961e9575495f2dd021d2 Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Tue, 10 Oct 2023 17:34:06 -0400 Subject: [PATCH 109/156] feat(ingest/tableau): Allow parsing of database name from fullName (#8981) --- .../src/datahub/ingestion/source/tableau.py | 74 ++------ .../ingestion/source/tableau_common.py | 162 +++++++++++++----- .../tableau/test_tableau_ingest.py | 34 ++-- 3 files changed, 151 insertions(+), 119 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/tableau.py b/metadata-ingestion/src/datahub/ingestion/source/tableau.py index e347cd26d245a..bad7ae49d325e 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/tableau.py +++ b/metadata-ingestion/src/datahub/ingestion/source/tableau.py @@ -77,6 +77,7 @@ FIELD_TYPE_MAPPING, MetadataQueryException, TableauLineageOverrides, + TableauUpstreamReference, clean_query, custom_sql_graphql_query, dashboard_graphql_query, @@ -85,7 +86,6 @@ get_overridden_info, get_unique_custom_sql, make_fine_grained_lineage_class, - make_table_urn, make_upstream_class, published_datasource_graphql_query, query_metadata, @@ -271,7 +271,7 @@ class TableauConfig( "You can change this if your Tableau projects contain slashes in their names, and you'd like to filter by project.", ) - default_schema_map: dict = Field( + default_schema_map: Dict[str, str] = Field( default={}, description="Default schema to use when schema is not found." ) ingest_tags: Optional[bool] = Field( @@ -997,41 +997,16 @@ def get_upstream_tables( ) continue - schema = table.get(tableau_constant.SCHEMA) or "" - table_name = table.get(tableau_constant.NAME) or "" - full_name = table.get(tableau_constant.FULL_NAME) or "" - upstream_db = ( - table[tableau_constant.DATABASE][tableau_constant.NAME] - if table.get(tableau_constant.DATABASE) - and table[tableau_constant.DATABASE].get(tableau_constant.NAME) - else "" - ) - logger.debug( - "Processing Table with Connection Type: {0} and id {1}".format( - table.get(tableau_constant.CONNECTION_TYPE) or "", - table.get(tableau_constant.ID) or "", + try: + ref = TableauUpstreamReference.create( + table, default_schema_map=self.config.default_schema_map ) - ) - schema = self._get_schema(schema, upstream_db, full_name) - # if the schema is included within the table name we omit it - if ( - schema - and table_name - and full_name - and table_name == full_name - and schema in table_name - ): - logger.debug( - f"Omitting schema for upstream table {table[tableau_constant.ID]}, schema included in table name" - ) - schema = "" + except Exception as e: + logger.info(f"Failed to generate upstream reference for {table}: {e}") + continue - table_urn = make_table_urn( + table_urn = ref.make_dataset_urn( self.config.env, - upstream_db, - table.get(tableau_constant.CONNECTION_TYPE) or "", - schema, - table_name, self.config.platform_instance_map, self.config.lineage_overrides, ) @@ -1052,7 +1027,7 @@ def get_upstream_tables( urn=table_urn, id=table[tableau_constant.ID], num_cols=num_tbl_cols, - paths=set([table_path]) if table_path else set(), + paths={table_path} if table_path else set(), ) else: self.database_tables[table_urn].update_table( @@ -2462,35 +2437,6 @@ def emit_embedded_datasources(self) -> Iterable[MetadataWorkUnit]: is_embedded_ds=True, ) - @lru_cache(maxsize=None) - def _get_schema(self, schema_provided: str, database: str, fullName: str) -> str: - # For some databases, the schema attribute in tableau api does not return - # correct schema name for the table. For more information, see - # https://help.tableau.com/current/api/metadata_api/en-us/docs/meta_api_model.html#schema_attribute. - # Hence we extract schema from fullName whenever fullName is available - schema = self._extract_schema_from_fullName(fullName) if fullName else "" - if not schema: - schema = schema_provided - elif schema != schema_provided: - logger.debug( - "Correcting schema, provided {0}, corrected {1}".format( - schema_provided, schema - ) - ) - - if not schema and database in self.config.default_schema_map: - schema = self.config.default_schema_map[database] - - return schema - - @lru_cache(maxsize=None) - def _extract_schema_from_fullName(self, fullName: str) -> str: - # fullName is observed to be in format [schemaName].[tableName] - # OR simply tableName OR [tableName] - if fullName.startswith("[") and "].[" in fullName: - return fullName[1 : fullName.index("]")] - return "" - @lru_cache(maxsize=None) def get_last_modified( self, creator: Optional[str], created_at: bytes, updated_at: bytes diff --git a/metadata-ingestion/src/datahub/ingestion/source/tableau_common.py b/metadata-ingestion/src/datahub/ingestion/source/tableau_common.py index 2c92285fdba77..7c4852042ce7c 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/tableau_common.py +++ b/metadata-ingestion/src/datahub/ingestion/source/tableau_common.py @@ -1,4 +1,6 @@ import html +import logging +from dataclasses import dataclass from functools import lru_cache from typing import Dict, List, Optional, Tuple @@ -6,6 +8,7 @@ import datahub.emitter.mce_builder as builder from datahub.configuration.common import ConfigModel +from datahub.ingestion.source import tableau_constant as tc from datahub.metadata.com.linkedin.pegasus2avro.dataset import ( DatasetLineageType, FineGrainedLineage, @@ -31,6 +34,8 @@ ) from datahub.utilities.sqlglot_lineage import ColumnLineageInfo, SqlParsingResult +logger = logging.getLogger(__name__) + class TableauLineageOverrides(ConfigModel): platform_override_map: Optional[Dict[str, str]] = Field( @@ -537,12 +542,12 @@ def get_fully_qualified_table_name( platform: str, upstream_db: str, schema: str, - full_name: str, + table_name: str, ) -> str: if platform == "athena": upstream_db = "" database_name = f"{upstream_db}." if upstream_db else "" - final_name = full_name.replace("[", "").replace("]", "") + final_name = table_name.replace("[", "").replace("]", "") schema_name = f"{schema}." if schema else "" @@ -573,17 +578,123 @@ def get_fully_qualified_table_name( return fully_qualified_table_name -def get_platform_instance( - platform: str, platform_instance_map: Optional[Dict[str, str]] -) -> Optional[str]: - if platform_instance_map is not None and platform in platform_instance_map.keys(): - return platform_instance_map[platform] +@dataclass +class TableauUpstreamReference: + database: Optional[str] + schema: Optional[str] + table: str + + connection_type: str + + @classmethod + def create( + cls, d: dict, default_schema_map: Optional[Dict[str, str]] = None + ) -> "TableauUpstreamReference": + # Values directly from `table` object from Tableau + database = t_database = d.get(tc.DATABASE, {}).get(tc.NAME) + schema = t_schema = d.get(tc.SCHEMA) + table = t_table = d.get(tc.NAME) or "" + t_full_name = d.get(tc.FULL_NAME) + t_connection_type = d[tc.CONNECTION_TYPE] # required to generate urn + t_id = d[tc.ID] + + parsed_full_name = cls.parse_full_name(t_full_name) + if parsed_full_name and len(parsed_full_name) == 3: + database, schema, table = parsed_full_name + elif parsed_full_name and len(parsed_full_name) == 2: + schema, table = parsed_full_name + else: + logger.debug( + f"Upstream urn generation ({t_id}):" + f" Did not parse full name {t_full_name}: unexpected number of values", + ) + + if not schema and default_schema_map and database in default_schema_map: + schema = default_schema_map[database] + + if database != t_database: + logger.debug( + f"Upstream urn generation ({t_id}):" + f" replacing database {t_database} with {database} from full name {t_full_name}" + ) + if schema != t_schema: + logger.debug( + f"Upstream urn generation ({t_id}):" + f" replacing schema {t_schema} with {schema} from full name {t_full_name}" + ) + if table != t_table: + logger.debug( + f"Upstream urn generation ({t_id}):" + f" replacing table {t_table} with {table} from full name {t_full_name}" + ) + + # TODO: See if we can remove this -- made for redshift + if ( + schema + and t_table + and t_full_name + and t_table == t_full_name + and schema in t_table + ): + logger.debug( + f"Omitting schema for upstream table {t_id}, schema included in table name" + ) + schema = "" + + return cls( + database=database, + schema=schema, + table=table, + connection_type=t_connection_type, + ) + + @staticmethod + def parse_full_name(full_name: Optional[str]) -> Optional[List[str]]: + # fullName is observed to be in formats: + # [database].[schema].[table] + # [schema].[table] + # [table] + # table + # schema + + # TODO: Validate the startswith check. Currently required for our integration tests + if full_name is None or not full_name.startswith("["): + return None + + return full_name.replace("[", "").replace("]", "").split(".") + + def make_dataset_urn( + self, + env: str, + platform_instance_map: Optional[Dict[str, str]], + lineage_overrides: Optional[TableauLineageOverrides] = None, + ) -> str: + ( + upstream_db, + platform_instance, + platform, + original_platform, + ) = get_overridden_info( + connection_type=self.connection_type, + upstream_db=self.database, + lineage_overrides=lineage_overrides, + platform_instance_map=platform_instance_map, + ) + + table_name = get_fully_qualified_table_name( + original_platform, + upstream_db or "", + self.schema, + self.table, + ) - return None + return builder.make_dataset_urn_with_platform_instance( + platform, table_name, platform_instance, env + ) def get_overridden_info( - connection_type: str, + connection_type: Optional[str], upstream_db: Optional[str], platform_instance_map: Optional[Dict[str, str]], lineage_overrides: Optional[TableauLineageOverrides] = None, @@ -605,7 +716,9 @@ def get_overridden_info( ): upstream_db = lineage_overrides.database_override_map[upstream_db] - platform_instance = get_platform_instance(original_platform, platform_instance_map) + platform_instance = ( + platform_instance_map.get(original_platform) if platform_instance_map else None + ) if original_platform in ("athena", "hive", "mysql"): # Two tier databases upstream_db = None @@ -613,35 +726,6 @@ def get_overridden_info( return upstream_db, platform_instance, platform, original_platform -def make_table_urn( - env: str, - upstream_db: Optional[str], - connection_type: str, - schema: str, - full_name: str, - platform_instance_map: Optional[Dict[str, str]], - lineage_overrides: Optional[TableauLineageOverrides] = None, -) -> str: - - upstream_db, platform_instance, platform, original_platform = get_overridden_info( - connection_type=connection_type, - upstream_db=upstream_db, - lineage_overrides=lineage_overrides, - platform_instance_map=platform_instance_map, - ) - - table_name = get_fully_qualified_table_name( - original_platform, - upstream_db if upstream_db is not None else "", - schema, - full_name, - ) - - return builder.make_dataset_urn_with_platform_instance( - platform, table_name, platform_instance, env - ) - - def make_description_from_params(description, formula): """ Generate column description diff --git a/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py b/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py index c31867f5aa904..0510f4a40f659 100644 --- a/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py +++ b/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py @@ -20,7 +20,7 @@ from datahub.ingestion.source.tableau import TableauConfig, TableauSource from datahub.ingestion.source.tableau_common import ( TableauLineageOverrides, - make_table_urn, + TableauUpstreamReference, ) from datahub.metadata.com.linkedin.pegasus2avro.dataset import ( DatasetLineageType, @@ -546,13 +546,13 @@ def test_lineage_overrides(): enable_logging() # Simple - specify platform instance to presto table assert ( - make_table_urn( - DEFAULT_ENV, + TableauUpstreamReference( "presto_catalog", - "presto", "test-schema", - "presto_catalog.test-schema.test-table", - platform_instance_map={"presto": "my_presto_instance"}, + "test-table", + "presto", + ).make_dataset_urn( + env=DEFAULT_ENV, platform_instance_map={"presto": "my_presto_instance"} ) == "urn:li:dataset:(urn:li:dataPlatform:presto,my_presto_instance.presto_catalog.test-schema.test-table,PROD)" ) @@ -560,12 +560,13 @@ def test_lineage_overrides(): # Transform presto urn to hive urn # resulting platform instance for hive = mapped platform instance + presto_catalog assert ( - make_table_urn( - DEFAULT_ENV, + TableauUpstreamReference( "presto_catalog", - "presto", "test-schema", - "presto_catalog.test-schema.test-table", + "test-table", + "presto", + ).make_dataset_urn( + env=DEFAULT_ENV, platform_instance_map={"presto": "my_instance"}, lineage_overrides=TableauLineageOverrides( platform_override_map={"presto": "hive"}, @@ -574,14 +575,15 @@ def test_lineage_overrides(): == "urn:li:dataset:(urn:li:dataPlatform:hive,my_instance.presto_catalog.test-schema.test-table,PROD)" ) - # tranform hive urn to presto urn + # transform hive urn to presto urn assert ( - make_table_urn( - DEFAULT_ENV, - "", - "hive", + TableauUpstreamReference( + None, "test-schema", - "test-schema.test-table", + "test-table", + "hive", + ).make_dataset_urn( + env=DEFAULT_ENV, platform_instance_map={"hive": "my_presto_instance.presto_catalog"}, lineage_overrides=TableauLineageOverrides( platform_override_map={"hive": "presto"}, From e2988017c23270acd95e25ec3289983ecc3895f7 Mon Sep 17 00:00:00 2001 From: Amanda Hernando <110099762+amanda-her@users.noreply.github.com> Date: Wed, 11 Oct 2023 01:36:01 +0200 Subject: [PATCH 110/156] feat(auth): add data platform instance field resolver provider (#8828) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Sergio Gómez Villamor Co-authored-by: Adrián Pertíñez --- .../authorization/ResolvedResourceSpec.java | 17 ++ .../authorization/ResourceFieldType.java | 6 +- .../DefaultResourceSpecResolver.java | 9 +- ...PlatformInstanceFieldResolverProvider.java | 70 +++++++ ...formInstanceFieldResolverProviderTest.java | 188 ++++++++++++++++++ 5 files changed, 286 insertions(+), 4 deletions(-) create mode 100644 metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProvider.java create mode 100644 metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProviderTest.java diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedResourceSpec.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedResourceSpec.java index 53dd0be44f963..8e429a8ca1b94 100644 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedResourceSpec.java +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedResourceSpec.java @@ -3,6 +3,7 @@ import java.util.Collections; import java.util.Map; import java.util.Set; +import javax.annotation.Nullable; import lombok.Getter; import lombok.RequiredArgsConstructor; import lombok.ToString; @@ -35,4 +36,20 @@ public Set getOwners() { } return fieldResolvers.get(ResourceFieldType.OWNER).getFieldValuesFuture().join().getValues(); } + + /** + * Fetch the platform instance for a Resolved Resource Spec + * @return a Platform Instance or null if one does not exist. + */ + @Nullable + public String getDataPlatformInstance() { + if (!fieldResolvers.containsKey(ResourceFieldType.DATA_PLATFORM_INSTANCE)) { + return null; + } + Set dataPlatformInstance = fieldResolvers.get(ResourceFieldType.DATA_PLATFORM_INSTANCE).getFieldValuesFuture().join().getValues(); + if (dataPlatformInstance.size() > 0) { + return dataPlatformInstance.stream().findFirst().get(); + } + return null; + } } diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceFieldType.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceFieldType.java index ee54d2bfbba1d..478522dc7c331 100644 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceFieldType.java +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceFieldType.java @@ -19,5 +19,9 @@ public enum ResourceFieldType { /** * Domains of resource */ - DOMAIN + DOMAIN, + /** + * Data platform instance of resource + */ + DATA_PLATFORM_INSTANCE } diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultResourceSpecResolver.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultResourceSpecResolver.java index cd4e0b0967829..64c43dc8aa591 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultResourceSpecResolver.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultResourceSpecResolver.java @@ -1,13 +1,15 @@ package com.datahub.authorization; -import com.datahub.authorization.fieldresolverprovider.EntityTypeFieldResolverProvider; -import com.datahub.authorization.fieldresolverprovider.OwnerFieldResolverProvider; import com.datahub.authentication.Authentication; +import com.datahub.authorization.fieldresolverprovider.DataPlatformInstanceFieldResolverProvider; import com.datahub.authorization.fieldresolverprovider.DomainFieldResolverProvider; +import com.datahub.authorization.fieldresolverprovider.EntityTypeFieldResolverProvider; import com.datahub.authorization.fieldresolverprovider.EntityUrnFieldResolverProvider; +import com.datahub.authorization.fieldresolverprovider.OwnerFieldResolverProvider; import com.datahub.authorization.fieldresolverprovider.ResourceFieldResolverProvider; import com.google.common.collect.ImmutableList; import com.linkedin.entity.client.EntityClient; + import java.util.List; import java.util.Map; import java.util.stream.Collectors; @@ -20,7 +22,8 @@ public DefaultResourceSpecResolver(Authentication systemAuthentication, EntityCl _resourceFieldResolverProviders = ImmutableList.of(new EntityTypeFieldResolverProvider(), new EntityUrnFieldResolverProvider(), new DomainFieldResolverProvider(entityClient, systemAuthentication), - new OwnerFieldResolverProvider(entityClient, systemAuthentication)); + new OwnerFieldResolverProvider(entityClient, systemAuthentication), + new DataPlatformInstanceFieldResolverProvider(entityClient, systemAuthentication)); } @Override diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProvider.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProvider.java new file mode 100644 index 0000000000000..cd838625c2ca1 --- /dev/null +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProvider.java @@ -0,0 +1,70 @@ +package com.datahub.authorization.fieldresolverprovider; + +import com.datahub.authentication.Authentication; +import com.datahub.authorization.FieldResolver; +import com.datahub.authorization.ResourceFieldType; +import com.datahub.authorization.ResourceSpec; +import com.linkedin.common.DataPlatformInstance; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.entity.client.EntityClient; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +import java.util.Collections; +import java.util.Objects; + +import static com.linkedin.metadata.Constants.*; + +/** + * Provides field resolver for domain given resourceSpec + */ +@Slf4j +@RequiredArgsConstructor +public class DataPlatformInstanceFieldResolverProvider implements ResourceFieldResolverProvider { + + private final EntityClient _entityClient; + private final Authentication _systemAuthentication; + + @Override + public ResourceFieldType getFieldType() { + return ResourceFieldType.DATA_PLATFORM_INSTANCE; + } + + @Override + public FieldResolver getFieldResolver(ResourceSpec resourceSpec) { + return FieldResolver.getResolverFromFunction(resourceSpec, this::getDataPlatformInstance); + } + + private FieldResolver.FieldValue getDataPlatformInstance(ResourceSpec resourceSpec) { + Urn entityUrn = UrnUtils.getUrn(resourceSpec.getResource()); + // In the case that the entity is a platform instance, the associated platform instance entity is the instance itself + if (entityUrn.getEntityType().equals(DATA_PLATFORM_INSTANCE_ENTITY_NAME)) { + return FieldResolver.FieldValue.builder() + .values(Collections.singleton(entityUrn.toString())) + .build(); + } + + EnvelopedAspect dataPlatformInstanceAspect; + try { + EntityResponse response = _entityClient.getV2(entityUrn.getEntityType(), entityUrn, + Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME), _systemAuthentication); + if (response == null || !response.getAspects().containsKey(DATA_PLATFORM_INSTANCE_ASPECT_NAME)) { + return FieldResolver.emptyFieldValue(); + } + dataPlatformInstanceAspect = response.getAspects().get(DATA_PLATFORM_INSTANCE_ASPECT_NAME); + } catch (Exception e) { + log.error("Error while retrieving platform instance aspect for urn {}", entityUrn, e); + return FieldResolver.emptyFieldValue(); + } + DataPlatformInstance dataPlatformInstance = new DataPlatformInstance(dataPlatformInstanceAspect.getValue().data()); + if (dataPlatformInstance.getInstance() == null) { + return FieldResolver.emptyFieldValue(); + } + return FieldResolver.FieldValue.builder() + .values(Collections.singleton(Objects.requireNonNull(dataPlatformInstance.getInstance()).toString())) + .build(); + } +} \ No newline at end of file diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProviderTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProviderTest.java new file mode 100644 index 0000000000000..e525c602c2620 --- /dev/null +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProviderTest.java @@ -0,0 +1,188 @@ +package com.datahub.authorization.fieldresolverprovider; + +import com.datahub.authentication.Authentication; +import com.datahub.authorization.ResourceFieldType; +import com.datahub.authorization.ResourceSpec; +import com.linkedin.common.DataPlatformInstance; +import com.linkedin.common.urn.Urn; +import com.linkedin.entity.Aspect; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.entity.EnvelopedAspectMap; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.r2.RemoteInvocationException; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import java.net.URISyntaxException; +import java.util.Collections; +import java.util.Set; + +import static com.linkedin.metadata.Constants.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.*; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +public class DataPlatformInstanceFieldResolverProviderTest { + + private static final String DATA_PLATFORM_INSTANCE_URN = + "urn:li:dataPlatformInstance:(urn:li:dataPlatform:s3,test-platform-instance)"; + private static final String RESOURCE_URN = + "urn:li:dataset:(urn:li:dataPlatform:s3,test-platform-instance.testDataset,PROD)"; + private static final ResourceSpec RESOURCE_SPEC = new ResourceSpec(DATASET_ENTITY_NAME, RESOURCE_URN); + + @Mock + private EntityClient entityClientMock; + @Mock + private Authentication systemAuthenticationMock; + + private DataPlatformInstanceFieldResolverProvider dataPlatformInstanceFieldResolverProvider; + + @BeforeMethod + public void setup() { + MockitoAnnotations.initMocks(this); + dataPlatformInstanceFieldResolverProvider = + new DataPlatformInstanceFieldResolverProvider(entityClientMock, systemAuthenticationMock); + } + + @Test + public void shouldReturnDataPlatformInstanceType() { + assertEquals(ResourceFieldType.DATA_PLATFORM_INSTANCE, dataPlatformInstanceFieldResolverProvider.getFieldType()); + } + + @Test + public void shouldReturnFieldValueWithResourceSpecIfTypeIsDataPlatformInstance() { + var resourceSpec = new ResourceSpec(DATA_PLATFORM_INSTANCE_ENTITY_NAME, DATA_PLATFORM_INSTANCE_URN); + + var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(resourceSpec); + + assertEquals(Set.of(DATA_PLATFORM_INSTANCE_URN), result.getFieldValuesFuture().join().getValues()); + verifyZeroInteractions(entityClientMock); + } + + @Test + public void shouldReturnEmptyFieldValueWhenResponseIsNull() throws RemoteInvocationException, URISyntaxException { + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenReturn(null); + + var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } + + @Test + public void shouldReturnEmptyFieldValueWhenResourceHasNoDataPlatformInstance() + throws RemoteInvocationException, URISyntaxException { + var entityResponseMock = mock(EntityResponse.class); + when(entityResponseMock.getAspects()).thenReturn(new EnvelopedAspectMap()); + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenReturn(entityResponseMock); + + var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } + + @Test + public void shouldReturnEmptyFieldValueWhenThereIsAnException() throws RemoteInvocationException, URISyntaxException { + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenThrow(new RemoteInvocationException()); + + var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } + + @Test + public void shouldReturnEmptyFieldValueWhenDataPlatformInstanceHasNoInstance() + throws RemoteInvocationException, URISyntaxException { + + var dataPlatform = new DataPlatformInstance() + .setPlatform(Urn.createFromString("urn:li:dataPlatform:s3")); + var entityResponseMock = mock(EntityResponse.class); + var envelopedAspectMap = new EnvelopedAspectMap(); + envelopedAspectMap.put(DATA_PLATFORM_INSTANCE_ASPECT_NAME, + new EnvelopedAspect().setValue(new Aspect(dataPlatform.data()))); + when(entityResponseMock.getAspects()).thenReturn(envelopedAspectMap); + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenReturn(entityResponseMock); + + var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } + + @Test + public void shouldReturnFieldValueWithDataPlatformInstanceOfTheResource() + throws RemoteInvocationException, URISyntaxException { + + var dataPlatformInstance = new DataPlatformInstance() + .setPlatform(Urn.createFromString("urn:li:dataPlatform:s3")) + .setInstance(Urn.createFromString(DATA_PLATFORM_INSTANCE_URN)); + var entityResponseMock = mock(EntityResponse.class); + var envelopedAspectMap = new EnvelopedAspectMap(); + envelopedAspectMap.put(DATA_PLATFORM_INSTANCE_ASPECT_NAME, + new EnvelopedAspect().setValue(new Aspect(dataPlatformInstance.data()))); + when(entityResponseMock.getAspects()).thenReturn(envelopedAspectMap); + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenReturn(entityResponseMock); + + var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertEquals(Set.of(DATA_PLATFORM_INSTANCE_URN), result.getFieldValuesFuture().join().getValues()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } +} From a17db676e37d90ec47f16a43ab95e0d562952939 Mon Sep 17 00:00:00 2001 From: siladitya <68184387+siladitya2@users.noreply.github.com> Date: Wed, 11 Oct 2023 02:43:36 +0200 Subject: [PATCH 111/156] feat(graphql): Added datafetcher for DataPlatformInstance entity (#8935) Co-authored-by: si-chakraborty Co-authored-by: John Joyce --- .../datahub/graphql/GmsGraphQLEngine.java | 1 + .../DataPlatformInstanceType.java | 34 ++++++++++++++++++- .../src/main/resources/entity.graphql | 5 +++ 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java index 3ba0cc1f747e3..ebb5c7d62c7d3 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java @@ -821,6 +821,7 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) { .dataFetcher("glossaryNode", getResolver(glossaryNodeType)) .dataFetcher("domain", getResolver((domainType))) .dataFetcher("dataPlatform", getResolver(dataPlatformType)) + .dataFetcher("dataPlatformInstance", getResolver(dataPlatformInstanceType)) .dataFetcher("mlFeatureTable", getResolver(mlFeatureTableType)) .dataFetcher("mlFeature", getResolver(mlFeatureType)) .dataFetcher("mlPrimaryKey", getResolver(mlPrimaryKeyType)) diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatforminstance/DataPlatformInstanceType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatforminstance/DataPlatformInstanceType.java index 2423fc31ea52e..87614e1332528 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatforminstance/DataPlatformInstanceType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatforminstance/DataPlatformInstanceType.java @@ -4,16 +4,25 @@ import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.generated.AutoCompleteResults; import com.linkedin.datahub.graphql.generated.DataPlatformInstance; import com.linkedin.datahub.graphql.generated.Entity; import com.linkedin.datahub.graphql.generated.EntityType; +import com.linkedin.datahub.graphql.generated.FacetFilterInput; +import com.linkedin.datahub.graphql.generated.SearchResults; import com.linkedin.datahub.graphql.types.dataplatforminstance.mappers.DataPlatformInstanceMapper; +import com.linkedin.datahub.graphql.types.mappers.AutoCompleteResultsMapper; +import com.linkedin.datahub.graphql.types.SearchableEntityType; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; +import com.linkedin.metadata.query.AutoCompleteResult; +import com.linkedin.metadata.query.filter.Filter; import graphql.execution.DataFetcherResult; +import org.apache.commons.lang3.NotImplementedException; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.util.ArrayList; import java.util.HashSet; import java.util.List; @@ -22,7 +31,10 @@ import java.util.function.Function; import java.util.stream.Collectors; -public class DataPlatformInstanceType implements com.linkedin.datahub.graphql.types.EntityType { +import static com.linkedin.metadata.Constants.DATA_PLATFORM_INSTANCE_ENTITY_NAME; + +public class DataPlatformInstanceType implements SearchableEntityType, + com.linkedin.datahub.graphql.types.EntityType { static final Set ASPECTS_TO_FETCH = ImmutableSet.of( Constants.DATA_PLATFORM_INSTANCE_KEY_ASPECT_NAME, @@ -84,4 +96,24 @@ public List> batchLoad(@Nonnull List filters, + int start, + int count, + @Nonnull final QueryContext context) throws Exception { + throw new NotImplementedException("Searchable type (deprecated) not implemented on DataPlatformInstance entity type"); + } + + @Override + public AutoCompleteResults autoComplete(@Nonnull String query, + @Nullable String field, + @Nullable Filter filters, + int limit, + @Nonnull final QueryContext context) throws Exception { + final AutoCompleteResult result = _entityClient.autoComplete(DATA_PLATFORM_INSTANCE_ENTITY_NAME, query, + filters, limit, context.getAuthentication()); + return AutoCompleteResultsMapper.map(result); + } + } diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql index 39f86948c77c4..0b15d7b875a9c 100644 --- a/datahub-graphql-core/src/main/resources/entity.graphql +++ b/datahub-graphql-core/src/main/resources/entity.graphql @@ -226,6 +226,11 @@ type Query { listOwnershipTypes( "Input required for listing custom ownership types" input: ListOwnershipTypesInput!): ListOwnershipTypesResult! + + """ + Fetch a Data Platform Instance by primary key (urn) + """ + dataPlatformInstance(urn: String!): DataPlatformInstance } """ From dfcea2441e75e1eef517c0f9a4765e6e7990f297 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20G=C3=B3mez=20Villamor?= Date: Wed, 11 Oct 2023 03:04:44 +0200 Subject: [PATCH 112/156] feat(config): configurable bootstrap policies file (#8812) Co-authored-by: John Joyce --- .../configuration/src/main/resources/application.yml | 4 ++++ .../boot/factories/BootstrapManagerFactory.java | 7 ++++++- .../metadata/boot/steps/IngestPoliciesStep.java | 10 +++++++--- 3 files changed, 17 insertions(+), 4 deletions(-) diff --git a/metadata-service/configuration/src/main/resources/application.yml b/metadata-service/configuration/src/main/resources/application.yml index 4dfd96ac75c6c..d22f92adca8f9 100644 --- a/metadata-service/configuration/src/main/resources/application.yml +++ b/metadata-service/configuration/src/main/resources/application.yml @@ -276,6 +276,10 @@ bootstrap: enabled: ${UPGRADE_DEFAULT_BROWSE_PATHS_ENABLED:false} # enable to run the upgrade to migrate legacy default browse paths to new ones backfillBrowsePathsV2: enabled: ${BACKFILL_BROWSE_PATHS_V2:false} # Enables running the backfill of browsePathsV2 upgrade step. There are concerns about the load of this step so hiding it behind a flag. Deprecating in favor of running through SystemUpdate + policies: + file: ${BOOTSTRAP_POLICIES_FILE:classpath:boot/policies.json} + # eg for local file + # file: "file:///datahub/datahub-gms/resources/custom-policies.json" servlets: waitTimeout: ${BOOTSTRAP_SERVLETS_WAITTIMEOUT:60} # Total waiting time in seconds for servlets to initialize diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/BootstrapManagerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/BootstrapManagerFactory.java index c490f00021201..3a761bd12647e 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/BootstrapManagerFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/BootstrapManagerFactory.java @@ -31,6 +31,7 @@ import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.SearchService; import com.linkedin.metadata.search.transformer.SearchDocumentTransformer; + import java.util.ArrayList; import java.util.List; import javax.annotation.Nonnull; @@ -41,6 +42,7 @@ import org.springframework.context.annotation.Configuration; import org.springframework.context.annotation.Import; import org.springframework.context.annotation.Scope; +import org.springframework.core.io.Resource; @Configuration @@ -89,13 +91,16 @@ public class BootstrapManagerFactory { @Value("${bootstrap.backfillBrowsePathsV2.enabled}") private Boolean _backfillBrowsePathsV2Enabled; + @Value("${bootstrap.policies.file}") + private Resource _policiesResource; + @Bean(name = "bootstrapManager") @Scope("singleton") @Nonnull protected BootstrapManager createInstance() { final IngestRootUserStep ingestRootUserStep = new IngestRootUserStep(_entityService); final IngestPoliciesStep ingestPoliciesStep = - new IngestPoliciesStep(_entityRegistry, _entityService, _entitySearchService, _searchDocumentTransformer); + new IngestPoliciesStep(_entityRegistry, _entityService, _entitySearchService, _searchDocumentTransformer, _policiesResource); final IngestRolesStep ingestRolesStep = new IngestRolesStep(_entityService, _entityRegistry); final IngestDataPlatformsStep ingestDataPlatformsStep = new IngestDataPlatformsStep(_entityService); final IngestDataPlatformInstancesStep ingestDataPlatformInstancesStep = diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java index 87dcfd736da40..cf29645214466 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java @@ -25,6 +25,7 @@ import com.linkedin.mxe.GenericAspect; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.policy.DataHubPolicyInfo; + import java.io.IOException; import java.net.URISyntaxException; import java.util.Collections; @@ -35,7 +36,8 @@ import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; -import org.springframework.core.io.ClassPathResource; +import org.springframework.core.io.Resource; + import static com.linkedin.metadata.Constants.*; @@ -52,6 +54,8 @@ public class IngestPoliciesStep implements BootstrapStep { private final EntitySearchService _entitySearchService; private final SearchDocumentTransformer _searchDocumentTransformer; + private final Resource _policiesResource; + @Override public String name() { return "IngestPoliciesStep"; @@ -66,10 +70,10 @@ public void execute() throws IOException, URISyntaxException { .maxStringLength(maxSize).build()); // 0. Execute preflight check to see whether we need to ingest policies - log.info("Ingesting default access policies..."); + log.info("Ingesting default access policies from: {}...", _policiesResource); // 1. Read from the file into JSON. - final JsonNode policiesObj = mapper.readTree(new ClassPathResource("./boot/policies.json").getFile()); + final JsonNode policiesObj = mapper.readTree(_policiesResource.getFile()); if (!policiesObj.isArray()) { throw new RuntimeException( From 10a190470e8c932b6d34cba49de7dbcba687a088 Mon Sep 17 00:00:00 2001 From: siddiquebagwan-gslab Date: Wed, 11 Oct 2023 08:54:08 +0530 Subject: [PATCH 113/156] feat(ingestion/redshift): CLL support in redshift (#8921) --- .../ingestion/source/redshift/config.py | 4 + .../ingestion/source/redshift/lineage.py | 215 +++++++++++++----- .../ingestion/source/redshift/redshift.py | 1 + .../tests/unit/test_redshift_lineage.py | 95 ++++++-- 4 files changed, 234 insertions(+), 81 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py index 804a14b0fe1cf..2789b800940db 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py @@ -132,6 +132,10 @@ class RedshiftConfig( description="Whether `schema_pattern` is matched against fully qualified schema name `.`.", ) + extract_column_level_lineage: bool = Field( + default=True, description="Whether to extract column level lineage." + ) + @root_validator(pre=True) def check_email_is_set_on_usage(cls, values): if values.get("include_usage_statistics"): diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/lineage.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/lineage.py index bbe52b5d98ba3..c9ddfbe92ab2a 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/redshift/lineage.py +++ b/metadata-ingestion/src/datahub/ingestion/source/redshift/lineage.py @@ -9,10 +9,12 @@ import humanfriendly import redshift_connector -from sqllineage.runner import LineageRunner +import datahub.emitter.mce_builder as builder +import datahub.utilities.sqlglot_lineage as sqlglot_l from datahub.emitter import mce_builder from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance +from datahub.ingestion.api.common import PipelineContext from datahub.ingestion.source.aws.s3_util import strip_s3_prefix from datahub.ingestion.source.redshift.common import get_db_name from datahub.ingestion.source.redshift.config import LineageMode, RedshiftConfig @@ -28,13 +30,19 @@ from datahub.ingestion.source.state.redundant_run_skip_handler import ( RedundantLineageRunSkipHandler, ) -from datahub.metadata.com.linkedin.pegasus2avro.dataset import UpstreamLineage +from datahub.metadata.com.linkedin.pegasus2avro.dataset import ( + FineGrainedLineage, + FineGrainedLineageDownstreamType, + FineGrainedLineageUpstreamType, + UpstreamLineage, +) from datahub.metadata.schema_classes import ( DatasetLineageTypeClass, UpstreamClass, UpstreamLineageClass, ) from datahub.utilities import memory_footprint +from datahub.utilities.urns import dataset_urn logger: logging.Logger = logging.getLogger(__name__) @@ -56,13 +64,14 @@ class LineageCollectorType(Enum): @dataclass(frozen=True, eq=True) class LineageDataset: platform: LineageDatasetPlatform - path: str + urn: str @dataclass() class LineageItem: dataset: LineageDataset upstreams: Set[LineageDataset] + cll: Optional[List[sqlglot_l.ColumnLineageInfo]] collector_type: LineageCollectorType dataset_lineage_type: str = field(init=False) @@ -83,10 +92,12 @@ def __init__( self, config: RedshiftConfig, report: RedshiftReport, + context: PipelineContext, redundant_run_skip_handler: Optional[RedundantLineageRunSkipHandler] = None, ): self.config = config self.report = report + self.context = context self._lineage_map: Dict[str, LineageItem] = defaultdict() self.redundant_run_skip_handler = redundant_run_skip_handler @@ -121,33 +132,37 @@ def _get_s3_path(self, path: str) -> str: return path - def _get_sources_from_query(self, db_name: str, query: str) -> List[LineageDataset]: + def _get_sources_from_query( + self, db_name: str, query: str + ) -> Tuple[List[LineageDataset], Optional[List[sqlglot_l.ColumnLineageInfo]]]: sources: List[LineageDataset] = list() - parser = LineageRunner(query) + parsed_result: Optional[ + sqlglot_l.SqlParsingResult + ] = sqlglot_l.create_lineage_sql_parsed_result( + query=query, + platform=LineageDatasetPlatform.REDSHIFT.value, + platform_instance=self.config.platform_instance, + database=db_name, + schema=str(self.config.default_schema), + graph=self.context.graph, + env=self.config.env, + ) - for table in parser.source_tables: - split = str(table).split(".") - if len(split) == 3: - db_name, source_schema, source_table = split - elif len(split) == 2: - source_schema, source_table = split - else: - raise ValueError( - f"Invalid table name {table} in query {query}. " - f"Expected format: [db_name].[schema].[table] or [schema].[table] or [table]." - ) + if parsed_result is None: + logger.debug(f"native query parsing failed for {query}") + return sources, None - if source_schema == "": - source_schema = str(self.config.default_schema) + logger.debug(f"parsed_result = {parsed_result}") + for table_urn in parsed_result.in_tables: source = LineageDataset( platform=LineageDatasetPlatform.REDSHIFT, - path=f"{db_name}.{source_schema}.{source_table}", + urn=table_urn, ) sources.append(source) - return sources + return sources, parsed_result.column_lineage def _build_s3_path_from_row(self, filename: str) -> str: path = filename.strip() @@ -165,9 +180,11 @@ def _get_sources( source_table: Optional[str], ddl: Optional[str], filename: Optional[str], - ) -> List[LineageDataset]: + ) -> Tuple[List[LineageDataset], Optional[List[sqlglot_l.ColumnLineageInfo]]]: sources: List[LineageDataset] = list() # Source + cll: Optional[List[sqlglot_l.ColumnLineageInfo]] = None + if ( lineage_type in { @@ -177,7 +194,7 @@ def _get_sources( and ddl is not None ): try: - sources = self._get_sources_from_query(db_name=db_name, query=ddl) + sources, cll = self._get_sources_from_query(db_name=db_name, query=ddl) except Exception as e: logger.warning( f"Error parsing query {ddl} for getting lineage. Error was {e}." @@ -192,22 +209,38 @@ def _get_sources( "Only s3 source supported with copy. The source was: {path}." ) self.report.num_lineage_dropped_not_support_copy_path += 1 - return sources + return sources, cll path = strip_s3_prefix(self._get_s3_path(path)) + urn = make_dataset_urn_with_platform_instance( + platform=platform.value, + name=path, + env=self.config.env, + platform_instance=self.config.platform_instance_map.get( + platform.value + ) + if self.config.platform_instance_map is not None + else None, + ) elif source_schema is not None and source_table is not None: platform = LineageDatasetPlatform.REDSHIFT path = f"{db_name}.{source_schema}.{source_table}" + urn = make_dataset_urn_with_platform_instance( + platform=platform.value, + platform_instance=self.config.platform_instance, + name=path, + env=self.config.env, + ) else: - return [] + return [], cll sources = [ LineageDataset( platform=platform, - path=path, + urn=urn, ) ] - return sources + return sources, cll def _populate_lineage_map( self, @@ -231,6 +264,7 @@ def _populate_lineage_map( :rtype: None """ try: + cll: Optional[List[sqlglot_l.ColumnLineageInfo]] = None raw_db_name = database alias_db_name = get_db_name(self.config) @@ -243,7 +277,7 @@ def _populate_lineage_map( if not target: continue - sources = self._get_sources( + sources, cll = self._get_sources( lineage_type, alias_db_name, source_schema=lineage_row.source_schema, @@ -251,6 +285,7 @@ def _populate_lineage_map( ddl=lineage_row.ddl, filename=lineage_row.filename, ) + target.cll = cll target.upstreams.update( self._get_upstream_lineages( @@ -262,20 +297,16 @@ def _populate_lineage_map( ) # Merging downstreams if dataset already exists and has downstreams - if target.dataset.path in self._lineage_map: - self._lineage_map[ - target.dataset.path - ].upstreams = self._lineage_map[ - target.dataset.path - ].upstreams.union( - target.upstreams - ) + if target.dataset.urn in self._lineage_map: + self._lineage_map[target.dataset.urn].upstreams = self._lineage_map[ + target.dataset.urn + ].upstreams.union(target.upstreams) else: - self._lineage_map[target.dataset.path] = target + self._lineage_map[target.dataset.urn] = target logger.debug( - f"Lineage[{target}]:{self._lineage_map[target.dataset.path]}" + f"Lineage[{target}]:{self._lineage_map[target.dataset.urn]}" ) except Exception as e: self.warn( @@ -308,17 +339,34 @@ def _get_target_lineage( target_platform = LineageDatasetPlatform.S3 # Following call requires 'filename' key in lineage_row target_path = self._build_s3_path_from_row(lineage_row.filename) + urn = make_dataset_urn_with_platform_instance( + platform=target_platform.value, + name=target_path, + env=self.config.env, + platform_instance=self.config.platform_instance_map.get( + target_platform.value + ) + if self.config.platform_instance_map is not None + else None, + ) except ValueError as e: self.warn(logger, "non-s3-lineage", str(e)) return None else: target_platform = LineageDatasetPlatform.REDSHIFT target_path = f"{alias_db_name}.{lineage_row.target_schema}.{lineage_row.target_table}" + urn = make_dataset_urn_with_platform_instance( + platform=target_platform.value, + platform_instance=self.config.platform_instance, + name=target_path, + env=self.config.env, + ) return LineageItem( - dataset=LineageDataset(platform=target_platform, path=target_path), + dataset=LineageDataset(platform=target_platform, urn=urn), upstreams=set(), collector_type=lineage_type, + cll=None, ) def _get_upstream_lineages( @@ -331,11 +379,22 @@ def _get_upstream_lineages( targe_source = [] for source in sources: if source.platform == LineageDatasetPlatform.REDSHIFT: - db, schema, table = source.path.split(".") + qualified_table_name = dataset_urn.DatasetUrn.create_from_string( + source.urn + ).get_entity_id()[1] + db, schema, table = qualified_table_name.split(".") if db == raw_db_name: db = alias_db_name path = f"{db}.{schema}.{table}" - source = LineageDataset(platform=source.platform, path=path) + source = LineageDataset( + platform=source.platform, + urn=make_dataset_urn_with_platform_instance( + platform=LineageDatasetPlatform.REDSHIFT.value, + platform_instance=self.config.platform_instance, + name=path, + env=self.config.env, + ), + ) # Filtering out tables which does not exist in Redshift # It was deleted in the meantime or query parser did not capture well the table name @@ -345,7 +404,7 @@ def _get_upstream_lineages( or not any(table == t.name for t in all_tables[db][schema]) ): logger.debug( - f"{source.path} missing table, dropping from lineage.", + f"{source.urn} missing table, dropping from lineage.", ) self.report.num_lineage_tables_dropped += 1 continue @@ -433,36 +492,73 @@ def populate_lineage( memory_footprint.total_size(self._lineage_map) ) + def make_fine_grained_lineage_class( + self, lineage_item: LineageItem, dataset_urn: str + ) -> List[FineGrainedLineage]: + fine_grained_lineages: List[FineGrainedLineage] = [] + + if ( + self.config.extract_column_level_lineage is False + or lineage_item.cll is None + ): + logger.debug("CLL extraction is disabled") + return fine_grained_lineages + + logger.debug("Extracting column level lineage") + + cll: List[sqlglot_l.ColumnLineageInfo] = lineage_item.cll + + for cll_info in cll: + downstream = ( + [builder.make_schema_field_urn(dataset_urn, cll_info.downstream.column)] + if cll_info.downstream is not None + and cll_info.downstream.column is not None + else [] + ) + + upstreams = [ + builder.make_schema_field_urn(column_ref.table, column_ref.column) + for column_ref in cll_info.upstreams + ] + + fine_grained_lineages.append( + FineGrainedLineage( + downstreamType=FineGrainedLineageDownstreamType.FIELD, + downstreams=downstream, + upstreamType=FineGrainedLineageUpstreamType.FIELD_SET, + upstreams=upstreams, + ) + ) + + logger.debug(f"Created fine_grained_lineage for {dataset_urn}") + + return fine_grained_lineages + def get_lineage( self, table: Union[RedshiftTable, RedshiftView], dataset_urn: str, schema: RedshiftSchema, ) -> Optional[Tuple[UpstreamLineageClass, Dict[str, str]]]: - dataset_key = mce_builder.dataset_urn_to_key(dataset_urn) - if dataset_key is None: - return None upstream_lineage: List[UpstreamClass] = [] - if dataset_key.name in self._lineage_map: - item = self._lineage_map[dataset_key.name] + cll_lineage: List[FineGrainedLineage] = [] + + if dataset_urn in self._lineage_map: + item = self._lineage_map[dataset_urn] for upstream in item.upstreams: upstream_table = UpstreamClass( - dataset=make_dataset_urn_with_platform_instance( - upstream.platform.value, - upstream.path, - platform_instance=self.config.platform_instance_map.get( - upstream.platform.value - ) - if self.config.platform_instance_map - else None, - env=self.config.env, - ), + dataset=upstream.urn, type=item.dataset_lineage_type, ) upstream_lineage.append(upstream_table) + cll_lineage = self.make_fine_grained_lineage_class( + lineage_item=item, + dataset_urn=dataset_urn, + ) + tablename = table.name if table.type == "EXTERNAL_TABLE": # external_db_params = schema.option @@ -489,7 +585,12 @@ def get_lineage( else: return None - return UpstreamLineage(upstreams=upstream_lineage), {} + return ( + UpstreamLineage( + upstreams=upstream_lineage, fineGrainedLineages=cll_lineage or None + ), + {}, + ) def report_status(self, step: str, status: bool) -> None: if self.redundant_run_skip_handler: diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/redshift.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/redshift.py index e8a8ff976afa6..a1b6333a3775d 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/redshift/redshift.py +++ b/metadata-ingestion/src/datahub/ingestion/source/redshift/redshift.py @@ -881,6 +881,7 @@ def extract_lineage( self.lineage_extractor = RedshiftLineageExtractor( config=self.config, report=self.report, + context=self.ctx, redundant_run_skip_handler=self.redundant_lineage_run_skip_handler, ) diff --git a/metadata-ingestion/tests/unit/test_redshift_lineage.py b/metadata-ingestion/tests/unit/test_redshift_lineage.py index c7d6ac18e044c..db5af3a71efb9 100644 --- a/metadata-ingestion/tests/unit/test_redshift_lineage.py +++ b/metadata-ingestion/tests/unit/test_redshift_lineage.py @@ -1,6 +1,8 @@ +from datahub.ingestion.api.common import PipelineContext from datahub.ingestion.source.redshift.config import RedshiftConfig from datahub.ingestion.source.redshift.lineage import RedshiftLineageExtractor from datahub.ingestion.source.redshift.report import RedshiftReport +from datahub.utilities.sqlglot_lineage import ColumnLineageInfo, DownstreamColumnRef def test_get_sources_from_query(): @@ -10,14 +12,20 @@ def test_get_sources_from_query(): test_query = """ select * from my_schema.my_table """ - lineage_extractor = RedshiftLineageExtractor(config, report) - lineage_datasets = lineage_extractor._get_sources_from_query( + lineage_extractor = RedshiftLineageExtractor( + config, report, PipelineContext(run_id="foo") + ) + lineage_datasets, _ = lineage_extractor._get_sources_from_query( db_name="test", query=test_query ) assert len(lineage_datasets) == 1 lineage = lineage_datasets[0] - assert lineage.path == "test.my_schema.my_table" + + assert ( + lineage.urn + == "urn:li:dataset:(urn:li:dataPlatform:redshift,test.my_schema.my_table,PROD)" + ) def test_get_sources_from_query_with_only_table_name(): @@ -27,14 +35,20 @@ def test_get_sources_from_query_with_only_table_name(): test_query = """ select * from my_table """ - lineage_extractor = RedshiftLineageExtractor(config, report) - lineage_datasets = lineage_extractor._get_sources_from_query( + lineage_extractor = RedshiftLineageExtractor( + config, report, PipelineContext(run_id="foo") + ) + lineage_datasets, _ = lineage_extractor._get_sources_from_query( db_name="test", query=test_query ) assert len(lineage_datasets) == 1 lineage = lineage_datasets[0] - assert lineage.path == "test.public.my_table" + + assert ( + lineage.urn + == "urn:li:dataset:(urn:li:dataPlatform:redshift,test.public.my_table,PROD)" + ) def test_get_sources_from_query_with_database(): @@ -44,14 +58,20 @@ def test_get_sources_from_query_with_database(): test_query = """ select * from test.my_schema.my_table """ - lineage_extractor = RedshiftLineageExtractor(config, report) - lineage_datasets = lineage_extractor._get_sources_from_query( + lineage_extractor = RedshiftLineageExtractor( + config, report, PipelineContext(run_id="foo") + ) + lineage_datasets, _ = lineage_extractor._get_sources_from_query( db_name="test", query=test_query ) assert len(lineage_datasets) == 1 lineage = lineage_datasets[0] - assert lineage.path == "test.my_schema.my_table" + + assert ( + lineage.urn + == "urn:li:dataset:(urn:li:dataPlatform:redshift,test.my_schema.my_table,PROD)" + ) def test_get_sources_from_query_with_non_default_database(): @@ -61,14 +81,20 @@ def test_get_sources_from_query_with_non_default_database(): test_query = """ select * from test2.my_schema.my_table """ - lineage_extractor = RedshiftLineageExtractor(config, report) - lineage_datasets = lineage_extractor._get_sources_from_query( + lineage_extractor = RedshiftLineageExtractor( + config, report, PipelineContext(run_id="foo") + ) + lineage_datasets, _ = lineage_extractor._get_sources_from_query( db_name="test", query=test_query ) assert len(lineage_datasets) == 1 lineage = lineage_datasets[0] - assert lineage.path == "test2.my_schema.my_table" + + assert ( + lineage.urn + == "urn:li:dataset:(urn:li:dataPlatform:redshift,test2.my_schema.my_table,PROD)" + ) def test_get_sources_from_query_with_only_table(): @@ -78,27 +104,48 @@ def test_get_sources_from_query_with_only_table(): test_query = """ select * from my_table """ - lineage_extractor = RedshiftLineageExtractor(config, report) - lineage_datasets = lineage_extractor._get_sources_from_query( + lineage_extractor = RedshiftLineageExtractor( + config, report, PipelineContext(run_id="foo") + ) + lineage_datasets, _ = lineage_extractor._get_sources_from_query( db_name="test", query=test_query ) assert len(lineage_datasets) == 1 lineage = lineage_datasets[0] - assert lineage.path == "test.public.my_table" + + assert ( + lineage.urn + == "urn:li:dataset:(urn:li:dataPlatform:redshift,test.public.my_table,PROD)" + ) -def test_get_sources_from_query_with_four_part_table_should_throw_exception(): +def test_cll(): config = RedshiftConfig(host_port="localhost:5439", database="test") report = RedshiftReport() test_query = """ - select * from database.schema.my_table.test + select a,b,c from db.public.customer inner join db.public.order on db.public.customer.id = db.public.order.customer_id """ - lineage_extractor = RedshiftLineageExtractor(config, report) - try: - lineage_extractor._get_sources_from_query(db_name="test", query=test_query) - except ValueError: - pass - - assert f"{test_query} should have thrown a ValueError exception but it didn't" + lineage_extractor = RedshiftLineageExtractor( + config, report, PipelineContext(run_id="foo") + ) + _, cll = lineage_extractor._get_sources_from_query(db_name="db", query=test_query) + + assert cll == [ + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="a"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="b"), + upstreams=[], + logic=None, + ), + ColumnLineageInfo( + downstream=DownstreamColumnRef(table=None, column="c"), + upstreams=[], + logic=None, + ), + ] From 4b6b941a2abf13854511c9af0e88a17d5acfd5e6 Mon Sep 17 00:00:00 2001 From: Harsha Mandadi <115464537+harsha-mandadi-4026@users.noreply.github.com> Date: Wed, 11 Oct 2023 19:01:46 +0100 Subject: [PATCH 114/156] fix(ingest): Fix postgres lineage within views (#8906) Co-authored-by: Harshal Sheth Co-authored-by: Maggie Hays --- .../datahub/ingestion/source/sql/postgres.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/postgres.py b/metadata-ingestion/src/datahub/ingestion/source/sql/postgres.py index ba8655b83446d..a6a9d8e2c8597 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/postgres.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/postgres.py @@ -217,14 +217,15 @@ def _get_view_lineage_elements( key = (lineage.dependent_view, lineage.dependent_schema) # Append the source table to the list. lineage_elements[key].append( - mce_builder.make_dataset_urn( - self.platform, - self.get_identifier( + mce_builder.make_dataset_urn_with_platform_instance( + platform=self.platform, + name=self.get_identifier( schema=lineage.source_schema, entity=lineage.source_table, inspector=inspector, ), - self.config.env, + platform_instance=self.config.platform_instance, + env=self.config.env, ) ) @@ -244,12 +245,13 @@ def _get_view_lineage_workunits( dependent_view, dependent_schema = key # Construct a lineage object. - urn = mce_builder.make_dataset_urn( - self.platform, - self.get_identifier( + urn = mce_builder.make_dataset_urn_with_platform_instance( + platform=self.platform, + name=self.get_identifier( schema=dependent_schema, entity=dependent_view, inspector=inspector ), - self.config.env, + platform_instance=self.config.platform_instance, + env=self.config.env, ) # use the mce_builder to ensure that the change proposal inherits From 932fbcddbf7c3201898e0918218e80c9246b0cd2 Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Wed, 11 Oct 2023 14:17:02 -0400 Subject: [PATCH 115/156] refactor(ingest/dbt): move dbt tests logic to dedicated file (#8984) --- .../src/datahub/ingestion/api/common.py | 9 + .../datahub/ingestion/source/csv_enricher.py | 8 +- .../datahub/ingestion/source/dbt/dbt_cloud.py | 3 +- .../ingestion/source/dbt/dbt_common.py | 278 +----------------- .../datahub/ingestion/source/dbt/dbt_core.py | 3 +- .../datahub/ingestion/source/dbt/dbt_tests.py | 261 ++++++++++++++++ 6 files changed, 288 insertions(+), 274 deletions(-) create mode 100644 metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_tests.py diff --git a/metadata-ingestion/src/datahub/ingestion/api/common.py b/metadata-ingestion/src/datahub/ingestion/api/common.py index 778bd119615e2..a6761a3c77d5e 100644 --- a/metadata-ingestion/src/datahub/ingestion/api/common.py +++ b/metadata-ingestion/src/datahub/ingestion/api/common.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, Generic, Iterable, Optional, Tuple, TypeVar +from datahub.configuration.common import ConfigurationError from datahub.emitter.mce_builder import set_dataset_urn_to_lower from datahub.ingestion.api.committable import Committable from datahub.ingestion.graph.client import DataHubGraph @@ -75,3 +76,11 @@ def register_checkpointer(self, committable: Committable) -> None: def get_committables(self) -> Iterable[Tuple[str, Committable]]: yield from self.checkpointers.items() + + def require_graph(self, operation: Optional[str] = None) -> DataHubGraph: + if not self.graph: + raise ConfigurationError( + f"{operation or 'This operation'} requires a graph, but none was provided. " + "To provide one, either use the datahub-rest sink or set the top-level datahub_api config in the recipe." + ) + return self.graph diff --git a/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py b/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py index 7cb487a86d931..611f0c5c52cc6 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py +++ b/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py @@ -129,11 +129,9 @@ def __init__(self, config: CSVEnricherConfig, ctx: PipelineContext): # Map from entity urn to a list of SubResourceRow. self.editable_schema_metadata_map: Dict[str, List[SubResourceRow]] = {} self.should_overwrite: bool = self.config.write_semantics == "OVERRIDE" - if not self.should_overwrite and not self.ctx.graph: - raise ConfigurationError( - "With PATCH semantics, the csv-enricher source requires a datahub_api to connect to. " - "Consider using the datahub-rest sink or provide a datahub_api: configuration on your ingestion recipe." - ) + + if not self.should_overwrite: + self.ctx.require_graph(operation="The csv-enricher's PATCH semantics flag") def get_resource_glossary_terms_work_unit( self, diff --git a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_cloud.py b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_cloud.py index af9769bc9d94c..da1ea8ecb4678 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_cloud.py +++ b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_cloud.py @@ -20,9 +20,8 @@ DBTCommonConfig, DBTNode, DBTSourceBase, - DBTTest, - DBTTestResult, ) +from datahub.ingestion.source.dbt.dbt_tests import DBTTest, DBTTestResult logger = logging.getLogger(__name__) diff --git a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py index 0f5c08eb6ac54..48d2118a9b091 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py +++ b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py @@ -1,11 +1,10 @@ -import json import logging import re from abc import abstractmethod from dataclasses import dataclass, field from datetime import datetime from enum import auto -from typing import Any, Callable, ClassVar, Dict, Iterable, List, Optional, Tuple, Union +from typing import Any, Dict, Iterable, List, Optional, Tuple import pydantic from pydantic import root_validator, validator @@ -34,6 +33,12 @@ from datahub.ingestion.api.source import MetadataWorkUnitProcessor from datahub.ingestion.api.workunit import MetadataWorkUnit from datahub.ingestion.source.common.subtypes import DatasetSubTypes +from datahub.ingestion.source.dbt.dbt_tests import ( + DBTTest, + DBTTestResult, + make_assertion_from_test, + make_assertion_result_from_test, +) from datahub.ingestion.source.sql.sql_types import ( ATHENA_SQL_TYPES_MAP, BIGQUERY_TYPES_MAP, @@ -81,20 +86,7 @@ TimeTypeClass, ) from datahub.metadata.schema_classes import ( - AssertionInfoClass, - AssertionResultClass, - AssertionResultTypeClass, - AssertionRunEventClass, - AssertionRunStatusClass, - AssertionStdAggregationClass, - AssertionStdOperatorClass, - AssertionStdParameterClass, - AssertionStdParametersClass, - AssertionStdParameterTypeClass, - AssertionTypeClass, DataPlatformInstanceClass, - DatasetAssertionInfoClass, - DatasetAssertionScopeClass, DatasetPropertiesClass, GlobalTagsClass, GlossaryTermsClass, @@ -551,134 +543,6 @@ def get_column_type( return SchemaFieldDataType(type=TypeClass()) -@dataclass -class AssertionParams: - scope: Union[DatasetAssertionScopeClass, str] - operator: Union[AssertionStdOperatorClass, str] - aggregation: Union[AssertionStdAggregationClass, str] - parameters: Optional[Callable[[Dict[str, str]], AssertionStdParametersClass]] = None - logic_fn: Optional[Callable[[Dict[str, str]], Optional[str]]] = None - - -def _get_name_for_relationship_test(kw_args: Dict[str, str]) -> Optional[str]: - """ - Try to produce a useful string for the name of a relationship constraint. - Return None if we fail to - """ - destination_ref = kw_args.get("to") - source_ref = kw_args.get("model") - column_name = kw_args.get("column_name") - dest_field_name = kw_args.get("field") - if not destination_ref or not source_ref or not column_name or not dest_field_name: - # base assertions are violated, bail early - return None - m = re.match(r"^ref\(\'(.*)\'\)$", destination_ref) - if m: - destination_table = m.group(1) - else: - destination_table = destination_ref - m = re.search(r"ref\(\'(.*)\'\)", source_ref) - if m: - source_table = m.group(1) - else: - source_table = source_ref - return f"{source_table}.{column_name} referential integrity to {destination_table}.{dest_field_name}" - - -@dataclass -class DBTTest: - qualified_test_name: str - column_name: Optional[str] - kw_args: dict - - TEST_NAME_TO_ASSERTION_MAP: ClassVar[Dict[str, AssertionParams]] = { - "not_null": AssertionParams( - scope=DatasetAssertionScopeClass.DATASET_COLUMN, - operator=AssertionStdOperatorClass.NOT_NULL, - aggregation=AssertionStdAggregationClass.IDENTITY, - ), - "unique": AssertionParams( - scope=DatasetAssertionScopeClass.DATASET_COLUMN, - operator=AssertionStdOperatorClass.EQUAL_TO, - aggregation=AssertionStdAggregationClass.UNIQUE_PROPOTION, - parameters=lambda _: AssertionStdParametersClass( - value=AssertionStdParameterClass( - value="1.0", - type=AssertionStdParameterTypeClass.NUMBER, - ) - ), - ), - "accepted_values": AssertionParams( - scope=DatasetAssertionScopeClass.DATASET_COLUMN, - operator=AssertionStdOperatorClass.IN, - aggregation=AssertionStdAggregationClass.IDENTITY, - parameters=lambda kw_args: AssertionStdParametersClass( - value=AssertionStdParameterClass( - value=json.dumps(kw_args.get("values")), - type=AssertionStdParameterTypeClass.SET, - ), - ), - ), - "relationships": AssertionParams( - scope=DatasetAssertionScopeClass.DATASET_COLUMN, - operator=AssertionStdOperatorClass._NATIVE_, - aggregation=AssertionStdAggregationClass.IDENTITY, - parameters=lambda kw_args: AssertionStdParametersClass( - value=AssertionStdParameterClass( - value=json.dumps(kw_args.get("values")), - type=AssertionStdParameterTypeClass.SET, - ), - ), - logic_fn=_get_name_for_relationship_test, - ), - "dbt_expectations.expect_column_values_to_not_be_null": AssertionParams( - scope=DatasetAssertionScopeClass.DATASET_COLUMN, - operator=AssertionStdOperatorClass.NOT_NULL, - aggregation=AssertionStdAggregationClass.IDENTITY, - ), - "dbt_expectations.expect_column_values_to_be_between": AssertionParams( - scope=DatasetAssertionScopeClass.DATASET_COLUMN, - operator=AssertionStdOperatorClass.BETWEEN, - aggregation=AssertionStdAggregationClass.IDENTITY, - parameters=lambda x: AssertionStdParametersClass( - minValue=AssertionStdParameterClass( - value=str(x.get("min_value", "unknown")), - type=AssertionStdParameterTypeClass.NUMBER, - ), - maxValue=AssertionStdParameterClass( - value=str(x.get("max_value", "unknown")), - type=AssertionStdParameterTypeClass.NUMBER, - ), - ), - ), - "dbt_expectations.expect_column_values_to_be_in_set": AssertionParams( - scope=DatasetAssertionScopeClass.DATASET_COLUMN, - operator=AssertionStdOperatorClass.IN, - aggregation=AssertionStdAggregationClass.IDENTITY, - parameters=lambda kw_args: AssertionStdParametersClass( - value=AssertionStdParameterClass( - value=json.dumps(kw_args.get("value_set")), - type=AssertionStdParameterTypeClass.SET, - ), - ), - ), - } - - -@dataclass -class DBTTestResult: - invocation_id: str - - status: str - execution_time: datetime - - native_results: Dict[str, str] - - -def string_map(input_map: Dict[str, Any]) -> Dict[str, str]: - return {k: str(v) for k, v in input_map.items()} - - @platform_name("dbt") @config_class(DBTCommonConfig) @support_status(SupportStatus.CERTIFIED) @@ -750,7 +614,7 @@ def create_test_entity_mcps( for upstream_urn in sorted(upstream_urns): if self.config.entities_enabled.can_emit_node_type("test"): - yield self._make_assertion_from_test( + yield make_assertion_from_test( custom_props, node, assertion_urn, @@ -759,133 +623,17 @@ def create_test_entity_mcps( if node.test_result: if self.config.entities_enabled.can_emit_test_results: - yield self._make_assertion_result_from_test( - node, assertion_urn, upstream_urn + yield make_assertion_result_from_test( + node, + assertion_urn, + upstream_urn, + test_warnings_are_errors=self.config.test_warnings_are_errors, ) else: logger.debug( f"Skipping test result {node.name} emission since it is turned off." ) - def _make_assertion_from_test( - self, - extra_custom_props: Dict[str, str], - node: DBTNode, - assertion_urn: str, - upstream_urn: str, - ) -> MetadataWorkUnit: - assert node.test_info - qualified_test_name = node.test_info.qualified_test_name - column_name = node.test_info.column_name - kw_args = node.test_info.kw_args - - if qualified_test_name in DBTTest.TEST_NAME_TO_ASSERTION_MAP: - assertion_params = DBTTest.TEST_NAME_TO_ASSERTION_MAP[qualified_test_name] - assertion_info = AssertionInfoClass( - type=AssertionTypeClass.DATASET, - customProperties=extra_custom_props, - datasetAssertion=DatasetAssertionInfoClass( - dataset=upstream_urn, - scope=assertion_params.scope, - operator=assertion_params.operator, - fields=[ - mce_builder.make_schema_field_urn(upstream_urn, column_name) - ] - if ( - assertion_params.scope - == DatasetAssertionScopeClass.DATASET_COLUMN - and column_name - ) - else [], - nativeType=node.name, - aggregation=assertion_params.aggregation, - parameters=assertion_params.parameters(kw_args) - if assertion_params.parameters - else None, - logic=assertion_params.logic_fn(kw_args) - if assertion_params.logic_fn - else None, - nativeParameters=string_map(kw_args), - ), - ) - elif column_name: - # no match with known test types, column-level test - assertion_info = AssertionInfoClass( - type=AssertionTypeClass.DATASET, - customProperties=extra_custom_props, - datasetAssertion=DatasetAssertionInfoClass( - dataset=upstream_urn, - scope=DatasetAssertionScopeClass.DATASET_COLUMN, - operator=AssertionStdOperatorClass._NATIVE_, - fields=[ - mce_builder.make_schema_field_urn(upstream_urn, column_name) - ], - nativeType=node.name, - logic=node.compiled_code or node.raw_code, - aggregation=AssertionStdAggregationClass._NATIVE_, - nativeParameters=string_map(kw_args), - ), - ) - else: - # no match with known test types, default to row-level test - assertion_info = AssertionInfoClass( - type=AssertionTypeClass.DATASET, - customProperties=extra_custom_props, - datasetAssertion=DatasetAssertionInfoClass( - dataset=upstream_urn, - scope=DatasetAssertionScopeClass.DATASET_ROWS, - operator=AssertionStdOperatorClass._NATIVE_, - logic=node.compiled_code or node.raw_code, - nativeType=node.name, - aggregation=AssertionStdAggregationClass._NATIVE_, - nativeParameters=string_map(kw_args), - ), - ) - - wu = MetadataChangeProposalWrapper( - entityUrn=assertion_urn, - aspect=assertion_info, - ).as_workunit() - - return wu - - def _make_assertion_result_from_test( - self, - node: DBTNode, - assertion_urn: str, - upstream_urn: str, - ) -> MetadataWorkUnit: - assert node.test_result - test_result = node.test_result - - assertionResult = AssertionRunEventClass( - timestampMillis=int(test_result.execution_time.timestamp() * 1000.0), - assertionUrn=assertion_urn, - asserteeUrn=upstream_urn, - runId=test_result.invocation_id, - result=AssertionResultClass( - type=AssertionResultTypeClass.SUCCESS - if test_result.status == "pass" - or ( - not self.config.test_warnings_are_errors - and test_result.status == "warn" - ) - else AssertionResultTypeClass.FAILURE, - nativeResults=test_result.native_results, - ), - status=AssertionRunStatusClass.COMPLETE, - ) - - event = MetadataChangeProposalWrapper( - entityUrn=assertion_urn, - aspect=assertionResult, - ) - wu = MetadataWorkUnit( - id=f"{assertion_urn}-assertionRunEvent-{upstream_urn}", - mcp=event, - ) - return wu - @abstractmethod def load_nodes(self) -> Tuple[List[DBTNode], Dict[str, Optional[str]]]: # return dbt nodes + global custom properties diff --git a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_core.py b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_core.py index c08295ed1dc59..dc3a84847beb2 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_core.py +++ b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_core.py @@ -26,9 +26,8 @@ DBTNode, DBTSourceBase, DBTSourceReport, - DBTTest, - DBTTestResult, ) +from datahub.ingestion.source.dbt.dbt_tests import DBTTest, DBTTestResult logger = logging.getLogger(__name__) diff --git a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_tests.py b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_tests.py new file mode 100644 index 0000000000000..721769d214d9e --- /dev/null +++ b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_tests.py @@ -0,0 +1,261 @@ +import json +import re +from dataclasses import dataclass +from datetime import datetime +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union + +from datahub.emitter import mce_builder +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.ingestion.api.workunit import MetadataWorkUnit +from datahub.metadata.schema_classes import ( + AssertionInfoClass, + AssertionResultClass, + AssertionResultTypeClass, + AssertionRunEventClass, + AssertionRunStatusClass, + AssertionStdAggregationClass, + AssertionStdOperatorClass, + AssertionStdParameterClass, + AssertionStdParametersClass, + AssertionStdParameterTypeClass, + AssertionTypeClass, + DatasetAssertionInfoClass, + DatasetAssertionScopeClass, +) + +if TYPE_CHECKING: + from datahub.ingestion.source.dbt.dbt_common import DBTNode + + +@dataclass +class DBTTest: + qualified_test_name: str + column_name: Optional[str] + kw_args: dict + + +@dataclass +class DBTTestResult: + invocation_id: str + + status: str + execution_time: datetime + + native_results: Dict[str, str] + + +def _get_name_for_relationship_test(kw_args: Dict[str, str]) -> Optional[str]: + """ + Try to produce a useful string for the name of a relationship constraint. + Return None if we fail to + """ + destination_ref = kw_args.get("to") + source_ref = kw_args.get("model") + column_name = kw_args.get("column_name") + dest_field_name = kw_args.get("field") + if not destination_ref or not source_ref or not column_name or not dest_field_name: + # base assertions are violated, bail early + return None + m = re.match(r"^ref\(\'(.*)\'\)$", destination_ref) + if m: + destination_table = m.group(1) + else: + destination_table = destination_ref + m = re.search(r"ref\(\'(.*)\'\)", source_ref) + if m: + source_table = m.group(1) + else: + source_table = source_ref + return f"{source_table}.{column_name} referential integrity to {destination_table}.{dest_field_name}" + + +@dataclass +class AssertionParams: + scope: Union[DatasetAssertionScopeClass, str] + operator: Union[AssertionStdOperatorClass, str] + aggregation: Union[AssertionStdAggregationClass, str] + parameters: Optional[Callable[[Dict[str, str]], AssertionStdParametersClass]] = None + logic_fn: Optional[Callable[[Dict[str, str]], Optional[str]]] = None + + +_DBT_TEST_NAME_TO_ASSERTION_MAP: Dict[str, AssertionParams] = { + "not_null": AssertionParams( + scope=DatasetAssertionScopeClass.DATASET_COLUMN, + operator=AssertionStdOperatorClass.NOT_NULL, + aggregation=AssertionStdAggregationClass.IDENTITY, + ), + "unique": AssertionParams( + scope=DatasetAssertionScopeClass.DATASET_COLUMN, + operator=AssertionStdOperatorClass.EQUAL_TO, + aggregation=AssertionStdAggregationClass.UNIQUE_PROPOTION, + parameters=lambda _: AssertionStdParametersClass( + value=AssertionStdParameterClass( + value="1.0", + type=AssertionStdParameterTypeClass.NUMBER, + ) + ), + ), + "accepted_values": AssertionParams( + scope=DatasetAssertionScopeClass.DATASET_COLUMN, + operator=AssertionStdOperatorClass.IN, + aggregation=AssertionStdAggregationClass.IDENTITY, + parameters=lambda kw_args: AssertionStdParametersClass( + value=AssertionStdParameterClass( + value=json.dumps(kw_args.get("values")), + type=AssertionStdParameterTypeClass.SET, + ), + ), + ), + "relationships": AssertionParams( + scope=DatasetAssertionScopeClass.DATASET_COLUMN, + operator=AssertionStdOperatorClass._NATIVE_, + aggregation=AssertionStdAggregationClass.IDENTITY, + parameters=lambda kw_args: AssertionStdParametersClass( + value=AssertionStdParameterClass( + value=json.dumps(kw_args.get("values")), + type=AssertionStdParameterTypeClass.SET, + ), + ), + logic_fn=_get_name_for_relationship_test, + ), + "dbt_expectations.expect_column_values_to_not_be_null": AssertionParams( + scope=DatasetAssertionScopeClass.DATASET_COLUMN, + operator=AssertionStdOperatorClass.NOT_NULL, + aggregation=AssertionStdAggregationClass.IDENTITY, + ), + "dbt_expectations.expect_column_values_to_be_between": AssertionParams( + scope=DatasetAssertionScopeClass.DATASET_COLUMN, + operator=AssertionStdOperatorClass.BETWEEN, + aggregation=AssertionStdAggregationClass.IDENTITY, + parameters=lambda x: AssertionStdParametersClass( + minValue=AssertionStdParameterClass( + value=str(x.get("min_value", "unknown")), + type=AssertionStdParameterTypeClass.NUMBER, + ), + maxValue=AssertionStdParameterClass( + value=str(x.get("max_value", "unknown")), + type=AssertionStdParameterTypeClass.NUMBER, + ), + ), + ), + "dbt_expectations.expect_column_values_to_be_in_set": AssertionParams( + scope=DatasetAssertionScopeClass.DATASET_COLUMN, + operator=AssertionStdOperatorClass.IN, + aggregation=AssertionStdAggregationClass.IDENTITY, + parameters=lambda kw_args: AssertionStdParametersClass( + value=AssertionStdParameterClass( + value=json.dumps(kw_args.get("value_set")), + type=AssertionStdParameterTypeClass.SET, + ), + ), + ), +} + + +def _string_map(input_map: Dict[str, Any]) -> Dict[str, str]: + return {k: str(v) for k, v in input_map.items()} + + +def make_assertion_from_test( + extra_custom_props: Dict[str, str], + node: "DBTNode", + assertion_urn: str, + upstream_urn: str, +) -> MetadataWorkUnit: + assert node.test_info + qualified_test_name = node.test_info.qualified_test_name + column_name = node.test_info.column_name + kw_args = node.test_info.kw_args + + if qualified_test_name in _DBT_TEST_NAME_TO_ASSERTION_MAP: + assertion_params = _DBT_TEST_NAME_TO_ASSERTION_MAP[qualified_test_name] + assertion_info = AssertionInfoClass( + type=AssertionTypeClass.DATASET, + customProperties=extra_custom_props, + datasetAssertion=DatasetAssertionInfoClass( + dataset=upstream_urn, + scope=assertion_params.scope, + operator=assertion_params.operator, + fields=[mce_builder.make_schema_field_urn(upstream_urn, column_name)] + if ( + assertion_params.scope == DatasetAssertionScopeClass.DATASET_COLUMN + and column_name + ) + else [], + nativeType=node.name, + aggregation=assertion_params.aggregation, + parameters=assertion_params.parameters(kw_args) + if assertion_params.parameters + else None, + logic=assertion_params.logic_fn(kw_args) + if assertion_params.logic_fn + else None, + nativeParameters=_string_map(kw_args), + ), + ) + elif column_name: + # no match with known test types, column-level test + assertion_info = AssertionInfoClass( + type=AssertionTypeClass.DATASET, + customProperties=extra_custom_props, + datasetAssertion=DatasetAssertionInfoClass( + dataset=upstream_urn, + scope=DatasetAssertionScopeClass.DATASET_COLUMN, + operator=AssertionStdOperatorClass._NATIVE_, + fields=[mce_builder.make_schema_field_urn(upstream_urn, column_name)], + nativeType=node.name, + logic=node.compiled_code or node.raw_code, + aggregation=AssertionStdAggregationClass._NATIVE_, + nativeParameters=_string_map(kw_args), + ), + ) + else: + # no match with known test types, default to row-level test + assertion_info = AssertionInfoClass( + type=AssertionTypeClass.DATASET, + customProperties=extra_custom_props, + datasetAssertion=DatasetAssertionInfoClass( + dataset=upstream_urn, + scope=DatasetAssertionScopeClass.DATASET_ROWS, + operator=AssertionStdOperatorClass._NATIVE_, + logic=node.compiled_code or node.raw_code, + nativeType=node.name, + aggregation=AssertionStdAggregationClass._NATIVE_, + nativeParameters=_string_map(kw_args), + ), + ) + + return MetadataChangeProposalWrapper( + entityUrn=assertion_urn, + aspect=assertion_info, + ).as_workunit() + + +def make_assertion_result_from_test( + node: "DBTNode", + assertion_urn: str, + upstream_urn: str, + test_warnings_are_errors: bool, +) -> MetadataWorkUnit: + assert node.test_result + test_result = node.test_result + + assertionResult = AssertionRunEventClass( + timestampMillis=int(test_result.execution_time.timestamp() * 1000.0), + assertionUrn=assertion_urn, + asserteeUrn=upstream_urn, + runId=test_result.invocation_id, + result=AssertionResultClass( + type=AssertionResultTypeClass.SUCCESS + if test_result.status == "pass" + or (not test_warnings_are_errors and test_result.status == "warn") + else AssertionResultTypeClass.FAILURE, + nativeResults=test_result.native_results, + ), + status=AssertionRunStatusClass.COMPLETE, + ) + + return MetadataChangeProposalWrapper( + entityUrn=assertion_urn, + aspect=assertionResult, + ).as_workunit() From 1b06c6a30c8d6c0ee57f75f75ee6a436aa6c13a7 Mon Sep 17 00:00:00 2001 From: Mayuri Nehate <33225191+mayurinehate@users.noreply.github.com> Date: Thu, 12 Oct 2023 00:31:42 +0530 Subject: [PATCH 116/156] fix(ingest/snowflake): fix sample fraction for very large tables (#8988) --- .../datahub/ingestion/source/snowflake/snowflake_profiler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_profiler.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_profiler.py index 24275dcdff34d..8e18d85d6f3ca 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_profiler.py +++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_profiler.py @@ -86,7 +86,7 @@ def get_batch_kwargs( # Fixed-size sampling can be slower than equivalent fraction-based sampling # as per https://docs.snowflake.com/en/sql-reference/constructs/sample#performance-considerations sample_pc = 100 * self.config.profiling.sample_size / table.rows_count - custom_sql = f'select * from "{db_name}"."{schema_name}"."{table.name}" TABLESAMPLE ({sample_pc:.3f})' + custom_sql = f'select * from "{db_name}"."{schema_name}"."{table.name}" TABLESAMPLE ({sample_pc:.8f})' return { **super().get_batch_kwargs(table, schema_name, db_name), # Lowercase/Mixedcase table names in Snowflake do not work by default. From 245284ec6c6b754b22943ba42d7139ddd5772377 Mon Sep 17 00:00:00 2001 From: jayasimhankv <145704974+jayasimhankv@users.noreply.github.com> Date: Wed, 11 Oct 2023 17:40:20 -0500 Subject: [PATCH 117/156] fix(): Display generic not found page for corp groups that do not exist (#8880) Co-authored-by: Jay Kadambi --- .../java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java | 3 ++- datahub-graphql-core/src/main/resources/entity.graphql | 5 +++++ datahub-web-react/src/app/entity/group/GroupProfile.tsx | 4 ++++ datahub-web-react/src/graphql/group.graphql | 1 + 4 files changed, 12 insertions(+), 1 deletion(-) diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java index ebb5c7d62c7d3..b99f712034fe0 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java @@ -1292,7 +1292,8 @@ private void configureCorpUserResolvers(final RuntimeWiring.Builder builder) { */ private void configureCorpGroupResolvers(final RuntimeWiring.Builder builder) { builder.type("CorpGroup", typeWiring -> typeWiring - .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))); + .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient)) + .dataFetcher("exists", new EntityExistsResolver(entityService))); builder.type("CorpGroupInfo", typeWiring -> typeWiring .dataFetcher("admins", new LoadableTypeBatchResolver<>(corpUserType, diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql index 0b15d7b875a9c..b37a8f34fa056 100644 --- a/datahub-graphql-core/src/main/resources/entity.graphql +++ b/datahub-graphql-core/src/main/resources/entity.graphql @@ -3788,6 +3788,11 @@ type CorpGroup implements Entity { Additional read only info about the group """ info: CorpGroupInfo @deprecated + + """ + Whether or not this entity exists on DataHub + """ + exists: Boolean } """ diff --git a/datahub-web-react/src/app/entity/group/GroupProfile.tsx b/datahub-web-react/src/app/entity/group/GroupProfile.tsx index d5e284af931df..53d2062277dec 100644 --- a/datahub-web-react/src/app/entity/group/GroupProfile.tsx +++ b/datahub-web-react/src/app/entity/group/GroupProfile.tsx @@ -11,6 +11,7 @@ import { RoutedTabs } from '../../shared/RoutedTabs'; import GroupInfoSidebar from './GroupInfoSideBar'; import { GroupAssets } from './GroupAssets'; import { ErrorSection } from '../../shared/error/ErrorSection'; +import NonExistentEntityPage from '../shared/entity/NonExistentEntityPage'; const messageStyle = { marginTop: '10%' }; @@ -110,6 +111,9 @@ export default function GroupProfile() { urn, }; + if (data?.corpGroup?.exists === false) { + return ; + } return ( <> {error && } diff --git a/datahub-web-react/src/graphql/group.graphql b/datahub-web-react/src/graphql/group.graphql index 9aa6e2b005f16..1007721e51a4e 100644 --- a/datahub-web-react/src/graphql/group.graphql +++ b/datahub-web-react/src/graphql/group.graphql @@ -3,6 +3,7 @@ query getGroup($urn: String!, $membersCount: Int!) { urn type name + exists origin { type externalType From 245c5c00087116d236acf7a9bbddbdb4dee15949 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20G=C3=B3mez=20Villamor?= Date: Thu, 12 Oct 2023 02:06:19 +0200 Subject: [PATCH 118/156] fix(ingest/looker): stop emitting tag owner (#8942) --- docs/how/updating-datahub.md | 2 + .../ingestion/source/looker/looker_common.py | 13 +----- .../looker/golden_looker_mces.json | 42 ------------------- .../looker/golden_test_allow_ingest.json | 42 ------------------- ...olden_test_external_project_view_mces.json | 42 ------------------- .../looker/golden_test_file_path_ingest.json | 42 ------------------- .../golden_test_independent_look_ingest.json | 42 ------------------- .../looker/golden_test_ingest.json | 42 ------------------- .../looker/golden_test_ingest_joins.json | 42 ------------------- .../golden_test_ingest_unaliased_joins.json | 42 ------------------- .../looker_mces_golden_deleted_stateful.json | 42 ------------------- .../looker/looker_mces_usage_history.json | 42 ------------------- .../lookml/lookml_mces_api_bigquery.json | 42 ------------------- .../lookml/lookml_mces_api_hive2.json | 42 ------------------- .../lookml/lookml_mces_badsql_parser.json | 42 ------------------- .../lookml/lookml_mces_offline.json | 42 ------------------- .../lookml_mces_offline_deny_pattern.json | 42 ------------------- ...lookml_mces_offline_platform_instance.json | 42 ------------------- .../lookml_mces_with_external_urls.json | 42 ------------------- .../lookml/lookml_reachable_views.json | 42 ------------------- 20 files changed, 3 insertions(+), 768 deletions(-) diff --git a/docs/how/updating-datahub.md b/docs/how/updating-datahub.md index 5d0ad5eaf8f7e..9cd4ad5c6f02d 100644 --- a/docs/how/updating-datahub.md +++ b/docs/how/updating-datahub.md @@ -7,6 +7,8 @@ This file documents any backwards-incompatible changes in DataHub and assists pe ### Breaking Changes - #8810 - Removed support for SQLAlchemy 1.3.x. Only SQLAlchemy 1.4.x is supported now. +- #8942 - Removed `urn:li:corpuser:datahub` owner for the `Measure`, `Dimension` and `Temporal` tags emitted + by Looker and LookML source connectors. - #8853 - The Airflow plugin no longer supports Airflow 2.0.x or Python 3.7. See the docs for more details. - #8853 - Introduced the Airflow plugin v2. If you're using Airflow 2.3+, the v2 plugin will be enabled by default, and so you'll need to switch your requirements to include `pip install 'acryl-datahub-airflow-plugin[plugin-v2]'`. To continue using the v1 plugin, set the `DATAHUB_AIRFLOW_PLUGIN_USE_V1_PLUGIN` environment variable to `true`. - #8943 The Unity Catalog ingestion source has a new option `include_metastore`, which will cause all urns to be changed when disabled. diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_common.py b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_common.py index 89b1e45695c57..30c38720dd96c 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_common.py +++ b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_common.py @@ -81,9 +81,6 @@ EnumTypeClass, FineGrainedLineageClass, GlobalTagsClass, - OwnerClass, - OwnershipClass, - OwnershipTypeClass, SchemaMetadataClass, StatusClass, SubTypesClass, @@ -453,17 +450,9 @@ def _get_schema( @staticmethod def _get_tag_mce_for_urn(tag_urn: str) -> MetadataChangeEvent: assert tag_urn in LookerUtil.tag_definitions - ownership = OwnershipClass( - owners=[ - OwnerClass( - owner="urn:li:corpuser:datahub", - type=OwnershipTypeClass.DATAOWNER, - ) - ] - ) return MetadataChangeEvent( proposedSnapshot=TagSnapshotClass( - urn=tag_urn, aspects=[ownership, LookerUtil.tag_definitions[tag_urn]] + urn=tag_urn, aspects=[LookerUtil.tag_definitions[tag_urn]] ) ) diff --git a/metadata-ingestion/tests/integration/looker/golden_looker_mces.json b/metadata-ingestion/tests/integration/looker/golden_looker_mces.json index dee85b40bb7a8..1da42b94e320c 100644 --- a/metadata-ingestion/tests/integration/looker/golden_looker_mces.json +++ b/metadata-ingestion/tests/integration/looker/golden_looker_mces.json @@ -533,20 +533,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -566,20 +552,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -599,20 +571,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/looker/golden_test_allow_ingest.json b/metadata-ingestion/tests/integration/looker/golden_test_allow_ingest.json index 72db36e63daf7..685a606a57c33 100644 --- a/metadata-ingestion/tests/integration/looker/golden_test_allow_ingest.json +++ b/metadata-ingestion/tests/integration/looker/golden_test_allow_ingest.json @@ -327,20 +327,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -360,20 +346,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -393,20 +365,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/looker/golden_test_external_project_view_mces.json b/metadata-ingestion/tests/integration/looker/golden_test_external_project_view_mces.json index e5508bdb06b9e..069788cb088ac 100644 --- a/metadata-ingestion/tests/integration/looker/golden_test_external_project_view_mces.json +++ b/metadata-ingestion/tests/integration/looker/golden_test_external_project_view_mces.json @@ -327,20 +327,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -360,20 +346,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -393,20 +365,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/looker/golden_test_file_path_ingest.json b/metadata-ingestion/tests/integration/looker/golden_test_file_path_ingest.json index b0f66e7b245c9..f1c932ebd5a70 100644 --- a/metadata-ingestion/tests/integration/looker/golden_test_file_path_ingest.json +++ b/metadata-ingestion/tests/integration/looker/golden_test_file_path_ingest.json @@ -335,20 +335,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -369,20 +355,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -403,20 +375,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/looker/golden_test_independent_look_ingest.json b/metadata-ingestion/tests/integration/looker/golden_test_independent_look_ingest.json index 91e13debfa028..9521c9af4bbdc 100644 --- a/metadata-ingestion/tests/integration/looker/golden_test_independent_look_ingest.json +++ b/metadata-ingestion/tests/integration/looker/golden_test_independent_look_ingest.json @@ -550,20 +550,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -583,20 +569,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -616,20 +588,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/looker/golden_test_ingest.json b/metadata-ingestion/tests/integration/looker/golden_test_ingest.json index e93079119e4f4..dbacd52fe83de 100644 --- a/metadata-ingestion/tests/integration/looker/golden_test_ingest.json +++ b/metadata-ingestion/tests/integration/looker/golden_test_ingest.json @@ -327,20 +327,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -360,20 +346,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -393,20 +365,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/looker/golden_test_ingest_joins.json b/metadata-ingestion/tests/integration/looker/golden_test_ingest_joins.json index a9c8efa7cdb98..aaa874d9ff348 100644 --- a/metadata-ingestion/tests/integration/looker/golden_test_ingest_joins.json +++ b/metadata-ingestion/tests/integration/looker/golden_test_ingest_joins.json @@ -351,20 +351,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -384,20 +370,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -417,20 +389,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/looker/golden_test_ingest_unaliased_joins.json b/metadata-ingestion/tests/integration/looker/golden_test_ingest_unaliased_joins.json index edd15624a14cd..be8db0722aea3 100644 --- a/metadata-ingestion/tests/integration/looker/golden_test_ingest_unaliased_joins.json +++ b/metadata-ingestion/tests/integration/looker/golden_test_ingest_unaliased_joins.json @@ -343,20 +343,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -376,20 +362,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -409,20 +381,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/looker/looker_mces_golden_deleted_stateful.json b/metadata-ingestion/tests/integration/looker/looker_mces_golden_deleted_stateful.json index aebc89b609a08..05b74f163ad45 100644 --- a/metadata-ingestion/tests/integration/looker/looker_mces_golden_deleted_stateful.json +++ b/metadata-ingestion/tests/integration/looker/looker_mces_golden_deleted_stateful.json @@ -327,20 +327,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -360,20 +346,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -393,20 +365,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/looker/looker_mces_usage_history.json b/metadata-ingestion/tests/integration/looker/looker_mces_usage_history.json index 34bded3cf691e..0778aa0050b00 100644 --- a/metadata-ingestion/tests/integration/looker/looker_mces_usage_history.json +++ b/metadata-ingestion/tests/integration/looker/looker_mces_usage_history.json @@ -279,20 +279,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -312,20 +298,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -345,20 +317,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/lookml/lookml_mces_api_bigquery.json b/metadata-ingestion/tests/integration/lookml/lookml_mces_api_bigquery.json index 238f4c2580cdf..5a0bd4e12fd3a 100644 --- a/metadata-ingestion/tests/integration/lookml/lookml_mces_api_bigquery.json +++ b/metadata-ingestion/tests/integration/lookml/lookml_mces_api_bigquery.json @@ -2121,20 +2121,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -2154,20 +2140,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -2187,20 +2159,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/lookml/lookml_mces_api_hive2.json b/metadata-ingestion/tests/integration/lookml/lookml_mces_api_hive2.json index 45d5d839e9d21..1b0ee3216383c 100644 --- a/metadata-ingestion/tests/integration/lookml/lookml_mces_api_hive2.json +++ b/metadata-ingestion/tests/integration/lookml/lookml_mces_api_hive2.json @@ -2121,20 +2121,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -2154,20 +2140,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -2187,20 +2159,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/lookml/lookml_mces_badsql_parser.json b/metadata-ingestion/tests/integration/lookml/lookml_mces_badsql_parser.json index 187cedaefb6b2..b960ba581e6b5 100644 --- a/metadata-ingestion/tests/integration/lookml/lookml_mces_badsql_parser.json +++ b/metadata-ingestion/tests/integration/lookml/lookml_mces_badsql_parser.json @@ -2004,20 +2004,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -2037,20 +2023,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -2070,20 +2042,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/lookml/lookml_mces_offline.json b/metadata-ingestion/tests/integration/lookml/lookml_mces_offline.json index c2c879e38f37b..e29292a44c949 100644 --- a/metadata-ingestion/tests/integration/lookml/lookml_mces_offline.json +++ b/metadata-ingestion/tests/integration/lookml/lookml_mces_offline.json @@ -2121,20 +2121,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -2154,20 +2140,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -2187,20 +2159,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/lookml/lookml_mces_offline_deny_pattern.json b/metadata-ingestion/tests/integration/lookml/lookml_mces_offline_deny_pattern.json index c1ac54b0fb588..04ecaecbd4afb 100644 --- a/metadata-ingestion/tests/integration/lookml/lookml_mces_offline_deny_pattern.json +++ b/metadata-ingestion/tests/integration/lookml/lookml_mces_offline_deny_pattern.json @@ -584,20 +584,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -617,20 +603,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -650,20 +622,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/lookml/lookml_mces_offline_platform_instance.json b/metadata-ingestion/tests/integration/lookml/lookml_mces_offline_platform_instance.json index f602ca37b3160..080931ae637bc 100644 --- a/metadata-ingestion/tests/integration/lookml/lookml_mces_offline_platform_instance.json +++ b/metadata-ingestion/tests/integration/lookml/lookml_mces_offline_platform_instance.json @@ -2121,20 +2121,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -2154,20 +2140,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -2187,20 +2159,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/lookml/lookml_mces_with_external_urls.json b/metadata-ingestion/tests/integration/lookml/lookml_mces_with_external_urls.json index 104bd365669e3..5826c4316b539 100644 --- a/metadata-ingestion/tests/integration/lookml/lookml_mces_with_external_urls.json +++ b/metadata-ingestion/tests/integration/lookml/lookml_mces_with_external_urls.json @@ -2134,20 +2134,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -2167,20 +2153,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -2200,20 +2172,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", diff --git a/metadata-ingestion/tests/integration/lookml/lookml_reachable_views.json b/metadata-ingestion/tests/integration/lookml/lookml_reachable_views.json index 37a6c94c6952e..53d1ec0229de1 100644 --- a/metadata-ingestion/tests/integration/lookml/lookml_reachable_views.json +++ b/metadata-ingestion/tests/integration/lookml/lookml_reachable_views.json @@ -681,20 +681,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Dimension", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Dimension", @@ -714,20 +700,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Temporal", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Temporal", @@ -747,20 +719,6 @@ "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { "urn": "urn:li:tag:Measure", "aspects": [ - { - "com.linkedin.pegasus2avro.common.Ownership": { - "owners": [ - { - "owner": "urn:li:corpuser:datahub", - "type": "DATAOWNER" - } - ], - "lastModified": { - "time": 0, - "actor": "urn:li:corpuser:unknown" - } - } - }, { "com.linkedin.pegasus2avro.tag.TagProperties": { "name": "Measure", From 84bba4dc446ee97f8991689fd17bfa6d14232601 Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Thu, 12 Oct 2023 01:31:17 -0400 Subject: [PATCH 119/156] feat(ingest): add output schema inference for sql parser (#8989) --- .../src/datahub/utilities/sqlglot_lineage.py | 119 ++++++++++++++++-- .../integration/powerbi/test_m_parser.py | 93 ++++---------- .../test_bigquery_create_view_with_cte.json | 32 ++++- ..._bigquery_from_sharded_table_wildcard.json | 16 ++- .../test_bigquery_nested_subqueries.json | 16 ++- ..._bigquery_sharded_table_normalization.json | 16 ++- .../test_bigquery_star_with_replace.json | 24 +++- .../test_bigquery_view_from_union.json | 16 ++- .../goldens/test_create_view_as_select.json | 16 ++- .../test_expand_select_star_basic.json | 80 ++++++++++-- .../goldens/test_insert_as_select.json | 36 +++++- ...est_select_ambiguous_column_no_schema.json | 12 +- .../goldens/test_select_count.json | 8 +- .../test_select_from_struct_subfields.json | 16 ++- .../goldens/test_select_from_union.json | 16 ++- .../sql_parsing/goldens/test_select_max.json | 4 +- .../goldens/test_select_with_ctes.json | 8 +- .../test_select_with_full_col_name.json | 12 +- .../test_snowflake_case_statement.json | 16 ++- .../goldens/test_snowflake_column_cast.json | 63 ++++++++++ .../test_snowflake_column_normalization.json | 32 ++++- ...t_snowflake_ctas_column_normalization.json | 32 ++++- .../test_snowflake_default_normalization.json | 48 ++++++- .../unit/sql_parsing/test_sqlglot_lineage.py | 21 ++++ 24 files changed, 604 insertions(+), 148 deletions(-) create mode 100644 metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_column_cast.json diff --git a/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py b/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py index 81c43884fdf7d..349eb40a5e865 100644 --- a/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py +++ b/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py @@ -5,12 +5,13 @@ import logging import pathlib from collections import defaultdict -from typing import Dict, List, Optional, Set, Tuple, Union +from typing import Any, Dict, List, Optional, Set, Tuple, Union import pydantic.dataclasses import sqlglot import sqlglot.errors import sqlglot.lineage +import sqlglot.optimizer.annotate_types import sqlglot.optimizer.qualify import sqlglot.optimizer.qualify_columns from pydantic import BaseModel @@ -23,7 +24,17 @@ from datahub.ingestion.api.closeable import Closeable from datahub.ingestion.graph.client import DataHubGraph from datahub.ingestion.source.bigquery_v2.bigquery_audit import BigqueryTableIdentifier -from datahub.metadata.schema_classes import OperationTypeClass, SchemaMetadataClass +from datahub.metadata.schema_classes import ( + ArrayTypeClass, + BooleanTypeClass, + DateTypeClass, + NumberTypeClass, + OperationTypeClass, + SchemaFieldDataTypeClass, + SchemaMetadataClass, + StringTypeClass, + TimeTypeClass, +) from datahub.utilities.file_backed_collections import ConnectionWrapper, FileBackedDict from datahub.utilities.urns.dataset_urn import DatasetUrn @@ -90,8 +101,18 @@ def get_query_type_of_sql(expression: sqlglot.exp.Expression) -> QueryType: return QueryType.UNKNOWN +class _ParserBaseModel( + BaseModel, + arbitrary_types_allowed=True, + json_encoders={ + SchemaFieldDataTypeClass: lambda v: v.to_obj(), + }, +): + pass + + @functools.total_ordering -class _FrozenModel(BaseModel, frozen=True): +class _FrozenModel(_ParserBaseModel, frozen=True): def __lt__(self, other: "_FrozenModel") -> bool: for field in self.__fields__: self_v = getattr(self, field) @@ -146,29 +167,42 @@ class _ColumnRef(_FrozenModel): column: str -class ColumnRef(BaseModel): +class ColumnRef(_ParserBaseModel): table: Urn column: str -class _DownstreamColumnRef(BaseModel): +class _DownstreamColumnRef(_ParserBaseModel): table: Optional[_TableName] column: str + column_type: Optional[sqlglot.exp.DataType] -class DownstreamColumnRef(BaseModel): +class DownstreamColumnRef(_ParserBaseModel): table: Optional[Urn] column: str + column_type: Optional[SchemaFieldDataTypeClass] + native_column_type: Optional[str] + + @pydantic.validator("column_type", pre=True) + def _load_column_type( + cls, v: Optional[Union[dict, SchemaFieldDataTypeClass]] + ) -> Optional[SchemaFieldDataTypeClass]: + if v is None: + return None + if isinstance(v, SchemaFieldDataTypeClass): + return v + return SchemaFieldDataTypeClass.from_obj(v) -class _ColumnLineageInfo(BaseModel): +class _ColumnLineageInfo(_ParserBaseModel): downstream: _DownstreamColumnRef upstreams: List[_ColumnRef] logic: Optional[str] -class ColumnLineageInfo(BaseModel): +class ColumnLineageInfo(_ParserBaseModel): downstream: DownstreamColumnRef upstreams: List[ColumnRef] @@ -176,7 +210,7 @@ class ColumnLineageInfo(BaseModel): logic: Optional[str] = pydantic.Field(default=None, exclude=True) -class SqlParsingDebugInfo(BaseModel, arbitrary_types_allowed=True): +class SqlParsingDebugInfo(_ParserBaseModel): confidence: float = 0.0 tables_discovered: int = 0 @@ -190,7 +224,7 @@ def error(self) -> Optional[Exception]: return self.table_error or self.column_error -class SqlParsingResult(BaseModel): +class SqlParsingResult(_ParserBaseModel): query_type: QueryType = QueryType.UNKNOWN in_tables: List[Urn] @@ -541,6 +575,15 @@ def _schema_aware_fuzzy_column_resolve( ) from e logger.debug("Qualified sql %s", statement.sql(pretty=True, dialect=dialect)) + # Try to figure out the types of the output columns. + try: + statement = sqlglot.optimizer.annotate_types.annotate_types( + statement, schema=sqlglot_db_schema + ) + except sqlglot.errors.OptimizeError as e: + # This is not a fatal error, so we can continue. + logger.debug("sqlglot failed to annotate types: %s", e) + column_lineage = [] try: @@ -553,7 +596,6 @@ def _schema_aware_fuzzy_column_resolve( logger.debug("output columns: %s", [col[0] for col in output_columns]) output_col: str for output_col, original_col_expression in output_columns: - # print(f"output column: {output_col}") if output_col == "*": # If schema information is available, the * will be expanded to the actual columns. # Otherwise, we can't process it. @@ -613,12 +655,19 @@ def _schema_aware_fuzzy_column_resolve( output_col = _schema_aware_fuzzy_column_resolve(output_table, output_col) + # Guess the output column type. + output_col_type = None + if original_col_expression.type: + output_col_type = original_col_expression.type + if not direct_col_upstreams: logger.debug(f' "{output_col}" has no upstreams') column_lineage.append( _ColumnLineageInfo( downstream=_DownstreamColumnRef( - table=output_table, column=output_col + table=output_table, + column=output_col, + column_type=output_col_type, ), upstreams=sorted(direct_col_upstreams), # logic=column_logic.sql(pretty=True, dialect=dialect), @@ -673,6 +722,42 @@ def _try_extract_select( return statement +def _translate_sqlglot_type( + sqlglot_type: sqlglot.exp.DataType.Type, +) -> Optional[SchemaFieldDataTypeClass]: + TypeClass: Any + if sqlglot_type in sqlglot.exp.DataType.TEXT_TYPES: + TypeClass = StringTypeClass + elif sqlglot_type in sqlglot.exp.DataType.NUMERIC_TYPES or sqlglot_type in { + sqlglot.exp.DataType.Type.DECIMAL, + }: + TypeClass = NumberTypeClass + elif sqlglot_type in { + sqlglot.exp.DataType.Type.BOOLEAN, + sqlglot.exp.DataType.Type.BIT, + }: + TypeClass = BooleanTypeClass + elif sqlglot_type in { + sqlglot.exp.DataType.Type.DATE, + }: + TypeClass = DateTypeClass + elif sqlglot_type in sqlglot.exp.DataType.TEMPORAL_TYPES: + TypeClass = TimeTypeClass + elif sqlglot_type in { + sqlglot.exp.DataType.Type.ARRAY, + }: + TypeClass = ArrayTypeClass + elif sqlglot_type in { + sqlglot.exp.DataType.Type.UNKNOWN, + }: + return None + else: + logger.debug("Unknown sqlglot type: %s", sqlglot_type) + return None + + return SchemaFieldDataTypeClass(type=TypeClass()) + + def _translate_internal_column_lineage( table_name_urn_mapping: Dict[_TableName, str], raw_column_lineage: _ColumnLineageInfo, @@ -684,6 +769,16 @@ def _translate_internal_column_lineage( downstream=DownstreamColumnRef( table=downstream_urn, column=raw_column_lineage.downstream.column, + column_type=_translate_sqlglot_type( + raw_column_lineage.downstream.column_type.this + ) + if raw_column_lineage.downstream.column_type + else None, + native_column_type=raw_column_lineage.downstream.column_type.sql() + if raw_column_lineage.downstream.column_type + and raw_column_lineage.downstream.column_type.this + != sqlglot.exp.DataType.Type.UNKNOWN + else None, ), upstreams=[ ColumnRef( diff --git a/metadata-ingestion/tests/integration/powerbi/test_m_parser.py b/metadata-ingestion/tests/integration/powerbi/test_m_parser.py index e3cc6c8101650..b6cb578217a2c 100644 --- a/metadata-ingestion/tests/integration/powerbi/test_m_parser.py +++ b/metadata-ingestion/tests/integration/powerbi/test_m_parser.py @@ -17,7 +17,6 @@ ) from datahub.ingestion.source.powerbi.m_query import parser, resolver, tree_function from datahub.ingestion.source.powerbi.m_query.resolver import DataPlatformTable, Lineage -from datahub.utilities.sqlglot_lineage import ColumnLineageInfo, DownstreamColumnRef pytestmark = pytest.mark.integration_batch_2 @@ -742,75 +741,25 @@ def test_sqlglot_parser(): == "urn:li:dataset:(urn:li:dataPlatform:snowflake,sales_deployment.operations_analytics.transformed_prod.v_sme_unit_targets,PROD)" ) - assert lineage[0].column_lineage == [ - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="client_director"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="tier"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column='upper("manager")'), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="team_type"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="date_target"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="monthid"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="target_team"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="seller_email"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="agent_key"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="sme_quota"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="revenue_quota"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="service_quota"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="bl_target"), - upstreams=[], - logic=None, - ), - ColumnLineageInfo( - downstream=DownstreamColumnRef(table=None, column="software_quota"), - upstreams=[], - logic=None, - ), + # TODO: None of these columns have upstreams? + # That doesn't seem right - we probably need to add fake schemas for the two tables above. + cols = [ + "client_director", + "tier", + 'upper("manager")', + "team_type", + "date_target", + "monthid", + "target_team", + "seller_email", + "agent_key", + "sme_quota", + "revenue_quota", + "service_quota", + "bl_target", + "software_quota", ] + for i, column in enumerate(cols): + assert lineage[0].column_lineage[i].downstream.table is None + assert lineage[0].column_lineage[i].downstream.column == column + assert lineage[0].column_lineage[i].upstreams == [] diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_create_view_with_cte.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_create_view_with_cte.json index e50d944ce72e3..f0175b4dc8892 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_create_view_with_cte.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_create_view_with_cte.json @@ -12,7 +12,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj-2.dataset.my_view,PROD)", - "column": "col5" + "column": "col5", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -24,7 +30,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj-2.dataset.my_view,PROD)", - "column": "col1" + "column": "col1", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -36,7 +48,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj-2.dataset.my_view,PROD)", - "column": "col2" + "column": "col2", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -48,7 +66,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj-2.dataset.my_view,PROD)", - "column": "col3" + "column": "col3", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_from_sharded_table_wildcard.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_from_sharded_table_wildcard.json index 78591286feb50..b7df5444987f2 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_from_sharded_table_wildcard.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_from_sharded_table_wildcard.json @@ -8,7 +8,13 @@ { "downstream": { "table": null, - "column": "col1" + "column": "col1", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -20,7 +26,13 @@ { "downstream": { "table": null, - "column": "col2" + "column": "col2", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_nested_subqueries.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_nested_subqueries.json index 0e93d31fbb6a6..67e306bebf545 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_nested_subqueries.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_nested_subqueries.json @@ -8,7 +8,13 @@ { "downstream": { "table": null, - "column": "col1" + "column": "col1", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -20,7 +26,13 @@ { "downstream": { "table": null, - "column": "col2" + "column": "col2", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_sharded_table_normalization.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_sharded_table_normalization.json index 78591286feb50..b7df5444987f2 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_sharded_table_normalization.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_sharded_table_normalization.json @@ -8,7 +8,13 @@ { "downstream": { "table": null, - "column": "col1" + "column": "col1", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -20,7 +26,13 @@ { "downstream": { "table": null, - "column": "col2" + "column": "col2", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_star_with_replace.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_star_with_replace.json index 17a801a63e3ff..b393b2445d6c4 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_star_with_replace.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_star_with_replace.json @@ -10,7 +10,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:bigquery,my-project.my-dataset.test_table,PROD)", - "column": "col1" + "column": "col1", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -22,7 +28,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:bigquery,my-project.my-dataset.test_table,PROD)", - "column": "col2" + "column": "col2", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -34,7 +46,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:bigquery,my-project.my-dataset.test_table,PROD)", - "column": "something" + "column": "something", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_view_from_union.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_view_from_union.json index fd8a586ac74ac..53fb94300e804 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_view_from_union.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_view_from_union.json @@ -11,7 +11,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:bigquery,my_view,PROD)", - "column": "col1" + "column": "col1", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -27,7 +33,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:bigquery,my_view,PROD)", - "column": "col2" + "column": "col2", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_create_view_as_select.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_create_view_as_select.json index 1ca56840531e4..ff452467aa5bd 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_create_view_as_select.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_create_view_as_select.json @@ -10,7 +10,9 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:oracle,vsal,PROD)", - "column": "Department" + "column": "Department", + "column_type": null, + "native_column_type": null }, "upstreams": [ { @@ -22,14 +24,22 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:oracle,vsal,PROD)", - "column": "Employees" + "column": "Employees", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "BIGINT" }, "upstreams": [] }, { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:oracle,vsal,PROD)", - "column": "Salary" + "column": "Salary", + "column_type": null, + "native_column_type": null }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_expand_select_star_basic.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_expand_select_star_basic.json index e241bdd08e243..eecb2265eaec5 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_expand_select_star_basic.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_expand_select_star_basic.json @@ -8,7 +8,13 @@ { "downstream": { "table": null, - "column": "total_agg" + "column": "total_agg", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DOUBLE" }, "upstreams": [ { @@ -20,7 +26,13 @@ { "downstream": { "table": null, - "column": "orderkey" + "column": "orderkey", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DECIMAL" }, "upstreams": [ { @@ -32,7 +44,13 @@ { "downstream": { "table": null, - "column": "custkey" + "column": "custkey", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DECIMAL" }, "upstreams": [ { @@ -44,7 +62,13 @@ { "downstream": { "table": null, - "column": "orderstatus" + "column": "orderstatus", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -56,7 +80,13 @@ { "downstream": { "table": null, - "column": "totalprice" + "column": "totalprice", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "FLOAT" }, "upstreams": [ { @@ -68,7 +98,13 @@ { "downstream": { "table": null, - "column": "orderdate" + "column": "orderdate", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.DateType": {} + } + }, + "native_column_type": "DATE" }, "upstreams": [ { @@ -80,7 +116,13 @@ { "downstream": { "table": null, - "column": "orderpriority" + "column": "orderpriority", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -92,7 +134,13 @@ { "downstream": { "table": null, - "column": "clerk" + "column": "clerk", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { @@ -104,7 +152,13 @@ { "downstream": { "table": null, - "column": "shippriority" + "column": "shippriority", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DECIMAL" }, "upstreams": [ { @@ -116,7 +170,13 @@ { "downstream": { "table": null, - "column": "comment" + "column": "comment", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "TEXT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_insert_as_select.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_insert_as_select.json index d7264fd2db6b2..326db47e7ab33 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_insert_as_select.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_insert_as_select.json @@ -18,21 +18,27 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:hive,query72,PROD)", - "column": "i_item_desc" + "column": "i_item_desc", + "column_type": null, + "native_column_type": null }, "upstreams": [] }, { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:hive,query72,PROD)", - "column": "w_warehouse_name" + "column": "w_warehouse_name", + "column_type": null, + "native_column_type": null }, "upstreams": [] }, { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:hive,query72,PROD)", - "column": "d_week_seq" + "column": "d_week_seq", + "column_type": null, + "native_column_type": null }, "upstreams": [ { @@ -44,7 +50,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:hive,query72,PROD)", - "column": "no_promo" + "column": "no_promo", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "BIGINT" }, "upstreams": [ { @@ -56,7 +68,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:hive,query72,PROD)", - "column": "promo" + "column": "promo", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "BIGINT" }, "upstreams": [ { @@ -68,7 +86,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:hive,query72,PROD)", - "column": "total_cnt" + "column": "total_cnt", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "BIGINT" }, "upstreams": [] } diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_ambiguous_column_no_schema.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_ambiguous_column_no_schema.json index 10f5ee20b0c1f..b5fd5eebeb1b1 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_ambiguous_column_no_schema.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_ambiguous_column_no_schema.json @@ -9,21 +9,27 @@ { "downstream": { "table": null, - "column": "a" + "column": "a", + "column_type": null, + "native_column_type": null }, "upstreams": [] }, { "downstream": { "table": null, - "column": "b" + "column": "b", + "column_type": null, + "native_column_type": null }, "upstreams": [] }, { "downstream": { "table": null, - "column": "c" + "column": "c", + "column_type": null, + "native_column_type": null }, "upstreams": [] } diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_count.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_count.json index 9f6eeae46c294..a67c944822138 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_count.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_count.json @@ -8,7 +8,13 @@ { "downstream": { "table": null, - "column": "COUNT(`fact_complaint_snapshot`.`etl_data_dt_id`)" + "column": "COUNT(`fact_complaint_snapshot`.`etl_data_dt_id`)", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "BIGINT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_from_struct_subfields.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_from_struct_subfields.json index 109de96180422..5ad847e252497 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_from_struct_subfields.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_from_struct_subfields.json @@ -8,7 +8,13 @@ { "downstream": { "table": null, - "column": "post_id" + "column": "post_id", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DECIMAL" }, "upstreams": [ { @@ -20,7 +26,9 @@ { "downstream": { "table": null, - "column": "id" + "column": "id", + "column_type": null, + "native_column_type": null }, "upstreams": [ { @@ -32,7 +40,9 @@ { "downstream": { "table": null, - "column": "min_metric" + "column": "min_metric", + "column_type": null, + "native_column_type": null }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_from_union.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_from_union.json index 2340b2e95b0d0..902aa010c8afc 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_from_union.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_from_union.json @@ -9,14 +9,26 @@ { "downstream": { "table": null, - "column": "label" + "column": "label", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "VARCHAR" }, "upstreams": [] }, { "downstream": { "table": null, - "column": "total_agg" + "column": "total_agg", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DOUBLE" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_max.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_max.json index 326c07d332c26..6ea88f45847ce 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_max.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_max.json @@ -8,7 +8,9 @@ { "downstream": { "table": null, - "column": "max_col" + "column": "max_col", + "column_type": null, + "native_column_type": null }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_with_ctes.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_with_ctes.json index 3e02314d6e8c3..67e9fd2d21a0e 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_with_ctes.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_with_ctes.json @@ -9,7 +9,9 @@ { "downstream": { "table": null, - "column": "COL1" + "column": "COL1", + "column_type": null, + "native_column_type": null }, "upstreams": [ { @@ -21,7 +23,9 @@ { "downstream": { "table": null, - "column": "COL3" + "column": "COL3", + "column_type": null, + "native_column_type": null }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_with_full_col_name.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_with_full_col_name.json index c12ad23b2f03b..6ee3d2e61c39b 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_with_full_col_name.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_select_with_full_col_name.json @@ -8,7 +8,13 @@ { "downstream": { "table": null, - "column": "post_id" + "column": "post_id", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DECIMAL" }, "upstreams": [ { @@ -20,7 +26,9 @@ { "downstream": { "table": null, - "column": "id" + "column": "id", + "column_type": null, + "native_column_type": null }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_case_statement.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_case_statement.json index 64cd80e9a2d69..a876824127ec1 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_case_statement.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_case_statement.json @@ -8,7 +8,13 @@ { "downstream": { "table": null, - "column": "total_price_category" + "column": "total_price_category", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "VARCHAR" }, "upstreams": [ { @@ -20,7 +26,13 @@ { "downstream": { "table": null, - "column": "total_price_success" + "column": "total_price_success", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "FLOAT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_column_cast.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_column_cast.json new file mode 100644 index 0000000000000..7545e2b3269dc --- /dev/null +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_column_cast.json @@ -0,0 +1,63 @@ +{ + "query_type": "SELECT", + "in_tables": [ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders,PROD)" + ], + "out_tables": [], + "column_lineage": [ + { + "downstream": { + "table": null, + "column": "orderkey", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DECIMAL(20, 0)" + }, + "upstreams": [ + { + "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders,PROD)", + "column": "o_orderkey" + } + ] + }, + { + "downstream": { + "table": null, + "column": "total_cast_int", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "INT" + }, + "upstreams": [ + { + "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders,PROD)", + "column": "o_totalprice" + } + ] + }, + { + "downstream": { + "table": null, + "column": "total_cast_float", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DECIMAL(16, 4)" + }, + "upstreams": [ + { + "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders,PROD)", + "column": "o_totalprice" + } + ] + } + ] +} \ No newline at end of file diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_column_normalization.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_column_normalization.json index 7b22a46757e39..84e6b053000f1 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_column_normalization.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_column_normalization.json @@ -8,7 +8,13 @@ { "downstream": { "table": null, - "column": "total_agg" + "column": "total_agg", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DOUBLE" }, "upstreams": [ { @@ -20,7 +26,13 @@ { "downstream": { "table": null, - "column": "total_avg" + "column": "total_avg", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DOUBLE" }, "upstreams": [ { @@ -32,7 +44,13 @@ { "downstream": { "table": null, - "column": "total_min" + "column": "total_min", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "FLOAT" }, "upstreams": [ { @@ -44,7 +62,13 @@ { "downstream": { "table": null, - "column": "total_max" + "column": "total_max", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "FLOAT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_ctas_column_normalization.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_ctas_column_normalization.json index c912d99a3a8a3..39c94cf83c561 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_ctas_column_normalization.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_ctas_column_normalization.json @@ -10,7 +10,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders_normalized,PROD)", - "column": "Total_Agg" + "column": "Total_Agg", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DOUBLE" }, "upstreams": [ { @@ -22,7 +28,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders_normalized,PROD)", - "column": "total_avg" + "column": "total_avg", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DOUBLE" }, "upstreams": [ { @@ -34,7 +46,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders_normalized,PROD)", - "column": "TOTAL_MIN" + "column": "TOTAL_MIN", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "FLOAT" }, "upstreams": [ { @@ -46,7 +64,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders_normalized,PROD)", - "column": "total_max" + "column": "total_max", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "FLOAT" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_default_normalization.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_default_normalization.json index 2af308ec60623..dbf5b1b9a4453 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_default_normalization.json +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_default_normalization.json @@ -11,7 +11,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.active_customer_ltv,PROD)", - "column": "user_fk" + "column": "user_fk", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DECIMAL(38, 0)" }, "upstreams": [ { @@ -23,7 +29,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.active_customer_ltv,PROD)", - "column": "email" + "column": "email", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.StringType": {} + } + }, + "native_column_type": "VARCHAR(16777216)" }, "upstreams": [ { @@ -35,7 +47,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.active_customer_ltv,PROD)", - "column": "last_purchase_date" + "column": "last_purchase_date", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.DateType": {} + } + }, + "native_column_type": "DATE" }, "upstreams": [ { @@ -47,7 +65,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.active_customer_ltv,PROD)", - "column": "lifetime_purchase_amount" + "column": "lifetime_purchase_amount", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DECIMAL" }, "upstreams": [ { @@ -59,7 +83,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.active_customer_ltv,PROD)", - "column": "lifetime_purchase_count" + "column": "lifetime_purchase_count", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "BIGINT" }, "upstreams": [ { @@ -71,7 +101,13 @@ { "downstream": { "table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.active_customer_ltv,PROD)", - "column": "average_purchase_amount" + "column": "average_purchase_amount", + "column_type": { + "type": { + "com.linkedin.pegasus2avro.schema.NumberType": {} + } + }, + "native_column_type": "DECIMAL" }, "upstreams": [ { diff --git a/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py b/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py index 2a965a9bb1e61..bb6e5f1581754 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py +++ b/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py @@ -608,4 +608,25 @@ def test_snowflake_default_normalization(): ) +def test_snowflake_column_cast(): + assert_sql_result( + """ +SELECT + o.o_orderkey::NUMBER(20,0) as orderkey, + CAST(o.o_totalprice AS INT) as total_cast_int, + CAST(o.o_totalprice AS NUMBER(16,4)) as total_cast_float +FROM snowflake_sample_data.tpch_sf1.orders o +LIMIT 10 +""", + dialect="snowflake", + schemas={ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders,PROD)": { + "orderkey": "NUMBER(38,0)", + "totalprice": "NUMBER(12,2)", + }, + }, + expected_file=RESOURCE_DIR / "test_snowflake_column_cast.json", + ) + + # TODO: Add a test for setting platform_instance or env From dd418de76d96fb41c9064261cdba37bc2af85309 Mon Sep 17 00:00:00 2001 From: Tamas Nemeth Date: Thu, 12 Oct 2023 13:10:59 +0200 Subject: [PATCH 120/156] fix(ingest/bigquery): Fix shard regexp to match without underscore as well (#8934) --- .../ingestion/source/bigquery_v2/bigquery.py | 1 + .../source/bigquery_v2/bigquery_audit.py | 27 ++++++++++++++----- .../ingestion/source/bigquery_v2/queries.py | 8 +++--- .../ingestion/source_config/bigquery.py | 8 +++++- .../tests/unit/test_bigquery_source.py | 10 ++++--- .../unit/test_bigqueryv2_usage_source.py | 4 +-- 6 files changed, 41 insertions(+), 17 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py index b4a04d96b532b..e577c2bac8bbd 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py @@ -1057,6 +1057,7 @@ def gen_schema_fields(self, columns: List[BigqueryColumn]) -> List[SchemaField]: ): field.description = col.comment schema_fields[idx] = field + break else: tags = [] if col.is_partition_column: diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit.py index b0ac77201b415..88060a9cdc91d 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit.py @@ -20,7 +20,13 @@ logger: logging.Logger = logging.getLogger(__name__) -_BIGQUERY_DEFAULT_SHARDED_TABLE_REGEX = "((.+)[_$])?(\\d{8})$" +# Regexp for sharded tables. +# A sharded table is a table that has a suffix of the form _yyyymmdd or yyyymmdd, where yyyymmdd is a date. +# The regexp checks for valid dates in the suffix (e.g. 20200101, 20200229, 20201231) and if the date is not valid +# then it is not a sharded table. +_BIGQUERY_DEFAULT_SHARDED_TABLE_REGEX = ( + "((.+\\D)[_$]?)?(\\d\\d\\d\\d(?:0[1-9]|1[0-2])(?:0[1-9]|[12][0-9]|3[01]))$" +) @dataclass(frozen=True, order=True) @@ -40,7 +46,7 @@ class BigqueryTableIdentifier: _BQ_SHARDED_TABLE_SUFFIX: str = "_yyyymmdd" @staticmethod - def get_table_and_shard(table_name: str) -> Tuple[str, Optional[str]]: + def get_table_and_shard(table_name: str) -> Tuple[Optional[str], Optional[str]]: """ Args: table_name: @@ -53,16 +59,25 @@ def get_table_and_shard(table_name: str) -> Tuple[str, Optional[str]]: In case of non-sharded tables, returns (, None) In case of sharded tables, returns (, shard) """ + new_table_name = table_name match = re.match( BigqueryTableIdentifier._BIGQUERY_DEFAULT_SHARDED_TABLE_REGEX, table_name, re.IGNORECASE, ) if match: - table_name = match.group(2) - shard = match.group(3) - return table_name, shard - return table_name, None + shard: str = match[3] + if shard: + if table_name.endswith(shard): + new_table_name = table_name[: -len(shard)] + + new_table_name = ( + new_table_name.rstrip("_") if new_table_name else new_table_name + ) + if new_table_name.endswith("."): + new_table_name = table_name + return (new_table_name, shard) if new_table_name else (None, shard) + return new_table_name, None @classmethod def from_string_name(cls, table: str) -> "BigqueryTableIdentifier": diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries.py index a87cb8c1cbfa5..67fcc33cdf218 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries.py @@ -51,8 +51,8 @@ class BigqueryQuery: p.max_partition_id, p.active_billable_bytes, p.long_term_billable_bytes, - REGEXP_EXTRACT(t.table_name, r".*_(\\d+)$") as table_suffix, - REGEXP_REPLACE(t.table_name, r"_(\\d+)$", "") as table_base + REGEXP_EXTRACT(t.table_name, r"(?:(?:.+\\D)[_$]?)(\\d\\d\\d\\d(?:0[1-9]|1[012])(?:0[1-9]|[12][0-9]|3[01]))$") as table_suffix, + REGEXP_REPLACE(t.table_name, r"(?:[_$]?)(\\d\\d\\d\\d(?:0[1-9]|1[012])(?:0[1-9]|[12][0-9]|3[01]))$", "") as table_base FROM `{{project_id}}`.`{{dataset_name}}`.INFORMATION_SCHEMA.TABLES t @@ -92,8 +92,8 @@ class BigqueryQuery: tos.OPTION_VALUE as comment, t.is_insertable_into, t.ddl, - REGEXP_EXTRACT(t.table_name, r".*_(\\d+)$") as table_suffix, - REGEXP_REPLACE(t.table_name, r"_(\\d+)$", "") as table_base + REGEXP_EXTRACT(t.table_name, r"(?:(?:.+\\D)[_$]?)(\\d\\d\\d\\d(?:0[1-9]|1[012])(?:0[1-9]|[12][0-9]|3[01]))$") as table_suffix, + REGEXP_REPLACE(t.table_name, r"(?:[_$]?)(\\d\\d\\d\\d(?:0[1-9]|1[012])(?:0[1-9]|[12][0-9]|3[01]))$", "") as table_base FROM `{{project_id}}`.`{{dataset_name}}`.INFORMATION_SCHEMA.TABLES t diff --git a/metadata-ingestion/src/datahub/ingestion/source_config/bigquery.py b/metadata-ingestion/src/datahub/ingestion/source_config/bigquery.py index 8ca1296d819c1..0a73bb5203e72 100644 --- a/metadata-ingestion/src/datahub/ingestion/source_config/bigquery.py +++ b/metadata-ingestion/src/datahub/ingestion/source_config/bigquery.py @@ -4,7 +4,13 @@ from datahub.configuration.common import ConfigModel, ConfigurationError -_BIGQUERY_DEFAULT_SHARDED_TABLE_REGEX: str = "((.+)[_$])?(\\d{8})$" +# Regexp for sharded tables. +# A sharded table is a table that has a suffix of the form _yyyymmdd or yyyymmdd, where yyyymmdd is a date. +# The regexp checks for valid dates in the suffix (e.g. 20200101, 20200229, 20201231) and if the date is not valid +# then it is not a sharded table. +_BIGQUERY_DEFAULT_SHARDED_TABLE_REGEX: str = ( + "((.+\\D)[_$]?)?(\\d\\d\\d\\d(?:0[1-9]|1[0-2])(?:0[1-9]|[12][0-9]|3[01]))$" +) class BigQueryBaseConfig(ConfigModel): diff --git a/metadata-ingestion/tests/unit/test_bigquery_source.py b/metadata-ingestion/tests/unit/test_bigquery_source.py index e9e91361f49f4..5a11a933c8595 100644 --- a/metadata-ingestion/tests/unit/test_bigquery_source.py +++ b/metadata-ingestion/tests/unit/test_bigquery_source.py @@ -765,11 +765,14 @@ def test_gen_view_dataset_workunits( ("project.dataset.table_20231215", "project.dataset.table", "20231215"), ("project.dataset.table_2023", "project.dataset.table_2023", None), # incorrectly handled special case where dataset itself is a sharded table if full name is specified - ("project.dataset.20231215", "project.dataset.20231215", None), + ("project.dataset.20231215", "project.dataset.20231215", "20231215"), + ("project1.dataset2.20231215", "project1.dataset2.20231215", "20231215"), # Cases with Just the table name as input ("table", "table", None), - ("table20231215", "table20231215", None), + ("table20231215", "table", "20231215"), ("table_20231215", "table", "20231215"), + ("table2_20231215", "table2", "20231215"), + ("table220231215", "table220231215", None), ("table_1624046611000_name", "table_1624046611000_name", None), ("table_1624046611000", "table_1624046611000", None), # Special case where dataset itself is a sharded table @@ -801,7 +804,6 @@ def test_get_table_and_shard_default( ("project.dataset.2023", "project.dataset.2023", None), # Cases with Just the table name as input ("table", "table", None), - ("table20231215", "table20231215", None), ("table_20231215", "table", "20231215"), ("table_2023", "table", "2023"), ("table_1624046611000_name", "table_1624046611000_name", None), @@ -842,7 +844,7 @@ def test_get_table_and_shard_custom_shard_pattern( "project.dataset.table_1624046611000_name", ), ("project.dataset.table_1624046611000", "project.dataset.table_1624046611000"), - ("project.dataset.table20231215", "project.dataset.table20231215"), + ("project.dataset.table20231215", "project.dataset.table"), ("project.dataset.table_*", "project.dataset.table"), ("project.dataset.table_2023*", "project.dataset.table"), ("project.dataset.table_202301*", "project.dataset.table"), diff --git a/metadata-ingestion/tests/unit/test_bigqueryv2_usage_source.py b/metadata-ingestion/tests/unit/test_bigqueryv2_usage_source.py index 4cf42da4395f9..44fd840f28d59 100644 --- a/metadata-ingestion/tests/unit/test_bigqueryv2_usage_source.py +++ b/metadata-ingestion/tests/unit/test_bigqueryv2_usage_source.py @@ -144,10 +144,10 @@ def test_bigquery_table_sanitasitation(): assert new_table_ref.dataset == "dataset-4567" table_ref = BigQueryTableRef( - BigqueryTableIdentifier("project-1234", "dataset-4567", "foo_20222110") + BigqueryTableIdentifier("project-1234", "dataset-4567", "foo_20221210") ) new_table_identifier = table_ref.table_identifier - assert new_table_identifier.table == "foo_20222110" + assert new_table_identifier.table == "foo_20221210" assert new_table_identifier.is_sharded_table() assert new_table_identifier.get_table_display_name() == "foo" assert new_table_identifier.project_id == "project-1234" From c381806110ae995dd2164305394ee4e1d131e033 Mon Sep 17 00:00:00 2001 From: Tamas Nemeth Date: Thu, 12 Oct 2023 13:56:30 +0200 Subject: [PATCH 121/156] feat(ingestion): Adding config option to auto lowercase dataset urns (#8928) --- .../datahub/configuration/source_common.py | 7 ++ .../src/datahub/ingestion/api/source.py | 24 +++++++ .../datahub/ingestion/api/source_helpers.py | 20 +++++- .../ingestion/source/bigquery_v2/bigquery.py | 3 - .../source/bigquery_v2/bigquery_config.py | 5 -- .../src/datahub/ingestion/source/kafka.py | 11 ++- .../ingestion/source/sql/sql_config.py | 11 ++- .../datahub/ingestion/source/unity/config.py | 6 +- .../src/datahub/utilities/urns/urn_iter.py | 33 +++++++-- .../api/source_helpers/test_source_helpers.py | 70 +++++++++++++++++++ 10 files changed, 170 insertions(+), 20 deletions(-) diff --git a/metadata-ingestion/src/datahub/configuration/source_common.py b/metadata-ingestion/src/datahub/configuration/source_common.py index a9f891ddb7b1e..80b6ceb576c1c 100644 --- a/metadata-ingestion/src/datahub/configuration/source_common.py +++ b/metadata-ingestion/src/datahub/configuration/source_common.py @@ -54,6 +54,13 @@ class DatasetSourceConfigMixin(PlatformInstanceConfigMixin, EnvConfigMixin): """ +class LowerCaseDatasetUrnConfigMixin(ConfigModel): + convert_urns_to_lowercase: bool = Field( + default=False, + description="Whether to convert dataset urns to lowercase.", + ) + + class DatasetLineageProviderConfigBase(EnvConfigMixin): """ Any non-Dataset source that produces lineage to Datasets should inherit this class. diff --git a/metadata-ingestion/src/datahub/ingestion/api/source.py b/metadata-ingestion/src/datahub/ingestion/api/source.py index 0bcc220cad49b..b86844b1c4c83 100644 --- a/metadata-ingestion/src/datahub/ingestion/api/source.py +++ b/metadata-ingestion/src/datahub/ingestion/api/source.py @@ -29,6 +29,7 @@ from datahub.ingestion.api.report import Report from datahub.ingestion.api.source_helpers import ( auto_browse_path_v2, + auto_lowercase_urns, auto_materialize_referenced_tags, auto_status_aspect, auto_workunit_reporter, @@ -192,7 +193,30 @@ def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]: self.ctx.pipeline_config.flags.generate_browse_path_v2_dry_run ) + auto_lowercase_dataset_urns: Optional[MetadataWorkUnitProcessor] = None + if ( + self.ctx.pipeline_config + and self.ctx.pipeline_config.source + and self.ctx.pipeline_config.source.config + and ( + ( + hasattr( + self.ctx.pipeline_config.source.config, + "convert_urns_to_lowercase", + ) + and self.ctx.pipeline_config.source.config.convert_urns_to_lowercase + ) + or ( + hasattr(self.ctx.pipeline_config.source.config, "get") + and self.ctx.pipeline_config.source.config.get( + "convert_urns_to_lowercase" + ) + ) + ) + ): + auto_lowercase_dataset_urns = auto_lowercase_urns return [ + auto_lowercase_dataset_urns, auto_status_aspect, auto_materialize_referenced_tags, browse_path_processor, diff --git a/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py b/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py index 7fc15cf829678..2ce9e07bc57bc 100644 --- a/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py +++ b/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py @@ -35,7 +35,7 @@ from datahub.utilities.urns.dataset_urn import DatasetUrn from datahub.utilities.urns.tag_urn import TagUrn from datahub.utilities.urns.urn import guess_entity_type -from datahub.utilities.urns.urn_iter import list_urns +from datahub.utilities.urns.urn_iter import list_urns, lowercase_dataset_urns if TYPE_CHECKING: from datahub.ingestion.api.source import SourceReport @@ -70,7 +70,6 @@ def auto_status_aspect( for wu in stream: urn = wu.get_urn() all_urns.add(urn) - if not wu.is_primary_source: # If this is a non-primary source, we pretend like we've seen the status # aspect so that we don't try to emit a removal for it. @@ -173,6 +172,23 @@ def auto_materialize_referenced_tags( ).as_workunit() +def auto_lowercase_urns( + stream: Iterable[MetadataWorkUnit], +) -> Iterable[MetadataWorkUnit]: + """Lowercase all dataset urns""" + + for wu in stream: + try: + old_urn = wu.get_urn() + lowercase_dataset_urns(wu.metadata) + wu.id = wu.id.replace(old_urn, wu.get_urn()) + + yield wu + except Exception as e: + logger.warning(f"Failed to lowercase urns for {wu}: {e}", exc_info=True) + yield wu + + def auto_browse_path_v2( stream: Iterable[MetadataWorkUnit], *, diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py index e577c2bac8bbd..552612f877b9a 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py @@ -16,7 +16,6 @@ make_dataplatform_instance_urn, make_dataset_urn, make_tag_urn, - set_dataset_urn_to_lower, ) from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.mcp_builder import BigQueryDatasetKey, ContainerKey, ProjectIdKey @@ -218,8 +217,6 @@ def __init__(self, ctx: PipelineContext, config: BigQueryV2Config): if self.config.enable_legacy_sharded_table_support: BigqueryTableIdentifier._BQ_SHARDED_TABLE_SUFFIX = "" - set_dataset_urn_to_lower(self.config.convert_urns_to_lowercase) - self.bigquery_data_dictionary = BigQuerySchemaApi( self.report.schema_api_perf, self.config.get_bigquery_client() ) diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py index 483355a85ac05..944814b6936a4 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py @@ -206,11 +206,6 @@ def validate_column_lineage(cls, v: bool, values: Dict[str, Any]) -> bool: description="This flag enables the data lineage extraction from Data Lineage API exposed by Google Data Catalog. NOTE: This extractor can't build views lineage. It's recommended to enable the view's DDL parsing. Read the docs to have more information about: https://cloud.google.com/data-catalog/docs/concepts/about-data-lineage", ) - convert_urns_to_lowercase: bool = Field( - default=False, - description="Convert urns to lowercase.", - ) - enable_legacy_sharded_table_support: bool = Field( default=True, description="Use the legacy sharded table urn suffix added.", diff --git a/metadata-ingestion/src/datahub/ingestion/source/kafka.py b/metadata-ingestion/src/datahub/ingestion/source/kafka.py index 566304e1999b7..d5039360da567 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/kafka.py +++ b/metadata-ingestion/src/datahub/ingestion/source/kafka.py @@ -18,7 +18,10 @@ from datahub.configuration.common import AllowDenyPattern from datahub.configuration.kafka import KafkaConsumerConnectionConfig -from datahub.configuration.source_common import DatasetSourceConfigMixin +from datahub.configuration.source_common import ( + DatasetSourceConfigMixin, + LowerCaseDatasetUrnConfigMixin, +) from datahub.emitter import mce_builder from datahub.emitter.mce_builder import ( make_data_platform_urn, @@ -76,7 +79,11 @@ class KafkaTopicConfigKeys(str, Enum): UNCLEAN_LEADER_ELECTION_CONFIG = "unclean.leader.election.enable" -class KafkaSourceConfig(StatefulIngestionConfigBase, DatasetSourceConfigMixin): +class KafkaSourceConfig( + StatefulIngestionConfigBase, + DatasetSourceConfigMixin, + LowerCaseDatasetUrnConfigMixin, +): connection: KafkaConsumerConnectionConfig = KafkaConsumerConnectionConfig() topic_patterns: AllowDenyPattern = AllowDenyPattern(allow=[".*"], deny=["^_.*"]) diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_config.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_config.py index 677d32c8bac08..08cc74aec3977 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_config.py @@ -7,7 +7,10 @@ from pydantic import Field from datahub.configuration.common import AllowDenyPattern, ConfigModel -from datahub.configuration.source_common import DatasetSourceConfigMixin +from datahub.configuration.source_common import ( + DatasetSourceConfigMixin, + LowerCaseDatasetUrnConfigMixin, +) from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig from datahub.ingestion.source.state.stale_entity_removal_handler import ( @@ -21,7 +24,11 @@ logger: logging.Logger = logging.getLogger(__name__) -class SQLCommonConfig(StatefulIngestionConfigBase, DatasetSourceConfigMixin): +class SQLCommonConfig( + StatefulIngestionConfigBase, + DatasetSourceConfigMixin, + LowerCaseDatasetUrnConfigMixin, +): options: dict = pydantic.Field( default_factory=dict, description="Any options specified here will be passed to [SQLAlchemy.create_engine](https://docs.sqlalchemy.org/en/14/core/engines.html#sqlalchemy.create_engine) as kwargs.", diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py index 51390873712d3..a57ee39848855 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py @@ -7,7 +7,10 @@ from pydantic import Field from datahub.configuration.common import AllowDenyPattern, ConfigModel -from datahub.configuration.source_common import DatasetSourceConfigMixin +from datahub.configuration.source_common import ( + DatasetSourceConfigMixin, + LowerCaseDatasetUrnConfigMixin, +) from datahub.configuration.validate_field_removal import pydantic_removed_field from datahub.configuration.validate_field_rename import pydantic_renamed_field from datahub.ingestion.source.state.stale_entity_removal_handler import ( @@ -91,6 +94,7 @@ class UnityCatalogSourceConfig( BaseUsageConfig, DatasetSourceConfigMixin, StatefulProfilingConfigMixin, + LowerCaseDatasetUrnConfigMixin, ): token: str = pydantic.Field(description="Databricks personal access token") workspace_url: str = pydantic.Field( diff --git a/metadata-ingestion/src/datahub/utilities/urns/urn_iter.py b/metadata-ingestion/src/datahub/utilities/urns/urn_iter.py index 261f95331af61..e13d439161064 100644 --- a/metadata-ingestion/src/datahub/utilities/urns/urn_iter.py +++ b/metadata-ingestion/src/datahub/utilities/urns/urn_iter.py @@ -3,7 +3,11 @@ from avro.schema import Field, RecordSchema from datahub.emitter.mcp import MetadataChangeProposalWrapper -from datahub.metadata.schema_classes import DictWrapper +from datahub.metadata.schema_classes import ( + DictWrapper, + MetadataChangeEventClass, + MetadataChangeProposalClass, +) from datahub.utilities.urns.dataset_urn import DatasetUrn from datahub.utilities.urns.urn import Urn, guess_entity_type @@ -32,7 +36,7 @@ def list_urns_with_path( if isinstance(model, MetadataChangeProposalWrapper): if model.entityUrn: - urns.append((model.entityUrn, ["urn"])) + urns.append((model.entityUrn, ["entityUrn"])) if model.entityKeyAspect: urns.extend( _add_prefix_to_paths( @@ -83,7 +87,15 @@ def list_urns(model: Union[DictWrapper, MetadataChangeProposalWrapper]) -> List[ return [urn for urn, _ in list_urns_with_path(model)] -def transform_urns(model: DictWrapper, func: Callable[[str], str]) -> None: +def transform_urns( + model: Union[ + DictWrapper, + MetadataChangeEventClass, + MetadataChangeProposalClass, + MetadataChangeProposalWrapper, + ], + func: Callable[[str], str], +) -> None: """ Rewrites all URNs in the given object according to the given function. """ @@ -95,7 +107,9 @@ def transform_urns(model: DictWrapper, func: Callable[[str], str]) -> None: def _modify_at_path( - model: Union[DictWrapper, list], path: _Path, new_value: str + model: Union[DictWrapper, MetadataChangeProposalWrapper, list], + path: _Path, + new_value: str, ) -> None: assert len(path) > 0 @@ -103,6 +117,8 @@ def _modify_at_path( if isinstance(path[0], int): assert isinstance(model, list) model[path[0]] = new_value + elif isinstance(model, MetadataChangeProposalWrapper): + setattr(model, path[0], new_value) else: assert isinstance(model, DictWrapper) model._inner_dict[path[0]] = new_value @@ -120,7 +136,14 @@ def _lowercase_dataset_urn(dataset_urn: str) -> str: return str(cur_urn) -def lowercase_dataset_urns(model: DictWrapper) -> None: +def lowercase_dataset_urns( + model: Union[ + DictWrapper, + MetadataChangeEventClass, + MetadataChangeProposalClass, + MetadataChangeProposalWrapper, + ] +) -> None: def modify_urn(urn: str) -> str: if guess_entity_type(urn) == "dataset": return _lowercase_dataset_urn(urn) diff --git a/metadata-ingestion/tests/unit/api/source_helpers/test_source_helpers.py b/metadata-ingestion/tests/unit/api/source_helpers/test_source_helpers.py index b6ec6ebce240c..b667af8bb41e9 100644 --- a/metadata-ingestion/tests/unit/api/source_helpers/test_source_helpers.py +++ b/metadata-ingestion/tests/unit/api/source_helpers/test_source_helpers.py @@ -16,6 +16,7 @@ from datahub.ingestion.api.source_helpers import ( auto_browse_path_v2, auto_empty_dataset_usage_statistics, + auto_lowercase_urns, auto_status_aspect, auto_workunit, ) @@ -275,6 +276,75 @@ def test_auto_browse_path_v2_legacy_browse_path(telemetry_ping_mock): assert paths["platform,dataset-2,PROD)"] == _make_browse_path_entries(["something"]) +def test_auto_lowercase_aspects(): + mcws = auto_workunit( + [ + MetadataChangeProposalWrapper( + entityUrn=make_dataset_urn( + "bigquery", "myProject.mySchema.myTable", "PROD" + ), + aspect=models.DatasetKeyClass( + "urn:li:dataPlatform:bigquery", "myProject.mySchema.myTable", "PROD" + ), + ), + MetadataChangeProposalWrapper( + entityUrn="urn:li:container:008e111aa1d250dd52e0fd5d4b307b1a", + aspect=models.ContainerPropertiesClass( + name="test", + ), + ), + models.MetadataChangeEventClass( + proposedSnapshot=models.DatasetSnapshotClass( + urn="urn:li:dataset:(urn:li:dataPlatform:bigquery,bigquery-Public-Data.Covid19_Aha.staffing,PROD)", + aspects=[ + models.DatasetPropertiesClass( + customProperties={ + "key": "value", + }, + ), + ], + ), + ), + ] + ) + + expected = [ + *list( + auto_workunit( + [ + MetadataChangeProposalWrapper( + entityUrn="urn:li:dataset:(urn:li:dataPlatform:bigquery,myproject.myschema.mytable,PROD)", + aspect=models.DatasetKeyClass( + "urn:li:dataPlatform:bigquery", + "myProject.mySchema.myTable", + "PROD", + ), + ), + MetadataChangeProposalWrapper( + entityUrn="urn:li:container:008e111aa1d250dd52e0fd5d4b307b1a", + aspect=models.ContainerPropertiesClass( + name="test", + ), + ), + models.MetadataChangeEventClass( + proposedSnapshot=models.DatasetSnapshotClass( + urn="urn:li:dataset:(urn:li:dataPlatform:bigquery,bigquery-public-data.covid19_aha.staffing,PROD)", + aspects=[ + models.DatasetPropertiesClass( + customProperties={ + "key": "value", + }, + ), + ], + ), + ), + ] + ) + ), + ] + assert list(auto_lowercase_urns(mcws)) == expected + + @patch("datahub.ingestion.api.source_helpers.telemetry.telemetry_instance.ping") def test_auto_browse_path_v2_container_over_legacy_browse_path(telemetry_ping_mock): structure = {"a": {"b": ["c"]}} From 8813ae2fb15a1f80d5f0ef433fce1f84e1a240b5 Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Thu, 12 Oct 2023 07:58:10 -0400 Subject: [PATCH 122/156] feat(ingest/s3): support .gzip and fix decompression bug (#8990) --- .../ingestion/source/data_lake_common/path_spec.py | 9 ++++++++- .../src/datahub/ingestion/source/s3/source.py | 8 +++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/data_lake_common/path_spec.py b/metadata-ingestion/src/datahub/ingestion/source/data_lake_common/path_spec.py index d1c949f48e2cd..a35fb94614f72 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/data_lake_common/path_spec.py +++ b/metadata-ingestion/src/datahub/ingestion/source/data_lake_common/path_spec.py @@ -18,7 +18,14 @@ logger: logging.Logger = logging.getLogger(__name__) SUPPORTED_FILE_TYPES: List[str] = ["csv", "tsv", "json", "parquet", "avro"] -SUPPORTED_COMPRESSIONS: List[str] = ["gz", "bz2"] + +# These come from the smart_open library. +SUPPORTED_COMPRESSIONS: List[str] = [ + "gz", + "bz2", + # We have a monkeypatch on smart_open that aliases .gzip to .gz. + "gzip", +] class PathSpec(ConfigModel): diff --git a/metadata-ingestion/src/datahub/ingestion/source/s3/source.py b/metadata-ingestion/src/datahub/ingestion/source/s3/source.py index ac4433b7eb1f0..eb49fcbb268c0 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/s3/source.py +++ b/metadata-ingestion/src/datahub/ingestion/source/s3/source.py @@ -10,6 +10,7 @@ from pathlib import PurePath from typing import Any, Dict, Iterable, List, Optional, Tuple +import smart_open.compression as so_compression from more_itertools import peekable from pyspark.conf import SparkConf from pyspark.sql import SparkSession @@ -120,6 +121,9 @@ } PAGE_SIZE = 1000 +# Hack to support the .gzip extension with smart_open. +so_compression.register_compressor(".gzip", so_compression._COMPRESSOR_REGISTRY[".gz"]) + def get_column_type( report: SourceReport, dataset_name: str, column_type: str @@ -407,7 +411,9 @@ def get_fields(self, table_data: TableData, path_spec: PathSpec) -> List: table_data.full_path, "rb", transport_params={"client": s3_client} ) else: - file = open(table_data.full_path, "rb") + # We still use smart_open here to take advantage of the compression + # capabilities of smart_open. + file = smart_open(table_data.full_path, "rb") fields = [] From f6e131206394e1f56e4f966689c8abd1e8641919 Mon Sep 17 00:00:00 2001 From: Pedro Silva Date: Thu, 12 Oct 2023 18:43:14 +0100 Subject: [PATCH 123/156] feat(ingestion): Adds support for memory profiling (#8856) Co-authored-by: Harshal Sheth --- docs-website/sidebars.js | 1 + .../docs/dev_guides/profiling_ingestions.md | 55 +++++++ metadata-ingestion/setup.py | 5 + .../src/datahub/ingestion/run/pipeline.py | 148 ++++++++++-------- .../datahub/ingestion/run/pipeline_config.py | 7 + 5 files changed, 148 insertions(+), 68 deletions(-) create mode 100644 metadata-ingestion/docs/dev_guides/profiling_ingestions.md diff --git a/docs-website/sidebars.js b/docs-website/sidebars.js index bdf3926c17e0d..21b3a1d3fe4d3 100644 --- a/docs-website/sidebars.js +++ b/docs-website/sidebars.js @@ -140,6 +140,7 @@ module.exports = { "metadata-ingestion/docs/dev_guides/classification", "metadata-ingestion/docs/dev_guides/add_stateful_ingestion_to_source", "metadata-ingestion/docs/dev_guides/sql_profiles", + "metadata-ingestion/docs/dev_guides/profiling_ingestions", ], }, ], diff --git a/metadata-ingestion/docs/dev_guides/profiling_ingestions.md b/metadata-ingestion/docs/dev_guides/profiling_ingestions.md new file mode 100644 index 0000000000000..d876d99b494f8 --- /dev/null +++ b/metadata-ingestion/docs/dev_guides/profiling_ingestions.md @@ -0,0 +1,55 @@ +import FeatureAvailability from '@site/src/components/FeatureAvailability'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Profiling ingestions + + + +**🤝 Version compatibility** +> Open Source DataHub: **0.11.1** | Acryl: **0.2.12** + +This page documents how to perform memory profiles of ingestion runs. +It is useful when trying to size the amount of resources necessary to ingest some source or when developing new features or sources. + +## How to use +Install the `debug` plugin for DataHub's CLI wherever the ingestion runs: + +```bash +pip install 'acryl-datahub[debug]' +``` + +This will install [memray](https://github.com/bloomberg/memray) in your python environment. + +Add a flag to your ingestion recipe to generate a memray memory dump of your ingestion: +```yaml +source: + ... + +sink: + ... + +flags: + generate_memory_profiles: "" +``` + +Once the ingestion run starts a binary file will be created and appended to during the execution of the ingestion. + +These files follow the pattern `file-.bin` for a unique identification. +Once the ingestion has finished you can use `memray` to analyze the memory dump in a flamegraph view using: + +```$ memray flamegraph file-None-file-2023_09_18-21_38_43.bin``` + +This will generate an interactive HTML file for analysis: + +

+ +

+ + +`memray` has an extensive set of features for memory investigation. Take a look at their [documentation](https://bloomberg.github.io/memray/overview.html) to see the full feature set. + + +## Questions + +If you've got any questions on configuring profiling, feel free to ping us on [our Slack](https://slack.datahubproject.io/)! diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py index fe8e3be4632c4..61e7b684682a4 100644 --- a/metadata-ingestion/setup.py +++ b/metadata-ingestion/setup.py @@ -431,6 +431,10 @@ deepdiff_dep = "deepdiff" test_api_requirements = {pytest_dep, deepdiff_dep, "PyYAML"} +debug_requirements = { + "memray" +} + base_dev_requirements = { *base_requirements, *framework_common, @@ -723,5 +727,6 @@ "dev": list(dev_requirements), "testing-utils": list(test_api_requirements), # To import `datahub.testing` "integration-tests": list(full_test_dev_requirements), + "debug": list(debug_requirements), }, ) diff --git a/metadata-ingestion/src/datahub/ingestion/run/pipeline.py b/metadata-ingestion/src/datahub/ingestion/run/pipeline.py index 79d959965e0dd..07b55e0e25a89 100644 --- a/metadata-ingestion/src/datahub/ingestion/run/pipeline.py +++ b/metadata-ingestion/src/datahub/ingestion/run/pipeline.py @@ -353,77 +353,89 @@ def _time_to_print(self) -> bool: return False def run(self) -> None: - self.final_status = "unknown" - self._notify_reporters_on_ingestion_start() - callback = None - try: - callback = ( - LoggingCallback() - if not self.config.failure_log.enabled - else DeadLetterQueueCallback( - self.ctx, self.config.failure_log.log_config - ) - ) - for wu in itertools.islice( - self.source.get_workunits(), - self.preview_workunits if self.preview_mode else None, - ): - try: - if self._time_to_print(): - self.pretty_print_summary(currently_running=True) - except Exception as e: - logger.warning(f"Failed to print summary {e}") - - if not self.dry_run: - self.sink.handle_work_unit_start(wu) - try: - record_envelopes = self.extractor.get_records(wu) - for record_envelope in self.transform(record_envelopes): - if not self.dry_run: - self.sink.write_record_async(record_envelope, callback) - - except RuntimeError: - raise - except SystemExit: - raise - except Exception as e: - logger.error( - "Failed to process some records. Continuing.", exc_info=e + with contextlib.ExitStack() as stack: + if self.config.flags.generate_memory_profiles: + import memray + + stack.enter_context( + memray.Tracker( + f"{self.config.flags.generate_memory_profiles}/{self.config.run_id}.bin" ) - # TODO: Transformer errors should cause the pipeline to fail. - - self.extractor.close() - if not self.dry_run: - self.sink.handle_work_unit_end(wu) - self.source.close() - # no more data is coming, we need to let the transformers produce any additional records if they are holding on to state - for record_envelope in self.transform( - [ - RecordEnvelope( - record=EndOfStream(), metadata={"workunit_id": "end-of-stream"} + ) + + self.final_status = "unknown" + self._notify_reporters_on_ingestion_start() + callback = None + try: + callback = ( + LoggingCallback() + if not self.config.failure_log.enabled + else DeadLetterQueueCallback( + self.ctx, self.config.failure_log.log_config ) - ] - ): - if not self.dry_run and not isinstance( - record_envelope.record, EndOfStream + ) + for wu in itertools.islice( + self.source.get_workunits(), + self.preview_workunits if self.preview_mode else None, + ): + try: + if self._time_to_print(): + self.pretty_print_summary(currently_running=True) + except Exception as e: + logger.warning(f"Failed to print summary {e}") + + if not self.dry_run: + self.sink.handle_work_unit_start(wu) + try: + record_envelopes = self.extractor.get_records(wu) + for record_envelope in self.transform(record_envelopes): + if not self.dry_run: + self.sink.write_record_async(record_envelope, callback) + + except RuntimeError: + raise + except SystemExit: + raise + except Exception as e: + logger.error( + "Failed to process some records. Continuing.", + exc_info=e, + ) + # TODO: Transformer errors should cause the pipeline to fail. + + self.extractor.close() + if not self.dry_run: + self.sink.handle_work_unit_end(wu) + self.source.close() + # no more data is coming, we need to let the transformers produce any additional records if they are holding on to state + for record_envelope in self.transform( + [ + RecordEnvelope( + record=EndOfStream(), + metadata={"workunit_id": "end-of-stream"}, + ) + ] ): - # TODO: propagate EndOfStream and other control events to sinks, to allow them to flush etc. - self.sink.write_record_async(record_envelope, callback) - - self.sink.close() - self.process_commits() - self.final_status = "completed" - except (SystemExit, RuntimeError, KeyboardInterrupt) as e: - self.final_status = "cancelled" - logger.error("Caught error", exc_info=e) - raise - finally: - clear_global_warnings() - - if callback and hasattr(callback, "close"): - callback.close() # type: ignore - - self._notify_reporters_on_ingestion_completion() + if not self.dry_run and not isinstance( + record_envelope.record, EndOfStream + ): + # TODO: propagate EndOfStream and other control events to sinks, to allow them to flush etc. + self.sink.write_record_async(record_envelope, callback) + + self.sink.close() + self.process_commits() + self.final_status = "completed" + except (SystemExit, RuntimeError, KeyboardInterrupt) as e: + self.final_status = "cancelled" + logger.error("Caught error", exc_info=e) + raise + finally: + clear_global_warnings() + + if callback and hasattr(callback, "close"): + callback.close() # type: ignore + + self._notify_reporters_on_ingestion_completion() def transform(self, records: Iterable[RecordEnvelope]) -> Iterable[RecordEnvelope]: """ diff --git a/metadata-ingestion/src/datahub/ingestion/run/pipeline_config.py b/metadata-ingestion/src/datahub/ingestion/run/pipeline_config.py index ff9a7a6f3d146..da3cee8ad9c1b 100644 --- a/metadata-ingestion/src/datahub/ingestion/run/pipeline_config.py +++ b/metadata-ingestion/src/datahub/ingestion/run/pipeline_config.py @@ -57,6 +57,13 @@ class FlagsConfig(ConfigModel): ), ) + generate_memory_profiles: Optional[str] = Field( + default=None, + description=( + "Generate memray memory dumps for ingestion process by providing a path to write the dump file in." + ), + ) + class PipelineConfig(ConfigModel): # Once support for discriminated unions gets merged into Pydantic, we can From c564abcbf049e5251f9cc25bf0e339956279649d Mon Sep 17 00:00:00 2001 From: Amanda Hernando <110099762+amanda-her@users.noreply.github.com> Date: Thu, 12 Oct 2023 20:38:42 +0200 Subject: [PATCH 124/156] feat(auth): add group membership field resolver provider (#8846) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Adrián Pertíñez Co-authored-by: Adrián Pertíñez --- .../authorization/AuthorizationUtils.java | 8 +- .../dataset/DatasetStatsSummaryResolver.java | 4 +- .../dataset/DatasetUsageStatsResolver.java | 4 +- .../load/TimeSeriesAspectResolver.java | 4 +- .../policy/GetGrantedPrivilegesResolver.java | 6 +- .../resolvers/glossary/GlossaryUtilsTest.java | 36 +-- .../query/CreateQueryResolverTest.java | 6 +- .../query/DeleteQueryResolverTest.java | 6 +- .../query/UpdateQueryResolverTest.java | 10 +- .../com/datahub/authorization/AuthUtil.java | 10 +- .../authorization/AuthorizationRequest.java | 2 +- .../authorization/AuthorizerContext.java | 4 +- .../authorization/EntityFieldType.java | 31 ++ .../com/datahub/authorization/EntitySpec.java | 23 ++ .../authorization/EntitySpecResolver.java | 11 + .../datahub/authorization/FieldResolver.java | 6 +- .../authorization/ResolvedEntitySpec.java | 66 ++++ .../authorization/ResolvedResourceSpec.java | 55 ---- .../authorization/ResourceFieldType.java | 27 -- .../datahub/authorization/ResourceSpec.java | 23 -- .../authorization/ResourceSpecResolver.java | 11 - .../auth/authorization/Authorizer.java | 4 +- .../authorization/AuthorizerChain.java | 2 +- .../authorization/DataHubAuthorizer.java | 42 ++- ...er.java => DefaultEntitySpecResolver.java} | 33 +- .../datahub/authorization/FilterUtils.java | 8 +- .../datahub/authorization/PolicyEngine.java | 206 +++++------- ...PlatformInstanceFieldResolverProvider.java | 28 +- .../DomainFieldResolverProvider.java | 20 +- .../EntityFieldResolverProvider.java | 22 ++ .../EntityTypeFieldResolverProvider.java | 16 +- .../EntityUrnFieldResolverProvider.java | 16 +- .../GroupMembershipFieldResolverProvider.java | 78 +++++ .../OwnerFieldResolverProvider.java | 20 +- .../ResourceFieldResolverProvider.java | 22 -- .../authorization/DataHubAuthorizerTest.java | 22 +- .../authorization/PolicyEngineTest.java | 304 ++++++++---------- ...formInstanceFieldResolverProviderTest.java | 37 ++- ...upMembershipFieldResolverProviderTest.java | 212 ++++++++++++ .../factory/auth/AuthorizerChainFactory.java | 14 +- .../delegates/EntityApiDelegateImpl.java | 9 +- .../openapi/entities/EntitiesController.java | 10 +- .../RelationshipsController.java | 6 +- .../openapi/timeline/TimelineController.java | 4 +- .../openapi/util/MappingUtil.java | 11 +- .../datahub/plugins/test/TestAuthorizer.java | 4 +- .../resources/entity/AspectResource.java | 13 +- .../entity/BatchIngestionRunResource.java | 6 +- .../resources/entity/EntityResource.java | 54 ++-- .../resources/entity/EntityV2Resource.java | 8 +- .../entity/EntityVersionedV2Resource.java | 6 +- .../resources/lineage/Relationships.java | 8 +- .../metadata/resources/operations/Utils.java | 6 +- .../resources/platform/PlatformResource.java | 4 +- .../resources/restli/RestliUtils.java | 6 +- .../metadata/resources/usage/UsageStats.java | 8 +- 56 files changed, 937 insertions(+), 685 deletions(-) create mode 100644 metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntityFieldType.java create mode 100644 metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntitySpec.java create mode 100644 metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntitySpecResolver.java create mode 100644 metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedEntitySpec.java delete mode 100644 metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedResourceSpec.java delete mode 100644 metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceFieldType.java delete mode 100644 metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceSpec.java delete mode 100644 metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceSpecResolver.java rename metadata-service/auth-impl/src/main/java/com/datahub/authorization/{DefaultResourceSpecResolver.java => DefaultEntitySpecResolver.java} (51%) create mode 100644 metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityFieldResolverProvider.java create mode 100644 metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/GroupMembershipFieldResolverProvider.java delete mode 100644 metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/ResourceFieldResolverProvider.java create mode 100644 metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/GroupMembershipFieldResolverProviderTest.java diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java index 3089b8c8fc2db..03e63c7fb472f 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java @@ -4,7 +4,7 @@ import com.datahub.plugins.auth.authorization.Authorizer; import com.datahub.authorization.ConjunctivePrivilegeGroup; import com.datahub.authorization.DisjunctivePrivilegeGroup; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.google.common.collect.ImmutableList; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; @@ -90,7 +90,7 @@ public static boolean canManageTags(@Nonnull QueryContext context) { } public static boolean canDeleteEntity(@Nonnull Urn entityUrn, @Nonnull QueryContext context) { - return isAuthorized(context, Optional.of(new ResourceSpec(entityUrn.getEntityType(), entityUrn.toString())), PoliciesConfig.DELETE_ENTITY_PRIVILEGE); + return isAuthorized(context, Optional.of(new EntitySpec(entityUrn.getEntityType(), entityUrn.toString())), PoliciesConfig.DELETE_ENTITY_PRIVILEGE); } public static boolean canManageUserCredentials(@Nonnull QueryContext context) { @@ -173,7 +173,7 @@ public static boolean canDeleteQuery(@Nonnull Urn entityUrn, @Nonnull List public static boolean isAuthorized( @Nonnull QueryContext context, - @Nonnull Optional resourceSpec, + @Nonnull Optional resourceSpec, @Nonnull PoliciesConfig.Privilege privilege) { final Authorizer authorizer = context.getAuthorizer(); final String actor = context.getActorUrn(); @@ -196,7 +196,7 @@ public static boolean isAuthorized( @Nonnull String resource, @Nonnull DisjunctivePrivilegeGroup privilegeGroup ) { - final ResourceSpec resourceSpec = new ResourceSpec(resourceType, resource); + final EntitySpec resourceSpec = new EntitySpec(resourceType, resource); return AuthUtil.isAuthorized(authorizer, actor, Optional.of(resourceSpec), privilegeGroup); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetStatsSummaryResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetStatsSummaryResolver.java index 23be49c7e7140..2873866bb34f7 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetStatsSummaryResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetStatsSummaryResolver.java @@ -1,6 +1,6 @@ package com.linkedin.datahub.graphql.resolvers.dataset; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; import com.linkedin.common.urn.Urn; @@ -104,7 +104,7 @@ private CorpUser createPartialUser(final Urn userUrn) { private boolean isAuthorized(final Urn resourceUrn, final QueryContext context) { return AuthorizationUtils.isAuthorized(context, - Optional.of(new ResourceSpec(resourceUrn.getEntityType(), resourceUrn.toString())), + Optional.of(new EntitySpec(resourceUrn.getEntityType(), resourceUrn.toString())), PoliciesConfig.VIEW_DATASET_USAGE_PRIVILEGE); } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetUsageStatsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetUsageStatsResolver.java index 20361830ad5a5..e4bec8e896fdf 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetUsageStatsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetUsageStatsResolver.java @@ -1,6 +1,6 @@ package com.linkedin.datahub.graphql.resolvers.dataset; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.datahub.graphql.QueryContext; @@ -52,7 +52,7 @@ public CompletableFuture get(DataFetchingEnvironment environme private boolean isAuthorized(final Urn resourceUrn, final QueryContext context) { return AuthorizationUtils.isAuthorized(context, - Optional.of(new ResourceSpec(resourceUrn.getEntityType(), resourceUrn.toString())), + Optional.of(new EntitySpec(resourceUrn.getEntityType(), resourceUrn.toString())), PoliciesConfig.VIEW_DATASET_USAGE_PRIVILEGE); } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/TimeSeriesAspectResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/TimeSeriesAspectResolver.java index 197ca8640559d..f13ebf8373e91 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/TimeSeriesAspectResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/TimeSeriesAspectResolver.java @@ -1,6 +1,6 @@ package com.linkedin.datahub.graphql.resolvers.load; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.authorization.AuthorizationUtils; import com.linkedin.datahub.graphql.generated.Entity; @@ -79,7 +79,7 @@ public TimeSeriesAspectResolver( private boolean isAuthorized(QueryContext context, String urn) { if (_entityName.equals(Constants.DATASET_ENTITY_NAME) && _aspectName.equals( Constants.DATASET_PROFILE_ASPECT_NAME)) { - return AuthorizationUtils.isAuthorized(context, Optional.of(new ResourceSpec(_entityName, urn)), + return AuthorizationUtils.isAuthorized(context, Optional.of(new EntitySpec(_entityName, urn)), PoliciesConfig.VIEW_DATASET_PROFILE_PRIVILEGE); } return true; diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/GetGrantedPrivilegesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/GetGrantedPrivilegesResolver.java index 2f20fdaf1e9b1..11f7793db82c8 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/GetGrantedPrivilegesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/GetGrantedPrivilegesResolver.java @@ -2,7 +2,7 @@ import com.datahub.authorization.AuthorizerChain; import com.datahub.authorization.DataHubAuthorizer; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.exception.AuthorizationException; import com.linkedin.datahub.graphql.generated.GetGrantedPrivilegesInput; @@ -33,8 +33,8 @@ public CompletableFuture get(final DataFetchingEnvironment environme if (!isAuthorized(context, actor)) { throw new AuthorizationException("Unauthorized to get privileges for the given author."); } - final Optional resourceSpec = Optional.ofNullable(input.getResourceSpec()) - .map(spec -> new ResourceSpec(EntityTypeMapper.getName(spec.getResourceType()), spec.getResourceUrn())); + final Optional resourceSpec = Optional.ofNullable(input.getResourceSpec()) + .map(spec -> new EntitySpec(EntityTypeMapper.getName(spec.getResourceType()), spec.getResourceUrn())); if (context.getAuthorizer() instanceof AuthorizerChain) { DataHubAuthorizer dataHubAuthorizer = ((AuthorizerChain) context.getAuthorizer()).getDefaultAuthorizer(); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GlossaryUtilsTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GlossaryUtilsTest.java index ccaab44f60dd4..8bfc32e1999ae 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GlossaryUtilsTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GlossaryUtilsTest.java @@ -5,7 +5,7 @@ import com.datahub.authorization.AuthorizationRequest; import com.datahub.authorization.AuthorizationResult; import com.datahub.plugins.auth.authorization.Authorizer; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.linkedin.common.urn.GlossaryNodeUrn; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; @@ -89,17 +89,17 @@ private void setUpTests() throws Exception { Mockito.any(Authentication.class) )).thenReturn(new EntityResponse().setAspects(new EnvelopedAspectMap(parentNode3Aspects))); - final ResourceSpec resourceSpec3 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn3.toString()); + final EntitySpec resourceSpec3 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn3.toString()); mockAuthRequest("MANAGE_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec3); - final ResourceSpec resourceSpec2 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn2.toString()); + final EntitySpec resourceSpec2 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn2.toString()); mockAuthRequest("MANAGE_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec2); - final ResourceSpec resourceSpec1 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn1.toString()); + final EntitySpec resourceSpec1 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn1.toString()); mockAuthRequest("MANAGE_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec1); } - private void mockAuthRequest(String privilege, AuthorizationResult.Type allowOrDeny, ResourceSpec resourceSpec) { + private void mockAuthRequest(String privilege, AuthorizationResult.Type allowOrDeny, EntitySpec resourceSpec) { final AuthorizationRequest authorizationRequest = new AuthorizationRequest( userUrn, privilege, @@ -150,7 +150,7 @@ public void testCanManageChildrenEntitiesAuthorized() throws Exception { // they do NOT have the MANAGE_GLOSSARIES platform privilege mockAuthRequest("MANAGE_GLOSSARIES", AuthorizationResult.Type.DENY, null); - final ResourceSpec resourceSpec = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn.toString()); + final EntitySpec resourceSpec = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn.toString()); mockAuthRequest("MANAGE_GLOSSARY_CHILDREN", AuthorizationResult.Type.ALLOW, resourceSpec); assertTrue(GlossaryUtils.canManageChildrenEntities(mockContext, parentNodeUrn, mockClient)); @@ -162,7 +162,7 @@ public void testCanManageChildrenEntitiesUnauthorized() throws Exception { // they do NOT have the MANAGE_GLOSSARIES platform privilege mockAuthRequest("MANAGE_GLOSSARIES", AuthorizationResult.Type.DENY, null); - final ResourceSpec resourceSpec = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn.toString()); + final EntitySpec resourceSpec = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn.toString()); mockAuthRequest("MANAGE_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec); @@ -175,13 +175,13 @@ public void testCanManageChildrenRecursivelyEntitiesAuthorized() throws Exceptio // they do NOT have the MANAGE_GLOSSARIES platform privilege mockAuthRequest("MANAGE_GLOSSARIES", AuthorizationResult.Type.DENY, null); - final ResourceSpec resourceSpec3 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn3.toString()); + final EntitySpec resourceSpec3 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn3.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.ALLOW, resourceSpec3); - final ResourceSpec resourceSpec2 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn2.toString()); + final EntitySpec resourceSpec2 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn2.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec2); - final ResourceSpec resourceSpec1 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn1.toString()); + final EntitySpec resourceSpec1 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn1.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec1); assertTrue(GlossaryUtils.canManageChildrenEntities(mockContext, parentNodeUrn1, mockClient)); @@ -193,13 +193,13 @@ public void testCanManageChildrenRecursivelyEntitiesUnauthorized() throws Except // they do NOT have the MANAGE_GLOSSARIES platform privilege mockAuthRequest("MANAGE_GLOSSARIES", AuthorizationResult.Type.DENY, null); - final ResourceSpec resourceSpec3 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn3.toString()); + final EntitySpec resourceSpec3 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn3.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec3); - final ResourceSpec resourceSpec2 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn2.toString()); + final EntitySpec resourceSpec2 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn2.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec2); - final ResourceSpec resourceSpec1 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn1.toString()); + final EntitySpec resourceSpec1 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn1.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec1); assertFalse(GlossaryUtils.canManageChildrenEntities(mockContext, parentNodeUrn1, mockClient)); @@ -211,10 +211,10 @@ public void testCanManageChildrenRecursivelyEntitiesAuthorizedLevel2() throws Ex // they do NOT have the MANAGE_GLOSSARIES platform privilege mockAuthRequest("MANAGE_GLOSSARIES", AuthorizationResult.Type.DENY, null); - final ResourceSpec resourceSpec2 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn2.toString()); + final EntitySpec resourceSpec2 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn2.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.ALLOW, resourceSpec2); - final ResourceSpec resourceSpec1 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn1.toString()); + final EntitySpec resourceSpec1 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn1.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec1); assertTrue(GlossaryUtils.canManageChildrenEntities(mockContext, parentNodeUrn1, mockClient)); @@ -226,10 +226,10 @@ public void testCanManageChildrenRecursivelyEntitiesUnauthorizedLevel2() throws // they do NOT have the MANAGE_GLOSSARIES platform privilege mockAuthRequest("MANAGE_GLOSSARIES", AuthorizationResult.Type.DENY, null); - final ResourceSpec resourceSpec3 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn3.toString()); + final EntitySpec resourceSpec3 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn3.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec3); - final ResourceSpec resourceSpec2 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn2.toString()); + final EntitySpec resourceSpec2 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn2.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec2); assertFalse(GlossaryUtils.canManageChildrenEntities(mockContext, parentNodeUrn2, mockClient)); @@ -241,7 +241,7 @@ public void testCanManageChildrenRecursivelyEntitiesNoLevel2() throws Exception // they do NOT have the MANAGE_GLOSSARIES platform privilege mockAuthRequest("MANAGE_GLOSSARIES", AuthorizationResult.Type.DENY, null); - final ResourceSpec resourceSpec3 = new ResourceSpec(parentNodeUrn.getEntityType(), parentNodeUrn3.toString()); + final EntitySpec resourceSpec3 = new EntitySpec(parentNodeUrn.getEntityType(), parentNodeUrn3.toString()); mockAuthRequest("MANAGE_ALL_GLOSSARY_CHILDREN", AuthorizationResult.Type.DENY, resourceSpec3); assertFalse(GlossaryUtils.canManageChildrenEntities(mockContext, parentNodeUrn3, mockClient)); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/CreateQueryResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/CreateQueryResolverTest.java index 196eb24b52bf8..9c04c67dd3a3b 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/CreateQueryResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/CreateQueryResolverTest.java @@ -5,7 +5,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authorization.AuthorizationRequest; import com.datahub.authorization.AuthorizationResult; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -201,7 +201,7 @@ private QueryContext getMockQueryContext(boolean allowEditEntityQueries) { TEST_ACTOR_URN.toString(), PoliciesConfig.EDIT_QUERIES_PRIVILEGE.getType(), Optional.of( - new ResourceSpec( + new EntitySpec( TEST_DATASET_URN.getEntityType(), TEST_DATASET_URN.toString())) ); @@ -210,7 +210,7 @@ private QueryContext getMockQueryContext(boolean allowEditEntityQueries) { TEST_ACTOR_URN.toString(), PoliciesConfig.EDIT_ENTITY_PRIVILEGE.getType(), Optional.of( - new ResourceSpec( + new EntitySpec( TEST_DATASET_URN.getEntityType(), TEST_DATASET_URN.toString())) ); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/DeleteQueryResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/DeleteQueryResolverTest.java index a6b4887b0e882..78c894f27cbc3 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/DeleteQueryResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/DeleteQueryResolverTest.java @@ -5,7 +5,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authorization.AuthorizationRequest; import com.datahub.authorization.AuthorizationResult; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; @@ -134,7 +134,7 @@ private QueryContext getMockQueryContext(boolean allowEditEntityQueries) { DeleteQueryResolverTest.TEST_ACTOR_URN.toString(), PoliciesConfig.EDIT_QUERIES_PRIVILEGE.getType(), Optional.of( - new ResourceSpec( + new EntitySpec( DeleteQueryResolverTest.TEST_DATASET_URN.getEntityType(), DeleteQueryResolverTest.TEST_DATASET_URN.toString())) ); @@ -143,7 +143,7 @@ private QueryContext getMockQueryContext(boolean allowEditEntityQueries) { TEST_ACTOR_URN.toString(), PoliciesConfig.EDIT_ENTITY_PRIVILEGE.getType(), Optional.of( - new ResourceSpec( + new EntitySpec( TEST_DATASET_URN.getEntityType(), TEST_DATASET_URN.toString())) ); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/UpdateQueryResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/UpdateQueryResolverTest.java index 7a76b6d6be5a4..9b500b5fb3936 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/UpdateQueryResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/UpdateQueryResolverTest.java @@ -5,7 +5,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authorization.AuthorizationRequest; import com.datahub.authorization.AuthorizationResult; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -206,7 +206,7 @@ private QueryContext getMockQueryContext(boolean allowEditEntityQueries) { TEST_ACTOR_URN.toString(), PoliciesConfig.EDIT_QUERIES_PRIVILEGE.getType(), Optional.of( - new ResourceSpec( + new EntitySpec( TEST_DATASET_URN.getEntityType(), TEST_DATASET_URN.toString())) ); @@ -215,7 +215,7 @@ private QueryContext getMockQueryContext(boolean allowEditEntityQueries) { TEST_ACTOR_URN.toString(), PoliciesConfig.EDIT_ENTITY_PRIVILEGE.getType(), Optional.of( - new ResourceSpec( + new EntitySpec( TEST_DATASET_URN.getEntityType(), TEST_DATASET_URN.toString())) ); @@ -224,7 +224,7 @@ private QueryContext getMockQueryContext(boolean allowEditEntityQueries) { TEST_ACTOR_URN.toString(), PoliciesConfig.EDIT_QUERIES_PRIVILEGE.getType(), Optional.of( - new ResourceSpec( + new EntitySpec( TEST_DATASET_URN_2.getEntityType(), TEST_DATASET_URN_2.toString())) ); @@ -233,7 +233,7 @@ private QueryContext getMockQueryContext(boolean allowEditEntityQueries) { TEST_ACTOR_URN.toString(), PoliciesConfig.EDIT_ENTITY_PRIVILEGE.getType(), Optional.of( - new ResourceSpec( + new EntitySpec( TEST_DATASET_URN_2.getEntityType(), TEST_DATASET_URN_2.toString())) ); diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthUtil.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthUtil.java index dfb936c61ee0c..e159993a8a243 100644 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthUtil.java +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthUtil.java @@ -11,7 +11,7 @@ public class AuthUtil { public static boolean isAuthorized( @Nonnull Authorizer authorizer, @Nonnull String actor, - @Nonnull Optional maybeResourceSpec, + @Nonnull Optional maybeResourceSpec, @Nonnull DisjunctivePrivilegeGroup privilegeGroup ) { for (ConjunctivePrivilegeGroup andPrivilegeGroup : privilegeGroup.getAuthorizedPrivilegeGroups()) { @@ -27,7 +27,7 @@ public static boolean isAuthorized( public static boolean isAuthorizedForResources( @Nonnull Authorizer authorizer, @Nonnull String actor, - @Nonnull List> resourceSpecs, + @Nonnull List> resourceSpecs, @Nonnull DisjunctivePrivilegeGroup privilegeGroup ) { for (ConjunctivePrivilegeGroup andPrivilegeGroup : privilegeGroup.getAuthorizedPrivilegeGroups()) { @@ -44,7 +44,7 @@ private static boolean isAuthorized( @Nonnull Authorizer authorizer, @Nonnull String actor, @Nonnull ConjunctivePrivilegeGroup requiredPrivileges, - @Nonnull Optional resourceSpec) { + @Nonnull Optional resourceSpec) { // Each privilege in a group _must_ all be true to permit the operation. for (final String privilege : requiredPrivileges.getRequiredPrivileges()) { // Create and evaluate an Authorization request. @@ -62,11 +62,11 @@ private static boolean isAuthorizedForResources( @Nonnull Authorizer authorizer, @Nonnull String actor, @Nonnull ConjunctivePrivilegeGroup requiredPrivileges, - @Nonnull List> resourceSpecs) { + @Nonnull List> resourceSpecs) { // Each privilege in a group _must_ all be true to permit the operation. for (final String privilege : requiredPrivileges.getRequiredPrivileges()) { // Create and evaluate an Authorization request. - for (Optional resourceSpec : resourceSpecs) { + for (Optional resourceSpec : resourceSpecs) { final AuthorizationRequest request = new AuthorizationRequest(actor, privilege, resourceSpec); final AuthorizationResult result = authorizer.authorize(request); if (AuthorizationResult.Type.DENY.equals(result.getType())) { diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizationRequest.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizationRequest.java index 084a455495551..9e75de3cbf44d 100644 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizationRequest.java +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizationRequest.java @@ -21,5 +21,5 @@ public class AuthorizationRequest { * The resource that the user is requesting for, if applicable. If the privilege is a platform privilege * this optional will be empty. */ - Optional resourceSpec; + Optional resourceSpec; } diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizerContext.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizerContext.java index f9940d171d5d4..b79a4fa20c7ea 100644 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizerContext.java +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizerContext.java @@ -18,9 +18,9 @@ public class AuthorizerContext { private final Map contextMap; /** - * A utility for resolving a {@link ResourceSpec} to resolved resource field values. + * A utility for resolving an {@link EntitySpec} to resolved entity field values. */ - private ResourceSpecResolver resourceSpecResolver; + private EntitySpecResolver entitySpecResolver; /** * diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntityFieldType.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntityFieldType.java new file mode 100644 index 0000000000000..46763f29a7040 --- /dev/null +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntityFieldType.java @@ -0,0 +1,31 @@ +package com.datahub.authorization; + +/** + * List of entity field types to fetch for a given entity + */ +public enum EntityFieldType { + /** + * Type of the entity (e.g. dataset, chart) + */ + TYPE, + /** + * Urn of the entity + */ + URN, + /** + * Owners of the entity + */ + OWNER, + /** + * Domains of the entity + */ + DOMAIN, + /** + * Groups of which the entity (only applies to corpUser) is a member + */ + GROUP_MEMBERSHIP, + /** + * Data platform instance of resource + */ + DATA_PLATFORM_INSTANCE +} diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntitySpec.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntitySpec.java new file mode 100644 index 0000000000000..656bec0f44fc2 --- /dev/null +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntitySpec.java @@ -0,0 +1,23 @@ +package com.datahub.authorization; + +import javax.annotation.Nonnull; +import lombok.Value; + + +/** + * Details about the entities involved in the authorization process. It models the actor and the resource being acted + * upon. Resource types currently supported can be found inside of {@link com.linkedin.metadata.authorization.PoliciesConfig} + */ +@Value +public class EntitySpec { + /** + * The entity type. (dataset, chart, dashboard, corpGroup, etc). + */ + @Nonnull + String type; + /** + * The entity identity. Most often, this corresponds to the raw entity urn. (urn:li:corpGroup:groupId) + */ + @Nonnull + String entity; +} \ No newline at end of file diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntitySpecResolver.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntitySpecResolver.java new file mode 100644 index 0000000000000..67347fbf87a87 --- /dev/null +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntitySpecResolver.java @@ -0,0 +1,11 @@ +package com.datahub.authorization; + +/** + * An Entity Spec Resolver is responsible for resolving a {@link EntitySpec} to a {@link ResolvedEntitySpec}. + */ +public interface EntitySpecResolver { + /** + Resolve a {@link EntitySpec} to a resolved entity spec. + **/ + ResolvedEntitySpec resolve(EntitySpec entitySpec); +} diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/FieldResolver.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/FieldResolver.java index 9318f5f8e7b96..955a06fd54cb9 100644 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/FieldResolver.java +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/FieldResolver.java @@ -33,9 +33,9 @@ public static FieldResolver getResolverFromValues(Set values) { /** * Helper function that returns FieldResolver given a fetchFieldValue function */ - public static FieldResolver getResolverFromFunction(ResourceSpec resourceSpec, - Function fetchFieldValue) { - return new FieldResolver(() -> CompletableFuture.supplyAsync(() -> fetchFieldValue.apply(resourceSpec))); + public static FieldResolver getResolverFromFunction(EntitySpec entitySpec, + Function fetchFieldValue) { + return new FieldResolver(() -> CompletableFuture.supplyAsync(() -> fetchFieldValue.apply(entitySpec))); } public static FieldValue emptyFieldValue() { diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedEntitySpec.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedEntitySpec.java new file mode 100644 index 0000000000000..7948766df5715 --- /dev/null +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedEntitySpec.java @@ -0,0 +1,66 @@ +package com.datahub.authorization; + +import java.util.Collections; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import lombok.Getter; +import lombok.RequiredArgsConstructor; +import lombok.ToString; + + +/** + * Wrapper around authorization request with field resolvers for lazily fetching the field values for each field type + */ +@RequiredArgsConstructor +@ToString +public class ResolvedEntitySpec { + @Getter + private final EntitySpec spec; + private final Map fieldResolvers; + + public Set getFieldValues(EntityFieldType entityFieldType) { + if (!fieldResolvers.containsKey(entityFieldType)) { + return Collections.emptySet(); + } + return fieldResolvers.get(entityFieldType).getFieldValuesFuture().join().getValues(); + } + + /** + * Fetch the owners for an entity. + * @return a set of owner urns, or empty set if none exist. + */ + public Set getOwners() { + if (!fieldResolvers.containsKey(EntityFieldType.OWNER)) { + return Collections.emptySet(); + } + return fieldResolvers.get(EntityFieldType.OWNER).getFieldValuesFuture().join().getValues(); + } + + /** + * Fetch the platform instance for a Resolved Resource Spec + * @return a Platform Instance or null if one does not exist. + */ + @Nullable + public String getDataPlatformInstance() { + if (!fieldResolvers.containsKey(EntityFieldType.DATA_PLATFORM_INSTANCE)) { + return null; + } + Set dataPlatformInstance = fieldResolvers.get(EntityFieldType.DATA_PLATFORM_INSTANCE).getFieldValuesFuture().join().getValues(); + if (dataPlatformInstance.size() > 0) { + return dataPlatformInstance.stream().findFirst().get(); + } + return null; + } + + /** + * Fetch the group membership for an entity. + * @return a set of groups urns, or empty set if none exist. + */ + public Set getGroupMembership() { + if (!fieldResolvers.containsKey(EntityFieldType.GROUP_MEMBERSHIP)) { + return Collections.emptySet(); + } + return fieldResolvers.get(EntityFieldType.GROUP_MEMBERSHIP).getFieldValuesFuture().join().getValues(); + } +} diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedResourceSpec.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedResourceSpec.java deleted file mode 100644 index 8e429a8ca1b94..0000000000000 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedResourceSpec.java +++ /dev/null @@ -1,55 +0,0 @@ -package com.datahub.authorization; - -import java.util.Collections; -import java.util.Map; -import java.util.Set; -import javax.annotation.Nullable; -import lombok.Getter; -import lombok.RequiredArgsConstructor; -import lombok.ToString; - - -/** - * Wrapper around authorization request with field resolvers for lazily fetching the field values for each field type - */ -@RequiredArgsConstructor -@ToString -public class ResolvedResourceSpec { - @Getter - private final ResourceSpec spec; - private final Map fieldResolvers; - - public Set getFieldValues(ResourceFieldType resourceFieldType) { - if (!fieldResolvers.containsKey(resourceFieldType)) { - return Collections.emptySet(); - } - return fieldResolvers.get(resourceFieldType).getFieldValuesFuture().join().getValues(); - } - - /** - * Fetch the owners for a resource. - * @return a set of owner urns, or empty set if none exist. - */ - public Set getOwners() { - if (!fieldResolvers.containsKey(ResourceFieldType.OWNER)) { - return Collections.emptySet(); - } - return fieldResolvers.get(ResourceFieldType.OWNER).getFieldValuesFuture().join().getValues(); - } - - /** - * Fetch the platform instance for a Resolved Resource Spec - * @return a Platform Instance or null if one does not exist. - */ - @Nullable - public String getDataPlatformInstance() { - if (!fieldResolvers.containsKey(ResourceFieldType.DATA_PLATFORM_INSTANCE)) { - return null; - } - Set dataPlatformInstance = fieldResolvers.get(ResourceFieldType.DATA_PLATFORM_INSTANCE).getFieldValuesFuture().join().getValues(); - if (dataPlatformInstance.size() > 0) { - return dataPlatformInstance.stream().findFirst().get(); - } - return null; - } -} diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceFieldType.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceFieldType.java deleted file mode 100644 index 478522dc7c331..0000000000000 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceFieldType.java +++ /dev/null @@ -1,27 +0,0 @@ -package com.datahub.authorization; - -/** - * List of resource field types to fetch for a given resource - */ -public enum ResourceFieldType { - /** - * Type of resource (e.g. dataset, chart) - */ - RESOURCE_TYPE, - /** - * Urn of resource - */ - RESOURCE_URN, - /** - * Owners of resource - */ - OWNER, - /** - * Domains of resource - */ - DOMAIN, - /** - * Data platform instance of resource - */ - DATA_PLATFORM_INSTANCE -} diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceSpec.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceSpec.java deleted file mode 100644 index c1bd53e31fe29..0000000000000 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceSpec.java +++ /dev/null @@ -1,23 +0,0 @@ -package com.datahub.authorization; - -import javax.annotation.Nonnull; -import lombok.Value; - - -/** - * Details about a specific resource being acted upon. Resource types currently supported - * can be found inside of {@link com.linkedin.metadata.authorization.PoliciesConfig} - */ -@Value -public class ResourceSpec { - /** - * The resource type. Most often, this corresponds to the entity type. (dataset, chart, dashboard, corpGroup, etc). - */ - @Nonnull - String type; - /** - * The resource identity. Most often, this corresponds to the raw entity urn. (urn:li:corpGroup:groupId) - */ - @Nonnull - String resource; -} \ No newline at end of file diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceSpecResolver.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceSpecResolver.java deleted file mode 100644 index 05c35f377b9a9..0000000000000 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceSpecResolver.java +++ /dev/null @@ -1,11 +0,0 @@ -package com.datahub.authorization; - -/** - * A Resource Spec Resolver is responsible for resolving a {@link ResourceSpec} to a {@link ResolvedResourceSpec}. - */ -public interface ResourceSpecResolver { - /** - Resolve a {@link ResourceSpec} to a resolved resource spec. - **/ - ResolvedResourceSpec resolve(ResourceSpec resourceSpec); -} diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/plugins/auth/authorization/Authorizer.java b/metadata-auth/auth-api/src/main/java/com/datahub/plugins/auth/authorization/Authorizer.java index ce7a3f22b3147..c731a3ec987c1 100644 --- a/metadata-auth/auth-api/src/main/java/com/datahub/plugins/auth/authorization/Authorizer.java +++ b/metadata-auth/auth-api/src/main/java/com/datahub/plugins/auth/authorization/Authorizer.java @@ -4,7 +4,7 @@ import com.datahub.authorization.AuthorizationResult; import com.datahub.authorization.AuthorizedActors; import com.datahub.authorization.AuthorizerContext; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.datahub.plugins.Plugin; import java.util.Map; import java.util.Optional; @@ -32,5 +32,5 @@ public interface Authorizer extends Plugin { * Retrieves the current list of actors authorized to for a particular privilege against * an optional resource */ - AuthorizedActors authorizedActors(final String privilege, final Optional resourceSpec); + AuthorizedActors authorizedActors(final String privilege, final Optional resourceSpec); } diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/AuthorizerChain.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/AuthorizerChain.java index d62c37160f816..f8eca541e1efb 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/AuthorizerChain.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/AuthorizerChain.java @@ -82,7 +82,7 @@ public AuthorizationResult authorize(@Nonnull final AuthorizationRequest request } @Override - public AuthorizedActors authorizedActors(String privilege, Optional resourceSpec) { + public AuthorizedActors authorizedActors(String privilege, Optional resourceSpec) { if (this.authorizers.isEmpty()) { return null; } diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DataHubAuthorizer.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DataHubAuthorizer.java index f653ccf72cf54..4553139e3ca54 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DataHubAuthorizer.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DataHubAuthorizer.java @@ -8,6 +8,8 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.policy.DataHubPolicyInfo; + +import java.net.URISyntaxException; import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -55,7 +57,7 @@ public enum AuthorizationMode { private final ScheduledExecutorService _refreshExecutorService = Executors.newScheduledThreadPool(1); private final PolicyRefreshRunnable _policyRefreshRunnable; private final PolicyEngine _policyEngine; - private ResourceSpecResolver _resourceSpecResolver; + private EntitySpecResolver _entitySpecResolver; private AuthorizationMode _mode; public static final String ALL = "ALL"; @@ -76,7 +78,7 @@ public DataHubAuthorizer( @Override public void init(@Nonnull Map authorizerConfig, @Nonnull AuthorizerContext ctx) { // Pass. No static config. - _resourceSpecResolver = Objects.requireNonNull(ctx.getResourceSpecResolver()); + _entitySpecResolver = Objects.requireNonNull(ctx.getEntitySpecResolver()); } public AuthorizationResult authorize(@Nonnull final AuthorizationRequest request) { @@ -86,7 +88,7 @@ public AuthorizationResult authorize(@Nonnull final AuthorizationRequest request return new AuthorizationResult(request, AuthorizationResult.Type.ALLOW, null); } - Optional resolvedResourceSpec = request.getResourceSpec().map(_resourceSpecResolver::resolve); + Optional resolvedResourceSpec = request.getResourceSpec().map(_entitySpecResolver::resolve); // 1. Fetch the policies relevant to the requested privilege. final List policiesToEvaluate = _policyCache.getOrDefault(request.getPrivilege(), new ArrayList<>()); @@ -102,14 +104,17 @@ public AuthorizationResult authorize(@Nonnull final AuthorizationRequest request return new AuthorizationResult(request, AuthorizationResult.Type.DENY, null); } - public List getGrantedPrivileges(final String actorUrn, final Optional resourceSpec) { + public List getGrantedPrivileges(final String actor, final Optional resourceSpec) { // 1. Fetch all policies final List policiesToEvaluate = _policyCache.getOrDefault(ALL, new ArrayList<>()); - Optional resolvedResourceSpec = resourceSpec.map(_resourceSpecResolver::resolve); + Urn actorUrn = UrnUtils.getUrn(actor); + final ResolvedEntitySpec resolvedActorSpec = _entitySpecResolver.resolve(new EntitySpec(actorUrn.getEntityType(), actor)); + + Optional resolvedResourceSpec = resourceSpec.map(_entitySpecResolver::resolve); - return _policyEngine.getGrantedPrivileges(policiesToEvaluate, UrnUtils.getUrn(actorUrn), resolvedResourceSpec); + return _policyEngine.getGrantedPrivileges(policiesToEvaluate, resolvedActorSpec, resolvedResourceSpec); } /** @@ -118,11 +123,11 @@ public List getGrantedPrivileges(final String actorUrn, final Optional resourceSpec) { + final Optional resourceSpec) { // Step 1: Find policies granting the privilege. final List policiesToEvaluate = _policyCache.getOrDefault(privilege, new ArrayList<>()); - Optional resolvedResourceSpec = resourceSpec.map(_resourceSpecResolver::resolve); + Optional resolvedResourceSpec = resourceSpec.map(_entitySpecResolver::resolve); final List authorizedUsers = new ArrayList<>(); final List authorizedGroups = new ArrayList<>(); @@ -180,19 +185,36 @@ private boolean isSystemRequest(final AuthorizationRequest request, final Authen /** * Returns true if a policy grants the requested privilege for a given actor and resource. */ - private boolean isRequestGranted(final DataHubPolicyInfo policy, final AuthorizationRequest request, final Optional resourceSpec) { + private boolean isRequestGranted(final DataHubPolicyInfo policy, final AuthorizationRequest request, final Optional resourceSpec) { if (AuthorizationMode.ALLOW_ALL.equals(mode())) { return true; } + + Optional actorUrn = getUrnFromRequestActor(request.getActorUrn()); + if (actorUrn.isEmpty()) { + return false; + } + + final ResolvedEntitySpec resolvedActorSpec = _entitySpecResolver.resolve( + new EntitySpec(actorUrn.get().getEntityType(), request.getActorUrn())); final PolicyEngine.PolicyEvaluationResult result = _policyEngine.evaluatePolicy( policy, - request.getActorUrn(), + resolvedActorSpec, request.getPrivilege(), resourceSpec ); return result.isGranted(); } + private Optional getUrnFromRequestActor(String actor) { + try { + return Optional.of(Urn.createFromString(actor)); + } catch (URISyntaxException e) { + log.error(String.format("Failed to bind actor %s to an URN. Actors must be URNs. Denying the authorization request", actor)); + return Optional.empty(); + } + } + /** * A {@link Runnable} used to periodically fetch a new instance of the policies Cache. * diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultResourceSpecResolver.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultEntitySpecResolver.java similarity index 51% rename from metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultResourceSpecResolver.java rename to metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultEntitySpecResolver.java index 64c43dc8aa591..4ad14ed59c9c0 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultResourceSpecResolver.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultEntitySpecResolver.java @@ -1,39 +1,40 @@ package com.datahub.authorization; -import com.datahub.authentication.Authentication; import com.datahub.authorization.fieldresolverprovider.DataPlatformInstanceFieldResolverProvider; -import com.datahub.authorization.fieldresolverprovider.DomainFieldResolverProvider; import com.datahub.authorization.fieldresolverprovider.EntityTypeFieldResolverProvider; -import com.datahub.authorization.fieldresolverprovider.EntityUrnFieldResolverProvider; import com.datahub.authorization.fieldresolverprovider.OwnerFieldResolverProvider; -import com.datahub.authorization.fieldresolverprovider.ResourceFieldResolverProvider; +import com.datahub.authentication.Authentication; +import com.datahub.authorization.fieldresolverprovider.DomainFieldResolverProvider; +import com.datahub.authorization.fieldresolverprovider.EntityUrnFieldResolverProvider; +import com.datahub.authorization.fieldresolverprovider.EntityFieldResolverProvider; +import com.datahub.authorization.fieldresolverprovider.GroupMembershipFieldResolverProvider; import com.google.common.collect.ImmutableList; import com.linkedin.entity.client.EntityClient; - import java.util.List; import java.util.Map; import java.util.stream.Collectors; -public class DefaultResourceSpecResolver implements ResourceSpecResolver { - private final List _resourceFieldResolverProviders; +public class DefaultEntitySpecResolver implements EntitySpecResolver { + private final List _entityFieldResolverProviders; - public DefaultResourceSpecResolver(Authentication systemAuthentication, EntityClient entityClient) { - _resourceFieldResolverProviders = + public DefaultEntitySpecResolver(Authentication systemAuthentication, EntityClient entityClient) { + _entityFieldResolverProviders = ImmutableList.of(new EntityTypeFieldResolverProvider(), new EntityUrnFieldResolverProvider(), new DomainFieldResolverProvider(entityClient, systemAuthentication), new OwnerFieldResolverProvider(entityClient, systemAuthentication), - new DataPlatformInstanceFieldResolverProvider(entityClient, systemAuthentication)); + new DataPlatformInstanceFieldResolverProvider(entityClient, systemAuthentication), + new GroupMembershipFieldResolverProvider(entityClient, systemAuthentication)); } @Override - public ResolvedResourceSpec resolve(ResourceSpec resourceSpec) { - return new ResolvedResourceSpec(resourceSpec, getFieldResolvers(resourceSpec)); + public ResolvedEntitySpec resolve(EntitySpec entitySpec) { + return new ResolvedEntitySpec(entitySpec, getFieldResolvers(entitySpec)); } - private Map getFieldResolvers(ResourceSpec resourceSpec) { - return _resourceFieldResolverProviders.stream() - .collect(Collectors.toMap(ResourceFieldResolverProvider::getFieldType, - hydrator -> hydrator.getFieldResolver(resourceSpec))); + private Map getFieldResolvers(EntitySpec entitySpec) { + return _entityFieldResolverProviders.stream() + .collect(Collectors.toMap(EntityFieldResolverProvider::getFieldType, + hydrator -> hydrator.getFieldResolver(entitySpec))); } } diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/FilterUtils.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/FilterUtils.java index 76ed18e2baf78..0dbb9cd132f8a 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/FilterUtils.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/FilterUtils.java @@ -26,7 +26,7 @@ private FilterUtils() { * Creates new PolicyMatchCriterion with field and value, using EQUAL PolicyMatchCondition. */ @Nonnull - public static PolicyMatchCriterion newCriterion(@Nonnull ResourceFieldType field, @Nonnull List values) { + public static PolicyMatchCriterion newCriterion(@Nonnull EntityFieldType field, @Nonnull List values) { return newCriterion(field, values, PolicyMatchCondition.EQUALS); } @@ -34,7 +34,7 @@ public static PolicyMatchCriterion newCriterion(@Nonnull ResourceFieldType field * Creates new PolicyMatchCriterion with field, value and PolicyMatchCondition. */ @Nonnull - public static PolicyMatchCriterion newCriterion(@Nonnull ResourceFieldType field, @Nonnull List values, + public static PolicyMatchCriterion newCriterion(@Nonnull EntityFieldType field, @Nonnull List values, @Nonnull PolicyMatchCondition policyMatchCondition) { return new PolicyMatchCriterion().setField(field.name()) .setValues(new StringArray(values)) @@ -45,7 +45,7 @@ public static PolicyMatchCriterion newCriterion(@Nonnull ResourceFieldType field * Creates new PolicyMatchFilter from a map of Criteria by removing null-valued Criteria and using EQUAL PolicyMatchCondition (default). */ @Nonnull - public static PolicyMatchFilter newFilter(@Nullable Map> params) { + public static PolicyMatchFilter newFilter(@Nullable Map> params) { if (params == null) { return EMPTY_FILTER; } @@ -61,7 +61,7 @@ public static PolicyMatchFilter newFilter(@Nullable Map values) { + public static PolicyMatchFilter newFilter(@Nonnull EntityFieldType field, @Nonnull List values) { return newFilter(Collections.singletonMap(field, values)); } } diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyEngine.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyEngine.java index 6a36fac7de4e0..f8c017ea74e1f 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyEngine.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyEngine.java @@ -1,7 +1,6 @@ package com.datahub.authorization; import com.datahub.authentication.Authentication; -import com.google.common.collect.ImmutableSet; import com.linkedin.common.Owner; import com.linkedin.common.Ownership; import com.linkedin.common.urn.Urn; @@ -11,8 +10,6 @@ import com.linkedin.entity.EnvelopedAspect; import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.entity.client.EntityClient; -import com.linkedin.identity.GroupMembership; -import com.linkedin.identity.NativeGroupMembership; import com.linkedin.identity.RoleMembership; import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; @@ -23,7 +20,7 @@ import com.linkedin.policy.PolicyMatchCriterion; import com.linkedin.policy.PolicyMatchCriterionArray; import com.linkedin.policy.PolicyMatchFilter; -import java.net.URISyntaxException; + import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; @@ -34,6 +31,7 @@ import java.util.stream.Collectors; import java.util.stream.Stream; import javax.annotation.Nullable; + import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -49,37 +47,22 @@ public class PolicyEngine { public PolicyEvaluationResult evaluatePolicy( final DataHubPolicyInfo policy, - final String actorStr, + final ResolvedEntitySpec resolvedActorSpec, final String privilege, - final Optional resource) { - try { - // Currently Actor must be an urn. Consider whether this contract should be pushed up. - final Urn actor = Urn.createFromString(actorStr); - return evaluatePolicy(policy, actor, privilege, resource); - } catch (URISyntaxException e) { - log.error(String.format("Failed to bind actor %s to an URN. Actors must be URNs. Denying the authorization request", actorStr)); - return PolicyEvaluationResult.DENIED; - } - } - - public PolicyEvaluationResult evaluatePolicy( - final DataHubPolicyInfo policy, - final Urn actor, - final String privilege, - final Optional resource) { + final Optional resource) { final PolicyEvaluationContext context = new PolicyEvaluationContext(); log.debug("Evaluating policy {}", policy.getDisplayName()); // If the privilege is not in scope, deny the request. - if (!isPrivilegeMatch(privilege, policy.getPrivileges(), context)) { + if (!isPrivilegeMatch(privilege, policy.getPrivileges())) { log.debug("Policy denied based on irrelevant privileges {} for {}", policy.getPrivileges(), privilege); return PolicyEvaluationResult.DENIED; } // If policy is not applicable, deny the request - if (!isPolicyApplicable(policy, actor, resource, context)) { - log.debug("Policy does not applicable for actor {} and resource {}", actor, resource); + if (!isPolicyApplicable(policy, resolvedActorSpec, resource, context)) { + log.debug("Policy does not applicable for actor {} and resource {}", resolvedActorSpec.getSpec().getEntity(), resource); return PolicyEvaluationResult.DENIED; } @@ -89,7 +72,7 @@ public PolicyEvaluationResult evaluatePolicy( public PolicyActors getMatchingActors( final DataHubPolicyInfo policy, - final Optional resource) { + final Optional resource) { final List users = new ArrayList<>(); final List groups = new ArrayList<>(); boolean allUsers = false; @@ -126,8 +109,8 @@ public PolicyActors getMatchingActors( private boolean isPolicyApplicable( final DataHubPolicyInfo policy, - final Urn actor, - final Optional resource, + final ResolvedEntitySpec resolvedActorSpec, + final Optional resource, final PolicyEvaluationContext context ) { @@ -137,25 +120,21 @@ private boolean isPolicyApplicable( } // If the resource is not in scope, deny the request. - if (!isResourceMatch(policy.getType(), policy.getResources(), resource, context)) { + if (!isResourceMatch(policy.getType(), policy.getResources(), resource)) { return false; } // If the actor does not match, deny the request. - if (!isActorMatch(actor, policy.getActors(), resource, context)) { - return false; - } - - return true; + return isActorMatch(resolvedActorSpec, policy.getActors(), resource, context); } public List getGrantedPrivileges( final List policies, - final Urn actor, - final Optional resource) { + final ResolvedEntitySpec resolvedActorSpec, + final Optional resource) { PolicyEvaluationContext context = new PolicyEvaluationContext(); return policies.stream() - .filter(policy -> isPolicyApplicable(policy, actor, resource, context)) + .filter(policy -> isPolicyApplicable(policy, resolvedActorSpec, resource, context)) .flatMap(policy -> policy.getPrivileges().stream()) .distinct() .collect(Collectors.toList()); @@ -168,9 +147,8 @@ public List getGrantedPrivileges( * If the policy is of type "METADATA", the resourceSpec parameter will be matched against the * resource filter defined on the policy. */ - public Boolean policyMatchesResource(final DataHubPolicyInfo policy, final Optional resourceSpec) { - return isResourceMatch(policy.getType(), policy.getResources(), resourceSpec, - new PolicyEvaluationContext()); + public Boolean policyMatchesResource(final DataHubPolicyInfo policy, final Optional resourceSpec) { + return isResourceMatch(policy.getType(), policy.getResources(), resourceSpec); } /** @@ -178,8 +156,7 @@ public Boolean policyMatchesResource(final DataHubPolicyInfo policy, final Optio */ private boolean isPrivilegeMatch( final String requestPrivilege, - final List policyPrivileges, - final PolicyEvaluationContext context) { + final List policyPrivileges) { return policyPrivileges.contains(requestPrivilege); } @@ -189,8 +166,7 @@ private boolean isPrivilegeMatch( private boolean isResourceMatch( final String policyType, final @Nullable DataHubResourceFilter policyResourceFilter, - final Optional requestResource, - final PolicyEvaluationContext context) { + final Optional requestResource) { if (PoliciesConfig.PLATFORM_POLICY_TYPE.equals(policyType)) { // Currently, platform policies have no associated resource. return true; @@ -199,7 +175,7 @@ private boolean isResourceMatch( // No resource defined on the policy. return true; } - if (!requestResource.isPresent()) { + if (requestResource.isEmpty()) { // Resource filter present in policy, but no resource spec provided. log.debug("Resource filter present in policy, but no resource spec provided."); return false; @@ -218,31 +194,31 @@ private PolicyMatchFilter getFilter(DataHubResourceFilter policyResourceFilter) } PolicyMatchCriterionArray criteria = new PolicyMatchCriterionArray(); if (policyResourceFilter.hasType()) { - criteria.add(new PolicyMatchCriterion().setField(ResourceFieldType.RESOURCE_TYPE.name()) + criteria.add(new PolicyMatchCriterion().setField(EntityFieldType.TYPE.name()) .setValues(new StringArray(Collections.singletonList(policyResourceFilter.getType())))); } if (policyResourceFilter.hasType() && policyResourceFilter.hasResources() && !policyResourceFilter.isAllResources()) { criteria.add( - new PolicyMatchCriterion().setField(ResourceFieldType.RESOURCE_URN.name()).setValues(policyResourceFilter.getResources())); + new PolicyMatchCriterion().setField(EntityFieldType.URN.name()).setValues(policyResourceFilter.getResources())); } return new PolicyMatchFilter().setCriteria(criteria); } - private boolean checkFilter(final PolicyMatchFilter filter, final ResolvedResourceSpec resource) { + private boolean checkFilter(final PolicyMatchFilter filter, final ResolvedEntitySpec resource) { return filter.getCriteria().stream().allMatch(criterion -> checkCriterion(criterion, resource)); } - private boolean checkCriterion(final PolicyMatchCriterion criterion, final ResolvedResourceSpec resource) { - ResourceFieldType resourceFieldType; + private boolean checkCriterion(final PolicyMatchCriterion criterion, final ResolvedEntitySpec resource) { + EntityFieldType entityFieldType; try { - resourceFieldType = ResourceFieldType.valueOf(criterion.getField().toUpperCase()); + entityFieldType = EntityFieldType.valueOf(criterion.getField().toUpperCase()); } catch (IllegalArgumentException e) { log.error("Unsupported field type {}", criterion.getField()); return false; } - Set fieldValues = resource.getFieldValues(resourceFieldType); + Set fieldValues = resource.getFieldValues(entityFieldType); return criterion.getValues() .stream() .anyMatch(filterValue -> checkCondition(fieldValues, filterValue, criterion.getCondition())); @@ -257,46 +233,51 @@ private boolean checkCondition(Set fieldValues, String filterValue, Poli } /** + * Returns true if the actor portion of a DataHub policy matches a the actor being evaluated, false otherwise. * Returns true if the actor portion of a DataHub policy matches a the actor being evaluated, false otherwise. */ private boolean isActorMatch( - final Urn actor, + final ResolvedEntitySpec resolvedActorSpec, final DataHubActorFilter actorFilter, - final Optional resourceSpec, + final Optional resourceSpec, final PolicyEvaluationContext context) { // 1. If the actor is a matching "User" in the actor filter, return true immediately. - if (isUserMatch(actor, actorFilter)) { + if (isUserMatch(resolvedActorSpec, actorFilter)) { return true; } // 2. If the actor is in a matching "Group" in the actor filter, return true immediately. - if (isGroupMatch(actor, actorFilter, context)) { + if (isGroupMatch(resolvedActorSpec, actorFilter, context)) { return true; } // 3. If the actor is the owner, either directly or indirectly via a group, return true immediately. - if (isOwnerMatch(actor, actorFilter, resourceSpec, context)) { + if (isOwnerMatch(resolvedActorSpec, actorFilter, resourceSpec, context)) { return true; } // 4. If the actor is in a matching "Role" in the actor filter, return true immediately. - return isRoleMatch(actor, actorFilter, context); + return isRoleMatch(resolvedActorSpec, actorFilter, context); } - private boolean isUserMatch(final Urn actor, final DataHubActorFilter actorFilter) { + private boolean isUserMatch(final ResolvedEntitySpec resolvedActorSpec, final DataHubActorFilter actorFilter) { // If the actor is a matching "User" in the actor filter, return true immediately. return actorFilter.isAllUsers() || (actorFilter.hasUsers() && Objects.requireNonNull(actorFilter.getUsers()) - .stream() - .anyMatch(user -> user.equals(actor))); + .stream().map(Urn::toString) + .anyMatch(user -> user.equals(resolvedActorSpec.getSpec().getEntity()))); } - private boolean isGroupMatch(final Urn actor, final DataHubActorFilter actorFilter, final PolicyEvaluationContext context) { + private boolean isGroupMatch( + final ResolvedEntitySpec resolvedActorSpec, + final DataHubActorFilter actorFilter, + final PolicyEvaluationContext context) { // If the actor is in a matching "Group" in the actor filter, return true immediately. if (actorFilter.isAllGroups() || actorFilter.hasGroups()) { - final Set groups = resolveGroups(actor, context); - return actorFilter.isAllGroups() || (actorFilter.hasGroups() && Objects.requireNonNull(actorFilter.getGroups()) - .stream() + final Set groups = resolveGroups(resolvedActorSpec, context); + return (actorFilter.isAllGroups() && !groups.isEmpty()) + || (actorFilter.hasGroups() && Objects.requireNonNull(actorFilter.getGroups()) + .stream().map(Urn::toString) .anyMatch(groups::contains)); } // If there are no groups on the policy, return false for the group match. @@ -304,24 +285,24 @@ private boolean isGroupMatch(final Urn actor, final DataHubActorFilter actorFilt } private boolean isOwnerMatch( - final Urn actor, + final ResolvedEntitySpec resolvedActorSpec, final DataHubActorFilter actorFilter, - final Optional requestResource, + final Optional requestResource, final PolicyEvaluationContext context) { // If the policy does not apply to owners, or there is no resource to own, return false immediately. - if (!actorFilter.isResourceOwners() || !requestResource.isPresent()) { + if (!actorFilter.isResourceOwners() || requestResource.isEmpty()) { return false; } List ownershipTypes = actorFilter.getResourceOwnersTypes(); - return isActorOwner(actor, requestResource.get(), ownershipTypes, context); + return isActorOwner(resolvedActorSpec, requestResource.get(), ownershipTypes, context); } - private Set getOwnersForType(ResourceSpec resourceSpec, List ownershipTypes) { - Urn entityUrn = UrnUtils.getUrn(resourceSpec.getResource()); + private Set getOwnersForType(EntitySpec resourceSpec, List ownershipTypes) { + Urn entityUrn = UrnUtils.getUrn(resourceSpec.getEntity()); EnvelopedAspect ownershipAspect; try { EntityResponse response = _entityClient.getV2(entityUrn.getEntityType(), entityUrn, - Collections.singleton(Constants.OWNERSHIP_ASPECT_NAME), _systemAuthentication); + Collections.singleton(Constants.OWNERSHIP_ASPECT_NAME), _systemAuthentication); if (response == null || !response.getAspects().containsKey(Constants.OWNERSHIP_ASPECT_NAME)) { return Collections.emptySet(); } @@ -338,50 +319,56 @@ private Set getOwnersForType(ResourceSpec resourceSpec, List owners return ownersStream.map(owner -> owner.getOwner().toString()).collect(Collectors.toSet()); } - private boolean isActorOwner(Urn actor, ResolvedResourceSpec resourceSpec, List ownershipTypes, PolicyEvaluationContext context) { + private boolean isActorOwner( + final ResolvedEntitySpec resolvedActorSpec, + ResolvedEntitySpec resourceSpec, List ownershipTypes, + PolicyEvaluationContext context) { Set owners = this.getOwnersForType(resourceSpec.getSpec(), ownershipTypes); - if (isUserOwner(actor, owners)) { - return true; - } - final Set groups = resolveGroups(actor, context); - if (isGroupOwner(groups, owners)) { + if (isUserOwner(resolvedActorSpec, owners)) { return true; } - return false; + final Set groups = resolveGroups(resolvedActorSpec, context); + + return isGroupOwner(groups, owners); } - private boolean isUserOwner(Urn actor, Set owners) { - return owners.contains(actor.toString()); + private boolean isUserOwner(final ResolvedEntitySpec resolvedActorSpec, Set owners) { + return owners.contains(resolvedActorSpec.getSpec().getEntity()); } - private boolean isGroupOwner(Set groups, Set owners) { - return groups.stream().anyMatch(group -> owners.contains(group.toString())); + private boolean isGroupOwner(Set groups, Set owners) { + return groups.stream().anyMatch(owners::contains); } - private boolean isRoleMatch(final Urn actor, final DataHubActorFilter actorFilter, + private boolean isRoleMatch( + final ResolvedEntitySpec resolvedActorSpec, + final DataHubActorFilter actorFilter, final PolicyEvaluationContext context) { // Can immediately return false if the actor filter does not have any roles if (!actorFilter.hasRoles()) { return false; } // If the actor has a matching "Role" in the actor filter, return true immediately. - Set actorRoles = resolveRoles(actor, context); + Set actorRoles = resolveRoles(resolvedActorSpec, context); return Objects.requireNonNull(actorFilter.getRoles()) .stream() .anyMatch(actorRoles::contains); } - private Set resolveRoles(Urn actor, PolicyEvaluationContext context) { + private Set resolveRoles(final ResolvedEntitySpec resolvedActorSpec, PolicyEvaluationContext context) { if (context.roles != null) { return context.roles; } + String actor = resolvedActorSpec.getSpec().getEntity(); + Set roles = new HashSet<>(); final EnvelopedAspectMap aspectMap; try { - final EntityResponse corpUser = _entityClient.batchGetV2(CORP_USER_ENTITY_NAME, Collections.singleton(actor), - Collections.singleton(ROLE_MEMBERSHIP_ASPECT_NAME), _systemAuthentication).get(actor); + Urn actorUrn = Urn.createFromString(actor); + final EntityResponse corpUser = _entityClient.batchGetV2(CORP_USER_ENTITY_NAME, Collections.singleton(actorUrn), + Collections.singleton(ROLE_MEMBERSHIP_ASPECT_NAME), _systemAuthentication).get(actorUrn); if (corpUser == null || !corpUser.hasAspects()) { return roles; } @@ -403,62 +390,25 @@ private Set resolveRoles(Urn actor, PolicyEvaluationContext context) { return roles; } - private Set resolveGroups(Urn actor, PolicyEvaluationContext context) { + private Set resolveGroups(ResolvedEntitySpec resolvedActorSpec, PolicyEvaluationContext context) { if (context.groups != null) { return context.groups; } - Set groups = new HashSet<>(); - final EnvelopedAspectMap aspectMap; - - try { - final EntityResponse corpUser = _entityClient.batchGetV2(CORP_USER_ENTITY_NAME, Collections.singleton(actor), - ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME), _systemAuthentication) - .get(actor); - if (corpUser == null || !corpUser.hasAspects()) { - return groups; - } - aspectMap = corpUser.getAspects(); - } catch (Exception e) { - throw new RuntimeException(String.format("Failed to fetch %s and %s for urn %s", GROUP_MEMBERSHIP_ASPECT_NAME, - NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME, actor), e); - } - - Optional maybeGroupMembership = resolveGroupMembership(aspectMap); - maybeGroupMembership.ifPresent(groupMembership -> groups.addAll(groupMembership.getGroups())); - - Optional maybeNativeGroupMembership = resolveNativeGroupMembership(aspectMap); - maybeNativeGroupMembership.ifPresent( - nativeGroupMembership -> groups.addAll(nativeGroupMembership.getNativeGroups())); + Set groups = resolvedActorSpec.getGroupMembership(); context.setGroups(groups); // Cache the groups. return groups; } - // TODO: Optimization - Cache the group membership. Refresh periodically. - private Optional resolveGroupMembership(final EnvelopedAspectMap aspectMap) { - if (aspectMap.containsKey(GROUP_MEMBERSHIP_ASPECT_NAME)) { - return Optional.of(new GroupMembership(aspectMap.get(GROUP_MEMBERSHIP_ASPECT_NAME).getValue().data())); - } - return Optional.empty(); - } - - private Optional resolveNativeGroupMembership(final EnvelopedAspectMap aspectMap) { - if (aspectMap.containsKey(NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)) { - return Optional.of( - new NativeGroupMembership(aspectMap.get(NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME).getValue().data())); - } - return Optional.empty(); - } - /** * Class used to store state across a single Policy evaluation. */ static class PolicyEvaluationContext { - private Set groups; + private Set groups; private Set roles; - public void setGroups(Set groups) { + public void setGroups(Set groups) { this.groups = groups; } diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProvider.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProvider.java index cd838625c2ca1..27cb8fcee8138 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProvider.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProvider.java @@ -1,45 +1,45 @@ package com.datahub.authorization.fieldresolverprovider; +import static com.linkedin.metadata.Constants.DATA_PLATFORM_INSTANCE_ASPECT_NAME; +import static com.linkedin.metadata.Constants.DATA_PLATFORM_INSTANCE_ENTITY_NAME; + import com.datahub.authentication.Authentication; +import com.datahub.authorization.EntityFieldType; +import com.datahub.authorization.EntitySpec; import com.datahub.authorization.FieldResolver; -import com.datahub.authorization.ResourceFieldType; -import com.datahub.authorization.ResourceSpec; import com.linkedin.common.DataPlatformInstance; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.EnvelopedAspect; import com.linkedin.entity.client.EntityClient; -import lombok.RequiredArgsConstructor; -import lombok.extern.slf4j.Slf4j; - import java.util.Collections; import java.util.Objects; - -import static com.linkedin.metadata.Constants.*; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; /** * Provides field resolver for domain given resourceSpec */ @Slf4j @RequiredArgsConstructor -public class DataPlatformInstanceFieldResolverProvider implements ResourceFieldResolverProvider { +public class DataPlatformInstanceFieldResolverProvider implements EntityFieldResolverProvider { private final EntityClient _entityClient; private final Authentication _systemAuthentication; @Override - public ResourceFieldType getFieldType() { - return ResourceFieldType.DATA_PLATFORM_INSTANCE; + public EntityFieldType getFieldType() { + return EntityFieldType.DATA_PLATFORM_INSTANCE; } @Override - public FieldResolver getFieldResolver(ResourceSpec resourceSpec) { - return FieldResolver.getResolverFromFunction(resourceSpec, this::getDataPlatformInstance); + public FieldResolver getFieldResolver(EntitySpec entitySpec) { + return FieldResolver.getResolverFromFunction(entitySpec, this::getDataPlatformInstance); } - private FieldResolver.FieldValue getDataPlatformInstance(ResourceSpec resourceSpec) { - Urn entityUrn = UrnUtils.getUrn(resourceSpec.getResource()); + private FieldResolver.FieldValue getDataPlatformInstance(EntitySpec entitySpec) { + Urn entityUrn = UrnUtils.getUrn(entitySpec.getEntity()); // In the case that the entity is a platform instance, the associated platform instance entity is the instance itself if (entityUrn.getEntityType().equals(DATA_PLATFORM_INSTANCE_ENTITY_NAME)) { return FieldResolver.FieldValue.builder() diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DomainFieldResolverProvider.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DomainFieldResolverProvider.java index 68c1dd4f644e5..25c2165f02b94 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DomainFieldResolverProvider.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DomainFieldResolverProvider.java @@ -2,8 +2,8 @@ import com.datahub.authentication.Authentication; import com.datahub.authorization.FieldResolver; -import com.datahub.authorization.ResourceFieldType; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntityFieldType; +import com.datahub.authorization.EntitySpec; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.domain.DomainProperties; @@ -27,23 +27,23 @@ /** - * Provides field resolver for domain given resourceSpec + * Provides field resolver for domain given entitySpec */ @Slf4j @RequiredArgsConstructor -public class DomainFieldResolverProvider implements ResourceFieldResolverProvider { +public class DomainFieldResolverProvider implements EntityFieldResolverProvider { private final EntityClient _entityClient; private final Authentication _systemAuthentication; @Override - public ResourceFieldType getFieldType() { - return ResourceFieldType.DOMAIN; + public EntityFieldType getFieldType() { + return EntityFieldType.DOMAIN; } @Override - public FieldResolver getFieldResolver(ResourceSpec resourceSpec) { - return FieldResolver.getResolverFromFunction(resourceSpec, this::getDomains); + public FieldResolver getFieldResolver(EntitySpec entitySpec) { + return FieldResolver.getResolverFromFunction(entitySpec, this::getDomains); } private Set getBatchedParentDomains(@Nonnull final Set urns) { @@ -78,8 +78,8 @@ private Set getBatchedParentDomains(@Nonnull final Set urns) { return parentUrns; } - private FieldResolver.FieldValue getDomains(ResourceSpec resourceSpec) { - final Urn entityUrn = UrnUtils.getUrn(resourceSpec.getResource()); + private FieldResolver.FieldValue getDomains(EntitySpec entitySpec) { + final Urn entityUrn = UrnUtils.getUrn(entitySpec.getEntity()); // In the case that the entity is a domain, the associated domain is the domain itself if (entityUrn.getEntityType().equals(DOMAIN_ENTITY_NAME)) { return FieldResolver.FieldValue.builder() diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityFieldResolverProvider.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityFieldResolverProvider.java new file mode 100644 index 0000000000000..a76db0ecb5102 --- /dev/null +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityFieldResolverProvider.java @@ -0,0 +1,22 @@ +package com.datahub.authorization.fieldresolverprovider; + +import com.datahub.authorization.FieldResolver; +import com.datahub.authorization.EntityFieldType; +import com.datahub.authorization.EntitySpec; + + +/** + * Base class for defining a class that provides the field resolver for the given field type + */ +public interface EntityFieldResolverProvider { + + /** + * Field that this hydrator is hydrating + */ + EntityFieldType getFieldType(); + + /** + * Return resolver for fetching the field values given the entity + */ + FieldResolver getFieldResolver(EntitySpec entitySpec); +} diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityTypeFieldResolverProvider.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityTypeFieldResolverProvider.java index 58e3d78ce8c3b..187f696904947 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityTypeFieldResolverProvider.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityTypeFieldResolverProvider.java @@ -1,22 +1,22 @@ package com.datahub.authorization.fieldresolverprovider; import com.datahub.authorization.FieldResolver; -import com.datahub.authorization.ResourceFieldType; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntityFieldType; +import com.datahub.authorization.EntitySpec; import java.util.Collections; /** - * Provides field resolver for entity type given resourceSpec + * Provides field resolver for entity type given entitySpec */ -public class EntityTypeFieldResolverProvider implements ResourceFieldResolverProvider { +public class EntityTypeFieldResolverProvider implements EntityFieldResolverProvider { @Override - public ResourceFieldType getFieldType() { - return ResourceFieldType.RESOURCE_TYPE; + public EntityFieldType getFieldType() { + return EntityFieldType.TYPE; } @Override - public FieldResolver getFieldResolver(ResourceSpec resourceSpec) { - return FieldResolver.getResolverFromValues(Collections.singleton(resourceSpec.getType())); + public FieldResolver getFieldResolver(EntitySpec entitySpec) { + return FieldResolver.getResolverFromValues(Collections.singleton(entitySpec.getType())); } } diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityUrnFieldResolverProvider.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityUrnFieldResolverProvider.java index b9d98f1dcbac0..2f5c4a7c6c961 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityUrnFieldResolverProvider.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/EntityUrnFieldResolverProvider.java @@ -1,22 +1,22 @@ package com.datahub.authorization.fieldresolverprovider; import com.datahub.authorization.FieldResolver; -import com.datahub.authorization.ResourceFieldType; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntityFieldType; +import com.datahub.authorization.EntitySpec; import java.util.Collections; /** - * Provides field resolver for entity urn given resourceSpec + * Provides field resolver for entity urn given entitySpec */ -public class EntityUrnFieldResolverProvider implements ResourceFieldResolverProvider { +public class EntityUrnFieldResolverProvider implements EntityFieldResolverProvider { @Override - public ResourceFieldType getFieldType() { - return ResourceFieldType.RESOURCE_URN; + public EntityFieldType getFieldType() { + return EntityFieldType.URN; } @Override - public FieldResolver getFieldResolver(ResourceSpec resourceSpec) { - return FieldResolver.getResolverFromValues(Collections.singleton(resourceSpec.getResource())); + public FieldResolver getFieldResolver(EntitySpec entitySpec) { + return FieldResolver.getResolverFromValues(Collections.singleton(entitySpec.getEntity())); } } diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/GroupMembershipFieldResolverProvider.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/GroupMembershipFieldResolverProvider.java new file mode 100644 index 0000000000000..8db029632d7e2 --- /dev/null +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/GroupMembershipFieldResolverProvider.java @@ -0,0 +1,78 @@ +package com.datahub.authorization.fieldresolverprovider; + +import com.datahub.authentication.Authentication; +import com.datahub.authorization.FieldResolver; +import com.datahub.authorization.EntityFieldType; +import com.datahub.authorization.EntitySpec; +import com.google.common.collect.ImmutableSet; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.identity.NativeGroupMembership; +import com.linkedin.metadata.Constants; +import com.linkedin.identity.GroupMembership; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Collectors; + +import static com.linkedin.metadata.Constants.GROUP_MEMBERSHIP_ASPECT_NAME; +import static com.linkedin.metadata.Constants.NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME; + + +/** + * Provides field resolver for owners given entitySpec + */ +@Slf4j +@RequiredArgsConstructor +public class GroupMembershipFieldResolverProvider implements EntityFieldResolverProvider { + + private final EntityClient _entityClient; + private final Authentication _systemAuthentication; + + @Override + public EntityFieldType getFieldType() { + return EntityFieldType.GROUP_MEMBERSHIP; + } + + @Override + public FieldResolver getFieldResolver(EntitySpec entitySpec) { + return FieldResolver.getResolverFromFunction(entitySpec, this::getGroupMembership); + } + + private FieldResolver.FieldValue getGroupMembership(EntitySpec entitySpec) { + Urn entityUrn = UrnUtils.getUrn(entitySpec.getEntity()); + EnvelopedAspect groupMembershipAspect; + EnvelopedAspect nativeGroupMembershipAspect; + List groups = new ArrayList<>(); + try { + EntityResponse response = _entityClient.getV2(entityUrn.getEntityType(), entityUrn, + ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME), _systemAuthentication); + if (response == null + || !(response.getAspects().containsKey(Constants.GROUP_MEMBERSHIP_ASPECT_NAME) + || response.getAspects().containsKey(Constants.NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME))) { + return FieldResolver.emptyFieldValue(); + } + if (response.getAspects().containsKey(Constants.GROUP_MEMBERSHIP_ASPECT_NAME)) { + groupMembershipAspect = response.getAspects().get(Constants.GROUP_MEMBERSHIP_ASPECT_NAME); + GroupMembership groupMembership = new GroupMembership(groupMembershipAspect.getValue().data()); + groups.addAll(groupMembership.getGroups()); + } + if (response.getAspects().containsKey(Constants.NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)) { + nativeGroupMembershipAspect = response.getAspects().get(Constants.NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME); + NativeGroupMembership nativeGroupMembership = new NativeGroupMembership(nativeGroupMembershipAspect.getValue().data()); + groups.addAll(nativeGroupMembership.getNativeGroups()); + } + } catch (Exception e) { + log.error("Error while retrieving group membership aspect for urn {}", entityUrn, e); + return FieldResolver.emptyFieldValue(); + } + return FieldResolver.FieldValue.builder() + .values(groups.stream().map(Urn::toString).collect(Collectors.toSet())) + .build(); + } +} diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/OwnerFieldResolverProvider.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/OwnerFieldResolverProvider.java index 20ec6a09377c8..bdd652d1d3871 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/OwnerFieldResolverProvider.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/OwnerFieldResolverProvider.java @@ -2,8 +2,8 @@ import com.datahub.authentication.Authentication; import com.datahub.authorization.FieldResolver; -import com.datahub.authorization.ResourceFieldType; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntityFieldType; +import com.datahub.authorization.EntitySpec; import com.linkedin.common.Ownership; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; @@ -18,27 +18,27 @@ /** - * Provides field resolver for owners given resourceSpec + * Provides field resolver for owners given entitySpec */ @Slf4j @RequiredArgsConstructor -public class OwnerFieldResolverProvider implements ResourceFieldResolverProvider { +public class OwnerFieldResolverProvider implements EntityFieldResolverProvider { private final EntityClient _entityClient; private final Authentication _systemAuthentication; @Override - public ResourceFieldType getFieldType() { - return ResourceFieldType.OWNER; + public EntityFieldType getFieldType() { + return EntityFieldType.OWNER; } @Override - public FieldResolver getFieldResolver(ResourceSpec resourceSpec) { - return FieldResolver.getResolverFromFunction(resourceSpec, this::getOwners); + public FieldResolver getFieldResolver(EntitySpec entitySpec) { + return FieldResolver.getResolverFromFunction(entitySpec, this::getOwners); } - private FieldResolver.FieldValue getOwners(ResourceSpec resourceSpec) { - Urn entityUrn = UrnUtils.getUrn(resourceSpec.getResource()); + private FieldResolver.FieldValue getOwners(EntitySpec entitySpec) { + Urn entityUrn = UrnUtils.getUrn(entitySpec.getEntity()); EnvelopedAspect ownershipAspect; try { EntityResponse response = _entityClient.getV2(entityUrn.getEntityType(), entityUrn, diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/ResourceFieldResolverProvider.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/ResourceFieldResolverProvider.java deleted file mode 100644 index 4ba4200f8035e..0000000000000 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/ResourceFieldResolverProvider.java +++ /dev/null @@ -1,22 +0,0 @@ -package com.datahub.authorization.fieldresolverprovider; - -import com.datahub.authorization.FieldResolver; -import com.datahub.authorization.ResourceFieldType; -import com.datahub.authorization.ResourceSpec; - - -/** - * Base class for defining a class that provides the field resolver for the given field type - */ -public interface ResourceFieldResolverProvider { - - /** - * Field that this hydrator is hydrating - */ - ResourceFieldType getFieldType(); - - /** - * Return resolver for fetching the field values given the resource - */ - FieldResolver getFieldResolver(ResourceSpec resourceSpec); -} diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java index 2e48123fb1813..24ecfa6fefc85 100644 --- a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java @@ -158,7 +158,7 @@ public void testSystemAuthentication() throws Exception { // Validate that the System Actor is authorized, even if there is no policy. - ResourceSpec resourceSpec = new ResourceSpec("dataset", "urn:li:dataset:test"); + EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:test"); AuthorizationRequest request = new AuthorizationRequest( new Actor(ActorType.USER, DATAHUB_SYSTEM_CLIENT_ID).toUrnStr(), @@ -172,7 +172,7 @@ public void testSystemAuthentication() throws Exception { @Test public void testAuthorizeGranted() throws Exception { - ResourceSpec resourceSpec = new ResourceSpec("dataset", "urn:li:dataset:test"); + EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:test"); AuthorizationRequest request = new AuthorizationRequest( "urn:li:corpuser:test", @@ -186,7 +186,7 @@ public void testAuthorizeGranted() throws Exception { @Test public void testAuthorizeNotGranted() throws Exception { - ResourceSpec resourceSpec = new ResourceSpec("dataset", "urn:li:dataset:test"); + EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:test"); // Policy for this privilege is inactive. AuthorizationRequest request = new AuthorizationRequest( @@ -203,7 +203,7 @@ public void testAllowAllMode() throws Exception { _dataHubAuthorizer.setMode(DataHubAuthorizer.AuthorizationMode.ALLOW_ALL); - ResourceSpec resourceSpec = new ResourceSpec("dataset", "urn:li:dataset:test"); + EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:test"); // Policy for this privilege is inactive. AuthorizationRequest request = new AuthorizationRequest( @@ -219,7 +219,7 @@ public void testAllowAllMode() throws Exception { public void testInvalidateCache() throws Exception { // First make sure that the default policies are as expected. - ResourceSpec resourceSpec = new ResourceSpec("dataset", "urn:li:dataset:test"); + EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:test"); AuthorizationRequest request = new AuthorizationRequest( "urn:li:corpuser:test", @@ -250,7 +250,7 @@ public void testInvalidateCache() throws Exception { public void testAuthorizedActorsActivePolicy() throws Exception { final AuthorizedActors actors = _dataHubAuthorizer.authorizedActors("EDIT_ENTITY_TAGS", // Should be inside the active policy. - Optional.of(new ResourceSpec("dataset", "urn:li:dataset:1"))); + Optional.of(new EntitySpec("dataset", "urn:li:dataset:1"))); assertTrue(actors.isAllUsers()); assertTrue(actors.isAllGroups()); @@ -272,7 +272,7 @@ public void testAuthorizedActorsActivePolicy() throws Exception { @Test public void testAuthorizationOnDomainWithPrivilegeIsAllowed() { - ResourceSpec resourceSpec = new ResourceSpec("dataset", "urn:li:dataset:test"); + EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:test"); AuthorizationRequest request = new AuthorizationRequest( "urn:li:corpuser:test", @@ -285,7 +285,7 @@ public void testAuthorizationOnDomainWithPrivilegeIsAllowed() { @Test public void testAuthorizationOnDomainWithParentPrivilegeIsAllowed() { - ResourceSpec resourceSpec = new ResourceSpec("dataset", "urn:li:dataset:test"); + EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:test"); AuthorizationRequest request = new AuthorizationRequest( "urn:li:corpuser:test", @@ -298,7 +298,7 @@ public void testAuthorizationOnDomainWithParentPrivilegeIsAllowed() { @Test public void testAuthorizationOnDomainWithoutPrivilegeIsDenied() { - ResourceSpec resourceSpec = new ResourceSpec("dataset", "urn:li:dataset:test"); + EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:test"); AuthorizationRequest request = new AuthorizationRequest( "urn:li:corpuser:test", @@ -334,7 +334,7 @@ private DataHubPolicyInfo createDataHubPolicyInfo(boolean active, List p resourceFilter.setType("dataset"); if (domain != null) { - resourceFilter.setFilter(FilterUtils.newFilter(ImmutableMap.of(ResourceFieldType.DOMAIN, Collections.singletonList(domain.toString())))); + resourceFilter.setFilter(FilterUtils.newFilter(ImmutableMap.of(EntityFieldType.DOMAIN, Collections.singletonList(domain.toString())))); } dataHubPolicyInfo.setResources(resourceFilter); @@ -398,6 +398,6 @@ private Map createDomainPropertiesBatchResponse(@Nullable f } private AuthorizerContext createAuthorizerContext(final Authentication systemAuthentication, final EntityClient entityClient) { - return new AuthorizerContext(Collections.emptyMap(), new DefaultResourceSpecResolver(systemAuthentication, entityClient)); + return new AuthorizerContext(Collections.emptyMap(), new DefaultEntitySpecResolver(systemAuthentication, entityClient)); } } diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/PolicyEngineTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/PolicyEngineTest.java index 99d8fee309d91..be8c948f8ef89 100644 --- a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/PolicyEngineTest.java +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/PolicyEngineTest.java @@ -11,15 +11,12 @@ import com.linkedin.common.OwnershipType; import com.linkedin.common.UrnArray; import com.linkedin.common.urn.Urn; -import com.linkedin.common.urn.UrnUtils; import com.linkedin.data.template.StringArray; import com.linkedin.entity.Aspect; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.EnvelopedAspect; import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.entity.client.EntityClient; -import com.linkedin.identity.CorpUserInfo; -import com.linkedin.identity.GroupMembership; import com.linkedin.identity.RoleMembership; import com.linkedin.metadata.Constants; import com.linkedin.policy.DataHubActorFilter; @@ -45,22 +42,19 @@ public class PolicyEngineTest { private static final String AUTHORIZED_PRINCIPAL = "urn:li:corpuser:datahub"; private static final String UNAUTHORIZED_PRINCIPAL = "urn:li:corpuser:unauthorized"; - private static final String AUTHORIZED_GROUP = "urn:li:corpGroup:authorizedGroup"; - private static final String RESOURCE_URN = "urn:li:dataset:test"; - private static final String DOMAIN_URN = "urn:li:domain:domain1"; - private static final String OWNERSHIP_TYPE_URN = "urn:li:ownershipType:__system__technical_owner"; - private static final String OTHER_OWNERSHIP_TYPE_URN = "urn:li:ownershipType:__system__data_steward"; private EntityClient _entityClient; private PolicyEngine _policyEngine; private Urn authorizedUserUrn; + private ResolvedEntitySpec resolvedAuthorizedUserSpec; private Urn unauthorizedUserUrn; + private ResolvedEntitySpec resolvedUnauthorizedUserSpec; private Urn resourceUrn; @BeforeMethod @@ -68,29 +62,34 @@ public void setupTest() throws Exception { _entityClient = Mockito.mock(EntityClient.class); _policyEngine = new PolicyEngine(Mockito.mock(Authentication.class), _entityClient); - // Init mocks. - EntityResponse authorizedEntityResponse = createAuthorizedEntityResponse(); authorizedUserUrn = Urn.createFromString(AUTHORIZED_PRINCIPAL); + resolvedAuthorizedUserSpec = buildEntityResolvers(CORP_USER_ENTITY_NAME, AUTHORIZED_PRINCIPAL, + Collections.emptySet(), Collections.emptySet(), Collections.singleton(AUTHORIZED_GROUP)); + unauthorizedUserUrn = Urn.createFromString(UNAUTHORIZED_PRINCIPAL); + resolvedUnauthorizedUserSpec = buildEntityResolvers(CORP_USER_ENTITY_NAME, UNAUTHORIZED_PRINCIPAL); + resourceUrn = Urn.createFromString(RESOURCE_URN); + + // Init role membership mocks. + EntityResponse authorizedEntityResponse = createAuthorizedEntityResponse(); authorizedEntityResponse.setUrn(authorizedUserUrn); Map authorizedEntityResponseMap = Collections.singletonMap(authorizedUserUrn, authorizedEntityResponse); - when(_entityClient.batchGetV2(eq(CORP_USER_ENTITY_NAME), eq(Collections.singleton(authorizedUserUrn)), any(), - any())).thenReturn(authorizedEntityResponseMap); + when(_entityClient.batchGetV2(eq(CORP_USER_ENTITY_NAME), eq(Collections.singleton(authorizedUserUrn)), + eq(Collections.singleton(ROLE_MEMBERSHIP_ASPECT_NAME)), any())).thenReturn(authorizedEntityResponseMap); EntityResponse unauthorizedEntityResponse = createUnauthorizedEntityResponse(); - unauthorizedUserUrn = Urn.createFromString(UNAUTHORIZED_PRINCIPAL); unauthorizedEntityResponse.setUrn(unauthorizedUserUrn); Map unauthorizedEntityResponseMap = Collections.singletonMap(unauthorizedUserUrn, unauthorizedEntityResponse); - when(_entityClient.batchGetV2(eq(CORP_USER_ENTITY_NAME), eq(Collections.singleton(unauthorizedUserUrn)), any(), - any())).thenReturn(unauthorizedEntityResponseMap); + when(_entityClient.batchGetV2(eq(CORP_USER_ENTITY_NAME), eq(Collections.singleton(unauthorizedUserUrn)), + eq(Collections.singleton(ROLE_MEMBERSHIP_ASPECT_NAME)), any())).thenReturn(unauthorizedEntityResponseMap); + // Init ownership type mocks. EntityResponse entityResponse = new EntityResponse(); EnvelopedAspectMap envelopedAspectMap = new EnvelopedAspectMap(); envelopedAspectMap.put(OWNERSHIP_ASPECT_NAME, new EnvelopedAspect().setValue(new com.linkedin.entity.Aspect(createOwnershipAspect(true, true).data()))); entityResponse.setAspects(envelopedAspectMap); - resourceUrn = Urn.createFromString(RESOURCE_URN); Map mockMap = mock(Map.class); when(_entityClient.batchGetV2(any(), eq(Collections.singleton(resourceUrn)), eq(Collections.singleton(OWNERSHIP_ASPECT_NAME)), any())).thenReturn(mockMap); @@ -120,9 +119,9 @@ public void testEvaluatePolicyInactivePolicyState() { resourceFilter.setAllResources(true); resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); PolicyEngine.PolicyEvaluationResult result = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertFalse(result.isGranted()); @@ -149,9 +148,9 @@ public void testEvaluatePolicyPrivilegeFilterNoMatch() throws Exception { resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); PolicyEngine.PolicyEvaluationResult result = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_OWNERS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_OWNERS", Optional.of(resourceSpec)); assertFalse(result.isGranted()); @@ -176,7 +175,8 @@ public void testEvaluatePlatformPolicyPrivilegeFilterMatch() throws Exception { dataHubPolicyInfo.setActors(actorFilter); PolicyEngine.PolicyEvaluationResult result = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "MANAGE_POLICIES", Optional.empty()); + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "MANAGE_POLICIES", + Optional.empty()); assertTrue(result.isGranted()); // Verify no network calls @@ -208,10 +208,10 @@ public void testEvaluatePolicyActorFilterUserMatch() throws Exception { resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); // Assert Authorized user can edit entity tags. PolicyEngine.PolicyEvaluationResult result1 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result1.isGranted()); @@ -245,10 +245,10 @@ public void testEvaluatePolicyActorFilterUserNoMatch() throws Exception { resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); // Assert unauthorized user cannot edit entity tags. PolicyEngine.PolicyEvaluationResult result2 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, "urn:li:corpuser:test", "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, buildEntityResolvers(CORP_USER_ENTITY_NAME, "urn:li:corpuser:test"), "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertFalse(result2.isGranted()); @@ -270,7 +270,7 @@ public void testEvaluatePolicyActorFilterGroupMatch() throws Exception { final DataHubActorFilter actorFilter = new DataHubActorFilter(); final UrnArray groupsUrnArray = new UrnArray(); - groupsUrnArray.add(Urn.createFromString("urn:li:corpGroup:authorizedGroup")); + groupsUrnArray.add(Urn.createFromString(AUTHORIZED_GROUP)); actorFilter.setGroups(groupsUrnArray); actorFilter.setResourceOwners(false); actorFilter.setAllUsers(false); @@ -282,16 +282,15 @@ public void testEvaluatePolicyActorFilterGroupMatch() throws Exception { resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); // Assert authorized user can edit entity tags, because of group membership. PolicyEngine.PolicyEvaluationResult result1 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result1.isGranted()); - // Verify we are only calling for group during these requests. - verify(_entityClient, times(1)).batchGetV2(eq(CORP_USER_ENTITY_NAME), eq(Collections.singleton(authorizedUserUrn)), - any(), any()); + // Verify no network calls + verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any()); } @Test @@ -307,7 +306,7 @@ public void testEvaluatePolicyActorFilterGroupNoMatch() throws Exception { final DataHubActorFilter actorFilter = new DataHubActorFilter(); final UrnArray groupsUrnArray = new UrnArray(); - groupsUrnArray.add(Urn.createFromString("urn:li:corpGroup:authorizedGroup")); + groupsUrnArray.add(Urn.createFromString(AUTHORIZED_GROUP)); actorFilter.setGroups(groupsUrnArray); actorFilter.setResourceOwners(false); actorFilter.setAllUsers(false); @@ -319,16 +318,15 @@ public void testEvaluatePolicyActorFilterGroupNoMatch() throws Exception { resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); // Assert unauthorized user cannot edit entity tags. PolicyEngine.PolicyEvaluationResult result2 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, UNAUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedUnauthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertFalse(result2.isGranted()); - // Verify we are only calling for group during these requests. - verify(_entityClient, times(1)).batchGetV2(eq(CORP_USER_ENTITY_NAME), - eq(Collections.singleton(unauthorizedUserUrn)), any(), any()); + // Verify no network calls + verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any()); } @Test @@ -357,17 +355,17 @@ public void testEvaluatePolicyActorFilterRoleMatch() throws Exception { resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); // Assert authorized user can edit entity tags. PolicyEngine.PolicyEvaluationResult authorizedResult = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(authorizedResult.isGranted()); // Verify we are only calling for roles during these requests. - verify(_entityClient, times(1)).batchGetV2(eq(CORP_USER_ENTITY_NAME), eq(Collections.singleton(authorizedUserUrn)), - any(), any()); + verify(_entityClient, times(1)).batchGetV2(eq(CORP_USER_ENTITY_NAME), + eq(Collections.singleton(authorizedUserUrn)), any(), any()); } @Test @@ -396,10 +394,10 @@ public void testEvaluatePolicyActorFilterNoRoleMatch() throws Exception { resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); // Assert authorized user can edit entity tags. PolicyEngine.PolicyEvaluationResult unauthorizedResult = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, UNAUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedUnauthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertFalse(unauthorizedResult.isGranted()); @@ -431,16 +429,16 @@ public void testEvaluatePolicyActorFilterAllUsersMatch() throws Exception { resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); // Assert authorized user can edit entity tags, because of group membership. PolicyEngine.PolicyEvaluationResult result1 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result1.isGranted()); // Assert unauthorized user cannot edit entity tags. PolicyEngine.PolicyEvaluationResult result2 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, UNAUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedUnauthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result2.isGranted()); @@ -470,24 +468,21 @@ public void testEvaluatePolicyActorFilterAllGroupsMatch() throws Exception { resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); // Assert authorized user can edit entity tags, because of group membership. PolicyEngine.PolicyEvaluationResult result1 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result1.isGranted()); // Assert unauthorized user cannot edit entity tags. PolicyEngine.PolicyEvaluationResult result2 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, UNAUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedUnauthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); - assertTrue(result2.isGranted()); + assertFalse(result2.isGranted()); - // Verify we are only calling for group during these requests. - verify(_entityClient, times(1)).batchGetV2(eq(CORP_USER_ENTITY_NAME), eq(Collections.singleton(authorizedUserUrn)), - any(), any()); - verify(_entityClient, times(1)).batchGetV2(eq(CORP_USER_ENTITY_NAME), - eq(Collections.singleton(unauthorizedUserUrn)), any(), any()); + // Verify no network calls + verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any()); } @Test @@ -519,17 +514,17 @@ public void testEvaluatePolicyActorFilterUserResourceOwnersMatch() throws Except when(_entityClient.getV2(eq(resourceUrn.getEntityType()), eq(resourceUrn), eq(Collections.singleton(Constants.OWNERSHIP_ASPECT_NAME)), any())).thenReturn(entityResponse); - ResolvedResourceSpec resourceSpec = - buildResourceResolvers("dataset", RESOURCE_URN, ImmutableSet.of(AUTHORIZED_PRINCIPAL), Collections.emptySet()); + ResolvedEntitySpec resourceSpec = + buildEntityResolvers("dataset", RESOURCE_URN, ImmutableSet.of(AUTHORIZED_PRINCIPAL), Collections.emptySet(), + Collections.emptySet()); // Assert authorized user can edit entity tags, because he is a user owner. PolicyEngine.PolicyEvaluationResult result1 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result1.isGranted()); - // Ensure no calls for group membership. - verify(_entityClient, times(0)).batchGetV2(eq(CORP_USER_ENTITY_NAME), eq(Collections.singleton(authorizedUserUrn)), - eq(null), any()); + // Verify no network calls + verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any()); } @Test @@ -562,13 +557,17 @@ public void testEvaluatePolicyActorFilterUserResourceOwnersTypeMatch() throws Ex when(_entityClient.getV2(eq(resourceUrn.getEntityType()), eq(resourceUrn), eq(Collections.singleton(Constants.OWNERSHIP_ASPECT_NAME)), any())).thenReturn(entityResponse); - ResolvedResourceSpec resourceSpec = - buildResourceResolvers("dataset", RESOURCE_URN, ImmutableSet.of(AUTHORIZED_PRINCIPAL), Collections.emptySet()); + ResolvedEntitySpec resourceSpec = + buildEntityResolvers("dataset", RESOURCE_URN, ImmutableSet.of(AUTHORIZED_PRINCIPAL), Collections.emptySet(), + Collections.emptySet()); PolicyEngine.PolicyEvaluationResult result1 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result1.isGranted()); + + // Verify no network calls + verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any()); } @Test @@ -601,13 +600,16 @@ public void testEvaluatePolicyActorFilterUserResourceOwnersTypeNoMatch() throws when(_entityClient.getV2(eq(resourceUrn.getEntityType()), eq(resourceUrn), eq(Collections.singleton(Constants.OWNERSHIP_ASPECT_NAME)), any())).thenReturn(entityResponse); - ResolvedResourceSpec resourceSpec = - buildResourceResolvers("dataset", RESOURCE_URN, ImmutableSet.of(AUTHORIZED_PRINCIPAL), Collections.emptySet()); + ResolvedEntitySpec resourceSpec = + buildEntityResolvers("dataset", RESOURCE_URN, ImmutableSet.of(AUTHORIZED_PRINCIPAL), Collections.emptySet(), Collections.emptySet()); PolicyEngine.PolicyEvaluationResult result1 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertFalse(result1.isGranted()); + + // Verify no network calls + verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any()); } @Test @@ -639,17 +641,17 @@ public void testEvaluatePolicyActorFilterGroupResourceOwnersMatch() throws Excep when(_entityClient.getV2(eq(resourceUrn.getEntityType()), eq(resourceUrn), eq(Collections.singleton(Constants.OWNERSHIP_ASPECT_NAME)), any())).thenReturn(entityResponse); - ResolvedResourceSpec resourceSpec = - buildResourceResolvers("dataset", RESOURCE_URN, ImmutableSet.of(AUTHORIZED_GROUP), Collections.emptySet()); + ResolvedEntitySpec resourceSpec = + buildEntityResolvers("dataset", RESOURCE_URN, ImmutableSet.of(AUTHORIZED_GROUP), Collections.emptySet(), + Collections.emptySet()); // Assert authorized user can edit entity tags, because he is a user owner. PolicyEngine.PolicyEvaluationResult result1 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result1.isGranted()); - // Ensure that caching of groups is working with 1 call to entity client for each principal. - verify(_entityClient, times(1)).batchGetV2(eq(CORP_USER_ENTITY_NAME), eq(Collections.singleton(authorizedUserUrn)), - any(), any()); + // Verify no network calls + verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any()); } @Test @@ -673,16 +675,15 @@ public void testEvaluatePolicyActorFilterGroupResourceOwnersNoMatch() throws Exc resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); // Assert unauthorized user cannot edit entity tags. PolicyEngine.PolicyEvaluationResult result2 = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, UNAUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedUnauthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertFalse(result2.isGranted()); - // Ensure that caching of groups is working with 1 call to entity client for each principal. - verify(_entityClient, times(1)).batchGetV2(eq(CORP_USER_ENTITY_NAME), - eq(Collections.singleton(unauthorizedUserUrn)), any(), any()); + // Verify no network calls + verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any()); } @Test @@ -706,10 +707,10 @@ public void testEvaluatePolicyResourceFilterAllResourcesMatch() throws Exception resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = - buildResourceResolvers("dataset", "urn:li:dataset:random"); // A dataset Authorized principal _does not own_. + ResolvedEntitySpec resourceSpec = + buildEntityResolvers("dataset", "urn:li:dataset:random"); // A dataset Authorized principal _does not own_. PolicyEngine.PolicyEvaluationResult result = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result.isGranted()); @@ -738,9 +739,9 @@ public void testEvaluatePolicyResourceFilterAllResourcesNoMatch() throws Excepti resourceFilter.setType("dataset"); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("chart", RESOURCE_URN); // Notice: Not a dataset. + ResolvedEntitySpec resourceSpec = buildEntityResolvers("chart", RESOURCE_URN); // Notice: Not a dataset. PolicyEngine.PolicyEvaluationResult result = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertFalse(result.isGranted()); @@ -773,9 +774,9 @@ public void testEvaluatePolicyResourceFilterSpecificResourceMatchLegacy() throws resourceFilter.setResources(resourceUrns); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); PolicyEngine.PolicyEvaluationResult result = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result.isGranted()); @@ -801,13 +802,13 @@ public void testEvaluatePolicyResourceFilterSpecificResourceMatch() throws Excep final DataHubResourceFilter resourceFilter = new DataHubResourceFilter(); resourceFilter.setFilter(FilterUtils.newFilter( - ImmutableMap.of(ResourceFieldType.RESOURCE_TYPE, Collections.singletonList("dataset"), - ResourceFieldType.RESOURCE_URN, Collections.singletonList(RESOURCE_URN)))); + ImmutableMap.of(EntityFieldType.TYPE, Collections.singletonList("dataset"), + EntityFieldType.URN, Collections.singletonList(RESOURCE_URN)))); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN); + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN); PolicyEngine.PolicyEvaluationResult result = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result.isGranted()); @@ -833,14 +834,14 @@ public void testEvaluatePolicyResourceFilterSpecificResourceNoMatch() throws Exc final DataHubResourceFilter resourceFilter = new DataHubResourceFilter(); resourceFilter.setFilter(FilterUtils.newFilter( - ImmutableMap.of(ResourceFieldType.RESOURCE_TYPE, Collections.singletonList("dataset"), - ResourceFieldType.RESOURCE_URN, Collections.singletonList(RESOURCE_URN)))); + ImmutableMap.of(EntityFieldType.TYPE, Collections.singletonList("dataset"), + EntityFieldType.URN, Collections.singletonList(RESOURCE_URN)))); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = - buildResourceResolvers("dataset", "urn:li:dataset:random"); // A resource not covered by the policy. + ResolvedEntitySpec resourceSpec = + buildEntityResolvers("dataset", "urn:li:dataset:random"); // A resource not covered by the policy. PolicyEngine.PolicyEvaluationResult result = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertFalse(result.isGranted()); @@ -866,14 +867,14 @@ public void testEvaluatePolicyResourceFilterSpecificResourceMatchDomain() throws final DataHubResourceFilter resourceFilter = new DataHubResourceFilter(); resourceFilter.setFilter(FilterUtils.newFilter( - ImmutableMap.of(ResourceFieldType.RESOURCE_TYPE, Collections.singletonList("dataset"), ResourceFieldType.DOMAIN, + ImmutableMap.of(EntityFieldType.TYPE, Collections.singletonList("dataset"), EntityFieldType.DOMAIN, Collections.singletonList(DOMAIN_URN)))); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = - buildResourceResolvers("dataset", RESOURCE_URN, Collections.emptySet(), Collections.singleton(DOMAIN_URN)); + ResolvedEntitySpec resourceSpec = + buildEntityResolvers("dataset", RESOURCE_URN, Collections.emptySet(), Collections.singleton(DOMAIN_URN), Collections.emptySet()); PolicyEngine.PolicyEvaluationResult result = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertTrue(result.isGranted()); @@ -899,14 +900,14 @@ public void testEvaluatePolicyResourceFilterSpecificResourceNoMatchDomain() thro final DataHubResourceFilter resourceFilter = new DataHubResourceFilter(); resourceFilter.setFilter(FilterUtils.newFilter( - ImmutableMap.of(ResourceFieldType.RESOURCE_TYPE, Collections.singletonList("dataset"), ResourceFieldType.DOMAIN, + ImmutableMap.of(EntityFieldType.TYPE, Collections.singletonList("dataset"), EntityFieldType.DOMAIN, Collections.singletonList(DOMAIN_URN)))); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN, Collections.emptySet(), - Collections.singleton("urn:li:domain:domain2")); // Domain doesn't match + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN, Collections.emptySet(), + Collections.singleton("urn:li:domain:domain2"), Collections.emptySet()); // Domain doesn't match PolicyEngine.PolicyEvaluationResult result = - _policyEngine.evaluatePolicy(dataHubPolicyInfo, AUTHORIZED_PRINCIPAL, "EDIT_ENTITY_TAGS", + _policyEngine.evaluatePolicy(dataHubPolicyInfo, resolvedAuthorizedUserSpec, "EDIT_ENTITY_TAGS", Optional.of(resourceSpec)); assertFalse(result.isGranted()); @@ -933,7 +934,7 @@ public void testGetGrantedPrivileges() throws Exception { final DataHubResourceFilter resourceFilter1 = new DataHubResourceFilter(); resourceFilter1.setFilter(FilterUtils.newFilter( - ImmutableMap.of(ResourceFieldType.RESOURCE_TYPE, Collections.singletonList("dataset"), ResourceFieldType.DOMAIN, + ImmutableMap.of(EntityFieldType.TYPE, Collections.singletonList("dataset"), EntityFieldType.DOMAIN, Collections.singletonList(DOMAIN_URN)))); dataHubPolicyInfo1.setResources(resourceFilter1); @@ -954,8 +955,8 @@ public void testGetGrantedPrivileges() throws Exception { final DataHubResourceFilter resourceFilter2 = new DataHubResourceFilter(); resourceFilter2.setFilter(FilterUtils.newFilter( - ImmutableMap.of(ResourceFieldType.RESOURCE_TYPE, Collections.singletonList("dataset"), - ResourceFieldType.RESOURCE_URN, Collections.singletonList(RESOURCE_URN)))); + ImmutableMap.of(EntityFieldType.TYPE, Collections.singletonList("dataset"), + EntityFieldType.URN, Collections.singletonList(RESOURCE_URN)))); dataHubPolicyInfo2.setResources(resourceFilter2); // Policy 3, match dataset type and owner (legacy resource filter) @@ -981,25 +982,25 @@ public void testGetGrantedPrivileges() throws Exception { final List policies = ImmutableList.of(dataHubPolicyInfo1, dataHubPolicyInfo2, dataHubPolicyInfo3); - assertEquals(_policyEngine.getGrantedPrivileges(policies, UrnUtils.getUrn(AUTHORIZED_PRINCIPAL), Optional.empty()), + assertEquals(_policyEngine.getGrantedPrivileges(policies, resolvedAuthorizedUserSpec, Optional.empty()), Collections.emptyList()); - ResolvedResourceSpec resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN, Collections.emptySet(), - Collections.singleton(DOMAIN_URN)); // Everything matches + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN, Collections.emptySet(), + Collections.singleton(DOMAIN_URN), Collections.emptySet()); // Everything matches assertEquals( - _policyEngine.getGrantedPrivileges(policies, UrnUtils.getUrn(AUTHORIZED_PRINCIPAL), Optional.of(resourceSpec)), + _policyEngine.getGrantedPrivileges(policies, resolvedAuthorizedUserSpec, Optional.of(resourceSpec)), ImmutableList.of("PRIVILEGE_1", "PRIVILEGE_2_1", "PRIVILEGE_2_2")); - resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN, Collections.emptySet(), - Collections.singleton("urn:li:domain:domain2")); // Domain doesn't match + resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN, Collections.emptySet(), + Collections.singleton("urn:li:domain:domain2"), Collections.emptySet()); // Domain doesn't match assertEquals( - _policyEngine.getGrantedPrivileges(policies, UrnUtils.getUrn(AUTHORIZED_PRINCIPAL), Optional.of(resourceSpec)), + _policyEngine.getGrantedPrivileges(policies, resolvedAuthorizedUserSpec, Optional.of(resourceSpec)), ImmutableList.of("PRIVILEGE_2_1", "PRIVILEGE_2_2")); - resourceSpec = buildResourceResolvers("dataset", "urn:li:dataset:random", Collections.emptySet(), - Collections.singleton(DOMAIN_URN)); // Resource doesn't match + resourceSpec = buildEntityResolvers("dataset", "urn:li:dataset:random", Collections.emptySet(), + Collections.singleton(DOMAIN_URN), Collections.emptySet()); // Resource doesn't match assertEquals( - _policyEngine.getGrantedPrivileges(policies, UrnUtils.getUrn(AUTHORIZED_PRINCIPAL), Optional.of(resourceSpec)), + _policyEngine.getGrantedPrivileges(policies, resolvedAuthorizedUserSpec, Optional.of(resourceSpec)), ImmutableList.of("PRIVILEGE_1")); final EntityResponse entityResponse = new EntityResponse(); @@ -1008,16 +1009,16 @@ public void testGetGrantedPrivileges() throws Exception { entityResponse.setAspects(aspectMap); when(_entityClient.getV2(eq(resourceUrn.getEntityType()), eq(resourceUrn), eq(Collections.singleton(Constants.OWNERSHIP_ASPECT_NAME)), any())).thenReturn(entityResponse); - resourceSpec = buildResourceResolvers("dataset", RESOURCE_URN, Collections.singleton(AUTHORIZED_PRINCIPAL), - Collections.singleton(DOMAIN_URN)); // Is owner + resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN, Collections.singleton(AUTHORIZED_PRINCIPAL), + Collections.singleton(DOMAIN_URN), Collections.emptySet()); // Is owner assertEquals( - _policyEngine.getGrantedPrivileges(policies, UrnUtils.getUrn(AUTHORIZED_PRINCIPAL), Optional.of(resourceSpec)), + _policyEngine.getGrantedPrivileges(policies, resolvedAuthorizedUserSpec, Optional.of(resourceSpec)), ImmutableList.of("PRIVILEGE_1", "PRIVILEGE_2_1", "PRIVILEGE_2_2", "PRIVILEGE_3")); - resourceSpec = buildResourceResolvers("chart", RESOURCE_URN, Collections.singleton(AUTHORIZED_PRINCIPAL), - Collections.singleton(DOMAIN_URN)); // Resource type doesn't match + resourceSpec = buildEntityResolvers("chart", RESOURCE_URN, Collections.singleton(AUTHORIZED_PRINCIPAL), + Collections.singleton(DOMAIN_URN), Collections.emptySet()); // Resource type doesn't match assertEquals( - _policyEngine.getGrantedPrivileges(policies, UrnUtils.getUrn(AUTHORIZED_PRINCIPAL), Optional.of(resourceSpec)), + _policyEngine.getGrantedPrivileges(policies, resolvedAuthorizedUserSpec, Optional.of(resourceSpec)), Collections.emptyList()); } @@ -1050,9 +1051,9 @@ public void testGetMatchingActorsResourceMatch() throws Exception { resourceFilter.setResources(resourceUrns); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = - buildResourceResolvers("dataset", RESOURCE_URN, ImmutableSet.of(AUTHORIZED_PRINCIPAL, AUTHORIZED_GROUP), - Collections.emptySet()); + ResolvedEntitySpec resourceSpec = + buildEntityResolvers("dataset", RESOURCE_URN, ImmutableSet.of(AUTHORIZED_PRINCIPAL, AUTHORIZED_GROUP), + Collections.emptySet(), Collections.emptySet()); PolicyEngine.PolicyActors actors = _policyEngine.getMatchingActors(dataHubPolicyInfo, Optional.of(resourceSpec)); assertTrue(actors.allUsers()); @@ -1101,8 +1102,8 @@ public void testGetMatchingActorsNoResourceMatch() throws Exception { resourceFilter.setResources(resourceUrns); dataHubPolicyInfo.setResources(resourceFilter); - ResolvedResourceSpec resourceSpec = - buildResourceResolvers("dataset", "urn:li:dataset:random"); // A resource not covered by the policy. + ResolvedEntitySpec resourceSpec = + buildEntityResolvers("dataset", "urn:li:dataset:random"); // A resource not covered by the policy. PolicyEngine.PolicyActors actors = _policyEngine.getMatchingActors(dataHubPolicyInfo, Optional.of(resourceSpec)); assertFalse(actors.allUsers()); @@ -1155,21 +1156,6 @@ private EntityResponse createAuthorizedEntityResponse() throws URISyntaxExceptio final EntityResponse entityResponse = new EntityResponse(); final EnvelopedAspectMap aspectMap = new EnvelopedAspectMap(); - final CorpUserInfo userInfo = new CorpUserInfo(); - userInfo.setActive(true); - userInfo.setFullName("Data Hub"); - userInfo.setFirstName("Data"); - userInfo.setLastName("Hub"); - userInfo.setEmail("datahub@gmail.com"); - userInfo.setTitle("Admin"); - aspectMap.put(CORP_USER_INFO_ASPECT_NAME, new EnvelopedAspect().setValue(new Aspect(userInfo.data()))); - - final GroupMembership groupsAspect = new GroupMembership(); - final UrnArray groups = new UrnArray(); - groups.add(Urn.createFromString("urn:li:corpGroup:authorizedGroup")); - groupsAspect.setGroups(groups); - aspectMap.put(GROUP_MEMBERSHIP_ASPECT_NAME, new EnvelopedAspect().setValue(new Aspect(groupsAspect.data()))); - final RoleMembership rolesAspect = new RoleMembership(); final UrnArray roles = new UrnArray(); roles.add(Urn.createFromString("urn:li:dataHubRole:admin")); @@ -1184,21 +1170,6 @@ private EntityResponse createUnauthorizedEntityResponse() throws URISyntaxExcept final EntityResponse entityResponse = new EntityResponse(); final EnvelopedAspectMap aspectMap = new EnvelopedAspectMap(); - final CorpUserInfo userInfo = new CorpUserInfo(); - userInfo.setActive(true); - userInfo.setFullName("Unauthorized User"); - userInfo.setFirstName("Unauthorized"); - userInfo.setLastName("User"); - userInfo.setEmail("Unauth"); - userInfo.setTitle("Engineer"); - aspectMap.put(CORP_USER_INFO_ASPECT_NAME, new EnvelopedAspect().setValue(new Aspect(userInfo.data()))); - - final GroupMembership groupsAspect = new GroupMembership(); - final UrnArray groups = new UrnArray(); - groups.add(Urn.createFromString("urn:li:corpGroup:unauthorizedGroup")); - groupsAspect.setGroups(groups); - aspectMap.put(GROUP_MEMBERSHIP_ASPECT_NAME, new EnvelopedAspect().setValue(new Aspect(groupsAspect.data()))); - final RoleMembership rolesAspect = new RoleMembership(); final UrnArray roles = new UrnArray(); roles.add(Urn.createFromString("urn:li:dataHubRole:reader")); @@ -1209,17 +1180,18 @@ private EntityResponse createUnauthorizedEntityResponse() throws URISyntaxExcept return entityResponse; } - public static ResolvedResourceSpec buildResourceResolvers(String entityType, String entityUrn) { - return buildResourceResolvers(entityType, entityUrn, Collections.emptySet(), Collections.emptySet()); + public static ResolvedEntitySpec buildEntityResolvers(String entityType, String entityUrn) { + return buildEntityResolvers(entityType, entityUrn, Collections.emptySet(), Collections.emptySet(), Collections.emptySet()); } - public static ResolvedResourceSpec buildResourceResolvers(String entityType, String entityUrn, Set owners, - Set domains) { - return new ResolvedResourceSpec(new ResourceSpec(entityType, entityUrn), - ImmutableMap.of(ResourceFieldType.RESOURCE_TYPE, - FieldResolver.getResolverFromValues(Collections.singleton(entityType)), ResourceFieldType.RESOURCE_URN, - FieldResolver.getResolverFromValues(Collections.singleton(entityUrn)), ResourceFieldType.OWNER, - FieldResolver.getResolverFromValues(owners), ResourceFieldType.DOMAIN, - FieldResolver.getResolverFromValues(domains))); + public static ResolvedEntitySpec buildEntityResolvers(String entityType, String entityUrn, Set owners, + Set domains, Set groups) { + return new ResolvedEntitySpec(new EntitySpec(entityType, entityUrn), + ImmutableMap.of(EntityFieldType.TYPE, + FieldResolver.getResolverFromValues(Collections.singleton(entityType)), EntityFieldType.URN, + FieldResolver.getResolverFromValues(Collections.singleton(entityUrn)), EntityFieldType.OWNER, + FieldResolver.getResolverFromValues(owners), EntityFieldType.DOMAIN, + FieldResolver.getResolverFromValues(domains), EntityFieldType.GROUP_MEMBERSHIP, + FieldResolver.getResolverFromValues(groups))); } } diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProviderTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProviderTest.java index e525c602c2620..b2343bbb01509 100644 --- a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProviderTest.java +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProviderTest.java @@ -1,8 +1,21 @@ package com.datahub.authorization.fieldresolverprovider; +import static com.linkedin.metadata.Constants.DATASET_ENTITY_NAME; +import static com.linkedin.metadata.Constants.DATA_PLATFORM_INSTANCE_ASPECT_NAME; +import static com.linkedin.metadata.Constants.DATA_PLATFORM_INSTANCE_ENTITY_NAME; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyZeroInteractions; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + import com.datahub.authentication.Authentication; -import com.datahub.authorization.ResourceFieldType; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntityFieldType; +import com.datahub.authorization.EntitySpec; import com.linkedin.common.DataPlatformInstance; import com.linkedin.common.urn.Urn; import com.linkedin.entity.Aspect; @@ -11,29 +24,21 @@ import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.entity.client.EntityClient; import com.linkedin.r2.RemoteInvocationException; +import java.net.URISyntaxException; +import java.util.Collections; +import java.util.Set; import org.mockito.Mock; import org.mockito.MockitoAnnotations; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; -import java.net.URISyntaxException; -import java.util.Collections; -import java.util.Set; - -import static com.linkedin.metadata.Constants.*; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.*; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertTrue; - public class DataPlatformInstanceFieldResolverProviderTest { private static final String DATA_PLATFORM_INSTANCE_URN = "urn:li:dataPlatformInstance:(urn:li:dataPlatform:s3,test-platform-instance)"; private static final String RESOURCE_URN = "urn:li:dataset:(urn:li:dataPlatform:s3,test-platform-instance.testDataset,PROD)"; - private static final ResourceSpec RESOURCE_SPEC = new ResourceSpec(DATASET_ENTITY_NAME, RESOURCE_URN); + private static final EntitySpec RESOURCE_SPEC = new EntitySpec(DATASET_ENTITY_NAME, RESOURCE_URN); @Mock private EntityClient entityClientMock; @@ -51,12 +56,12 @@ public void setup() { @Test public void shouldReturnDataPlatformInstanceType() { - assertEquals(ResourceFieldType.DATA_PLATFORM_INSTANCE, dataPlatformInstanceFieldResolverProvider.getFieldType()); + assertEquals(EntityFieldType.DATA_PLATFORM_INSTANCE, dataPlatformInstanceFieldResolverProvider.getFieldType()); } @Test public void shouldReturnFieldValueWithResourceSpecIfTypeIsDataPlatformInstance() { - var resourceSpec = new ResourceSpec(DATA_PLATFORM_INSTANCE_ENTITY_NAME, DATA_PLATFORM_INSTANCE_URN); + var resourceSpec = new EntitySpec(DATA_PLATFORM_INSTANCE_ENTITY_NAME, DATA_PLATFORM_INSTANCE_URN); var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(resourceSpec); diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/GroupMembershipFieldResolverProviderTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/GroupMembershipFieldResolverProviderTest.java new file mode 100644 index 0000000000000..54675045b4413 --- /dev/null +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/GroupMembershipFieldResolverProviderTest.java @@ -0,0 +1,212 @@ +package com.datahub.authorization.fieldresolverprovider; + +import com.datahub.authentication.Authentication; +import com.datahub.authorization.EntityFieldType; +import com.datahub.authorization.EntitySpec; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.linkedin.common.UrnArray; +import com.linkedin.common.urn.Urn; +import com.linkedin.entity.Aspect; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.entity.EnvelopedAspectMap; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.identity.GroupMembership; +import com.linkedin.identity.NativeGroupMembership; +import com.linkedin.r2.RemoteInvocationException; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import java.net.URISyntaxException; +import java.util.Set; + +import static com.linkedin.metadata.Constants.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.*; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +public class GroupMembershipFieldResolverProviderTest { + + private static final String CORPGROUP_URN = "urn:li:corpGroup:groupname"; + private static final String NATIVE_CORPGROUP_URN = "urn:li:corpGroup:nativegroupname"; + private static final String RESOURCE_URN = "urn:li:dataset:(urn:li:dataPlatform:testPlatform,testDataset,PROD)"; + private static final EntitySpec RESOURCE_SPEC = new EntitySpec(DATASET_ENTITY_NAME, RESOURCE_URN); + + @Mock + private EntityClient entityClientMock; + @Mock + private Authentication systemAuthenticationMock; + + private GroupMembershipFieldResolverProvider groupMembershipFieldResolverProvider; + + @BeforeMethod + public void setup() { + MockitoAnnotations.initMocks(this); + groupMembershipFieldResolverProvider = + new GroupMembershipFieldResolverProvider(entityClientMock, systemAuthenticationMock); + } + + @Test + public void shouldReturnGroupsMembershipType() { + assertEquals(EntityFieldType.GROUP_MEMBERSHIP, groupMembershipFieldResolverProvider.getFieldType()); + } + + @Test + public void shouldReturnEmptyFieldValueWhenResponseIsNull() throws RemoteInvocationException, URISyntaxException { + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenReturn(null); + + var result = groupMembershipFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } + + @Test + public void shouldReturnEmptyFieldValueWhenResourceDoesNotBelongToAnyGroup() + throws RemoteInvocationException, URISyntaxException { + var entityResponseMock = mock(EntityResponse.class); + when(entityResponseMock.getAspects()).thenReturn(new EnvelopedAspectMap()); + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenReturn(entityResponseMock); + + var result = groupMembershipFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } + + @Test + public void shouldReturnEmptyFieldValueWhenThereIsAnException() throws RemoteInvocationException, URISyntaxException { + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenThrow(new RemoteInvocationException()); + + var result = groupMembershipFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } + + @Test + public void shouldReturnFieldValueWithOnlyGroupsOfTheResource() + throws RemoteInvocationException, URISyntaxException { + + var groupMembership = new GroupMembership().setGroups( + new UrnArray(ImmutableList.of(Urn.createFromString(CORPGROUP_URN)))); + var entityResponseMock = mock(EntityResponse.class); + var envelopedAspectMap = new EnvelopedAspectMap(); + envelopedAspectMap.put(GROUP_MEMBERSHIP_ASPECT_NAME, + new EnvelopedAspect().setValue(new Aspect(groupMembership.data()))); + when(entityResponseMock.getAspects()).thenReturn(envelopedAspectMap); + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenReturn(entityResponseMock); + + var result = groupMembershipFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertEquals(Set.of(CORPGROUP_URN), result.getFieldValuesFuture().join().getValues()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } + + @Test + public void shouldReturnFieldValueWithOnlyNativeGroupsOfTheResource() + throws RemoteInvocationException, URISyntaxException { + + var nativeGroupMembership = new NativeGroupMembership().setNativeGroups( + new UrnArray(ImmutableList.of(Urn.createFromString(NATIVE_CORPGROUP_URN)))); + var entityResponseMock = mock(EntityResponse.class); + var envelopedAspectMap = new EnvelopedAspectMap(); + envelopedAspectMap.put(NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME, + new EnvelopedAspect().setValue(new Aspect(nativeGroupMembership.data()))); + when(entityResponseMock.getAspects()).thenReturn(envelopedAspectMap); + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenReturn(entityResponseMock); + + var result = groupMembershipFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertEquals(Set.of(NATIVE_CORPGROUP_URN), result.getFieldValuesFuture().join().getValues()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } + + @Test + public void shouldReturnFieldValueWithGroupsAndNativeGroupsOfTheResource() + throws RemoteInvocationException, URISyntaxException { + + var groupMembership = new GroupMembership().setGroups( + new UrnArray(ImmutableList.of(Urn.createFromString(CORPGROUP_URN)))); + var nativeGroupMembership = new NativeGroupMembership().setNativeGroups( + new UrnArray(ImmutableList.of(Urn.createFromString(NATIVE_CORPGROUP_URN)))); + var entityResponseMock = mock(EntityResponse.class); + var envelopedAspectMap = new EnvelopedAspectMap(); + envelopedAspectMap.put(GROUP_MEMBERSHIP_ASPECT_NAME, + new EnvelopedAspect().setValue(new Aspect(groupMembership.data()))); + envelopedAspectMap.put(NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME, + new EnvelopedAspect().setValue(new Aspect(nativeGroupMembership.data()))); + when(entityResponseMock.getAspects()).thenReturn(envelopedAspectMap); + when(entityClientMock.getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + )).thenReturn(entityResponseMock); + + var result = groupMembershipFieldResolverProvider.getFieldResolver(RESOURCE_SPEC); + + assertEquals(Set.of(CORPGROUP_URN, NATIVE_CORPGROUP_URN), result.getFieldValuesFuture().join().getValues()); + verify(entityClientMock, times(1)).getV2( + eq(DATASET_ENTITY_NAME), + any(Urn.class), + eq(ImmutableSet.of(GROUP_MEMBERSHIP_ASPECT_NAME, NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)), + eq(systemAuthenticationMock) + ); + } +} \ No newline at end of file diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/AuthorizerChainFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/AuthorizerChainFactory.java index bf50a0c7b6473..b90257870a8b2 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/AuthorizerChainFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/AuthorizerChainFactory.java @@ -2,12 +2,12 @@ import com.datahub.authorization.AuthorizerChain; import com.datahub.authorization.DataHubAuthorizer; -import com.datahub.authorization.DefaultResourceSpecResolver; +import com.datahub.authorization.DefaultEntitySpecResolver; import com.datahub.plugins.PluginConstant; import com.datahub.authentication.Authentication; import com.datahub.plugins.auth.authorization.Authorizer; import com.datahub.authorization.AuthorizerContext; -import com.datahub.authorization.ResourceSpecResolver; +import com.datahub.authorization.EntitySpecResolver; import com.datahub.plugins.common.PluginConfig; import com.datahub.plugins.common.PluginPermissionManager; import com.datahub.plugins.common.PluginType; @@ -64,7 +64,7 @@ public class AuthorizerChainFactory { @Scope("singleton") @Nonnull protected AuthorizerChain getInstance() { - final ResourceSpecResolver resolver = initResolver(); + final EntitySpecResolver resolver = initResolver(); // Extract + initialize customer authorizers from application configs. final List authorizers = new ArrayList<>(initCustomAuthorizers(resolver)); @@ -79,11 +79,11 @@ protected AuthorizerChain getInstance() { return new AuthorizerChain(authorizers, dataHubAuthorizer); } - private ResourceSpecResolver initResolver() { - return new DefaultResourceSpecResolver(systemAuthentication, entityClient); + private EntitySpecResolver initResolver() { + return new DefaultEntitySpecResolver(systemAuthentication, entityClient); } - private List initCustomAuthorizers(ResourceSpecResolver resolver) { + private List initCustomAuthorizers(EntitySpecResolver resolver) { final List customAuthorizers = new ArrayList<>(); Path pluginBaseDirectory = Paths.get(configurationProvider.getDatahub().getPlugin().getAuth().getPath()); @@ -99,7 +99,7 @@ private List initCustomAuthorizers(ResourceSpecResolver resolver) { return customAuthorizers; } - private void registerAuthorizer(List customAuthorizers, ResourceSpecResolver resolver, Config config) { + private void registerAuthorizer(List customAuthorizers, EntitySpecResolver resolver, Config config) { PluginConfigFactory authorizerPluginPluginConfigFactory = new PluginConfigFactory(config); // Load only Authorizer configuration from plugin config factory List authorizers = diff --git a/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/delegates/EntityApiDelegateImpl.java b/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/delegates/EntityApiDelegateImpl.java index ade49c876f168..207c2284e2673 100644 --- a/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/delegates/EntityApiDelegateImpl.java +++ b/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/delegates/EntityApiDelegateImpl.java @@ -45,8 +45,7 @@ import io.datahubproject.openapi.util.OpenApiEntitiesUtil; import com.datahub.authorization.ConjunctivePrivilegeGroup; import com.datahub.authorization.DisjunctivePrivilegeGroup; -import com.linkedin.metadata.models.EntitySpec; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.linkedin.metadata.authorization.PoliciesConfig; import com.google.common.collect.ImmutableList; import com.datahub.authorization.AuthUtil; @@ -377,7 +376,7 @@ public ResponseEntity scroll(@Valid Boolean systemMetadata, @Valid List sort, @Valid SortOrder sortOrder, @Valid String query) { Authentication authentication = AuthenticationContext.getAuthentication(); - EntitySpec entitySpec = OpenApiEntitiesUtil.responseClassToEntitySpec(_entityRegistry, _respClazz); + com.linkedin.metadata.models.EntitySpec entitySpec = OpenApiEntitiesUtil.responseClassToEntitySpec(_entityRegistry, _respClazz); checkScrollAuthorized(authentication, entitySpec); // TODO multi-field sort @@ -410,12 +409,12 @@ public ResponseEntity scroll(@Valid Boolean systemMetadata, @Valid List> resourceSpecs = List.of(Optional.of(new ResourceSpec(entitySpec.getName(), ""))); + List> resourceSpecs = List.of(Optional.of(new EntitySpec(entitySpec.getName(), ""))); if (_restApiAuthorizationEnabled && !AuthUtil.isAuthorizedForResources(_authorizationChain, actorUrnStr, resourceSpecs, orGroup)) { throw new UnauthorizedException(actorUrnStr + " is unauthorized to get entities."); } diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/entities/EntitiesController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/entities/EntitiesController.java index 6439e2f31f7b0..898f768cf999a 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/entities/EntitiesController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/entities/EntitiesController.java @@ -8,7 +8,7 @@ import com.datahub.authorization.AuthorizerChain; import com.datahub.authorization.ConjunctivePrivilegeGroup; import com.datahub.authorization.DisjunctivePrivilegeGroup; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; @@ -93,8 +93,8 @@ public ResponseEntity getEntities( ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE.getType()) ))); - List> resourceSpecs = entityUrns.stream() - .map(urn -> Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))) + List> resourceSpecs = entityUrns.stream() + .map(urn -> Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))) .collect(Collectors.toList()); if (restApiAuthorizationEnabled && !AuthUtil.isAuthorizedForResources(_authorizerChain, actorUrnStr, resourceSpecs, orGroup)) { throw new UnauthorizedException(actorUrnStr + " is unauthorized to get entities."); @@ -175,8 +175,8 @@ public ResponseEntity> deleteEntities( .map(URLDecoder::decode) .map(UrnUtils::getUrn).collect(Collectors.toSet()); - List> resourceSpecs = entityUrns.stream() - .map(urn -> Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))) + List> resourceSpecs = entityUrns.stream() + .map(urn -> Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))) .collect(Collectors.toList()); if (restApiAuthorizationEnabled && !AuthUtil.isAuthorizedForResources(_authorizerChain, actorUrnStr, resourceSpecs, orGroup)) { UnauthorizedException unauthorizedException = new UnauthorizedException(actorUrnStr + " is unauthorized to delete entities."); diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/relationships/RelationshipsController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/relationships/RelationshipsController.java index 1e37170f37b3b..4641fed3a8610 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/relationships/RelationshipsController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/relationships/RelationshipsController.java @@ -8,7 +8,7 @@ import com.datahub.authorization.AuthorizerChain; import com.datahub.authorization.ConjunctivePrivilegeGroup; import com.datahub.authorization.DisjunctivePrivilegeGroup; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; @@ -131,8 +131,8 @@ public ResponseEntity getRelationships( // Re-using GET_ENTITY_PRIVILEGE here as it doesn't make sense to split the privileges between these APIs. ))); - List> resourceSpecs = - Collections.singletonList(Optional.of(new ResourceSpec(entityUrn.getEntityType(), entityUrn.toString()))); + List> resourceSpecs = + Collections.singletonList(Optional.of(new EntitySpec(entityUrn.getEntityType(), entityUrn.toString()))); if (restApiAuthorizationEnabled && !AuthUtil.isAuthorizedForResources(_authorizerChain, actorUrnStr, resourceSpecs, orGroup)) { throw new UnauthorizedException(actorUrnStr + " is unauthorized to get relationships."); diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/timeline/TimelineController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/timeline/TimelineController.java index 5a0ce2e314e1b..fbde9e8072002 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/timeline/TimelineController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/timeline/TimelineController.java @@ -6,7 +6,7 @@ import com.datahub.authorization.AuthorizerChain; import com.datahub.authorization.ConjunctivePrivilegeGroup; import com.datahub.authorization.DisjunctivePrivilegeGroup; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.fasterxml.jackson.core.JsonProcessingException; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; @@ -67,7 +67,7 @@ public ResponseEntity> getTimeline( Urn urn = Urn.createFromString(rawUrn); Authentication authentication = AuthenticationContext.getAuthentication(); String actorUrnStr = authentication.getActor().toUrnStr(); - ResourceSpec resourceSpec = new ResourceSpec(urn.getEntityType(), rawUrn); + EntitySpec resourceSpec = new EntitySpec(urn.getEntityType(), rawUrn); DisjunctivePrivilegeGroup orGroup = new DisjunctivePrivilegeGroup( ImmutableList.of(new ConjunctivePrivilegeGroup(ImmutableList.of(PoliciesConfig.GET_TIMELINE_PRIVILEGE.getType())))); if (restApiAuthorizationEnabled && !AuthUtil.isAuthorized(_authorizerChain, actorUrnStr, Optional.of(resourceSpec), orGroup)) { diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/util/MappingUtil.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/util/MappingUtil.java index 2b3e84e2df20f..21dc5a4c8a0d6 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/util/MappingUtil.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/util/MappingUtil.java @@ -5,7 +5,7 @@ import com.datahub.authorization.AuthUtil; import com.datahub.plugins.auth.authorization.Authorizer; import com.datahub.authorization.DisjunctivePrivilegeGroup; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; @@ -27,7 +27,6 @@ import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.entity.transactions.AspectsBatch; import com.linkedin.metadata.entity.validation.ValidationException; -import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.entity.AspectUtils; import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.metrics.MetricUtils; @@ -378,11 +377,11 @@ public static GenericAspect convertGenericAspect(@Nonnull io.datahubproject.open public static boolean authorizeProposals(List proposals, EntityService entityService, Authorizer authorizer, String actorUrnStr, DisjunctivePrivilegeGroup orGroup) { - List> resourceSpecs = proposals.stream() + List> resourceSpecs = proposals.stream() .map(proposal -> { - EntitySpec entitySpec = entityService.getEntityRegistry().getEntitySpec(proposal.getEntityType()); + com.linkedin.metadata.models.EntitySpec entitySpec = entityService.getEntityRegistry().getEntitySpec(proposal.getEntityType()); Urn entityUrn = EntityKeyUtils.getUrnFromProposal(proposal, entitySpec.getKeyAspectSpec()); - return Optional.of(new ResourceSpec(proposal.getEntityType(), entityUrn.toString())); + return Optional.of(new EntitySpec(proposal.getEntityType(), entityUrn.toString())); }) .collect(Collectors.toList()); return AuthUtil.isAuthorizedForResources(authorizer, actorUrnStr, resourceSpecs, orGroup); @@ -513,7 +512,7 @@ public static RollbackRunResultDto mapRollbackRunResult(RollbackRunResult rollba } public static UpsertAspectRequest createStatusRemoval(Urn urn, EntityService entityService) { - EntitySpec entitySpec = entityService.getEntityRegistry().getEntitySpec(urn.getEntityType()); + com.linkedin.metadata.models.EntitySpec entitySpec = entityService.getEntityRegistry().getEntitySpec(urn.getEntityType()); if (entitySpec == null || !entitySpec.getAspectSpecMap().containsKey(STATUS_ASPECT_NAME)) { throw new IllegalArgumentException("Entity type is not valid for soft deletes: " + urn.getEntityType()); } diff --git a/metadata-service/plugin/src/test/sample-test-plugins/src/main/java/com/datahub/plugins/test/TestAuthorizer.java b/metadata-service/plugin/src/test/sample-test-plugins/src/main/java/com/datahub/plugins/test/TestAuthorizer.java index b6bc282f10b65..442ac1b0d287b 100644 --- a/metadata-service/plugin/src/test/sample-test-plugins/src/main/java/com/datahub/plugins/test/TestAuthorizer.java +++ b/metadata-service/plugin/src/test/sample-test-plugins/src/main/java/com/datahub/plugins/test/TestAuthorizer.java @@ -4,7 +4,7 @@ import com.datahub.authorization.AuthorizationResult; import com.datahub.authorization.AuthorizedActors; import com.datahub.authorization.AuthorizerContext; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.datahub.plugins.PluginConstant; import com.datahub.plugins.auth.authorization.Authorizer; import java.io.BufferedReader; @@ -74,7 +74,7 @@ public AuthorizationResult authorize(@Nonnull AuthorizationRequest request) { } @Override - public AuthorizedActors authorizedActors(String privilege, Optional resourceSpec) { + public AuthorizedActors authorizedActors(String privilege, Optional resourceSpec) { return new AuthorizedActors("ALL", null, null, true, true); } } diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java index 936c8bb67e645..af76af90ce77f 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java @@ -3,7 +3,7 @@ import com.codahale.metrics.MetricRegistry; import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; @@ -20,7 +20,6 @@ import com.linkedin.metadata.entity.AspectUtils; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.validation.ValidationException; -import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.restli.RestliUtil; @@ -123,7 +122,7 @@ public Task get(@Nonnull String urnStr, @QueryParam("aspect") @Option Authentication authentication = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(authentication, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), - new ResourceSpec(urn.getEntityType(), urn.toString()))) { + new EntitySpec(urn.getEntityType(), urn.toString()))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get aspect for " + urn); } final VersionedAspect aspect = _entityService.getVersionedAspect(urn, aspectName, version); @@ -154,7 +153,7 @@ public Task getTimeseriesAspectValues( Authentication authentication = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(authentication, _authorizer, ImmutableList.of(PoliciesConfig.GET_TIMESERIES_ASPECT_PRIVILEGE), - new ResourceSpec(urn.getEntityType(), urn.toString()))) { + new EntitySpec(urn.getEntityType(), urn.toString()))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get timeseries aspect for " + urn); } GetTimeseriesAspectValuesResponse response = new GetTimeseriesAspectValuesResponse(); @@ -193,11 +192,11 @@ public Task ingestProposal( } Authentication authentication = AuthenticationContext.getAuthentication(); - EntitySpec entitySpec = _entityService.getEntityRegistry().getEntitySpec(metadataChangeProposal.getEntityType()); + com.linkedin.metadata.models.EntitySpec entitySpec = _entityService.getEntityRegistry().getEntitySpec(metadataChangeProposal.getEntityType()); Urn urn = EntityKeyUtils.getUrnFromProposal(metadataChangeProposal, entitySpec.getKeyAspectSpec()); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(authentication, _authorizer, ImmutableList.of(PoliciesConfig.EDIT_ENTITY_PRIVILEGE), - new ResourceSpec(urn.getEntityType(), urn.toString()))) { + new EntitySpec(urn.getEntityType(), urn.toString()))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to modify entity " + urn); } String actorUrnStr = authentication.getActor().toUrnStr(); @@ -249,7 +248,7 @@ public Task getCount(@ActionParam(PARAM_ASPECT) @Nonnull String aspectN Authentication authentication = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(authentication, _authorizer, ImmutableList.of(PoliciesConfig.GET_COUNTS_PRIVILEGE), - (ResourceSpec) null)) { + (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get aspect counts."); } return _entityService.getCountAspect(aspectName, urnLike); diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java index 3ff22fb767676..9bab846d1bdcc 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java @@ -4,7 +4,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; import com.datahub.plugins.auth.authorization.Authorizer; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.google.common.collect.ImmutableList; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; @@ -123,9 +123,9 @@ public Task rollback(@ActionParam("runId") @Nonnull String run List aspectRowsToDelete; aspectRowsToDelete = _systemMetadataService.findByRunId(runId, doHardDelete, 0, ESUtils.MAX_RESULT_SIZE); Set urns = aspectRowsToDelete.stream().collect(Collectors.groupingBy(AspectRowSummary::getUrn)).keySet(); - List> resourceSpecs = urns.stream() + List> resourceSpecs = urns.stream() .map(UrnUtils::getUrn) - .map(urn -> java.util.Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))) + .map(urn -> java.util.Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))) .collect(Collectors.toList()); Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java index f6dedfb9a07c6..3ee98b3244718 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java @@ -3,7 +3,7 @@ import com.codahale.metrics.MetricRegistry; import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.linkedin.common.AuditStamp; @@ -173,7 +173,7 @@ public Task get(@Nonnull String urnStr, final Urn urn = Urn.createFromString(urnStr); Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), new ResourceSpec(urn.getEntityType(), urnStr))) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), new EntitySpec(urn.getEntityType(), urnStr))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get entity " + urn); } @@ -198,8 +198,8 @@ public Task> batchGet(@Nonnull Set urnStrs, for (final String urnStr : urnStrs) { urns.add(Urn.createFromString(urnStr)); } - List> resourceSpecs = urns.stream() - .map(urn -> java.util.Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))) + List> resourceSpecs = urns.stream() + .map(urn -> java.util.Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))) .collect(Collectors.toList()); Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) @@ -242,7 +242,7 @@ public Task ingest(@ActionParam(PARAM_ENTITY) @Nonnull Entity entity, final Urn urn = com.datahub.util.ModelUtils.getUrnFromSnapshotUnion(entity.getValue()); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(authentication, _authorizer, ImmutableList.of(PoliciesConfig.EDIT_ENTITY_PRIVILEGE), - new ResourceSpec(urn.getEntityType(), urn.toString()))) { + new EntitySpec(urn.getEntityType(), urn.toString()))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to edit entity " + urn); } @@ -273,10 +273,10 @@ public Task batchIngest(@ActionParam(PARAM_ENTITIES) @Nonnull Entity[] ent Authentication authentication = AuthenticationContext.getAuthentication(); String actorUrnStr = authentication.getActor().toUrnStr(); - List> resourceSpecs = Arrays.stream(entities) + List> resourceSpecs = Arrays.stream(entities) .map(Entity::getValue) .map(com.datahub.util.ModelUtils::getUrnFromSnapshotUnion) - .map(urn -> java.util.Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))) + .map(urn -> java.util.Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))) .collect(Collectors.toList()); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(authentication, _authorizer, ImmutableList.of(PoliciesConfig.EDIT_ENTITY_PRIVILEGE), resourceSpecs)) { @@ -322,7 +322,7 @@ public Task search(@ActionParam(PARAM_ENTITY) @Nonnull String enti @Optional @Nullable @ActionParam(PARAM_SEARCH_FLAGS) SearchFlags searchFlags) { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } @@ -347,7 +347,7 @@ public Task searchAcrossEntities(@ActionParam(PARAM_ENTITIES) @Opt @ActionParam(PARAM_COUNT) int count, @ActionParam(PARAM_SEARCH_FLAGS) @Optional SearchFlags searchFlags) { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } @@ -391,7 +391,7 @@ public Task searchAcrossLineage(@ActionParam(PARAM_URN) @No @Optional @Nullable @ActionParam(PARAM_SEARCH_FLAGS) SearchFlags searchFlags) throws URISyntaxException { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } @@ -443,7 +443,7 @@ public Task list(@ActionParam(PARAM_ENTITY) @Nonnull String entityNa Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } @@ -462,7 +462,7 @@ public Task autocomplete(@ActionParam(PARAM_ENTITY) @Nonnull Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } @@ -479,7 +479,7 @@ public Task browse(@ActionParam(PARAM_ENTITY) @Nonnull String enti Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } @@ -497,7 +497,7 @@ public Task getBrowsePaths( Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), - new ResourceSpec(urn.getEntityType(), urn.toString()))) { + new EntitySpec(urn.getEntityType(), urn.toString()))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get entity: " + urn); } @@ -546,9 +546,9 @@ public Task deleteEntities(@ActionParam("registryId") @Optiona log.info("found {} rows to delete...", stringifyRowCount(aspectRowsToDelete.size())); response.setAspectsAffected(aspectRowsToDelete.size()); Set urns = aspectRowsToDelete.stream().collect(Collectors.groupingBy(AspectRowSummary::getUrn)).keySet(); - List> resourceSpecs = urns.stream() + List> resourceSpecs = urns.stream() .map(UrnUtils::getUrn) - .map(urn -> java.util.Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))) + .map(urn -> java.util.Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))) .collect(Collectors.toList()); Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) @@ -590,7 +590,7 @@ public Task deleteEntity(@ActionParam(PARAM_URN) @Nonnull Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.DELETE_ENTITY_PRIVILEGE), - Collections.singletonList(java.util.Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))))) { + Collections.singletonList(java.util.Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to delete entity: " + urnStr); } @@ -638,7 +638,7 @@ private Long deleteTimeseriesAspects(@Nonnull Urn urn, @Nullable Long startTimeM Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.DELETE_ENTITY_PRIVILEGE), - new ResourceSpec(urn.getEntityType(), urn.toString()))) { + new EntitySpec(urn.getEntityType(), urn.toString()))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to delete entity " + urn); } @@ -678,7 +678,7 @@ public Task deleteReferencesTo(@ActionParam(PARAM_URN) Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.DELETE_ENTITY_PRIVILEGE), - new ResourceSpec(urn.getEntityType(), urnStr))) { + new EntitySpec(urn.getEntityType(), urnStr))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to delete entity " + urnStr); } @@ -695,7 +695,7 @@ public Task deleteReferencesTo(@ActionParam(PARAM_URN) public Task setWriteable(@ActionParam(PARAM_VALUE) @Optional("true") @Nonnull Boolean value) { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SET_WRITEABLE_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SET_WRITEABLE_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to enable and disable write mode."); } @@ -712,7 +712,7 @@ public Task setWriteable(@ActionParam(PARAM_VALUE) @Optional("true") @Nonn public Task getTotalEntityCount(@ActionParam(PARAM_ENTITY) @Nonnull String entityName) { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_COUNTS_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_COUNTS_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get entity counts."); } @@ -725,7 +725,7 @@ public Task getTotalEntityCount(@ActionParam(PARAM_ENTITY) @Nonnull String public Task batchGetTotalEntityCount(@ActionParam(PARAM_ENTITIES) @Nonnull String[] entityNames) { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_COUNTS_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_COUNTS_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get entity counts."); } @@ -739,7 +739,7 @@ public Task listUrns(@ActionParam(PARAM_ENTITY) @Nonnull String @ActionParam(PARAM_START) int start, @ActionParam(PARAM_COUNT) int count) throws URISyntaxException { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } @@ -757,10 +757,10 @@ public Task applyRetention(@ActionParam(PARAM_START) @Optional @Nullable @ActionParam(PARAM_URN) @Optional @Nullable String urn ) { Authentication auth = AuthenticationContext.getAuthentication(); - ResourceSpec resourceSpec = null; + EntitySpec resourceSpec = null; if (StringUtils.isNotBlank(urn)) { Urn resource = UrnUtils.getUrn(urn); - resourceSpec = new ResourceSpec(resource.getEntityType(), resource.toString()); + resourceSpec = new EntitySpec(resource.getEntityType(), resource.toString()); } if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.APPLY_RETENTION_PRIVILEGE), resourceSpec)) { @@ -781,7 +781,7 @@ public Task filter(@ActionParam(PARAM_ENTITY) @Nonnull String enti Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.SEARCH_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } @@ -799,7 +799,7 @@ public Task exists(@ActionParam(PARAM_URN) @Nonnull String urnStr) thro Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), - new ResourceSpec(urn.getEntityType(), urnStr))) { + new EntitySpec(urn.getEntityType(), urnStr))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized get entity: " + urnStr); } diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityV2Resource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityV2Resource.java index 7efb93c0f50e6..0c3e93273b863 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityV2Resource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityV2Resource.java @@ -4,7 +4,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; import com.datahub.plugins.auth.authorization.Authorizer; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; import com.linkedin.entity.EntityResponse; @@ -68,7 +68,7 @@ public Task get(@Nonnull String urnStr, final Urn urn = Urn.createFromString(urnStr); Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), new ResourceSpec(urn.getEntityType(), urnStr))) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), new EntitySpec(urn.getEntityType(), urnStr))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get entity " + urn); } @@ -96,8 +96,8 @@ public Task> batchGet(@Nonnull Set urnStrs, urns.add(Urn.createFromString(urnStr)); } Authentication auth = AuthenticationContext.getAuthentication(); - List> resourceSpecs = urns.stream() - .map(urn -> java.util.Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))) + List> resourceSpecs = urns.stream() + .map(urn -> java.util.Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))) .collect(Collectors.toList()); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), resourceSpecs)) { diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityVersionedV2Resource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityVersionedV2Resource.java index fd5c3507b5408..05b7e6b3ff24b 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityVersionedV2Resource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityVersionedV2Resource.java @@ -4,7 +4,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; import com.datahub.plugins.auth.authorization.Authorizer; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.google.common.collect.ImmutableList; import com.linkedin.common.VersionedUrn; import com.linkedin.common.urn.Urn; @@ -65,9 +65,9 @@ public Task> batchGetVersioned( @QueryParam(PARAM_ENTITY_TYPE) @Nonnull String entityType, @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { Authentication auth = AuthenticationContext.getAuthentication(); - List> resourceSpecs = versionedUrnStrs.stream() + List> resourceSpecs = versionedUrnStrs.stream() .map(versionedUrn -> UrnUtils.getUrn(versionedUrn.getUrn())) - .map(urn -> java.util.Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))) + .map(urn -> java.util.Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))) .collect(Collectors.toList()); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), resourceSpecs)) { diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/lineage/Relationships.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/lineage/Relationships.java index 313d16333f9e9..4a8e74c89039a 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/lineage/Relationships.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/lineage/Relationships.java @@ -4,7 +4,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; import com.datahub.plugins.auth.authorization.Authorizer; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.google.common.collect.ImmutableList; import com.linkedin.common.EntityRelationship; import com.linkedin.common.EntityRelationshipArray; @@ -107,7 +107,7 @@ public Task get(@QueryParam("urn") @Nonnull String rawUrn, Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), - Collections.singletonList(java.util.Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))))) { + Collections.singletonList(java.util.Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get entity lineage: " + rawUrn); } @@ -142,7 +142,7 @@ public UpdateResponse delete(@QueryParam("urn") @Nonnull String rawUrn) throws E Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.DELETE_ENTITY_PRIVILEGE), - Collections.singletonList(java.util.Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))))) { + Collections.singletonList(java.util.Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to delete entity: " + rawUrn); } @@ -162,7 +162,7 @@ public Task getLineage(@ActionParam(PARAM_URN) @Nonnull Str Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE), - Collections.singletonList(java.util.Optional.of(new ResourceSpec(urn.getEntityType(), urn.toString()))))) { + Collections.singletonList(java.util.Optional.of(new EntitySpec(urn.getEntityType(), urn.toString()))))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get entity lineage: " + urnStr); } diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/operations/Utils.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/operations/Utils.java index 188e5ae18ee8f..12586b66495a9 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/operations/Utils.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/operations/Utils.java @@ -2,7 +2,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; @@ -37,10 +37,10 @@ public static String restoreIndices( @Nonnull EntityService entityService ) { Authentication authentication = AuthenticationContext.getAuthentication(); - ResourceSpec resourceSpec = null; + EntitySpec resourceSpec = null; if (StringUtils.isNotBlank(urn)) { Urn resource = UrnUtils.getUrn(urn); - resourceSpec = new ResourceSpec(resource.getEntityType(), resource.toString()); + resourceSpec = new EntitySpec(resource.getEntityType(), resource.toString()); } if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(authentication, authorizer, ImmutableList.of(PoliciesConfig.RESTORE_INDICES_PRIVILEGE), diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/platform/PlatformResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/platform/PlatformResource.java index f36841bb4abae..a8018074497c4 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/platform/PlatformResource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/platform/PlatformResource.java @@ -3,7 +3,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; import com.datahub.plugins.auth.authorization.Authorizer; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.google.common.collect.ImmutableList; import com.linkedin.entity.Entity; import com.linkedin.metadata.authorization.PoliciesConfig; @@ -54,7 +54,7 @@ public Task producePlatformEvent( @ActionParam("event") @Nonnull PlatformEvent event) { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.PRODUCE_PLATFORM_EVENT_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.PRODUCE_PLATFORM_EVENT_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to produce platform events."); } diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/restli/RestliUtils.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/restli/RestliUtils.java index 5c3b90a84aec1..9949556c99b81 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/restli/RestliUtils.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/restli/RestliUtils.java @@ -4,7 +4,7 @@ import com.datahub.authorization.AuthUtil; import com.datahub.authorization.ConjunctivePrivilegeGroup; import com.datahub.authorization.DisjunctivePrivilegeGroup; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.linkedin.metadata.authorization.PoliciesConfig; @@ -82,13 +82,13 @@ public static RestLiServiceException invalidArgumentsException(@Nullable String } public static boolean isAuthorized(@Nonnull Authentication authentication, @Nonnull Authorizer authorizer, - @Nonnull final List privileges, @Nonnull final List> resources) { + @Nonnull final List privileges, @Nonnull final List> resources) { DisjunctivePrivilegeGroup orGroup = convertPrivilegeGroup(privileges); return AuthUtil.isAuthorizedForResources(authorizer, authentication.getActor().toUrnStr(), resources, orGroup); } public static boolean isAuthorized(@Nonnull Authentication authentication, @Nonnull Authorizer authorizer, - @Nonnull final List privileges, @Nullable final ResourceSpec resource) { + @Nonnull final List privileges, @Nullable final EntitySpec resource) { DisjunctivePrivilegeGroup orGroup = convertPrivilegeGroup(privileges); return AuthUtil.isAuthorized(authorizer, authentication.getActor().toUrnStr(), java.util.Optional.ofNullable(resource), orGroup); } diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/usage/UsageStats.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/usage/UsageStats.java index be70cf9c494ef..02d413301f3b4 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/usage/UsageStats.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/usage/UsageStats.java @@ -4,7 +4,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; import com.datahub.plugins.auth.authorization.Authorizer; -import com.datahub.authorization.ResourceSpec; +import com.datahub.authorization.EntitySpec; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.StreamReadConstraints; import com.fasterxml.jackson.databind.JsonNode; @@ -125,7 +125,7 @@ public Task batchIngest(@ActionParam(PARAM_BUCKETS) @Nonnull UsageAggregat return RestliUtil.toTask(() -> { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) - && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.EDIT_ENTITY_PRIVILEGE), (ResourceSpec) null)) { + && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.EDIT_ENTITY_PRIVILEGE), (EntitySpec) null)) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to edit entities."); } @@ -323,7 +323,7 @@ public Task query(@ActionParam(PARAM_RESOURCE) @Nonnull String Urn resourceUrn = UrnUtils.getUrn(resource); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.VIEW_DATASET_USAGE_PRIVILEGE), - new ResourceSpec(resourceUrn.getEntityType(), resourceUrn.toString()))) { + new EntitySpec(resourceUrn.getEntityType(), resourceUrn.toString()))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to query usage."); } @@ -383,7 +383,7 @@ public Task queryRange(@ActionParam(PARAM_RESOURCE) @Nonnull S Urn resourceUrn = UrnUtils.getUrn(resource); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized(auth, _authorizer, ImmutableList.of(PoliciesConfig.VIEW_DATASET_USAGE_PRIVILEGE), - new ResourceSpec(resourceUrn.getEntityType(), resourceUrn.toString()))) { + new EntitySpec(resourceUrn.getEntityType(), resourceUrn.toString()))) { throw new RestLiServiceException(HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to query usage."); } From d04d25bf428aa442b08a4011fcac81b3d1526a86 Mon Sep 17 00:00:00 2001 From: Kos Korchak <97058061+kkorchak@users.noreply.github.com> Date: Thu, 12 Oct 2023 15:50:20 -0400 Subject: [PATCH 125/156] smoke test(): Query plus filter search test (#8993) --- .../e2e/search/query_and_filter_search.js | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 smoke-test/tests/cypress/cypress/e2e/search/query_and_filter_search.js diff --git a/smoke-test/tests/cypress/cypress/e2e/search/query_and_filter_search.js b/smoke-test/tests/cypress/cypress/e2e/search/query_and_filter_search.js new file mode 100644 index 0000000000000..4637310b86496 --- /dev/null +++ b/smoke-test/tests/cypress/cypress/e2e/search/query_and_filter_search.js @@ -0,0 +1,57 @@ +describe("auto-complete dropdown, filter plus query search test", () => { + + const platformQuerySearch = (query,test_id,active_filter) => { + cy.visit("/"); + cy.get("input[data-testid=search-input]").type(query); + cy.get(`[data-testid="quick-filter-urn:li:dataPlatform:${test_id}"]`).click(); + cy.focused().type("{enter}").wait(3000); + cy.url().should( + "include", + `?filter_platform___false___EQUAL___0=urn%3Ali%3AdataPlatform%3A${test_id}` + ); + cy.get('[data-testid="search-input"]').should("have.value", query); + cy.get(`[data-testid="active-filter-${active_filter}"]`).should("be.visible"); + cy.contains("of 0 results").should("not.exist"); + cy.contains(/of [0-9]+ results/); + } + + const entityQuerySearch = (query,test_id,active_filter) => { + cy.visit("/"); + cy.get("input[data-testid=search-input]").type(query); + cy.get(`[data-testid="quick-filter-${test_id}"]`).click(); + cy.focused().type("{enter}").wait(3000); + cy.url().should( + "include", + `?filter__entityType___false___EQUAL___0=${test_id}` + ); + cy.get('[data-testid="search-input"]').should("have.value", query); + cy.get(`[data-testid="active-filter-${active_filter}"]`).should("be.visible"); + cy.contains("of 0 results").should("not.exist"); + cy.contains(/of [0-9]+ results/); + } + + it("verify the 'filter by' section + query (result in search page with query applied + filter applied)", () => { + // Platform query plus filter test + cy.loginWithCredentials(); + // Airflow + platformQuerySearch ("cypress","airflow","Airflow"); + // BigQuery + platformQuerySearch ("cypress","bigquery","BigQuery"); + // dbt + platformQuerySearch ("cypress","dbt","dbt"); + // Hive + platformQuerySearch ("cypress","hive","Hive"); + + // Entity type query plus filter test + // Datasets + entityQuerySearch ("cypress","DATASET","Datasets"); + // Dashboards + entityQuerySearch ("cypress","DASHBOARD","Dashboards"); + // Pipelines + entityQuerySearch ("cypress","DATA_FLOW","Pipelines"); + // Domains + entityQuerySearch ("Marketing","DOMAIN","Domains"); + // Glossary Terms + entityQuerySearch ("cypress","GLOSSARY_TERM","Glossary Terms"); + }); +}); \ No newline at end of file From a8f0080c08b5c816f0dae9d3bef07ea00220541e Mon Sep 17 00:00:00 2001 From: Tamas Nemeth Date: Fri, 13 Oct 2023 00:14:45 +0200 Subject: [PATCH 126/156] feat(ingest/teradata): Teradata source (#8977) --- .../docs/sources/teradata/teradata_pre.md | 28 +++ .../docs/sources/teradata/teradata_recipe.yml | 17 ++ metadata-ingestion/setup.py | 3 + .../datahub/ingestion/source/sql/teradata.py | 228 ++++++++++++++++++ .../testing/check_sql_parser_result.py | 5 +- .../src/datahub/utilities/sqlglot_lineage.py | 5 + .../test_teradata_default_normalization.json | 38 +++ .../unit/sql_parsing/test_sqlglot_lineage.py | 42 ++++ 8 files changed, 365 insertions(+), 1 deletion(-) create mode 100644 metadata-ingestion/docs/sources/teradata/teradata_pre.md create mode 100644 metadata-ingestion/docs/sources/teradata/teradata_recipe.yml create mode 100644 metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py create mode 100644 metadata-ingestion/tests/unit/sql_parsing/goldens/test_teradata_default_normalization.json diff --git a/metadata-ingestion/docs/sources/teradata/teradata_pre.md b/metadata-ingestion/docs/sources/teradata/teradata_pre.md new file mode 100644 index 0000000000000..eb59caa29eb52 --- /dev/null +++ b/metadata-ingestion/docs/sources/teradata/teradata_pre.md @@ -0,0 +1,28 @@ +### Prerequisites +1. Create a user which has access to the database you want to ingest. + ```sql + CREATE USER datahub FROM AS PASSWORD = PERM = 20000000; + ``` +2. Create a user with the following privileges: + ```sql + GRANT SELECT ON dbc.columns TO datahub; + GRANT SELECT ON dbc.databases TO datahub; + GRANT SELECT ON dbc.tables TO datahub; + GRANT SELECT ON DBC.All_RI_ChildrenV TO datahub; + GRANT SELECT ON DBC.ColumnsV TO datahub; + GRANT SELECT ON DBC.IndicesV TO datahub; + GRANT SELECT ON dbc.TableTextV TO datahub; + GRANT SELECT ON dbc.TablesV TO datahub; + GRANT SELECT ON dbc.dbqlogtbl TO datahub; -- if lineage or usage extraction is enabled + ``` + + If you want to run profiling, you need to grant select permission on all the tables you want to profile. + +3. If linege or usage extraction is enabled, please, check if query logging is enabled and it is set to size which +will fit for your queries (the default query text size Teradata captures is max 200 chars) + An example how you can set it for all users: + ```sql + REPLACE QUERY LOGGING LIMIT SQLTEXT=2000 ON ALL; + ``` + See more here about query logging: + [https://docs.teradata.com/r/Teradata-VantageCloud-Lake/Database-Reference/Database-Administration/Tracking-Query-Behavior-with-Database-Query-Logging-Operational-DBAs]() diff --git a/metadata-ingestion/docs/sources/teradata/teradata_recipe.yml b/metadata-ingestion/docs/sources/teradata/teradata_recipe.yml new file mode 100644 index 0000000000000..8cf07ba4c3a01 --- /dev/null +++ b/metadata-ingestion/docs/sources/teradata/teradata_recipe.yml @@ -0,0 +1,17 @@ +pipeline_name: my-teradata-ingestion-pipeline +source: + type: teradata + config: + host_port: "myteradatainstance.teradata.com:1025" + #platform_instance: "myteradatainstance" + username: myuser + password: mypassword + #database_pattern: + # allow: + # - "demo_user" + # ignoreCase: true + include_table_lineage: true + include_usage_statistics: true + stateful_ingestion: + enabled: true +sink: diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py index 61e7b684682a4..3ea9a2ea61d74 100644 --- a/metadata-ingestion/setup.py +++ b/metadata-ingestion/setup.py @@ -373,6 +373,7 @@ # FIXME: I don't think tableau uses sqllineage anymore so we should be able # to remove that dependency. "tableau": {"tableauserverclient>=0.17.0"} | sqllineage_lib | sqlglot_lib, + "teradata": sql_common | {"teradatasqlalchemy>=17.20.0.0"}, "trino": sql_common | trino, "starburst-trino-usage": sql_common | usage_common | trino, "nifi": {"requests", "packaging", "requests-gssapi"}, @@ -499,6 +500,7 @@ "s3", "snowflake", "tableau", + "teradata", "trino", "hive", "starburst-trino-usage", @@ -597,6 +599,7 @@ "tableau = datahub.ingestion.source.tableau:TableauSource", "openapi = datahub.ingestion.source.openapi:OpenApiSource", "metabase = datahub.ingestion.source.metabase:MetabaseSource", + "teradata = datahub.ingestion.source.sql.teradata:TeradataSource", "trino = datahub.ingestion.source.sql.trino:TrinoSource", "starburst-trino-usage = datahub.ingestion.source.usage.starburst_trino_usage:TrinoUsageSource", "nifi = datahub.ingestion.source.nifi:NifiSource", diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py b/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py new file mode 100644 index 0000000000000..dd11cd840bed9 --- /dev/null +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py @@ -0,0 +1,228 @@ +import logging +from dataclasses import dataclass +from typing import Iterable, Optional, Set, Union + +# This import verifies that the dependencies are available. +import teradatasqlalchemy # noqa: F401 +import teradatasqlalchemy.types as custom_types +from pydantic.fields import Field +from sqlalchemy import create_engine +from sqlalchemy.engine import Engine + +from datahub.configuration.common import AllowDenyPattern +from datahub.configuration.time_window_config import BaseTimeWindowConfig +from datahub.emitter.sql_parsing_builder import SqlParsingBuilder +from datahub.ingestion.api.common import PipelineContext +from datahub.ingestion.api.decorators import ( + SourceCapability, + SupportStatus, + capability, + config_class, + platform_name, + support_status, +) +from datahub.ingestion.api.workunit import MetadataWorkUnit +from datahub.ingestion.graph.client import DataHubGraph +from datahub.ingestion.source.sql.sql_common import SqlWorkUnit, register_custom_type +from datahub.ingestion.source.sql.sql_generic_profiler import ProfilingSqlReport +from datahub.ingestion.source.sql.two_tier_sql_source import ( + TwoTierSQLAlchemyConfig, + TwoTierSQLAlchemySource, +) +from datahub.ingestion.source.usage.usage_common import BaseUsageConfig +from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport +from datahub.ingestion.source_report.time_window import BaseTimeWindowReport +from datahub.metadata.com.linkedin.pegasus2avro.schema import ( + BytesTypeClass, + TimeTypeClass, +) +from datahub.utilities.sqlglot_lineage import SchemaResolver, sqlglot_lineage + +logger: logging.Logger = logging.getLogger(__name__) + +register_custom_type(custom_types.JSON, BytesTypeClass) +register_custom_type(custom_types.INTERVAL_DAY, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_DAY_TO_SECOND, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_DAY_TO_MINUTE, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_DAY_TO_HOUR, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_SECOND, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_MINUTE, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_MINUTE_TO_SECOND, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_HOUR, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_HOUR_TO_MINUTE, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_HOUR_TO_SECOND, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_MONTH, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_YEAR, TimeTypeClass) +register_custom_type(custom_types.INTERVAL_YEAR_TO_MONTH, TimeTypeClass) +register_custom_type(custom_types.MBB, BytesTypeClass) +register_custom_type(custom_types.MBR, BytesTypeClass) +register_custom_type(custom_types.GEOMETRY, BytesTypeClass) +register_custom_type(custom_types.TDUDT, BytesTypeClass) +register_custom_type(custom_types.XML, BytesTypeClass) + + +@dataclass +class TeradataReport(ProfilingSqlReport, IngestionStageReport, BaseTimeWindowReport): + num_queries_parsed: int = 0 + num_table_parse_failures: int = 0 + + +class BaseTeradataConfig(TwoTierSQLAlchemyConfig): + scheme = Field(default="teradatasql", description="database scheme") + + +class TeradataConfig(BaseTeradataConfig, BaseTimeWindowConfig): + database_pattern = Field( + default=AllowDenyPattern(deny=["dbc"]), + description="Regex patterns for databases to filter in ingestion.", + ) + include_table_lineage = Field( + default=False, + description="Whether to include table lineage in the ingestion. " + "This requires to have the table lineage feature enabled.", + ) + + usage: BaseUsageConfig = Field( + description="The usage config to use when generating usage statistics", + default=BaseUsageConfig(), + ) + + use_schema_resolver: bool = Field( + default=True, + description="Read SchemaMetadata aspects from DataHub to aid in SQL parsing. Turn off only for testing.", + hidden_from_docs=True, + ) + + default_db: Optional[str] = Field( + default=None, + description="The default database to use for unqualified table names", + ) + + include_usage_statistics: bool = Field( + default=False, + description="Generate usage statistic.", + ) + + +@platform_name("Teradata") +@config_class(TeradataConfig) +@support_status(SupportStatus.TESTING) +@capability(SourceCapability.DOMAINS, "Enabled by default") +@capability(SourceCapability.CONTAINERS, "Enabled by default") +@capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default") +@capability(SourceCapability.DELETION_DETECTION, "Optionally enabled via configuration") +@capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration") +@capability(SourceCapability.LINEAGE_COARSE, "Optionally enabled via configuration") +@capability(SourceCapability.LINEAGE_FINE, "Optionally enabled via configuration") +@capability(SourceCapability.USAGE_STATS, "Optionally enabled via configuration") +class TeradataSource(TwoTierSQLAlchemySource): + """ + This plugin extracts the following: + + - Metadata for databases, schemas, views, and tables + - Column types associated with each table + - Table, row, and column statistics via optional SQL profiling + """ + + config: TeradataConfig + + LINEAGE_QUERY: str = """SELECT ProcID, UserName as "user", StartTime AT TIME ZONE 'GMT' as "timestamp", DefaultDatabase as default_database, QueryText as query + FROM "DBC".DBQLogTbl + where ErrorCode = 0 + and QueryText like 'create table demo_user.test_lineage%' + and "timestamp" >= TIMESTAMP '{start_time}' + and "timestamp" < TIMESTAMP '{end_time}' + """ + urns: Optional[Set[str]] + + def __init__(self, config: TeradataConfig, ctx: PipelineContext): + super().__init__(config, ctx, "teradata") + + self.report: TeradataReport = TeradataReport() + self.graph: Optional[DataHubGraph] = ctx.graph + + if self.graph: + if self.config.use_schema_resolver: + self.schema_resolver = ( + self.graph.initialize_schema_resolver_from_datahub( + platform=self.platform, + platform_instance=self.config.platform_instance, + env=self.config.env, + ) + ) + self.urns = self.schema_resolver.get_urns() + else: + self.schema_resolver = self.graph._make_schema_resolver( + platform=self.platform, + platform_instance=self.config.platform_instance, + env=self.config.env, + ) + self.urns = None + else: + self.schema_resolver = SchemaResolver( + platform=self.platform, + platform_instance=self.config.platform_instance, + graph=None, + env=self.config.env, + ) + self.urns = None + + self.builder: SqlParsingBuilder = SqlParsingBuilder( + usage_config=self.config.usage + if self.config.include_usage_statistics + else None, + generate_lineage=self.config.include_table_lineage, + generate_usage_statistics=self.config.include_usage_statistics, + generate_operations=self.config.usage.include_operational_stats, + ) + + @classmethod + def create(cls, config_dict, ctx): + config = TeradataConfig.parse_obj(config_dict) + return cls(config, ctx) + + def get_audit_log_mcps(self) -> Iterable[MetadataWorkUnit]: + engine = self.get_metadata_engine() + for entry in engine.execute( + self.LINEAGE_QUERY.format( + start_time=self.config.start_time, end_time=self.config.end_time + ) + ): + self.report.num_queries_parsed += 1 + if self.report.num_queries_parsed % 1000 == 0: + logger.info(f"Parsed {self.report.num_queries_parsed} queries") + + result = sqlglot_lineage( + sql=entry.query, + schema_resolver=self.schema_resolver, + default_db=None, + default_schema=entry.default_database + if entry.default_database + else self.config.default_db, + ) + if result.debug_info.table_error: + logger.debug( + f"Error parsing table lineage, {result.debug_info.table_error}" + ) + self.report.num_table_parse_failures += 1 + continue + + yield from self.builder.process_sql_parsing_result( + result, + query=entry.query, + query_timestamp=entry.timestamp, + user=f"urn:li:corpuser:{entry.user}", + include_urns=self.urns, + ) + + def get_metadata_engine(self) -> Engine: + url = self.config.get_sql_alchemy_url() + logger.debug(f"sql_alchemy_url={url}") + return create_engine(url, **self.config.options) + + def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]: + yield from super().get_workunits_internal() + if self.config.include_table_lineage or self.config.include_usage_statistics: + self.report.report_ingestion_stage_start("audit log extraction") + yield from self.get_audit_log_mcps() + yield from self.builder.gen_workunits() diff --git a/metadata-ingestion/src/datahub/testing/check_sql_parser_result.py b/metadata-ingestion/src/datahub/testing/check_sql_parser_result.py index 8516a7054a9cd..b3b1331db768b 100644 --- a/metadata-ingestion/src/datahub/testing/check_sql_parser_result.py +++ b/metadata-ingestion/src/datahub/testing/check_sql_parser_result.py @@ -70,11 +70,14 @@ def assert_sql_result( sql: str, *, dialect: str, + platform_instance: Optional[str] = None, expected_file: pathlib.Path, schemas: Optional[Dict[str, SchemaInfo]] = None, **kwargs: Any, ) -> None: - schema_resolver = SchemaResolver(platform=dialect) + schema_resolver = SchemaResolver( + platform=dialect, platform_instance=platform_instance + ) if schemas: for urn, schema in schemas.items(): schema_resolver.add_raw_schema_info(urn, schema) diff --git a/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py b/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py index 349eb40a5e865..c830ec8c02fd4 100644 --- a/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py +++ b/metadata-ingestion/src/datahub/utilities/sqlglot_lineage.py @@ -482,6 +482,11 @@ def _column_level_lineage( # noqa: C901 # Our snowflake source lowercases column identifiers, so we are forced # to do fuzzy (case-insensitive) resolution instead of exact resolution. "snowflake", + # Teradata column names are case-insensitive. + # A name, even when enclosed in double quotation marks, is not case sensitive. For example, CUSTOMER and Customer are the same. + # See more below: + # https://documentation.sas.com/doc/en/pgmsascdc/9.4_3.5/acreldb/n0ejgx4895bofnn14rlguktfx5r3.htm + "teradata", } sqlglot_db_schema = sqlglot.MappingSchema( diff --git a/metadata-ingestion/tests/unit/sql_parsing/goldens/test_teradata_default_normalization.json b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_teradata_default_normalization.json new file mode 100644 index 0000000000000..b0351a7e07ad2 --- /dev/null +++ b/metadata-ingestion/tests/unit/sql_parsing/goldens/test_teradata_default_normalization.json @@ -0,0 +1,38 @@ +{ + "query_type": "CREATE", + "in_tables": [ + "urn:li:dataset:(urn:li:dataPlatform:teradata,myteradata.demo_user.pima_patient_diagnoses,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:teradata,myteradata.demo_user.pima_patient_features,PROD)" + ], + "out_tables": [ + "urn:li:dataset:(urn:li:dataPlatform:teradata,myteradata.demo_user.test_lineage2,PROD)" + ], + "column_lineage": [ + { + "downstream": { + "table": "urn:li:dataset:(urn:li:dataPlatform:teradata,myteradata.demo_user.test_lineage2,PROD)", + "column": "PatientId", + "native_column_type": "INTEGER()" + }, + "upstreams": [ + { + "table": "urn:li:dataset:(urn:li:dataPlatform:teradata,myteradata.demo_user.pima_patient_diagnoses,PROD)", + "column": "PatientId" + } + ] + }, + { + "downstream": { + "table": "urn:li:dataset:(urn:li:dataPlatform:teradata,myteradata.demo_user.test_lineage2,PROD)", + "column": "BMI", + "native_column_type": "FLOAT()" + }, + "upstreams": [ + { + "table": "urn:li:dataset:(urn:li:dataPlatform:teradata,myteradata.demo_user.pima_patient_features,PROD)", + "column": "BMI" + } + ] + } + ] +} \ No newline at end of file diff --git a/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py b/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py index bb6e5f1581754..059add8db67e4 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py +++ b/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py @@ -630,3 +630,45 @@ def test_snowflake_column_cast(): # TODO: Add a test for setting platform_instance or env + + +def test_teradata_default_normalization(): + assert_sql_result( + """ +create table demo_user.test_lineage2 as + ( + select + ppd.PatientId, + ppf.bmi + from + demo_user.pima_patient_features ppf + join demo_user.pima_patient_diagnoses ppd on + ppd.PatientId = ppf.PatientId + ) with data; +""", + dialect="teradata", + default_schema="dbc", + platform_instance="myteradata", + schemas={ + "urn:li:dataset:(urn:li:dataPlatform:teradata,myteradata.demo_user.pima_patient_diagnoses,PROD)": { + "HasDiabetes": "INTEGER()", + "PatientId": "INTEGER()", + }, + "urn:li:dataset:(urn:li:dataPlatform:teradata,myteradata.demo_user.pima_patient_features,PROD)": { + "Age": "INTEGER()", + "BMI": "FLOAT()", + "BloodP": "INTEGER()", + "DiPedFunc": "FLOAT()", + "NumTimesPrg": "INTEGER()", + "PatientId": "INTEGER()", + "PlGlcConc": "INTEGER()", + "SkinThick": "INTEGER()", + "TwoHourSerIns": "INTEGER()", + }, + "urn:li:dataset:(urn:li:dataPlatform:teradata,myteradata.demo_user.test_lineage2,PROD)": { + "BMI": "FLOAT()", + "PatientId": "INTEGER()", + }, + }, + expected_file=RESOURCE_DIR / "test_teradata_default_normalization.json", + ) From 71c9bd3a495c1f3663d2268088f04d56dd8c37c9 Mon Sep 17 00:00:00 2001 From: Aseem Bansal Date: Fri, 13 Oct 2023 11:48:22 +0530 Subject: [PATCH 127/156] ci(ingest): update base requirements (#8995) --- .../base-requirements.txt | 398 +++++++++--------- 1 file changed, 205 insertions(+), 193 deletions(-) diff --git a/docker/datahub-ingestion-base/base-requirements.txt b/docker/datahub-ingestion-base/base-requirements.txt index 82d9a93a9a2c3..eb082d50b3020 100644 --- a/docker/datahub-ingestion-base/base-requirements.txt +++ b/docker/datahub-ingestion-base/base-requirements.txt @@ -2,62 +2,58 @@ # pyspark==3.0.3 # pydeequ==1.0.1 -acryl-datahub-classify==0.0.6 -acryl-iceberg-legacy==0.0.4 -acryl-PyHive==0.6.13 -aenum==3.1.12 -aiohttp==3.8.4 +acryl-datahub-classify==0.0.8 +acryl-PyHive==0.6.14 +acryl-sqlglot==18.5.2.dev45 +aenum==3.1.15 +aiohttp==3.8.6 aiosignal==1.3.1 -alembic==1.11.1 +alembic==1.12.0 altair==4.2.0 -anyio==3.7.0 -apache-airflow==2.6.1 -apache-airflow-providers-common-sql==1.5.1 -apache-airflow-providers-ftp==3.4.1 -apache-airflow-providers-http==4.4.1 -apache-airflow-providers-imap==3.2.1 -apache-airflow-providers-sqlite==3.4.1 -apispec==5.2.2 +anyio==3.7.1 +apache-airflow==2.7.2 +apache-airflow-providers-common-sql==1.7.2 +apache-airflow-providers-ftp==3.5.2 +apache-airflow-providers-http==4.5.2 +apache-airflow-providers-imap==3.3.2 +apache-airflow-providers-sqlite==3.4.3 +apispec==6.3.0 appdirs==1.4.4 appnope==0.1.3 -argcomplete==3.0.8 -argon2-cffi==21.3.0 +argcomplete==3.1.2 +argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 asgiref==3.7.2 asn1crypto==1.5.1 -asttokens==2.2.1 -async-timeout==4.0.2 +asttokens==2.4.0 +async-timeout==4.0.3 asynch==0.2.2 attrs==23.1.0 avro==1.10.2 -avro-gen3==0.7.10 -azure-core==1.26.4 -azure-identity==1.10.0 -azure-storage-blob==12.16.0 -azure-storage-file-datalake==12.11.0 -Babel==2.12.1 +avro-gen3==0.7.11 +Babel==2.13.0 backcall==0.2.0 backoff==2.2.1 beautifulsoup4==4.12.2 -bleach==6.0.0 -blinker==1.6.2 -blis==0.7.9 -boto3==1.26.142 -botocore==1.29.142 +bleach==6.1.0 +blinker==1.6.3 +blis==0.7.11 +boto3==1.28.62 +botocore==1.31.62 bowler==0.9.0 -bracex==2.3.post1 +bracex==2.4 cached-property==1.5.2 cachelib==0.9.0 cachetools==5.3.1 -catalogue==2.0.8 -cattrs==22.2.0 -certifi==2023.5.7 -cffi==1.15.1 -chardet==5.1.0 -charset-normalizer==2.1.1 +catalogue==2.0.10 +cattrs==23.1.2 +certifi==2023.7.22 +cffi==1.16.0 +chardet==5.2.0 +charset-normalizer==3.3.0 ciso8601==2.3.0 -click==8.1.3 -click-default-group==1.2.2 +click==8.1.7 +click-default-group==1.2.4 click-spinner==0.1.10 clickclick==20.10.2 clickhouse-cityhash==1.0.2.4 @@ -66,205 +62,217 @@ clickhouse-sqlalchemy==0.2.4 cloudpickle==2.2.1 colorama==0.4.6 colorlog==4.8.0 -confection==0.0.4 +comm==0.1.4 +confection==0.1.3 ConfigUpdater==3.1.1 confluent-kafka==1.8.2 connexion==2.14.2 cron-descriptor==1.4.0 -croniter==1.3.15 -cryptography==37.0.4 +croniter==2.0.1 +cryptography==41.0.4 cx-Oracle==8.3.0 -cymem==2.0.7 -dask==2023.5.1 -databricks-cli==0.17.7 +cymem==2.0.8 +dask==2023.9.3 +databricks-cli==0.18.0 databricks-dbapi==0.6.0 -databricks-sdk==0.1.8 -debugpy==1.6.7 +databricks-sdk==0.10.0 +debugpy==1.8.0 decorator==5.1.1 defusedxml==0.7.1 -deltalake==0.9.0 +deltalake==0.11.0 Deprecated==1.2.14 -dill==0.3.6 -dnspython==2.3.0 -docker==6.1.2 +dill==0.3.7 +dnspython==2.4.2 +docker==6.1.3 docutils==0.20.1 ecdsa==0.18.0 elasticsearch==7.13.4 email-validator==1.3.1 entrypoints==0.4 et-xmlfile==1.1.0 -exceptiongroup==1.1.1 -executing==1.2.0 -expandvars==0.9.0 -fastapi==0.95.2 -fastavro==1.7.4 -fastjsonschema==2.17.1 -feast==0.29.0 -filelock==3.12.0 +exceptiongroup==1.1.3 +executing==2.0.0 +expandvars==0.11.0 +fastapi==0.103.2 +fastavro==1.8.4 +fastjsonschema==2.18.1 +feast==0.31.1 +filelock==3.12.4 fissix==21.11.13 Flask==2.2.5 flatdict==4.0.1 -frozenlist==1.3.3 -fsspec==2023.5.0 +frozenlist==1.4.0 +fsspec==2023.9.2 future==0.18.3 -GeoAlchemy2==0.13.3 +GeoAlchemy2==0.14.1 gitdb==4.0.10 -GitPython==3.1.31 -google-api-core==2.11.0 -google-auth==2.19.0 -google-cloud-appengine-logging==1.3.0 +GitPython==3.1.37 +google-api-core==2.12.0 +google-auth==2.23.3 +google-cloud-appengine-logging==1.3.2 google-cloud-audit-log==0.2.5 -google-cloud-bigquery==3.10.0 -google-cloud-bigquery-storage==2.19.1 -google-cloud-core==2.3.2 +google-cloud-bigquery==3.12.0 +google-cloud-core==2.3.3 google-cloud-datacatalog-lineage==0.2.2 google-cloud-logging==3.5.0 google-crc32c==1.5.0 -google-resumable-media==2.5.0 -googleapis-common-protos==1.59.0 +google-re2==1.1 +google-resumable-media==2.6.0 +googleapis-common-protos==1.60.0 gql==3.4.1 graphql-core==3.2.3 graphviz==0.20.1 great-expectations==0.15.50 -greenlet==2.0.2 +greenlet==3.0.0 grpc-google-iam-v1==0.12.6 -grpcio==1.54.2 -grpcio-reflection==1.54.2 -grpcio-status==1.54.2 -grpcio-tools==1.54.2 -gssapi==1.8.2 -gunicorn==20.1.0 +grpcio==1.59.0 +grpcio-reflection==1.59.0 +grpcio-status==1.59.0 +grpcio-tools==1.59.0 +gssapi==1.8.3 +gunicorn==21.2.0 h11==0.14.0 -hmsclient==0.1.1 -httpcore==0.17.2 -httptools==0.5.0 -httpx==0.24.1 +httpcore==0.18.0 +httptools==0.6.0 +httpx==0.25.0 humanfriendly==10.0 idna==3.4 -ijson==3.2.0.post0 -importlib-metadata==6.6.0 -importlib-resources==5.12.0 +ijson==3.2.3 +importlib-metadata==6.8.0 +importlib-resources==6.1.0 inflection==0.5.1 ipaddress==1.0.23 ipykernel==6.17.1 -ipython==8.13.2 +ipython==8.16.1 ipython-genutils==0.2.0 -ipywidgets==8.0.6 +ipywidgets==8.1.1 iso3166==2.1.1 isodate==0.6.1 itsdangerous==2.1.2 -jedi==0.18.2 +jedi==0.19.1 Jinja2==3.1.2 jmespath==1.0.1 JPype1==1.4.1 -jsonlines==3.1.0 -jsonpatch==1.32 -jsonpointer==2.3 +jsonlines==4.0.0 +jsonpatch==1.33 +jsonpointer==2.4 jsonref==1.1.0 -jsonschema==4.17.3 +jsonschema==4.19.1 +jsonschema-specifications==2023.7.1 jupyter-server==1.24.0 jupyter_client==7.4.9 jupyter_core==4.12.0 jupyterlab-pygments==0.2.2 -jupyterlab-widgets==3.0.7 +jupyterlab-widgets==3.0.9 langcodes==3.3.0 lark==1.1.4 lazy-object-proxy==1.9.0 leb128==1.0.5 -limits==3.5.0 +limits==3.6.0 linear-tsv==1.1.0 linkify-it-py==2.0.2 lkml==1.3.1 locket==1.0.0 lockfile==0.12.2 looker-sdk==23.0.0 -lxml==4.9.2 +lxml==4.9.3 lz4==4.3.2 makefun==1.15.1 Mako==1.2.4 -Markdown==3.4.3 -markdown-it-py==2.2.0 -MarkupSafe==2.1.2 -marshmallow==3.19.0 -marshmallow-enum==1.5.1 +Markdown==3.5 +markdown-it-py==3.0.0 +MarkupSafe==2.1.3 +marshmallow==3.20.1 marshmallow-oneofschema==3.0.1 marshmallow-sqlalchemy==0.26.1 matplotlib-inline==0.1.6 -mdit-py-plugins==0.3.5 +mdit-py-plugins==0.4.0 mdurl==0.1.2 -mistune==2.0.5 +mistune==3.0.2 mixpanel==4.10.0 -mmh3==4.0.0 -more-itertools==9.1.0 +mlflow-skinny==2.7.1 +mmh3==4.0.1 +mmhash3==3.0.1 +more-itertools==10.1.0 moreorless==0.4.0 -moto==4.1.10 -msal==1.16.0 -msal-extensions==1.0.0 +moto==4.2.5 +msal==1.22.0 multidict==6.0.4 -murmurhash==1.0.9 -mypy==1.3.0 +murmurhash==1.0.10 +mypy==1.6.0 mypy-extensions==1.0.0 nbclassic==1.0.0 nbclient==0.6.3 -nbconvert==7.4.0 -nbformat==5.8.0 -nest-asyncio==1.5.6 +nbconvert==7.9.2 +nbformat==5.9.1 +nest-asyncio==1.5.8 networkx==3.1 -notebook==6.5.4 +notebook==6.5.6 notebook_shim==0.2.3 -numpy==1.24.3 +numpy==1.26.0 oauthlib==3.2.2 okta==1.7.0 +openlineage-airflow==1.2.0 +openlineage-integration-common==1.2.0 +openlineage-python==1.2.0 +openlineage_sql==1.2.0 openpyxl==3.1.2 +opentelemetry-api==1.20.0 +opentelemetry-exporter-otlp==1.20.0 +opentelemetry-exporter-otlp-proto-common==1.20.0 +opentelemetry-exporter-otlp-proto-grpc==1.20.0 +opentelemetry-exporter-otlp-proto-http==1.20.0 +opentelemetry-proto==1.20.0 +opentelemetry-sdk==1.20.0 +opentelemetry-semantic-conventions==0.41b0 ordered-set==4.1.0 oscrypto==1.3.0 -packaging==23.1 +packaging==23.2 pandas==1.5.3 pandavro==1.5.2 pandocfilters==1.5.0 -parse==1.19.0 +parse==1.19.1 parso==0.8.3 -partd==1.4.0 -pathspec==0.9.0 -pathy==0.10.1 +partd==1.4.1 +pathspec==0.11.2 +pathy==0.10.2 pendulum==2.1.2 pexpect==4.8.0 phonenumbers==8.13.0 pickleshare==0.7.5 -platformdirs==3.5.1 -pluggy==1.0.0 -portalocker==2.7.0 -preshed==3.0.8 +platformdirs==3.11.0 +pluggy==1.3.0 +preshed==3.0.9 prison==0.2.1 progressbar2==4.2.0 -prometheus-client==0.17.0 -prompt-toolkit==3.0.38 -proto-plus==1.22.2 -protobuf==4.23.2 +prometheus-client==0.17.1 +prompt-toolkit==3.0.39 +proto-plus==1.22.3 +protobuf==4.24.4 psutil==5.9.5 -psycopg2-binary==2.9.6 +psycopg2-binary==2.9.9 ptyprocess==0.7.0 pure-eval==0.2.2 pure-sasl==0.6.2 -py-partiql-parser==0.3.0 -pyarrow==8.0.0 +py-partiql-parser==0.3.7 +pyarrow==11.0.0 pyasn1==0.5.0 pyasn1-modules==0.3.0 pyathena==2.4.1 pycountry==22.3.5 pycparser==2.21 -pycryptodome==3.18.0 -pycryptodomex==3.18.0 -pydantic==1.10.8 -pydash==7.0.3 +pycryptodome==3.19.0 +pycryptodomex==3.19.0 +pydantic==1.10.13 +pydash==7.0.6 pydruid==0.6.5 -Pygments==2.15.1 -pymongo==4.3.3 -PyMySQL==1.0.3 -pyOpenSSL==22.0.0 +Pygments==2.16.1 +pyiceberg==0.4.0 +pymongo==4.5.0 +PyMySQL==1.1.0 +pyOpenSSL==23.2.0 pyparsing==3.0.9 -pyrsistent==0.19.3 -pyspnego==0.9.0 +pyspnego==0.10.2 python-daemon==3.0.1 python-dateutil==2.8.2 python-dotenv==1.0.0 @@ -272,111 +280,115 @@ python-jose==3.3.0 python-ldap==3.4.3 python-nvd3==0.15.0 python-slugify==8.0.1 -python-stdnum==1.18 -python-tds==1.12.0 -python-utils==3.6.0 +python-stdnum==1.19 +python-tds==1.13.0 +python-utils==3.8.1 python3-openid==3.2.0 -pytz==2023.3 +pytz==2023.3.post1 pytzdata==2020.1 -PyYAML==6.0 -pyzmq==25.1.0 +PyYAML==6.0.1 +pyzmq==24.0.1 ratelimiter==1.2.0.post0 redash-toolbelt==0.1.9 -redshift-connector==2.0.910 -regex==2023.5.5 -requests==2.28.2 +redshift-connector==2.0.914 +referencing==0.30.2 +regex==2023.10.3 +requests==2.31.0 requests-file==1.5.1 requests-gssapi==1.2.3 requests-ntlm==1.2.0 requests-toolbelt==0.10.1 -responses==0.23.1 -retrying==1.3.4 +responses==0.23.3 rfc3339-validator==0.1.4 rfc3986==2.0.0 -rich==13.3.5 -rich_argparse==1.1.0 +rich==13.6.0 +rich-argparse==1.3.0 +rpds-py==0.10.6 rsa==4.9 ruamel.yaml==0.17.17 -s3transfer==0.6.1 -sasl3==0.2.11 -schwifty==2023.3.0 -scipy==1.10.1 +ruamel.yaml.clib==0.2.8 +s3transfer==0.7.0 +schwifty==2023.9.0 +scipy==1.11.3 scramp==1.4.4 Send2Trash==1.8.2 -setproctitle==1.3.2 -simple-salesforce==1.12.4 +sentry-sdk==1.32.0 +setproctitle==1.3.3 +simple-salesforce==1.12.5 six==1.16.0 -smart-open==6.3.0 -smmap==5.0.0 +smart-open==6.4.0 +smmap==5.0.1 sniffio==1.3.0 -snowflake-connector-python==2.9.0 -snowflake-sqlalchemy==1.4.7 -soupsieve==2.4.1 +snowflake-connector-python==3.2.1 +snowflake-sqlalchemy==1.5.0 +sortedcontainers==2.4.0 +soupsieve==2.5 spacy==3.4.3 spacy-legacy==3.0.12 -spacy-loggers==1.0.4 +spacy-loggers==1.0.5 sql-metadata==2.2.2 -SQLAlchemy==1.4.41 -sqlalchemy-bigquery==1.6.1 +SQLAlchemy==1.4.44 +sqlalchemy-bigquery==1.8.0 SQLAlchemy-JSONField==1.0.1.post0 sqlalchemy-pytds==0.3.5 sqlalchemy-redshift==0.8.14 SQLAlchemy-Utils==0.41.1 -sqlalchemy2-stubs==0.0.2a34 -sqllineage==1.3.6 -sqlparse==0.4.3 -srsly==2.4.6 -stack-data==0.6.2 +sqlalchemy2-stubs==0.0.2a35 +sqllineage==1.3.8 +sqlparse==0.4.4 +srsly==2.4.8 +stack-data==0.6.3 starlette==0.27.0 +strictyaml==1.7.3 tableauserverclient==0.25 tableschema==1.20.2 tabulate==0.9.0 tabulator==1.53.5 -tenacity==8.2.2 +tenacity==8.2.3 termcolor==2.3.0 terminado==0.17.1 text-unidecode==1.3 -thinc==8.1.10 -thrift==0.16.0 +thinc==8.1.12 +thrift==0.13.0 thrift-sasl==0.4.3 tinycss2==1.2.1 toml==0.10.2 tomli==2.0.1 +tomlkit==0.12.1 toolz==0.12.0 -tornado==6.3.2 -tqdm==4.65.0 +tornado==6.3.3 +tqdm==4.66.1 traitlets==5.2.1.post0 -trino==0.324.0 +trino==0.327.0 typeguard==2.13.3 typer==0.7.0 -types-PyYAML==6.0.12.10 +types-PyYAML==6.0.12.12 typing-inspect==0.9.0 -typing_extensions==4.5.0 -tzlocal==5.0.1 +typing_extensions==4.8.0 +tzlocal==5.1 uc-micro-py==1.0.2 -ujson==5.7.0 +ujson==5.8.0 unicodecsv==0.14.1 -urllib3==1.26.16 -uvicorn==0.22.0 +urllib3==1.26.17 +uvicorn==0.23.2 uvloop==0.17.0 -vertica-python==1.3.2 -vertica-sqlalchemy-dialect==0.0.1 +vertica-python==1.3.5 +vertica-sqlalchemy-dialect==0.0.8 vininfo==1.7.0 volatile==2.1.0 wasabi==0.10.1 -watchfiles==0.19.0 -wcmatch==8.4.1 -wcwidth==0.2.6 +watchfiles==0.20.0 +wcmatch==8.5 +wcwidth==0.2.8 webencodings==0.5.1 -websocket-client==1.5.2 +websocket-client==1.6.4 websockets==11.0.3 Werkzeug==2.2.3 -widgetsnbextension==4.0.7 +widgetsnbextension==4.0.9 wrapt==1.15.0 -WTForms==3.0.1 +WTForms==3.1.0 xlrd==2.0.1 xmltodict==0.13.0 yarl==1.9.2 zeep==4.2.1 -zipp==3.15.0 -zstd==1.5.5.1 +zstd==1.5.5.1 \ No newline at end of file From c02cbb31e2896f9b596bc329af2e86459057b37e Mon Sep 17 00:00:00 2001 From: Aseem Bansal Date: Fri, 13 Oct 2023 17:52:53 +0530 Subject: [PATCH 128/156] docs(Acryl DataHub): release notes for 0.2.12 (#9006) --- docs-website/sidebars.js | 1 + .../managed-datahub/release-notes/v_0_2_11.md | 2 +- .../managed-datahub/release-notes/v_0_2_12.md | 30 +++++++++++++++++++ 3 files changed, 32 insertions(+), 1 deletion(-) create mode 100644 docs/managed-datahub/release-notes/v_0_2_12.md diff --git a/docs-website/sidebars.js b/docs-website/sidebars.js index 21b3a1d3fe4d3..4fa73c995157a 100644 --- a/docs-website/sidebars.js +++ b/docs-website/sidebars.js @@ -608,6 +608,7 @@ module.exports = { }, { "Managed DataHub Release History": [ + "docs/managed-datahub/release-notes/v_0_2_12", "docs/managed-datahub/release-notes/v_0_2_11", "docs/managed-datahub/release-notes/v_0_2_10", "docs/managed-datahub/release-notes/v_0_2_9", diff --git a/docs/managed-datahub/release-notes/v_0_2_11.md b/docs/managed-datahub/release-notes/v_0_2_11.md index 1f42090848712..c99d10201e097 100644 --- a/docs/managed-datahub/release-notes/v_0_2_11.md +++ b/docs/managed-datahub/release-notes/v_0_2_11.md @@ -7,7 +7,7 @@ Release Availability Date Recommended CLI/SDK --- -- `v0.11.0` with release notes at https://github.com/acryldata/datahub/releases/tag/v0.10.5.5 +- `v0.11.0` with release notes at https://github.com/acryldata/datahub/releases/tag/v0.11.0 - [Deprecation] In LDAP ingestor, the manager_pagination_enabled changed to general pagination_enabled If you are using an older CLI/SDK version then please upgrade it. This applies for all CLI/SDK usages, if you are using it through your terminal, github actions, airflow, in python SDK somewhere, Java SKD etc. This is a strong recommendation to upgrade as we keep on pushing fixes in the CLI and it helps us support you better. diff --git a/docs/managed-datahub/release-notes/v_0_2_12.md b/docs/managed-datahub/release-notes/v_0_2_12.md new file mode 100644 index 0000000000000..b13f471d9bf63 --- /dev/null +++ b/docs/managed-datahub/release-notes/v_0_2_12.md @@ -0,0 +1,30 @@ +# v0.2.12 +--- + +Release Availability Date +--- +13-Oct-2023 + +Recommended CLI/SDK +--- +- `v0.11.0.4` with release notes at https://github.com/acryldata/datahub/releases/tag/v0.11.0.4 +- [breaking] Removed support for SQLAlchemy 1.3.x. Only SQLAlchemy 1.4.x is supported now. +- [breaking] Removed `urn:li:corpuser:datahub` owner for the `Measure`, `Dimension` and `Temporal` tags emitted by Looker and LookML source connectors. +- [breaking] The Airflow plugin no longer supports Airflow 2.0.x or Python 3.7. +- [breaking] Introduced the Airflow plugin v2. If you're using Airflow 2.3+, the v2 plugin will be enabled by default, and so you'll need to switch your requirements to include `pip install 'acryl-datahub-airflow-plugin[plugin-v2]'`. To continue using the v1 plugin, set the `DATAHUB_AIRFLOW_PLUGIN_USE_V1_PLUGIN` environment variable to `true`. +- [breaking] The Unity Catalog ingestion source has a new option `include_metastore`, which will cause all urns to be changed when disabled. +This is currently enabled by default to preserve compatibility, but will be disabled by default and then removed in the future. +If stateful ingestion is enabled, simply setting `include_metastore: false` will perform all required cleanup. +Otherwise, we recommend soft deleting all databricks data via the DataHub CLI: +`datahub delete --platform databricks --soft` and then reingesting with `include_metastore: false`. + + +If you are using an older CLI/SDK version then please upgrade it. This applies for all CLI/SDK usages, if you are using it through your terminal, github actions, airflow, in python SDK somewhere, Java SKD etc. This is a strong recommendation to upgrade as we keep on pushing fixes in the CLI and it helps us support you better. + + +## Release Changelog +--- +- Since `v0.2.11` these changes from OSS DataHub https://github.com/datahub-project/datahub/compare/75252a3d9f6a576904be5a0790d644b9ae2df6ac...10a190470e8c932b6d34cba49de7dbcba687a088 have been pulled in. + +## Some notable features in this SaaS release +- Nested Domains available in this release From 6bc742535379f6cc4558daa67b6561d549d6e607 Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Fri, 13 Oct 2023 12:36:18 -0400 Subject: [PATCH 129/156] feat(cli/datacontract): Add data quality assertion support (#8968) Co-authored-by: Mayuri Nehate <33225191+mayurinehate@users.noreply.github.com> Co-authored-by: Harshal Sheth Co-authored-by: Aseem Bansal --- .../api/entities/datacontract/assertion.py | 7 + .../datacontract/assertion_operator.py | 162 ++++++++++++++++++ .../datacontract/data_quality_assertion.py | 60 ++++--- .../api/entities/datacontract/datacontract.py | 2 +- .../datacontract/freshness_assertion.py | 54 +++--- .../entities/datacontract/schema_assertion.py | 17 +- .../api/entities/datacontract/__init__.py | 0 .../test_data_quality_assertion.py | 55 ++++++ 8 files changed, 292 insertions(+), 65 deletions(-) create mode 100644 metadata-ingestion/src/datahub/api/entities/datacontract/assertion.py create mode 100644 metadata-ingestion/src/datahub/api/entities/datacontract/assertion_operator.py create mode 100644 metadata-ingestion/tests/unit/api/entities/datacontract/__init__.py create mode 100644 metadata-ingestion/tests/unit/api/entities/datacontract/test_data_quality_assertion.py diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/assertion.py b/metadata-ingestion/src/datahub/api/entities/datacontract/assertion.py new file mode 100644 index 0000000000000..c45d4ddc92458 --- /dev/null +++ b/metadata-ingestion/src/datahub/api/entities/datacontract/assertion.py @@ -0,0 +1,7 @@ +from typing import Optional + +from datahub.configuration import ConfigModel + + +class BaseAssertion(ConfigModel): + description: Optional[str] = None diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/assertion_operator.py b/metadata-ingestion/src/datahub/api/entities/datacontract/assertion_operator.py new file mode 100644 index 0000000000000..a41b0f7aafd9f --- /dev/null +++ b/metadata-ingestion/src/datahub/api/entities/datacontract/assertion_operator.py @@ -0,0 +1,162 @@ +from typing import Optional, Union + +from typing_extensions import Literal, Protocol + +from datahub.configuration import ConfigModel +from datahub.metadata.schema_classes import ( + AssertionStdOperatorClass, + AssertionStdParameterClass, + AssertionStdParametersClass, + AssertionStdParameterTypeClass, +) + + +class Operator(Protocol): + """Specification for an assertion operator. + + This class exists only for documentation (not used in typing checking). + """ + + operator: str + + def id(self) -> str: + ... + + def generate_parameters(self) -> AssertionStdParametersClass: + ... + + +def _generate_assertion_std_parameter( + value: Union[str, int, float] +) -> AssertionStdParameterClass: + if isinstance(value, str): + return AssertionStdParameterClass( + value=value, type=AssertionStdParameterTypeClass.STRING + ) + elif isinstance(value, (int, float)): + return AssertionStdParameterClass( + value=str(value), type=AssertionStdParameterTypeClass.NUMBER + ) + else: + raise ValueError( + f"Unsupported assertion parameter {value} of type {type(value)}" + ) + + +Param = Union[str, int, float] + + +def _generate_assertion_std_parameters( + value: Optional[Param] = None, + min_value: Optional[Param] = None, + max_value: Optional[Param] = None, +) -> AssertionStdParametersClass: + return AssertionStdParametersClass( + value=_generate_assertion_std_parameter(value) if value else None, + minValue=_generate_assertion_std_parameter(min_value) if min_value else None, + maxValue=_generate_assertion_std_parameter(max_value) if max_value else None, + ) + + +class EqualToOperator(ConfigModel): + type: Literal["equal_to"] + value: Union[str, int, float] + + operator: str = AssertionStdOperatorClass.EQUAL_TO + + def id(self) -> str: + return f"{self.type}-{self.value}" + + def generate_parameters(self) -> AssertionStdParametersClass: + return _generate_assertion_std_parameters(value=self.value) + + +class BetweenOperator(ConfigModel): + type: Literal["between"] + min: Union[int, float] + max: Union[int, float] + + operator: str = AssertionStdOperatorClass.BETWEEN + + def id(self) -> str: + return f"{self.type}-{self.min}-{self.max}" + + def generate_parameters(self) -> AssertionStdParametersClass: + return _generate_assertion_std_parameters( + min_value=self.min, max_value=self.max + ) + + +class LessThanOperator(ConfigModel): + type: Literal["less_than"] + value: Union[int, float] + + operator: str = AssertionStdOperatorClass.LESS_THAN + + def id(self) -> str: + return f"{self.type}-{self.value}" + + def generate_parameters(self) -> AssertionStdParametersClass: + return _generate_assertion_std_parameters(value=self.value) + + +class GreaterThanOperator(ConfigModel): + type: Literal["greater_than"] + value: Union[int, float] + + operator: str = AssertionStdOperatorClass.GREATER_THAN + + def id(self) -> str: + return f"{self.type}-{self.value}" + + def generate_parameters(self) -> AssertionStdParametersClass: + return _generate_assertion_std_parameters(value=self.value) + + +class LessThanOrEqualToOperator(ConfigModel): + type: Literal["less_than_or_equal_to"] + value: Union[int, float] + + operator: str = AssertionStdOperatorClass.LESS_THAN_OR_EQUAL_TO + + def id(self) -> str: + return f"{self.type}-{self.value}" + + def generate_parameters(self) -> AssertionStdParametersClass: + return _generate_assertion_std_parameters(value=self.value) + + +class GreaterThanOrEqualToOperator(ConfigModel): + type: Literal["greater_than_or_equal_to"] + value: Union[int, float] + + operator: str = AssertionStdOperatorClass.GREATER_THAN_OR_EQUAL_TO + + def id(self) -> str: + return f"{self.type}-{self.value}" + + def generate_parameters(self) -> AssertionStdParametersClass: + return _generate_assertion_std_parameters(value=self.value) + + +class NotNullOperator(ConfigModel): + type: Literal["not_null"] + + operator: str = AssertionStdOperatorClass.NOT_NULL + + def id(self) -> str: + return f"{self.type}" + + def generate_parameters(self) -> AssertionStdParametersClass: + return _generate_assertion_std_parameters() + + +Operators = Union[ + EqualToOperator, + BetweenOperator, + LessThanOperator, + LessThanOrEqualToOperator, + GreaterThanOperator, + GreaterThanOrEqualToOperator, + NotNullOperator, +] diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/data_quality_assertion.py b/metadata-ingestion/src/datahub/api/entities/datacontract/data_quality_assertion.py index a665e95e93c43..6a3944ba36baf 100644 --- a/metadata-ingestion/src/datahub/api/entities/datacontract/data_quality_assertion.py +++ b/metadata-ingestion/src/datahub/api/entities/datacontract/data_quality_assertion.py @@ -4,6 +4,8 @@ from typing_extensions import Literal import datahub.emitter.mce_builder as builder +from datahub.api.entities.datacontract.assertion import BaseAssertion +from datahub.api.entities.datacontract.assertion_operator import Operators from datahub.configuration.common import ConfigModel from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.metadata.schema_classes import ( @@ -14,12 +16,15 @@ AssertionStdParametersClass, AssertionStdParameterTypeClass, AssertionTypeClass, + AssertionValueChangeTypeClass, DatasetAssertionInfoClass, DatasetAssertionScopeClass, + SqlAssertionInfoClass, + SqlAssertionTypeClass, ) -class IdConfigMixin(ConfigModel): +class IdConfigMixin(BaseAssertion): id_raw: Optional[str] = pydantic.Field( default=None, alias="id", @@ -30,25 +35,32 @@ def generate_default_id(self) -> str: raise NotImplementedError -class CustomSQLAssertion(IdConfigMixin, ConfigModel): +class CustomSQLAssertion(IdConfigMixin, BaseAssertion): type: Literal["custom_sql"] - sql: str + operator: Operators = pydantic.Field(discriminator="type") - def generate_dataset_assertion_info( - self, entity_urn: str - ) -> DatasetAssertionInfoClass: - return DatasetAssertionInfoClass( - dataset=entity_urn, - scope=DatasetAssertionScopeClass.UNKNOWN, - fields=[], - operator=AssertionStdOperatorClass._NATIVE_, - aggregation=AssertionStdAggregationClass._NATIVE_, - logic=self.sql, + def generate_default_id(self) -> str: + return f"{self.type}-{self.sql}-{self.operator.id()}" + + def generate_assertion_info(self, entity_urn: str) -> AssertionInfoClass: + sql_assertion_info = SqlAssertionInfoClass( + entity=entity_urn, + statement=self.sql, + operator=self.operator.operator, + parameters=self.operator.generate_parameters(), + # TODO: Support other types of assertions + type=SqlAssertionTypeClass.METRIC, + changeType=AssertionValueChangeTypeClass.ABSOLUTE, + ) + return AssertionInfoClass( + type=AssertionTypeClass.SQL, + sqlAssertion=sql_assertion_info, + description=self.description, ) -class ColumnUniqueAssertion(IdConfigMixin, ConfigModel): +class ColumnUniqueAssertion(IdConfigMixin, BaseAssertion): type: Literal["unique"] # TODO: support multiple columns? @@ -57,10 +69,8 @@ class ColumnUniqueAssertion(IdConfigMixin, ConfigModel): def generate_default_id(self) -> str: return f"{self.type}-{self.column}" - def generate_dataset_assertion_info( - self, entity_urn: str - ) -> DatasetAssertionInfoClass: - return DatasetAssertionInfoClass( + def generate_assertion_info(self, entity_urn: str) -> AssertionInfoClass: + dataset_assertion_info = DatasetAssertionInfoClass( dataset=entity_urn, scope=DatasetAssertionScopeClass.DATASET_COLUMN, fields=[builder.make_schema_field_urn(entity_urn, self.column)], @@ -72,6 +82,11 @@ def generate_dataset_assertion_info( ) ), ) + return AssertionInfoClass( + type=AssertionTypeClass.DATASET, + datasetAssertion=dataset_assertion_info, + description=self.description, + ) class DataQualityAssertion(ConfigModel): @@ -92,16 +107,9 @@ def id(self) -> str: def generate_mcp( self, assertion_urn: str, entity_urn: str ) -> List[MetadataChangeProposalWrapper]: - dataset_assertion_info = self.__root__.generate_dataset_assertion_info( - entity_urn - ) - return [ MetadataChangeProposalWrapper( entityUrn=assertion_urn, - aspect=AssertionInfoClass( - type=AssertionTypeClass.DATASET, - datasetAssertion=dataset_assertion_info, - ), + aspect=self.__root__.generate_assertion_info(entity_urn), ) ] diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/datacontract.py b/metadata-ingestion/src/datahub/api/entities/datacontract/datacontract.py index 2df446623a9d6..f3c6be55e5fea 100644 --- a/metadata-ingestion/src/datahub/api/entities/datacontract/datacontract.py +++ b/metadata-ingestion/src/datahub/api/entities/datacontract/datacontract.py @@ -54,7 +54,7 @@ class DataContract(ConfigModel): freshness: Optional[FreshnessAssertion] = pydantic.Field(default=None) # TODO: Add a validator to ensure that ids are unique - data_quality: Optional[List[DataQualityAssertion]] = None + data_quality: Optional[List[DataQualityAssertion]] = pydantic.Field(default=None) _original_yaml_dict: Optional[dict] = None diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/freshness_assertion.py b/metadata-ingestion/src/datahub/api/entities/datacontract/freshness_assertion.py index ee8fa1181e614..71741d76b22fc 100644 --- a/metadata-ingestion/src/datahub/api/entities/datacontract/freshness_assertion.py +++ b/metadata-ingestion/src/datahub/api/entities/datacontract/freshness_assertion.py @@ -6,6 +6,7 @@ import pydantic from typing_extensions import Literal +from datahub.api.entities.datacontract.assertion import BaseAssertion from datahub.configuration.common import ConfigModel from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.metadata.schema_classes import ( @@ -21,7 +22,7 @@ ) -class CronFreshnessAssertion(ConfigModel): +class CronFreshnessAssertion(BaseAssertion): type: Literal["cron"] cron: str = pydantic.Field( @@ -32,12 +33,30 @@ class CronFreshnessAssertion(ConfigModel): description="The timezone to use for the cron schedule. Defaults to UTC.", ) + def generate_freshness_assertion_schedule(self) -> FreshnessAssertionScheduleClass: + return FreshnessAssertionScheduleClass( + type=FreshnessAssertionScheduleTypeClass.CRON, + cron=FreshnessCronScheduleClass( + cron=self.cron, + timezone=self.timezone, + ), + ) + -class FixedIntervalFreshnessAssertion(ConfigModel): +class FixedIntervalFreshnessAssertion(BaseAssertion): type: Literal["interval"] interval: timedelta + def generate_freshness_assertion_schedule(self) -> FreshnessAssertionScheduleClass: + return FreshnessAssertionScheduleClass( + type=FreshnessAssertionScheduleTypeClass.FIXED_INTERVAL, + fixedInterval=FixedIntervalScheduleClass( + unit=CalendarIntervalClass.SECOND, + multiple=int(self.interval.total_seconds()), + ), + ) + class FreshnessAssertion(ConfigModel): __root__: Union[ @@ -51,36 +70,13 @@ def id(self): def generate_mcp( self, assertion_urn: str, entity_urn: str ) -> List[MetadataChangeProposalWrapper]: - freshness = self.__root__ - - if isinstance(freshness, CronFreshnessAssertion): - schedule = FreshnessAssertionScheduleClass( - type=FreshnessAssertionScheduleTypeClass.CRON, - cron=FreshnessCronScheduleClass( - cron=freshness.cron, - timezone=freshness.timezone, - ), - ) - elif isinstance(freshness, FixedIntervalFreshnessAssertion): - schedule = FreshnessAssertionScheduleClass( - type=FreshnessAssertionScheduleTypeClass.FIXED_INTERVAL, - fixedInterval=FixedIntervalScheduleClass( - unit=CalendarIntervalClass.SECOND, - multiple=int(freshness.interval.total_seconds()), - ), - ) - else: - raise ValueError(f"Unknown freshness type {freshness}") - - assertionInfo = AssertionInfoClass( + aspect = AssertionInfoClass( type=AssertionTypeClass.FRESHNESS, freshnessAssertion=FreshnessAssertionInfoClass( entity=entity_urn, type=FreshnessAssertionTypeClass.DATASET_CHANGE, - schedule=schedule, + schedule=self.__root__.generate_freshness_assertion_schedule(), ), + description=self.__root__.description, ) - - return [ - MetadataChangeProposalWrapper(entityUrn=assertion_urn, aspect=assertionInfo) - ] + return [MetadataChangeProposalWrapper(entityUrn=assertion_urn, aspect=aspect)] diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/schema_assertion.py b/metadata-ingestion/src/datahub/api/entities/datacontract/schema_assertion.py index b5b592e01f58f..b62f94e0592fc 100644 --- a/metadata-ingestion/src/datahub/api/entities/datacontract/schema_assertion.py +++ b/metadata-ingestion/src/datahub/api/entities/datacontract/schema_assertion.py @@ -6,6 +6,7 @@ import pydantic from typing_extensions import Literal +from datahub.api.entities.datacontract.assertion import BaseAssertion from datahub.configuration.common import ConfigModel from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.extractor.json_schema_util import get_schema_metadata @@ -19,7 +20,7 @@ ) -class JsonSchemaContract(ConfigModel): +class JsonSchemaContract(BaseAssertion): type: Literal["json-schema"] json_schema: dict = pydantic.Field(alias="json-schema") @@ -36,7 +37,7 @@ def _init_private_attributes(self) -> None: ) -class FieldListSchemaContract(ConfigModel, arbitrary_types_allowed=True): +class FieldListSchemaContract(BaseAssertion, arbitrary_types_allowed=True): type: Literal["field-list"] fields: List[SchemaFieldClass] @@ -67,15 +68,13 @@ def id(self): def generate_mcp( self, assertion_urn: str, entity_urn: str ) -> List[MetadataChangeProposalWrapper]: - schema_metadata = self.__root__._schema_metadata - - assertionInfo = AssertionInfoClass( + aspect = AssertionInfoClass( type=AssertionTypeClass.DATA_SCHEMA, schemaAssertion=SchemaAssertionInfoClass( - entity=entity_urn, schema=schema_metadata + entity=entity_urn, + schema=self.__root__._schema_metadata, ), + description=self.__root__.description, ) - return [ - MetadataChangeProposalWrapper(entityUrn=assertion_urn, aspect=assertionInfo) - ] + return [MetadataChangeProposalWrapper(entityUrn=assertion_urn, aspect=aspect)] diff --git a/metadata-ingestion/tests/unit/api/entities/datacontract/__init__.py b/metadata-ingestion/tests/unit/api/entities/datacontract/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/metadata-ingestion/tests/unit/api/entities/datacontract/test_data_quality_assertion.py b/metadata-ingestion/tests/unit/api/entities/datacontract/test_data_quality_assertion.py new file mode 100644 index 0000000000000..7be8b667a500b --- /dev/null +++ b/metadata-ingestion/tests/unit/api/entities/datacontract/test_data_quality_assertion.py @@ -0,0 +1,55 @@ +from datahub.api.entities.datacontract.data_quality_assertion import ( + DataQualityAssertion, +) +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.metadata.schema_classes import ( + AssertionInfoClass, + AssertionStdOperatorClass, + AssertionStdParameterClass, + AssertionStdParametersClass, + AssertionStdParameterTypeClass, + AssertionTypeClass, + AssertionValueChangeTypeClass, + SqlAssertionInfoClass, + SqlAssertionTypeClass, +) + + +def test_parse_sql_assertion(): + assertion_urn = "urn:li:assertion:a" + entity_urn = "urn:li:dataset:d" + statement = "SELECT COUNT(*) FROM my_table WHERE value IS NOT NULL" + + d = { + "type": "custom_sql", + "sql": statement, + "operator": {"type": "between", "min": 5, "max": 10}, + } + + assert DataQualityAssertion.parse_obj(d).generate_mcp( + assertion_urn, entity_urn + ) == [ + MetadataChangeProposalWrapper( + entityUrn=assertion_urn, + aspect=AssertionInfoClass( + type=AssertionTypeClass.SQL, + sqlAssertion=SqlAssertionInfoClass( + type=SqlAssertionTypeClass.METRIC, + changeType=AssertionValueChangeTypeClass.ABSOLUTE, + entity=entity_urn, + statement="SELECT COUNT(*) FROM my_table WHERE value IS NOT NULL", + operator=AssertionStdOperatorClass.BETWEEN, + parameters=AssertionStdParametersClass( + minValue=AssertionStdParameterClass( + value="5", + type=AssertionStdParameterTypeClass.NUMBER, + ), + maxValue=AssertionStdParameterClass( + value="10", + type=AssertionStdParameterTypeClass.NUMBER, + ), + ), + ), + ), + ) + ] From 1007204cda802f02a5639e074d95b634b2be9ddf Mon Sep 17 00:00:00 2001 From: Tamas Nemeth Date: Fri, 13 Oct 2023 21:07:19 +0200 Subject: [PATCH 130/156] feat(ingest/teradata): view parsing (#9005) --- .../docs/sources/teradata/teradata_pre.md | 2 +- .../docs/sources/teradata/teradata_recipe.yml | 3 +- .../datahub/ingestion/source/sql/teradata.py | 156 ++++++++++++------ 3 files changed, 106 insertions(+), 55 deletions(-) diff --git a/metadata-ingestion/docs/sources/teradata/teradata_pre.md b/metadata-ingestion/docs/sources/teradata/teradata_pre.md index eb59caa29eb52..7263a59f5ea3d 100644 --- a/metadata-ingestion/docs/sources/teradata/teradata_pre.md +++ b/metadata-ingestion/docs/sources/teradata/teradata_pre.md @@ -18,7 +18,7 @@ If you want to run profiling, you need to grant select permission on all the tables you want to profile. -3. If linege or usage extraction is enabled, please, check if query logging is enabled and it is set to size which +3. If lineage or usage extraction is enabled, please, check if query logging is enabled and it is set to size which will fit for your queries (the default query text size Teradata captures is max 200 chars) An example how you can set it for all users: ```sql diff --git a/metadata-ingestion/docs/sources/teradata/teradata_recipe.yml b/metadata-ingestion/docs/sources/teradata/teradata_recipe.yml index 8cf07ba4c3a01..cc94de20110fe 100644 --- a/metadata-ingestion/docs/sources/teradata/teradata_recipe.yml +++ b/metadata-ingestion/docs/sources/teradata/teradata_recipe.yml @@ -3,12 +3,11 @@ source: type: teradata config: host_port: "myteradatainstance.teradata.com:1025" - #platform_instance: "myteradatainstance" username: myuser password: mypassword #database_pattern: # allow: - # - "demo_user" + # - "my_database" # ignoreCase: true include_table_lineage: true include_usage_statistics: true diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py b/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py index dd11cd840bed9..6080cf7b371e3 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py @@ -1,5 +1,6 @@ import logging from dataclasses import dataclass +from datetime import datetime from typing import Iterable, Optional, Set, Union # This import verifies that the dependencies are available. @@ -11,6 +12,7 @@ from datahub.configuration.common import AllowDenyPattern from datahub.configuration.time_window_config import BaseTimeWindowConfig +from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.sql_parsing_builder import SqlParsingBuilder from datahub.ingestion.api.common import PipelineContext from datahub.ingestion.api.decorators import ( @@ -32,11 +34,18 @@ from datahub.ingestion.source.usage.usage_common import BaseUsageConfig from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport from datahub.ingestion.source_report.time_window import BaseTimeWindowReport +from datahub.metadata._schema_classes import ( + MetadataChangeEventClass, + SchemaMetadataClass, + ViewPropertiesClass, +) from datahub.metadata.com.linkedin.pegasus2avro.schema import ( BytesTypeClass, TimeTypeClass, ) +from datahub.utilities.file_backed_collections import FileBackedDict from datahub.utilities.sqlglot_lineage import SchemaResolver, sqlglot_lineage +from datahub.utilities.urns.dataset_urn import DatasetUrn logger: logging.Logger = logging.getLogger(__name__) @@ -64,6 +73,7 @@ @dataclass class TeradataReport(ProfilingSqlReport, IngestionStageReport, BaseTimeWindowReport): num_queries_parsed: int = 0 + num_view_ddl_parsed: int = 0 num_table_parse_failures: int = 0 @@ -82,17 +92,16 @@ class TeradataConfig(BaseTeradataConfig, BaseTimeWindowConfig): "This requires to have the table lineage feature enabled.", ) + include_view_lineage = Field( + default=True, + description="Whether to include view lineage in the ingestion. " + "This requires to have the view lineage feature enabled.", + ) usage: BaseUsageConfig = Field( description="The usage config to use when generating usage statistics", default=BaseUsageConfig(), ) - use_schema_resolver: bool = Field( - default=True, - description="Read SchemaMetadata aspects from DataHub to aid in SQL parsing. Turn off only for testing.", - hidden_from_docs=True, - ) - default_db: Optional[str] = Field( default=None, description="The default database to use for unqualified table names", @@ -141,46 +150,47 @@ def __init__(self, config: TeradataConfig, ctx: PipelineContext): self.report: TeradataReport = TeradataReport() self.graph: Optional[DataHubGraph] = ctx.graph - if self.graph: - if self.config.use_schema_resolver: - self.schema_resolver = ( - self.graph.initialize_schema_resolver_from_datahub( - platform=self.platform, - platform_instance=self.config.platform_instance, - env=self.config.env, - ) - ) - self.urns = self.schema_resolver.get_urns() - else: - self.schema_resolver = self.graph._make_schema_resolver( - platform=self.platform, - platform_instance=self.config.platform_instance, - env=self.config.env, - ) - self.urns = None - else: - self.schema_resolver = SchemaResolver( - platform=self.platform, - platform_instance=self.config.platform_instance, - graph=None, - env=self.config.env, - ) - self.urns = None - self.builder: SqlParsingBuilder = SqlParsingBuilder( usage_config=self.config.usage if self.config.include_usage_statistics else None, - generate_lineage=self.config.include_table_lineage, + generate_lineage=True, generate_usage_statistics=self.config.include_usage_statistics, generate_operations=self.config.usage.include_operational_stats, ) + self.schema_resolver = SchemaResolver( + platform=self.platform, + platform_instance=self.config.platform_instance, + graph=None, + env=self.config.env, + ) + + self._view_definition_cache: FileBackedDict[str] = FileBackedDict() + @classmethod def create(cls, config_dict, ctx): config = TeradataConfig.parse_obj(config_dict) return cls(config, ctx) + def get_view_lineage(self) -> Iterable[MetadataWorkUnit]: + for key in self._view_definition_cache.keys(): + view_definition = self._view_definition_cache[key] + dataset_urn = DatasetUrn.create_from_string(key) + + db_name: Optional[str] = None + # We need to get the default db from the dataset urn otherwise the builder generates the wrong urns + if "." in dataset_urn.get_dataset_name(): + db_name = dataset_urn.get_dataset_name().split(".", 1)[0] + + self.report.num_view_ddl_parsed += 1 + if self.report.num_view_ddl_parsed % 1000 == 0: + logger.info(f"Parsed {self.report.num_queries_parsed} view ddl") + + yield from self.gen_lineage_from_query( + query=view_definition, default_database=db_name, is_view_ddl=True + ) + def get_audit_log_mcps(self) -> Iterable[MetadataWorkUnit]: engine = self.get_metadata_engine() for entry in engine.execute( @@ -192,27 +202,43 @@ def get_audit_log_mcps(self) -> Iterable[MetadataWorkUnit]: if self.report.num_queries_parsed % 1000 == 0: logger.info(f"Parsed {self.report.num_queries_parsed} queries") - result = sqlglot_lineage( - sql=entry.query, - schema_resolver=self.schema_resolver, - default_db=None, - default_schema=entry.default_database - if entry.default_database - else self.config.default_db, + yield from self.gen_lineage_from_query( + query=entry.query, + default_database=entry.default_database, + timestamp=entry.timestamp, + user=entry.user, + is_view_ddl=False, ) - if result.debug_info.table_error: - logger.debug( - f"Error parsing table lineage, {result.debug_info.table_error}" - ) - self.report.num_table_parse_failures += 1 - continue + def gen_lineage_from_query( + self, + query: str, + default_database: Optional[str] = None, + timestamp: Optional[datetime] = None, + user: Optional[str] = None, + is_view_ddl: bool = False, + ) -> Iterable[MetadataWorkUnit]: + result = sqlglot_lineage( + sql=query, + schema_resolver=self.schema_resolver, + default_db=None, + default_schema=default_database + if default_database + else self.config.default_db, + ) + if result.debug_info.table_error: + logger.debug( + f"Error parsing table lineage, {result.debug_info.table_error}" + ) + self.report.num_table_parse_failures += 1 + else: yield from self.builder.process_sql_parsing_result( result, - query=entry.query, - query_timestamp=entry.timestamp, - user=f"urn:li:corpuser:{entry.user}", - include_urns=self.urns, + query=query, + is_view_ddl=is_view_ddl, + query_timestamp=timestamp, + user=f"urn:li:corpuser:{user}", + include_urns=self.schema_resolver.get_urns(), ) def get_metadata_engine(self) -> Engine: @@ -221,8 +247,34 @@ def get_metadata_engine(self) -> Engine: return create_engine(url, **self.config.options) def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]: - yield from super().get_workunits_internal() + # Add all schemas to the schema resolver + for wu in super().get_workunits_internal(): + if isinstance(wu.metadata, MetadataChangeEventClass): + if wu.metadata.proposedSnapshot: + for aspect in wu.metadata.proposedSnapshot.aspects: + if isinstance(aspect, SchemaMetadataClass): + self.schema_resolver.add_schema_metadata( + wu.metadata.proposedSnapshot.urn, + aspect, + ) + break + if isinstance(wu.metadata, MetadataChangeProposalWrapper): + if ( + wu.metadata.entityUrn + and isinstance(wu.metadata.aspect, ViewPropertiesClass) + and wu.metadata.aspect.viewLogic + ): + self._view_definition_cache[ + wu.metadata.entityUrn + ] = wu.metadata.aspect.viewLogic + yield wu + + if self.config.include_view_lineage: + self.report.report_ingestion_stage_start("view lineage extraction") + yield from self.get_view_lineage() + if self.config.include_table_lineage or self.config.include_usage_statistics: self.report.report_ingestion_stage_start("audit log extraction") yield from self.get_audit_log_mcps() - yield from self.builder.gen_workunits() + + yield from self.builder.gen_workunits() From c2e8041d771db1a20889255372312791fb6d911c Mon Sep 17 00:00:00 2001 From: Tamas Nemeth Date: Fri, 13 Oct 2023 22:59:18 +0200 Subject: [PATCH 131/156] Adding missing sqlparser libs to setup.py (#9015) --- metadata-ingestion/setup.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py index 3ea9a2ea61d74..545cafca9d4df 100644 --- a/metadata-ingestion/setup.py +++ b/metadata-ingestion/setup.py @@ -373,7 +373,10 @@ # FIXME: I don't think tableau uses sqllineage anymore so we should be able # to remove that dependency. "tableau": {"tableauserverclient>=0.17.0"} | sqllineage_lib | sqlglot_lib, - "teradata": sql_common | {"teradatasqlalchemy>=17.20.0.0"}, + "teradata": sql_common + | usage_common + | sqlglot_lib + | {"teradatasqlalchemy>=17.20.0.0"}, "trino": sql_common | trino, "starburst-trino-usage": sql_common | usage_common | trino, "nifi": {"requests", "packaging", "requests-gssapi"}, @@ -432,9 +435,7 @@ deepdiff_dep = "deepdiff" test_api_requirements = {pytest_dep, deepdiff_dep, "PyYAML"} -debug_requirements = { - "memray" -} +debug_requirements = {"memray"} base_dev_requirements = { *base_requirements, From 78b342f441b340189e4eab60574daa60074457e0 Mon Sep 17 00:00:00 2001 From: Indy Prentice Date: Fri, 13 Oct 2023 19:04:44 -0300 Subject: [PATCH 132/156] feat(graphql): support filtering based on greater than/less than criteria (#9001) Co-authored-by: Indy Prentice --- .../src/main/resources/search.graphql | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/datahub-graphql-core/src/main/resources/search.graphql b/datahub-graphql-core/src/main/resources/search.graphql index 4cabdb04afe77..e0cde5a2db9f9 100644 --- a/datahub-graphql-core/src/main/resources/search.graphql +++ b/datahub-graphql-core/src/main/resources/search.graphql @@ -458,6 +458,26 @@ enum FilterOperator { Represents the relation: The field exists. If the field is an array, the field is either not present or empty. """ EXISTS + + """ + Represent the relation greater than, e.g. ownerCount > 5 + """ + GREATER_THAN + + """ + Represent the relation greater than or equal to, e.g. ownerCount >= 5 + """ + GREATER_THAN_OR_EQUAL_TO + + """ + Represent the relation less than, e.g. ownerCount < 3 + """ + LESS_THAN + + """ + Represent the relation less than or equal to, e.g. ownerCount <= 3 + """ + LESS_THAN_OR_EQUAL_TO } """ From c81a339bfc3a57161e433c64bd331ca6af4f6f2d Mon Sep 17 00:00:00 2001 From: Mayuri Nehate <33225191+mayurinehate@users.noreply.github.com> Date: Mon, 16 Oct 2023 21:57:57 +0530 Subject: [PATCH 133/156] build(ingest): remove ratelimiter dependency (#9008) --- metadata-ingestion/setup.py | 1 - .../bigquery_v2/bigquery_audit_log_api.py | 2 +- .../src/datahub/utilities/ratelimiter.py | 56 +++++++++++++++++++ .../tests/unit/utilities/test_ratelimiter.py | 20 +++++++ 4 files changed, 77 insertions(+), 2 deletions(-) create mode 100644 metadata-ingestion/src/datahub/utilities/ratelimiter.py create mode 100644 metadata-ingestion/tests/unit/utilities/test_ratelimiter.py diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py index 545cafca9d4df..1f4f0a0bad9b2 100644 --- a/metadata-ingestion/setup.py +++ b/metadata-ingestion/setup.py @@ -38,7 +38,6 @@ "progressbar2", "termcolor>=1.0.0", "psutil>=5.8.0", - "ratelimiter", "Deprecated", "humanfriendly", "packaging", diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit_log_api.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit_log_api.py index 03b12c61ee5c6..db552c09cd0a7 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit_log_api.py +++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit_log_api.py @@ -4,7 +4,6 @@ from google.cloud import bigquery from google.cloud.logging_v2.client import Client as GCPLoggingClient -from ratelimiter import RateLimiter from datahub.ingestion.source.bigquery_v2.bigquery_audit import ( AuditLogEntry, @@ -17,6 +16,7 @@ BQ_DATE_SHARD_FORMAT, BQ_DATETIME_FORMAT, ) +from datahub.utilities.ratelimiter import RateLimiter logger: logging.Logger = logging.getLogger(__name__) diff --git a/metadata-ingestion/src/datahub/utilities/ratelimiter.py b/metadata-ingestion/src/datahub/utilities/ratelimiter.py new file mode 100644 index 0000000000000..3d47d25e14c49 --- /dev/null +++ b/metadata-ingestion/src/datahub/utilities/ratelimiter.py @@ -0,0 +1,56 @@ +import collections +import threading +import time +from contextlib import AbstractContextManager +from typing import Any, Deque + + +# Modified version of https://github.com/RazerM/ratelimiter/blob/master/ratelimiter/_sync.py +class RateLimiter(AbstractContextManager): + + """Provides rate limiting for an operation with a configurable number of + requests for a time period. + """ + + def __init__(self, max_calls: int, period: float = 1.0) -> None: + """Initialize a RateLimiter object which enforces as much as max_calls + operations on period (eventually floating) number of seconds. + """ + if period <= 0: + raise ValueError("Rate limiting period should be > 0") + if max_calls <= 0: + raise ValueError("Rate limiting number of calls should be > 0") + + # We're using a deque to store the last execution timestamps, not for + # its maxlen attribute, but to allow constant time front removal. + self.calls: Deque = collections.deque() + + self.period = period + self.max_calls = max_calls + self._lock = threading.Lock() + + def __enter__(self) -> "RateLimiter": + with self._lock: + # We want to ensure that no more than max_calls were run in the allowed + # period. For this, we store the last timestamps of each call and run + # the rate verification upon each __enter__ call. + if len(self.calls) >= self.max_calls: + until = time.time() + self.period - self._timespan + sleeptime = until - time.time() + if sleeptime > 0: + time.sleep(sleeptime) + return self + + def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> None: + with self._lock: + # Store the last operation timestamp. + self.calls.append(time.time()) + + # Pop the timestamp list front (ie: the older calls) until the sum goes + # back below the period. This is our 'sliding period' window. + while self._timespan >= self.period: + self.calls.popleft() + + @property + def _timespan(self) -> float: + return self.calls[-1] - self.calls[0] diff --git a/metadata-ingestion/tests/unit/utilities/test_ratelimiter.py b/metadata-ingestion/tests/unit/utilities/test_ratelimiter.py new file mode 100644 index 0000000000000..0384e1f918881 --- /dev/null +++ b/metadata-ingestion/tests/unit/utilities/test_ratelimiter.py @@ -0,0 +1,20 @@ +from collections import defaultdict +from datetime import datetime +from typing import Dict + +from datahub.utilities.ratelimiter import RateLimiter + + +def test_rate_is_limited(): + MAX_CALLS_PER_SEC = 5 + TOTAL_CALLS = 18 + actual_calls: Dict[float, int] = defaultdict(lambda: 0) + + ratelimiter = RateLimiter(max_calls=MAX_CALLS_PER_SEC, period=1) + for _ in range(TOTAL_CALLS): + with ratelimiter: + actual_calls[datetime.now().replace(microsecond=0).timestamp()] += 1 + + assert len(actual_calls) == round(TOTAL_CALLS / MAX_CALLS_PER_SEC) + assert all(calls <= MAX_CALLS_PER_SEC for calls in actual_calls.values()) + assert sum(actual_calls.values()) == TOTAL_CALLS From 9ccd1d4f5da8f3c93cb9aaacdb5de66600c99c99 Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Mon, 16 Oct 2023 14:34:15 -0400 Subject: [PATCH 134/156] build(ingest/redshift): Add sqlglot dependency (#9021) --- metadata-ingestion/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py index 1f4f0a0bad9b2..7be565d51260d 100644 --- a/metadata-ingestion/setup.py +++ b/metadata-ingestion/setup.py @@ -353,7 +353,7 @@ | {"psycopg2-binary", "pymysql>=1.0.2"}, "pulsar": {"requests"}, "redash": {"redash-toolbelt", "sql-metadata"} | sqllineage_lib, - "redshift": sql_common | redshift_common | usage_common | {"redshift-connector"}, + "redshift": sql_common | redshift_common | usage_common | sqlglot_lib | {"redshift-connector"}, "redshift-legacy": sql_common | redshift_common, "redshift-usage-legacy": sql_common | usage_common | redshift_common, "s3": {*s3_base, *data_lake_profiling}, From 6366b63e48d37de883af61fb801632e9a43d6e48 Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Mon, 16 Oct 2023 19:13:23 -0400 Subject: [PATCH 135/156] feat(ingest/teradata): Add option to not use file backed dict for view definitions (#9024) --- .../datahub/ingestion/source/sql/teradata.py | 47 ++++++++----------- 1 file changed, 20 insertions(+), 27 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py b/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py index 6080cf7b371e3..e628e4dbd3446 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py @@ -1,7 +1,7 @@ import logging from dataclasses import dataclass from datetime import datetime -from typing import Iterable, Optional, Set, Union +from typing import Iterable, MutableMapping, Optional, Union # This import verifies that the dependencies are available. import teradatasqlalchemy # noqa: F401 @@ -12,7 +12,6 @@ from datahub.configuration.common import AllowDenyPattern from datahub.configuration.time_window_config import BaseTimeWindowConfig -from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.sql_parsing_builder import SqlParsingBuilder from datahub.ingestion.api.common import PipelineContext from datahub.ingestion.api.decorators import ( @@ -34,11 +33,7 @@ from datahub.ingestion.source.usage.usage_common import BaseUsageConfig from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport from datahub.ingestion.source_report.time_window import BaseTimeWindowReport -from datahub.metadata._schema_classes import ( - MetadataChangeEventClass, - SchemaMetadataClass, - ViewPropertiesClass, -) +from datahub.metadata._schema_classes import SchemaMetadataClass, ViewPropertiesClass from datahub.metadata.com.linkedin.pegasus2avro.schema import ( BytesTypeClass, TimeTypeClass, @@ -112,6 +107,11 @@ class TeradataConfig(BaseTeradataConfig, BaseTimeWindowConfig): description="Generate usage statistic.", ) + use_file_backed_cache: bool = Field( + default=True, + description="Whether to use a file backed cache for the view definitions.", + ) + @platform_name("Teradata") @config_class(TeradataConfig) @@ -142,7 +142,8 @@ class TeradataSource(TwoTierSQLAlchemySource): and "timestamp" >= TIMESTAMP '{start_time}' and "timestamp" < TIMESTAMP '{end_time}' """ - urns: Optional[Set[str]] + + _view_definition_cache: MutableMapping[str, str] def __init__(self, config: TeradataConfig, ctx: PipelineContext): super().__init__(config, ctx, "teradata") @@ -166,7 +167,10 @@ def __init__(self, config: TeradataConfig, ctx: PipelineContext): env=self.config.env, ) - self._view_definition_cache: FileBackedDict[str] = FileBackedDict() + if self.config.use_file_backed_cache: + self._view_definition_cache = FileBackedDict[str]() + else: + self._view_definition_cache = {} @classmethod def create(cls, config_dict, ctx): @@ -249,24 +253,13 @@ def get_metadata_engine(self) -> Engine: def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]: # Add all schemas to the schema resolver for wu in super().get_workunits_internal(): - if isinstance(wu.metadata, MetadataChangeEventClass): - if wu.metadata.proposedSnapshot: - for aspect in wu.metadata.proposedSnapshot.aspects: - if isinstance(aspect, SchemaMetadataClass): - self.schema_resolver.add_schema_metadata( - wu.metadata.proposedSnapshot.urn, - aspect, - ) - break - if isinstance(wu.metadata, MetadataChangeProposalWrapper): - if ( - wu.metadata.entityUrn - and isinstance(wu.metadata.aspect, ViewPropertiesClass) - and wu.metadata.aspect.viewLogic - ): - self._view_definition_cache[ - wu.metadata.entityUrn - ] = wu.metadata.aspect.viewLogic + urn = wu.get_urn() + schema_metadata = wu.get_aspect_of_type(SchemaMetadataClass) + if schema_metadata: + self.schema_resolver.add_schema_metadata(urn, schema_metadata) + view_properties = wu.get_aspect_of_type(ViewPropertiesClass) + if view_properties and self.config.include_view_lineage: + self._view_definition_cache[urn] = view_properties.viewLogic yield wu if self.config.include_view_lineage: From 9fec6024fb177a321860e49f3c9977b41bb9e65f Mon Sep 17 00:00:00 2001 From: Andrew Sikowitz Date: Tue, 17 Oct 2023 09:58:38 -0400 Subject: [PATCH 136/156] feat(ingest/unity-catalog): Support external S3 lineage (#9025) --- .../datahub/ingestion/source/aws/s3_util.py | 11 +++++-- .../source/snowflake/snowflake_lineage_v2.py | 6 ++-- .../datahub/ingestion/source/unity/config.py | 8 +++++ .../datahub/ingestion/source/unity/proxy.py | 8 +++++ .../ingestion/source/unity/proxy_types.py | 31 +++++++++++++++++++ .../datahub/ingestion/source/unity/report.py | 2 ++ .../datahub/ingestion/source/unity/source.py | 23 ++++++++++++++ 7 files changed, 84 insertions(+), 5 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/aws/s3_util.py b/metadata-ingestion/src/datahub/ingestion/source/aws/s3_util.py index 501162455cc45..878b8dd1bb9a5 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/aws/s3_util.py +++ b/metadata-ingestion/src/datahub/ingestion/source/aws/s3_util.py @@ -34,21 +34,26 @@ def get_bucket_relative_path(s3_uri: str) -> str: return "/".join(strip_s3_prefix(s3_uri).split("/")[1:]) -def make_s3_urn(s3_uri: str, env: str) -> str: +def make_s3_urn(s3_uri: str, env: str, remove_extension: bool = True) -> str: s3_name = strip_s3_prefix(s3_uri) if s3_name.endswith("/"): s3_name = s3_name[:-1] name, extension = os.path.splitext(s3_name) - - if extension != "": + if remove_extension and extension != "": extension = extension[1:] # remove the dot return f"urn:li:dataset:(urn:li:dataPlatform:s3,{name}_{extension},{env})" return f"urn:li:dataset:(urn:li:dataPlatform:s3,{s3_name},{env})" +def make_s3_urn_for_lineage(s3_uri: str, env: str) -> str: + # Ideally this is the implementation for all S3 URNs + # Don't feel comfortable changing `make_s3_urn` for glue, sagemaker, and athena + return make_s3_urn(s3_uri, env, remove_extension=False) + + def get_bucket_name(s3_uri: str) -> str: if not is_s3_uri(s3_uri): raise ValueError( diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_lineage_v2.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_lineage_v2.py index 9a993f5774032..0a15c352fc842 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_lineage_v2.py +++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_lineage_v2.py @@ -21,7 +21,7 @@ import datahub.emitter.mce_builder as builder from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.ingestion.api.workunit import MetadataWorkUnit -from datahub.ingestion.source.aws.s3_util import make_s3_urn +from datahub.ingestion.source.aws.s3_util import make_s3_urn_for_lineage from datahub.ingestion.source.snowflake.constants import ( LINEAGE_PERMISSION_ERROR, SnowflakeEdition, @@ -652,7 +652,9 @@ def get_external_upstreams(self, external_lineage: Set[str]) -> List[UpstreamCla # For now, populate only for S3 if external_lineage_entry.startswith("s3://"): external_upstream_table = UpstreamClass( - dataset=make_s3_urn(external_lineage_entry, self.config.env), + dataset=make_s3_urn_for_lineage( + external_lineage_entry, self.config.env + ), type=DatasetLineageTypeClass.COPY, ) external_upstreams.append(external_upstream_table) diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py index a57ee39848855..16820c37d546e 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py @@ -166,6 +166,14 @@ class UnityCatalogSourceConfig( description="Option to enable/disable lineage generation.", ) + include_external_lineage: bool = pydantic.Field( + default=True, + description=( + "Option to enable/disable lineage generation for external tables." + " Only external S3 tables are supported at the moment." + ), + ) + include_notebooks: bool = pydantic.Field( default=False, description="Ingest notebooks, represented as DataHub datasets.", diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py index 9bcdb200f180e..3fb77ce512ed2 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py @@ -33,6 +33,7 @@ ALLOWED_STATEMENT_TYPES, Catalog, Column, + ExternalTableReference, Metastore, Notebook, Query, @@ -248,6 +249,13 @@ def table_lineage(self, table: Table, include_entity_lineage: bool) -> None: ) if table_ref: table.upstreams[table_ref] = {} + elif "fileInfo" in item: + external_ref = ExternalTableReference.create_from_lineage( + item["fileInfo"] + ) + if external_ref: + table.external_upstreams.add(external_ref) + for notebook in item.get("notebookInfos") or []: table.upstream_notebooks.add(notebook["notebook_id"]) diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py index 18ac2475b51e0..315c1c0d20186 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy_types.py @@ -10,6 +10,7 @@ CatalogType, ColumnTypeName, DataSourceFormat, + SecurableType, TableType, ) from databricks.sdk.service.sql import QueryStatementType @@ -176,6 +177,35 @@ def external_path(self) -> str: return f"{self.catalog}/{self.schema}/{self.table}" +@dataclass(frozen=True, order=True) +class ExternalTableReference: + path: str + has_permission: bool + name: Optional[str] + type: Optional[SecurableType] + storage_location: Optional[str] + + @classmethod + def create_from_lineage(cls, d: dict) -> Optional["ExternalTableReference"]: + try: + securable_type: Optional[SecurableType] + try: + securable_type = SecurableType(d.get("securable_type", "").lower()) + except ValueError: + securable_type = None + + return cls( + path=d["path"], + has_permission=d.get("has_permission") or True, + name=d.get("securable_name"), + type=securable_type, + storage_location=d.get("storage_location"), + ) + except Exception as e: + logger.warning(f"Failed to create ExternalTableReference from {d}: {e}") + return None + + @dataclass class Table(CommonProperty): schema: Schema @@ -193,6 +223,7 @@ class Table(CommonProperty): view_definition: Optional[str] properties: Dict[str, str] upstreams: Dict[TableReference, Dict[str, List[str]]] = field(default_factory=dict) + external_upstreams: Set[ExternalTableReference] = field(default_factory=set) upstream_notebooks: Set[NotebookId] = field(default_factory=set) downstream_notebooks: Set[NotebookId] = field(default_factory=set) diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/report.py b/metadata-ingestion/src/datahub/ingestion/source/unity/report.py index fa61571fa92cb..4153d9dd88eb8 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/report.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/report.py @@ -19,6 +19,8 @@ class UnityCatalogReport(IngestionStageReport, StaleEntityRemovalSourceReport): notebooks: EntityFilterReport = EntityFilterReport.field(type="notebook") num_column_lineage_skipped_column_count: int = 0 + num_external_upstreams_lacking_permissions: int = 0 + num_external_upstreams_unsupported: int = 0 num_queries: int = 0 num_queries_dropped_parse_failure: int = 0 diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py index 27c1f341aa84d..b63cf65d55dc8 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py @@ -41,6 +41,7 @@ TestConnectionReport, ) from datahub.ingestion.api.workunit import MetadataWorkUnit +from datahub.ingestion.source.aws.s3_util import make_s3_urn_for_lineage from datahub.ingestion.source.common.subtypes import ( DatasetContainerSubTypes, DatasetSubTypes, @@ -455,6 +456,28 @@ def _generate_lineage_aspect( ) ) + if self.config.include_external_lineage: + for external_ref in table.external_upstreams: + if not external_ref.has_permission or not external_ref.path: + self.report.num_external_upstreams_lacking_permissions += 1 + logger.warning( + f"Lacking permissions for external file upstream on {table.ref}" + ) + elif external_ref.path.startswith("s3://"): + upstreams.append( + UpstreamClass( + dataset=make_s3_urn_for_lineage( + external_ref.path, self.config.env + ), + type=DatasetLineageTypeClass.COPY, + ) + ) + else: + self.report.num_external_upstreams_unsupported += 1 + logger.warning( + f"Unsupported external file upstream on {table.ref}: {external_ref.path}" + ) + if upstreams: return UpstreamLineageClass( upstreams=upstreams, From 10eb205cb8d455639c6d09dcc0c8f3853264f96f Mon Sep 17 00:00:00 2001 From: Tamas Nemeth Date: Tue, 17 Oct 2023 16:16:25 +0200 Subject: [PATCH 137/156] fix(ingest) - Fix file backed collection temp directory removal (#9027) --- .../src/datahub/utilities/file_backed_collections.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/metadata-ingestion/src/datahub/utilities/file_backed_collections.py b/metadata-ingestion/src/datahub/utilities/file_backed_collections.py index c04d2138bc116..18493edded4b7 100644 --- a/metadata-ingestion/src/datahub/utilities/file_backed_collections.py +++ b/metadata-ingestion/src/datahub/utilities/file_backed_collections.py @@ -3,6 +3,7 @@ import logging import pathlib import pickle +import shutil import sqlite3 import tempfile from dataclasses import dataclass, field @@ -56,15 +57,15 @@ class ConnectionWrapper: conn: sqlite3.Connection filename: pathlib.Path - _temp_directory: Optional[tempfile.TemporaryDirectory] + _temp_directory: Optional[str] def __init__(self, filename: Optional[pathlib.Path] = None): self._temp_directory = None # Warning: If filename is provided, the file will not be automatically cleaned up. if not filename: - self._temp_directory = tempfile.TemporaryDirectory() - filename = pathlib.Path(self._temp_directory.name) / _DEFAULT_FILE_NAME + self._temp_directory = tempfile.mkdtemp() + filename = pathlib.Path(self._temp_directory) / _DEFAULT_FILE_NAME self.conn = sqlite3.connect(filename, isolation_level=None) self.conn.row_factory = sqlite3.Row @@ -101,7 +102,8 @@ def executemany( def close(self) -> None: self.conn.close() if self._temp_directory: - self._temp_directory.cleanup() + shutil.rmtree(self._temp_directory) + self._temp_directory = None def __enter__(self) -> "ConnectionWrapper": return self From e7c662a0aca0be97e34bec55161766ea84036ced Mon Sep 17 00:00:00 2001 From: ethan-cartwright Date: Tue, 17 Oct 2023 10:54:07 -0400 Subject: [PATCH 138/156] add dependency level to scrollAcrossLineage search results (#9016) --- datahub-web-react/src/graphql/scroll.graphql | 1 + 1 file changed, 1 insertion(+) diff --git a/datahub-web-react/src/graphql/scroll.graphql b/datahub-web-react/src/graphql/scroll.graphql index 18274c50c2166..1031fed7b9e13 100644 --- a/datahub-web-react/src/graphql/scroll.graphql +++ b/datahub-web-react/src/graphql/scroll.graphql @@ -408,6 +408,7 @@ fragment downloadScrollAcrossLineageResult on ScrollAcrossLineageResults { count total searchResults { + degree entity { ...downloadSearchResults } From ae5fd90c73ff29e00f4b8e20735ce0b72e7b823b Mon Sep 17 00:00:00 2001 From: ethan-cartwright Date: Tue, 17 Oct 2023 10:55:07 -0400 Subject: [PATCH 139/156] add create dataproduct example (#9009) --- .../examples/library/create_dataproduct.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 metadata-ingestion/examples/library/create_dataproduct.py diff --git a/metadata-ingestion/examples/library/create_dataproduct.py b/metadata-ingestion/examples/library/create_dataproduct.py new file mode 100644 index 0000000000000..245395b602480 --- /dev/null +++ b/metadata-ingestion/examples/library/create_dataproduct.py @@ -0,0 +1,25 @@ +from datahub.api.entities.dataproduct.dataproduct import DataProduct +from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph + +gms_endpoint = "http://localhost:8080" +graph = DataHubGraph(DatahubClientConfig(server=gms_endpoint)) + +data_product = DataProduct( + id="pet_of_the_week", + display_name="Pet of the Week Campagin", + domain="urn:li:domain:ef39e99a-9d61-406d-b4a8-c70b16380206", + description="This campaign includes Pet of the Week data.", + assets=[ + "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.pet_details,PROD)", + "urn:li:dashboard:(looker,baz)", + "urn:li:dataFlow:(airflow,dag_abc,PROD)", + ], + owners=[{"id": "urn:li:corpuser:jdoe", "type": "BUSINESS_OWNER"}], + terms=["urn:li:glossaryTerm:ClientsAndAccounts.AccountBalance"], + tags=["urn:li:tag:adoption"], + properties={"lifecycle": "production", "sla": "7am every day"}, + external_url="https://en.wikipedia.org/wiki/Sloth", +) + +for mcp in data_product.generate_mcp(upsert=False): + graph.emit(mcp) From 75108ceb2ff125af52fb1e37f7f6d371a77de3b7 Mon Sep 17 00:00:00 2001 From: Kos Korchak <97058061+kkorchak@users.noreply.github.com> Date: Tue, 17 Oct 2023 14:13:31 -0400 Subject: [PATCH 140/156] Download Lineage Results Cypress Test (#9017) --- .../styled/search/DownloadAsCsvModal.tsx | 2 + .../styled/search/SearchExtendedMenu.tsx | 4 +- .../e2e/lineage/download_lineage_results.js | 80 +++++++++++++++++++ 3 files changed, 84 insertions(+), 2 deletions(-) create mode 100644 smoke-test/tests/cypress/cypress/e2e/lineage/download_lineage_results.js diff --git a/datahub-web-react/src/app/entity/shared/components/styled/search/DownloadAsCsvModal.tsx b/datahub-web-react/src/app/entity/shared/components/styled/search/DownloadAsCsvModal.tsx index 452658583cf61..92e859ee1b329 100644 --- a/datahub-web-react/src/app/entity/shared/components/styled/search/DownloadAsCsvModal.tsx +++ b/datahub-web-react/src/app/entity/shared/components/styled/search/DownloadAsCsvModal.tsx @@ -130,6 +130,7 @@ export default function DownloadAsCsvModal({ Close - 
{`${logs}${!showExpandedLogs && isOutputExpandable ? '...' : ''}`}
- {isOutputExpandable && ( + 
{`${logs}${!showExpandedLogs && areLogsExpandable ? '...' : ''}`}
+ {areLogsExpandable && ( setShowExpandedLogs(!showExpandedLogs)}> {showExpandedLogs ? 'Hide' : 'Show More'} )} + {recipe && ( + + Recipe + + + The recipe used for this ingestion run. + + + + 
{`${recipe}${!showExpandedRecipe && isRecipeExpandable ? '\n...' : ''}`}
+ + {isRecipeExpandable && ( + setShowExpandedRecipe((v) => !v)}> + {showExpandedRecipe ? 'Hide' : 'Show More'} + + )} + + )} ); diff --git a/datahub-web-react/src/app/ingest/source/utils.ts b/datahub-web-react/src/app/ingest/source/utils.ts index c372388e958b7..f789ed8434721 100644 --- a/datahub-web-react/src/app/ingest/source/utils.ts +++ b/datahub-web-react/src/app/ingest/source/utils.ts @@ -1,17 +1,19 @@ -import YAML from 'yamljs'; import { CheckCircleOutlined, ClockCircleOutlined, CloseCircleOutlined, + ExclamationCircleOutlined, LoadingOutlined, + StopOutlined, WarningOutlined, } from '@ant-design/icons'; -import { ANTD_GRAY, REDESIGN_COLORS } from '../../entity/shared/constants'; +import YAML from 'yamljs'; +import { ListIngestionSourcesDocument, ListIngestionSourcesQuery } from '../../../graphql/ingestion.generated'; import { EntityType, FacetMetadata } from '../../../types.generated'; -import { capitalizeFirstLetterOnly, pluralize } from '../../shared/textUtil'; import EntityRegistry from '../../entity/EntityRegistry'; +import { ANTD_GRAY, REDESIGN_COLORS } from '../../entity/shared/constants'; +import { capitalizeFirstLetterOnly, pluralize } from '../../shared/textUtil'; import { SourceConfig } from './builder/types'; -import { ListIngestionSourcesDocument, ListIngestionSourcesQuery } from '../../../graphql/ingestion.generated'; export const getSourceConfigs = (ingestionSources: SourceConfig[], sourceType: string) => { const sourceConfigs = ingestionSources.find((source) => source.name === sourceType); @@ -40,7 +42,9 @@ export function getPlaceholderRecipe(ingestionSources: SourceConfig[], type?: st export const RUNNING = 'RUNNING'; export const SUCCESS = 'SUCCESS'; +export const WARNING = 'WARNING'; export const FAILURE = 'FAILURE'; +export const CONNECTION_FAILURE = 'CONNECTION_FAILURE'; export const CANCELLED = 'CANCELLED'; export const UP_FOR_RETRY = 'UP_FOR_RETRY'; export const ROLLING_BACK = 'ROLLING_BACK'; @@ -56,8 +60,10 @@ export const getExecutionRequestStatusIcon = (status: string) => { return ( (status === RUNNING && LoadingOutlined) || (status === SUCCESS && CheckCircleOutlined) || + (status === WARNING && ExclamationCircleOutlined) || (status === FAILURE && CloseCircleOutlined) || - (status === CANCELLED && CloseCircleOutlined) || + (status === CONNECTION_FAILURE && CloseCircleOutlined) || + (status === CANCELLED && StopOutlined) || (status === UP_FOR_RETRY && ClockCircleOutlined) || (status === ROLLED_BACK && WarningOutlined) || (status === ROLLING_BACK && LoadingOutlined) || @@ -70,7 +76,9 @@ export const getExecutionRequestStatusDisplayText = (status: string) => { return ( (status === RUNNING && 'Running') || (status === SUCCESS && 'Succeeded') || + (status === WARNING && 'Completed') || (status === FAILURE && 'Failed') || + (status === CONNECTION_FAILURE && 'Connection Failed') || (status === CANCELLED && 'Cancelled') || (status === UP_FOR_RETRY && 'Up for Retry') || (status === ROLLED_BACK && 'Rolled Back') || @@ -83,21 +91,25 @@ export const getExecutionRequestStatusDisplayText = (status: string) => { export const getExecutionRequestSummaryText = (status: string) => { switch (status) { case RUNNING: - return 'Ingestion is running'; + return 'Ingestion is running...'; case SUCCESS: - return 'Ingestion successfully completed'; + return 'Ingestion succeeded with no errors or suspected missing data.'; + case WARNING: + return 'Ingestion completed with minor or intermittent errors.'; case FAILURE: - return 'Ingestion completed with errors'; + return 'Ingestion failed to complete, or completed with serious errors.'; + case CONNECTION_FAILURE: + return 'Ingestion failed due to network, authentication, or permission issues.'; case CANCELLED: - return 'Ingestion was cancelled'; + return 'Ingestion was cancelled.'; case ROLLED_BACK: - return 'Ingestion was rolled back'; + return 'Ingestion was rolled back.'; case ROLLING_BACK: - return 'Ingestion is in the process of rolling back'; + return 'Ingestion is in the process of rolling back.'; case ROLLBACK_FAILED: - return 'Ingestion rollback failed'; + return 'Ingestion rollback failed.'; default: - return 'Ingestion status not recognized'; + return 'Ingestion status not recognized.'; } }; @@ -105,7 +117,9 @@ export const getExecutionRequestStatusDisplayColor = (status: string) => { return ( (status === RUNNING && REDESIGN_COLORS.BLUE) || (status === SUCCESS && 'green') || + (status === WARNING && 'orangered') || (status === FAILURE && 'red') || + (status === CONNECTION_FAILURE && 'crimson') || (status === UP_FOR_RETRY && 'orange') || (status === CANCELLED && ANTD_GRAY[9]) || (status === ROLLED_BACK && 'orange') || diff --git a/datahub-web-react/src/graphql/ingestion.graphql b/datahub-web-react/src/graphql/ingestion.graphql index 80f66642fe11f..c127e9ec03f9a 100644 --- a/datahub-web-react/src/graphql/ingestion.graphql +++ b/datahub-web-react/src/graphql/ingestion.graphql @@ -90,6 +90,10 @@ query getIngestionExecutionRequest($urn: String!) { source { type } + arguments { + key + value + } } result { status From 1b737243b266843136918ec92f6d20573b999272 Mon Sep 17 00:00:00 2001 From: RyanHolstien Date: Wed, 18 Oct 2023 13:45:46 -0500 Subject: [PATCH 151/156] feat(avro): upgrade avro to 1.11 (#9031) --- build.gradle | 7 +++---- buildSrc/build.gradle | 9 ++++++++- docker/datahub-frontend/start.sh | 1 + metadata-dao-impl/kafka-producer/build.gradle | 4 ++-- metadata-events/{mxe-avro-1.7 => mxe-avro}/.gitignore | 0 metadata-events/{mxe-avro-1.7 => mxe-avro}/build.gradle | 6 +++--- metadata-events/mxe-registration/build.gradle | 2 +- metadata-events/mxe-schemas/build.gradle | 2 +- .../{mxe-utils-avro-1.7 => mxe-utils-avro}/.gitignore | 0 .../{mxe-utils-avro-1.7 => mxe-utils-avro}/build.gradle | 2 +- .../src/main/java/com/linkedin/metadata/EventUtils.java | 0 .../test/java/com/linkedin/metadata/EventUtilsTests.java | 0 .../src/test/resources/test-avro2pegasus-mae.json | 0 .../src/test/resources/test-avro2pegasus-mce.json | 0 .../src/test/resources/test-pegasus2avro-fmce.json | 0 .../src/test/resources/test-pegasus2avro-mae.json | 0 .../src/test/resources/test-pegasus2avro-mce.json | 0 metadata-integration/java/datahub-client/build.gradle | 2 +- .../main/java/datahub/client/kafka/AvroSerializer.java | 4 +++- metadata-io/build.gradle | 4 ++-- metadata-jobs/mae-consumer/build.gradle | 4 ++-- metadata-jobs/mce-consumer/build.gradle | 4 ++-- metadata-jobs/pe-consumer/build.gradle | 4 ++-- metadata-service/restli-servlet-impl/build.gradle | 2 +- metadata-service/services/build.gradle | 4 ++-- metadata-utils/build.gradle | 6 +++--- settings.gradle | 4 ++-- 27 files changed, 40 insertions(+), 31 deletions(-) rename metadata-events/{mxe-avro-1.7 => mxe-avro}/.gitignore (100%) rename metadata-events/{mxe-avro-1.7 => mxe-avro}/build.gradle (81%) rename metadata-events/{mxe-utils-avro-1.7 => mxe-utils-avro}/.gitignore (100%) rename metadata-events/{mxe-utils-avro-1.7 => mxe-utils-avro}/build.gradle (95%) rename metadata-events/{mxe-utils-avro-1.7 => mxe-utils-avro}/src/main/java/com/linkedin/metadata/EventUtils.java (100%) rename metadata-events/{mxe-utils-avro-1.7 => mxe-utils-avro}/src/test/java/com/linkedin/metadata/EventUtilsTests.java (100%) rename metadata-events/{mxe-utils-avro-1.7 => mxe-utils-avro}/src/test/resources/test-avro2pegasus-mae.json (100%) rename metadata-events/{mxe-utils-avro-1.7 => mxe-utils-avro}/src/test/resources/test-avro2pegasus-mce.json (100%) rename metadata-events/{mxe-utils-avro-1.7 => mxe-utils-avro}/src/test/resources/test-pegasus2avro-fmce.json (100%) rename metadata-events/{mxe-utils-avro-1.7 => mxe-utils-avro}/src/test/resources/test-pegasus2avro-mae.json (100%) rename metadata-events/{mxe-utils-avro-1.7 => mxe-utils-avro}/src/test/resources/test-pegasus2avro-mce.json (100%) diff --git a/build.gradle b/build.gradle index 025c588da2b52..cf55a59cfe694 100644 --- a/build.gradle +++ b/build.gradle @@ -27,7 +27,7 @@ buildscript { dependencies { classpath 'com.linkedin.pegasus:gradle-plugins:' + pegasusVersion classpath 'com.github.node-gradle:gradle-node-plugin:2.2.4' - classpath 'io.acryl.gradle.plugin:gradle-avro-plugin:0.8.1' + classpath 'io.acryl.gradle.plugin:gradle-avro-plugin:0.2.0' classpath 'org.springframework.boot:spring-boot-gradle-plugin:' + springBootVersion classpath "io.codearte.gradle.nexus:gradle-nexus-staging-plugin:0.30.0" classpath "com.palantir.gradle.gitversion:gradle-git-version:3.0.0" @@ -67,8 +67,8 @@ project.ext.externalDependency = [ 'antlr4Runtime': 'org.antlr:antlr4-runtime:4.7.2', 'antlr4': 'org.antlr:antlr4:4.7.2', 'assertJ': 'org.assertj:assertj-core:3.11.1', - 'avro_1_7': 'org.apache.avro:avro:1.7.7', - 'avroCompiler_1_7': 'org.apache.avro:avro-compiler:1.7.7', + 'avro': 'org.apache.avro:avro:1.11.3', + 'avroCompiler': 'org.apache.avro:avro-compiler:1.11.3', 'awsGlueSchemaRegistrySerde': 'software.amazon.glue:schema-registry-serde:1.1.10', 'awsMskIamAuth': 'software.amazon.msk:aws-msk-iam-auth:1.1.1', 'awsSecretsManagerJdbc': 'com.amazonaws.secretsmanager:aws-secretsmanager-jdbc:1.0.8', @@ -127,7 +127,6 @@ project.ext.externalDependency = [ 'jgrapht': 'org.jgrapht:jgrapht-core:1.5.1', 'jna': 'net.java.dev.jna:jna:5.12.1', 'jsonPatch': 'com.github.java-json-tools:json-patch:1.13', - 'jsonSchemaAvro': 'com.github.fge:json-schema-avro:0.1.4', 'jsonSimple': 'com.googlecode.json-simple:json-simple:1.1.1', 'jsonSmart': 'net.minidev:json-smart:2.4.9', 'json': 'org.json:json:20230227', diff --git a/buildSrc/build.gradle b/buildSrc/build.gradle index 65b3780431db9..1f9d30d520171 100644 --- a/buildSrc/build.gradle +++ b/buildSrc/build.gradle @@ -5,7 +5,14 @@ buildscript { } dependencies { - implementation('io.acryl:json-schema-avro:0.1.5') { + /** + * Forked version of abandoned repository: https://github.com/fge/json-schema-avro + * Maintainer last active 2014, we maintain an active fork of this repository to utilize mapping Avro schemas to Json Schemas, + * repository is as close to official library for this as you can get. Original maintainer is one of the authors of Json Schema spec. + * Other companies are also separately maintaining forks (like: https://github.com/java-json-tools/json-schema-avro). + * We have built several customizations on top of it for various bug fixes, especially around union scheams + */ + implementation('io.acryl:json-schema-avro:0.2.2') { exclude group: 'com.fasterxml.jackson.core', module: 'jackson-databind' exclude group: 'com.google.guava', module: 'guava' } diff --git a/docker/datahub-frontend/start.sh b/docker/datahub-frontend/start.sh index 9dc1514144bb1..430982aa2456b 100755 --- a/docker/datahub-frontend/start.sh +++ b/docker/datahub-frontend/start.sh @@ -50,6 +50,7 @@ export JAVA_OPTS="-Xms512m \ -Djava.security.auth.login.config=datahub-frontend/conf/jaas.conf \ -Dlogback.configurationFile=datahub-frontend/conf/logback.xml \ -Dlogback.debug=false \ + -agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=5005 \ ${PROMETHEUS_AGENT:-} ${OTEL_AGENT:-} \ ${TRUSTSTORE_FILE:-} ${TRUSTSTORE_TYPE:-} ${TRUSTSTORE_PASSWORD:-} \ ${HTTP_PROXY:-} ${HTTPS_PROXY:-} ${NO_PROXY:-} \ diff --git a/metadata-dao-impl/kafka-producer/build.gradle b/metadata-dao-impl/kafka-producer/build.gradle index 393b10b0e9d24..bc3415b2ccc8c 100644 --- a/metadata-dao-impl/kafka-producer/build.gradle +++ b/metadata-dao-impl/kafka-producer/build.gradle @@ -1,9 +1,9 @@ apply plugin: 'java' dependencies { - implementation project(':metadata-events:mxe-avro-1.7') + implementation project(':metadata-events:mxe-avro') implementation project(':metadata-events:mxe-registration') - implementation project(':metadata-events:mxe-utils-avro-1.7') + implementation project(':metadata-events:mxe-utils-avro') implementation project(':entity-registry') implementation project(':metadata-io') diff --git a/metadata-events/mxe-avro-1.7/.gitignore b/metadata-events/mxe-avro/.gitignore similarity index 100% rename from metadata-events/mxe-avro-1.7/.gitignore rename to metadata-events/mxe-avro/.gitignore diff --git a/metadata-events/mxe-avro-1.7/build.gradle b/metadata-events/mxe-avro/build.gradle similarity index 81% rename from metadata-events/mxe-avro-1.7/build.gradle rename to metadata-events/mxe-avro/build.gradle index 8c0a26d22dc7d..9d11eeb160ff0 100644 --- a/metadata-events/mxe-avro-1.7/build.gradle +++ b/metadata-events/mxe-avro/build.gradle @@ -6,8 +6,8 @@ apply plugin: 'io.acryl.gradle.plugin.avro' apply plugin: 'java-library' dependencies { - api externalDependency.avro_1_7 - implementation(externalDependency.avroCompiler_1_7) { + api externalDependency.avro + implementation(externalDependency.avroCompiler) { exclude group: 'org.apache.velocity', module: 'velocity' } constraints { @@ -21,7 +21,7 @@ dependencies { def genDir = file("src/generated/java") -task avroCodeGen(type: com.commercehub.gradle.plugin.avro.GenerateAvroJavaTask, dependsOn: configurations.avsc) { +task avroCodeGen(type: com.github.davidmc24.gradle.plugin.avro.GenerateAvroJavaTask, dependsOn: configurations.avsc) { source("$rootDir/metadata-events/mxe-schemas/src/renamed/avro") outputDir = genDir dependsOn(':metadata-events:mxe-schemas:renameNamespace') diff --git a/metadata-events/mxe-registration/build.gradle b/metadata-events/mxe-registration/build.gradle index 60e0da59616d9..032870d93329f 100644 --- a/metadata-events/mxe-registration/build.gradle +++ b/metadata-events/mxe-registration/build.gradle @@ -5,7 +5,7 @@ configurations { } dependencies { - implementation project(':metadata-events:mxe-avro-1.7') + implementation project(':metadata-events:mxe-avro') implementation project(':metadata-models') implementation spec.product.pegasus.dataAvro1_6 diff --git a/metadata-events/mxe-schemas/build.gradle b/metadata-events/mxe-schemas/build.gradle index fe46601fb68b7..8dc8b71bd1cd8 100644 --- a/metadata-events/mxe-schemas/build.gradle +++ b/metadata-events/mxe-schemas/build.gradle @@ -1,4 +1,4 @@ -apply plugin: 'java' +apply plugin: 'java-library' apply plugin: 'pegasus' dependencies { diff --git a/metadata-events/mxe-utils-avro-1.7/.gitignore b/metadata-events/mxe-utils-avro/.gitignore similarity index 100% rename from metadata-events/mxe-utils-avro-1.7/.gitignore rename to metadata-events/mxe-utils-avro/.gitignore diff --git a/metadata-events/mxe-utils-avro-1.7/build.gradle b/metadata-events/mxe-utils-avro/build.gradle similarity index 95% rename from metadata-events/mxe-utils-avro-1.7/build.gradle rename to metadata-events/mxe-utils-avro/build.gradle index 3b137965d6c19..a7bf287ab224d 100644 --- a/metadata-events/mxe-utils-avro-1.7/build.gradle +++ b/metadata-events/mxe-utils-avro/build.gradle @@ -1,7 +1,7 @@ apply plugin: 'java-library' dependencies { - api project(':metadata-events:mxe-avro-1.7') + api project(':metadata-events:mxe-avro') api project(':metadata-models') api spec.product.pegasus.dataAvro1_6 diff --git a/metadata-events/mxe-utils-avro-1.7/src/main/java/com/linkedin/metadata/EventUtils.java b/metadata-events/mxe-utils-avro/src/main/java/com/linkedin/metadata/EventUtils.java similarity index 100% rename from metadata-events/mxe-utils-avro-1.7/src/main/java/com/linkedin/metadata/EventUtils.java rename to metadata-events/mxe-utils-avro/src/main/java/com/linkedin/metadata/EventUtils.java diff --git a/metadata-events/mxe-utils-avro-1.7/src/test/java/com/linkedin/metadata/EventUtilsTests.java b/metadata-events/mxe-utils-avro/src/test/java/com/linkedin/metadata/EventUtilsTests.java similarity index 100% rename from metadata-events/mxe-utils-avro-1.7/src/test/java/com/linkedin/metadata/EventUtilsTests.java rename to metadata-events/mxe-utils-avro/src/test/java/com/linkedin/metadata/EventUtilsTests.java diff --git a/metadata-events/mxe-utils-avro-1.7/src/test/resources/test-avro2pegasus-mae.json b/metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mae.json similarity index 100% rename from metadata-events/mxe-utils-avro-1.7/src/test/resources/test-avro2pegasus-mae.json rename to metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mae.json diff --git a/metadata-events/mxe-utils-avro-1.7/src/test/resources/test-avro2pegasus-mce.json b/metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mce.json similarity index 100% rename from metadata-events/mxe-utils-avro-1.7/src/test/resources/test-avro2pegasus-mce.json rename to metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mce.json diff --git a/metadata-events/mxe-utils-avro-1.7/src/test/resources/test-pegasus2avro-fmce.json b/metadata-events/mxe-utils-avro/src/test/resources/test-pegasus2avro-fmce.json similarity index 100% rename from metadata-events/mxe-utils-avro-1.7/src/test/resources/test-pegasus2avro-fmce.json rename to metadata-events/mxe-utils-avro/src/test/resources/test-pegasus2avro-fmce.json diff --git a/metadata-events/mxe-utils-avro-1.7/src/test/resources/test-pegasus2avro-mae.json b/metadata-events/mxe-utils-avro/src/test/resources/test-pegasus2avro-mae.json similarity index 100% rename from metadata-events/mxe-utils-avro-1.7/src/test/resources/test-pegasus2avro-mae.json rename to metadata-events/mxe-utils-avro/src/test/resources/test-pegasus2avro-mae.json diff --git a/metadata-events/mxe-utils-avro-1.7/src/test/resources/test-pegasus2avro-mce.json b/metadata-events/mxe-utils-avro/src/test/resources/test-pegasus2avro-mce.json similarity index 100% rename from metadata-events/mxe-utils-avro-1.7/src/test/resources/test-pegasus2avro-mce.json rename to metadata-events/mxe-utils-avro/src/test/resources/test-pegasus2avro-mce.json diff --git a/metadata-integration/java/datahub-client/build.gradle b/metadata-integration/java/datahub-client/build.gradle index 95de3cdb3c526..e6210f1f073f6 100644 --- a/metadata-integration/java/datahub-client/build.gradle +++ b/metadata-integration/java/datahub-client/build.gradle @@ -30,7 +30,7 @@ dependencies { implementation(externalDependency.kafkaAvroSerializer) { exclude group: "org.apache.avro" } - implementation externalDependency.avro_1_7 + implementation externalDependency.avro constraints { implementation('commons-collections:commons-collections:3.2.2') { because 'Vulnerability Issue' diff --git a/metadata-integration/java/datahub-client/src/main/java/datahub/client/kafka/AvroSerializer.java b/metadata-integration/java/datahub-client/src/main/java/datahub/client/kafka/AvroSerializer.java index ee0d459aaa7d3..6212e57470be4 100644 --- a/metadata-integration/java/datahub-client/src/main/java/datahub/client/kafka/AvroSerializer.java +++ b/metadata-integration/java/datahub-client/src/main/java/datahub/client/kafka/AvroSerializer.java @@ -16,12 +16,14 @@ class AvroSerializer { private final Schema _recordSchema; private final Schema _genericAspectSchema; + private final Schema _changeTypeEnumSchema; private final EventFormatter _eventFormatter; public AvroSerializer() throws IOException { _recordSchema = new Schema.Parser() .parse(this.getClass().getClassLoader().getResourceAsStream("MetadataChangeProposal.avsc")); _genericAspectSchema = this._recordSchema.getField("aspect").schema().getTypes().get(1); + _changeTypeEnumSchema = this._recordSchema.getField("changeType").schema(); _eventFormatter = new EventFormatter(EventFormatter.Format.PEGASUS_JSON); } @@ -43,7 +45,7 @@ public GenericRecord serialize(MetadataChangeProposal mcp) throws IOException { genericRecord.put("aspect", genericAspect); genericRecord.put("aspectName", mcp.getAspectName()); genericRecord.put("entityType", mcp.getEntityType()); - genericRecord.put("changeType", mcp.getChangeType()); + genericRecord.put("changeType", new GenericData.EnumSymbol(_changeTypeEnumSchema, mcp.getChangeType())); return genericRecord; } } \ No newline at end of file diff --git a/metadata-io/build.gradle b/metadata-io/build.gradle index ad54cf6524398..740fed61f13d5 100644 --- a/metadata-io/build.gradle +++ b/metadata-io/build.gradle @@ -8,9 +8,9 @@ configurations { dependencies { implementation project(':entity-registry') api project(':metadata-utils') - api project(':metadata-events:mxe-avro-1.7') + api project(':metadata-events:mxe-avro') api project(':metadata-events:mxe-registration') - api project(':metadata-events:mxe-utils-avro-1.7') + api project(':metadata-events:mxe-utils-avro') api project(':metadata-models') api project(':metadata-service:restli-client') api project(':metadata-service:configuration') diff --git a/metadata-jobs/mae-consumer/build.gradle b/metadata-jobs/mae-consumer/build.gradle index d36fd0de40d03..fcb8b62e4ac9d 100644 --- a/metadata-jobs/mae-consumer/build.gradle +++ b/metadata-jobs/mae-consumer/build.gradle @@ -21,9 +21,9 @@ dependencies { implementation project(':ingestion-scheduler') implementation project(':metadata-utils') implementation project(":entity-registry") - implementation project(':metadata-events:mxe-avro-1.7') + implementation project(':metadata-events:mxe-avro') implementation project(':metadata-events:mxe-registration') - implementation project(':metadata-events:mxe-utils-avro-1.7') + implementation project(':metadata-events:mxe-utils-avro') implementation project(':datahub-graphql-core') implementation externalDependency.elasticSearchRest diff --git a/metadata-jobs/mce-consumer/build.gradle b/metadata-jobs/mce-consumer/build.gradle index 0bca55e0e5f92..97eec9fcff051 100644 --- a/metadata-jobs/mce-consumer/build.gradle +++ b/metadata-jobs/mce-consumer/build.gradle @@ -17,9 +17,9 @@ dependencies { } implementation project(':metadata-utils') implementation project(':metadata-events:mxe-schemas') - implementation project(':metadata-events:mxe-avro-1.7') + implementation project(':metadata-events:mxe-avro') implementation project(':metadata-events:mxe-registration') - implementation project(':metadata-events:mxe-utils-avro-1.7') + implementation project(':metadata-events:mxe-utils-avro') implementation project(':metadata-io') implementation project(':metadata-service:restli-client') implementation spec.product.pegasus.restliClient diff --git a/metadata-jobs/pe-consumer/build.gradle b/metadata-jobs/pe-consumer/build.gradle index 1899a4de15635..81e8b8c9971f0 100644 --- a/metadata-jobs/pe-consumer/build.gradle +++ b/metadata-jobs/pe-consumer/build.gradle @@ -10,9 +10,9 @@ configurations { dependencies { avro project(path: ':metadata-models', configuration: 'avroSchema') implementation project(':li-utils') - implementation project(':metadata-events:mxe-avro-1.7') + implementation project(':metadata-events:mxe-avro') implementation project(':metadata-events:mxe-registration') - implementation project(':metadata-events:mxe-utils-avro-1.7') + implementation project(':metadata-events:mxe-utils-avro') implementation(project(':metadata-service:factories')) { exclude group: 'org.neo4j.test' } diff --git a/metadata-service/restli-servlet-impl/build.gradle b/metadata-service/restli-servlet-impl/build.gradle index cb307863748c3..de6fb6690e693 100644 --- a/metadata-service/restli-servlet-impl/build.gradle +++ b/metadata-service/restli-servlet-impl/build.gradle @@ -48,7 +48,7 @@ dependencies { implementation externalDependency.dropwizardMetricsCore implementation externalDependency.dropwizardMetricsJmx - compileOnly externalDependency.lombok + implementation externalDependency.lombok implementation externalDependency.neo4jJavaDriver implementation externalDependency.opentelemetryAnnotations diff --git a/metadata-service/services/build.gradle b/metadata-service/services/build.gradle index 22c62af324c12..b6af3d330d185 100644 --- a/metadata-service/services/build.gradle +++ b/metadata-service/services/build.gradle @@ -9,9 +9,9 @@ dependencies { implementation externalDependency.jsonPatch implementation project(':entity-registry') implementation project(':metadata-utils') - implementation project(':metadata-events:mxe-avro-1.7') + implementation project(':metadata-events:mxe-avro') implementation project(':metadata-events:mxe-registration') - implementation project(':metadata-events:mxe-utils-avro-1.7') + implementation project(':metadata-events:mxe-utils-avro') implementation project(':metadata-models') implementation project(':metadata-service:restli-client') implementation project(':metadata-service:configuration') diff --git a/metadata-utils/build.gradle b/metadata-utils/build.gradle index 1c1c368611488..7bc6aa2d43442 100644 --- a/metadata-utils/build.gradle +++ b/metadata-utils/build.gradle @@ -1,7 +1,7 @@ apply plugin: 'java-library' dependencies { - api externalDependency.avro_1_7 + api externalDependency.avro implementation externalDependency.commonsLang api externalDependency.dropwizardMetricsCore implementation externalDependency.dropwizardMetricsJmx @@ -16,8 +16,8 @@ dependencies { api project(':li-utils') api project(':entity-registry') - api project(':metadata-events:mxe-avro-1.7') - api project(':metadata-events:mxe-utils-avro-1.7') + api project(':metadata-events:mxe-avro') + api project(':metadata-events:mxe-utils-avro') implementation externalDependency.slf4jApi compileOnly externalDependency.lombok diff --git a/settings.gradle b/settings.gradle index d6777b07b3fb3..52de461383b5e 100644 --- a/settings.gradle +++ b/settings.gradle @@ -20,10 +20,10 @@ include 'metadata-service:openapi-analytics-servlet' include 'metadata-service:plugin' include 'metadata-service:plugin:src:test:sample-test-plugins' include 'metadata-dao-impl:kafka-producer' -include 'metadata-events:mxe-avro-1.7' +include 'metadata-events:mxe-avro' include 'metadata-events:mxe-registration' include 'metadata-events:mxe-schemas' -include 'metadata-events:mxe-utils-avro-1.7' +include 'metadata-events:mxe-utils-avro' include 'metadata-ingestion' include 'metadata-jobs:mae-consumer' include 'metadata-jobs:mce-consumer' From aae1347efce9edf1b5c4512ba3c72569e165947d Mon Sep 17 00:00:00 2001 From: Indy Prentice Date: Wed, 18 Oct 2023 16:26:24 -0300 Subject: [PATCH 152/156] fix(search): Detect field type for use in defining the sort order (#8992) Co-authored-by: Indy Prentice --- .../indexbuilder/MappingsBuilder.java | 48 +++++------- .../query/request/SearchRequestHandler.java | 8 +- .../metadata/search/utils/ESUtils.java | 74 ++++++++++++++++++- .../fixtures/SampleDataFixtureTestBase.java | 64 ++++++++++++++-- 4 files changed, 154 insertions(+), 40 deletions(-) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/MappingsBuilder.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/MappingsBuilder.java index 004b2e0a2adc4..1edc77bbd214c 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/MappingsBuilder.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/MappingsBuilder.java @@ -5,6 +5,7 @@ import com.linkedin.metadata.models.SearchScoreFieldSpec; import com.linkedin.metadata.models.SearchableFieldSpec; import com.linkedin.metadata.models.annotation.SearchableAnnotation.FieldType; +import com.linkedin.metadata.search.utils.ESUtils; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -31,15 +32,6 @@ public static Map getPartialNgramConfigWithOverrides(Map KEYWORD_TYPE_MAP = ImmutableMap.of(TYPE, KEYWORD); - // Field Types - public static final String BOOLEAN = "boolean"; - public static final String DATE = "date"; - public static final String DOUBLE = "double"; - public static final String LONG = "long"; - public static final String OBJECT = "object"; - public static final String TEXT = "text"; - public static final String TOKEN_COUNT = "token_count"; - // Subfields public static final String DELIMITED = "delimited"; public static final String LENGTH = "length"; @@ -74,7 +66,7 @@ public static Map getMappings(@Nonnull final EntitySpec entitySp private static Map getMappingsForUrn() { Map subFields = new HashMap<>(); subFields.put(DELIMITED, ImmutableMap.of( - TYPE, TEXT, + TYPE, ESUtils.TEXT_FIELD_TYPE, ANALYZER, URN_ANALYZER, SEARCH_ANALYZER, URN_SEARCH_ANALYZER, SEARCH_QUOTE_ANALYZER, CUSTOM_QUOTE_ANALYZER) @@ -85,13 +77,13 @@ private static Map getMappingsForUrn() { ) )); return ImmutableMap.builder() - .put(TYPE, KEYWORD) + .put(TYPE, ESUtils.KEYWORD_FIELD_TYPE) .put(FIELDS, subFields) .build(); } private static Map getMappingsForRunId() { - return ImmutableMap.builder().put(TYPE, KEYWORD).build(); + return ImmutableMap.builder().put(TYPE, ESUtils.KEYWORD_FIELD_TYPE).build(); } private static Map getMappingsForField(@Nonnull final SearchableFieldSpec searchableFieldSpec) { @@ -104,23 +96,23 @@ private static Map getMappingsForField(@Nonnull final Searchable } else if (fieldType == FieldType.TEXT || fieldType == FieldType.TEXT_PARTIAL || fieldType == FieldType.WORD_GRAM) { mappingForField.putAll(getMappingsForSearchText(fieldType)); } else if (fieldType == FieldType.BROWSE_PATH) { - mappingForField.put(TYPE, TEXT); + mappingForField.put(TYPE, ESUtils.TEXT_FIELD_TYPE); mappingForField.put(FIELDS, ImmutableMap.of(LENGTH, ImmutableMap.of( - TYPE, TOKEN_COUNT, + TYPE, ESUtils.TOKEN_COUNT_FIELD_TYPE, ANALYZER, SLASH_PATTERN_ANALYZER))); mappingForField.put(ANALYZER, BROWSE_PATH_HIERARCHY_ANALYZER); mappingForField.put(FIELDDATA, true); } else if (fieldType == FieldType.BROWSE_PATH_V2) { - mappingForField.put(TYPE, TEXT); + mappingForField.put(TYPE, ESUtils.TEXT_FIELD_TYPE); mappingForField.put(FIELDS, ImmutableMap.of(LENGTH, ImmutableMap.of( - TYPE, TOKEN_COUNT, + TYPE, ESUtils.TOKEN_COUNT_FIELD_TYPE, ANALYZER, UNIT_SEPARATOR_PATTERN_ANALYZER))); mappingForField.put(ANALYZER, BROWSE_PATH_V2_HIERARCHY_ANALYZER); mappingForField.put(FIELDDATA, true); } else if (fieldType == FieldType.URN || fieldType == FieldType.URN_PARTIAL) { - mappingForField.put(TYPE, TEXT); + mappingForField.put(TYPE, ESUtils.TEXT_FIELD_TYPE); mappingForField.put(ANALYZER, URN_ANALYZER); mappingForField.put(SEARCH_ANALYZER, URN_SEARCH_ANALYZER); mappingForField.put(SEARCH_QUOTE_ANALYZER, CUSTOM_QUOTE_ANALYZER); @@ -135,13 +127,13 @@ private static Map getMappingsForField(@Nonnull final Searchable subFields.put(KEYWORD, KEYWORD_TYPE_MAP); mappingForField.put(FIELDS, subFields); } else if (fieldType == FieldType.BOOLEAN) { - mappingForField.put(TYPE, BOOLEAN); + mappingForField.put(TYPE, ESUtils.BOOLEAN_FIELD_TYPE); } else if (fieldType == FieldType.COUNT) { - mappingForField.put(TYPE, LONG); + mappingForField.put(TYPE, ESUtils.LONG_FIELD_TYPE); } else if (fieldType == FieldType.DATETIME) { - mappingForField.put(TYPE, DATE); + mappingForField.put(TYPE, ESUtils.DATE_FIELD_TYPE); } else if (fieldType == FieldType.OBJECT) { - mappingForField.put(TYPE, OBJECT); + mappingForField.put(TYPE, ESUtils.DATE_FIELD_TYPE); } else { log.info("FieldType {} has no mappings implemented", fieldType); } @@ -149,10 +141,10 @@ private static Map getMappingsForField(@Nonnull final Searchable searchableFieldSpec.getSearchableAnnotation() .getHasValuesFieldName() - .ifPresent(fieldName -> mappings.put(fieldName, ImmutableMap.of(TYPE, BOOLEAN))); + .ifPresent(fieldName -> mappings.put(fieldName, ImmutableMap.of(TYPE, ESUtils.BOOLEAN_FIELD_TYPE))); searchableFieldSpec.getSearchableAnnotation() .getNumValuesFieldName() - .ifPresent(fieldName -> mappings.put(fieldName, ImmutableMap.of(TYPE, LONG))); + .ifPresent(fieldName -> mappings.put(fieldName, ImmutableMap.of(TYPE, ESUtils.LONG_FIELD_TYPE))); mappings.putAll(getMappingsForFieldNameAliases(searchableFieldSpec)); return mappings; @@ -160,7 +152,7 @@ private static Map getMappingsForField(@Nonnull final Searchable private static Map getMappingsForKeyword() { Map mappingForField = new HashMap<>(); - mappingForField.put(TYPE, KEYWORD); + mappingForField.put(TYPE, ESUtils.KEYWORD_FIELD_TYPE); mappingForField.put(NORMALIZER, KEYWORD_NORMALIZER); // Add keyword subfield without lowercase filter mappingForField.put(FIELDS, ImmutableMap.of(KEYWORD, KEYWORD_TYPE_MAP)); @@ -169,7 +161,7 @@ private static Map getMappingsForKeyword() { private static Map getMappingsForSearchText(FieldType fieldType) { Map mappingForField = new HashMap<>(); - mappingForField.put(TYPE, KEYWORD); + mappingForField.put(TYPE, ESUtils.KEYWORD_FIELD_TYPE); mappingForField.put(NORMALIZER, KEYWORD_NORMALIZER); Map subFields = new HashMap<>(); if (fieldType == FieldType.TEXT_PARTIAL || fieldType == FieldType.WORD_GRAM) { @@ -186,14 +178,14 @@ private static Map getMappingsForSearchText(FieldType fieldType) String fieldName = entry.getKey(); String analyzerName = entry.getValue(); subFields.put(fieldName, ImmutableMap.of( - TYPE, TEXT, + TYPE, ESUtils.TEXT_FIELD_TYPE, ANALYZER, analyzerName )); } } } subFields.put(DELIMITED, ImmutableMap.of( - TYPE, TEXT, + TYPE, ESUtils.TEXT_FIELD_TYPE, ANALYZER, TEXT_ANALYZER, SEARCH_ANALYZER, TEXT_SEARCH_ANALYZER, SEARCH_QUOTE_ANALYZER, CUSTOM_QUOTE_ANALYZER)); @@ -206,7 +198,7 @@ private static Map getMappingsForSearchText(FieldType fieldType) private static Map getMappingsForSearchScoreField( @Nonnull final SearchScoreFieldSpec searchScoreFieldSpec) { return ImmutableMap.of(searchScoreFieldSpec.getSearchScoreAnnotation().getFieldName(), - ImmutableMap.of(TYPE, DOUBLE)); + ImmutableMap.of(TYPE, ESUtils.DOUBLE_FIELD_TYPE)); } private static Map getMappingsForFieldNameAliases(@Nonnull final SearchableFieldSpec searchableFieldSpec) { diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java index 5fcc10b7af5cf..c06907e800d5e 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java @@ -202,7 +202,7 @@ public SearchRequest getSearchRequest(@Nonnull String input, @Nullable Filter fi if (!finalSearchFlags.isSkipHighlighting()) { searchSourceBuilder.highlighter(_highlights); } - ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion); + ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, _entitySpecs); if (finalSearchFlags.isGetSuggestions()) { ESUtils.buildNameSuggestions(searchSourceBuilder, input); @@ -243,7 +243,7 @@ public SearchRequest getSearchRequest(@Nonnull String input, @Nullable Filter fi searchSourceBuilder.query(QueryBuilders.boolQuery().must(getQuery(input, finalSearchFlags.isFulltext())).filter(filterQuery)); _aggregationQueryBuilder.getAggregations().forEach(searchSourceBuilder::aggregation); searchSourceBuilder.highlighter(getHighlights()); - ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion); + ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, _entitySpecs); searchRequest.source(searchSourceBuilder); log.debug("Search request is: " + searchRequest); searchRequest.indicesOptions(null); @@ -270,7 +270,7 @@ public SearchRequest getFilterRequest(@Nullable Filter filters, @Nullable SortCr final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchSourceBuilder.query(filterQuery); searchSourceBuilder.from(from).size(size); - ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion); + ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, _entitySpecs); searchRequest.source(searchSourceBuilder); return searchRequest; @@ -301,7 +301,7 @@ public SearchRequest getFilterRequest(@Nullable Filter filters, @Nullable SortCr searchSourceBuilder.size(size); ESUtils.setSearchAfter(searchSourceBuilder, sort, pitId, keepAlive); - ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion); + ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, _entitySpecs); searchRequest.source(searchSourceBuilder); return searchRequest; diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java index 9a7d9a1b4c420..53765acb8e29e 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java @@ -2,6 +2,9 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.metadata.models.SearchableFieldSpec; +import com.linkedin.metadata.models.annotation.SearchableAnnotation; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.Criterion; @@ -49,7 +52,28 @@ public class ESUtils { public static final int MAX_RESULT_SIZE = 10000; public static final String OPAQUE_ID_HEADER = "X-Opaque-Id"; public static final String HEADER_VALUE_DELIMITER = "|"; - public static final String KEYWORD_TYPE = "keyword"; + + // Field types + public static final String KEYWORD_FIELD_TYPE = "keyword"; + public static final String BOOLEAN_FIELD_TYPE = "boolean"; + public static final String DATE_FIELD_TYPE = "date"; + public static final String DOUBLE_FIELD_TYPE = "double"; + public static final String LONG_FIELD_TYPE = "long"; + public static final String OBJECT_FIELD_TYPE = "object"; + public static final String TEXT_FIELD_TYPE = "text"; + public static final String TOKEN_COUNT_FIELD_TYPE = "token_count"; + // End of field types + + public static final Set FIELD_TYPES_STORED_AS_KEYWORD = Set.of( + SearchableAnnotation.FieldType.KEYWORD, + SearchableAnnotation.FieldType.TEXT, + SearchableAnnotation.FieldType.TEXT_PARTIAL, + SearchableAnnotation.FieldType.WORD_GRAM); + public static final Set FIELD_TYPES_STORED_AS_TEXT = Set.of( + SearchableAnnotation.FieldType.BROWSE_PATH, + SearchableAnnotation.FieldType.BROWSE_PATH_V2, + SearchableAnnotation.FieldType.URN, + SearchableAnnotation.FieldType.URN_PARTIAL); public static final String ENTITY_NAME_FIELD = "_entityName"; public static final String NAME_SUGGESTION = "nameSuggestion"; @@ -174,6 +198,25 @@ public static QueryBuilder getQueryBuilderFromCriterion(@Nonnull final Criterion return getQueryBuilderFromCriterionForSingleField(criterion, isTimeseries); } + public static String getElasticTypeForFieldType(SearchableAnnotation.FieldType fieldType) { + if (FIELD_TYPES_STORED_AS_KEYWORD.contains(fieldType)) { + return KEYWORD_FIELD_TYPE; + } else if (FIELD_TYPES_STORED_AS_TEXT.contains(fieldType)) { + return TEXT_FIELD_TYPE; + } else if (fieldType == SearchableAnnotation.FieldType.BOOLEAN) { + return BOOLEAN_FIELD_TYPE; + } else if (fieldType == SearchableAnnotation.FieldType.COUNT) { + return LONG_FIELD_TYPE; + } else if (fieldType == SearchableAnnotation.FieldType.DATETIME) { + return DATE_FIELD_TYPE; + } else if (fieldType == SearchableAnnotation.FieldType.OBJECT) { + return OBJECT_FIELD_TYPE; + } else { + log.warn("FieldType {} has no mappings implemented", fieldType); + return null; + } + } + /** * Populates source field of search query with the sort order as per the criterion provided. * @@ -189,14 +232,39 @@ public static QueryBuilder getQueryBuilderFromCriterion(@Nonnull final Criterion * @param sortCriterion {@link SortCriterion} to be applied to the search results */ public static void buildSortOrder(@Nonnull SearchSourceBuilder searchSourceBuilder, - @Nullable SortCriterion sortCriterion) { + @Nullable SortCriterion sortCriterion, List entitySpecs) { if (sortCriterion == null) { searchSourceBuilder.sort(new ScoreSortBuilder().order(SortOrder.DESC)); } else { + Optional fieldTypeForDefault = Optional.empty(); + for (EntitySpec entitySpec : entitySpecs) { + List fieldSpecs = entitySpec.getSearchableFieldSpecs(); + for (SearchableFieldSpec fieldSpec : fieldSpecs) { + SearchableAnnotation annotation = fieldSpec.getSearchableAnnotation(); + if (annotation.getFieldName().equals(sortCriterion.getField()) + || annotation.getFieldNameAliases().contains(sortCriterion.getField())) { + fieldTypeForDefault = Optional.of(fieldSpec.getSearchableAnnotation().getFieldType()); + break; + } + } + if (fieldTypeForDefault.isPresent()) { + break; + } + } + if (fieldTypeForDefault.isEmpty()) { + log.warn("Sort criterion field " + sortCriterion.getField() + " was not found in any entity spec to be searched"); + } final SortOrder esSortOrder = (sortCriterion.getOrder() == com.linkedin.metadata.query.filter.SortOrder.ASCENDING) ? SortOrder.ASC : SortOrder.DESC; - searchSourceBuilder.sort(new FieldSortBuilder(sortCriterion.getField()).order(esSortOrder).unmappedType(KEYWORD_TYPE)); + FieldSortBuilder sortBuilder = new FieldSortBuilder(sortCriterion.getField()).order(esSortOrder); + if (fieldTypeForDefault.isPresent()) { + String esFieldtype = getElasticTypeForFieldType(fieldTypeForDefault.get()); + if (esFieldtype != null) { + sortBuilder.unmappedType(esFieldtype); + } + } + searchSourceBuilder.sort(sortBuilder); } if (sortCriterion == null || !sortCriterion.getField().equals(DEFAULT_SEARCH_RESULTS_SORT_BY_FIELD)) { searchSourceBuilder.sort(new FieldSortBuilder(DEFAULT_SEARCH_RESULTS_SORT_BY_FIELD).order(SortOrder.ASC)); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/SampleDataFixtureTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/SampleDataFixtureTestBase.java index 1660504810296..69dd5c80bef1d 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/SampleDataFixtureTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/SampleDataFixtureTestBase.java @@ -22,12 +22,15 @@ import com.linkedin.metadata.query.filter.Criterion; import com.linkedin.metadata.query.filter.CriterionArray; import com.linkedin.metadata.query.filter.Filter; +import com.linkedin.metadata.query.filter.SortCriterion; +import com.linkedin.metadata.query.filter.SortOrder; import com.linkedin.metadata.search.AggregationMetadata; import com.linkedin.metadata.search.ScrollResult; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.search.SearchService; import com.linkedin.metadata.search.elasticsearch.query.request.SearchFieldConfig; +import com.linkedin.metadata.search.utils.ESUtils; import com.linkedin.r2.RemoteInvocationException; import org.junit.Assert; import org.opensearch.client.RequestOptions; @@ -36,6 +39,9 @@ import org.opensearch.client.indices.AnalyzeResponse; import org.opensearch.client.indices.GetMappingsRequest; import org.opensearch.client.indices.GetMappingsResponse; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.sort.FieldSortBuilder; +import org.opensearch.search.sort.SortBuilder; import org.springframework.test.context.testng.AbstractTestNGSpringContextTests; import org.testng.annotations.Test; @@ -54,11 +60,7 @@ import static com.linkedin.metadata.Constants.DATA_JOB_ENTITY_NAME; import static com.linkedin.metadata.search.elasticsearch.query.request.SearchQueryBuilder.STRUCTURED_QUERY_PREFIX; import static com.linkedin.metadata.utils.SearchUtil.AGGREGATION_SEPARATOR_CHAR; -import static io.datahubproject.test.search.SearchTestUtils.autocomplete; -import static io.datahubproject.test.search.SearchTestUtils.scroll; -import static io.datahubproject.test.search.SearchTestUtils.search; -import static io.datahubproject.test.search.SearchTestUtils.searchAcrossEntities; -import static io.datahubproject.test.search.SearchTestUtils.searchStructured; +import static io.datahubproject.test.search.SearchTestUtils.*; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotNull; @@ -174,6 +176,48 @@ public void testSearchFieldConfig() throws IOException { } } + @Test + public void testGetSortOrder() { + String dateFieldName = "lastOperationTime"; + List entityNamesToTestSearch = List.of("dataset", "chart", "corpgroup"); + List entitySpecs = entityNamesToTestSearch.stream().map( + name -> getEntityRegistry().getEntitySpec(name)) + .collect(Collectors.toList()); + SearchSourceBuilder builder = new SearchSourceBuilder(); + SortCriterion sortCriterion = new SortCriterion().setOrder(SortOrder.DESCENDING).setField(dateFieldName); + ESUtils.buildSortOrder(builder, sortCriterion, entitySpecs); + List> sorts = builder.sorts(); + assertEquals(sorts.size(), 2); // sort by last modified and then by urn + for (SortBuilder sort : sorts) { + assertTrue(sort instanceof FieldSortBuilder); + FieldSortBuilder fieldSortBuilder = (FieldSortBuilder) sort; + if (fieldSortBuilder.getFieldName().equals(dateFieldName)) { + assertEquals(fieldSortBuilder.order(), org.opensearch.search.sort.SortOrder.DESC); + assertEquals(fieldSortBuilder.unmappedType(), "date"); + } else { + assertEquals(fieldSortBuilder.getFieldName(), "urn"); + } + } + + // Test alias field + String entityNameField = "_entityName"; + SearchSourceBuilder nameBuilder = new SearchSourceBuilder(); + SortCriterion nameCriterion = new SortCriterion().setOrder(SortOrder.ASCENDING).setField(entityNameField); + ESUtils.buildSortOrder(nameBuilder, nameCriterion, entitySpecs); + sorts = nameBuilder.sorts(); + assertEquals(sorts.size(), 2); + for (SortBuilder sort : sorts) { + assertTrue(sort instanceof FieldSortBuilder); + FieldSortBuilder fieldSortBuilder = (FieldSortBuilder) sort; + if (fieldSortBuilder.getFieldName().equals(entityNameField)) { + assertEquals(fieldSortBuilder.order(), org.opensearch.search.sort.SortOrder.ASC); + assertEquals(fieldSortBuilder.unmappedType(), "keyword"); + } else { + assertEquals(fieldSortBuilder.getFieldName(), "urn"); + } + } + } + @Test public void testDatasetHasTags() throws IOException { GetMappingsRequest req = new GetMappingsRequest() @@ -1454,6 +1498,16 @@ public void testColumnExactMatch() { "Expected table with column name exact match first"); } + @Test + public void testSortOrdering() { + String query = "unit_data"; + SortCriterion criterion = new SortCriterion().setOrder(SortOrder.ASCENDING).setField("lastOperationTime"); + SearchResult result = getSearchService().searchAcrossEntities(SEARCHABLE_ENTITIES, query, null, criterion, 0, + 100, new SearchFlags().setFulltext(true).setSkipCache(true), null); + assertTrue(result.getEntities().size() > 2, + String.format("%s - Expected search results to have at least two results", query)); + } + private Stream getTokens(AnalyzeRequest request) throws IOException { return getSearchClient().indices().analyze(request, RequestOptions.DEFAULT).getTokens().stream(); } From 7855fb60a7e96e6d04d8d96f7505f8b4dd62a7c4 Mon Sep 17 00:00:00 2001 From: Indy Prentice Date: Wed, 18 Oct 2023 17:19:10 -0300 Subject: [PATCH 153/156] fix(api): Add preceding / to get index sizes path (#9043) Co-authored-by: Indy Prentice --- .../ElasticSearchTimeseriesAspectService.java | 2 +- .../search/TimeseriesAspectServiceTestBase.java | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/ElasticSearchTimeseriesAspectService.java b/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/ElasticSearchTimeseriesAspectService.java index a496fc427138e..3e8f83a531b59 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/ElasticSearchTimeseriesAspectService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/ElasticSearchTimeseriesAspectService.java @@ -169,7 +169,7 @@ public List getIndexSizes() { List res = new ArrayList<>(); try { String indicesPattern = _indexConvention.getAllTimeseriesAspectIndicesPattern(); - Response r = _searchClient.getLowLevelClient().performRequest(new Request("GET", indicesPattern + "/_stats")); + Response r = _searchClient.getLowLevelClient().performRequest(new Request("GET", "/" + indicesPattern + "/_stats")); JsonNode body = new ObjectMapper().readTree(r.getEntity().getContent()); body.get("indices").fields().forEachRemaining(entry -> { TimeseriesIndexSizeResult elemResult = new TimeseriesIndexSizeResult(); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceTestBase.java index cc60ba8679e1f..f9b8f84b10ad2 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceTestBase.java @@ -45,6 +45,7 @@ import com.linkedin.timeseries.GroupingBucket; import com.linkedin.timeseries.GroupingBucketType; import com.linkedin.timeseries.TimeWindowSize; +import com.linkedin.timeseries.TimeseriesIndexSizeResult; import org.opensearch.client.RestHighLevelClient; import org.springframework.test.context.testng.AbstractTestNGSpringContextTests; import org.testng.annotations.BeforeClass; @@ -884,4 +885,19 @@ public void testCountByFilterAfterDelete() throws InterruptedException { _elasticSearchTimeseriesAspectService.countByFilter(ENTITY_NAME, ASPECT_NAME, urnAndTimeFilter); assertEquals(count, 0L); } + + @Test(groups = {"getAggregatedStats"}, dependsOnGroups = {"upsert"}) + public void testGetIndexSizes() { + List result = _elasticSearchTimeseriesAspectService.getIndexSizes(); + /* + Example result: + {aspectName=testentityprofile, sizeMb=52.234, indexName=es_timeseries_aspect_service_test_testentity_testentityprofileaspect_v1, entityName=testentity} + {aspectName=testentityprofile, sizeMb=0.208, indexName=es_timeseries_aspect_service_test_testentitywithouttests_testentityprofileaspect_v1, entityName=testentitywithouttests} + */ + // There may be other indices in there from other tests, so just make sure that index for entity + aspect is in there + assertTrue(result.size() > 1); + assertTrue( + result.stream().anyMatch(idxSizeResult -> idxSizeResult.getIndexName().equals( + "es_timeseries_aspect_service_test_testentitywithouttests_testentityprofileaspect_v1"))); + } } From 409f981fd3e12a1d470a79cb091ac92e1a4a2c46 Mon Sep 17 00:00:00 2001 From: Indy Prentice Date: Wed, 18 Oct 2023 18:25:54 -0300 Subject: [PATCH 154/156] fix(search): Apply SearchFlags passed in through to scroll queries (#9041) Co-authored-by: Indy Prentice --- .../client/CachingEntitySearchService.java | 13 ++++++---- .../elasticsearch/ElasticSearchService.java | 13 ++++++---- .../query/request/SearchRequestHandler.java | 4 +++- .../search/LineageServiceTestBase.java | 16 ++++++++++--- .../request/SearchRequestHandlerTest.java | 24 +++++++++++++++++++ .../metadata/search/EntitySearchService.java | 6 +++-- 6 files changed, 60 insertions(+), 16 deletions(-) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/client/CachingEntitySearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/client/CachingEntitySearchService.java index 13a7d16b723a7..ceaf37a1289d9 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/client/CachingEntitySearchService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/client/CachingEntitySearchService.java @@ -256,13 +256,13 @@ public ScrollResult getCachedScrollResults( cacheAccess.stop(); if (result == null) { Timer.Context cacheMiss = MetricUtils.timer(this.getClass(), "scroll_cache_miss").time(); - result = getRawScrollResults(entities, query, filters, sortCriterion, scrollId, keepAlive, size, isFullText); + result = getRawScrollResults(entities, query, filters, sortCriterion, scrollId, keepAlive, size, isFullText, flags); cache.put(cacheKey, toJsonString(result)); cacheMiss.stop(); MetricUtils.counter(this.getClass(), "scroll_cache_miss_count").inc(); } } else { - result = getRawScrollResults(entities, query, filters, sortCriterion, scrollId, keepAlive, size, isFullText); + result = getRawScrollResults(entities, query, filters, sortCriterion, scrollId, keepAlive, size, isFullText, flags); } return result; } @@ -328,7 +328,8 @@ private ScrollResult getRawScrollResults( @Nullable final String scrollId, @Nullable final String keepAlive, final int count, - final boolean fulltext) { + final boolean fulltext, + @Nullable final SearchFlags searchFlags) { if (fulltext) { return entitySearchService.fullTextScroll( entities, @@ -337,7 +338,8 @@ private ScrollResult getRawScrollResults( sortCriterion, scrollId, keepAlive, - count); + count, + searchFlags); } else { return entitySearchService.structuredScroll(entities, input, @@ -345,7 +347,8 @@ private ScrollResult getRawScrollResults( sortCriterion, scrollId, keepAlive, - count); + count, + searchFlags); } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java index ef5a555e95ba8..024cf2b0abec2 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java @@ -175,23 +175,26 @@ public List getBrowsePaths(@Nonnull String entityName, @Nonnull Urn urn) @Nonnull @Override public ScrollResult fullTextScroll(@Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, - @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, int size) { + @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, int size, @Nullable SearchFlags searchFlags) { log.debug(String.format( "Scrolling Structured Search documents entities: %s, input: %s, postFilters: %s, sortCriterion: %s, scrollId: %s, size: %s", entities, input, postFilters, sortCriterion, scrollId, size)); + SearchFlags flags = Optional.ofNullable(searchFlags).orElse(new SearchFlags()); + flags.setFulltext(true); return esSearchDAO.scroll(entities, input, postFilters, sortCriterion, scrollId, keepAlive, size, - new SearchFlags().setFulltext(true)); + flags); } @Nonnull @Override public ScrollResult structuredScroll(@Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, - @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, int size) { + @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, int size, @Nullable SearchFlags searchFlags) { log.debug(String.format( "Scrolling FullText Search documents entities: %s, input: %s, postFilters: %s, sortCriterion: %s, scrollId: %s, size: %s", entities, input, postFilters, sortCriterion, scrollId, size)); - return esSearchDAO.scroll(entities, input, postFilters, sortCriterion, scrollId, keepAlive, size, - new SearchFlags().setFulltext(false)); + SearchFlags flags = Optional.ofNullable(searchFlags).orElse(new SearchFlags()); + flags.setFulltext(false); + return esSearchDAO.scroll(entities, input, postFilters, sortCriterion, scrollId, keepAlive, size, flags); } public Optional raw(@Nonnull String indexName, @Nullable String jsonQuery) { diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java index c06907e800d5e..49571a60d5f21 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java @@ -242,7 +242,9 @@ public SearchRequest getSearchRequest(@Nonnull String input, @Nullable Filter fi BoolQueryBuilder filterQuery = getFilterQuery(filter); searchSourceBuilder.query(QueryBuilders.boolQuery().must(getQuery(input, finalSearchFlags.isFulltext())).filter(filterQuery)); _aggregationQueryBuilder.getAggregations().forEach(searchSourceBuilder::aggregation); - searchSourceBuilder.highlighter(getHighlights()); + if (!finalSearchFlags.isSkipHighlighting()) { + searchSourceBuilder.highlighter(_highlights); + } ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, _entitySpecs); searchRequest.source(searchSourceBuilder); log.debug("Search request is: " + searchRequest); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java index 461a146022446..696e3b62834bd 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java @@ -47,8 +47,10 @@ import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; import org.junit.Assert; +import org.mockito.ArgumentCaptor; import org.mockito.Mockito; import org.opensearch.client.RestHighLevelClient; +import org.opensearch.action.search.SearchRequest; import org.springframework.cache.CacheManager; import org.springframework.cache.concurrent.ConcurrentMapCacheManager; import org.springframework.test.context.testng.AbstractTestNGSpringContextTests; @@ -108,6 +110,7 @@ abstract public class LineageServiceTestBase extends AbstractTestNGSpringContext private GraphService _graphService; private CacheManager _cacheManager; private LineageSearchService _lineageSearchService; + private RestHighLevelClient _searchClientSpy; private static final String ENTITY_NAME = "testEntity"; private static final Urn TEST_URN = TestEntityUtil.getTestEntityUrn(); @@ -162,10 +165,11 @@ private ElasticSearchService buildEntitySearchService() { EntityIndexBuilders indexBuilders = new EntityIndexBuilders(getIndexBuilder(), _entityRegistry, _indexConvention, _settingsBuilder); - ESSearchDAO searchDAO = new ESSearchDAO(_entityRegistry, getSearchClient(), _indexConvention, false, + _searchClientSpy = spy(getSearchClient()); + ESSearchDAO searchDAO = new ESSearchDAO(_entityRegistry, _searchClientSpy, _indexConvention, false, ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH, getSearchConfiguration(), null); - ESBrowseDAO browseDAO = new ESBrowseDAO(_entityRegistry, getSearchClient(), _indexConvention, getSearchConfiguration(), getCustomSearchConfiguration()); - ESWriteDAO writeDAO = new ESWriteDAO(_entityRegistry, getSearchClient(), _indexConvention, getBulkProcessor(), 1); + ESBrowseDAO browseDAO = new ESBrowseDAO(_entityRegistry, _searchClientSpy, _indexConvention, getSearchConfiguration(), getCustomSearchConfiguration()); + ESWriteDAO writeDAO = new ESWriteDAO(_entityRegistry, _searchClientSpy, _indexConvention, getBulkProcessor(), 1); return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO); } @@ -246,9 +250,15 @@ public void testSearchService() throws Exception { _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); syncAfterWrite(getBulkProcessor()); + Mockito.reset(_searchClientSpy); searchResult = searchAcrossLineage(null, TEST1); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); + // Verify that highlighting was turned off in the query + ArgumentCaptor searchRequestCaptor = ArgumentCaptor.forClass(SearchRequest.class); + Mockito.verify(_searchClientSpy, times(1)).search(searchRequestCaptor.capture(), any()); + SearchRequest capturedRequest = searchRequestCaptor.getValue(); + assertNull(capturedRequest.source().highlighter()); clearCache(false); when(_graphService.getLineage(eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), anyInt(), diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java index 90c6c523c588f..0ea035a10f91d 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java @@ -97,6 +97,30 @@ public void testDatasetFieldsAndHighlights() { ), "unexpected lineage fields in highlights: " + highlightFields); } + @Test + public void testSearchRequestHandlerHighlightingTurnedOff() { + SearchRequestHandler requestHandler = SearchRequestHandler.getBuilder(TestEntitySpecBuilder.getSpec(), testQueryConfig, null); + SearchRequest searchRequest = requestHandler.getSearchRequest("testQuery", null, null, 0, + 10, new SearchFlags().setFulltext(false).setSkipHighlighting(true), null); + SearchSourceBuilder sourceBuilder = searchRequest.source(); + assertEquals(sourceBuilder.from(), 0); + assertEquals(sourceBuilder.size(), 10); + // Filters + Collection aggBuilders = sourceBuilder.aggregations().getAggregatorFactories(); + // Expect 2 aggregations: textFieldOverride and _index + assertEquals(aggBuilders.size(), 2); + for (AggregationBuilder aggBuilder : aggBuilders) { + if (aggBuilder.getName().equals("textFieldOverride")) { + TermsAggregationBuilder filterPanelBuilder = (TermsAggregationBuilder) aggBuilder; + assertEquals(filterPanelBuilder.field(), "textFieldOverride.keyword"); + } else if (!aggBuilder.getName().equals("_entityType")) { + fail("Found unexepected aggregation: " + aggBuilder.getName()); + } + } + // Highlights should not be present + assertNull(sourceBuilder.highlighter()); + } + @Test public void testSearchRequestHandler() { SearchRequestHandler requestHandler = SearchRequestHandler.getBuilder(TestEntitySpecBuilder.getSpec(), testQueryConfig, null); diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java index a46b58aabfb0b..64f59780b887f 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java @@ -188,11 +188,12 @@ BrowseResult browse(@Nonnull String entityName, @Nonnull String path, @Nullable * @param sortCriterion {@link SortCriterion} to be applied to search results * @param scrollId opaque scroll identifier to pass to search service * @param size the number of search hits to return + * @param searchFlags flags controlling search options * @return a {@link ScrollResult} that contains a list of matched documents and related search result metadata */ @Nonnull ScrollResult fullTextScroll(@Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, - @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nonnull String keepAlive, int size); + @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nonnull String keepAlive, int size, @Nullable SearchFlags searchFlags); /** * Gets a list of documents that match given search request. The results are aggregated and filters are applied to the @@ -204,11 +205,12 @@ ScrollResult fullTextScroll(@Nonnull List entities, @Nonnull String inpu * @param sortCriterion {@link SortCriterion} to be applied to search results * @param scrollId opaque scroll identifier to pass to search service * @param size the number of search hits to return + * @param searchFlags flags controlling search options * @return a {@link ScrollResult} that contains a list of matched documents and related search result metadata */ @Nonnull ScrollResult structuredScroll(@Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, - @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nonnull String keepAlive, int size); + @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nonnull String keepAlive, int size, @Nullable SearchFlags searchFlags); /** * Max result size returned by the underlying search backend From 269c4eac7ef09d73224050e432bfbf60727e4d65 Mon Sep 17 00:00:00 2001 From: Pedro Silva Date: Thu, 19 Oct 2023 01:43:05 +0100 Subject: [PATCH 155/156] fix(ownership): Corrects validation of ownership type and makes it consistent across graphQL calls (#9044) Co-authored-by: Ellie O'Neil --- .../resolvers/mutate/AddOwnerResolver.java | 27 ++- .../resolvers/mutate/AddOwnersResolver.java | 2 +- .../mutate/BatchAddOwnersResolver.java | 3 +- .../resolvers/mutate/util/OwnerUtils.java | 65 +++----- .../owner/AddOwnersResolverTest.java | 157 ++++++++++++++++-- 5 files changed, 183 insertions(+), 71 deletions(-) diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnerResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnerResolver.java index 5ca7007d98e43..3f2dab0a5ba71 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnerResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnerResolver.java @@ -2,14 +2,11 @@ import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.CorpuserUrn; - import com.linkedin.common.urn.Urn; import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.exception.AuthorizationException; import com.linkedin.datahub.graphql.generated.AddOwnerInput; -import com.linkedin.datahub.graphql.generated.OwnerEntityType; import com.linkedin.datahub.graphql.generated.OwnerInput; -import com.linkedin.datahub.graphql.generated.OwnershipType; import com.linkedin.datahub.graphql.generated.ResourceRefInput; import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils; import com.linkedin.metadata.entity.EntityService; @@ -20,7 +17,6 @@ import lombok.extern.slf4j.Slf4j; import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*; -import static com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils.*; @Slf4j @@ -32,30 +28,33 @@ public class AddOwnerResolver implements DataFetcher> @Override public CompletableFuture get(DataFetchingEnvironment environment) throws Exception { final AddOwnerInput input = bindArgument(environment.getArgument("input"), AddOwnerInput.class); - Urn ownerUrn = Urn.createFromString(input.getOwnerUrn()); - OwnerEntityType ownerEntityType = input.getOwnerEntityType(); - OwnershipType type = input.getType() == null ? OwnershipType.NONE : input.getType(); - String ownershipUrn = input.getOwnershipTypeUrn() == null ? mapOwnershipTypeToEntity(type.name()) : input.getOwnershipTypeUrn(); Urn targetUrn = Urn.createFromString(input.getResourceUrn()); + OwnerInput.Builder ownerInputBuilder = OwnerInput.builder(); + ownerInputBuilder.setOwnerUrn(input.getOwnerUrn()); + ownerInputBuilder.setOwnerEntityType(input.getOwnerEntityType()); + if (input.getType() != null) { + ownerInputBuilder.setType(input.getType()); + } + if (input.getOwnershipTypeUrn() != null) { + ownerInputBuilder.setOwnershipTypeUrn(input.getOwnershipTypeUrn()); + } + OwnerInput ownerInput = ownerInputBuilder.build(); if (!OwnerUtils.isAuthorizedToUpdateOwners(environment.getContext(), targetUrn)) { throw new AuthorizationException("Unauthorized to perform this action. Please contact your DataHub administrator."); } return CompletableFuture.supplyAsync(() -> { - OwnerUtils.validateAddInput( - ownerUrn, input.getOwnershipTypeUrn(), ownerEntityType, - targetUrn, - _entityService - ); + OwnerUtils.validateAddOwnerInput(ownerInput, ownerUrn, _entityService); + try { log.debug("Adding Owner. input: {}", input); Urn actor = CorpuserUrn.createFromString(((QueryContext) environment.getContext()).getActorUrn()); OwnerUtils.addOwnersToResources( - ImmutableList.of(new OwnerInput(input.getOwnerUrn(), ownerEntityType, type, ownershipUrn)), + ImmutableList.of(ownerInput), ImmutableList.of(new ResourceRefInput(input.getResourceUrn(), null, null)), actor, _entityService diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnersResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnersResolver.java index 06424efa83819..4e5b5bdb2a651 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnersResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnersResolver.java @@ -39,7 +39,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) throw throw new AuthorizationException("Unauthorized to perform this action. Please contact your DataHub administrator."); } - OwnerUtils.validateAddInput( + OwnerUtils.validateAddOwnerInput( owners, targetUrn, _entityService diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchAddOwnersResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchAddOwnersResolver.java index 019c044d81ab3..5beaeecae673f 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchAddOwnersResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchAddOwnersResolver.java @@ -53,8 +53,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) throw private void validateOwners(List owners) { for (OwnerInput ownerInput : owners) { - OwnerUtils.validateOwner(UrnUtils.getUrn(ownerInput.getOwnerUrn()), ownerInput.getOwnerEntityType(), - UrnUtils.getUrn(ownerInput.getOwnershipTypeUrn()), _entityService); + OwnerUtils.validateOwner(ownerInput, _entityService); } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java index d2f7f896e5953..7233995804423 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java @@ -50,7 +50,7 @@ public static void addOwnersToResources( ) { final List changes = new ArrayList<>(); for (ResourceRefInput resource : resources) { - changes.add(buildAddOwnersProposal(owners, UrnUtils.getUrn(resource.getResourceUrn()), actor, entityService)); + changes.add(buildAddOwnersProposal(owners, UrnUtils.getUrn(resource.getResourceUrn()), entityService)); } EntityUtils.ingestChangeProposals(changes, entityService, actor, false); } @@ -69,7 +69,7 @@ public static void removeOwnersFromResources( } - private static MetadataChangeProposal buildAddOwnersProposal(List owners, Urn resourceUrn, Urn actor, EntityService entityService) { + static MetadataChangeProposal buildAddOwnersProposal(List owners, Urn resourceUrn, EntityService entityService) { Ownership ownershipAspect = (Ownership) EntityUtils.getAspectFromEntity( resourceUrn.toString(), Constants.OWNERSHIP_ASPECT_NAME, entityService, @@ -181,18 +181,13 @@ public static boolean isAuthorizedToUpdateOwners(@Nonnull QueryContext context, orPrivilegeGroups); } - public static Boolean validateAddInput( + public static Boolean validateAddOwnerInput( List owners, Urn resourceUrn, EntityService entityService ) { for (OwnerInput owner : owners) { - boolean result = validateAddInput( - UrnUtils.getUrn(owner.getOwnerUrn()), - owner.getOwnershipTypeUrn(), - owner.getOwnerEntityType(), - resourceUrn, - entityService); + boolean result = validateAddOwnerInput(owner, resourceUrn, entityService); if (!result) { return false; } @@ -200,44 +195,29 @@ public static Boolean validateAddInput( return true; } - public static Boolean validateAddInput( - Urn ownerUrn, - String ownershipEntityUrn, - OwnerEntityType ownerEntityType, + public static Boolean validateAddOwnerInput( + OwnerInput owner, Urn resourceUrn, EntityService entityService ) { - if (OwnerEntityType.CORP_GROUP.equals(ownerEntityType) && !Constants.CORP_GROUP_ENTITY_NAME.equals(ownerUrn.getEntityType())) { - throw new IllegalArgumentException(String.format("Failed to change ownership for resource %s. Expected a corp group urn.", resourceUrn)); - } - - if (OwnerEntityType.CORP_USER.equals(ownerEntityType) && !Constants.CORP_USER_ENTITY_NAME.equals(ownerUrn.getEntityType())) { - throw new IllegalArgumentException(String.format("Failed to change ownership for resource %s. Expected a corp user urn.", resourceUrn)); - } - if (!entityService.exists(resourceUrn)) { throw new IllegalArgumentException(String.format("Failed to change ownership for resource %s. Resource does not exist.", resourceUrn)); } - if (!entityService.exists(ownerUrn)) { - throw new IllegalArgumentException(String.format("Failed to change ownership for resource %s. Owner %s does not exist.", resourceUrn, ownerUrn)); - } - - if (ownershipEntityUrn != null && !entityService.exists(UrnUtils.getUrn(ownershipEntityUrn))) { - throw new IllegalArgumentException(String.format("Failed to change ownership type for resource %s. Ownership Type " - + "%s does not exist.", resourceUrn, ownershipEntityUrn)); - } + validateOwner(owner, entityService); return true; } public static void validateOwner( - Urn ownerUrn, - OwnerEntityType ownerEntityType, - Urn ownershipEntityUrn, + OwnerInput owner, EntityService entityService ) { + + OwnerEntityType ownerEntityType = owner.getOwnerEntityType(); + Urn ownerUrn = UrnUtils.getUrn(owner.getOwnerUrn()); + if (OwnerEntityType.CORP_GROUP.equals(ownerEntityType) && !Constants.CORP_GROUP_ENTITY_NAME.equals(ownerUrn.getEntityType())) { throw new IllegalArgumentException( String.format("Failed to change ownership for resource(s). Expected a corp group urn, found %s", ownerUrn)); @@ -252,9 +232,14 @@ public static void validateOwner( throw new IllegalArgumentException(String.format("Failed to change ownership for resource(s). Owner with urn %s does not exist.", ownerUrn)); } - if (!entityService.exists(ownershipEntityUrn)) { - throw new IllegalArgumentException(String.format("Failed to change ownership for resource(s). Ownership type with " - + "urn %s does not exist.", ownershipEntityUrn)); + if (owner.getOwnershipTypeUrn() != null && !entityService.exists(UrnUtils.getUrn(owner.getOwnershipTypeUrn()))) { + throw new IllegalArgumentException(String.format("Failed to change ownership for resource(s). Custom Ownership type with " + + "urn %s does not exist.", owner.getOwnershipTypeUrn())); + } + + if (owner.getType() == null && owner.getOwnershipTypeUrn() == null) { + throw new IllegalArgumentException("Failed to change ownership for resource(s). Expected either " + + "type or ownershipTypeUrn to be specified."); } } @@ -269,11 +254,11 @@ public static Boolean validateRemoveInput( } public static void addCreatorAsOwner( - QueryContext context, - String urn, - OwnerEntityType ownerEntityType, - OwnershipType ownershipType, - EntityService entityService) { + QueryContext context, + String urn, + OwnerEntityType ownerEntityType, + OwnershipType ownershipType, + EntityService entityService) { try { Urn actorUrn = CorpuserUrn.createFromString(context.getActorUrn()); String ownershipTypeUrn = mapOwnershipTypeToEntity(ownershipType.name()); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/AddOwnersResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/AddOwnersResolverTest.java index efc0c5dfcf36d..329d71ec125db 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/AddOwnersResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/AddOwnersResolverTest.java @@ -2,6 +2,11 @@ import com.google.common.collect.ImmutableList; import com.linkedin.common.AuditStamp; +import com.linkedin.common.Owner; +import com.linkedin.common.OwnerArray; +import com.linkedin.common.Ownership; +import com.linkedin.common.OwnershipSource; +import com.linkedin.common.OwnershipSourceType; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.datahub.graphql.QueryContext; @@ -28,6 +33,7 @@ public class AddOwnersResolverTest { private static final String TEST_ENTITY_URN = "urn:li:dataset:(urn:li:dataPlatform:mysql,my-test,PROD)"; private static final String TEST_OWNER_1_URN = "urn:li:corpuser:test-id-1"; private static final String TEST_OWNER_2_URN = "urn:li:corpuser:test-id-2"; + private static final String TEST_OWNER_3_URN = "urn:li:corpGroup:test-id-3"; @Test public void testGetSuccessNoExistingOwners() throws Exception { @@ -75,33 +81,41 @@ public void testGetSuccessNoExistingOwners() throws Exception { } @Test - public void testGetSuccessExistingOwners() throws Exception { + public void testGetSuccessExistingOwnerNewType() throws Exception { EntityService mockService = getMockEntityService(); + com.linkedin.common.Ownership oldOwnership = new Ownership().setOwners(new OwnerArray( + ImmutableList.of(new Owner() + .setOwner(UrnUtils.getUrn(TEST_OWNER_1_URN)) + .setType(com.linkedin.common.OwnershipType.NONE) + .setSource(new OwnershipSource().setType(OwnershipSourceType.MANUAL)) + ))); + Mockito.when(mockService.getAspect( - Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), - Mockito.eq(Constants.OWNERSHIP_ASPECT_NAME), - Mockito.eq(0L))) - .thenReturn(null); + Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), + Mockito.eq(Constants.OWNERSHIP_ASPECT_NAME), + Mockito.eq(0L))) + .thenReturn(oldOwnership); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(true); Mockito.when(mockService.exists(Urn.createFromString(TEST_OWNER_1_URN))).thenReturn(true); - Mockito.when(mockService.exists(Urn.createFromString(TEST_OWNER_2_URN))).thenReturn(true); Mockito.when(mockService.exists(Urn.createFromString( - OwnerUtils.mapOwnershipTypeToEntity(com.linkedin.datahub.graphql.generated.OwnershipType.TECHNICAL_OWNER.name())))) - .thenReturn(true); + OwnerUtils.mapOwnershipTypeToEntity(com.linkedin.datahub.graphql.generated.OwnershipType.TECHNICAL_OWNER.name())))) + .thenReturn(true); AddOwnersResolver resolver = new AddOwnersResolver(mockService); // Execute resolver QueryContext mockContext = getMockAllowContext(); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + AddOwnersInput input = new AddOwnersInput(ImmutableList.of( - new OwnerInput(TEST_OWNER_1_URN, OwnerEntityType.CORP_USER, OwnershipType.TECHNICAL_OWNER, - OwnerUtils.mapOwnershipTypeToEntity(OwnershipType.TECHNICAL_OWNER.name())), - new OwnerInput(TEST_OWNER_2_URN, OwnerEntityType.CORP_USER, OwnershipType.TECHNICAL_OWNER, - OwnerUtils.mapOwnershipTypeToEntity(OwnershipType.TECHNICAL_OWNER.name())) + OwnerInput.builder() + .setOwnerUrn(TEST_OWNER_1_URN) + .setOwnershipTypeUrn(OwnerUtils.mapOwnershipTypeToEntity(OwnershipType.TECHNICAL_OWNER.name())) + .setOwnerEntityType(OwnerEntityType.CORP_USER) + .build() ), TEST_ENTITY_URN); Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input); Mockito.when(mockEnv.getContext()).thenReturn(mockContext); @@ -111,11 +125,126 @@ public void testGetSuccessExistingOwners() throws Exception { verifyIngestProposal(mockService, 1); Mockito.verify(mockService, Mockito.times(1)).exists( - Mockito.eq(Urn.createFromString(TEST_OWNER_1_URN)) + Mockito.eq(Urn.createFromString(TEST_OWNER_1_URN)) ); + } + + @Test + public void testGetSuccessDeprecatedTypeToOwnershipType() throws Exception { + EntityService mockService = getMockEntityService(); + + com.linkedin.common.Ownership oldOwnership = new Ownership().setOwners(new OwnerArray( + ImmutableList.of(new Owner() + .setOwner(UrnUtils.getUrn(TEST_OWNER_1_URN)) + .setType(com.linkedin.common.OwnershipType.TECHNICAL_OWNER) + .setSource(new OwnershipSource().setType(OwnershipSourceType.MANUAL)) + ))); + + Mockito.when(mockService.getAspect( + Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), + Mockito.eq(Constants.OWNERSHIP_ASPECT_NAME), + Mockito.eq(0L))) + .thenReturn(oldOwnership); + + Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(true); + Mockito.when(mockService.exists(Urn.createFromString(TEST_OWNER_1_URN))).thenReturn(true); + + Mockito.when(mockService.exists(Urn.createFromString( + OwnerUtils.mapOwnershipTypeToEntity(com.linkedin.datahub.graphql.generated.OwnershipType.TECHNICAL_OWNER.name())))) + .thenReturn(true); + + AddOwnersResolver resolver = new AddOwnersResolver(mockService); + + // Execute resolver + QueryContext mockContext = getMockAllowContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + + AddOwnersInput input = new AddOwnersInput(ImmutableList.of(OwnerInput.builder() + .setOwnerUrn(TEST_OWNER_1_URN) + .setOwnershipTypeUrn(OwnerUtils.mapOwnershipTypeToEntity(OwnershipType.TECHNICAL_OWNER.name())) + .setOwnerEntityType(OwnerEntityType.CORP_USER) + .build() + ), TEST_ENTITY_URN); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + assertTrue(resolver.get(mockEnv).get()); + + // Unable to easily validate exact payload due to the injected timestamp + verifyIngestProposal(mockService, 1); Mockito.verify(mockService, Mockito.times(1)).exists( - Mockito.eq(Urn.createFromString(TEST_OWNER_2_URN)) + Mockito.eq(Urn.createFromString(TEST_OWNER_1_URN)) + ); + } + + @Test + public void testGetSuccessMultipleOwnerTypes() throws Exception { + EntityService mockService = getMockEntityService(); + + com.linkedin.common.Ownership oldOwnership = new Ownership().setOwners(new OwnerArray( + ImmutableList.of(new Owner() + .setOwner(UrnUtils.getUrn(TEST_OWNER_1_URN)) + .setType(com.linkedin.common.OwnershipType.NONE) + .setSource(new OwnershipSource().setType(OwnershipSourceType.MANUAL)) + ))); + + Mockito.when(mockService.getAspect( + Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), + Mockito.eq(Constants.OWNERSHIP_ASPECT_NAME), + Mockito.eq(0L))) + .thenReturn(oldOwnership); + + Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(true); + Mockito.when(mockService.exists(Urn.createFromString(TEST_OWNER_1_URN))).thenReturn(true); + Mockito.when(mockService.exists(Urn.createFromString(TEST_OWNER_2_URN))).thenReturn(true); + Mockito.when(mockService.exists(Urn.createFromString(TEST_OWNER_3_URN))).thenReturn(true); + + Mockito.when(mockService.exists(Urn.createFromString( + OwnerUtils.mapOwnershipTypeToEntity(com.linkedin.datahub.graphql.generated.OwnershipType.TECHNICAL_OWNER.name())))) + .thenReturn(true); + Mockito.when(mockService.exists(Urn.createFromString( + OwnerUtils.mapOwnershipTypeToEntity(com.linkedin.datahub.graphql.generated.OwnershipType.BUSINESS_OWNER.name())))) + .thenReturn(true); + + AddOwnersResolver resolver = new AddOwnersResolver(mockService); + + // Execute resolver + QueryContext mockContext = getMockAllowContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + + AddOwnersInput input = new AddOwnersInput(ImmutableList.of(OwnerInput.builder() + .setOwnerUrn(TEST_OWNER_1_URN) + .setOwnershipTypeUrn(OwnerUtils.mapOwnershipTypeToEntity(OwnershipType.TECHNICAL_OWNER.name())) + .setOwnerEntityType(OwnerEntityType.CORP_USER) + .build(), + OwnerInput.builder() + .setOwnerUrn(TEST_OWNER_2_URN) + .setOwnershipTypeUrn(OwnerUtils.mapOwnershipTypeToEntity(OwnershipType.BUSINESS_OWNER.name())) + .setOwnerEntityType(OwnerEntityType.CORP_USER) + .build(), + OwnerInput.builder() + .setOwnerUrn(TEST_OWNER_3_URN) + .setOwnershipTypeUrn(OwnerUtils.mapOwnershipTypeToEntity(OwnershipType.TECHNICAL_OWNER.name())) + .setOwnerEntityType(OwnerEntityType.CORP_GROUP) + .build() + ), TEST_ENTITY_URN); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + assertTrue(resolver.get(mockEnv).get()); + + // Unable to easily validate exact payload due to the injected timestamp + verifyIngestProposal(mockService, 1); + + Mockito.verify(mockService, Mockito.times(1)).exists( + Mockito.eq(Urn.createFromString(TEST_OWNER_1_URN)) + ); + + Mockito.verify(mockService, Mockito.times(1)).exists( + Mockito.eq(Urn.createFromString(TEST_OWNER_2_URN)) + ); + + Mockito.verify(mockService, Mockito.times(1)).exists( + Mockito.eq(Urn.createFromString(TEST_OWNER_3_URN)) ); } From 75b36c41ee4fd74891b1bfe37885b4cd840e2906 Mon Sep 17 00:00:00 2001 From: Ellie O'Neil <110510035+eboneil@users.noreply.github.com> Date: Thu, 19 Oct 2023 08:32:24 -0700 Subject: [PATCH 156/156] docs(protobuf) Update messaging around nesting messages (#9048) --- metadata-integration/java/datahub-protobuf/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/metadata-integration/java/datahub-protobuf/README.md b/metadata-integration/java/datahub-protobuf/README.md index daea8d438679c..29b82aa3e68f5 100644 --- a/metadata-integration/java/datahub-protobuf/README.md +++ b/metadata-integration/java/datahub-protobuf/README.md @@ -1,6 +1,6 @@ # Protobuf Schemas -The `datahub-protobuf` module is designed to be used with the Java Emitter, the input is a compiled protobuf binary `*.protoc` files and optionally the corresponding `*.proto` source code. In addition, you can supply the root message in cases where a single protobuf source file includes multiple non-nested messages. +The `datahub-protobuf` module is designed to be used with the Java Emitter, the input is a compiled protobuf binary `*.protoc` files and optionally the corresponding `*.proto` source code. You can supply a file with multiple nested messages to be processed. If you have a file with multiple non-nested messages, you will need to separate them out into different files or supply the root message, as otherwise we will only process the first one. ## Supported Features