From f561efee8387587e74f5471ac2fcd8d6e912df6f Mon Sep 17 00:00:00 2001 From: Tamas Nemeth Date: Wed, 28 Feb 2024 15:00:14 +0100 Subject: [PATCH 01/13] fix(ingest/mysql): Adding support for bit type (#9950) --- .../src/datahub/ingestion/source/sql/mysql.py | 8 +- .../mysql/mysql_mces_no_db_golden.json | 273 +++++++++++++++++- .../mysql/mysql_mces_with_db_golden.json | 62 +++- .../mysql/mysql_table_level_only.json | 109 +++++-- .../mysql_table_row_count_estimate_only.json | 62 +++- 5 files changed, 485 insertions(+), 29 deletions(-) diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/mysql.py b/metadata-ingestion/src/datahub/ingestion/source/sql/mysql.py index 2126717f835a26..9b482beba924f9 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/mysql.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/mysql.py @@ -3,7 +3,7 @@ import pymysql # noqa: F401 from pydantic.fields import Field from sqlalchemy import util -from sqlalchemy.dialects.mysql import base +from sqlalchemy.dialects.mysql import BIT, base from sqlalchemy.dialects.mysql.enumerated import SET from sqlalchemy.engine.reflection import Inspector @@ -24,6 +24,7 @@ TwoTierSQLAlchemyConfig, TwoTierSQLAlchemySource, ) +from datahub.metadata._schema_classes import BytesTypeClass SET.__repr__ = util.generic_repr # type:ignore @@ -38,6 +39,7 @@ register_custom_type(LINESTRING) register_custom_type(POLYGON) register_custom_type(DECIMAL128) +register_custom_type(BIT, BytesTypeClass) base.ischema_names["geometry"] = GEOMETRY base.ischema_names["point"] = POINT @@ -92,5 +94,5 @@ def add_profile_metadata(self, inspector: Inspector) -> None: "SELECT table_schema, table_name, data_length from information_schema.tables" ): self.profile_metadata_info.dataset_name_to_storage_bytes[ - f"{row.table_schema}.{row.table_name}" - ] = row.data_length + f"{row.TABLE_SCHEMA}.{row.TABLE_NAME}" + ] = row.DATA_LENGTH diff --git a/metadata-ingestion/tests/integration/mysql/mysql_mces_no_db_golden.json b/metadata-ingestion/tests/integration/mysql/mysql_mces_no_db_golden.json index a86ed53406e40d..cb22a6cb0a346c 100644 --- a/metadata-ingestion/tests/integration/mysql/mysql_mces_no_db_golden.json +++ b/metadata-ingestion/tests/integration/mysql/mysql_mces_no_db_golden.json @@ -562,7 +562,8 @@ "1989-08-24" ] } - ] + ], + "sizeInBytes": 16384 } }, "systemMetadata": { @@ -978,7 +979,8 @@ "1993-02-09" ] } - ] + ], + "sizeInBytes": 16384 } }, "systemMetadata": { @@ -2227,7 +2229,8 @@ "3.8" ] } - ] + ], + "sizeInBytes": 16384 } }, "systemMetadata": { @@ -2266,7 +2269,8 @@ "uniqueCount": 0, "nullCount": 0 } - ] + ], + "sizeInBytes": 16384 } }, "systemMetadata": { @@ -2614,7 +2618,8 @@ "uniqueCount": 0, "nullCount": 0 } - ] + ], + "sizeInBytes": 16384 } }, "systemMetadata": { @@ -2643,7 +2648,8 @@ "uniqueCount": 0, "nullCount": 0 } - ] + ], + "sizeInBytes": 16384 } }, "systemMetadata": { @@ -2722,5 +2728,258 @@ "runId": "mysql-test", "lastRunId": "no-run-id-provided" } +}, +{ + "entityType": "container", + "entityUrn": "urn:li:container:0f72a1bc79da282eb614cc089c0ba302", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,dataCharmer.employees,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:0f72a1bc79da282eb614cc089c0ba302", + "urn": "urn:li:container:0f72a1bc79da282eb614cc089c0ba302" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,dataCharmer.salaries,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:0f72a1bc79da282eb614cc089c0ba302", + "urn": "urn:li:container:0f72a1bc79da282eb614cc089c0ba302" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "container", + "entityUrn": "urn:li:container:17751259af32dd0385cad799df608c40", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,metagalaxy.metadata_aspect,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:17751259af32dd0385cad799df608c40", + "urn": "urn:li:container:17751259af32dd0385cad799df608c40" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,metagalaxy.metadata_index,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:17751259af32dd0385cad799df608c40", + "urn": "urn:li:container:17751259af32dd0385cad799df608c40" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,metagalaxy.metadata_index_view,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:17751259af32dd0385cad799df608c40", + "urn": "urn:li:container:17751259af32dd0385cad799df608c40" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "container", + "entityUrn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,northwind.customers,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "urn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,northwind.orders,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "urn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "container", + "entityUrn": "urn:li:container:28176129fe1c0e526e1803250ec124ef", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,test_cases.myset,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:28176129fe1c0e526e1803250ec124ef", + "urn": "urn:li:container:28176129fe1c0e526e1803250ec124ef" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,test_cases.test_empty,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:28176129fe1c0e526e1803250ec124ef", + "urn": "urn:li:container:28176129fe1c0e526e1803250ec124ef" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } } -] +] \ No newline at end of file diff --git a/metadata-ingestion/tests/integration/mysql/mysql_mces_with_db_golden.json b/metadata-ingestion/tests/integration/mysql/mysql_mces_with_db_golden.json index b5ebca424d9a2d..f24220b4dbf596 100644 --- a/metadata-ingestion/tests/integration/mysql/mysql_mces_with_db_golden.json +++ b/metadata-ingestion/tests/integration/mysql/mysql_mces_with_db_golden.json @@ -583,7 +583,8 @@ "3.8" ] } - ] + ], + "sizeInBytes": 16384 } }, "systemMetadata": { @@ -622,6 +623,65 @@ "uniqueCount": 0, "nullCount": 0 } + ], + "sizeInBytes": 16384 + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "container", + "entityUrn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,northwind.customers,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "urn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-test", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,northwind.orders,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "urn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f" + } ] } }, diff --git a/metadata-ingestion/tests/integration/mysql/mysql_table_level_only.json b/metadata-ingestion/tests/integration/mysql/mysql_table_level_only.json index cbe5f4e9e314d5..b8dfd7d9efc37d 100644 --- a/metadata-ingestion/tests/integration/mysql/mysql_table_level_only.json +++ b/metadata-ingestion/tests/integration/mysql/mysql_table_level_only.json @@ -16,7 +16,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -31,7 +32,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -46,7 +48,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -63,7 +66,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -78,7 +82,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -93,7 +98,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -213,7 +219,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -230,7 +237,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -250,7 +258,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -265,7 +274,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -361,7 +371,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -378,7 +389,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -398,7 +410,8 @@ }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -415,12 +428,14 @@ }, "rowCount": 5, "columnCount": 6, - "fieldProfiles": [] + "fieldProfiles": [], + "sizeInBytes": 16384 } }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } }, { @@ -437,12 +452,72 @@ }, "rowCount": 0, "columnCount": 3, - "fieldProfiles": [] + "fieldProfiles": [], + "sizeInBytes": 16384 } }, "systemMetadata": { "lastObserved": 1586847600000, - "runId": "mysql-2020_04_14-07_00_00" + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "container", + "entityUrn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,northwind.customers,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "urn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,northwind.orders,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "urn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" } } ] \ No newline at end of file diff --git a/metadata-ingestion/tests/integration/mysql/mysql_table_row_count_estimate_only.json b/metadata-ingestion/tests/integration/mysql/mysql_table_row_count_estimate_only.json index 634e04984986d6..34b18089aeebf4 100644 --- a/metadata-ingestion/tests/integration/mysql/mysql_table_row_count_estimate_only.json +++ b/metadata-ingestion/tests/integration/mysql/mysql_table_row_count_estimate_only.json @@ -464,7 +464,8 @@ "uniqueProportion": 0.75, "nullCount": 0 } - ] + ], + "sizeInBytes": 16384 } }, "systemMetadata": { @@ -503,6 +504,65 @@ "uniqueCount": 0, "nullCount": 0 } + ], + "sizeInBytes": 16384 + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "container", + "entityUrn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,northwind.customers,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "urn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f" + } + ] + } + }, + "systemMetadata": { + "lastObserved": 1586847600000, + "runId": "mysql-2020_04_14-07_00_00", + "lastRunId": "no-run-id-provided" + } +}, +{ + "entityType": "dataset", + "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:mysql,northwind.orders,PROD)", + "changeType": "UPSERT", + "aspectName": "browsePathsV2", + "aspect": { + "json": { + "path": [ + { + "id": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f", + "urn": "urn:li:container:dc2ae101b66746b9c2b6df8ee89ca88f" + } ] } }, From f399a872ad221a00f9cf04d3ef900b6ee13bb39e Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Wed, 28 Feb 2024 06:03:09 -0800 Subject: [PATCH 02/13] feat(ingest/patch): add helper for auto-quoting (#9938) --- metadata-ingestion/build.gradle | 10 ++++----- .../src/datahub/emitter/mcp_patch_builder.py | 9 ++++++-- .../src/datahub/specific/chart.py | 4 +--- .../src/datahub/specific/custom_properties.py | 8 +++---- .../src/datahub/specific/dashboard.py | 8 +++---- .../src/datahub/specific/datajob.py | 22 +++++++++---------- .../src/datahub/specific/dataproduct.py | 4 +--- .../src/datahub/specific/dataset.py | 10 ++++----- .../src/datahub/specific/ownership.py | 8 +++---- .../datahub/specific/structured_properties.py | 12 +++++----- 10 files changed, 46 insertions(+), 49 deletions(-) diff --git a/metadata-ingestion/build.gradle b/metadata-ingestion/build.gradle index 1aae02ab8bede9..269424d3de72b2 100644 --- a/metadata-ingestion/build.gradle +++ b/metadata-ingestion/build.gradle @@ -42,8 +42,8 @@ task installPackageOnly(type: Exec, dependsOn: runPreFlightScript) { def sentinel_file = "${venv_name}/.build_install_package_only_sentinel" inputs.file file('setup.py') outputs.file(sentinel_file) - commandLine 'bash', '-x', '-c', - "source ${venv_name}/bin/activate && " + + commandLine 'bash', '-c', + "source ${venv_name}/bin/activate && set -x && " + "uv pip install -e . &&" + "touch ${sentinel_file}" } @@ -52,8 +52,8 @@ task installPackage(type: Exec, dependsOn: installPackageOnly) { def sentinel_file = "${venv_name}/.build_install_package_sentinel" inputs.file file('setup.py') outputs.file(sentinel_file) - commandLine 'bash', '-x', '-c', - "source ${venv_name}/bin/activate && " + + commandLine 'bash', '-c', + "source ${venv_name}/bin/activate && set -x && " + "uv pip install -e . ${extra_pip_requirements} && " + "touch ${sentinel_file}" } @@ -71,7 +71,7 @@ task customPackageGenerate(type: Exec, dependsOn: [environmentSetup, installPack def package_name = project.findProperty('package_name') def package_version = project.findProperty('package_version') commandLine 'bash', '-c', - "source ${venv_name}/bin/activate && " + + "source ${venv_name}/bin/activate && set -x && " + "uv pip install build && " + "./scripts/custom_package_codegen.sh '${package_name}' '${package_version}'" } diff --git a/metadata-ingestion/src/datahub/emitter/mcp_patch_builder.py b/metadata-ingestion/src/datahub/emitter/mcp_patch_builder.py index f5ef276d5f0349..37903995394379 100644 --- a/metadata-ingestion/src/datahub/emitter/mcp_patch_builder.py +++ b/metadata-ingestion/src/datahub/emitter/mcp_patch_builder.py @@ -1,7 +1,7 @@ import json from collections import defaultdict from dataclasses import dataclass -from typing import Any, Dict, Iterable, List, Optional +from typing import Any, Dict, Iterable, List, Optional, Sequence, Union from datahub.emitter.aspect import JSON_PATCH_CONTENT_TYPE from datahub.emitter.serialization_helper import pre_json_transform @@ -62,7 +62,12 @@ def __init__( def quote(cls, value: str) -> str: return value.replace("~", "~0").replace("/", "~1") - def _add_patch(self, aspect_name: str, op: str, path: str, value: Any) -> None: + def _add_patch( + self, aspect_name: str, op: str, path: Union[str, Sequence[str]], value: Any + ) -> None: + if not isinstance(path, str): + path = "/" + "/".join(self.quote(p) for p in path) + # TODO: Validate that aspectName is a valid aspect for this entityType self.patches[aspect_name].append(_Patch(op, path, value)) diff --git a/metadata-ingestion/src/datahub/specific/chart.py b/metadata-ingestion/src/datahub/specific/chart.py index 2e61d899efb131..51994ad1d063eb 100644 --- a/metadata-ingestion/src/datahub/specific/chart.py +++ b/metadata-ingestion/src/datahub/specific/chart.py @@ -1,5 +1,5 @@ import time -from typing import Dict, List, Optional, TypeVar, Union +from typing import Dict, List, Optional, Union from datahub.emitter.mcp_patch_builder import MetadataPatchProposal from datahub.metadata.schema_classes import ( @@ -20,8 +20,6 @@ from datahub.utilities.urns.tag_urn import TagUrn from datahub.utilities.urns.urn import Urn -T = TypeVar("T", bound=MetadataPatchProposal) - class ChartPatchBuilder(MetadataPatchProposal): def __init__( diff --git a/metadata-ingestion/src/datahub/specific/custom_properties.py b/metadata-ingestion/src/datahub/specific/custom_properties.py index 2d810f4c2b584f..d399a448cc0c23 100644 --- a/metadata-ingestion/src/datahub/specific/custom_properties.py +++ b/metadata-ingestion/src/datahub/specific/custom_properties.py @@ -2,20 +2,20 @@ from datahub.emitter.mcp_patch_builder import MetadataPatchProposal -T = TypeVar("T", bound=MetadataPatchProposal) +_Parent = TypeVar("_Parent", bound=MetadataPatchProposal) -class CustomPropertiesPatchHelper(Generic[T]): +class CustomPropertiesPatchHelper(Generic[_Parent]): def __init__( self, - parent: T, + parent: _Parent, aspect_name: str, ) -> None: self.aspect_name = aspect_name self._parent = parent self.aspect_field = "customProperties" - def parent(self) -> T: + def parent(self) -> _Parent: return self._parent def add_property(self, key: str, value: str) -> "CustomPropertiesPatchHelper": diff --git a/metadata-ingestion/src/datahub/specific/dashboard.py b/metadata-ingestion/src/datahub/specific/dashboard.py index cf6f0085a00cda..e6d911b5986550 100644 --- a/metadata-ingestion/src/datahub/specific/dashboard.py +++ b/metadata-ingestion/src/datahub/specific/dashboard.py @@ -1,5 +1,5 @@ import time -from typing import Dict, List, Optional, TypeVar, Union +from typing import Dict, List, Optional, Union from datahub.emitter.mcp_patch_builder import MetadataPatchProposal from datahub.metadata.schema_classes import ( @@ -20,8 +20,6 @@ from datahub.utilities.urns.tag_urn import TagUrn from datahub.utilities.urns.urn import Urn -T = TypeVar("T", bound=MetadataPatchProposal) - class DashboardPatchBuilder(MetadataPatchProposal): def __init__( @@ -165,7 +163,7 @@ def add_dataset_edge( self._add_patch( DashboardInfo.ASPECT_NAME, "add", - path=f"/datasetEdges/{MetadataPatchProposal.quote(dataset_urn)}", + path=f"/datasetEdges/{self.quote(dataset_urn)}", value=dataset_edge, ) return self @@ -248,7 +246,7 @@ def add_chart_edge(self, chart: Union[Edge, Urn, str]) -> "DashboardPatchBuilder self._add_patch( DashboardInfo.ASPECT_NAME, "add", - path=f"/chartEdges/{MetadataPatchProposal.quote(chart_urn)}", + path=f"/chartEdges/{self.quote(chart_urn)}", value=chart_edge, ) return self diff --git a/metadata-ingestion/src/datahub/specific/datajob.py b/metadata-ingestion/src/datahub/specific/datajob.py index 984c555a53ffd1..acbc1a860968bd 100644 --- a/metadata-ingestion/src/datahub/specific/datajob.py +++ b/metadata-ingestion/src/datahub/specific/datajob.py @@ -1,5 +1,5 @@ import time -from typing import Dict, List, Optional, TypeVar, Union +from typing import Dict, List, Optional, Union from datahub.emitter.mcp_patch_builder import MetadataPatchProposal from datahub.metadata.schema_classes import ( @@ -21,8 +21,6 @@ from datahub.utilities.urns.tag_urn import TagUrn from datahub.utilities.urns.urn import Urn -T = TypeVar("T", bound=MetadataPatchProposal) - class DataJobPatchBuilder(MetadataPatchProposal): def __init__( @@ -164,7 +162,7 @@ def add_input_datajob(self, input: Union[Edge, Urn, str]) -> "DataJobPatchBuilde self._add_patch( DataJobInputOutput.ASPECT_NAME, "add", - path=f"/inputDatajobEdges/{MetadataPatchProposal.quote(input_urn)}", + path=f"/inputDatajobEdges/{self.quote(input_urn)}", value=input_edge, ) return self @@ -247,7 +245,7 @@ def add_input_dataset(self, input: Union[Edge, Urn, str]) -> "DataJobPatchBuilde self._add_patch( DataJobInputOutput.ASPECT_NAME, "add", - path=f"/inputDatasetEdges/{MetadataPatchProposal.quote(input_urn)}", + path=f"/inputDatasetEdges/{self.quote(input_urn)}", value=input_edge, ) return self @@ -265,7 +263,7 @@ def remove_input_dataset(self, input: Union[str, Urn]) -> "DataJobPatchBuilder": self._add_patch( DataJobInputOutput.ASPECT_NAME, "remove", - path=f"/inputDatasetEdges/{MetadataPatchProposal.quote(str(input))}", + path=f"/inputDatasetEdges/{self.quote(str(input))}", value={}, ) return self @@ -332,7 +330,7 @@ def add_output_dataset( self._add_patch( DataJobInputOutput.ASPECT_NAME, "add", - path=f"/outputDatasetEdges/{MetadataPatchProposal.quote(str(input))}", + path=f"/outputDatasetEdges/{self.quote(str(input))}", value=output_edge, ) return self @@ -350,7 +348,7 @@ def remove_output_dataset(self, output: Union[str, Urn]) -> "DataJobPatchBuilder self._add_patch( DataJobInputOutput.ASPECT_NAME, "remove", - path=f"/outputDatasetEdges/{output}", + path=f"/outputDatasetEdges/{self.quote(str(output))}", value={}, ) return self @@ -417,7 +415,7 @@ def add_input_dataset_field( self._add_patch( DataJobInputOutput.ASPECT_NAME, "add", - path=f"/inputDatasetFields/{MetadataPatchProposal.quote(input_urn)}", + path=f"/inputDatasetFields/{self.quote(input_urn)}", value=input_edge, ) return self @@ -438,7 +436,7 @@ def remove_input_dataset_field( self._add_patch( DataJobInputOutput.ASPECT_NAME, "remove", - path=f"/inputDatasetFields/{MetadataPatchProposal.quote(input_urn)}", + path=f"/inputDatasetFields/{self.quote(input_urn)}", value={}, ) return self @@ -505,7 +503,7 @@ def add_output_dataset_field( self._add_patch( DataJobInputOutput.ASPECT_NAME, "add", - path=f"/outputDatasetFields/{MetadataPatchProposal.quote(output_urn)}", + path=f"/outputDatasetFields/{self.quote(output_urn)}", value=output_edge, ) return self @@ -526,7 +524,7 @@ def remove_output_dataset_field( self._add_patch( DataJobInputOutput.ASPECT_NAME, "remove", - path=f"/outputDatasetFields/{MetadataPatchProposal.quote(output_urn)}", + path=f"/outputDatasetFields/{self.quote(output_urn)}", value={}, ) return self diff --git a/metadata-ingestion/src/datahub/specific/dataproduct.py b/metadata-ingestion/src/datahub/specific/dataproduct.py index c698c511fd9b58..6b7e695b4d57e7 100644 --- a/metadata-ingestion/src/datahub/specific/dataproduct.py +++ b/metadata-ingestion/src/datahub/specific/dataproduct.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, TypeVar, Union +from typing import Dict, List, Optional, Union from datahub.emitter.mcp_patch_builder import MetadataPatchProposal from datahub.metadata.schema_classes import ( @@ -18,8 +18,6 @@ from datahub.utilities.urns.tag_urn import TagUrn from datahub.utilities.urns.urn import Urn -T = TypeVar("T", bound=MetadataPatchProposal) - class DataProductPatchBuilder(MetadataPatchProposal): def __init__( diff --git a/metadata-ingestion/src/datahub/specific/dataset.py b/metadata-ingestion/src/datahub/specific/dataset.py index c5439e4d28ec6a..daf5975f74b42b 100644 --- a/metadata-ingestion/src/datahub/specific/dataset.py +++ b/metadata-ingestion/src/datahub/specific/dataset.py @@ -24,17 +24,17 @@ from datahub.utilities.urns.tag_urn import TagUrn from datahub.utilities.urns.urn import Urn -T = TypeVar("T", bound=MetadataPatchProposal) +_Parent = TypeVar("_Parent", bound=MetadataPatchProposal) -class FieldPatchHelper(Generic[T]): +class FieldPatchHelper(Generic[_Parent]): def __init__( self, - parent: T, + parent: _Parent, field_path: str, editable: bool = True, ) -> None: - self._parent: T = parent + self._parent: _Parent = parent self.field_path = field_path self.aspect_name = ( EditableSchemaMetadata.ASPECT_NAME @@ -83,7 +83,7 @@ def remove_term(self, term: Union[str, Urn]) -> "FieldPatchHelper": ) return self - def parent(self) -> T: + def parent(self) -> _Parent: return self._parent diff --git a/metadata-ingestion/src/datahub/specific/ownership.py b/metadata-ingestion/src/datahub/specific/ownership.py index c2a3874a3a33f3..b377a8814f38a0 100644 --- a/metadata-ingestion/src/datahub/specific/ownership.py +++ b/metadata-ingestion/src/datahub/specific/ownership.py @@ -7,15 +7,15 @@ OwnershipTypeClass, ) -T = TypeVar("T", bound=MetadataPatchProposal) +_Parent = TypeVar("_Parent", bound=MetadataPatchProposal) -class OwnershipPatchHelper(Generic[T]): - def __init__(self, parent: T) -> None: +class OwnershipPatchHelper(Generic[_Parent]): + def __init__(self, parent: _Parent) -> None: self._parent = parent self.aspect_field = OwnershipClass.ASPECT_NAME - def parent(self) -> T: + def parent(self) -> _Parent: return self._parent def add_owner(self, owner: OwnerClass) -> "OwnershipPatchHelper": diff --git a/metadata-ingestion/src/datahub/specific/structured_properties.py b/metadata-ingestion/src/datahub/specific/structured_properties.py index 6b2592bf1cbba2..17d896249c4746 100644 --- a/metadata-ingestion/src/datahub/specific/structured_properties.py +++ b/metadata-ingestion/src/datahub/specific/structured_properties.py @@ -6,20 +6,20 @@ make_structured_property_urn, ) -T = TypeVar("T", bound=MetadataPatchProposal) +_Parent = TypeVar("_Parent", bound=MetadataPatchProposal) -class StructuredPropertiesPatchHelper(Generic[T]): +class StructuredPropertiesPatchHelper(Generic[_Parent]): def __init__( self, - parent: T, + parent: _Parent, aspect_name: str = "structuredProperties", ) -> None: self.aspect_name = aspect_name self._parent = parent self.aspect_field = "properties" - def parent(self) -> T: + def parent(self) -> _Parent: return self._parent def set_property( @@ -33,7 +33,7 @@ def remove_property(self, key: str) -> "StructuredPropertiesPatchHelper": self._parent._add_patch( self.aspect_name, "remove", - path=f"/{self.aspect_field}/{make_structured_property_urn(key)}", + path=(self.aspect_field, make_structured_property_urn(key)), value={}, ) return self @@ -44,7 +44,7 @@ def add_property( self._parent._add_patch( self.aspect_name, "add", - path=f"/{self.aspect_field}/{make_structured_property_urn(key)}", + path=(self.aspect_field, make_structured_property_urn(key)), value=StructuredPropertyValueAssignmentClass( propertyUrn=make_structured_property_urn(key), values=value if isinstance(value, list) else [value], From 92b1cfa1949b653bd60a14f56e7813768715302f Mon Sep 17 00:00:00 2001 From: Aditya Malik Date: Wed, 28 Feb 2024 19:35:30 +0530 Subject: [PATCH 03/13] feat(ingest): Support for JSONL in s3 source with max_rows support (#9921) Co-authored-by: Aditya Co-authored-by: Harshal Sheth --- metadata-ingestion/docs/sources/gcs/README.md | 3 ++- metadata-ingestion/docs/sources/s3/README.md | 3 ++- .../src/datahub/ingestion/source/s3/source.py | 6 ++++- .../ingestion/source/schema_inference/json.py | 25 +++++++++++++++---- .../unit/data_lake/test_schema_inference.py | 13 ++++++++++ 5 files changed, 42 insertions(+), 8 deletions(-) diff --git a/metadata-ingestion/docs/sources/gcs/README.md b/metadata-ingestion/docs/sources/gcs/README.md index d6bb8147f076da..47a6b731dca1ec 100644 --- a/metadata-ingestion/docs/sources/gcs/README.md +++ b/metadata-ingestion/docs/sources/gcs/README.md @@ -21,13 +21,14 @@ Supported file types are as follows: - CSV - TSV +- JSONL - JSON - Parquet - Apache Avro Schemas for Parquet and Avro files are extracted as provided. -Schemas for schemaless formats (CSV, TSV, JSON) are inferred. For CSV and TSV files, we consider the first 100 rows by default, which can be controlled via the `max_rows` recipe parameter (see [below](#config-details)) +Schemas for schemaless formats (CSV, TSV, JSONL, JSON) are inferred. For CSV, TSV and JSONL files, we consider the first 100 rows by default, which can be controlled via the `max_rows` recipe parameter (see [below](#config-details)) JSON file schemas are inferred on the basis of the entire file (given the difficulty in extracting only the first few objects of the file), which may impact performance. We are working on using iterator-based JSON parsers to avoid reading in the entire JSON object. diff --git a/metadata-ingestion/docs/sources/s3/README.md b/metadata-ingestion/docs/sources/s3/README.md index 8d65e1cf8b943e..7944f78280a428 100644 --- a/metadata-ingestion/docs/sources/s3/README.md +++ b/metadata-ingestion/docs/sources/s3/README.md @@ -19,13 +19,14 @@ Supported file types are as follows: - CSV (*.csv) - TSV (*.tsv) +- JSONL (*.jsonl) - JSON (*.json) - Parquet (*.parquet) - Apache Avro (*.avro) Schemas for Parquet and Avro files are extracted as provided. -Schemas for schemaless formats (CSV, TSV, JSON) are inferred. For CSV and TSV files, we consider the first 100 rows by default, which can be controlled via the `max_rows` recipe parameter (see [below](#config-details)) +Schemas for schemaless formats (CSV, TSV, JSONL, JSON) are inferred. For CSV, TSV and JSONL files, we consider the first 100 rows by default, which can be controlled via the `max_rows` recipe parameter (see [below](#config-details)) JSON file schemas are inferred on the basis of the entire file (given the difficulty in extracting only the first few objects of the file), which may impact performance. We are working on using iterator-based JSON parsers to avoid reading in the entire JSON object. diff --git a/metadata-ingestion/src/datahub/ingestion/source/s3/source.py b/metadata-ingestion/src/datahub/ingestion/source/s3/source.py index 41fc5782352c94..43a1bcd06d3f32 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/s3/source.py +++ b/metadata-ingestion/src/datahub/ingestion/source/s3/source.py @@ -377,7 +377,7 @@ def read_file_spark(self, file: str, ext: str) -> Optional[DataFrame]: ignoreLeadingWhiteSpace=True, ignoreTrailingWhiteSpace=True, ) - elif ext.endswith(".json"): + elif ext.endswith(".json") or ext.endswith(".jsonl"): df = self.spark.read.json(file) elif ext.endswith(".avro"): try: @@ -441,6 +441,10 @@ def get_fields(self, table_data: TableData, path_spec: PathSpec) -> List: fields = csv_tsv.TsvInferrer( max_rows=self.source_config.max_rows ).infer_schema(file) + elif extension == ".jsonl": + fields = json.JsonInferrer( + max_rows=self.source_config.max_rows, format="jsonl" + ).infer_schema(file) elif extension == ".json": fields = json.JsonInferrer().infer_schema(file) elif extension == ".avro": diff --git a/metadata-ingestion/src/datahub/ingestion/source/schema_inference/json.py b/metadata-ingestion/src/datahub/ingestion/source/schema_inference/json.py index 251d136fe92ee7..1f2c73a2522d04 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/schema_inference/json.py +++ b/metadata-ingestion/src/datahub/ingestion/source/schema_inference/json.py @@ -1,3 +1,4 @@ +import itertools import logging from typing import IO, Dict, List, Type, Union @@ -33,14 +34,28 @@ class JsonInferrer(SchemaInferenceBase): + def __init__(self, max_rows: int = 100, format: str = "json"): + self.max_rows = max_rows + self.format = format + def infer_schema(self, file: IO[bytes]) -> List[SchemaField]: - try: - datastore = ujson.load(file) - except ujson.JSONDecodeError as e: - logger.info(f"Got ValueError: {e}. Retry with jsonlines") + if self.format == "jsonl": file.seek(0) reader = jsl.Reader(file) - datastore = [obj for obj in reader.iter(type=dict, skip_invalid=True)] + datastore = [ + obj + for obj in itertools.islice( + reader.iter(type=dict, skip_invalid=True), self.max_rows + ) + ] + else: + try: + datastore = ujson.load(file) + except ujson.JSONDecodeError as e: + logger.info(f"Got ValueError: {e}. Retry with jsonlines") + file.seek(0) + reader = jsl.Reader(file) + datastore = [obj for obj in reader.iter(type=dict, skip_invalid=True)] if not isinstance(datastore, list): datastore = [datastore] diff --git a/metadata-ingestion/tests/unit/data_lake/test_schema_inference.py b/metadata-ingestion/tests/unit/data_lake/test_schema_inference.py index de88deec9b9cb0..a1ef02c27ea540 100644 --- a/metadata-ingestion/tests/unit/data_lake/test_schema_inference.py +++ b/metadata-ingestion/tests/unit/data_lake/test_schema_inference.py @@ -74,6 +74,19 @@ def test_infer_schema_tsv(): assert_field_types_match(fields, expected_field_types) +def test_infer_schema_jsonl(): + with tempfile.TemporaryFile(mode="w+b") as file: + file.write( + bytes(test_table.to_json(orient="records", lines=True), encoding="utf-8") + ) + file.seek(0) + + fields = json.JsonInferrer(max_rows=100, format="jsonl").infer_schema(file) + + assert_field_paths_match(fields, expected_field_paths) + assert_field_types_match(fields, expected_field_types) + + def test_infer_schema_json(): with tempfile.TemporaryFile(mode="w+b") as file: file.write(bytes(test_table.to_json(orient="records"), encoding="utf-8")) From 1736edf8f5d8dd36ed005e19a2ded80b645099a7 Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Wed, 28 Feb 2024 06:06:33 -0800 Subject: [PATCH 04/13] feat(ingest): fix bugs in SqlParsingAggregator (#9926) --- .../datahub/sql_parsing/schema_resolver.py | 2 +- .../sql_parsing/sql_parsing_aggregator.py | 19 +++++++++----- .../test_known_lineage_mapping.json | 9 +++---- .../aggregator_goldens/test_temp_table.json | 3 +-- .../unit/sql_parsing/test_schemaresolver.py | 25 +++++++++++++++++++ 5 files changed, 43 insertions(+), 15 deletions(-) diff --git a/metadata-ingestion/src/datahub/sql_parsing/schema_resolver.py b/metadata-ingestion/src/datahub/sql_parsing/schema_resolver.py index ce43fb5da57c08..ec52e839212c54 100644 --- a/metadata-ingestion/src/datahub/sql_parsing/schema_resolver.py +++ b/metadata-ingestion/src/datahub/sql_parsing/schema_resolver.py @@ -80,7 +80,7 @@ def get_urns(self) -> Set[str]: def schema_count(self) -> int: return int( self._schema_cache.sql_query( - f"SELECT COUNT(*) FROM {self._schema_cache.tablename} WHERE is_missing" + f"SELECT COUNT(*) FROM {self._schema_cache.tablename} WHERE NOT is_missing" )[0][0] ) diff --git a/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py b/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py index cecaef33efcd7b..49b58ddd22b83f 100644 --- a/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py +++ b/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py @@ -115,6 +115,7 @@ class SqlAggregatorReport(Report): _aggregator: "SqlParsingAggregator" query_log_path: Optional[str] = None + # Observed queries. num_observed_queries: int = 0 num_observed_queries_failed: int = 0 num_observed_queries_column_timeout: int = 0 @@ -123,6 +124,7 @@ class SqlAggregatorReport(Report): default_factory=LossyList ) + # Views. num_view_definitions: int = 0 num_views_failed: int = 0 num_views_column_timeout: int = 0 @@ -131,28 +133,30 @@ class SqlAggregatorReport(Report): default_factory=LossyDict ) + # Other lineage loading metrics. num_known_query_lineage: int = 0 num_known_mapping_lineage: int = 0 num_table_renames: int = 0 - num_queries_with_temp_tables_in_session: int = 0 - - num_unique_query_fingerprints: Optional[int] = None - - # Lineage-related. - num_urns_with_lineage: Optional[int] = None + # Temp tables. num_temp_sessions: Optional[int] = None num_inferred_temp_schemas: Optional[int] = None + num_queries_with_temp_tables_in_session: int = 0 queries_with_temp_upstreams: LossyDict[QueryId, LossyList] = dataclasses.field( default_factory=LossyDict ) + # Lineage-related. + schema_resolver_count: Optional[int] = None + num_unique_query_fingerprints: Optional[int] = None + num_urns_with_lineage: Optional[int] = None num_queries_entities_generated: int = 0 # Usage-related. usage_skipped_missing_timestamp: int = 0 def compute_stats(self) -> None: + self.schema_resolver_count = self._aggregator._schema_resolver.schema_count() self.num_unique_query_fingerprints = len(self._aggregator._query_map) self.num_urns_with_lineage = len(self._aggregator._lineage_map) @@ -865,6 +869,9 @@ def _gen_lineage_for_downstream( confidenceScore=queries_map[query_id].confidence_score, ) ) + upstream_aspect.fineGrainedLineages = ( + upstream_aspect.fineGrainedLineages or None + ) yield MetadataChangeProposalWrapper( entityUrn=downstream_urn, diff --git a/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_known_lineage_mapping.json b/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_known_lineage_mapping.json index ab210c6f701b3f..86d5ef34a87569 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_known_lineage_mapping.json +++ b/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_known_lineage_mapping.json @@ -19,8 +19,7 @@ "dataset": "urn:li:dataset:(urn:li:dataPlatform:s3,bucket1/key1,PROD)", "type": "COPY" } - ], - "fineGrainedLineages": [] + ] } } }, @@ -44,8 +43,7 @@ "dataset": "urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.bar,PROD)", "type": "COPY" } - ], - "fineGrainedLineages": [] + ] } } }, @@ -69,8 +67,7 @@ "dataset": "urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.foo,PROD)", "type": "COPY" } - ], - "fineGrainedLineages": [] + ] } } } diff --git a/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_temp_table.json b/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_temp_table.json index b93e7e0f5260fe..5e61fb2b6a20f2 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_temp_table.json +++ b/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_temp_table.json @@ -205,8 +205,7 @@ "type": "TRANSFORMED", "query": "urn:li:query:3e85e6f353c7fa33d6514cb090482852064d23df6491c9a8ae28be0d990a3c71" } - ], - "fineGrainedLineages": [] + ] } } }, diff --git a/metadata-ingestion/tests/unit/sql_parsing/test_schemaresolver.py b/metadata-ingestion/tests/unit/sql_parsing/test_schemaresolver.py index 5786c135a8d8c5..5a33034f274dc2 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/test_schemaresolver.py +++ b/metadata-ingestion/tests/unit/sql_parsing/test_schemaresolver.py @@ -1,6 +1,31 @@ from datahub.sql_parsing.schema_resolver import SchemaResolver, _TableName +def test_basic_schema_resolver(): + schema_resolver = SchemaResolver( + platform="redshift", + env="PROD", + graph=None, + ) + + schema_resolver.add_raw_schema_info( + urn="urn:li:dataset:(urn:li:dataPlatform:redshift,my_db.public.test_table,PROD)", + schema_info={"name": "STRING"}, + ) + + urn, schema = schema_resolver.resolve_table( + _TableName(database="my_db", db_schema="public", table="test_table") + ) + assert ( + urn + == "urn:li:dataset:(urn:li:dataPlatform:redshift,my_db.public.test_table,PROD)" + ) + assert schema + assert schema["name"] + + assert schema_resolver.schema_count() == 1 + + def test_get_urn_for_table_lowercase(): schema_resolver = SchemaResolver( platform="mssql", From 3b735556fa2ddb7f29fd8b7f52e6ea6be21f92a0 Mon Sep 17 00:00:00 2001 From: Aseem Bansal Date: Wed, 28 Feb 2024 20:05:30 +0530 Subject: [PATCH 05/13] fix(ui/schema): error handling add (#9952) --- .../src/app/entity/shared/tabs/Dataset/Schema/SchemaTab.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datahub-web-react/src/app/entity/shared/tabs/Dataset/Schema/SchemaTab.tsx b/datahub-web-react/src/app/entity/shared/tabs/Dataset/Schema/SchemaTab.tsx index 28dc3ba5c6ce54..e869b227fd5149 100644 --- a/datahub-web-react/src/app/entity/shared/tabs/Dataset/Schema/SchemaTab.tsx +++ b/datahub-web-react/src/app/entity/shared/tabs/Dataset/Schema/SchemaTab.tsx @@ -78,7 +78,7 @@ export const SchemaTab = ({ properties }: { properties?: any }) => { const hasProperties = useMemo( () => - entityWithSchema?.schemaMetadata?.fields.some( + entityWithSchema?.schemaMetadata?.fields?.some( (schemaField) => !!schemaField.schemaFieldEntity?.structuredProperties?.properties?.length, ), [entityWithSchema], From b45c04f286fd1c07c842dec6cb72b1550177dbf8 Mon Sep 17 00:00:00 2001 From: Kunal-kankriya <127090035+Kunal-kankriya@users.noreply.github.com> Date: Wed, 28 Feb 2024 21:19:31 +0530 Subject: [PATCH 06/13] fix(tests): test_group_upsert smoke test updated (#9888) --- smoke-test/tests/cli/user_groups_cmd/test_group_cmd.py | 3 +-- smoke-test/tests/consistency_utils.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) 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 555687c98ed3e2..feabcc5f9d6558 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 @@ -3,7 +3,6 @@ import tempfile from typing import Any, Dict, Iterable, List -import pytest import yaml from click.testing import CliRunner, Result from datahub.api.entities.corpgroup.corpgroup import CorpGroup @@ -42,6 +41,7 @@ def gen_datahub_groups(num_groups: int) -> Iterable[CorpGroup]: description=f"The Group {i}", picture_link=f"https://images.google.com/group{i}.jpg", slack=f"@group{i}", + owners=["user1"], members=["user2"], ) yield group @@ -79,7 +79,6 @@ def get_group_membership(user_urn: str) -> List[str]: return [entity.urn for entity in entities] -@pytest.mark.skip(reason="Functionality and test needs to be validated for correctness") def test_group_upsert(wait_for_healthchecks: Any) -> None: num_groups: int = 10 for i, datahub_group in enumerate(gen_datahub_groups(num_groups)): diff --git a/smoke-test/tests/consistency_utils.py b/smoke-test/tests/consistency_utils.py index 6a04efb4228c68..5ffc642d494697 100644 --- a/smoke-test/tests/consistency_utils.py +++ b/smoke-test/tests/consistency_utils.py @@ -31,7 +31,7 @@ def wait_for_writes_to_sync(max_timeout_in_sec: int = 120) -> None: result = str(completed_process.stdout) lines = result.splitlines() lag_values = [int(line) for line in lines if line != ""] - maximum_lag = max(lag_values) + maximum_lag = max(lag_values) if lag_values else 0 if maximum_lag == 0: lag_zero = True From 55bc955304c4c192c04a0393a47355a295f5770a Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Wed, 28 Feb 2024 09:49:58 -0800 Subject: [PATCH 07/13] feat(ci): skip smoke-test lint if there are no changes (#9945) --- .github/actions/ci-optimization/action.yml | 7 ++- .github/workflows/docker-unified.yml | 50 ++++++++++------------ 2 files changed, 29 insertions(+), 28 deletions(-) diff --git a/.github/actions/ci-optimization/action.yml b/.github/actions/ci-optimization/action.yml index f6160fdbcff675..6bb389318c8574 100644 --- a/.github/actions/ci-optimization/action.yml +++ b/.github/actions/ci-optimization/action.yml @@ -35,6 +35,9 @@ outputs: elasticsearch-setup-change: description: "Elasticsearch setup docker change" value: ${{ steps.filter.outputs.elasticsearch-setup == 'true' }} + smoke-test-change: + description: "Smoke test change" + value: ${{ steps.filter.outputs.smoke-test == 'true' }} runs: using: "composite" steps: @@ -76,4 +79,6 @@ runs: postgres-setup: - "docker/postgres-setup/**" elasticsearch-setup: - - "docker/elasticsearch-setup/**" \ No newline at end of file + - "docker/elasticsearch-setup/**" + smoke-test: + - "smoke-test/**" diff --git a/.github/workflows/docker-unified.yml b/.github/workflows/docker-unified.yml index 7c4742dbf8fc5a..38ffa3484c0bff 100644 --- a/.github/workflows/docker-unified.yml +++ b/.github/workflows/docker-unified.yml @@ -29,35 +29,8 @@ env: DATAHUB_INGESTION_IMAGE: "acryldata/datahub-ingestion" jobs: - check_lint: - runs-on: ubuntu-latest - steps: - - name: Check out the repo - uses: acryldata/sane-checkout-action@v3 - - uses: actions/setup-python@v4 - with: - python-version: "3.10" - cache: "pip" - - uses: actions/cache@v4 - with: - path: | - ~/.cache/uv - key: ${{ runner.os }}-uv-${{ hashFiles('**/requirements.txt') }} - - name: Set up JDK 17 - uses: actions/setup-java@v3 - with: - distribution: "zulu" - java-version: 17 - - uses: gradle/gradle-build-action@v2 - - name: Ensure packages are correct - run: | - python ./.github/scripts/check_python_package.py - - name: Run lint on smoke test - run: | - ./gradlew :smoke-test:lint setup: runs-on: ubuntu-latest - needs: check_lint outputs: tag: ${{ steps.tag.outputs.tag }} slim_tag: ${{ steps.tag.outputs.slim_tag }} @@ -106,6 +79,29 @@ jobs: echo "publish=${{ env.ENABLE_PUBLISH }}" >> $GITHUB_OUTPUT - uses: ./.github/actions/ci-optimization id: ci-optimize + - uses: actions/setup-python@v4 + if: ${{ steps.ci-optimize.outputs.smoke-test-change == 'true' }} + with: + python-version: "3.10" + cache: "pip" + - uses: actions/cache@v4 + if: ${{ steps.ci-optimize.outputs.smoke-test-change == 'true' }} + with: + path: | + ~/.cache/uv + key: ${{ runner.os }}-uv-${{ hashFiles('**/requirements.txt') }} + - name: Set up JDK 17 + uses: actions/setup-java@v3 + if: ${{ steps.ci-optimize.outputs.smoke-test-change == 'true' }} + with: + distribution: "zulu" + java-version: 17 + - uses: gradle/gradle-build-action@v2 + - name: Run lint on smoke test + if: ${{ steps.ci-optimize.outputs.smoke-test-change == 'true' }} + run: | + python ./.github/scripts/check_python_package.py + ./gradlew :smoke-test:lint gms_build: name: Build and Push DataHub GMS Docker Image From ed10a8d8cca3b17e982db6d14ea435833c5a87ea Mon Sep 17 00:00:00 2001 From: david-leifker <114954101+david-leifker@users.noreply.github.com> Date: Wed, 28 Feb 2024 16:57:26 -0600 Subject: [PATCH 08/13] feat(search): search access controls (#9892) Co-authored-by: Chris Collins --- build.gradle | 6 +- datahub-frontend/app/auth/AuthModule.java | 38 +- .../app/config/ConfigurationProvider.java | 4 + datahub-frontend/play.gradle | 4 +- .../linkedin/datahub/graphql/Constants.java | 1 - .../datahub/graphql/GmsGraphQLEngine.java | 91 ++- .../datahub/graphql/GmsGraphQLEngineArgs.java | 2 + .../datahub/graphql/QueryContext.java | 6 + .../GetMetadataAnalyticsResolver.java | 21 +- .../authorization/AuthorizationUtils.java | 15 + .../auth/ListAccessTokensResolver.java | 8 +- .../resolvers/chart/BrowseV2Resolver.java | 5 +- .../container/ContainerEntitiesResolver.java | 4 +- .../ListDataProductAssetsResolver.java | 11 +- .../UpdateDeprecationResolver.java | 2 +- .../domain/DomainEntitiesResolver.java | 4 +- .../resolvers/domain/ListDomainsResolver.java | 6 +- .../glossary/CreateGlossaryTermResolver.java | 2 +- .../GetRootGlossaryNodesResolver.java | 4 +- .../GetRootGlossaryTermsResolver.java | 4 +- .../resolvers/group/EntityCountsResolver.java | 4 +- .../resolvers/group/ListGroupsResolver.java | 8 +- .../health/EntityHealthResolver.java | 2 +- .../incident/EntityIncidentsResolver.java | 4 +- ...estionSourceExecutionRequestsResolver.java | 4 +- .../ingest/secret/ListSecretsResolver.java | 8 +- .../source/ListIngestionSourcesResolver.java | 8 +- .../resolvers/jobs/DataJobRunsResolver.java | 4 +- .../resolvers/jobs/EntityRunsResolver.java | 4 +- .../load/EntityLineageResultResolver.java | 77 ++- .../resolvers/mutate/util/DomainUtils.java | 4 +- .../ownership/ListOwnershipTypesResolver.java | 6 +- .../policy/ListPoliciesResolver.java | 2 +- .../resolvers/post/ListPostsResolver.java | 6 +- .../resolvers/query/ListQueriesResolver.java | 9 +- .../ListRecommendationsResolver.java | 4 +- .../role/CreateInviteTokenResolver.java | 3 +- .../role/GetInviteTokenResolver.java | 3 +- .../resolvers/role/ListRolesResolver.java | 6 +- .../AggregateAcrossEntitiesResolver.java | 3 +- .../search/GetQuickFiltersResolver.java | 21 +- .../search/ScrollAcrossEntitiesResolver.java | 11 +- .../search/ScrollAcrossLineageResolver.java | 11 +- .../search/SearchAcrossEntitiesResolver.java | 5 +- .../search/SearchAcrossLineageResolver.java | 10 +- .../resolvers/search/SearchResolver.java | 7 +- .../resolvers/test/ListTestsResolver.java | 8 +- .../resolvers/user/ListUsersResolver.java | 8 +- .../view/ListGlobalViewsResolver.java | 6 +- .../resolvers/view/ListMyViewsResolver.java | 6 +- .../graphql/types/chart/ChartType.java | 15 +- .../mappers/SearchFlagsInputMapper.java | 6 + .../common/mappers/UrnToEntityMapper.java | 6 + .../types/container/ContainerType.java | 9 +- .../types/corpgroup/CorpGroupType.java | 9 +- .../graphql/types/corpuser/CorpUserType.java | 9 +- .../types/dashboard/DashboardType.java | 16 +- .../graphql/types/dataflow/DataFlowType.java | 16 +- .../graphql/types/datajob/DataJobType.java | 15 +- .../DataPlatformInstanceType.java | 6 +- .../types/dataproduct/DataProductType.java | 2 +- .../graphql/types/dataset/DatasetType.java | 16 +- .../graphql/types/domain/DomainType.java | 2 +- .../types/entitytype/EntityTypeMapper.java | 1 + .../types/glossary/GlossaryTermType.java | 15 +- .../types/mlmodel/MLFeatureTableType.java | 15 +- .../graphql/types/mlmodel/MLFeatureType.java | 9 +- .../types/mlmodel/MLModelGroupType.java | 15 +- .../graphql/types/mlmodel/MLModelType.java | 15 +- .../types/mlmodel/MLPrimaryKeyType.java | 8 +- .../graphql/types/notebook/NotebookType.java | 15 +- .../types/restricted/RestrictedMapper.java | 32 + .../types/restricted/RestrictedType.java | 103 +++ .../graphql/types/rolemetadata/RoleType.java | 8 +- .../datahub/graphql/types/tag/TagType.java | 8 +- .../src/main/resources/entity.graphql | 31 + .../src/main/resources/search.graphql | 10 + .../linkedin/datahub/graphql/TestUtils.java | 21 +- .../auth/ListAccessTokensResolverTest.java | 7 +- .../browse/BrowseV2ResolverTest.java | 6 +- .../ContainerEntitiesResolverTest.java | 16 +- .../domain/CreateDomainResolverTest.java | 13 +- .../domain/DeleteDomainResolverTest.java | 9 +- .../domain/DomainEntitiesResolverTest.java | 8 +- .../domain/ListDomainsResolverTest.java | 23 +- .../domain/MoveDomainResolverTest.java | 5 +- .../CreateGlossaryTermResolverTest.java | 9 +- .../GetRootGlossaryNodesResolverTest.java | 8 +- .../GetRootGlossaryTermsResolverTest.java | 8 +- .../glossary/UpdateNameResolverTest.java | 5 +- .../incident/EntityIncidentsResolverTest.java | 7 +- .../resolvers/ingest/IngestTestUtils.java | 3 + ...onSourceExecutionRequestsResolverTest.java | 27 +- .../secret/ListSecretsResolverTest.java | 15 +- .../ListIngestionSourceResolverTest.java | 15 +- .../ListOwnershipTypesResolverTest.java | 12 +- .../resolvers/post/ListPostsResolverTest.java | 11 +- .../query/ListQueriesResolverTest.java | 18 +- .../role/CreateInviteTokenResolverTest.java | 3 +- .../role/GetInviteTokenResolverTest.java | 7 +- .../resolvers/role/ListRolesResolverTest.java | 10 +- .../AggregateAcrossEntitiesResolverTest.java | 11 +- .../AutoCompleteForMultipleResolverTest.java | 13 +- .../search/GetQuickFiltersResolverTest.java | 10 +- .../SearchAcrossEntitiesResolverTest.java | 15 +- .../SearchAcrossLineageResolverTest.java | 6 +- .../resolvers/search/SearchResolverTest.java | 20 +- .../resolvers/test/ListTestsResolverTest.java | 26 +- .../view/ListGlobalViewsResolverTest.java | 18 +- .../view/ListMyViewsResolverTest.java | 23 +- .../linkedin/datahub/upgrade/UpgradeCli.java | 35 +- .../config/BackfillBrowsePathsV2Config.java | 11 +- .../config/BackfillOwnershipTypesConfig.java | 28 + .../config/BackfillPolicyFieldsConfig.java | 8 +- .../upgrade/config/BuildIndicesConfig.java | 5 +- .../upgrade/config/CleanIndicesConfig.java | 7 +- .../ReindexDataJobViaNodesCLLConfig.java | 7 +- .../upgrade/config/SystemUpdateCondition.java | 36 +- .../upgrade/config/SystemUpdateConfig.java | 70 +- .../upgrade/system/BlockingSystemUpgrade.java | 5 + .../system/NonBlockingSystemUpgrade.java | 5 + .../datahub/upgrade/system/SystemUpdate.java | 70 +- .../upgrade/system/SystemUpdateBlocking.java | 16 + .../system/SystemUpdateNonBlocking.java | 16 + .../BackfillBrowsePathsV2.java | 10 +- .../BackfillBrowsePathsV2Step.java | 24 +- .../system/elasticsearch/BuildIndices.java | 4 +- .../system/elasticsearch/CleanIndices.java | 4 +- .../system/ownershiptypes/OwnershipTypes.java | 41 ++ .../ownershiptypes/OwnershipTypesStep.java | 276 ++++++++ .../BackfillPolicyFields.java | 10 +- .../BackfillPolicyFieldsStep.java | 25 +- .../ReindexDataJobViaNodesCLL.java | 8 +- .../ReindexDataJobViaNodesCLLStep.java | 5 +- .../upgrade/UpgradeCliApplicationTest.java | 5 +- .../src/app/buildEntityRegistry.ts | 2 + .../entity/restricted/RestrictedEntity.tsx | 90 +++ .../restricted/RestrictedEntityProfile.tsx | 30 + .../containers/profile/header/EntityName.tsx | 2 +- .../PlatformContent/PlatformContentView.tsx | 8 +- .../src/app/lineage/LineageEntityNode.tsx | 64 +- .../src/app/lineage/LineageExplorer.tsx | 8 +- datahub-web-react/src/graphql/lineage.graphql | 4 + datahub-web-react/src/images/restricted.svg | 2 + .../profiles/docker-compose.prerequisites.yml | 4 +- .../metadata/aspect/batch/ChangeMCP.java | 12 +- .../metadata/aspect/hooks/OwnerTypeMap.java | 111 +++ .../metadata/models/FieldSpecUtils.java | 13 +- .../RelationshipFieldSpecExtractor.java | 3 +- .../models/SearchableFieldSpecExtractor.java | 6 +- .../annotation/SearchableAnnotation.java | 5 +- .../aspect/hooks/OwnerTypeMapTest.java | 220 ++++++ gradle.properties | 4 +- .../ingestion/IngestionScheduler.java | 132 ++-- .../ingestion/IngestionSchedulerTest.java | 46 +- .../java/com/linkedin/metadata/Constants.java | 4 + .../ConjunctivePrivilegeGroup.java | 8 +- .../SearchAuthorizationConfiguration.java | 24 + .../auth/authorization/Authorizer.java | 46 +- .../test/resources/test-avro2pegasus-mae.json | 1 + .../test/resources/test-avro2pegasus-mce.json | 1 + .../datahub/ingestion/source/csv_enricher.py | 2 +- .../tests/unit/test_rest_sink.py | 2 +- metadata-io/build.gradle | 2 + .../client/EntityClientAspectRetriever.java | 3 +- .../metadata/client/JavaEntityClient.java | 160 +++-- .../client/SystemJavaEntityClient.java | 17 +- .../entity/ebean/batch/AspectsBatchImpl.java | 3 + .../candidatesource/MostPopularSource.java | 36 +- .../candidatesource/RecentlyEditedSource.java | 12 +- .../candidatesource/RecentlyViewedSource.java | 7 +- .../search/EntityLineageResultCacheKey.java | 5 +- .../metadata/search/LineageSearchService.java | 90 ++- .../metadata/search/SearchService.java | 50 +- .../search/cache/CacheableSearcher.java | 13 +- .../search/cache/EntityDocCountCache.java | 40 +- .../client/CachingEntitySearchService.java | 106 +-- .../elasticsearch/ElasticSearchService.java | 155 ++++- .../indexbuilder/MappingsBuilder.java | 5 +- .../indexbuilder/ReindexConfig.java | 51 +- .../elasticsearch/query/ESBrowseDAO.java | 116 +++- .../elasticsearch/query/ESSearchDAO.java | 62 +- .../request/AggregationQueryBuilder.java | 5 +- .../request/AutocompleteRequestHandler.java | 21 +- .../query/request/SearchFieldConfig.java | 3 +- .../query/request/SearchRequestHandler.java | 109 +-- .../SearchDocumentTransformer.java | 23 +- .../search/utils/ESAccessControlUtil.java | 260 +++++++ .../metadata/search/utils/ESUtils.java | 49 +- .../metadata/search/utils/SearchUtils.java | 8 + .../elastic/indexbuilder/MappingsBuilder.java | 3 + .../metadata/client/JavaEntityClientTest.java | 4 + .../search/SearchGraphServiceTestBase.java | 2 +- .../RecommendationsServiceTest.java | 15 +- ...ySearchAggregationCandidateSourceTest.java | 51 +- .../RecommendationUtilsTest.java | 15 +- .../candidatesource/TestSource.java | 6 +- .../LineageSearchResultCacheKeyTest.java | 12 +- .../search/LineageServiceTestBase.java | 75 +- .../search/SearchServiceTestBase.java | 75 +- .../metadata/search/TestEntityTestBase.java | 209 +++++- .../search/cache/CacheableSearcherTest.java | 54 +- .../LineageDataFixtureElasticSearchTest.java | 8 + .../search/fixtures/GoldenTestBase.java | 49 +- .../fixtures/LineageDataFixtureTestBase.java | 29 +- .../fixtures/SampleDataFixtureTestBase.java | 186 +++-- .../LineageDataFixtureOpenSearchTest.java | 10 + .../search/query/SearchDAOTestBase.java | 22 +- .../AutocompleteRequestHandlerTest.java | 11 +- .../request/SearchRequestHandlerTest.java | 44 +- .../search/utils/ESAccessControlUtilTest.java | 646 ++++++++++++++++++ .../search/utils/SearchUtilsTest.java | 56 +- .../SampleDataFixtureConfiguration.java | 4 + .../SearchLineageFixtureConfiguration.java | 5 + .../test/search/SearchTestUtils.java | 107 +-- .../config/SearchCommonTestConfiguration.java | 7 + .../kafka/MaeConsumerApplication.java | 3 +- .../kafka/MaeConsumerApplicationTest.java | 8 +- .../hook/siblings/SiblingAssociationHook.java | 46 +- .../siblings/SiblingAssociationHookTest.java | 51 +- .../spring/MCLSpringTestConfiguration.java | 18 + .../kafka/MceConsumerApplication.java | 3 +- ...eConsumerApplicationTestConfiguration.java | 11 +- .../pegasus/com/linkedin/common/Ownership.pdl | 11 + .../linkedin/metadata/query/SearchFlags.pdl | 10 + .../linkedin/metadata/search/SearchEntity.pdl | 6 + .../src/main/resources/entity-registry.yml | 11 +- metadata-operation-context/build.gradle | 15 + .../metadata/context/ActorContext.java | 71 ++ .../metadata/context/AuthorizerContext.java | 28 + .../metadata/context/ContextInterface.java | 21 + .../context/EntityRegistryContext.java | 22 + .../metadata/context/OperationContext.java | 285 ++++++++ .../context/OperationContextConfig.java | 24 + .../metadata/context/SearchContext.java | 86 +++ .../context/TestOperationContexts.java | 69 ++ .../metadata/context/ActorContextTest.java | 129 ++++ .../context/OperationContextTest.java | 67 ++ .../metadata/context/SearchContextTest.java | 79 +++ .../AuthorizationConfiguration.java | 3 + .../invite/InviteTokenService.java | 16 +- .../authorization/AuthorizerChain.java | 25 + .../authorization/DataHubAuthorizer.java | 62 +- .../datahub/authorization/PolicyEngine.java | 7 +- .../datahub/authorization/PolicyFetcher.java | 40 +- .../invite/InviteTokenServiceTest.java | 58 +- .../authorization/DataHubAuthorizerTest.java | 29 +- .../src/main/resources/application.yml | 10 + .../auth/DataHubAuthorizerFactory.java | 8 +- .../auth/RestrictedServiceFactory.java | 28 + .../SystemOperationContextFactory.java | 44 ++ .../entityclient/JavaEntityClientFactory.java | 10 +- .../RestliEntityClientFactory.java | 8 +- .../gms/factory/form/FormServiceFactory.java | 9 +- .../factory/graphql/GraphQLEngineFactory.java | 6 + .../ingestion/IngestionSchedulerFactory.java | 7 +- .../DomainsCandidateSourceFactory.java | 4 +- .../gms/factory/usage/UsageClientFactory.java | 6 +- .../factories/BootstrapManagerFactory.java | 11 +- .../boot/steps/BackfillBrowsePathsV2Step.java | 14 +- .../boot/steps/IngestPoliciesStep.java | 4 +- .../boot/steps/RestoreGlossaryIndices.java | 33 +- .../gms/factory/search/CacheTest.java | 17 +- .../steps/BackfillBrowsePathsV2StepTest.java | 20 +- .../steps/RestoreGlossaryIndicesTest.java | 73 +- .../datahub/graphql/GraphQLController.java | 11 +- .../datahub/graphql/SpringQueryContext.java | 25 +- .../v2/delegates/EntityApiDelegateImpl.java | 17 +- .../JavaSpring/apiController.mustache | 7 +- .../OpenAPIEntityTestConfiguration.java | 9 +- .../elastic/OperationsController.java | 18 +- .../v2/controller/EntityController.java | 15 +- ...com.linkedin.entity.entities.restspec.json | 8 + .../com.linkedin.entity.aspects.snapshot.json | 18 + ...com.linkedin.entity.entities.snapshot.json | 46 ++ .../com.linkedin.entity.runs.snapshot.json | 18 + ...nkedin.operations.operations.snapshot.json | 18 + ...m.linkedin.platform.platform.snapshot.json | 18 + metadata-service/restli-client/build.gradle | 1 + .../linkedin/entity/client/EntityClient.java | 148 ++-- .../entity/client/EntityClientCache.java | 137 ++-- .../entity/client/RestliEntityClient.java | 142 ++-- .../entity/client/SystemEntityClient.java | 66 +- .../client/SystemRestliEntityClient.java | 14 +- .../java/com/linkedin/usage/UsageClient.java | 14 +- .../com/linkedin/usage/UsageClientCache.java | 23 +- .../resources/entity/EntityResource.java | 93 ++- .../linkedin/metadata/entity/AspectUtils.java | 9 - .../RecommendationsService.java | 15 +- .../DomainsCandidateSource.java | 4 +- .../EntitySearchAggregationSource.java | 5 +- .../RecentlySearchedSource.java | 7 +- .../candidatesource/RecommendationSource.java | 19 +- .../candidatesource/RecommendationUtils.java | 6 +- .../candidatesource/TopPlatformsSource.java | 3 +- .../candidatesource/TopTagsSource.java | 4 +- .../candidatesource/TopTermsSource.java | 4 +- .../ranker/RecommendationModuleRanker.java | 8 +- .../ranker/SimpleRecommendationRanker.java | 6 +- .../metadata/search/EntitySearchService.java | 42 +- .../metadata/service/FormService.java | 11 +- .../metadata/service/LineageService.java | 1 + .../metadata/service/RestrictedService.java | 29 + .../SearchBasedFormAssignmentManager.java | 13 +- .../util/SearchBasedFormAssignmentRunner.java | 8 +- .../service/RestrictedServiceTest.java | 37 + .../gms/servlet/ConfigSearchExport.java | 17 +- .../war/src/main/resources/boot/policies.json | 6 + .../authorization/PoliciesConfig.java | 15 + .../metadata/utils/AuditStampUtils.java | 22 +- .../linkedin/metadata/utils/SearchUtil.java | 30 +- .../elasticsearch/IndexConventionImpl.java | 2 + .../IndexConventionImplTest.java | 4 +- settings.gradle | 1 + smoke-test/cypress-dev.sh | 2 + .../cli/user_groups_cmd/test_group_cmd.py | 5 + smoke-test/tests/consistency_utils.py | 35 +- 317 files changed, 6771 insertions(+), 2112 deletions(-) create mode 100644 datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/restricted/RestrictedMapper.java create mode 100644 datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/restricted/RestrictedType.java create mode 100644 datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillOwnershipTypesConfig.java create mode 100644 datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/BlockingSystemUpgrade.java create mode 100644 datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/NonBlockingSystemUpgrade.java create mode 100644 datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdateBlocking.java create mode 100644 datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdateNonBlocking.java rename datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/{entity/steps => browsepaths}/BackfillBrowsePathsV2.java (66%) rename datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/{entity/steps => browsepaths}/BackfillBrowsePathsV2Step.java (92%) create mode 100644 datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/ownershiptypes/OwnershipTypes.java create mode 100644 datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/ownershiptypes/OwnershipTypesStep.java rename datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/{entity/steps => policyfields}/BackfillPolicyFields.java (66%) rename datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/{entity/steps => policyfields}/BackfillPolicyFieldsStep.java (92%) rename datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/{via => vianodes}/ReindexDataJobViaNodesCLL.java (80%) rename datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/{via => vianodes}/ReindexDataJobViaNodesCLLStep.java (94%) create mode 100644 datahub-web-react/src/app/entity/restricted/RestrictedEntity.tsx create mode 100644 datahub-web-react/src/app/entity/restricted/RestrictedEntityProfile.tsx create mode 100644 datahub-web-react/src/images/restricted.svg create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/hooks/OwnerTypeMap.java create mode 100644 entity-registry/src/test/java/com/linkedin/metadata/aspect/hooks/OwnerTypeMapTest.java create mode 100644 metadata-auth/auth-api/src/main/java/com/datahub/authorization/config/SearchAuthorizationConfiguration.java create mode 100644 metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESAccessControlUtil.java create mode 100644 metadata-io/src/test/java/com/linkedin/metadata/search/utils/ESAccessControlUtilTest.java create mode 100644 metadata-operation-context/build.gradle create mode 100644 metadata-operation-context/src/main/java/io/datahubproject/metadata/context/ActorContext.java create mode 100644 metadata-operation-context/src/main/java/io/datahubproject/metadata/context/AuthorizerContext.java create mode 100644 metadata-operation-context/src/main/java/io/datahubproject/metadata/context/ContextInterface.java create mode 100644 metadata-operation-context/src/main/java/io/datahubproject/metadata/context/EntityRegistryContext.java create mode 100644 metadata-operation-context/src/main/java/io/datahubproject/metadata/context/OperationContext.java create mode 100644 metadata-operation-context/src/main/java/io/datahubproject/metadata/context/OperationContextConfig.java create mode 100644 metadata-operation-context/src/main/java/io/datahubproject/metadata/context/SearchContext.java create mode 100644 metadata-operation-context/src/main/java/io/datahubproject/test/metadata/context/TestOperationContexts.java create mode 100644 metadata-operation-context/src/test/java/io/datahubproject/metadata/context/ActorContextTest.java create mode 100644 metadata-operation-context/src/test/java/io/datahubproject/metadata/context/OperationContextTest.java create mode 100644 metadata-operation-context/src/test/java/io/datahubproject/metadata/context/SearchContextTest.java create mode 100644 metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/RestrictedServiceFactory.java create mode 100644 metadata-service/factories/src/main/java/com/linkedin/gms/factory/context/SystemOperationContextFactory.java create mode 100644 metadata-service/services/src/main/java/com/linkedin/metadata/service/RestrictedService.java create mode 100644 metadata-service/services/src/test/java/com/linkedin/metadata/service/RestrictedServiceTest.java diff --git a/build.gradle b/build.gradle index 17b21757f0380c..228a8a9f5ff0a0 100644 --- a/build.gradle +++ b/build.gradle @@ -17,7 +17,7 @@ buildscript { ext.javaClassVersion = { p -> // If Spring 6 is present, hard dependency on jdk17 - if (p.configurations.any { it.getDependencies().any{ + if (p.configurations.any { it.getDependencies().any { (it.getGroup().equals("org.springframework") && it.getVersion().startsWith("6.")) || (it.getGroup().equals("org.springframework.boot") && it.getVersion().startsWith("3.") && !it.getName().equals("spring-boot-starter-test")) }}) { @@ -43,7 +43,7 @@ buildscript { ext.elasticsearchVersion = '2.9.0' // ES 7.10, Opensearch 1.x, 2.x ext.jacksonVersion = '2.15.3' ext.jettyVersion = '11.0.19' - ext.playVersion = '2.8.18' + ext.playVersion = '2.8.21' ext.log4jVersion = '2.19.0' ext.slf4jVersion = '1.7.36' ext.logbackClassic = '1.4.14' @@ -132,7 +132,7 @@ project.ext.externalDependency = [ 'graphqlJavaScalars': 'com.graphql-java:graphql-java-extended-scalars:21.0', 'gson': 'com.google.code.gson:gson:2.8.9', 'guice': 'com.google.inject:guice:7.0.0', - 'guice4': 'com.google.inject:guice:4.2.3', // Used for frontend while still on old Play version + 'guicePlay': 'com.google.inject:guice:5.0.1', // Used for frontend while still on old Play version 'guava': 'com.google.guava:guava:32.1.2-jre', 'h2': 'com.h2database:h2:2.2.224', 'hadoopCommon':'org.apache.hadoop:hadoop-common:2.7.2', diff --git a/datahub-frontend/app/auth/AuthModule.java b/datahub-frontend/app/auth/AuthModule.java index 699dd360fa5232..c929d80328a311 100644 --- a/datahub-frontend/app/auth/AuthModule.java +++ b/datahub-frontend/app/auth/AuthModule.java @@ -8,9 +8,11 @@ import com.datahub.authentication.Actor; import com.datahub.authentication.ActorType; import com.datahub.authentication.Authentication; +import com.datahub.plugins.auth.authorization.Authorizer; import com.google.inject.AbstractModule; import com.google.inject.Provides; import com.google.inject.Singleton; +import com.google.inject.name.Named; import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.metadata.restli.DefaultRestliClientFactory; @@ -20,6 +22,13 @@ import controllers.SsoCallbackController; import java.nio.charset.StandardCharsets; import java.util.Collections; + +import io.datahubproject.metadata.context.ActorContext; +import io.datahubproject.metadata.context.AuthorizerContext; +import io.datahubproject.metadata.context.EntityRegistryContext; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.metadata.context.OperationContextConfig; +import io.datahubproject.metadata.context.SearchContext; import lombok.extern.slf4j.Slf4j; import org.apache.commons.codec.digest.DigestUtils; import org.apache.http.impl.client.CloseableHttpClient; @@ -152,6 +161,31 @@ protected Authentication provideSystemAuthentication() { Collections.emptyMap()); } + @Provides + @Singleton + @Named("systemOperationContext") + protected OperationContext provideOperationContext(final Authentication systemAuthentication, + final ConfigurationProvider configurationProvider) { + ActorContext systemActorContext = + ActorContext.builder() + .systemAuth(true) + .authentication(systemAuthentication) + .build(); + OperationContextConfig systemConfig = OperationContextConfig.builder() + .searchAuthorizationConfiguration(configurationProvider.getAuthorization().getSearch()) + .allowSystemAuthentication(true) + .build(); + + return OperationContext.builder() + .operationContextConfig(systemConfig) + .systemActorContext(systemActorContext) + .searchContext(SearchContext.EMPTY) + .entityRegistryContext(EntityRegistryContext.EMPTY) + // Authorizer.EMPTY doesn't actually apply to system auth + .authorizerContext(AuthorizerContext.builder().authorizer(Authorizer.EMPTY).build()) + .build(systemAuthentication); + } + @Provides @Singleton protected ConfigurationProvider provideConfigurationProvider() { @@ -163,13 +197,13 @@ protected ConfigurationProvider provideConfigurationProvider() { @Provides @Singleton protected SystemEntityClient provideEntityClient( - final Authentication systemAuthentication, + @Named("systemOperationContext") final OperationContext systemOperationContext, final ConfigurationProvider configurationProvider) { return new SystemRestliEntityClient( + systemOperationContext, buildRestliClient(), new ExponentialBackoff(_configs.getInt(ENTITY_CLIENT_RETRY_INTERVAL)), _configs.getInt(ENTITY_CLIENT_NUM_RETRIES), - systemAuthentication, configurationProvider.getCache().getClient().getEntityClient()); } diff --git a/datahub-frontend/app/config/ConfigurationProvider.java b/datahub-frontend/app/config/ConfigurationProvider.java index 0f2945d5d2393b..9f548b104e8fe0 100644 --- a/datahub-frontend/app/config/ConfigurationProvider.java +++ b/datahub-frontend/app/config/ConfigurationProvider.java @@ -1,5 +1,6 @@ package config; +import com.datahub.authorization.AuthorizationConfiguration; import com.linkedin.metadata.config.VisualConfiguration; import com.linkedin.metadata.config.cache.CacheConfiguration; import com.linkedin.metadata.config.kafka.KafkaConfiguration; @@ -26,4 +27,7 @@ public class ConfigurationProvider { /** Configuration for the view layer */ private VisualConfiguration visualConfig; + + /** Configuration for authorization */ + private AuthorizationConfiguration authorization; } diff --git a/datahub-frontend/play.gradle b/datahub-frontend/play.gradle index 9bd77e5279a91e..b14962e5900cd2 100644 --- a/datahub-frontend/play.gradle +++ b/datahub-frontend/play.gradle @@ -76,7 +76,7 @@ dependencies { implementation externalDependency.slf4jApi compileOnly externalDependency.lombok - runtimeOnly externalDependency.guice4 + runtimeOnly externalDependency.guicePlay runtimeOnly (externalDependency.playDocs) { exclude group: 'com.typesafe.akka', module: 'akka-http-core_2.12' } @@ -90,7 +90,7 @@ dependencies { play { platform { - playVersion = '2.8.18' + playVersion = '2.8.21' scalaVersion = '2.12' javaVersion = JavaVersion.VERSION_11 } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java index a647d0ae4e3bb8..5f555b45d3b09c 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java @@ -4,7 +4,6 @@ public class Constants { private Constants() {} - ; public static final String URN_FIELD_NAME = "urn"; public static final String URNS_FIELD_NAME = "urns"; 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 fdb91a749b2268..e9d94d313b70ec 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 @@ -321,6 +321,7 @@ import com.linkedin.datahub.graphql.types.ownership.OwnershipType; import com.linkedin.datahub.graphql.types.policy.DataHubPolicyType; import com.linkedin.datahub.graphql.types.query.QueryType; +import com.linkedin.datahub.graphql.types.restricted.RestrictedType; import com.linkedin.datahub.graphql.types.role.DataHubRoleType; import com.linkedin.datahub.graphql.types.rolemetadata.RoleType; import com.linkedin.datahub.graphql.types.schemafield.SchemaFieldType; @@ -349,6 +350,7 @@ import com.linkedin.metadata.service.LineageService; import com.linkedin.metadata.service.OwnershipTypeService; import com.linkedin.metadata.service.QueryService; +import com.linkedin.metadata.service.RestrictedService; import com.linkedin.metadata.service.SettingsService; import com.linkedin.metadata.service.ViewService; import com.linkedin.metadata.timeline.TimelineService; @@ -416,6 +418,7 @@ public class GmsGraphQLEngine { private final QueryService queryService; private final DataProductService dataProductService; private final FormService formService; + private final RestrictedService restrictedService; private final FeatureFlags featureFlags; @@ -468,6 +471,7 @@ public class GmsGraphQLEngine { private final EntityTypeType entityTypeType; private final FormType formType; private final IncidentType incidentType; + private final RestrictedType restrictedType; private final int graphQLQueryComplexityLimit; private final int graphQLQueryDepthLimit; @@ -527,6 +531,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) { this.queryService = args.queryService; this.dataProductService = args.dataProductService; this.formService = args.formService; + this.restrictedService = args.restrictedService; this.ingestionConfiguration = Objects.requireNonNull(args.ingestionConfiguration); this.authenticationConfiguration = Objects.requireNonNull(args.authenticationConfiguration); @@ -576,6 +581,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) { this.entityTypeType = new EntityTypeType(entityClient); this.formType = new FormType(entityClient); this.incidentType = new IncidentType(entityClient); + this.restrictedType = new RestrictedType(entityClient, restrictedService); this.graphQLQueryComplexityLimit = args.graphQLQueryComplexityLimit; this.graphQLQueryDepthLimit = args.graphQLQueryDepthLimit; @@ -619,7 +625,8 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) { dataTypeType, entityTypeType, formType, - incidentType); + incidentType, + restrictedType); this.loadableTypes = new ArrayList<>(entityTypes); // Extend loadable types with types from the plugins // This allows us to offer search and browse capabilities out of the box for those types @@ -709,6 +716,7 @@ public void configureRuntimeWiring(final RuntimeWiring.Builder builder) { configureStructuredPropertyResolvers(builder); configureFormResolvers(builder); configureIncidentResolvers(builder); + configureRestrictedResolvers(builder); } private void configureOrganisationRoleResolvers(RuntimeWiring.Builder builder) { @@ -1456,7 +1464,12 @@ private void configureDatasetResolvers(final RuntimeWiring.Builder builder) { .dataFetcher( "relationships", new EntityRelationshipsResultResolver(graphClient)) .dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.datasetType)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, + restrictedService, + this.authorizationConfiguration)) .dataFetcher( "platform", new LoadableTypeResolver<>( @@ -1823,7 +1836,10 @@ private void configureDashboardResolvers(final RuntimeWiring.Builder builder) { typeWiring .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient)) .dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.dashboardType)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, restrictedService, this.authorizationConfiguration)) .dataFetcher( "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)) .dataFetcher( @@ -1950,7 +1966,10 @@ private void configureChartResolvers(final RuntimeWiring.Builder builder) { .dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.chartType)) .dataFetcher( "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, restrictedService, this.authorizationConfiguration)) .dataFetcher( "platform", new LoadableTypeResolver<>( @@ -2072,7 +2091,12 @@ private void configureDataJobResolvers(final RuntimeWiring.Builder builder) { .dataFetcher( "relationships", new EntityRelationshipsResultResolver(graphClient)) .dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.dataJobType)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, + restrictedService, + this.authorizationConfiguration)) .dataFetcher( "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)) .dataFetcher( @@ -2144,7 +2168,10 @@ private void configureDataFlowResolvers(final RuntimeWiring.Builder builder) { typeWiring .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient)) .dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.dataFlowType)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, restrictedService, this.authorizationConfiguration)) .dataFetcher( "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)) .dataFetcher( @@ -2188,7 +2215,12 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde "browsePaths", new EntityBrowsePathsResolver(this.mlFeatureTableType)) .dataFetcher( "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, + restrictedService, + this.authorizationConfiguration)) .dataFetcher("exists", new EntityExistsResolver(entityService)) .dataFetcher( "platform", @@ -2271,7 +2303,12 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde .dataFetcher( "relationships", new EntityRelationshipsResultResolver(graphClient)) .dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.mlModelType)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, + restrictedService, + this.authorizationConfiguration)) .dataFetcher("exists", new EntityExistsResolver(entityService)) .dataFetcher( "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)) @@ -2316,7 +2353,12 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde "browsePaths", new EntityBrowsePathsResolver(this.mlModelGroupType)) .dataFetcher( "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, + restrictedService, + this.authorizationConfiguration)) .dataFetcher( "platform", new LoadableTypeResolver<>( @@ -2339,7 +2381,12 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde typeWiring .dataFetcher( "relationships", new EntityRelationshipsResultResolver(graphClient)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, + restrictedService, + this.authorizationConfiguration)) .dataFetcher( "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)) .dataFetcher("exists", new EntityExistsResolver(entityService)) @@ -2359,7 +2406,12 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde typeWiring .dataFetcher( "relationships", new EntityRelationshipsResultResolver(graphClient)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, + restrictedService, + this.authorizationConfiguration)) .dataFetcher( "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)) .dataFetcher("exists", new EntityExistsResolver(entityService)) @@ -2656,7 +2708,10 @@ private void configureDataProcessInstanceResolvers(final RuntimeWiring.Builder b typeWiring -> typeWiring .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient)) - .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService)) + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, restrictedService, this.authorizationConfiguration)) .dataFetcher( "state", new TimeSeriesAspectResolver( @@ -2759,4 +2814,16 @@ private void configureIncidentResolvers(final RuntimeWiring.Builder builder) { typeWiring.dataFetcher("incidents", new EntityIncidentsResolver(entityClient))); } } + + private void configureRestrictedResolvers(final RuntimeWiring.Builder builder) { + builder.type( + "Restricted", + typeWiring -> + typeWiring + .dataFetcher( + "lineage", + new EntityLineageResultResolver( + siblingGraphService, restrictedService, this.authorizationConfiguration)) + .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))); + } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java index c70e28811a1865..df32530129b040 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java @@ -29,6 +29,7 @@ import com.linkedin.metadata.service.LineageService; import com.linkedin.metadata.service.OwnershipTypeService; import com.linkedin.metadata.service.QueryService; +import com.linkedin.metadata.service.RestrictedService; import com.linkedin.metadata.service.SettingsService; import com.linkedin.metadata.service.ViewService; import com.linkedin.metadata.timeline.TimelineService; @@ -75,6 +76,7 @@ public class GmsGraphQLEngineArgs { FeatureFlags featureFlags; DataProductService dataProductService; FormService formService; + RestrictedService restrictedService; int graphQLQueryComplexityLimit; int graphQLQueryDepthLimit; diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/QueryContext.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/QueryContext.java index 9f110e713ed574..7dffd90cf2d7cc 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/QueryContext.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/QueryContext.java @@ -3,6 +3,7 @@ import com.datahub.authentication.Actor; import com.datahub.authentication.Authentication; import com.datahub.plugins.auth.authorization.Authorizer; +import io.datahubproject.metadata.context.OperationContext; /** Provided as input to GraphQL resolvers; used to carry information about GQL request context. */ public interface QueryContext { @@ -25,4 +26,9 @@ default String getActorUrn() { /** Returns the authorizer used to authorize specific actions. */ Authorizer getAuthorizer(); + + /** + * @return Returns the operational context + */ + OperationContext getOperationContext(); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/resolver/GetMetadataAnalyticsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/resolver/GetMetadataAnalyticsResolver.java index de389a358d9368..c4c353f6eb8dbd 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/resolver/GetMetadataAnalyticsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/resolver/GetMetadataAnalyticsResolver.java @@ -2,9 +2,9 @@ import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; +import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.analytics.service.AnalyticsUtil; import com.linkedin.datahub.graphql.generated.AnalyticsChart; import com.linkedin.datahub.graphql.generated.AnalyticsChartGroup; @@ -12,7 +12,6 @@ import com.linkedin.datahub.graphql.generated.BarSegment; import com.linkedin.datahub.graphql.generated.MetadataAnalyticsInput; import com.linkedin.datahub.graphql.generated.NamedBar; -import com.linkedin.datahub.graphql.resolvers.ResolverUtils; import com.linkedin.datahub.graphql.types.entitytype.EntityTypeMapper; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; @@ -22,6 +21,7 @@ import com.linkedin.metadata.search.utils.QueryUtils; import graphql.schema.DataFetcher; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -41,7 +41,7 @@ public final class GetMetadataAnalyticsResolver implements DataFetcher get(DataFetchingEnvironment environment) throws Exception { - final Authentication authentication = ResolverUtils.getAuthentication(environment); + final QueryContext context = environment.getContext(); final MetadataAnalyticsInput input = bindArgument(environment.getArgument("input"), MetadataAnalyticsInput.class); @@ -49,7 +49,7 @@ public final List get(DataFetchingEnvironment environment) final AnalyticsChartGroup group = new AnalyticsChartGroup(); group.setGroupId("FilteredMetadataAnalytics"); group.setTitle(""); - group.setCharts(getCharts(input, authentication)); + group.setCharts(getCharts(input, context.getOperationContext())); return ImmutableList.of(group); } catch (Exception e) { log.error("Failed to retrieve metadata analytics!", e); @@ -57,8 +57,8 @@ public final List get(DataFetchingEnvironment environment) } } - private List getCharts( - MetadataAnalyticsInput input, Authentication authentication) throws Exception { + private List getCharts(MetadataAnalyticsInput input, OperationContext opContext) + throws Exception { final List charts = new ArrayList<>(); List entities = Collections.emptyList(); @@ -77,8 +77,7 @@ private List getCharts( } SearchResult searchResult = - _entityClient.searchAcrossEntities( - entities, query, filter, 0, 0, null, null, authentication); + _entityClient.searchAcrossEntities(opContext, entities, query, filter, 0, 0, null, null); List aggregationMetadataList = searchResult.getMetadata().getAggregations(); @@ -96,7 +95,7 @@ private List getCharts( Constants.DOMAIN_ENTITY_NAME, ImmutableSet.of(Constants.DOMAIN_PROPERTIES_ASPECT_NAME), AnalyticsUtil::getDomainName, - authentication); + opContext.getSessionAuthentication()); charts.add(BarChart.builder().setTitle("Entities by Domain").setBars(domainChart).build()); } @@ -113,7 +112,7 @@ private List getCharts( Constants.DATA_PLATFORM_ENTITY_NAME, ImmutableSet.of(Constants.DATA_PLATFORM_INFO_ASPECT_NAME), AnalyticsUtil::getPlatformName, - authentication); + opContext.getSessionAuthentication()); charts.add( BarChart.builder().setTitle("Entities by Platform").setBars(platformChart).build()); } @@ -132,7 +131,7 @@ private List getCharts( ImmutableSet.of( Constants.GLOSSARY_TERM_KEY_ASPECT_NAME, Constants.GLOSSARY_TERM_INFO_ASPECT_NAME), AnalyticsUtil::getTermName, - authentication); + opContext.getSessionAuthentication()); charts.add(BarChart.builder().setTitle("Entities by Term").setBars(termChart).build()); } 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 4490bbcec03fc8..1a935d530505be 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 @@ -2,6 +2,7 @@ import static com.linkedin.datahub.graphql.resolvers.AuthUtils.*; import static com.linkedin.metadata.Constants.*; +import static com.linkedin.metadata.authorization.PoliciesConfig.VIEW_ENTITY_PRIVILEGES; import com.datahub.authorization.AuthUtil; import com.datahub.authorization.ConjunctivePrivilegeGroup; @@ -190,6 +191,20 @@ public static boolean canDeleteQuery( return canEditEntityQueries(subjectUrns, context); } + public static boolean canViewEntity(@Nonnull Urn entityUrn, @Nonnull QueryContext context) { + final DisjunctivePrivilegeGroup orGroup = + new DisjunctivePrivilegeGroup( + ImmutableList.of(new ConjunctivePrivilegeGroup(VIEW_ENTITY_PRIVILEGES))); + + final Authorizer authorizer = context.getAuthorizer(); + final String actor = context.getActorUrn(); + final String entityType = entityUrn.getEntityType(); + final Optional resourceSpec = + Optional.of(new EntitySpec(entityType, entityUrn.toString())); + + return AuthUtil.isAuthorized(authorizer, actor, resourceSpec, orGroup); + } + public static boolean isAuthorized( @Nonnull QueryContext context, @Nonnull Optional resourceSpec, diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolver.java index 5cfa80e394c5ff..83789ec488e64c 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolver.java @@ -13,7 +13,6 @@ import com.linkedin.datahub.graphql.generated.ListAccessTokenResult; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; import com.linkedin.metadata.search.SearchResult; @@ -65,14 +64,15 @@ public CompletableFuture get(DataFetchingEnvironment envi .setOrder(SortOrder.DESCENDING); final SearchResult searchResult = _entityClient.search( + context + .getOperationContext() + .withSearchFlags(flags -> flags.setFulltext(true)), Constants.ACCESS_TOKEN_ENTITY_NAME, "", buildFilter(filters, Collections.emptyList()), sortCriterion, start, - count, - getAuthentication(environment), - new SearchFlags().setFulltext(true)); + count); final List tokens = searchResult.getEntities().stream() diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/chart/BrowseV2Resolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/chart/BrowseV2Resolver.java index 7bcde0ea9bdc12..b40c6a3fd0f788 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/chart/BrowseV2Resolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/chart/BrowseV2Resolver.java @@ -76,6 +76,7 @@ public CompletableFuture get(DataFetchingEnvironment environmen BrowseResultV2 browseResults = _entityClient.browseV2( + context.getOperationContext().withSearchFlags(flags -> searchFlags), entityNames, pathStr, maybeResolvedView != null @@ -84,9 +85,7 @@ public CompletableFuture get(DataFetchingEnvironment environmen : inputFilter, sanitizedQuery, start, - count, - context.getAuthentication(), - searchFlags); + count); return mapBrowseResults(browseResults); } catch (Exception e) { throw new RuntimeException("Failed to execute browse V2", e); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/container/ContainerEntitiesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/container/ContainerEntitiesResolver.java index 58f7715c3e627e..b757a8aa2aab9e 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/container/ContainerEntitiesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/container/ContainerEntitiesResolver.java @@ -79,6 +79,7 @@ public CompletableFuture get(final DataFetchingEnvironment enviro return UrnSearchResultsMapper.map( _entityClient.searchAcrossEntities( + context.getOperationContext(), CONTAINABLE_ENTITY_NAMES, query, new Filter() @@ -90,8 +91,7 @@ public CompletableFuture get(final DataFetchingEnvironment enviro start, count, null, - null, - context.getAuthentication())); + null)); } catch (Exception e) { throw new RuntimeException( diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataproduct/ListDataProductAssetsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataproduct/ListDataProductAssetsResolver.java index 72912087190c05..29729a8799e790 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataproduct/ListDataProductAssetsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataproduct/ListDataProductAssetsResolver.java @@ -134,10 +134,12 @@ public CompletableFuture get(DataFetchingEnvironment environment) ResolverUtils.buildFilter(input.getFilters(), input.getOrFilters()); final Filter finalFilter = buildFilterWithUrns(new HashSet<>(assetUrns), baseFilter); - SearchFlags searchFlags = null; + final SearchFlags searchFlags; com.linkedin.datahub.graphql.generated.SearchFlags inputFlags = input.getSearchFlags(); if (inputFlags != null) { searchFlags = SearchFlagsInputMapper.INSTANCE.apply(inputFlags); + } else { + searchFlags = null; } try { @@ -151,14 +153,15 @@ public CompletableFuture get(DataFetchingEnvironment environment) return UrnSearchResultsMapper.map( _entityClient.searchAcrossEntities( + context + .getOperationContext() + .withSearchFlags(flags -> searchFlags != null ? searchFlags : flags), finalEntityNames, sanitizedQuery, finalFilter, start, count, - searchFlags, - null, - ResolverUtils.getAuthentication(environment))); + null)); } catch (Exception e) { log.error( "Failed to execute search for data product assets: entity types {}, query {}, filters: {}, start: {}, count: {}", diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/deprecation/UpdateDeprecationResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/deprecation/UpdateDeprecationResolver.java index be887d845f3853..e251e36a3e15f2 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/deprecation/UpdateDeprecationResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/deprecation/UpdateDeprecationResolver.java @@ -50,7 +50,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) throw return CompletableFuture.supplyAsync( () -> { - if (!isAuthorizedToUpdateDeprecationForEntity(environment.getContext(), entityUrn)) { + if (!isAuthorizedToUpdateDeprecationForEntity(context, entityUrn)) { throw new AuthorizationException( "Unauthorized to perform this action. Please contact your DataHub administrator."); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolver.java index 6229e38954163d..829b2f903833d4 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolver.java @@ -86,6 +86,7 @@ public CompletableFuture get(final DataFetchingEnvironment enviro return UrnSearchResultsMapper.map( _entityClient.searchAcrossEntities( + context.getOperationContext(), SEARCHABLE_ENTITY_TYPES.stream() .map(EntityTypeMapper::getName) .collect(Collectors.toList()), @@ -97,8 +98,7 @@ public CompletableFuture get(final DataFetchingEnvironment enviro start, count, null, - null, - context.getAuthentication())); + null)); } catch (Exception e) { throw new RuntimeException( diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/ListDomainsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/ListDomainsResolver.java index 5453603f4cc9f9..fe4a7f23cfaab6 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/ListDomainsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/ListDomainsResolver.java @@ -13,7 +13,6 @@ import com.linkedin.datahub.graphql.resolvers.mutate.util.DomainUtils; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; @@ -62,6 +61,7 @@ public CompletableFuture get(final DataFetchingEnvironment en // First, get all domain Urns. final SearchResult gmsResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), Constants.DOMAIN_ENTITY_NAME, query, filter, @@ -69,9 +69,7 @@ public CompletableFuture get(final DataFetchingEnvironment en .setField(DOMAIN_CREATED_TIME_INDEX_FIELD_NAME) .setOrder(SortOrder.DESCENDING), start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); // Now that we have entities we can bind this to a result. final ListDomainsResult result = new ListDomainsResult(); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolver.java index 147663059aa82c..54fd7ef1fe04d0 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolver.java @@ -140,7 +140,7 @@ private Map getTermsWithSameParent(Urn parentNode, QueryCon final Filter filter = buildParentNodeFilter(parentNode); final SearchResult searchResult = _entityClient.filter( - GLOSSARY_TERM_ENTITY_NAME, filter, null, 0, 1000, context.getAuthentication()); + context.getOperationContext(), GLOSSARY_TERM_ENTITY_NAME, filter, null, 0, 1000); final List termUrns = searchResult.getEntities().stream() diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryNodesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryNodesResolver.java index e7990b1a343d83..700a38d50b317d 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryNodesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryNodesResolver.java @@ -53,12 +53,12 @@ public CompletableFuture get( final Filter filter = buildGlossaryEntitiesFilter(); final SearchResult gmsNodesResult = _entityClient.filter( + context.getOperationContext(), Constants.GLOSSARY_NODE_ENTITY_NAME, filter, null, start, - count, - context.getAuthentication()); + count); final List glossaryNodeUrns = gmsNodesResult.getEntities().stream() diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryTermsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryTermsResolver.java index 40e4363dcff938..9669d406344e71 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryTermsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryTermsResolver.java @@ -53,12 +53,12 @@ public CompletableFuture get( final Filter filter = buildGlossaryEntitiesFilter(); final SearchResult gmsTermsResult = _entityClient.filter( + context.getOperationContext(), Constants.GLOSSARY_TERM_ENTITY_NAME, filter, null, start, - count, - context.getAuthentication()); + count); final List glossaryTermUrns = gmsTermsResult.getEntities().stream() diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/EntityCountsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/EntityCountsResolver.java index 1f8c17ee728847..8abe2378982930 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/EntityCountsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/EntityCountsResolver.java @@ -41,10 +41,10 @@ public CompletableFuture get(final DataFetchingEnvironment e // First, get all counts Map gmsResult = _entityClient.batchGetTotalEntityCount( + context.getOperationContext(), input.getTypes().stream() .map(EntityTypeMapper::getName) - .collect(Collectors.toList()), - context.getAuthentication()); + .collect(Collectors.toList())); // bind to a result. List resultList = diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/ListGroupsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/ListGroupsResolver.java index a6ad8698679f02..70be478d65c5c7 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/ListGroupsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/ListGroupsResolver.java @@ -13,7 +13,6 @@ import com.linkedin.datahub.graphql.generated.ListGroupsResult; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; import com.linkedin.metadata.search.SearchEntity; @@ -58,6 +57,9 @@ public CompletableFuture get(final DataFetchingEnvironment env // First, get all group Urns. final SearchResult gmsResult = _entityClient.search( + context + .getOperationContext() + .withSearchFlags(flags -> flags.setFulltext(true)), CORP_GROUP_ENTITY_NAME, query, null, @@ -65,9 +67,7 @@ public CompletableFuture get(final DataFetchingEnvironment env .setField(CORP_GROUP_CREATED_TIME_INDEX_FIELD_NAME) .setOrder(SortOrder.DESCENDING), start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); // Then, get hydrate all groups. final Map entities = diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/health/EntityHealthResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/health/EntityHealthResolver.java index 79585035562742..ad36621d20c668 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/health/EntityHealthResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/health/EntityHealthResolver.java @@ -138,7 +138,7 @@ private Health computeIncidentsHealthForAsset( final Filter filter = buildIncidentsEntityFilter(entityUrn, IncidentState.ACTIVE.toString()); final SearchResult searchResult = _entityClient.filter( - Constants.INCIDENT_ENTITY_NAME, filter, null, 0, 1, context.getAuthentication()); + context.getOperationContext(), Constants.INCIDENT_ENTITY_NAME, filter, null, 0, 1); final Integer activeIncidentCount = searchResult.getNumEntities(); if (activeIncidentCount > 0) { // There are active incidents. diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/incident/EntityIncidentsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/incident/EntityIncidentsResolver.java index c797044d1b224f..c0c3217fd056db 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/incident/EntityIncidentsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/incident/EntityIncidentsResolver.java @@ -62,12 +62,12 @@ public CompletableFuture get(DataFetchingEnvironment envi final SortCriterion sortCriterion = buildIncidentsSortCriterion(); final SearchResult searchResult = _entityClient.filter( + context.getOperationContext(), Constants.INCIDENT_ENTITY_NAME, filter, sortCriterion, start, - count, - context.getAuthentication()); + count); final List incidentUrns = searchResult.getEntities().stream() diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/execution/IngestionSourceExecutionRequestsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/execution/IngestionSourceExecutionRequestsResolver.java index 01100a24d6b15c..4c8a06e2d585a1 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/execution/IngestionSourceExecutionRequestsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/execution/IngestionSourceExecutionRequestsResolver.java @@ -67,6 +67,7 @@ public CompletableFuture get( final SearchResult executionsSearchResult = _entityClient.filter( + context.getOperationContext(), Constants.EXECUTION_REQUEST_ENTITY_NAME, new Filter() .setOr( @@ -78,8 +79,7 @@ public CompletableFuture get( .setField(REQUEST_TIME_MS_FIELD_NAME) .setOrder(SortOrder.DESCENDING), start, - count, - context.getAuthentication()); + count); // 2. Batch fetch the related ExecutionRequests final Set relatedExecRequests = diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolver.java index eb054295af09b3..b3c7db20f45374 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolver.java @@ -17,7 +17,6 @@ import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; import com.linkedin.metadata.search.SearchEntity; @@ -67,6 +66,9 @@ public CompletableFuture get(final DataFetchingEnvironment en // First, get all secrets final SearchResult gmsResult = _entityClient.search( + context + .getOperationContext() + .withSearchFlags(flags -> flags.setFulltext(true)), Constants.SECRETS_ENTITY_NAME, query, null, @@ -74,9 +76,7 @@ public CompletableFuture get(final DataFetchingEnvironment en .setField(DOMAIN_CREATED_TIME_INDEX_FIELD_NAME) .setOrder(SortOrder.DESCENDING), start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); // Then, resolve all secrets final Map entities = diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/source/ListIngestionSourcesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/source/ListIngestionSourcesResolver.java index 51c9e30aadcce1..d2387820ca7ab7 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/source/ListIngestionSourcesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/source/ListIngestionSourcesResolver.java @@ -14,7 +14,6 @@ import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchResult; import graphql.schema.DataFetcher; @@ -63,14 +62,15 @@ public CompletableFuture get( // First, get all ingestion sources Urns. final SearchResult gmsResult = _entityClient.search( + context + .getOperationContext() + .withSearchFlags(flags -> flags.setFulltext(true)), Constants.INGESTION_SOURCE_ENTITY_NAME, query, buildFilter(filters, Collections.emptyList()), null, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); // Then, resolve all ingestion sources final Map entities = diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/jobs/DataJobRunsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/jobs/DataJobRunsResolver.java index 06bad27e270620..325e804327b72f 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/jobs/DataJobRunsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/jobs/DataJobRunsResolver.java @@ -63,12 +63,12 @@ public CompletableFuture get(DataFetchingEnvironment final SortCriterion sortCriterion = buildTaskRunsSortCriterion(); final SearchResult gmsResult = _entityClient.filter( + context.getOperationContext(), Constants.DATA_PROCESS_INSTANCE_ENTITY_NAME, filter, sortCriterion, start, - count, - context.getAuthentication()); + count); final List dataProcessInstanceUrns = gmsResult.getEntities().stream() .map(SearchEntity::getEntity) diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/jobs/EntityRunsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/jobs/EntityRunsResolver.java index d595b1e513d75e..1c6fa352fecb66 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/jobs/EntityRunsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/jobs/EntityRunsResolver.java @@ -69,12 +69,12 @@ public CompletableFuture get(DataFetchingEnvironment final SortCriterion sortCriterion = buildTaskRunsSortCriterion(); final SearchResult gmsResult = _entityClient.filter( + context.getOperationContext(), Constants.DATA_PROCESS_INSTANCE_ENTITY_NAME, filter, sortCriterion, start, - count, - context.getAuthentication()); + count); final List dataProcessInstanceUrns = gmsResult.getEntities().stream() .map(SearchEntity::getEntity) diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityLineageResultResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityLineageResultResolver.java index c63ec819e8f6a1..6f56bfed942403 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityLineageResultResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityLineageResultResolver.java @@ -2,18 +2,26 @@ import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*; +import com.datahub.authorization.AuthorizationConfiguration; import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.authorization.AuthorizationUtils; import com.linkedin.datahub.graphql.generated.Entity; import com.linkedin.datahub.graphql.generated.EntityLineageResult; +import com.linkedin.datahub.graphql.generated.EntityType; import com.linkedin.datahub.graphql.generated.LineageDirection; import com.linkedin.datahub.graphql.generated.LineageInput; import com.linkedin.datahub.graphql.generated.LineageRelationship; +import com.linkedin.datahub.graphql.generated.Restricted; import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper; +import com.linkedin.metadata.Constants; import com.linkedin.metadata.graph.SiblingGraphService; +import com.linkedin.metadata.service.RestrictedService; import graphql.schema.DataFetcher; import graphql.schema.DataFetchingEnvironment; -import java.net.URISyntaxException; import java.util.HashSet; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.stream.Collectors; import javax.annotation.Nullable; @@ -29,16 +37,28 @@ public class EntityLineageResultResolver implements DataFetcher> { private final SiblingGraphService _siblingGraphService; + private final RestrictedService _restrictedService; + private final AuthorizationConfiguration _authorizationConfiguration; - public EntityLineageResultResolver(final SiblingGraphService siblingGraphService) { + public EntityLineageResultResolver( + final SiblingGraphService siblingGraphService, + final RestrictedService restrictedService, + final AuthorizationConfiguration authorizationConfiguration) { _siblingGraphService = siblingGraphService; + _restrictedService = restrictedService; + _authorizationConfiguration = authorizationConfiguration; } @Override public CompletableFuture get(DataFetchingEnvironment environment) { - final String urn = ((Entity) environment.getSource()).getUrn(); + final QueryContext context = environment.getContext(); + Urn urn = UrnUtils.getUrn(((Entity) environment.getSource()).getUrn()); final LineageInput input = bindArgument(environment.getArgument("input"), LineageInput.class); + if (urn.getEntityType().equals(Constants.RESTRICTED_ENTITY_NAME)) { + urn = _restrictedService.decryptRestrictedUrn(urn); + } + final LineageDirection lineageDirection = input.getDirection(); @Nullable final Integer start = input.getStart(); // Optional! @Nullable final Integer count = input.getCount(); // Optional! @@ -49,12 +69,13 @@ public CompletableFuture get(DataFetchingEnvironment enviro com.linkedin.metadata.graph.LineageDirection resolvedDirection = com.linkedin.metadata.graph.LineageDirection.valueOf(lineageDirection.toString()); + final Urn finalUrn = urn; return CompletableFuture.supplyAsync( () -> { try { - return mapEntityRelationships( + com.linkedin.metadata.graph.EntityLineageResult entityLineageResult = _siblingGraphService.getLineage( - Urn.createFromString(urn), + finalUrn, resolvedDirection, start != null ? start : 0, count != null ? count : 100, @@ -62,16 +83,31 @@ public CompletableFuture get(DataFetchingEnvironment enviro separateSiblings != null ? input.getSeparateSiblings() : false, new HashSet<>(), startTimeMillis, - endTimeMillis)); - } catch (URISyntaxException e) { - log.error("Failed to fetch lineage for {}", urn); - throw new RuntimeException(String.format("Failed to fetch lineage for {}", urn), e); + endTimeMillis); + + Set restrictedUrns = new HashSet<>(); + entityLineageResult + .getRelationships() + .forEach( + rel -> { + if (_authorizationConfiguration.getSearch().isEnabled() + && !AuthorizationUtils.canViewEntity(rel.getEntity(), context)) { + restrictedUrns.add(rel.getEntity()); + } + }); + + return mapEntityRelationships(entityLineageResult, restrictedUrns); + } catch (Exception e) { + log.error("Failed to fetch lineage for {}", finalUrn); + throw new RuntimeException( + String.format("Failed to fetch lineage for {}", finalUrn), e); } }); } private EntityLineageResult mapEntityRelationships( - final com.linkedin.metadata.graph.EntityLineageResult entityLineageResult) { + final com.linkedin.metadata.graph.EntityLineageResult entityLineageResult, + final Set restrictedUrns) { final EntityLineageResult result = new EntityLineageResult(); result.setStart(entityLineageResult.getStart()); result.setCount(entityLineageResult.getCount()); @@ -79,17 +115,28 @@ private EntityLineageResult mapEntityRelationships( result.setFiltered(entityLineageResult.getFiltered()); result.setRelationships( entityLineageResult.getRelationships().stream() - .map(this::mapEntityRelationship) + .map(r -> mapEntityRelationship(r, restrictedUrns)) .collect(Collectors.toList())); return result; } private LineageRelationship mapEntityRelationship( - final com.linkedin.metadata.graph.LineageRelationship lineageRelationship) { + final com.linkedin.metadata.graph.LineageRelationship lineageRelationship, + final Set restrictedUrns) { final LineageRelationship result = new LineageRelationship(); - final Entity partialEntity = UrnToEntityMapper.map(lineageRelationship.getEntity()); - if (partialEntity != null) { - result.setEntity(partialEntity); + if (restrictedUrns.contains(lineageRelationship.getEntity())) { + final Restricted restrictedEntity = new Restricted(); + restrictedEntity.setType(EntityType.RESTRICTED); + String restrictedUrnString = + _restrictedService.encryptRestrictedUrn(lineageRelationship.getEntity()).toString(); + + restrictedEntity.setUrn(restrictedUrnString); + result.setEntity(restrictedEntity); + } else { + final Entity partialEntity = UrnToEntityMapper.map(lineageRelationship.getEntity()); + if (partialEntity != null) { + result.setEntity(partialEntity); + } } result.setType(lineageRelationship.getType()); result.setDegree(lineageRelationship.getDegree()); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java index 5dbd282580c870..29447e6e7ef225 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java @@ -212,7 +212,7 @@ public static boolean hasChildDomains( // Limit count to 1 for existence check final SearchResult searchResult = entityClient.filter( - DOMAIN_ENTITY_NAME, parentDomainFilter, null, 0, 1, context.getAuthentication()); + context.getOperationContext(), DOMAIN_ENTITY_NAME, parentDomainFilter, null, 0, 1); return (searchResult.getNumEntities() > 0); } @@ -226,7 +226,7 @@ private static Map getDomainsByNameAndParent( final SearchResult searchResult = entityClient.filter( - DOMAIN_ENTITY_NAME, filter, null, 0, 1000, context.getAuthentication()); + context.getOperationContext(), DOMAIN_ENTITY_NAME, filter, null, 0, 1000); final Set domainUrns = searchResult.getEntities().stream() diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ownership/ListOwnershipTypesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ownership/ListOwnershipTypesResolver.java index 1c8f43a4901737..aec3848f966406 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ownership/ListOwnershipTypesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ownership/ListOwnershipTypesResolver.java @@ -11,7 +11,6 @@ import com.linkedin.datahub.graphql.generated.OwnershipTypeEntity; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; import com.linkedin.metadata.search.SearchEntity; @@ -60,14 +59,13 @@ public CompletableFuture get(DataFetchingEnvironment e final SearchResult gmsResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), Constants.OWNERSHIP_TYPE_ENTITY_NAME, query, buildFilter(filters, Collections.emptyList()), DEFAULT_SORT_CRITERION, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); final ListOwnershipTypesResult result = new ListOwnershipTypesResult(); result.setStart(gmsResult.getFrom()); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/ListPoliciesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/ListPoliciesResolver.java index c1e4bb7f833167..1d26aab85f4632 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/ListPoliciesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/ListPoliciesResolver.java @@ -61,7 +61,7 @@ public CompletableFuture get(final DataFetchingEnvironment e final Filter filter = ResolverUtils.buildFilter(facetFilters, Collections.emptyList()); return _policyFetcher - .fetchPolicies(start, query, count, filter, context.getAuthentication()) + .fetchPolicies(context.getOperationContext(), start, query, count, filter) .thenApply( policyFetchResult -> { final ListPoliciesResult result = new ListPoliciesResult(); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/post/ListPostsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/post/ListPostsResolver.java index 5292adbe3aac39..34ef616c61e411 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/post/ListPostsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/post/ListPostsResolver.java @@ -11,7 +11,6 @@ import com.linkedin.datahub.graphql.types.post.PostMapper; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; import com.linkedin.metadata.search.SearchEntity; @@ -57,14 +56,13 @@ public CompletableFuture get(final DataFetchingEnvironment envi // First, get all Post Urns. final SearchResult gmsResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), POST_ENTITY_NAME, query, null, sortCriterion, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); // Then, get and hydrate all Posts. final Map entities = diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/query/ListQueriesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/query/ListQueriesResolver.java index fec5bb120eebae..6fcc0fee763038 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/query/ListQueriesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/query/ListQueriesResolver.java @@ -14,7 +14,6 @@ import com.linkedin.datahub.graphql.generated.ListQueriesResult; import com.linkedin.datahub.graphql.generated.QueryEntity; import com.linkedin.entity.client.EntityClient; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; @@ -66,14 +65,16 @@ public CompletableFuture get(final DataFetchingEnvironment en // First, get all Query Urns. final SearchResult gmsResult = _entityClient.search( + context + .getOperationContext() + .withSearchFlags( + flags -> flags.setFulltext(true).setSkipHighlighting(true)), QUERY_ENTITY_NAME, query, buildFilters(input), sortCriterion, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true).setSkipHighlighting(true)); + count); final ListQueriesResult result = new ListQueriesResult(); result.setStart(gmsResult.getFrom()); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/recommendation/ListRecommendationsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/recommendation/ListRecommendationsResolver.java index e65666117b4fac..86d1b8bab669c2 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/recommendation/ListRecommendationsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/recommendation/ListRecommendationsResolver.java @@ -4,6 +4,7 @@ import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; +import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.ContentParams; import com.linkedin.datahub.graphql.generated.EntityProfileParams; import com.linkedin.datahub.graphql.generated.FacetFilter; @@ -46,6 +47,7 @@ public class ListRecommendationsResolver @WithSpan @Override public CompletableFuture get(DataFetchingEnvironment environment) { + final QueryContext context = environment.getContext(); final ListRecommendationsInput input = bindArgument(environment.getArgument("input"), ListRecommendationsInput.class); @@ -55,7 +57,7 @@ public CompletableFuture get(DataFetchingEnvironment log.debug("Listing recommendations for input {}", input); List modules = _recommendationsService.listRecommendations( - Urn.createFromString(input.getUserUrn()), + context.getOperationContext(), mapRequestContext(input.getRequestContext()), input.getLimit()); return ListRecommendationsResult.builder() diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/CreateInviteTokenResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/CreateInviteTokenResolver.java index 61ecf09fc91a51..3bf11b9febc638 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/CreateInviteTokenResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/CreateInviteTokenResolver.java @@ -38,7 +38,8 @@ public CompletableFuture get(final DataFetchingEnvironment environm () -> { try { return new InviteToken( - _inviteTokenService.getInviteToken(roleUrnStr, true, authentication)); + _inviteTokenService.getInviteToken( + context.getOperationContext(), roleUrnStr, true)); } catch (Exception e) { throw new RuntimeException( String.format("Failed to create invite token for role %s", roleUrnStr), e); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/GetInviteTokenResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/GetInviteTokenResolver.java index 066753c4f7559f..039a1730e7e67a 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/GetInviteTokenResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/GetInviteTokenResolver.java @@ -38,7 +38,8 @@ public CompletableFuture get(final DataFetchingEnvironment environm () -> { try { return new InviteToken( - _inviteTokenService.getInviteToken(roleUrnStr, false, authentication)); + _inviteTokenService.getInviteToken( + context.getOperationContext(), roleUrnStr, false)); } catch (Exception e) { throw new RuntimeException( String.format("Failed to get invite token for role %s", roleUrnStr), e); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolver.java index a1dd9219f6549c..5c0ea6651f67ed 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolver.java @@ -11,7 +11,6 @@ import com.linkedin.datahub.graphql.types.role.mappers.DataHubRoleMapper; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchResult; import graphql.schema.DataFetcher; @@ -53,13 +52,12 @@ public CompletableFuture get(final DataFetchingEnvironment envi // First, get all role Urns. final SearchResult gmsResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), DATAHUB_ROLE_ENTITY_NAME, query, Collections.emptyMap(), start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); // Then, get and hydrate all users. final Map entities = diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AggregateAcrossEntitiesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AggregateAcrossEntitiesResolver.java index b54987dc0e9b01..44b998219f01ea 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AggregateAcrossEntitiesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AggregateAcrossEntitiesResolver.java @@ -70,6 +70,7 @@ public CompletableFuture get(DataFetchingEnvironment environme try { return mapAggregateResults( _entityClient.searchAcrossEntities( + context.getOperationContext().withSearchFlags(flags -> searchFlags), maybeResolvedView != null ? SearchUtils.intersectEntityTypes( entityNames, maybeResolvedView.getDefinition().getEntityTypes()) @@ -81,9 +82,7 @@ public CompletableFuture get(DataFetchingEnvironment environme : inputFilter, 0, 0, // 0 entity count because we don't want resolved entities - searchFlags, null, - ResolverUtils.getAuthentication(environment), facets)); } catch (Exception e) { log.error( diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/GetQuickFiltersResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/GetQuickFiltersResolver.java index 1a380781385c34..21007bf228a707 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/GetQuickFiltersResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/GetQuickFiltersResolver.java @@ -5,13 +5,12 @@ import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.SEARCHABLE_ENTITY_TYPES; import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.resolveView; -import com.datahub.authentication.Authentication; import com.linkedin.common.urn.UrnUtils; +import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.Entity; import com.linkedin.datahub.graphql.generated.GetQuickFiltersInput; import com.linkedin.datahub.graphql.generated.GetQuickFiltersResult; import com.linkedin.datahub.graphql.generated.QuickFilter; -import com.linkedin.datahub.graphql.resolvers.ResolverUtils; import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper; import com.linkedin.datahub.graphql.types.entitytype.EntityTypeMapper; import com.linkedin.entity.client.EntityClient; @@ -23,6 +22,7 @@ import com.linkedin.view.DataHubViewInfo; import graphql.schema.DataFetcher; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import java.util.ArrayList; import java.util.Comparator; import java.util.List; @@ -48,6 +48,7 @@ public class GetQuickFiltersResolver public CompletableFuture get(final DataFetchingEnvironment environment) throws Exception { + final QueryContext context = environment.getContext(); final GetQuickFiltersInput input = bindArgument(environment.getArgument("input"), GetQuickFiltersInput.class); @@ -58,7 +59,7 @@ public CompletableFuture get(final DataFetchingEnvironmen try { final SearchResult searchResult = - getSearchResults(ResolverUtils.getAuthentication(environment), input); + getSearchResults(context.getOperationContext(), input); final AggregationMetadataArray aggregations = searchResult.getMetadata().getAggregations(); @@ -74,13 +75,17 @@ public CompletableFuture get(final DataFetchingEnvironmen }); } - /** Do a star search with view filter applied to get info about all data in this instance. */ + /** + * Do a star search with view filter applied to get info about all data in this instance. Include + * aggregations. + */ private SearchResult getSearchResults( - @Nonnull final Authentication authentication, @Nonnull final GetQuickFiltersInput input) + @Nonnull final OperationContext opContext, @Nonnull final GetQuickFiltersInput input) throws Exception { final DataHubViewInfo maybeResolvedView = (input.getViewUrn() != null) - ? resolveView(_viewService, UrnUtils.getUrn(input.getViewUrn()), authentication) + ? resolveView( + _viewService, UrnUtils.getUrn(input.getViewUrn()), opContext.getAuthentication()) : null; final List entityNames = SEARCHABLE_ENTITY_TYPES.stream() @@ -88,6 +93,7 @@ private SearchResult getSearchResults( .collect(Collectors.toList()); return _entityClient.searchAcrossEntities( + opContext.withSearchFlags(flags -> flags.setSkipAggregates(false)), maybeResolvedView != null ? SearchUtils.intersectEntityTypes( entityNames, maybeResolvedView.getDefinition().getEntityTypes()) @@ -99,8 +105,7 @@ private SearchResult getSearchResults( 0, 0, null, - null, - authentication); + null); } /** diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossEntitiesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossEntitiesResolver.java index 658138ae6e3dc0..79101c9b6a48f5 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossEntitiesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossEntitiesResolver.java @@ -72,10 +72,12 @@ public CompletableFuture get(DataFetchingEnvironment environment) : null; final Filter baseFilter = ResolverUtils.buildFilter(null, input.getOrFilters()); - SearchFlags searchFlags = null; + final SearchFlags searchFlags; com.linkedin.datahub.graphql.generated.SearchFlags inputFlags = input.getSearchFlags(); if (inputFlags != null) { searchFlags = SearchFlagsInputMapper.INSTANCE.apply(inputFlags); + } else { + searchFlags = null; } try { @@ -90,6 +92,9 @@ public CompletableFuture get(DataFetchingEnvironment environment) return UrnScrollResultsMapper.map( _entityClient.scrollAcrossEntities( + context + .getOperationContext() + .withSearchFlags(flags -> searchFlags != null ? searchFlags : flags), maybeResolvedView != null ? SearchUtils.intersectEntityTypes( entityNames, maybeResolvedView.getDefinition().getEntityTypes()) @@ -101,9 +106,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) : baseFilter, scrollId, keepAlive, - count, - searchFlags, - ResolverUtils.getAuthentication(environment))); + count)); } catch (Exception e) { log.error( "Failed to execute search for multiple entities: entity types {}, query {}, filters: {}, searchAfter: {}, count: {}", diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossLineageResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossLineageResolver.java index 0af0a3827b1bb4..fa82cad6fffdef 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossLineageResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossLineageResolver.java @@ -93,7 +93,7 @@ public CompletableFuture get(DataFetchingEnvironment scrollId, count); - SearchFlags searchFlags = null; + final SearchFlags searchFlags; final com.linkedin.datahub.graphql.generated.SearchFlags inputFlags = input.getSearchFlags(); if (inputFlags != null) { @@ -102,9 +102,14 @@ public CompletableFuture get(DataFetchingEnvironment .setSkipCache(inputFlags.getSkipCache()) .setFulltext(inputFlags.getFulltext()) .setMaxAggValues(inputFlags.getMaxAggValues()); + } else { + searchFlags = null; } return UrnScrollAcrossLineageResultsMapper.map( _entityClient.scrollAcrossLineage( + context + .getOperationContext() + .withSearchFlags(flags -> searchFlags != null ? searchFlags : flags), urn, resolvedDirection, entityNames, @@ -116,9 +121,7 @@ public CompletableFuture get(DataFetchingEnvironment keepAlive, count, startTimeMillis, - endTimeMillis, - searchFlags, - ResolverUtils.getAuthentication(environment))); + endTimeMillis)); } catch (RemoteInvocationException e) { log.error( "Failed to execute scroll across relationships: source urn {}, direction {}, entity types {}, query {}, filters: {}, start: {}, count: {}", diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossEntitiesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossEntitiesResolver.java index f8178e3b396cb5..6f4bcf937f2fe7 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossEntitiesResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossEntitiesResolver.java @@ -77,6 +77,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) return UrnSearchResultsMapper.map( _entityClient.searchAcrossEntities( + context.getOperationContext().withSearchFlags(flags -> searchFlags), maybeResolvedView != null ? SearchUtils.intersectEntityTypes( entityNames, maybeResolvedView.getDefinition().getEntityTypes()) @@ -88,9 +89,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) : baseFilter, start, count, - searchFlags, - sortCriterion, - ResolverUtils.getAuthentication(environment))); + sortCriterion)); } catch (Exception e) { log.error( "Failed to execute search for multiple entities: entity types {}, query {}, filters: {}, start: {}, count: {}", diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolver.java index 4a0eacaf09671a..9937dac2447f71 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolver.java @@ -6,6 +6,7 @@ import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; +import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.AndFilterInput; import com.linkedin.datahub.graphql.generated.EntityType; import com.linkedin.datahub.graphql.generated.FacetFilterInput; @@ -78,6 +79,8 @@ private List getEntityNamesFromInput(List inputTypes) { public CompletableFuture get(DataFetchingEnvironment environment) throws URISyntaxException { log.debug("Entering search across lineage graphql resolver"); + final QueryContext context = environment.getContext(); + final SearchAcrossLineageInput input = bindArgument(environment.getArgument("input"), SearchAcrossLineageInput.class); @@ -125,7 +128,7 @@ public CompletableFuture get(DataFetchingEnvironment final Filter filter = ResolverUtils.buildFilter(input.getFilters(), input.getOrFilters()); - SearchFlags searchFlags = null; + final SearchFlags searchFlags; com.linkedin.datahub.graphql.generated.SearchFlags inputFlags = input.getSearchFlags(); if (inputFlags != null) { searchFlags = SearchFlagsInputMapper.INSTANCE.apply(inputFlags); @@ -137,6 +140,7 @@ public CompletableFuture get(DataFetchingEnvironment } LineageSearchResult salResults = _entityClient.searchAcrossLineage( + context.getOperationContext().withSearchFlags(flags -> searchFlags), urn, resolvedDirection, entityNames, @@ -147,9 +151,7 @@ public CompletableFuture get(DataFetchingEnvironment start, count, startTimeMillis, - endTimeMillis, - searchFlags, - getAuthentication(environment)); + endTimeMillis); return UrnSearchAcrossLineageResultsMapper.map(salResults); } catch (RemoteInvocationException e) { diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolver.java index 7428207034f5dd..7d1b7d6ef28387 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolver.java @@ -4,6 +4,7 @@ import static com.linkedin.metadata.Constants.*; import static com.linkedin.metadata.search.utils.SearchUtils.applyDefaultSearchFlags; +import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.SearchInput; import com.linkedin.datahub.graphql.generated.SearchResults; import com.linkedin.datahub.graphql.resolvers.ResolverUtils; @@ -48,6 +49,7 @@ public class SearchResolver implements DataFetcher get(DataFetchingEnvironment environment) { + final QueryContext context = environment.getContext(); final SearchInput input = bindArgument(environment.getArgument("input"), SearchInput.class); final String entityName = EntityTypeMapper.getName(input.getType()); // escape forward slash since it is a reserved character in Elasticsearch @@ -78,14 +80,13 @@ public CompletableFuture get(DataFetchingEnvironment environment) return UrnSearchResultsMapper.map( _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> searchFlags), entityName, sanitizedQuery, ResolverUtils.buildFilter(input.getFilters(), input.getOrFilters()), null, start, - count, - ResolverUtils.getAuthentication(environment), - searchFlags)); + count)); } catch (Exception e) { log.error( "Failed to execute search: entity type {}, query {}, filters: {}, orFilters: {}, start: {}, count: {}, searchFlags: {}", diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/test/ListTestsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/test/ListTestsResolver.java index f345d9ceb21e52..3f4a0367af05ad 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/test/ListTestsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/test/ListTestsResolver.java @@ -12,7 +12,6 @@ import com.linkedin.datahub.graphql.generated.Test; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; @@ -57,13 +56,14 @@ public CompletableFuture get(final DataFetchingEnvironment envi // First, get all group Urns. final SearchResult gmsResult = _entityClient.search( + context + .getOperationContext() + .withSearchFlags(flags -> flags.setFulltext(true)), Constants.TEST_ENTITY_NAME, query, Collections.emptyMap(), start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); // Now that we have entities we can bind this to a result. final ListTestsResult result = new ListTestsResult(); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/user/ListUsersResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/user/ListUsersResolver.java index 215d53299c8ac1..ef03a67a55c82a 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/user/ListUsersResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/user/ListUsersResolver.java @@ -13,7 +13,6 @@ import com.linkedin.datahub.graphql.types.corpuser.mappers.CorpUserMapper; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchResult; import graphql.schema.DataFetcher; @@ -57,13 +56,14 @@ public CompletableFuture get(final DataFetchingEnvironment envi // First, get all policy Urns. final SearchResult gmsResult = _entityClient.search( + context + .getOperationContext() + .withSearchFlags(flags -> flags.setFulltext(true)), CORP_USER_ENTITY_NAME, query, Collections.emptyMap(), start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); // Then, get hydrate all users. final Map entities = diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/view/ListGlobalViewsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/view/ListGlobalViewsResolver.java index caa37f82648544..80d33b84b4c763 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/view/ListGlobalViewsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/view/ListGlobalViewsResolver.java @@ -15,7 +15,6 @@ import com.linkedin.datahub.graphql.generated.ListViewsResult; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; @@ -68,14 +67,13 @@ public CompletableFuture get(final DataFetchingEnvironment envi final SearchResult gmsResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), Constants.DATAHUB_VIEW_ENTITY_NAME, query, buildFilters(), DEFAULT_SORT_CRITERION, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); final ListViewsResult result = new ListViewsResult(); result.setStart(gmsResult.getFrom()); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/view/ListMyViewsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/view/ListMyViewsResolver.java index 945d2d50bcc3e1..fd029f9d6d3b2c 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/view/ListMyViewsResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/view/ListMyViewsResolver.java @@ -14,7 +14,6 @@ import com.linkedin.datahub.graphql.generated.ListViewsResult; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; @@ -71,14 +70,13 @@ public CompletableFuture get(final DataFetchingEnvironment envi final SearchResult gmsResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), Constants.DATAHUB_VIEW_ENTITY_NAME, query, buildFilters(viewType, context.getActorUrn()), DEFAULT_SORT_CRITERION, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); final ListViewsResult result = new ListViewsResult(); result.setStart(gmsResult.getFrom()); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/ChartType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/ChartType.java index 77168356682c20..ffd5bffdd43d9d 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/ChartType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/ChartType.java @@ -39,7 +39,6 @@ import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import com.linkedin.mxe.MetadataChangeProposal; @@ -154,13 +153,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "chart", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -173,7 +171,7 @@ public AutoCompleteResults autoComplete( @Nonnull QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete("chart", query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete(context.getOperationContext(), "chart", query, filters, limit); return AutoCompleteResultsMapper.map(result); } @@ -190,7 +188,12 @@ public BrowseResults browse( path.size() > 0 ? BROWSE_PATH_DELIMITER + String.join(BROWSE_PATH_DELIMITER, path) : ""; final BrowseResult result = _entityClient.browse( - "chart", pathStr, facetFilters, start, count, context.getAuthentication()); + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), + "chart", + pathStr, + facetFilters, + start, + count); return BrowseResultMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java index faede5cf9bb1be..16f1909af09f6a 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java @@ -45,6 +45,12 @@ public com.linkedin.metadata.query.SearchFlags apply(@Nonnull final SearchFlags if (searchFlags.getGetSuggestions() != null) { result.setGetSuggestions(searchFlags.getGetSuggestions()); } + if (searchFlags.getIncludeSoftDeleted() != null) { + result.setIncludeSoftDeleted(searchFlags.getIncludeSoftDeleted()); + } + if (searchFlags.getIncludeRestricted() != null) { + result.setIncludeRestricted(searchFlags.getIncludeRestricted()); + } if (searchFlags.getGroupingSpec() != null && searchFlags.getGroupingSpec().getGroupingCriteria() != null) { result.setGroupingSpec( diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java index 3ca018ea6f5c77..a859cd6c79e80d 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java @@ -31,6 +31,7 @@ import com.linkedin.datahub.graphql.generated.Notebook; import com.linkedin.datahub.graphql.generated.OwnershipTypeEntity; import com.linkedin.datahub.graphql.generated.QueryEntity; +import com.linkedin.datahub.graphql.generated.Restricted; import com.linkedin.datahub.graphql.generated.Role; import com.linkedin.datahub.graphql.generated.SchemaFieldEntity; import com.linkedin.datahub.graphql.generated.StructuredPropertyEntity; @@ -204,6 +205,11 @@ public Entity apply(Urn input) { ((QueryEntity) partialEntity).setUrn(input.toString()); ((QueryEntity) partialEntity).setType(EntityType.QUERY); } + if (input.getEntityType().equals(RESTRICTED_ENTITY_NAME)) { + partialEntity = new Restricted(); + ((Restricted) partialEntity).setUrn(input.toString()); + ((Restricted) partialEntity).setType(EntityType.RESTRICTED); + } return partialEntity; } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/container/ContainerType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/container/ContainerType.java index c6f2303fec7928..6408e2d6779fc5 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/container/ContainerType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/container/ContainerType.java @@ -18,7 +18,6 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import graphql.execution.DataFetcherResult; @@ -128,13 +127,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), ENTITY_NAME, query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -147,7 +145,8 @@ public AutoCompleteResults autoComplete( @Nonnull final QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete(ENTITY_NAME, query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete( + context.getOperationContext(), ENTITY_NAME, query, filters, limit); return AutoCompleteResultsMapper.map(result); } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpgroup/CorpGroupType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpgroup/CorpGroupType.java index 371cf6b280c20e..4eb038632c6c69 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpgroup/CorpGroupType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpgroup/CorpGroupType.java @@ -30,7 +30,6 @@ import com.linkedin.identity.CorpGroupEditableInfo; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import com.linkedin.mxe.MetadataChangeProposal; @@ -116,13 +115,12 @@ public SearchResults search( throws Exception { final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "corpGroup", query, Collections.emptyMap(), start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -135,7 +133,8 @@ public AutoCompleteResults autoComplete( @Nonnull final QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete("corpGroup", query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete( + context.getOperationContext(), "corpGroup", query, filters, limit); return AutoCompleteResultsMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpuser/CorpUserType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpuser/CorpUserType.java index 5749eef970fce8..3dab88fcc300b9 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpuser/CorpUserType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpuser/CorpUserType.java @@ -32,7 +32,6 @@ import com.linkedin.identity.CorpUserEditableInfo; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import com.linkedin.mxe.MetadataChangeProposal; @@ -116,13 +115,12 @@ public SearchResults search( throws Exception { final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "corpuser", query, Collections.emptyMap(), start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -135,7 +133,8 @@ public AutoCompleteResults autoComplete( @Nonnull final QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete("corpuser", query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete( + context.getOperationContext(), "corpuser", query, filters, limit); return AutoCompleteResultsMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/DashboardType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/DashboardType.java index 9f12fe7c0eccdd..4efcb42cf8e3cf 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/DashboardType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/DashboardType.java @@ -40,7 +40,6 @@ import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import com.linkedin.mxe.MetadataChangeProposal; @@ -153,13 +152,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "dashboard", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -172,7 +170,8 @@ public AutoCompleteResults autoComplete( @Nonnull QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete("dashboard", query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete( + context.getOperationContext(), "dashboard", query, filters, limit); return AutoCompleteResultsMapper.map(result); } @@ -189,7 +188,12 @@ public BrowseResults browse( path.size() > 0 ? BROWSE_PATH_DELIMITER + String.join(BROWSE_PATH_DELIMITER, path) : ""; final BrowseResult result = _entityClient.browse( - "dashboard", pathStr, facetFilters, start, count, context.getAuthentication()); + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), + "dashboard", + pathStr, + facetFilters, + start, + count); return BrowseResultMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/DataFlowType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/DataFlowType.java index fdc60635679f7a..fdb0bd603b27af 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/DataFlowType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/DataFlowType.java @@ -40,7 +40,6 @@ import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import com.linkedin.mxe.MetadataChangeProposal; @@ -147,13 +146,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "dataFlow", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -166,7 +164,8 @@ public AutoCompleteResults autoComplete( @Nonnull final QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete("dataFlow", query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete( + context.getOperationContext(), "dataFlow", query, filters, limit); return AutoCompleteResultsMapper.map(result); } @@ -183,7 +182,12 @@ public BrowseResults browse( path.size() > 0 ? BROWSE_PATH_DELIMITER + String.join(BROWSE_PATH_DELIMITER, path) : ""; final BrowseResult result = _entityClient.browse( - "dataFlow", pathStr, facetFilters, start, count, context.getAuthentication()); + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), + "dataFlow", + pathStr, + facetFilters, + start, + count); return BrowseResultMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/DataJobType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/DataJobType.java index 7593aede32a303..5127bee3f8a8c5 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/DataJobType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/DataJobType.java @@ -40,7 +40,6 @@ import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import com.linkedin.mxe.MetadataChangeProposal; @@ -149,13 +148,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "dataJob", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -168,7 +166,7 @@ public AutoCompleteResults autoComplete( @Nonnull final QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete("dataJob", query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete(context.getOperationContext(), "dataJob", query, filters, limit); return AutoCompleteResultsMapper.map(result); } @@ -185,7 +183,12 @@ public BrowseResults browse( path.size() > 0 ? BROWSE_PATH_DELIMITER + String.join(BROWSE_PATH_DELIMITER, path) : ""; final BrowseResult result = _entityClient.browse( - "dataJob", pathStr, facetFilters, start, count, context.getAuthentication()); + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), + "dataJob", + pathStr, + facetFilters, + start, + count); return BrowseResultMapper.map(result); } 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 6519a493f39917..c45cec34e5e79e 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 @@ -121,7 +121,11 @@ public AutoCompleteResults autoComplete( throws Exception { final AutoCompleteResult result = _entityClient.autoComplete( - DATA_PLATFORM_INSTANCE_ENTITY_NAME, query, filters, limit, context.getAuthentication()); + context.getOperationContext(), + DATA_PLATFORM_INSTANCE_ENTITY_NAME, + query, + filters, + limit); return AutoCompleteResultsMapper.map(result); } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/DataProductType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/DataProductType.java index 253a84ca34945d..a63c4bbbbf1d20 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/DataProductType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/DataProductType.java @@ -113,7 +113,7 @@ public AutoCompleteResults autoComplete( throws Exception { final AutoCompleteResult result = _entityClient.autoComplete( - DATA_PRODUCT_ENTITY_NAME, query, filters, limit, context.getAuthentication()); + context.getOperationContext(), DATA_PRODUCT_ENTITY_NAME, query, filters, limit); return AutoCompleteResultsMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java index fd31e1d394a927..0ae41eef6b1b1e 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java @@ -40,7 +40,6 @@ import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import com.linkedin.mxe.MetadataChangeProposal; @@ -168,13 +167,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), ENTITY_NAME, query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -187,7 +185,8 @@ public AutoCompleteResults autoComplete( @Nonnull final QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete(ENTITY_NAME, query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete( + context.getOperationContext(), ENTITY_NAME, query, filters, limit); return AutoCompleteResultsMapper.map(result); } @@ -204,7 +203,12 @@ public BrowseResults browse( path.size() > 0 ? BROWSE_PATH_DELIMITER + String.join(BROWSE_PATH_DELIMITER, path) : ""; final BrowseResult result = _entityClient.browse( - "dataset", pathStr, facetFilters, start, count, context.getAuthentication()); + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), + "dataset", + pathStr, + facetFilters, + start, + count); return BrowseResultMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/domain/DomainType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/domain/DomainType.java index cfacf0031c3185..d18633c763eedf 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/domain/DomainType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/domain/DomainType.java @@ -115,7 +115,7 @@ public AutoCompleteResults autoComplete( throws Exception { final AutoCompleteResult result = _entityClient.autoComplete( - Constants.DOMAIN_ENTITY_NAME, query, filters, limit, context.getAuthentication()); + context.getOperationContext(), Constants.DOMAIN_ENTITY_NAME, query, filters, limit); return AutoCompleteResultsMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeMapper.java index 23e793782e8dc6..3ecd01e99056b4 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeMapper.java @@ -43,6 +43,7 @@ public class EntityTypeMapper { .put(EntityType.SCHEMA_FIELD, "schemaField") .put(EntityType.STRUCTURED_PROPERTY, Constants.STRUCTURED_PROPERTY_ENTITY_NAME) .put(EntityType.ASSERTION, Constants.ASSERTION_ENTITY_NAME) + .put(EntityType.RESTRICTED, Constants.RESTRICTED_ENTITY_NAME) .build(); private static final Map ENTITY_NAME_TO_TYPE = diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/glossary/GlossaryTermType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/glossary/GlossaryTermType.java index 9114ff2670906f..b6b813cddf99be 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/glossary/GlossaryTermType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/glossary/GlossaryTermType.java @@ -28,7 +28,6 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import graphql.execution.DataFetcherResult; @@ -126,13 +125,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "glossaryTerm", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -146,7 +144,7 @@ public AutoCompleteResults autoComplete( throws Exception { final AutoCompleteResult result = _entityClient.autoComplete( - "glossaryTerm", query, filters, limit, context.getAuthentication()); + context.getOperationContext(), "glossaryTerm", query, filters, limit); return AutoCompleteResultsMapper.map(result); } @@ -163,7 +161,12 @@ public BrowseResults browse( path.size() > 0 ? BROWSE_PATH_DELIMITER + String.join(BROWSE_PATH_DELIMITER, path) : ""; final BrowseResult result = _entityClient.browse( - "glossaryTerm", pathStr, facetFilters, start, count, context.getAuthentication()); + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), + "glossaryTerm", + pathStr, + facetFilters, + start, + count); return BrowseResultMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLFeatureTableType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLFeatureTableType.java index da3ddd1115437f..b55f5740453934 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLFeatureTableType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLFeatureTableType.java @@ -28,7 +28,6 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import graphql.execution.DataFetcherResult; @@ -111,13 +110,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "mlFeatureTable", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -131,7 +129,7 @@ public AutoCompleteResults autoComplete( throws Exception { final AutoCompleteResult result = _entityClient.autoComplete( - "mlFeatureTable", query, filters, limit, context.getAuthentication()); + context.getOperationContext(), "mlFeatureTable", query, filters, limit); return AutoCompleteResultsMapper.map(result); } @@ -148,7 +146,12 @@ public BrowseResults browse( path.size() > 0 ? BROWSE_PATH_DELIMITER + String.join(BROWSE_PATH_DELIMITER, path) : ""; final BrowseResult result = _entityClient.browse( - "mlFeatureTable", pathStr, facetFilters, start, count, context.getAuthentication()); + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), + "mlFeatureTable", + pathStr, + facetFilters, + start, + count); return BrowseResultMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLFeatureType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLFeatureType.java index 6f94ea44cd476e..7046ee0f94eebf 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLFeatureType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLFeatureType.java @@ -20,7 +20,6 @@ import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import graphql.execution.DataFetcherResult; @@ -101,13 +100,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "mlFeature", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -120,7 +118,8 @@ public AutoCompleteResults autoComplete( @Nonnull final QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete("mlFeature", query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete( + context.getOperationContext(), "mlFeature", query, filters, limit); return AutoCompleteResultsMapper.map(result); } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLModelGroupType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLModelGroupType.java index d505b70effdd4c..8865d2acce12df 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLModelGroupType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLModelGroupType.java @@ -28,7 +28,6 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import graphql.execution.DataFetcherResult; @@ -111,13 +110,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "mlModelGroup", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -131,7 +129,7 @@ public AutoCompleteResults autoComplete( throws Exception { final AutoCompleteResult result = _entityClient.autoComplete( - "mlModelGroup", query, filters, limit, context.getAuthentication()); + context.getOperationContext(), "mlModelGroup", query, filters, limit); return AutoCompleteResultsMapper.map(result); } @@ -148,7 +146,12 @@ public BrowseResults browse( path.size() > 0 ? BROWSE_PATH_DELIMITER + String.join(BROWSE_PATH_DELIMITER, path) : ""; final BrowseResult result = _entityClient.browse( - "mlModelGroup", pathStr, facetFilters, start, count, context.getAuthentication()); + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), + "mlModelGroup", + pathStr, + facetFilters, + start, + count); return BrowseResultMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLModelType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLModelType.java index 27b791d78e78ea..a1c689d9f5c1dc 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLModelType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLModelType.java @@ -28,7 +28,6 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import graphql.execution.DataFetcherResult; @@ -106,13 +105,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "mlModel", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -125,7 +123,7 @@ public AutoCompleteResults autoComplete( @Nonnull final QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete("mlModel", query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete(context.getOperationContext(), "mlModel", query, filters, limit); return AutoCompleteResultsMapper.map(result); } @@ -142,7 +140,12 @@ public BrowseResults browse( path.size() > 0 ? BROWSE_PATH_DELIMITER + String.join(BROWSE_PATH_DELIMITER, path) : ""; final BrowseResult result = _entityClient.browse( - "mlModel", pathStr, facetFilters, start, count, context.getAuthentication()); + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), + "mlModel", + pathStr, + facetFilters, + start, + count); return BrowseResultMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLPrimaryKeyType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLPrimaryKeyType.java index 10cfe181dd292f..cccb05e8fa0f55 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLPrimaryKeyType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLPrimaryKeyType.java @@ -20,7 +20,6 @@ import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import graphql.execution.DataFetcherResult; @@ -101,13 +100,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "mlPrimaryKey", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -121,7 +119,7 @@ public AutoCompleteResults autoComplete( throws Exception { final AutoCompleteResult result = _entityClient.autoComplete( - "mlPrimaryKey", query, filters, limit, context.getAuthentication()); + context.getOperationContext(), "mlPrimaryKey", query, filters, limit); return AutoCompleteResultsMapper.map(result); } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/notebook/NotebookType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/notebook/NotebookType.java index b6990c3816b53f..a8e964581dfd59 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/notebook/NotebookType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/notebook/NotebookType.java @@ -38,7 +38,6 @@ import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import com.linkedin.mxe.MetadataChangeProposal; @@ -95,13 +94,12 @@ public SearchResults search( final Map facetFilters = Collections.emptyMap(); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), NOTEBOOK_ENTITY_NAME, query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -115,7 +113,7 @@ public AutoCompleteResults autoComplete( throws Exception { final AutoCompleteResult result = _entityClient.autoComplete( - NOTEBOOK_ENTITY_NAME, query, filters, limit, context.getAuthentication()); + context.getOperationContext(), NOTEBOOK_ENTITY_NAME, query, filters, limit); return AutoCompleteResultsMapper.map(result); } @@ -135,7 +133,12 @@ public BrowseResults browse( path.size() > 0 ? BROWSE_PATH_DELIMITER + String.join(BROWSE_PATH_DELIMITER, path) : ""; final BrowseResult result = _entityClient.browse( - NOTEBOOK_ENTITY_NAME, pathStr, facetFilters, start, count, context.getAuthentication()); + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), + NOTEBOOK_ENTITY_NAME, + pathStr, + facetFilters, + start, + count); return BrowseResultMapper.map(result); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/restricted/RestrictedMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/restricted/RestrictedMapper.java new file mode 100644 index 00000000000000..61186fc9f77e5a --- /dev/null +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/restricted/RestrictedMapper.java @@ -0,0 +1,32 @@ +package com.linkedin.datahub.graphql.types.restricted; + +import com.linkedin.common.urn.Urn; +import com.linkedin.datahub.graphql.generated.EntityType; +import com.linkedin.datahub.graphql.generated.Restricted; +import com.linkedin.entity.EntityResponse; +import com.linkedin.metadata.service.RestrictedService; +import javax.annotation.Nonnull; + +public class RestrictedMapper { + + public static final RestrictedMapper INSTANCE = new RestrictedMapper(); + + public static Restricted map( + @Nonnull final EntityResponse entityResponse, + @Nonnull final RestrictedService restrictedService) { + return INSTANCE.apply(entityResponse, restrictedService); + } + + public Restricted apply( + @Nonnull final EntityResponse entityResponse, + @Nonnull final RestrictedService restrictedService) { + final Restricted result = new Restricted(); + Urn entityUrn = entityResponse.getUrn(); + String restrictedUrnString = restrictedService.encryptRestrictedUrn(entityUrn).toString(); + + result.setUrn(restrictedUrnString); + result.setType(EntityType.RESTRICTED); + + return result; + } +} diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/restricted/RestrictedType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/restricted/RestrictedType.java new file mode 100644 index 00000000000000..a2030bb596d105 --- /dev/null +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/restricted/RestrictedType.java @@ -0,0 +1,103 @@ +package com.linkedin.datahub.graphql.types.restricted; + +import com.google.common.collect.ImmutableSet; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.generated.Entity; +import com.linkedin.datahub.graphql.generated.Restricted; +import com.linkedin.datahub.graphql.types.EntityType; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.metadata.service.RestrictedService; +import graphql.execution.DataFetcherResult; +import java.util.*; +import java.util.function.Function; +import java.util.stream.Collectors; +import javax.annotation.Nonnull; +import lombok.RequiredArgsConstructor; + +@RequiredArgsConstructor +public class RestrictedType implements EntityType { + public static final Set ASPECTS_TO_FETCH = ImmutableSet.of(); + + private final EntityClient _entityClient; + private final RestrictedService _restrictedService; + + @Override + public com.linkedin.datahub.graphql.generated.EntityType type() { + return com.linkedin.datahub.graphql.generated.EntityType.RESTRICTED; + } + + @Override + public Function getKeyProvider() { + return Entity::getUrn; + } + + @Override + public Class objectClass() { + return Restricted.class; + } + + @Override + public List> batchLoad( + @Nonnull List urns, @Nonnull QueryContext context) throws Exception { + final List restrictedUrns = + urns.stream().map(UrnUtils::getUrn).collect(Collectors.toList()); + final List entityUrns = + restrictedUrns.stream() + .map(_restrictedService::decryptRestrictedUrn) + .collect(Collectors.toList()); + + // Create a map for entityType: entityUrns so we can fetch by entity type below + final Map> entityTypeToUrns = createEntityTypeToUrnsMap(entityUrns); + + try { + // Fetch from the DB for each entity type and add to one result map + final Map entities = new HashMap<>(); + entityTypeToUrns + .keySet() + .forEach( + entityType -> { + try { + entities.putAll( + _entityClient.batchGetV2( + entityType, + new HashSet<>(entityTypeToUrns.get(entityType)), + ASPECTS_TO_FETCH, + context.getAuthentication())); + } catch (Exception e) { + throw new RuntimeException("Failed to fetch restricted entities", e); + } + }); + + final List gmsResults = new ArrayList<>(); + for (Urn urn : entityUrns) { + gmsResults.add(entities.getOrDefault(urn, null)); + } + return gmsResults.stream() + .map( + gmsResult -> + gmsResult == null + ? null + : DataFetcherResult.newResult() + .data(RestrictedMapper.map(gmsResult, _restrictedService)) + .build()) + .collect(Collectors.toList()); + } catch (Exception e) { + throw new RuntimeException("Failed to batch load Queries", e); + } + } + + private Map> createEntityTypeToUrnsMap(final List urns) { + final Map> entityTypeToUrns = new HashMap<>(); + urns.forEach( + urn -> { + String entityType = urn.getEntityType(); + List existingUrns = + entityTypeToUrns.computeIfAbsent(entityType, k -> new ArrayList<>()); + existingUrns.add(urn); + }); + return entityTypeToUrns; + } +} diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/rolemetadata/RoleType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/rolemetadata/RoleType.java index d51e0d06c0fdaa..cc77ee46d65dcc 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/rolemetadata/RoleType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/rolemetadata/RoleType.java @@ -18,7 +18,6 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import graphql.execution.DataFetcherResult; @@ -104,13 +103,12 @@ public SearchResults search( throws Exception { final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), Constants.ROLE_ENTITY_NAME, query, Collections.emptyMap(), start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -124,7 +122,7 @@ public AutoCompleteResults autoComplete( throws Exception { final AutoCompleteResult result = _entityClient.autoComplete( - Constants.ROLE_ENTITY_NAME, query, filters, limit, context.getAuthentication()); + context.getOperationContext(), Constants.ROLE_ENTITY_NAME, query, filters, limit); return AutoCompleteResultsMapper.map(result); } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/tag/TagType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/tag/TagType.java index c56833cc817eb7..a3c9bc380bdcfe 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/tag/TagType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/tag/TagType.java @@ -28,7 +28,6 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.SearchResult; import com.linkedin.mxe.MetadataChangeProposal; @@ -116,13 +115,12 @@ public SearchResults search( final Map facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS); final SearchResult searchResult = _entityClient.search( + context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), "tag", query, facetFilters, start, - count, - context.getAuthentication(), - new SearchFlags().setFulltext(true)); + count); return UrnSearchResultsMapper.map(searchResult); } @@ -135,7 +133,7 @@ public AutoCompleteResults autoComplete( @Nonnull QueryContext context) throws Exception { final AutoCompleteResult result = - _entityClient.autoComplete("tag", query, filters, limit, context.getAuthentication()); + _entityClient.autoComplete(context.getOperationContext(), "tag", query, filters, limit); 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 d8dedf20e3c0c1..855cf19f0cb3b8 100644 --- a/datahub-graphql-core/src/main/resources/entity.graphql +++ b/datahub-graphql-core/src/main/resources/entity.graphql @@ -961,6 +961,11 @@ enum EntityType { """ ENTITY_TYPE + """" + A type of entity that is restricted to the user + """ + RESTRICTED + """ Another entity type - refer to a provided entity type urn. """ @@ -11834,3 +11839,29 @@ type EntityTypeInfo { """ description: String } + +""" +A restricted entity that the user does not have full permissions to view. +This entity type does not relate to an entity type in the database. +""" +type Restricted implements Entity & EntityWithRelationships { + """ + The primary key of the restricted entity + """ + urn: String! + + """ + The standard Entity Type + """ + type: EntityType! + + """ + Edges extending from this entity + """ + relationships(input: RelationshipsInput!): EntityRelationshipsResult + + """ + Edges extending from this entity grouped by direction in the lineage graph + """ + lineage(input: LineageInput!): EntityLineageResult +} diff --git a/datahub-graphql-core/src/main/resources/search.graphql b/datahub-graphql-core/src/main/resources/search.graphql index a906362cee1854..60b3d73d77fdb1 100644 --- a/datahub-graphql-core/src/main/resources/search.graphql +++ b/datahub-graphql-core/src/main/resources/search.graphql @@ -152,6 +152,16 @@ input SearchFlags { Note: This is an experimental feature and is subject to change. """ groupingSpec: GroupingSpec + + """ + Whether to include soft deleted entities + """ + includeSoftDeleted: Boolean + + """ + Whether to include restricted entities + """ + includeRestricted: Boolean } """ 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 ac07053e59d75a..e8c968ab768f8c 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 @@ -17,6 +17,8 @@ import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.mxe.MetadataChangeProposal; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.util.List; import org.mockito.Mockito; @@ -48,10 +50,15 @@ public static QueryContext getMockAllowContext(String actorUrn) { Mockito.when(mockAuthorizer.authorize(Mockito.any())).thenReturn(result); Mockito.when(mockContext.getAuthorizer()).thenReturn(mockAuthorizer); - Mockito.when(mockContext.getAuthentication()) - .thenReturn( - new Authentication( - new Actor(ActorType.USER, UrnUtils.getUrn(actorUrn).getId()), "creds")); + Authentication authentication = + new Authentication(new Actor(ActorType.USER, UrnUtils.getUrn(actorUrn).getId()), "creds"); + Mockito.when(mockContext.getAuthentication()).thenReturn(authentication); + OperationContext operationContext = + TestOperationContexts.userContextNoSearchAuthorization( + mock(EntityRegistry.class), mockAuthorizer, authentication); + Mockito.when(mockContext.getOperationContext()).thenReturn(operationContext); + Mockito.when(mockContext.getOperationContext()) + .thenReturn(Mockito.mock(OperationContext.class)); return mockContext; } @@ -69,6 +76,8 @@ public static QueryContext getMockAllowContext(String actorUrn, AuthorizationReq .thenReturn( new Authentication( new Actor(ActorType.USER, UrnUtils.getUrn(actorUrn).getId()), "creds")); + Mockito.when(mockContext.getOperationContext()) + .thenReturn(Mockito.mock(OperationContext.class)); return mockContext; } @@ -90,6 +99,8 @@ public static QueryContext getMockDenyContext(String actorUrn) { .thenReturn( new Authentication( new Actor(ActorType.USER, UrnUtils.getUrn(actorUrn).getId()), "creds")); + Mockito.when(mockContext.getOperationContext()) + .thenReturn(Mockito.mock(OperationContext.class)); return mockContext; } @@ -107,6 +118,8 @@ public static QueryContext getMockDenyContext(String actorUrn, AuthorizationRequ .thenReturn( new Authentication( new Actor(ActorType.USER, UrnUtils.getUrn(actorUrn).getId()), "creds")); + Mockito.when(mockContext.getOperationContext()) + .thenReturn(Mockito.mock(OperationContext.class)); return mockContext; } diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolverTest.java index 419eb71d5e143d..ad30e48d8361b8 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolverTest.java @@ -1,6 +1,7 @@ package com.linkedin.datahub.graphql.resolvers.auth; import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*; +import static org.mockito.ArgumentMatchers.any; import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; @@ -11,7 +12,6 @@ import com.linkedin.datahub.graphql.generated.ListAccessTokenResult; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; @@ -43,14 +43,13 @@ public void testGetSuccess() throws Exception { final Authentication testAuth = getAuthentication(mockEnv); Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.ACCESS_TOKEN_ENTITY_NAME), Mockito.eq(""), Mockito.eq(buildFilter(filters, Collections.emptyList())), Mockito.any(SortCriterion.class), Mockito.eq(input.getStart()), - Mockito.eq(input.getCount()), - Mockito.eq(testAuth), - Mockito.any(SearchFlags.class))) + Mockito.eq(input.getCount()))) .thenReturn( new SearchResult() .setFrom(0) diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/browse/BrowseV2ResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/browse/BrowseV2ResolverTest.java index 41797fac636f1d..892ba4e1ebb3eb 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/browse/BrowseV2ResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/browse/BrowseV2ResolverTest.java @@ -21,7 +21,6 @@ import com.linkedin.metadata.browse.BrowseResultGroupV2Array; import com.linkedin.metadata.browse.BrowseResultMetadata; import com.linkedin.metadata.browse.BrowseResultV2; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; import com.linkedin.metadata.query.filter.Criterion; @@ -257,14 +256,13 @@ private static EntityClient initMockEntityClient( EntityClient client = Mockito.mock(EntityClient.class); Mockito.when( client.browseV2( + Mockito.any(), Mockito.eq(ImmutableList.of(entityName)), Mockito.eq(path), Mockito.eq(filter), Mockito.eq(query), Mockito.eq(start), - Mockito.eq(limit), - Mockito.any(Authentication.class), - Mockito.nullable(SearchFlags.class))) + Mockito.eq(limit))) .thenReturn(result); return client; } diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/container/ContainerEntitiesResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/container/ContainerEntitiesResolverTest.java index 1203f4e22bdc23..c63c9bccab68b5 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/container/ContainerEntitiesResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/container/ContainerEntitiesResolverTest.java @@ -1,5 +1,7 @@ package com.linkedin.datahub.graphql.resolvers.container; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -22,6 +24,7 @@ import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.search.SearchResultMetadata; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import java.util.Collections; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -34,7 +37,7 @@ public class ContainerEntitiesResolverTest { @Test public void testGetSuccess() throws Exception { // Create resolver - EntityClient mockClient = Mockito.mock(EntityClient.class); + EntityClient mockClient = mock(EntityClient.class); final String childUrn = "urn:li:dataset:(test,test,test)"; final String containerUrn = "urn:li:container:test-container"; @@ -47,6 +50,7 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.searchAcrossEntities( + any(), Mockito.eq(ContainerEntitiesResolver.CONTAINABLE_ENTITY_NAMES), Mockito.eq("*"), Mockito.eq( @@ -59,8 +63,7 @@ public void testGetSuccess() throws Exception { Mockito.eq(0), Mockito.eq(20), Mockito.eq(null), - Mockito.eq(null), - Mockito.any(Authentication.class))) + Mockito.eq(null))) .thenReturn( new SearchResult() .setFrom(0) @@ -76,9 +79,10 @@ public void testGetSuccess() throws Exception { ContainerEntitiesResolver resolver = new ContainerEntitiesResolver(mockClient); // Execute resolver - QueryContext mockContext = Mockito.mock(QueryContext.class); - Mockito.when(mockContext.getAuthentication()).thenReturn(Mockito.mock(Authentication.class)); - DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + QueryContext mockContext = mock(QueryContext.class); + Mockito.when(mockContext.getAuthentication()).thenReturn(mock(Authentication.class)); + Mockito.when(mockContext.getOperationContext()).thenReturn(mock(OperationContext.class)); + DataFetchingEnvironment mockEnv = mock(DataFetchingEnvironment.class); Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); Mockito.when(mockEnv.getContext()).thenReturn(mockContext); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/CreateDomainResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/CreateDomainResolverTest.java index 6184760abfabda..33cd1cb63d621e 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/CreateDomainResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/CreateDomainResolverTest.java @@ -2,6 +2,7 @@ import static com.linkedin.datahub.graphql.TestUtils.*; import static com.linkedin.metadata.Constants.DOMAIN_PROPERTIES_ASPECT_NAME; +import static org.mockito.ArgumentMatchers.any; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -71,14 +72,14 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.filter( + any(), Mockito.eq(Constants.DOMAIN_ENTITY_NAME), Mockito.eq( DomainUtils.buildNameAndParentDomainFilter( TEST_INPUT.getName(), TEST_PARENT_DOMAIN_URN)), Mockito.eq(null), Mockito.any(Integer.class), - Mockito.any(Integer.class), - Mockito.any(Authentication.class))) + Mockito.any(Integer.class))) .thenReturn(new SearchResult().setEntities(new SearchEntityArray())); resolver.get(mockEnv).get(); @@ -121,12 +122,12 @@ public void testGetSuccessNoParentDomain() throws Exception { Mockito.when( mockClient.filter( + any(), Mockito.eq(Constants.DOMAIN_ENTITY_NAME), Mockito.eq(DomainUtils.buildNameAndParentDomainFilter(TEST_INPUT.getName(), null)), Mockito.eq(null), Mockito.any(Integer.class), - Mockito.any(Integer.class), - Mockito.any(Authentication.class))) + Mockito.any(Integer.class))) .thenReturn(new SearchResult().setEntities(new SearchEntityArray())); resolver.get(mockEnv).get(); @@ -194,14 +195,14 @@ public void testGetNameConflict() throws Exception { Mockito.when( mockClient.filter( + any(), Mockito.eq(Constants.DOMAIN_ENTITY_NAME), Mockito.eq( DomainUtils.buildNameAndParentDomainFilter( TEST_INPUT.getName(), TEST_PARENT_DOMAIN_URN)), Mockito.eq(null), Mockito.any(Integer.class), - Mockito.any(Integer.class), - Mockito.any(Authentication.class))) + Mockito.any(Integer.class))) .thenReturn( new SearchResult() .setEntities(new SearchEntityArray(new SearchEntity().setEntity(TEST_DOMAIN_URN)))); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/DeleteDomainResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/DeleteDomainResolverTest.java index 5632654a26ad92..69fb98fbf9e310 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/DeleteDomainResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/DeleteDomainResolverTest.java @@ -1,6 +1,7 @@ package com.linkedin.datahub.graphql.resolvers.domain; import static com.linkedin.datahub.graphql.TestUtils.*; +import static org.mockito.ArgumentMatchers.any; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -31,12 +32,12 @@ public void testGetSuccess() throws Exception { // Domain has 0 child domains Mockito.when( mockClient.filter( + any(), Mockito.eq("domain"), Mockito.any(), Mockito.any(), Mockito.eq(0), - Mockito.eq(1), - Mockito.any())) + Mockito.eq(1))) .thenReturn(new SearchResult().setNumEntities(0)); assertTrue(resolver.get(mockEnv).get()); @@ -60,12 +61,12 @@ public void testDeleteWithChildDomains() throws Exception { // Domain has child domains Mockito.when( mockClient.filter( + any(), Mockito.eq("domain"), Mockito.any(), Mockito.any(), Mockito.eq(0), - Mockito.eq(1), - Mockito.any())) + Mockito.eq(1))) .thenReturn(new SearchResult().setNumEntities(1)); assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolverTest.java index c6e6cdc7f018e2..f970f9e2ea431d 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolverTest.java @@ -1,6 +1,8 @@ package com.linkedin.datahub.graphql.resolvers.domain; import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -24,6 +26,7 @@ import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.search.SearchResultMetadata; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import java.util.Collections; import java.util.stream.Collectors; import org.mockito.Mockito; @@ -50,6 +53,7 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.searchAcrossEntities( + any(), Mockito.eq( SEARCHABLE_ENTITY_TYPES.stream() .map(EntityTypeMapper::getName) @@ -65,8 +69,7 @@ public void testGetSuccess() throws Exception { Mockito.eq(0), Mockito.eq(20), Mockito.eq(null), - Mockito.eq(null), - Mockito.any(Authentication.class))) + Mockito.eq(null))) .thenReturn( new SearchResult() .setFrom(0) @@ -84,6 +87,7 @@ public void testGetSuccess() throws Exception { // Execute resolver QueryContext mockContext = Mockito.mock(QueryContext.class); Mockito.when(mockContext.getAuthentication()).thenReturn(Mockito.mock(Authentication.class)); + Mockito.when(mockContext.getOperationContext()).thenReturn(mock(OperationContext.class)); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); Mockito.when(mockEnv.getContext()).thenReturn(mockContext); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/ListDomainsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/ListDomainsResolverTest.java index ffc3e823d83510..53a16ed5f6cc8a 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/ListDomainsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/ListDomainsResolverTest.java @@ -2,10 +2,10 @@ import static com.linkedin.datahub.graphql.TestUtils.*; import static com.linkedin.metadata.Constants.*; +import static org.mockito.ArgumentMatchers.any; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertThrows; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; import com.linkedin.datahub.graphql.QueryContext; @@ -13,7 +13,6 @@ import com.linkedin.datahub.graphql.resolvers.mutate.util.DomainUtils; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; import com.linkedin.metadata.search.SearchEntity; @@ -43,6 +42,7 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.DOMAIN_ENTITY_NAME), Mockito.eq(""), Mockito.eq(DomainUtils.buildParentDomainFilter(TEST_PARENT_DOMAIN_URN)), @@ -51,9 +51,7 @@ public void testGetSuccess() throws Exception { .setField(DOMAIN_CREATED_TIME_INDEX_FIELD_NAME) .setOrder(SortOrder.DESCENDING)), Mockito.eq(0), - Mockito.eq(20), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true)))) + Mockito.eq(20))) .thenReturn( new SearchResult() .setFrom(0) @@ -87,6 +85,7 @@ public void testGetSuccessNoParentDomain() throws Exception { Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.DOMAIN_ENTITY_NAME), Mockito.eq(""), Mockito.eq(DomainUtils.buildParentDomainFilter(null)), @@ -95,9 +94,7 @@ public void testGetSuccessNoParentDomain() throws Exception { .setField(DOMAIN_CREATED_TIME_INDEX_FIELD_NAME) .setOrder(SortOrder.DESCENDING)), Mockito.eq(0), - Mockito.eq(20), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true)))) + Mockito.eq(20))) .thenReturn( new SearchResult() .setFrom(0) @@ -139,13 +136,12 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockClient, Mockito.times(0)) .search( + any(), Mockito.any(), Mockito.eq("*"), Mockito.anyMap(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + Mockito.anyInt()); } @Test @@ -155,13 +151,12 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RemoteInvocationException.class) .when(mockClient) .search( + any(), Mockito.any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + Mockito.anyInt()); ListDomainsResolver resolver = new ListDomainsResolver(mockClient); // Execute resolver diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/MoveDomainResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/MoveDomainResolverTest.java index 1aa7f5aef467ce..83ebe481708b52 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/MoveDomainResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/MoveDomainResolverTest.java @@ -2,6 +2,7 @@ import static com.linkedin.datahub.graphql.TestUtils.*; import static com.linkedin.metadata.Constants.*; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.testng.Assert.assertThrows; import static org.testng.Assert.assertTrue; @@ -53,14 +54,14 @@ private MetadataChangeProposal setupTests( Mockito.when( mockClient.filter( + any(), Mockito.eq(Constants.DOMAIN_ENTITY_NAME), Mockito.eq( DomainUtils.buildNameAndParentDomainFilter( name, Urn.createFromString(PARENT_DOMAIN_URN))), Mockito.eq(null), Mockito.any(Integer.class), - Mockito.any(Integer.class), - Mockito.any(Authentication.class))) + Mockito.any(Integer.class))) .thenReturn(new SearchResult().setEntities(new SearchEntityArray())); DomainProperties properties = new DomainProperties(); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolverTest.java index 6653b19d6ef2bd..72937cb650368f 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolverTest.java @@ -3,6 +3,7 @@ import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; import static com.linkedin.datahub.graphql.TestUtils.getMockEntityService; import static com.linkedin.metadata.Constants.*; +import static org.mockito.ArgumentMatchers.any; import static org.testng.Assert.assertThrows; import com.datahub.authentication.Authentication; @@ -127,12 +128,12 @@ public void testGetFailureExistingTermSameName() throws Exception { Mockito.when( mockClient.filter( + any(), Mockito.eq(GLOSSARY_TERM_ENTITY_NAME), Mockito.any(), Mockito.eq(null), Mockito.eq(0), - Mockito.eq(1000), - Mockito.any())) + Mockito.eq(1000))) .thenReturn( new SearchResult() .setEntities( @@ -177,12 +178,12 @@ private EntityClient initMockClient() throws Exception { Mockito.when( mockClient.filter( + any(), Mockito.eq(GLOSSARY_TERM_ENTITY_NAME), Mockito.any(), Mockito.eq(null), Mockito.eq(0), - Mockito.eq(1000), - Mockito.any())) + Mockito.eq(1000))) .thenReturn(new SearchResult().setEntities(new SearchEntityArray())); Mockito.when( mockClient.batchGetV2( diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryNodesResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryNodesResolverTest.java index b879baf1e65dcd..60787fc47c88a5 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryNodesResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryNodesResolverTest.java @@ -1,5 +1,7 @@ package com.linkedin.datahub.graphql.resolvers.glossary; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -21,6 +23,7 @@ import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -33,6 +36,7 @@ public class GetRootGlossaryNodesResolverTest { public void testGetSuccess() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); QueryContext mockContext = Mockito.mock(QueryContext.class); + Mockito.when(mockContext.getOperationContext()).thenReturn(mock(OperationContext.class)); Mockito.when(mockContext.getAuthentication()).thenReturn(Mockito.mock(Authentication.class)); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); @@ -41,12 +45,12 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.filter( + any(), Mockito.eq(Constants.GLOSSARY_NODE_ENTITY_NAME), Mockito.eq(buildGlossaryEntitiesFilter()), Mockito.eq(null), Mockito.eq(0), - Mockito.eq(100), - Mockito.any(Authentication.class))) + Mockito.eq(100))) .thenReturn( new SearchResult() .setEntities( diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryTermsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryTermsResolverTest.java index 201bea752d53f0..51760ff9d37f25 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryTermsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/GetRootGlossaryTermsResolverTest.java @@ -1,5 +1,7 @@ package com.linkedin.datahub.graphql.resolvers.glossary; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -21,6 +23,7 @@ import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -33,6 +36,7 @@ public class GetRootGlossaryTermsResolverTest { public void testGetSuccess() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); QueryContext mockContext = Mockito.mock(QueryContext.class); + Mockito.when(mockContext.getOperationContext()).thenReturn(mock(OperationContext.class)); Mockito.when(mockContext.getAuthentication()).thenReturn(Mockito.mock(Authentication.class)); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); @@ -41,12 +45,12 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.filter( + any(), Mockito.eq(Constants.GLOSSARY_TERM_ENTITY_NAME), Mockito.eq(buildGlossaryEntitiesFilter()), Mockito.eq(null), Mockito.eq(0), - Mockito.eq(100), - Mockito.any(Authentication.class))) + Mockito.eq(100))) .thenReturn( new SearchResult() .setEntities( diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateNameResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateNameResolverTest.java index 062c1da5e038dc..5b858b810657ab 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateNameResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateNameResolverTest.java @@ -2,6 +2,7 @@ import static com.linkedin.datahub.graphql.TestUtils.*; import static com.linkedin.metadata.Constants.*; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.testng.Assert.assertThrows; import static org.testng.Assert.assertTrue; @@ -125,13 +126,13 @@ public void testGetSuccessForDomain() throws Exception { Mockito.when( mockClient.filter( + any(), Mockito.eq(Constants.DOMAIN_ENTITY_NAME), Mockito.eq( DomainUtils.buildNameAndParentDomainFilter(INPUT_FOR_DOMAIN.getName(), null)), Mockito.eq(null), Mockito.any(Integer.class), - Mockito.any(Integer.class), - Mockito.any(Authentication.class))) + Mockito.any(Integer.class))) .thenReturn(new SearchResult().setEntities(new SearchEntityArray())); DomainProperties properties = new DomainProperties(); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/incident/EntityIncidentsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/incident/EntityIncidentsResolverTest.java index a3f4b508dfc3e3..86c7b86978127d 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/incident/EntityIncidentsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/incident/EntityIncidentsResolverTest.java @@ -1,6 +1,7 @@ package com.linkedin.datahub.graphql.resolvers.incident; import static com.linkedin.datahub.graphql.resolvers.incident.EntityIncidentsResolver.*; +import static org.mockito.Mockito.mock; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -34,6 +35,7 @@ import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.search.utils.QueryUtils; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import java.util.HashMap; import java.util.Map; import org.mockito.Mockito; @@ -86,12 +88,12 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.filter( + Mockito.any(), Mockito.eq(Constants.INCIDENT_ENTITY_NAME), Mockito.eq(expectedFilter), Mockito.eq(expectedSort), Mockito.eq(0), - Mockito.eq(10), - Mockito.any(Authentication.class))) + Mockito.eq(10))) .thenReturn( new SearchResult() .setFrom(0) @@ -120,6 +122,7 @@ public void testGetSuccess() throws Exception { // Execute resolver QueryContext mockContext = Mockito.mock(QueryContext.class); Mockito.when(mockContext.getAuthentication()).thenReturn(Mockito.mock(Authentication.class)); + Mockito.when(mockContext.getOperationContext()).thenReturn(mock(OperationContext.class)); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); Mockito.when(mockEnv.getArgumentOrDefault(Mockito.eq("start"), Mockito.eq(0))).thenReturn(0); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/IngestTestUtils.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/IngestTestUtils.java index e5cb43c4dab617..e0555f5886b8bb 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/IngestTestUtils.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/IngestTestUtils.java @@ -1,5 +1,6 @@ package com.linkedin.datahub.graphql.resolvers.ingest; +import static org.mockito.Mockito.mock; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -21,6 +22,7 @@ import com.linkedin.ingestion.DataHubIngestionSourceSchedule; import com.linkedin.metadata.Constants; import com.linkedin.secret.DataHubSecretValue; +import io.datahubproject.metadata.context.OperationContext; import org.mockito.Mockito; public class IngestTestUtils { @@ -43,6 +45,7 @@ public static QueryContext getMockAllowContext() { Mockito.when(mockContext.getAuthorizer()).thenReturn(mockAuthorizer); Mockito.when(mockContext.getAuthentication()).thenReturn(Mockito.mock(Authentication.class)); + Mockito.when(mockContext.getOperationContext()).thenReturn(mock(OperationContext.class)); return mockContext; } diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/execution/IngestionSourceExecutionRequestsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/execution/IngestionSourceExecutionRequestsResolverTest.java index fdb150e6924417..d64a41d59b30ec 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/execution/IngestionSourceExecutionRequestsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/execution/IngestionSourceExecutionRequestsResolverTest.java @@ -1,6 +1,8 @@ package com.linkedin.datahub.graphql.resolvers.ingest.execution; import static com.linkedin.datahub.graphql.resolvers.ingest.IngestTestUtils.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -27,6 +29,7 @@ import com.linkedin.metadata.search.SearchResult; import com.linkedin.r2.RemoteInvocationException; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import java.util.HashSet; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -36,17 +39,17 @@ public class IngestionSourceExecutionRequestsResolverTest { @Test public void testGetSuccess() throws Exception { // Create resolver - EntityClient mockClient = Mockito.mock(EntityClient.class); + EntityClient mockClient = mock(EntityClient.class); // Mock filter response Mockito.when( mockClient.filter( + any(), Mockito.eq(Constants.EXECUTION_REQUEST_ENTITY_NAME), Mockito.any(Filter.class), Mockito.any(SortCriterion.class), Mockito.eq(0), - Mockito.eq(10), - Mockito.any(Authentication.class))) + Mockito.eq(10))) .thenReturn( new SearchResult() .setFrom(0) @@ -101,7 +104,8 @@ public void testGetSuccess() throws Exception { // Execute resolver QueryContext mockContext = getMockAllowContext(); - DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockContext.getOperationContext()).thenReturn(mock(OperationContext.class)); + DataFetchingEnvironment mockEnv = mock(DataFetchingEnvironment.class); Mockito.when(mockEnv.getArgument(Mockito.eq("start"))).thenReturn(0); Mockito.when(mockEnv.getArgument(Mockito.eq("count"))).thenReturn(10); Mockito.when(mockEnv.getContext()).thenReturn(mockContext); @@ -121,12 +125,12 @@ public void testGetSuccess() throws Exception { @Test public void testGetUnauthorized() throws Exception { // Create resolver - EntityClient mockClient = Mockito.mock(EntityClient.class); + EntityClient mockClient = mock(EntityClient.class); IngestionSourceExecutionRequestsResolver resolver = new IngestionSourceExecutionRequestsResolver(mockClient); // Execute resolver - DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + DataFetchingEnvironment mockEnv = mock(DataFetchingEnvironment.class); QueryContext mockContext = getMockDenyContext(); Mockito.when(mockEnv.getArgument(Mockito.eq("start"))).thenReturn(0); Mockito.when(mockEnv.getArgument(Mockito.eq("count"))).thenReturn(10); @@ -140,18 +144,13 @@ public void testGetUnauthorized() throws Exception { .batchGetV2( Mockito.any(), Mockito.anySet(), Mockito.anySet(), Mockito.any(Authentication.class)); Mockito.verify(mockClient, Mockito.times(0)) - .list( - Mockito.any(), - Mockito.anyMap(), - Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class)); + .list(Mockito.any(), Mockito.any(), Mockito.anyMap(), Mockito.anyInt(), Mockito.anyInt()); } @Test public void testGetEntityClientException() throws Exception { // Create resolver - EntityClient mockClient = Mockito.mock(EntityClient.class); + EntityClient mockClient = mock(EntityClient.class); Mockito.doThrow(RemoteInvocationException.class) .when(mockClient) .batchGetV2( @@ -160,7 +159,7 @@ public void testGetEntityClientException() throws Exception { new IngestionSourceExecutionRequestsResolver(mockClient); // Execute resolver - DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + DataFetchingEnvironment mockEnv = mock(DataFetchingEnvironment.class); QueryContext mockContext = getMockAllowContext(); Mockito.when(mockEnv.getArgument(Mockito.eq("start"))).thenReturn(0); Mockito.when(mockEnv.getArgument(Mockito.eq("count"))).thenReturn(10); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolverTest.java index 7d89f4aafa01a5..7a1876466573d4 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolverTest.java @@ -1,6 +1,8 @@ package com.linkedin.datahub.graphql.resolvers.ingest.secret; import static com.linkedin.datahub.graphql.resolvers.ingest.IngestTestUtils.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -14,7 +16,6 @@ import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; @@ -22,6 +23,7 @@ import com.linkedin.r2.RemoteInvocationException; import com.linkedin.secret.DataHubSecretValue; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import java.util.HashSet; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -39,14 +41,13 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.SECRETS_ENTITY_NAME), Mockito.eq(""), Mockito.eq(null), Mockito.any(SortCriterion.class), Mockito.eq(0), - Mockito.eq(20), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true)))) + Mockito.eq(20))) .thenReturn( new SearchResult() .setFrom(0) @@ -101,6 +102,7 @@ public void testGetUnauthorized() throws Exception { QueryContext mockContext = getMockDenyContext(); Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + Mockito.when(mockContext.getOperationContext()).thenReturn(mock(OperationContext.class)); assertThrows(RuntimeException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockClient, Mockito.times(0)) @@ -108,14 +110,13 @@ public void testGetUnauthorized() throws Exception { Mockito.any(), Mockito.anySet(), Mockito.anySet(), Mockito.any(Authentication.class)); Mockito.verify(mockClient, Mockito.times(0)) .search( + any(), Mockito.any(), Mockito.eq(""), Mockito.eq(null), Mockito.any(SortCriterion.class), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + Mockito.anyInt()); } @Test diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/source/ListIngestionSourceResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/source/ListIngestionSourceResolverTest.java index a86d67fcd15c18..4dfce0e0c2ee83 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/source/ListIngestionSourceResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/source/ListIngestionSourceResolverTest.java @@ -1,6 +1,8 @@ package com.linkedin.datahub.graphql.resolvers.ingest.source; import static com.linkedin.datahub.graphql.resolvers.ingest.IngestTestUtils.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -16,12 +18,12 @@ import com.linkedin.ingestion.DataHubIngestionSourceInfo; import com.linkedin.metadata.Constants; import com.linkedin.metadata.key.DataHubIngestionSourceKey; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; import com.linkedin.r2.RemoteInvocationException; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import java.util.HashSet; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -42,14 +44,13 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.INGESTION_SOURCE_ENTITY_NAME), Mockito.eq(""), Mockito.any(), Mockito.any(), Mockito.eq(0), - Mockito.eq(20), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true)))) + Mockito.eq(20))) .thenReturn( new SearchResult() .setFrom(0) @@ -85,6 +86,7 @@ public void testGetSuccess() throws Exception { // Execute resolver QueryContext mockContext = getMockAllowContext(); + Mockito.when(mockContext.getOperationContext()).thenReturn(mock(OperationContext.class)); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); Mockito.when(mockEnv.getContext()).thenReturn(mockContext); @@ -116,13 +118,12 @@ public void testGetUnauthorized() throws Exception { Mockito.any(), Mockito.anySet(), Mockito.anySet(), Mockito.any(Authentication.class)); Mockito.verify(mockClient, Mockito.times(0)) .search( + any(), Mockito.any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + Mockito.anyInt()); } @Test diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ownership/ListOwnershipTypesResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ownership/ListOwnershipTypesResolverTest.java index fd7baf6af74691..d18bc3aa31f898 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ownership/ListOwnershipTypesResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ownership/ListOwnershipTypesResolverTest.java @@ -1,6 +1,7 @@ package com.linkedin.datahub.graphql.resolvers.ownership; import static com.linkedin.datahub.graphql.TestUtils.*; +import static org.mockito.ArgumentMatchers.any; import static org.testng.Assert.*; import com.datahub.authentication.Authentication; @@ -11,7 +12,6 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; import com.linkedin.metadata.key.OwnershipTypeKey; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; @@ -40,14 +40,13 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.OWNERSHIP_TYPE_ENTITY_NAME), Mockito.eq(""), Mockito.eq(null), Mockito.any(), Mockito.eq(0), - Mockito.eq(20), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true)))) + Mockito.eq(20))) .thenReturn( new SearchResult() .setFrom(0) @@ -90,13 +89,12 @@ public void testGetUnauthorized() throws Exception { Mockito.any(), Mockito.anySet(), Mockito.anySet(), Mockito.any(Authentication.class)); Mockito.verify(mockClient, Mockito.times(0)) .search( + any(), Mockito.any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + Mockito.anyInt()); } @Test diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/post/ListPostsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/post/ListPostsResolverTest.java index 6c475cdc7f5a85..340e8cbf8514c1 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/post/ListPostsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/post/ListPostsResolverTest.java @@ -20,7 +20,6 @@ import com.linkedin.entity.EnvelopedAspect; import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.entity.client.EntityClient; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; @@ -33,7 +32,6 @@ import graphql.schema.DataFetchingEnvironment; import java.net.URISyntaxException; import java.util.Map; -import org.mockito.Mockito; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -120,14 +118,7 @@ public void testListPosts() throws Exception { ImmutableList.of(new SearchEntity().setEntity(Urn.createFromString(POST_URN_STRING))))); when(_entityClient.search( - eq(POST_ENTITY_NAME), - any(), - eq(null), - any(), - anyInt(), - anyInt(), - eq(_authentication), - Mockito.eq(new SearchFlags().setFulltext(true)))) + any(), eq(POST_ENTITY_NAME), any(), eq(null), any(), anyInt(), anyInt())) .thenReturn(roleSearchResult); when(_entityClient.batchGetV2(eq(POST_ENTITY_NAME), any(), any(), any())) .thenReturn(_entityResponseMap); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/ListQueriesResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/ListQueriesResolverTest.java index 8a56b142e5b5ef..9ed1d5001b75c3 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/ListQueriesResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/query/ListQueriesResolverTest.java @@ -1,9 +1,9 @@ package com.linkedin.datahub.graphql.resolvers.query; import static com.linkedin.datahub.graphql.TestUtils.*; +import static org.mockito.ArgumentMatchers.any; import static org.testng.Assert.*; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; @@ -17,7 +17,6 @@ import com.linkedin.datahub.graphql.resolvers.ResolverUtils; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; @@ -62,6 +61,7 @@ public void testGetSuccess(final ListQueriesInput input) throws Exception { Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.QUERY_ENTITY_NAME), Mockito.eq( input.getQuery() == null @@ -73,9 +73,7 @@ public void testGetSuccess(final ListQueriesInput input) throws Exception { .setField(ListQueriesResolver.CREATED_AT_FIELD) .setOrder(SortOrder.DESCENDING)), Mockito.eq(input.getStart()), - Mockito.eq(input.getCount()), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true).setSkipHighlighting(true)))) + Mockito.eq(input.getCount()))) .thenReturn( new SearchResult() .setFrom(0) @@ -116,13 +114,12 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockClient, Mockito.times(0)) .search( + any(), Mockito.any(), Mockito.eq("*"), Mockito.anyMap(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true).setSkipHighlighting(true))); + Mockito.anyInt()); } @Test @@ -132,13 +129,12 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RemoteInvocationException.class) .when(mockClient) .search( + any(), Mockito.any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true).setSkipHighlighting(true))); + Mockito.anyInt()); ListQueriesResolver resolver = new ListQueriesResolver(mockClient); // Execute resolver diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/CreateInviteTokenResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/CreateInviteTokenResolverTest.java index 9197d1b18c0c9c..e9d5ef00e74dd7 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/CreateInviteTokenResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/CreateInviteTokenResolverTest.java @@ -9,6 +9,7 @@ import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.CreateInviteTokenInput; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -42,7 +43,7 @@ public void testPasses() throws Exception { QueryContext mockContext = getMockAllowContext(); when(_dataFetchingEnvironment.getContext()).thenReturn(mockContext); when(mockContext.getAuthentication()).thenReturn(_authentication); - when(_inviteTokenService.getInviteToken(any(), eq(true), eq(_authentication))) + when(_inviteTokenService.getInviteToken(any(OperationContext.class), any(), eq(true))) .thenReturn(INVITE_TOKEN_STRING); CreateInviteTokenInput input = new CreateInviteTokenInput(); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/GetInviteTokenResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/GetInviteTokenResolverTest.java index 8e761454cb06c3..78d848e882b6bf 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/GetInviteTokenResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/GetInviteTokenResolverTest.java @@ -9,6 +9,7 @@ import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.GetInviteTokenInput; import graphql.schema.DataFetchingEnvironment; +import io.datahubproject.metadata.context.OperationContext; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -19,13 +20,15 @@ public class GetInviteTokenResolverTest { private GetInviteTokenResolver _resolver; private DataFetchingEnvironment _dataFetchingEnvironment; private Authentication _authentication; + private OperationContext opContext; @BeforeMethod public void setupTest() throws Exception { _inviteTokenService = mock(InviteTokenService.class); _dataFetchingEnvironment = mock(DataFetchingEnvironment.class); _authentication = mock(Authentication.class); - + opContext = mock(OperationContext.class); + when(opContext.getAuthentication()).thenReturn(_authentication); _resolver = new GetInviteTokenResolver(_inviteTokenService); } @@ -42,7 +45,7 @@ public void testPasses() throws Exception { QueryContext mockContext = getMockAllowContext(); when(_dataFetchingEnvironment.getContext()).thenReturn(mockContext); when(mockContext.getAuthentication()).thenReturn(_authentication); - when(_inviteTokenService.getInviteToken(any(), eq(false), eq(_authentication))) + when(_inviteTokenService.getInviteToken(any(OperationContext.class), any(), eq(false))) .thenReturn(INVITE_TOKEN_STRING); GetInviteTokenInput input = new GetInviteTokenInput(); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolverTest.java index d956295faa180f..ab2f852d83040a 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolverTest.java @@ -19,7 +19,6 @@ import com.linkedin.entity.EnvelopedAspect; import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.entity.client.EntityClient; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; @@ -27,7 +26,6 @@ import com.linkedin.policy.DataHubRoleInfo; import graphql.schema.DataFetchingEnvironment; import java.util.Map; -import org.mockito.Mockito; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -102,13 +100,7 @@ public void testListRoles() throws Exception { new SearchEntity().setEntity(Urn.createFromString(EDITOR_ROLE_URN_STRING))))); when(_entityClient.search( - eq(DATAHUB_ROLE_ENTITY_NAME), - any(), - any(), - anyInt(), - anyInt(), - any(), - Mockito.eq(new SearchFlags().setFulltext(true)))) + any(), eq(DATAHUB_ROLE_ENTITY_NAME), any(), any(), anyInt(), anyInt())) .thenReturn(roleSearchResult); when(_entityClient.batchGetV2(eq(DATAHUB_ROLE_ENTITY_NAME), any(), any(), any())) .thenReturn(_entityResponseMap); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/AggregateAcrossEntitiesResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/AggregateAcrossEntitiesResolverTest.java index 4d56cc3d52af8b..58fbadf7e0d7fd 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/AggregateAcrossEntitiesResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/AggregateAcrossEntitiesResolverTest.java @@ -2,6 +2,7 @@ import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.SEARCHABLE_ENTITY_TYPES; +import static org.mockito.ArgumentMatchers.any; import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; @@ -318,14 +319,14 @@ public static void testErrorFetchingResults() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when( mockClient.searchAcrossEntities( + any(), Mockito.anyList(), Mockito.anyString(), Mockito.any(), Mockito.anyInt(), Mockito.anyInt(), Mockito.eq(null), - Mockito.eq(null), - Mockito.any(Authentication.class))) + Mockito.eq(null))) .thenThrow(new RemoteInvocationException()); final AggregateAcrossEntitiesResolver resolver = @@ -392,14 +393,13 @@ private static EntityClient initMockEntityClient( EntityClient client = Mockito.mock(EntityClient.class); Mockito.when( client.searchAcrossEntities( + any(), Mockito.eq(entityTypes), Mockito.eq(query), Mockito.eq(filter), Mockito.eq(start), Mockito.eq(limit), Mockito.eq(null), - Mockito.eq(null), - Mockito.any(Authentication.class), Mockito.eq(facets))) .thenReturn(result); return client; @@ -416,14 +416,13 @@ private static void verifyMockEntityClient( throws Exception { Mockito.verify(mockClient, Mockito.times(1)) .searchAcrossEntities( + any(), Mockito.eq(entityTypes), Mockito.eq(query), Mockito.eq(filter), Mockito.eq(start), Mockito.eq(limit), Mockito.eq(null), - Mockito.eq(null), - Mockito.any(Authentication.class), Mockito.eq(facets)); } diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/AutoCompleteForMultipleResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/AutoCompleteForMultipleResolverTest.java index 3b69337acfbd0e..ea0765ba9377c4 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/AutoCompleteForMultipleResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/AutoCompleteForMultipleResolverTest.java @@ -1,6 +1,7 @@ package com.linkedin.datahub.graphql.resolvers.search; import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; +import static org.mockito.ArgumentMatchers.any; import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; @@ -192,11 +193,11 @@ public static void testAutoCompleteResolverWithViewEntityFilter() throws Excepti // types Mockito.verify(mockClient, Mockito.times(0)) .autoComplete( + any(), Mockito.eq(Constants.DATASET_ENTITY_NAME), Mockito.eq("test"), Mockito.eq(viewInfo.getDefinition().getFilter()), - Mockito.eq(10), - Mockito.any(Authentication.class)); + Mockito.eq(10)); } @Test @@ -225,11 +226,11 @@ private static EntityClient initMockEntityClient( EntityClient client = Mockito.mock(EntityClient.class); Mockito.when( client.autoComplete( + any(), Mockito.eq(entityName), Mockito.eq(query), Mockito.eq(filters), - Mockito.eq(limit), - Mockito.any(Authentication.class))) + Mockito.eq(limit))) .thenReturn(result); return client; } @@ -246,11 +247,11 @@ private static void verifyMockEntityClient( throws Exception { Mockito.verify(mockClient, Mockito.times(1)) .autoComplete( + any(), Mockito.eq(entityName), Mockito.eq(query), Mockito.eq(filters), - Mockito.eq(limit), - Mockito.any(Authentication.class)); + Mockito.eq(limit)); } private static DataHubViewInfo createViewInfo(StringArray entityNames) { diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/GetQuickFiltersResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/GetQuickFiltersResolverTest.java index f5accdfb02043a..1f038427c9aaa2 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/GetQuickFiltersResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/GetQuickFiltersResolverTest.java @@ -2,8 +2,8 @@ import static com.linkedin.datahub.graphql.TestUtils.*; import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.SEARCHABLE_ENTITY_TYPES; +import static org.mockito.ArgumentMatchers.any; -import com.datahub.authentication.Authentication; import com.linkedin.common.urn.UrnUtils; import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.GetQuickFiltersInput; @@ -108,14 +108,14 @@ public static void testGetQuickFiltersFailure() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when( mockClient.searchAcrossEntities( + any(), Mockito.anyList(), Mockito.anyString(), Mockito.any(), Mockito.anyInt(), Mockito.anyInt(), Mockito.eq(null), - Mockito.eq(null), - Mockito.any(Authentication.class))) + Mockito.eq(null))) .thenThrow(new RemoteInvocationException()); final GetQuickFiltersResolver resolver = new GetQuickFiltersResolver(mockClient, mockService); @@ -294,14 +294,14 @@ private static EntityClient initMockEntityClient( EntityClient client = Mockito.mock(EntityClient.class); Mockito.when( client.searchAcrossEntities( + any(), Mockito.eq(entityTypes), Mockito.eq(query), Mockito.eq(filter), Mockito.eq(start), Mockito.eq(limit), Mockito.eq(null), - Mockito.eq(null), - Mockito.any(Authentication.class))) + Mockito.eq(null))) .thenReturn(result); return client; } diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossEntitiesResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossEntitiesResolverTest.java index 0b8c1f1aeb83fa..1ef44bbed4cbcd 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossEntitiesResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossEntitiesResolverTest.java @@ -2,6 +2,7 @@ import static com.linkedin.datahub.graphql.TestUtils.*; import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.*; +import static org.mockito.ArgumentMatchers.any; import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; @@ -431,14 +432,14 @@ public static void testApplyViewErrorFetchingView() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when( mockClient.searchAcrossEntities( + any(), Mockito.anyList(), Mockito.anyString(), Mockito.any(), Mockito.anyInt(), Mockito.anyInt(), Mockito.eq(null), - Mockito.eq(null), - Mockito.any(Authentication.class))) + Mockito.eq(null))) .thenThrow(new RemoteInvocationException()); final SearchAcrossEntitiesResolver resolver = @@ -480,14 +481,13 @@ private static EntityClient initMockEntityClient( EntityClient client = Mockito.mock(EntityClient.class); Mockito.when( client.searchAcrossEntities( + any(), Mockito.eq(entityTypes), Mockito.eq(query), Mockito.eq(filter), Mockito.eq(start), Mockito.eq(limit), - Mockito.eq(null), - Mockito.eq(null), - Mockito.any(Authentication.class))) + Mockito.eq(null))) .thenReturn(result); return client; } @@ -502,14 +502,13 @@ private static void verifyMockEntityClient( throws Exception { Mockito.verify(mockClient, Mockito.times(1)) .searchAcrossEntities( + any(), Mockito.eq(entityTypes), Mockito.eq(query), Mockito.eq(filter), Mockito.eq(start), Mockito.eq(limit), - Mockito.eq(null), - Mockito.eq(null), - Mockito.any(Authentication.class)); + Mockito.eq(null)); } private static void verifyMockViewService(ViewService mockService, Urn viewUrn) { diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolverTest.java index a50591b7fc3991..b5b7e78aec4846 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolverTest.java @@ -16,7 +16,6 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.AggregationMetadataArray; import com.linkedin.metadata.search.LineageSearchEntity; import com.linkedin.metadata.search.LineageSearchEntityArray; @@ -106,6 +105,7 @@ public void testSearchAcrossLineage() throws Exception { lineageSearchResult.setEntities(new LineageSearchEntityArray(lineageSearchEntity)); when(_entityClient.searchAcrossLineage( + any(), eq(UrnUtils.getUrn(SOURCE_URN_STRING)), eq(com.linkedin.metadata.graph.LineageDirection.DOWNSTREAM), anyList(), @@ -116,9 +116,7 @@ public void testSearchAcrossLineage() throws Exception { eq(START), eq(COUNT), eq(START_TIMESTAMP_MILLIS), - eq(END_TIMESTAMP_MILLIS), - eq(new SearchFlags().setFulltext(true).setSkipHighlighting(true)), - eq(_authentication))) + eq(END_TIMESTAMP_MILLIS))) .thenReturn(lineageSearchResult); final SearchAcrossLineageResults results = _resolver.get(_dataFetchingEnvironment).join(); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolverTest.java index 9716799628a453..a5310a052f613c 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolverTest.java @@ -2,8 +2,8 @@ import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; import static com.linkedin.metadata.Constants.*; +import static org.mockito.ArgumentMatchers.any; -import com.datahub.authentication.Authentication; import com.linkedin.data.template.SetMode; import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.EntityType; @@ -65,7 +65,9 @@ public void testDefaultSearchFlags() throws Exception { .setSkipAggregates(false) .setSkipHighlighting(true) // empty/wildcard .setMaxAggValues(20) - .setSkipCache(false), + .setSkipCache(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), true)); } @@ -136,7 +138,9 @@ public void testNonWildCardSearchFlags() throws Exception { .setSkipAggregates(false) .setSkipHighlighting(false) // empty/wildcard .setMaxAggValues(20) - .setSkipCache(false), + .setSkipCache(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), true)); } @@ -144,14 +148,13 @@ private EntityClient initMockSearchEntityClient() throws Exception { EntityClient client = Mockito.mock(EntityClient.class); Mockito.when( client.search( + any(), Mockito.anyString(), Mockito.anyString(), Mockito.any(), Mockito.any(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.any())) + Mockito.anyInt())) .thenReturn( new SearchResult() .setEntities(new SearchEntityArray()) @@ -174,14 +177,13 @@ private void verifyMockSearchEntityClient( throws Exception { Mockito.verify(mockClient, Mockito.times(1)) .search( + any(), Mockito.eq(entityName), Mockito.eq(query), Mockito.eq(filter), Mockito.eq(sortCriterion), Mockito.eq(start), - Mockito.eq(limit), - Mockito.any(Authentication.class), - Mockito.eq(searchFlags)); + Mockito.eq(limit)); } private SearchResolverTest() {} diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/test/ListTestsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/test/ListTestsResolverTest.java index 6075425d09c050..5e3cd539cade76 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/test/ListTestsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/test/ListTestsResolverTest.java @@ -1,16 +1,15 @@ package com.linkedin.datahub.graphql.resolvers.test; import static com.linkedin.datahub.graphql.TestUtils.*; +import static org.mockito.ArgumentMatchers.any; import static org.testng.Assert.*; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.ListTestsInput; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; @@ -34,13 +33,12 @@ public void testGetSuccess() throws Exception { Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.TEST_ENTITY_NAME), Mockito.eq(""), Mockito.eq(Collections.emptyMap()), Mockito.eq(0), - Mockito.eq(20), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true)))) + Mockito.eq(20))) .thenReturn( new SearchResult() .setFrom(0) @@ -80,14 +78,7 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockClient, Mockito.times(0)) - .search( - Mockito.any(), - Mockito.eq(""), - Mockito.anyMap(), - Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + .search(any(), any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), Mockito.anyInt()); } @Test @@ -96,14 +87,7 @@ public void testGetEntityClientException() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.doThrow(RemoteInvocationException.class) .when(mockClient) - .search( - Mockito.any(), - Mockito.eq(""), - Mockito.anyMap(), - Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + .search(any(), any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), Mockito.anyInt()); ListTestsResolver resolver = new ListTestsResolver(mockClient); // Execute resolver diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/view/ListGlobalViewsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/view/ListGlobalViewsResolverTest.java index 8c30c17201bc65..a3b9e25e992259 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/view/ListGlobalViewsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/view/ListGlobalViewsResolverTest.java @@ -1,9 +1,9 @@ package com.linkedin.datahub.graphql.resolvers.view; import static com.linkedin.datahub.graphql.TestUtils.*; +import static org.mockito.ArgumentMatchers.any; import static org.testng.Assert.*; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; @@ -15,7 +15,6 @@ import com.linkedin.datahub.graphql.generated.ListViewsResult; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; @@ -44,6 +43,7 @@ public void testGetSuccessInput() throws Exception { Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.DATAHUB_VIEW_ENTITY_NAME), Mockito.eq(""), Mockito.eq( @@ -67,9 +67,7 @@ public void testGetSuccessInput() throws Exception { .setNegated(false)))))))), Mockito.any(), Mockito.eq(0), - Mockito.eq(20), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true)))) + Mockito.eq(20))) .thenReturn( new SearchResult() .setFrom(0) @@ -112,13 +110,12 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockClient, Mockito.times(0)) .search( + any(), Mockito.any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + Mockito.anyInt()); } @Test @@ -128,13 +125,12 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RemoteInvocationException.class) .when(mockClient) .search( + any(), Mockito.any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + Mockito.anyInt()); ListMyViewsResolver resolver = new ListMyViewsResolver(mockClient); // Execute resolver diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/view/ListMyViewsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/view/ListMyViewsResolverTest.java index 85e20cd656fcd3..99b0e76976748e 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/view/ListMyViewsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/view/ListMyViewsResolverTest.java @@ -1,9 +1,9 @@ package com.linkedin.datahub.graphql.resolvers.view; import static com.linkedin.datahub.graphql.TestUtils.*; +import static org.mockito.ArgumentMatchers.any; import static org.testng.Assert.*; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; @@ -14,7 +14,6 @@ import com.linkedin.datahub.graphql.generated.ListMyViewsInput; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; @@ -46,6 +45,7 @@ public void testGetSuccessInput1() throws Exception { Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.DATAHUB_VIEW_ENTITY_NAME), Mockito.eq(""), Mockito.eq( @@ -78,9 +78,7 @@ public void testGetSuccessInput1() throws Exception { .setNegated(false)))))))), Mockito.any(), Mockito.eq(0), - Mockito.eq(20), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true)))) + Mockito.eq(20))) .thenReturn( new SearchResult() .setFrom(0) @@ -113,6 +111,7 @@ public void testGetSuccessInput2() throws Exception { Mockito.when( mockClient.search( + any(), Mockito.eq(Constants.DATAHUB_VIEW_ENTITY_NAME), Mockito.eq(""), Mockito.eq( @@ -135,9 +134,7 @@ public void testGetSuccessInput2() throws Exception { .setNegated(false)))))))), Mockito.any(), Mockito.eq(0), - Mockito.eq(20), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true)))) + Mockito.eq(20))) .thenReturn( new SearchResult() .setFrom(0) @@ -178,13 +175,12 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockClient, Mockito.times(0)) .search( + any(), Mockito.any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + Mockito.anyInt()); } @Test @@ -194,13 +190,12 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RemoteInvocationException.class) .when(mockClient) .search( + any(), Mockito.any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), - Mockito.anyInt(), - Mockito.any(Authentication.class), - Mockito.eq(new SearchFlags().setFulltext(true))); + Mockito.anyInt()); ListMyViewsResolver resolver = new ListMyViewsResolver(mockClient); // Execute resolver diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/UpgradeCli.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/UpgradeCli.java index eee27096e22388..126df9187bc232 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/UpgradeCli.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/UpgradeCli.java @@ -7,12 +7,13 @@ import com.linkedin.datahub.upgrade.restorebackup.RestoreBackup; import com.linkedin.datahub.upgrade.restoreindices.RestoreIndices; import com.linkedin.datahub.upgrade.system.SystemUpdate; -import com.linkedin.datahub.upgrade.system.elasticsearch.BuildIndices; -import com.linkedin.datahub.upgrade.system.elasticsearch.CleanIndices; +import com.linkedin.datahub.upgrade.system.SystemUpdateBlocking; +import com.linkedin.datahub.upgrade.system.SystemUpdateNonBlocking; import java.util.List; import javax.inject.Inject; import javax.inject.Named; import lombok.extern.slf4j.Slf4j; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.CommandLineRunner; import org.springframework.stereotype.Component; import picocli.CommandLine; @@ -51,18 +52,18 @@ private static final class Args { @Named("removeUnknownAspects") private RemoveUnknownAspects removeUnknownAspects; - @Inject - @Named("buildIndices") - private BuildIndices buildIndices; - - @Inject - @Named("cleanIndices") - private CleanIndices cleanIndices; - - @Inject + @Autowired(required = false) @Named("systemUpdate") private SystemUpdate systemUpdate; + @Autowired(required = false) + @Named("systemUpdateBlocking") + private SystemUpdateBlocking systemUpdateBlocking; + + @Autowired(required = false) + @Named("systemUpdateNonBlocking") + private SystemUpdateNonBlocking systemUpdateNonBlocking; + @Override public void run(String... cmdLineArgs) { _upgradeManager.register(noCodeUpgrade); @@ -70,9 +71,15 @@ public void run(String... cmdLineArgs) { _upgradeManager.register(restoreIndices); _upgradeManager.register(restoreBackup); _upgradeManager.register(removeUnknownAspects); - _upgradeManager.register(buildIndices); - _upgradeManager.register(cleanIndices); - _upgradeManager.register(systemUpdate); + if (systemUpdate != null) { + _upgradeManager.register(systemUpdate); + } + if (systemUpdateBlocking != null) { + _upgradeManager.register(systemUpdateBlocking); + } + if (systemUpdateNonBlocking != null) { + _upgradeManager.register(systemUpdateNonBlocking); + } final Args args = new Args(); new CommandLine(args).setCaseInsensitiveEnumValuesAllowed(true).parseArgs(cmdLineArgs); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillBrowsePathsV2Config.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillBrowsePathsV2Config.java index 2b2f4648f76e73..a33722d7761cc4 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillBrowsePathsV2Config.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillBrowsePathsV2Config.java @@ -1,23 +1,28 @@ package com.linkedin.datahub.upgrade.config; -import com.linkedin.datahub.upgrade.system.entity.steps.BackfillBrowsePathsV2; +import com.linkedin.datahub.upgrade.system.NonBlockingSystemUpgrade; +import com.linkedin.datahub.upgrade.system.browsepaths.BackfillBrowsePathsV2; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.search.SearchService; +import io.datahubproject.metadata.context.OperationContext; import org.springframework.beans.factory.annotation.Value; import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Conditional; import org.springframework.context.annotation.Configuration; @Configuration +@Conditional(SystemUpdateCondition.NonBlockingSystemUpdateCondition.class) public class BackfillBrowsePathsV2Config { @Bean - public BackfillBrowsePathsV2 backfillBrowsePathsV2( + public NonBlockingSystemUpgrade backfillBrowsePathsV2( + final OperationContext opContext, EntityService entityService, SearchService searchService, @Value("${systemUpdate.browsePathsV2.enabled}") final boolean enabled, @Value("${systemUpdate.browsePathsV2.reprocess.enabled}") final boolean reprocessEnabled, @Value("${systemUpdate.browsePathsV2.batchSize}") final Integer batchSize) { return new BackfillBrowsePathsV2( - entityService, searchService, enabled, reprocessEnabled, batchSize); + opContext, entityService, searchService, enabled, reprocessEnabled, batchSize); } } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillOwnershipTypesConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillOwnershipTypesConfig.java new file mode 100644 index 00000000000000..3ca397a8ce268c --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillOwnershipTypesConfig.java @@ -0,0 +1,28 @@ +package com.linkedin.datahub.upgrade.config; + +import com.linkedin.datahub.upgrade.system.NonBlockingSystemUpgrade; +import com.linkedin.datahub.upgrade.system.ownershiptypes.OwnershipTypes; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.search.SearchService; +import io.datahubproject.metadata.context.OperationContext; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Conditional; +import org.springframework.context.annotation.Configuration; + +@Configuration +@Conditional(SystemUpdateCondition.NonBlockingSystemUpdateCondition.class) +public class BackfillOwnershipTypesConfig { + + @Bean + public NonBlockingSystemUpgrade backfillOwnershipTypes( + final OperationContext opContext, + final EntityService entityService, + final SearchService searchService, + @Value("${systemUpdate.ownershipTypes.enabled}") final boolean enabled, + @Value("${systemUpdate.ownershipTypes.reprocess.enabled}") final boolean reprocessEnabled, + @Value("${systemUpdate.ownershipTypes.batchSize}") final Integer batchSize) { + return new OwnershipTypes( + opContext, entityService, searchService, enabled, reprocessEnabled, batchSize); + } +} diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillPolicyFieldsConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillPolicyFieldsConfig.java index 6da85a5c16979f..7226ec267dbbc5 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillPolicyFieldsConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BackfillPolicyFieldsConfig.java @@ -1,23 +1,27 @@ package com.linkedin.datahub.upgrade.config; -import com.linkedin.datahub.upgrade.system.entity.steps.BackfillPolicyFields; +import com.linkedin.datahub.upgrade.system.policyfields.BackfillPolicyFields; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.search.SearchService; +import io.datahubproject.metadata.context.OperationContext; import org.springframework.beans.factory.annotation.Value; import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Conditional; import org.springframework.context.annotation.Configuration; @Configuration +@Conditional(SystemUpdateCondition.NonBlockingSystemUpdateCondition.class) public class BackfillPolicyFieldsConfig { @Bean public BackfillPolicyFields backfillPolicyFields( + final OperationContext opContext, EntityService entityService, SearchService searchService, @Value("${systemUpdate.policyFields.enabled}") final boolean enabled, @Value("${systemUpdate.policyFields.reprocess.enabled}") final boolean reprocessEnabled, @Value("${systemUpdate.policyFields.batchSize}") final Integer batchSize) { return new BackfillPolicyFields( - entityService, searchService, enabled, reprocessEnabled, batchSize); + opContext, entityService, searchService, enabled, reprocessEnabled, batchSize); } } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BuildIndicesConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BuildIndicesConfig.java index caa45988733df0..3510fa513b3b9c 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BuildIndicesConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/BuildIndicesConfig.java @@ -1,5 +1,6 @@ package com.linkedin.datahub.upgrade.config; +import com.linkedin.datahub.upgrade.system.BlockingSystemUpgrade; import com.linkedin.datahub.upgrade.system.elasticsearch.BuildIndices; import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.gms.factory.search.BaseElasticSearchComponentsFactory; @@ -10,12 +11,14 @@ import com.linkedin.metadata.systemmetadata.SystemMetadataService; import com.linkedin.metadata.timeseries.TimeseriesAspectService; import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Conditional; import org.springframework.context.annotation.Configuration; @Configuration +@Conditional(SystemUpdateCondition.BlockingSystemUpdateCondition.class) public class BuildIndicesConfig { @Bean(name = "buildIndices") - public BuildIndices buildIndices( + public BlockingSystemUpgrade buildIndices( final SystemMetadataService systemMetadataService, final TimeseriesAspectService timeseriesAspectService, final EntitySearchService entitySearchService, diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/CleanIndicesConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/CleanIndicesConfig.java index 5bd7244a92e45a..4f54b01459625d 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/CleanIndicesConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/CleanIndicesConfig.java @@ -1,5 +1,6 @@ package com.linkedin.datahub.upgrade.config; +import com.linkedin.datahub.upgrade.system.NonBlockingSystemUpgrade; import com.linkedin.datahub.upgrade.system.elasticsearch.CleanIndices; import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.gms.factory.search.BaseElasticSearchComponentsFactory; @@ -8,12 +9,14 @@ import com.linkedin.metadata.systemmetadata.SystemMetadataService; import com.linkedin.metadata.timeseries.TimeseriesAspectService; import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Conditional; import org.springframework.context.annotation.Configuration; @Configuration +@Conditional(SystemUpdateCondition.NonBlockingSystemUpdateCondition.class) public class CleanIndicesConfig { - @Bean(name = "cleanIndices") - public CleanIndices cleanIndices( + @Bean + public NonBlockingSystemUpgrade cleanIndices( final SystemMetadataService systemMetadataService, final TimeseriesAspectService timeseriesAspectService, final EntitySearchService entitySearchService, diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/ReindexDataJobViaNodesCLLConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/ReindexDataJobViaNodesCLLConfig.java index 83dad80944f5f3..0281ff4f4169b5 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/ReindexDataJobViaNodesCLLConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/ReindexDataJobViaNodesCLLConfig.java @@ -1,16 +1,19 @@ package com.linkedin.datahub.upgrade.config; -import com.linkedin.datahub.upgrade.system.via.ReindexDataJobViaNodesCLL; +import com.linkedin.datahub.upgrade.system.NonBlockingSystemUpgrade; +import com.linkedin.datahub.upgrade.system.vianodes.ReindexDataJobViaNodesCLL; import com.linkedin.metadata.entity.EntityService; import org.springframework.beans.factory.annotation.Value; import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Conditional; import org.springframework.context.annotation.Configuration; @Configuration +@Conditional(SystemUpdateCondition.NonBlockingSystemUpdateCondition.class) public class ReindexDataJobViaNodesCLLConfig { @Bean - public ReindexDataJobViaNodesCLL _reindexDataJobViaNodesCLL( + public NonBlockingSystemUpgrade reindexDataJobViaNodesCLL( EntityService entityService, @Value("${systemUpdate.dataJobNodeCLL.enabled}") final boolean enabled, @Value("${systemUpdate.dataJobNodeCLL.batchSize}") final Integer batchSize) { diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/SystemUpdateCondition.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/SystemUpdateCondition.java index ea432dfa9f7df5..0d65af742a5925 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/SystemUpdateCondition.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/SystemUpdateCondition.java @@ -1,14 +1,48 @@ package com.linkedin.datahub.upgrade.config; +import java.util.Objects; +import java.util.Set; import org.springframework.boot.ApplicationArguments; import org.springframework.context.annotation.Condition; import org.springframework.context.annotation.ConditionContext; import org.springframework.core.type.AnnotatedTypeMetadata; public class SystemUpdateCondition implements Condition { + public static final String SYSTEM_UPDATE_ARG = "SystemUpdate"; + public static final String BLOCKING_SYSTEM_UPDATE_ARG = SYSTEM_UPDATE_ARG + "Blocking"; + public static final String NONBLOCKING_SYSTEM_UPDATE_ARG = SYSTEM_UPDATE_ARG + "NonBlocking"; + public static final Set SYSTEM_UPDATE_ARGS = + Set.of(SYSTEM_UPDATE_ARG, BLOCKING_SYSTEM_UPDATE_ARG, NONBLOCKING_SYSTEM_UPDATE_ARG); + @Override public boolean matches(ConditionContext context, AnnotatedTypeMetadata metadata) { return context.getBeanFactory().getBean(ApplicationArguments.class).getNonOptionArgs().stream() - .anyMatch("SystemUpdate"::equals); + .filter(Objects::nonNull) + .anyMatch(SYSTEM_UPDATE_ARGS::contains); + } + + public static class BlockingSystemUpdateCondition implements Condition { + @Override + public boolean matches(ConditionContext context, AnnotatedTypeMetadata metadata) { + return context + .getBeanFactory() + .getBean(ApplicationArguments.class) + .getNonOptionArgs() + .stream() + .anyMatch(arg -> SYSTEM_UPDATE_ARG.equals(arg) || BLOCKING_SYSTEM_UPDATE_ARG.equals(arg)); + } + } + + public static class NonBlockingSystemUpdateCondition implements Condition { + @Override + public boolean matches(ConditionContext context, AnnotatedTypeMetadata metadata) { + return context + .getBeanFactory() + .getBean(ApplicationArguments.class) + .getNonOptionArgs() + .stream() + .anyMatch( + arg -> SYSTEM_UPDATE_ARG.equals(arg) || NONBLOCKING_SYSTEM_UPDATE_ARG.equals(arg)); + } } } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/SystemUpdateConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/SystemUpdateConfig.java index 17ad56ec80bac5..bea38b616f86fb 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/SystemUpdateConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/SystemUpdateConfig.java @@ -1,11 +1,11 @@ package com.linkedin.datahub.upgrade.config; +import com.linkedin.datahub.upgrade.system.BlockingSystemUpgrade; +import com.linkedin.datahub.upgrade.system.NonBlockingSystemUpgrade; import com.linkedin.datahub.upgrade.system.SystemUpdate; -import com.linkedin.datahub.upgrade.system.elasticsearch.BuildIndices; -import com.linkedin.datahub.upgrade.system.elasticsearch.CleanIndices; -import com.linkedin.datahub.upgrade.system.entity.steps.BackfillBrowsePathsV2; -import com.linkedin.datahub.upgrade.system.entity.steps.BackfillPolicyFields; -import com.linkedin.datahub.upgrade.system.via.ReindexDataJobViaNodesCLL; +import com.linkedin.datahub.upgrade.system.SystemUpdateBlocking; +import com.linkedin.datahub.upgrade.system.SystemUpdateNonBlocking; +import com.linkedin.datahub.upgrade.system.elasticsearch.steps.DataHubStartupStep; import com.linkedin.gms.factory.common.TopicConventionFactory; import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.gms.factory.kafka.DataHubKafkaProducerFactory; @@ -14,8 +14,12 @@ import com.linkedin.metadata.config.kafka.KafkaConfiguration; import com.linkedin.metadata.dao.producer.KafkaEventProducer; import com.linkedin.metadata.dao.producer.KafkaHealthChecker; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.version.GitVersion; import com.linkedin.mxe.TopicConvention; +import java.util.List; +import javax.annotation.PostConstruct; import lombok.extern.slf4j.Slf4j; import org.apache.avro.generic.IndexedRecord; import org.apache.kafka.clients.producer.KafkaProducer; @@ -32,27 +36,28 @@ @Slf4j @Configuration +@Conditional(SystemUpdateCondition.class) public class SystemUpdateConfig { + @Bean(name = "systemUpdate") public SystemUpdate systemUpdate( - final BuildIndices buildIndices, - final CleanIndices cleanIndices, - @Qualifier("duheKafkaEventProducer") final KafkaEventProducer kafkaEventProducer, - final GitVersion gitVersion, - @Qualifier("revision") String revision, - final BackfillBrowsePathsV2 backfillBrowsePathsV2, - final ReindexDataJobViaNodesCLL reindexDataJobViaNodesCLL, - final BackfillPolicyFields backfillPolicyFields) { + final List blockingSystemUpgrades, + final List nonBlockingSystemUpgrades, + final DataHubStartupStep dataHubStartupStep) { + return new SystemUpdate(blockingSystemUpgrades, nonBlockingSystemUpgrades, dataHubStartupStep); + } - String version = String.format("%s-%s", gitVersion.getVersion(), revision); - return new SystemUpdate( - buildIndices, - cleanIndices, - kafkaEventProducer, - version, - backfillBrowsePathsV2, - reindexDataJobViaNodesCLL, - backfillPolicyFields); + @Bean(name = "systemUpdateBlocking") + public SystemUpdateBlocking systemUpdateBlocking( + final List blockingSystemUpgrades, + final DataHubStartupStep dataHubStartupStep) { + return new SystemUpdateBlocking(blockingSystemUpgrades, List.of(), dataHubStartupStep); + } + + @Bean(name = "systemUpdateNonBlocking") + public SystemUpdateNonBlocking systemUpdateNonBlocking( + final List nonBlockingSystemUpgrades) { + return new SystemUpdateNonBlocking(List.of(), nonBlockingSystemUpgrades, null); } @Value("#{systemEnvironment['DATAHUB_REVISION'] ?: '0'}") @@ -63,6 +68,15 @@ public String getRevision() { return revision; } + @Bean + public DataHubStartupStep dataHubStartupStep( + @Qualifier("duheKafkaEventProducer") final KafkaEventProducer kafkaEventProducer, + final GitVersion gitVersion, + @Qualifier("revision") String revision) { + return new DataHubStartupStep( + kafkaEventProducer, String.format("%s-%s", gitVersion.getVersion(), revision)); + } + @Autowired @Qualifier(TopicConventionFactory.TOPIC_CONVENTION_BEAN) private TopicConvention topicConvention; @@ -92,7 +106,6 @@ protected KafkaEventProducer duheKafkaEventProducer( */ @Primary @Bean(name = "kafkaEventProducer") - @Conditional(SystemUpdateCondition.class) @ConditionalOnProperty( name = "kafka.schemaRegistry.type", havingValue = InternalSchemaRegistryFactory.TYPE) @@ -100,4 +113,15 @@ protected KafkaEventProducer kafkaEventProducer( @Qualifier("duheKafkaEventProducer") KafkaEventProducer kafkaEventProducer) { return kafkaEventProducer; } + + @Configuration + public static class SystemUpdateSetup { + @Autowired private EntityService entityService; + @Autowired private EntitySearchService entitySearchService; + + @PostConstruct + protected void postConstruct() { + entitySearchService.postConstruct(entityService); + } + } } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/BlockingSystemUpgrade.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/BlockingSystemUpgrade.java new file mode 100644 index 00000000000000..4fae5b2239d11e --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/BlockingSystemUpgrade.java @@ -0,0 +1,5 @@ +package com.linkedin.datahub.upgrade.system; + +import com.linkedin.datahub.upgrade.Upgrade; + +public interface BlockingSystemUpgrade extends Upgrade {} diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/NonBlockingSystemUpgrade.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/NonBlockingSystemUpgrade.java new file mode 100644 index 00000000000000..fd83f1544a0982 --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/NonBlockingSystemUpgrade.java @@ -0,0 +1,5 @@ +package com.linkedin.datahub.upgrade.system; + +import com.linkedin.datahub.upgrade.Upgrade; + +public interface NonBlockingSystemUpgrade extends Upgrade {} diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdate.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdate.java index f02c820066c1cb..ad1c6c98fa3fd1 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdate.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdate.java @@ -3,58 +3,48 @@ import com.linkedin.datahub.upgrade.Upgrade; import com.linkedin.datahub.upgrade.UpgradeCleanupStep; import com.linkedin.datahub.upgrade.UpgradeStep; -import com.linkedin.datahub.upgrade.system.elasticsearch.BuildIndices; -import com.linkedin.datahub.upgrade.system.elasticsearch.CleanIndices; import com.linkedin.datahub.upgrade.system.elasticsearch.steps.DataHubStartupStep; -import com.linkedin.datahub.upgrade.system.entity.steps.BackfillBrowsePathsV2; -import com.linkedin.datahub.upgrade.system.entity.steps.BackfillPolicyFields; -import com.linkedin.datahub.upgrade.system.via.ReindexDataJobViaNodesCLL; -import com.linkedin.metadata.dao.producer.KafkaEventProducer; +import java.util.LinkedList; import java.util.List; -import java.util.stream.Collectors; -import java.util.stream.Stream; +import javax.annotation.Nullable; +import lombok.Getter; +import lombok.NonNull; +import lombok.experimental.Accessors; import lombok.extern.slf4j.Slf4j; +@Getter @Slf4j +@Accessors(fluent = true) public class SystemUpdate implements Upgrade { - private final List _preStartupUpgrades; - private final List _postStartupUpgrades; - private final List _steps; + private final List steps; + private final List cleanupSteps; public SystemUpdate( - final BuildIndices buildIndicesJob, - final CleanIndices cleanIndicesJob, - final KafkaEventProducer kafkaEventProducer, - final String version, - final BackfillBrowsePathsV2 backfillBrowsePathsV2, - final ReindexDataJobViaNodesCLL upgradeViaNodeCll, - final BackfillPolicyFields backfillPolicyFields) { - - _preStartupUpgrades = List.of(buildIndicesJob); - _steps = List.of(new DataHubStartupStep(kafkaEventProducer, version)); - _postStartupUpgrades = - List.of(cleanIndicesJob, backfillBrowsePathsV2, upgradeViaNodeCll, backfillPolicyFields); - } + @NonNull final List blockingSystemUpgrades, + @NonNull final List nonBlockingSystemUpgrades, + @Nullable final DataHubStartupStep dataHubStartupStep) { - @Override - public String id() { - return "SystemUpdate"; - } + steps = new LinkedList<>(); + cleanupSteps = new LinkedList<>(); - @Override - public List steps() { - return Stream.concat( - Stream.concat( - _preStartupUpgrades.stream().flatMap(up -> up.steps().stream()), _steps.stream()), - _postStartupUpgrades.stream().flatMap(up -> up.steps().stream())) - .collect(Collectors.toList()); + // blocking upgrades + steps.addAll(blockingSystemUpgrades.stream().flatMap(up -> up.steps().stream()).toList()); + cleanupSteps.addAll( + blockingSystemUpgrades.stream().flatMap(up -> up.cleanupSteps().stream()).toList()); + + // emit system update message if blocking upgrade(s) present + if (dataHubStartupStep != null && !blockingSystemUpgrades.isEmpty()) { + steps.add(dataHubStartupStep); + } + + // add non-blocking upgrades last + steps.addAll(nonBlockingSystemUpgrades.stream().flatMap(up -> up.steps().stream()).toList()); + cleanupSteps.addAll( + nonBlockingSystemUpgrades.stream().flatMap(up -> up.cleanupSteps().stream()).toList()); } @Override - public List cleanupSteps() { - return Stream.concat( - _preStartupUpgrades.stream().flatMap(up -> up.cleanupSteps().stream()), - _postStartupUpgrades.stream().flatMap(up -> up.cleanupSteps().stream())) - .collect(Collectors.toList()); + public String id() { + return getClass().getSimpleName(); } } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdateBlocking.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdateBlocking.java new file mode 100644 index 00000000000000..32841149c467b3 --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdateBlocking.java @@ -0,0 +1,16 @@ +package com.linkedin.datahub.upgrade.system; + +import com.linkedin.datahub.upgrade.system.elasticsearch.steps.DataHubStartupStep; +import java.util.List; +import lombok.NonNull; +import org.jetbrains.annotations.Nullable; + +public class SystemUpdateBlocking extends SystemUpdate { + + public SystemUpdateBlocking( + @NonNull List blockingSystemUpgrades, + @NonNull List nonBlockingSystemUpgrades, + @Nullable DataHubStartupStep dataHubStartupStep) { + super(blockingSystemUpgrades, nonBlockingSystemUpgrades, dataHubStartupStep); + } +} diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdateNonBlocking.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdateNonBlocking.java new file mode 100644 index 00000000000000..3309babc1f6cf2 --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/SystemUpdateNonBlocking.java @@ -0,0 +1,16 @@ +package com.linkedin.datahub.upgrade.system; + +import com.linkedin.datahub.upgrade.system.elasticsearch.steps.DataHubStartupStep; +import java.util.List; +import lombok.NonNull; +import org.jetbrains.annotations.Nullable; + +public class SystemUpdateNonBlocking extends SystemUpdate { + + public SystemUpdateNonBlocking( + @NonNull List blockingSystemUpgrades, + @NonNull List nonBlockingSystemUpgrades, + @Nullable DataHubStartupStep dataHubStartupStep) { + super(blockingSystemUpgrades, nonBlockingSystemUpgrades, dataHubStartupStep); + } +} diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillBrowsePathsV2.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/browsepaths/BackfillBrowsePathsV2.java similarity index 66% rename from datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillBrowsePathsV2.java rename to datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/browsepaths/BackfillBrowsePathsV2.java index 9b023e1e239a2d..16c039e2a64abd 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillBrowsePathsV2.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/browsepaths/BackfillBrowsePathsV2.java @@ -1,17 +1,19 @@ -package com.linkedin.datahub.upgrade.system.entity.steps; +package com.linkedin.datahub.upgrade.system.browsepaths; import com.google.common.collect.ImmutableList; -import com.linkedin.datahub.upgrade.Upgrade; import com.linkedin.datahub.upgrade.UpgradeStep; +import com.linkedin.datahub.upgrade.system.NonBlockingSystemUpgrade; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.search.SearchService; +import io.datahubproject.metadata.context.OperationContext; import java.util.List; -public class BackfillBrowsePathsV2 implements Upgrade { +public class BackfillBrowsePathsV2 implements NonBlockingSystemUpgrade { private final List _steps; public BackfillBrowsePathsV2( + OperationContext opContext, EntityService entityService, SearchService searchService, boolean enabled, @@ -21,7 +23,7 @@ public BackfillBrowsePathsV2( _steps = ImmutableList.of( new BackfillBrowsePathsV2Step( - entityService, searchService, reprocessEnabled, batchSize)); + opContext, entityService, searchService, reprocessEnabled, batchSize)); } else { _steps = ImmutableList.of(); } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillBrowsePathsV2Step.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/browsepaths/BackfillBrowsePathsV2Step.java similarity index 92% rename from datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillBrowsePathsV2Step.java rename to datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/browsepaths/BackfillBrowsePathsV2Step.java index 2d64e0052ae82a..30674ecc3d00eb 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillBrowsePathsV2Step.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/browsepaths/BackfillBrowsePathsV2Step.java @@ -1,4 +1,4 @@ -package com.linkedin.datahub.upgrade.system.entity.steps; +package com.linkedin.datahub.upgrade.system.browsepaths; import static com.linkedin.metadata.Constants.*; @@ -18,7 +18,6 @@ import com.linkedin.metadata.aspect.utils.DefaultAspectsUtil; import com.linkedin.metadata.boot.BootstrapStep; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; @@ -31,6 +30,7 @@ import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; +import io.datahubproject.metadata.context.OperationContext; import java.util.Set; import java.util.function.Function; import lombok.extern.slf4j.Slf4j; @@ -54,6 +54,7 @@ public class BackfillBrowsePathsV2Step implements UpgradeStep { Constants.ML_FEATURE_TABLE_ENTITY_NAME, Constants.ML_FEATURE_ENTITY_NAME); + private final OperationContext opContext; private final EntityService entityService; private final SearchService searchService; @@ -61,10 +62,12 @@ public class BackfillBrowsePathsV2Step implements UpgradeStep { private final Integer batchSize; public BackfillBrowsePathsV2Step( + OperationContext opContext, EntityService entityService, SearchService searchService, boolean reprocessEnabled, Integer batchSize) { + this.opContext = opContext; this.searchService = searchService; this.entityService = entityService; this.reprocessEnabled = reprocessEnabled; @@ -110,18 +113,20 @@ private String backfillBrowsePathsV2(String entityType, AuditStamp auditStamp, S final ScrollResult scrollResult = searchService.scrollAcrossEntities( + opContext.withSearchFlags( + flags -> + flags + .setFulltext(true) + .setSkipCache(true) + .setSkipHighlighting(true) + .setSkipAggregates(true)), ImmutableList.of(entityType), "*", filter, null, scrollId, null, - batchSize, - new SearchFlags() - .setFulltext(true) - .setSkipCache(true) - .setSkipHighlighting(true) - .setSkipAggregates(true)); + batchSize); if (scrollResult.getNumEntities() == 0 || scrollResult.getEntities().size() == 0) { return null; @@ -234,7 +239,8 @@ public boolean skip(UpgradeContext context) { return false; } - boolean previouslyRun = entityService.exists(UPGRADE_ID_URN, true); + boolean previouslyRun = + entityService.exists(UPGRADE_ID_URN, DATA_HUB_UPGRADE_RESULT_ASPECT_NAME, true); if (previouslyRun) { log.info("{} was already run. Skipping.", id()); } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/BuildIndices.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/BuildIndices.java index 970d67be337a64..fea0479876a2e9 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/BuildIndices.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/BuildIndices.java @@ -1,7 +1,7 @@ package com.linkedin.datahub.upgrade.system.elasticsearch; -import com.linkedin.datahub.upgrade.Upgrade; import com.linkedin.datahub.upgrade.UpgradeStep; +import com.linkedin.datahub.upgrade.system.BlockingSystemUpgrade; import com.linkedin.datahub.upgrade.system.elasticsearch.steps.BuildIndicesPostStep; import com.linkedin.datahub.upgrade.system.elasticsearch.steps.BuildIndicesPreStep; import com.linkedin.datahub.upgrade.system.elasticsearch.steps.BuildIndicesStep; @@ -19,7 +19,7 @@ import java.util.stream.Collectors; import java.util.stream.Stream; -public class BuildIndices implements Upgrade { +public class BuildIndices implements BlockingSystemUpgrade { private final List _steps; diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/CleanIndices.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/CleanIndices.java index ad68386622b216..e316481e2b07e6 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/CleanIndices.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/CleanIndices.java @@ -1,7 +1,7 @@ package com.linkedin.datahub.upgrade.system.elasticsearch; -import com.linkedin.datahub.upgrade.Upgrade; import com.linkedin.datahub.upgrade.UpgradeStep; +import com.linkedin.datahub.upgrade.system.NonBlockingSystemUpgrade; import com.linkedin.datahub.upgrade.system.elasticsearch.steps.CleanIndicesStep; import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.gms.factory.search.BaseElasticSearchComponentsFactory; @@ -16,7 +16,7 @@ import lombok.extern.slf4j.Slf4j; @Slf4j -public class CleanIndices implements Upgrade { +public class CleanIndices implements NonBlockingSystemUpgrade { private final List _steps; public CleanIndices( diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/ownershiptypes/OwnershipTypes.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/ownershiptypes/OwnershipTypes.java new file mode 100644 index 00000000000000..63aacde7ef8abf --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/ownershiptypes/OwnershipTypes.java @@ -0,0 +1,41 @@ +package com.linkedin.datahub.upgrade.system.ownershiptypes; + +import com.google.common.collect.ImmutableList; +import com.linkedin.datahub.upgrade.UpgradeStep; +import com.linkedin.datahub.upgrade.system.NonBlockingSystemUpgrade; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.search.SearchService; +import io.datahubproject.metadata.context.OperationContext; +import java.util.List; + +public class OwnershipTypes implements NonBlockingSystemUpgrade { + + private final List _steps; + + public OwnershipTypes( + OperationContext opContext, + EntityService entityService, + SearchService searchService, + boolean enabled, + boolean reprocessEnabled, + Integer batchSize) { + if (enabled) { + _steps = + ImmutableList.of( + new OwnershipTypesStep( + opContext, entityService, searchService, enabled, reprocessEnabled, batchSize)); + } else { + _steps = ImmutableList.of(); + } + } + + @Override + public String id() { + return getClass().getSimpleName(); + } + + @Override + public List steps() { + return _steps; + } +} diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/ownershiptypes/OwnershipTypesStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/ownershiptypes/OwnershipTypesStep.java new file mode 100644 index 00000000000000..4c55f4ddcb31db --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/ownershiptypes/OwnershipTypesStep.java @@ -0,0 +1,276 @@ +package com.linkedin.datahub.upgrade.system.ownershiptypes; + +import static com.linkedin.metadata.Constants.DATA_HUB_UPGRADE_RESULT_ASPECT_NAME; +import static com.linkedin.metadata.Constants.DEFAULT_RUN_ID; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.linkedin.common.AuditStamp; +import com.linkedin.common.urn.Urn; +import com.linkedin.datahub.upgrade.UpgradeContext; +import com.linkedin.datahub.upgrade.UpgradeStep; +import com.linkedin.datahub.upgrade.UpgradeStepResult; +import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult; +import com.linkedin.entity.Aspect; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.Constants; +import com.linkedin.metadata.aspect.batch.AspectsBatch; +import com.linkedin.metadata.boot.BootstrapStep; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.query.filter.Condition; +import com.linkedin.metadata.query.filter.ConjunctiveCriterion; +import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; +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.search.ScrollResult; +import com.linkedin.metadata.search.SearchEntity; +import com.linkedin.metadata.search.SearchEntityArray; +import com.linkedin.metadata.search.SearchService; +import com.linkedin.metadata.utils.AuditStampUtils; +import com.linkedin.metadata.utils.GenericRecordUtils; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.mxe.SystemMetadata; +import io.datahubproject.metadata.context.OperationContext; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Function; +import java.util.stream.Collectors; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +public class OwnershipTypesStep implements UpgradeStep { + + private static final String UPGRADE_ID = OwnershipTypes.class.getSimpleName(); + private static final Urn UPGRADE_ID_URN = BootstrapStep.getUpgradeUrn(UPGRADE_ID); + + private static final Set ENTITY_TYPES_TO_MIGRATE = + ImmutableSet.of( + Constants.DATASET_ENTITY_NAME, + Constants.DASHBOARD_ENTITY_NAME, + Constants.CHART_ENTITY_NAME, + Constants.DATA_JOB_ENTITY_NAME, + Constants.DATA_FLOW_ENTITY_NAME, + Constants.ML_MODEL_ENTITY_NAME, + Constants.ML_MODEL_GROUP_ENTITY_NAME, + Constants.ML_FEATURE_TABLE_ENTITY_NAME, + Constants.ML_FEATURE_ENTITY_NAME, + Constants.ML_PRIMARY_KEY_ENTITY_NAME, + Constants.GLOSSARY_TERM_ENTITY_NAME, + Constants.GLOSSARY_NODE_ENTITY_NAME, + Constants.TAG_ENTITY_NAME, + Constants.ROLE_ENTITY_NAME, + Constants.CORP_GROUP_ENTITY_NAME, + Constants.CORP_USER_ENTITY_NAME, + Constants.CONTAINER_ENTITY_NAME, + Constants.DOMAIN_ENTITY_NAME, + Constants.DATA_PRODUCT_ENTITY_NAME, + Constants.NOTEBOOK_ENTITY_NAME); + + private final OperationContext opContext; + private final EntityService entityService; + private final SearchService searchService; + private final boolean enabled; + private final boolean reprocessEnabled; + private final Integer batchSize; + + public OwnershipTypesStep( + OperationContext opContext, + EntityService entityService, + SearchService searchService, + boolean enabled, + boolean reprocessEnabled, + Integer batchSize) { + this.opContext = opContext; + this.entityService = entityService; + this.searchService = searchService; + this.enabled = enabled; + this.reprocessEnabled = reprocessEnabled; + this.batchSize = batchSize; + } + + @Override + public Function executable() { + return (context) -> { + final AuditStamp auditStamp = AuditStampUtils.createDefaultAuditStamp(); + + String scrollId = null; + for (String entityType : ENTITY_TYPES_TO_MIGRATE) { + int migratedCount = 0; + do { + log.info( + String.format( + "Upgrading batch %s-%s of browse paths for entity type %s", + migratedCount, migratedCount + batchSize, entityType)); + scrollId = ownershipTypes(entityType, auditStamp, scrollId); + migratedCount += batchSize; + } while (scrollId != null); + } + + BootstrapStep.setUpgradeResult(UPGRADE_ID_URN, entityService); + + return new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.SUCCEEDED); + }; + } + + private String ownershipTypes(String entityType, AuditStamp auditStamp, String scrollId) { + + final Filter filter; + + if (reprocessEnabled) { + filter = backfillDefaultOwnershipTypesFilter(); + } else { + filter = backfillOwnershipTypesFilter(); + } + + final ScrollResult scrollResult = + searchService.scrollAcrossEntities( + opContext.withSearchFlags( + flags -> + flags + .setFulltext(true) + .setSkipCache(true) + .setSkipHighlighting(true) + .setSkipAggregates(true)), + ImmutableList.of(entityType), + "*", + filter, + null, + scrollId, + null, + batchSize); + + if (scrollResult.getNumEntities() == 0 || scrollResult.getEntities().size() == 0) { + return null; + } + + try { + ingestOwnershipTypes(scrollResult.getEntities(), auditStamp); + } catch (Exception e) { + // don't stop the whole step because of one bad urn or one bad ingestion + log.error( + String.format( + "Error ingesting ownership aspect for urn %s", + scrollResult.getEntities().stream() + .map(SearchEntity::getEntity) + .collect(Collectors.toList())), + e); + } + + return scrollResult.getScrollId(); + } + + private Filter backfillOwnershipTypesFilter() { + // Condition: has `owners` AND does NOT have `ownershipTypes` + Criterion hasOwners = new Criterion(); + hasOwners.setCondition(Condition.EXISTS); + hasOwners.setField("owners"); + // Excludes entities with ownershipTypes + Criterion missingOwnershipTypes = new Criterion(); + missingOwnershipTypes.setCondition(Condition.IS_NULL); + missingOwnershipTypes.setField("ownershipTypes"); + + CriterionArray criterionArray = new CriterionArray(); + criterionArray.add(hasOwners); + criterionArray.add(missingOwnershipTypes); + + ConjunctiveCriterion conjunctiveCriterion = new ConjunctiveCriterion(); + conjunctiveCriterion.setAnd(criterionArray); + + ConjunctiveCriterionArray conjunctiveCriterionArray = new ConjunctiveCriterionArray(); + conjunctiveCriterionArray.add(conjunctiveCriterion); + + Filter filter = new Filter(); + filter.setOr(conjunctiveCriterionArray); + return filter; + } + + private Filter backfillDefaultOwnershipTypesFilter() { + // Condition: has `owners` + Criterion hasOwners = new Criterion(); + hasOwners.setCondition(Condition.EXISTS); + hasOwners.setField("owners"); + + CriterionArray criterionArray = new CriterionArray(); + criterionArray.add(hasOwners); + + ConjunctiveCriterion conjunctiveCriterion = new ConjunctiveCriterion(); + conjunctiveCriterion.setAnd(criterionArray); + + ConjunctiveCriterionArray conjunctiveCriterionArray = new ConjunctiveCriterionArray(); + conjunctiveCriterionArray.add(conjunctiveCriterion); + + Filter filter = new Filter(); + filter.setOr(conjunctiveCriterionArray); + return filter; + } + + private void ingestOwnershipTypes(SearchEntityArray searchBatch, AuditStamp auditStamp) + throws Exception { + Map> existing = + entityService.getLatestAspectObjects( + searchBatch.stream().map(SearchEntity::getEntity).collect(Collectors.toSet()), + Set.of(Constants.OWNERSHIP_ASPECT_NAME)); + + List mcps = + existing.entrySet().stream() + .filter(result -> result.getValue().containsKey(Constants.OWNERSHIP_ASPECT_NAME)) + .map( + result -> { + MetadataChangeProposal proposal = new MetadataChangeProposal(); + proposal.setEntityUrn(result.getKey()); + proposal.setEntityType(result.getKey().getEntityType()); + proposal.setAspectName(Constants.OWNERSHIP_ASPECT_NAME); + proposal.setChangeType(ChangeType.UPSERT); + proposal.setSystemMetadata( + new SystemMetadata() + .setRunId(DEFAULT_RUN_ID) + .setLastObserved(System.currentTimeMillis())); + proposal.setAspect( + GenericRecordUtils.serializeAspect( + result.getValue().get(Constants.OWNERSHIP_ASPECT_NAME))); + return proposal; + }) + .collect(Collectors.toList()); + + log.debug(String.format("Reingesting ownership for %s urns", mcps.size())); + AspectsBatch batch = AspectsBatchImpl.builder().mcps(mcps, auditStamp, entityService).build(); + + entityService.ingestProposal(batch, false); + } + + @Override + public String id() { + return UPGRADE_ID; + } + + /** + * Returns whether the upgrade should proceed if the step fails after exceeding the maximum + * retries. + */ + @Override + public boolean isOptional() { + return true; + } + + @Override + /** + * Returns whether the upgrade should be skipped. Uses previous run history or the environment + * variables to determine whether to skip. + */ + public boolean skip(UpgradeContext context) { + if (reprocessEnabled && enabled) { + return false; + } + + boolean previouslyRun = + entityService.exists(UPGRADE_ID_URN, DATA_HUB_UPGRADE_RESULT_ASPECT_NAME, true); + + if (previouslyRun) { + log.info("{} was already run. Skipping.", id()); + } + return (previouslyRun || !enabled); + } +} diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillPolicyFields.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/policyfields/BackfillPolicyFields.java similarity index 66% rename from datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillPolicyFields.java rename to datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/policyfields/BackfillPolicyFields.java index 3e1d385b87e45e..ca568e91928951 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillPolicyFields.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/policyfields/BackfillPolicyFields.java @@ -1,16 +1,18 @@ -package com.linkedin.datahub.upgrade.system.entity.steps; +package com.linkedin.datahub.upgrade.system.policyfields; import com.google.common.collect.ImmutableList; -import com.linkedin.datahub.upgrade.Upgrade; import com.linkedin.datahub.upgrade.UpgradeStep; +import com.linkedin.datahub.upgrade.system.NonBlockingSystemUpgrade; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.search.SearchService; +import io.datahubproject.metadata.context.OperationContext; import java.util.List; -public class BackfillPolicyFields implements Upgrade { +public class BackfillPolicyFields implements NonBlockingSystemUpgrade { private final List _steps; public BackfillPolicyFields( + OperationContext opContext, EntityService entityService, SearchService searchService, boolean enabled, @@ -20,7 +22,7 @@ public BackfillPolicyFields( _steps = ImmutableList.of( new BackfillPolicyFieldsStep( - entityService, searchService, reprocessEnabled, batchSize)); + opContext, entityService, searchService, reprocessEnabled, batchSize)); } else { _steps = ImmutableList.of(); } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillPolicyFieldsStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/policyfields/BackfillPolicyFieldsStep.java similarity index 92% rename from datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillPolicyFieldsStep.java rename to datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/policyfields/BackfillPolicyFieldsStep.java index 27d48aa5e05555..a9b8060f02c108 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillPolicyFieldsStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/policyfields/BackfillPolicyFieldsStep.java @@ -1,4 +1,4 @@ -package com.linkedin.datahub.upgrade.system.entity.steps; +package com.linkedin.datahub.upgrade.system.policyfields; import static com.linkedin.metadata.Constants.*; @@ -16,7 +16,6 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.boot.BootstrapStep; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; @@ -30,6 +29,7 @@ import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; import com.linkedin.policy.DataHubPolicyInfo; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.Collections; import java.util.function.Function; @@ -44,16 +44,20 @@ public class BackfillPolicyFieldsStep implements UpgradeStep { private static final String UPGRADE_ID = "BackfillPolicyFieldsStep"; private static final Urn UPGRADE_ID_URN = BootstrapStep.getUpgradeUrn(UPGRADE_ID); + + private final OperationContext opContext; private final boolean reprocessEnabled; private final Integer batchSize; private final EntityService entityService; private final SearchService _searchService; public BackfillPolicyFieldsStep( + OperationContext opContext, EntityService entityService, SearchService searchService, boolean reprocessEnabled, Integer batchSize) { + this.opContext = opContext; this.entityService = entityService; this._searchService = searchService; this.reprocessEnabled = reprocessEnabled; @@ -108,7 +112,8 @@ public boolean skip(UpgradeContext context) { return false; } - boolean previouslyRun = entityService.exists(UPGRADE_ID_URN, true); + boolean previouslyRun = + entityService.exists(UPGRADE_ID_URN, DATA_HUB_UPGRADE_RESULT_ASPECT_NAME, true); if (previouslyRun) { log.info("{} was already run. Skipping.", id()); } @@ -120,18 +125,20 @@ private String backfillPolicies(AuditStamp auditStamp, String scrollId) { final Filter filter = backfillPolicyFieldFilter(); final ScrollResult scrollResult = _searchService.scrollAcrossEntities( + opContext.withSearchFlags( + flags -> + flags + .setFulltext(true) + .setSkipCache(true) + .setSkipHighlighting(true) + .setSkipAggregates(true)), ImmutableList.of(Constants.POLICY_ENTITY_NAME), "*", filter, null, scrollId, null, - batchSize, - new SearchFlags() - .setFulltext(true) - .setSkipCache(true) - .setSkipHighlighting(true) - .setSkipAggregates(true)); + batchSize); if (scrollResult.getNumEntities() == 0 || scrollResult.getEntities().isEmpty()) { return null; diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/via/ReindexDataJobViaNodesCLL.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/vianodes/ReindexDataJobViaNodesCLL.java similarity index 80% rename from datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/via/ReindexDataJobViaNodesCLL.java rename to datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/vianodes/ReindexDataJobViaNodesCLL.java index 59975693322d10..c997aa15df9899 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/via/ReindexDataJobViaNodesCLL.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/vianodes/ReindexDataJobViaNodesCLL.java @@ -1,10 +1,8 @@ -package com.linkedin.datahub.upgrade.system.via; - -import static com.linkedin.metadata.Constants.*; +package com.linkedin.datahub.upgrade.system.vianodes; import com.google.common.collect.ImmutableList; -import com.linkedin.datahub.upgrade.Upgrade; import com.linkedin.datahub.upgrade.UpgradeStep; +import com.linkedin.datahub.upgrade.system.NonBlockingSystemUpgrade; import com.linkedin.metadata.entity.EntityService; import java.util.List; import lombok.extern.slf4j.Slf4j; @@ -14,7 +12,7 @@ * required to index column-level lineage correctly using via nodes. */ @Slf4j -public class ReindexDataJobViaNodesCLL implements Upgrade { +public class ReindexDataJobViaNodesCLL implements NonBlockingSystemUpgrade { private final List _steps; diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/via/ReindexDataJobViaNodesCLLStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/vianodes/ReindexDataJobViaNodesCLLStep.java similarity index 94% rename from datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/via/ReindexDataJobViaNodesCLLStep.java rename to datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/vianodes/ReindexDataJobViaNodesCLLStep.java index 56166caf5b57ea..a6a0331072a112 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/via/ReindexDataJobViaNodesCLLStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/vianodes/ReindexDataJobViaNodesCLLStep.java @@ -1,4 +1,4 @@ -package com.linkedin.datahub.upgrade.system.via; +package com.linkedin.datahub.upgrade.system.vianodes; import static com.linkedin.metadata.Constants.*; @@ -68,7 +68,8 @@ public boolean isOptional() { * variable SKIP_REINDEX_DATA_JOB_INPUT_OUTPUT to determine whether to skip. */ public boolean skip(UpgradeContext context) { - boolean previouslyRun = entityService.exists(UPGRADE_ID_URN, true); + boolean previouslyRun = + entityService.exists(UPGRADE_ID_URN, DATA_HUB_UPGRADE_RESULT_ASPECT_NAME, true); boolean envFlagRecommendsSkip = Boolean.parseBoolean(System.getenv("SKIP_REINDEX_DATA_JOB_INPUT_OUTPUT")); if (previouslyRun) { diff --git a/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeCliApplicationTest.java b/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeCliApplicationTest.java index 3e655be900bf28..dc4c3073ee351c 100644 --- a/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeCliApplicationTest.java +++ b/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeCliApplicationTest.java @@ -3,7 +3,7 @@ import static org.testng.AssertJUnit.*; import com.linkedin.datahub.upgrade.restoreindices.RestoreIndices; -import com.linkedin.datahub.upgrade.system.elasticsearch.BuildIndices; +import com.linkedin.datahub.upgrade.system.BlockingSystemUpgrade; import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder; import javax.inject.Named; import org.springframework.beans.factory.annotation.Autowired; @@ -14,6 +14,7 @@ @ActiveProfiles("test") @SpringBootTest( + args = {"-u", "SystemUpdate"}, classes = {UpgradeCliApplication.class, UpgradeCliApplicationTestConfiguration.class}) public class UpgradeCliApplicationTest extends AbstractTestNGSpringContextTests { @@ -23,7 +24,7 @@ public class UpgradeCliApplicationTest extends AbstractTestNGSpringContextTests @Autowired @Named("buildIndices") - private BuildIndices buildIndices; + private BlockingSystemUpgrade buildIndices; @Autowired private ESIndexBuilder esIndexBuilder; diff --git a/datahub-web-react/src/app/buildEntityRegistry.ts b/datahub-web-react/src/app/buildEntityRegistry.ts index 4f746815708029..fdd8e0afa77d1d 100644 --- a/datahub-web-react/src/app/buildEntityRegistry.ts +++ b/datahub-web-react/src/app/buildEntityRegistry.ts @@ -20,6 +20,7 @@ import { DataPlatformEntity } from './entity/dataPlatform/DataPlatformEntity'; import { DataProductEntity } from './entity/dataProduct/DataProductEntity'; import { DataPlatformInstanceEntity } from './entity/dataPlatformInstance/DataPlatformInstanceEntity'; import { RoleEntity } from './entity/Access/RoleEntity'; +import { RestrictedEntity } from './entity/restricted/RestrictedEntity'; export default function buildEntityRegistry() { const registry = new EntityRegistry(); @@ -44,5 +45,6 @@ export default function buildEntityRegistry() { registry.register(new DataPlatformEntity()); registry.register(new DataProductEntity()); registry.register(new DataPlatformInstanceEntity()); + registry.register(new RestrictedEntity()); return registry; } \ No newline at end of file diff --git a/datahub-web-react/src/app/entity/restricted/RestrictedEntity.tsx b/datahub-web-react/src/app/entity/restricted/RestrictedEntity.tsx new file mode 100644 index 00000000000000..482709c110d6b9 --- /dev/null +++ b/datahub-web-react/src/app/entity/restricted/RestrictedEntity.tsx @@ -0,0 +1,90 @@ +import React from 'react'; +import { QuestionOutlined } from '@ant-design/icons'; +import { EntityType, Restricted, SearchResult } from '../../../types.generated'; +import { Entity, IconStyleType, PreviewType } from '../Entity'; +import { getDataForEntityType } from '../shared/containers/profile/utils'; +import RestrictedIcon from '../../../images/restricted.svg'; +import { RestrictedEntityProfile } from './RestrictedEntityProfile'; + +/** + * Definition of the DataHub Data Product entity. + */ +export class RestrictedEntity implements Entity { + type: EntityType = EntityType.Restricted; + + icon = (fontSize: number, styleType: IconStyleType, color?: string) => { + if (styleType === IconStyleType.TAB_VIEW) { + return ; + } + + if (styleType === IconStyleType.HIGHLIGHT) { + return ; + } + + return ( + + ); + }; + + isSearchEnabled = () => false; + + isBrowseEnabled = () => false; + + isLineageEnabled = () => true; + + getAutoCompleteFieldName = () => 'name'; + + getPathName = () => 'restricted'; + + getEntityName = () => 'Restricted'; + + getCollectionName = () => 'Restricted Assets'; + + renderProfile = (_: string) => ; + + renderPreview = (_: PreviewType, _data: Restricted) => { + return ; + }; + + renderSearch = (_result: SearchResult) => { + return ; + }; + + getLineageVizConfig = (entity: Restricted) => { + return { + urn: entity?.urn, + name: 'Restricted Asset', + type: EntityType.Restricted, + icon: RestrictedIcon, + }; + }; + + displayName = (_data: Restricted) => { + return 'Restricted Asset'; + }; + + getOverridePropertiesFromEntity = (_data: Restricted) => { + return {}; + }; + + getGenericEntityProperties = (data: Restricted) => { + return getDataForEntityType({ + data, + entityType: this.type, + getOverrideProperties: this.getOverridePropertiesFromEntity, + }); + }; + + supportedCapabilities = () => { + return new Set([]); + }; + + getGraphName = () => { + return 'restricted'; + }; +} diff --git a/datahub-web-react/src/app/entity/restricted/RestrictedEntityProfile.tsx b/datahub-web-react/src/app/entity/restricted/RestrictedEntityProfile.tsx new file mode 100644 index 00000000000000..cf8aa5935e42f9 --- /dev/null +++ b/datahub-web-react/src/app/entity/restricted/RestrictedEntityProfile.tsx @@ -0,0 +1,30 @@ +import React from 'react'; +import styled from 'styled-components'; +import { + LogoIcon, + PlatformContentWrapper, + PlatformText, + PreviewImage, +} from '../shared/containers/profile/header/PlatformContent/PlatformContentView'; +import RestrictedIcon from '../../../images/restricted.svg'; +import { EntityTitle } from '../shared/containers/profile/header/EntityName'; + +const SubHeader = styled.div` + margin-top: 8px; + font-size: 14px; +`; + +export function RestrictedEntityProfile() { + return ( + <> + + + + + Restricted + + Restricted Asset + This asset is Restricted. Please request access to see more. + + ); +} diff --git a/datahub-web-react/src/app/entity/shared/containers/profile/header/EntityName.tsx b/datahub-web-react/src/app/entity/shared/containers/profile/header/EntityName.tsx index 762bd5f9111a0e..702f780f1aa111 100644 --- a/datahub-web-react/src/app/entity/shared/containers/profile/header/EntityName.tsx +++ b/datahub-web-react/src/app/entity/shared/containers/profile/header/EntityName.tsx @@ -7,7 +7,7 @@ import { useEntityRegistry } from '../../../../../useEntityRegistry'; import { useEntityData, useRefetch } from '../../../EntityContext'; import { useGlossaryEntityData } from '../../../GlossaryEntityContext'; -const EntityTitle = styled(Typography.Title)` +export const EntityTitle = styled(Typography.Title)` margin-right: 10px; &&& { diff --git a/datahub-web-react/src/app/entity/shared/containers/profile/header/PlatformContent/PlatformContentView.tsx b/datahub-web-react/src/app/entity/shared/containers/profile/header/PlatformContent/PlatformContentView.tsx index 1090dac501d0b2..52cb656f54eb99 100644 --- a/datahub-web-react/src/app/entity/shared/containers/profile/header/PlatformContent/PlatformContentView.tsx +++ b/datahub-web-react/src/app/entity/shared/containers/profile/header/PlatformContent/PlatformContentView.tsx @@ -13,20 +13,20 @@ import { } from './ParentNodesView'; import ParentEntities from '../../../../../../search/filters/ParentEntities'; -const LogoIcon = styled.span` +export const LogoIcon = styled.span` display: flex; gap: 4px; margin-right: 8px; `; -const PreviewImage = styled(Image)` +export const PreviewImage = styled(Image)` max-height: 17px; width: auto; object-fit: contain; background-color: transparent; `; -const PlatformContentWrapper = styled.div` +export const PlatformContentWrapper = styled.div` display: flex; align-items: center; margin: 0 8px 6px 0; @@ -34,7 +34,7 @@ const PlatformContentWrapper = styled.div` flex: 1; `; -const PlatformText = styled(Typography.Text)` +export const PlatformText = styled(Typography.Text)` font-size: 12px; line-height: 20px; font-weight: 700; diff --git a/datahub-web-react/src/app/lineage/LineageEntityNode.tsx b/datahub-web-react/src/app/lineage/LineageEntityNode.tsx index d9f7d8c2c69de4..fcf2198e7e0047 100644 --- a/datahub-web-react/src/app/lineage/LineageEntityNode.tsx +++ b/datahub-web-react/src/app/lineage/LineageEntityNode.tsx @@ -18,6 +18,7 @@ import { convertInputFieldsToSchemaFields } from './utils/columnLineageUtils'; import ManageLineageMenu from './manage/ManageLineageMenu'; import { useGetLineageTimeParams } from './utils/useGetLineageTimeParams'; import { EntityHealth } from '../entity/shared/containers/profile/header/EntityHealth'; +import { EntityType } from '../../types.generated'; const CLICK_DELAY_THRESHOLD = 1000; const DRAG_DISTANCE_THRESHOLD = 20; @@ -71,6 +72,7 @@ export default function LineageEntityNode({ const [getAsyncEntityLineage, { data: asyncLineageData, loading }] = useGetEntityLineageLazyQuery(); const isHideSiblingMode = useIsSeparateSiblingsMode(); const areColumnsCollapsed = !!collapsedColumnsNodes[node?.data?.urn || 'noop']; + const isRestricted = node.data.type === EntityType.Restricted; function fetchEntityLineage() { if (node.data.urn) { @@ -92,6 +94,12 @@ export default function LineageEntityNode({ } } + const centerEntity = () => { + if (!isRestricted) { + onEntityCenter({ urn: node.data.urn, type: node.data.type }); + } + }; + useEffect(() => { if (asyncLineageData && asyncLineageData.entity && !hasExpanded && !loading) { const entityAndType = { @@ -231,7 +239,7 @@ export default function LineageEntityNode({ ))} onEntityCenter({ urn: node.data.urn, type: node.data.type })} + onDoubleClick={centerEntity} onClick={(event) => { if ( event.timeStamp < lastMouseDownCoordinates.ts + CLICK_DELAY_THRESHOLD && @@ -311,25 +319,27 @@ export default function LineageEntityNode({ {entityRegistry.getIcon(node.data.type, 16, IconStyleType.SVG)} )} - e.stopPropagation()} - > - onEntityCenter({ urn: node.data.urn, type: node.data.type })} - entityType={node.data.type} - entityPlatform={node.data.platform?.name} - canEditLineage={node.data.canEditLineage} - /> - + {!isRestricted && ( + e.stopPropagation()} + > + onEntityCenter({ urn: node.data.urn, type: node.data.type })} + entityType={node.data.type} + entityPlatform={node.data.platform?.name} + canEditLineage={node.data.canEditLineage} + /> + + )} - {getShortenedTitle(platformDisplayText || '', width)} - - {' '} - |{' '} - + {platformDisplayText && ( + <> + {getShortenedTitle(platformDisplayText || '', width)} + + {' '} + |{' '} + + + )} {entityName} diff --git a/datahub-web-react/src/app/lineage/LineageExplorer.tsx b/datahub-web-react/src/app/lineage/LineageExplorer.tsx index a03f62f93abeb0..26ffaa26a6ca22 100644 --- a/datahub-web-react/src/app/lineage/LineageExplorer.tsx +++ b/datahub-web-react/src/app/lineage/LineageExplorer.tsx @@ -220,9 +220,11 @@ export default function LineageExplorer({ urn, type }: Props) { - + {selectedEntity.type !== EntityType.Restricted && ( + + )} ) } diff --git a/datahub-web-react/src/graphql/lineage.graphql b/datahub-web-react/src/graphql/lineage.graphql index 8d84168d4c67e6..b73a99488cf8af 100644 --- a/datahub-web-react/src/graphql/lineage.graphql +++ b/datahub-web-react/src/graphql/lineage.graphql @@ -269,6 +269,10 @@ fragment lineageNodeProperties on EntityWithRelationships { ... on MLPrimaryKey { ...nonRecursiveMLPrimaryKey } + ... on Restricted { + urn + type + } } fragment lineageFields on EntityWithRelationships { diff --git a/datahub-web-react/src/images/restricted.svg b/datahub-web-react/src/images/restricted.svg new file mode 100644 index 00000000000000..7537f7fb83fb46 --- /dev/null +++ b/datahub-web-react/src/images/restricted.svg @@ -0,0 +1,2 @@ + + \ No newline at end of file diff --git a/docker/profiles/docker-compose.prerequisites.yml b/docker/profiles/docker-compose.prerequisites.yml index 142c983182054e..d58b7493d99a0d 100644 --- a/docker/profiles/docker-compose.prerequisites.yml +++ b/docker/profiles/docker-compose.prerequisites.yml @@ -290,7 +290,7 @@ services: memory: 1G healthcheck: test: curl -sS --fail http://search:$${DATAHUB_ELASTIC_PORT:-9200}/_cluster/health?wait_for_status=yellow&timeout=0s - start_period: 20s + start_period: 30s interval: 1s retries: 3 timeout: 5s @@ -325,7 +325,7 @@ services: memory: 1G healthcheck: test: curl -sS --fail http://search:$${DATAHUB_ELASTIC_PORT:-9200}/_cluster/health?wait_for_status=yellow&timeout=0s - start_period: 20s + start_period: 30s interval: 1s retries: 3 timeout: 5s diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/ChangeMCP.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/ChangeMCP.java index 94e8bbab3ceeb6..19896e2b035443 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/ChangeMCP.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/ChangeMCP.java @@ -1,6 +1,7 @@ package com.linkedin.metadata.aspect.batch; import com.linkedin.data.DataMap; +import com.linkedin.data.template.RecordTemplate; import com.linkedin.metadata.aspect.SystemAspect; import java.lang.reflect.InvocationTargetException; import javax.annotation.Nonnull; @@ -23,6 +24,14 @@ public interface ChangeMCP extends MCPItem { void setNextAspectVersion(long nextAspectVersion); + @Nullable + default RecordTemplate getPreviousRecordTemplate() { + if (getPreviousSystemAspect() != null) { + return getPreviousSystemAspect().getRecordTemplate(); + } + return null; + } + default T getPreviousAspect(Class clazz) { if (getPreviousSystemAspect() != null) { try { @@ -35,8 +44,7 @@ default T getPreviousAspect(Class clazz) { | NoSuchMethodException e) { throw new RuntimeException(e); } - } else { - return null; } + return null; } } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/hooks/OwnerTypeMap.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/hooks/OwnerTypeMap.java new file mode 100644 index 00000000000000..45e92801993300 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/hooks/OwnerTypeMap.java @@ -0,0 +1,111 @@ +package com.linkedin.metadata.aspect.hooks; + +import static com.linkedin.metadata.Constants.DEFAULT_OWNERSHIP_TYPE_URN; +import static com.linkedin.metadata.Constants.OWNERSHIP_ASPECT_NAME; + +import com.linkedin.common.Owner; +import com.linkedin.common.Ownership; +import com.linkedin.common.UrnArray; +import com.linkedin.common.UrnArrayMap; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.ChangeMCP; +import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.aspect.plugins.hooks.MutationHook; +import com.linkedin.util.Pair; +import java.util.Collection; +import java.util.Collections; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +/** Hook to populate the ownerType map within the ownership aspect */ +public class OwnerTypeMap extends MutationHook { + public OwnerTypeMap(AspectPluginConfig aspectPluginConfig) { + super(aspectPluginConfig); + } + + @Override + protected Stream> writeMutation( + @Nonnull Collection changeMCPS, @Nonnull AspectRetriever aspectRetriever) { + + List> results = new LinkedList<>(); + + for (ChangeMCP item : changeMCPS) { + if (OWNERSHIP_ASPECT_NAME.equals(item.getAspectName()) && item.getRecordTemplate() != null) { + final Map> oldTypeOwner = + groupByOwnerType(item.getPreviousRecordTemplate()); + final Map> newTypeOwner = groupByOwnerType(item.getRecordTemplate()); + + Set removedTypes = + oldTypeOwner.keySet().stream() + .filter(typeUrn -> !newTypeOwner.containsKey(typeUrn)) + .collect(Collectors.toSet()); + + Set updatedTypes = newTypeOwner.keySet(); + + Map typeOwners = + Stream.concat(removedTypes.stream(), updatedTypes.stream()) + .map( + typeUrn -> { + final String typeFieldName = encodeFieldName(typeUrn.toString()); + if (removedTypes.contains(typeUrn)) { + // removed + return Pair.of(typeFieldName, new UrnArray()); + } + // updated + return Pair.of( + typeFieldName, + new UrnArray( + newTypeOwner.getOrDefault(typeUrn, Collections.emptySet()).stream() + .map(Owner::getOwner) + .collect(Collectors.toSet()))); + }) + .collect(Collectors.toMap(Pair::getFirst, Pair::getSecond)); + + if (!typeOwners.isEmpty()) { + item.getAspect(Ownership.class).setOwnerTypes(new UrnArrayMap(typeOwners)); + results.add(Pair.of(item, true)); + continue; + } + } + + // no op + results.add(Pair.of(item, false)); + } + + return results.stream(); + } + + private static Map> groupByOwnerType( + @Nullable RecordTemplate ownershipRecordTemplate) { + if (ownershipRecordTemplate != null) { + Ownership ownership = new Ownership(ownershipRecordTemplate.data()); + if (!ownership.getOwners().isEmpty()) { + return ownership.getOwners().stream() + .collect( + Collectors.groupingBy( + owner -> + owner.getTypeUrn() != null + ? owner.getTypeUrn() + : DEFAULT_OWNERSHIP_TYPE_URN, + Collectors.toSet())); + } + } + return Collections.emptyMap(); + } + + public static String encodeFieldName(String value) { + return value.replaceAll("[.]", "%2E"); + } + + public static String decodeFieldName(String value) { + return value.replaceAll("%2E", "."); + } +} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/FieldSpecUtils.java b/entity-registry/src/main/java/com/linkedin/metadata/models/FieldSpecUtils.java index 53a689602f27c0..271f62128c70d2 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/FieldSpecUtils.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/FieldSpecUtils.java @@ -22,10 +22,15 @@ public static String getSchemaFieldName(PathSpec pathSpec) { return lastComponent; } - public static Map getResolvedProperties(final DataSchema schema) { - return !schema.getResolvedProperties().isEmpty() - ? schema.getResolvedProperties() - : schema.getProperties(); + public static Map getResolvedProperties( + final DataSchema schema, Map fallback) { + if (!schema.getResolvedProperties().isEmpty()) { + return schema.getResolvedProperties(); + } else if (!schema.getProperties().isEmpty()) { + return schema.getProperties(); + } else { + return fallback; + } } public static Optional getPathSpecWithAspectName(TraverserContext context) { diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/RelationshipFieldSpecExtractor.java b/entity-registry/src/main/java/com/linkedin/metadata/models/RelationshipFieldSpecExtractor.java index ad32b315f6b1ab..855fba2ad46fb1 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/RelationshipFieldSpecExtractor.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/RelationshipFieldSpecExtractor.java @@ -9,6 +9,7 @@ import com.linkedin.data.schema.annotation.TraverserContext; import com.linkedin.metadata.models.annotation.RelationshipAnnotation; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; @@ -46,7 +47,7 @@ public void callbackOnContext(TraverserContext context, DataSchemaTraverse.Order // Next, check resolved properties for annotations on primitives. final Map resolvedProperties = - FieldSpecUtils.getResolvedProperties(currentSchema); + FieldSpecUtils.getResolvedProperties(currentSchema, Collections.emptyMap()); final Object resolvedAnnotationObj = resolvedProperties.get(RelationshipAnnotation.ANNOTATION_NAME); diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/SearchableFieldSpecExtractor.java b/entity-registry/src/main/java/com/linkedin/metadata/models/SearchableFieldSpecExtractor.java index a20fd36f0d70c5..0c5e1a4c31598f 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/SearchableFieldSpecExtractor.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/SearchableFieldSpecExtractor.java @@ -109,10 +109,12 @@ private Object getAnnotationObj(TraverserContext context) { .equals("com.linkedin.common.urn.Urn"); final Map resolvedProperties = - FieldSpecUtils.getResolvedProperties(currentSchema); + FieldSpecUtils.getResolvedProperties(currentSchema, properties); // if primary doesn't have an annotation, then ignore secondary urns - if (isUrn && primaryAnnotationObj != null) { + if (isUrn + && primaryAnnotationObj != null + && resolvedProperties.containsKey(SearchableAnnotation.ANNOTATION_NAME)) { DataMap annotationMap = (DataMap) resolvedProperties.get(SearchableAnnotation.ANNOTATION_NAME); Map result = new HashMap<>(annotationMap); diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/annotation/SearchableAnnotation.java b/entity-registry/src/main/java/com/linkedin/metadata/models/annotation/SearchableAnnotation.java index 4808b7ee3b5ac8..4ec8702efde709 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/annotation/SearchableAnnotation.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/annotation/SearchableAnnotation.java @@ -20,6 +20,8 @@ public class SearchableAnnotation { public static final String FIELD_NAME_ALIASES = "fieldNameAliases"; public static final String ANNOTATION_NAME = "Searchable"; + public static final Set OBJECT_FIELD_TYPES = + ImmutableSet.of(FieldType.OBJECT, FieldType.MAP_ARRAY); private static final Set DEFAULT_QUERY_FIELD_TYPES = ImmutableSet.of( FieldType.TEXT, @@ -71,7 +73,8 @@ public enum FieldType { OBJECT, BROWSE_PATH_V2, WORD_GRAM, - DOUBLE + DOUBLE, + MAP_ARRAY } @Nonnull diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/hooks/OwnerTypeMapTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/hooks/OwnerTypeMapTest.java new file mode 100644 index 00000000000000..895744bb182eb6 --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/hooks/OwnerTypeMapTest.java @@ -0,0 +1,220 @@ +package com.linkedin.metadata.aspect.hooks; + +import static com.linkedin.metadata.Constants.DEFAULT_OWNERSHIP_TYPE_URN; +import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertEquals; + +import com.linkedin.common.Owner; +import com.linkedin.common.OwnerArray; +import com.linkedin.common.Ownership; +import com.linkedin.common.UrnArray; +import com.linkedin.common.UrnArrayMap; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.ChangeMCP; +import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.test.metadata.aspect.TestEntityRegistry; +import com.linkedin.test.metadata.aspect.batch.TestMCP; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import javax.annotation.Nullable; +import org.testng.annotations.Test; + +public class OwnerTypeMapTest { + private static final AspectRetriever ASPECT_RETRIEVER = mock(AspectRetriever.class); + private static final EntityRegistry ENTITY_REGISTRY = new TestEntityRegistry(); + private static final AspectPluginConfig ASPECT_PLUGIN_CONFIG = + AspectPluginConfig.builder() + .className("some class") + .enabled(true) + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("*") + .aspectName("ownership") + .build())) + .build(); + private static final Urn TEST_ENTITY_URN = + UrnUtils.getUrn( + "urn:li:dataset:(urn:li:dataPlatform:bigquery,calm-pagoda-323403.jaffle_shop.orders,PROD)"); + private static final Urn TEST_USER_A = UrnUtils.getUrn("urn:li:corpUser:a"); + private static final Urn TEST_USER_B = UrnUtils.getUrn("urn:li:corpUser:b"); + private static final Urn TEST_GROUP_A = UrnUtils.getUrn("urn:li:corpGroup:a"); + private static final Urn TEST_GROUP_B = UrnUtils.getUrn("urn:li:corpGroup:b"); + private static final Urn TECH_OWNER = + UrnUtils.getUrn("urn:li:ownershipType:__system__technical_owner"); + private static final Urn BUS_OWNER = + UrnUtils.getUrn("urn:li:ownershipType:__system__business_owner"); + + @Test + public void ownershipTypeMutationNoneType() { + OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG); + Ownership ownership = buildOwnership(Map.of(TEST_USER_A, List.of(), TEST_GROUP_A, List.of())); + testHook.writeMutation(buildMCP(null, ownership), ASPECT_RETRIEVER); + + assertEquals( + ownership.getOwnerTypes(), + new UrnArrayMap( + Map.of( + DEFAULT_OWNERSHIP_TYPE_URN.toString(), + new UrnArray(List.of(TEST_USER_A, TEST_GROUP_A)))), + "Expected generic owners to be grouped by `none` ownership type."); + } + + @Test + public void ownershipTypeMutationNoneTypeAdd() { + OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG); + Ownership oldOwnership = buildOwnership(Map.of(TEST_USER_A, List.of())); + Ownership newOwnership = + buildOwnership(Map.of(TEST_USER_A, List.of(), TEST_GROUP_A, List.of())); + testHook.writeMutation(buildMCP(oldOwnership, newOwnership), ASPECT_RETRIEVER); + + assertEquals( + newOwnership.getOwnerTypes(), + new UrnArrayMap( + Map.of( + DEFAULT_OWNERSHIP_TYPE_URN.toString(), + new UrnArray(List.of(TEST_USER_A, TEST_GROUP_A)))), + "Expected generic owners to be grouped by `none` ownership type."); + } + + @Test + public void ownershipTypeMutationNoneTypeRemove() { + OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG); + Ownership oldOwnership = + buildOwnership(Map.of(TEST_USER_A, List.of(), TEST_GROUP_A, List.of())); + Ownership newOwnership = buildOwnership(Map.of(TEST_USER_A, List.of())); + testHook.writeMutation(buildMCP(oldOwnership, newOwnership), ASPECT_RETRIEVER); + + assertEquals( + newOwnership.getOwnerTypes(), + new UrnArrayMap( + Map.of(DEFAULT_OWNERSHIP_TYPE_URN.toString(), new UrnArray(List.of(TEST_USER_A)))), + "Expected generic owners to be grouped by `none` ownership type."); + } + + @Test + public void ownershipTypeMutationMixedType() { + OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG); + Ownership ownership = + buildOwnership( + Map.of( + TEST_USER_A, + List.of(), + TEST_GROUP_A, + List.of(), + TEST_USER_B, + List.of(BUS_OWNER), + TEST_GROUP_B, + List.of(TECH_OWNER))); + testHook.writeMutation(buildMCP(null, ownership), ASPECT_RETRIEVER); + + assertEquals( + ownership.getOwnerTypes(), + new UrnArrayMap( + Map.of( + DEFAULT_OWNERSHIP_TYPE_URN.toString(), + new UrnArray(List.of(TEST_USER_A, TEST_GROUP_A)), + BUS_OWNER.toString(), + new UrnArray(List.of(TEST_USER_B)), + TECH_OWNER.toString(), + new UrnArray(List.of(TEST_GROUP_B)))), + "Expected generic owners to be grouped by `none` ownership type as well as specified types."); + } + + @Test + public void ownershipTypeMutationMixedTypeAdd() { + OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG); + Ownership oldOwnership = + buildOwnership(Map.of(TEST_USER_A, List.of(), TEST_USER_B, List.of(BUS_OWNER))); + Ownership newOwnership = + buildOwnership( + Map.of( + TEST_USER_A, + List.of(), + TEST_GROUP_A, + List.of(), + TEST_USER_B, + List.of(BUS_OWNER), + TEST_GROUP_B, + List.of(TECH_OWNER))); + testHook.writeMutation(buildMCP(oldOwnership, newOwnership), ASPECT_RETRIEVER); + + assertEquals( + newOwnership.getOwnerTypes(), + new UrnArrayMap( + Map.of( + DEFAULT_OWNERSHIP_TYPE_URN.toString(), + new UrnArray(List.of(TEST_USER_A, TEST_GROUP_A)), + BUS_OWNER.toString(), + new UrnArray(List.of(TEST_USER_B)), + TECH_OWNER.toString(), + new UrnArray(List.of(TEST_GROUP_B)))), + "Expected generic owners to be grouped by `none` ownership type as well as specified types."); + } + + @Test + public void ownershipTypeMutationMixedTypeRemove() { + OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG); + Ownership oldOwnership = + buildOwnership( + Map.of( + TEST_USER_A, + List.of(), + TEST_GROUP_A, + List.of(), + TEST_USER_B, + List.of(BUS_OWNER), + TEST_GROUP_B, + List.of(TECH_OWNER))); + Ownership newOwnership = + buildOwnership(Map.of(TEST_GROUP_A, List.of(), TEST_GROUP_B, List.of(TECH_OWNER))); + testHook.writeMutation(buildMCP(oldOwnership, newOwnership), ASPECT_RETRIEVER); + + assertEquals( + newOwnership.getOwnerTypes(), + new UrnArrayMap( + Map.of( + DEFAULT_OWNERSHIP_TYPE_URN.toString(), + new UrnArray(List.of(TEST_GROUP_A)), + BUS_OWNER.toString(), + new UrnArray(), + TECH_OWNER.toString(), + new UrnArray(List.of(TEST_GROUP_B)))), + "Expected generic owners to be grouped by `none` ownership type as well as specified types."); + } + + private static Ownership buildOwnership(Map> ownershipTypes) { + Ownership ownership = new Ownership(); + ownership.setOwners( + ownershipTypes.entrySet().stream() + .flatMap( + entry -> { + if (entry.getValue().isEmpty()) { + Owner owner = new Owner(); + owner.setOwner(entry.getKey()); + return Stream.of(owner); + } else { + return entry.getValue().stream() + .map( + typeUrn -> { + Owner owner = new Owner(); + owner.setOwner(entry.getKey()); + owner.setTypeUrn(typeUrn); + return owner; + }); + } + }) + .collect(Collectors.toCollection(OwnerArray::new))); + return ownership; + } + + private static Set buildMCP(@Nullable Ownership oldOwnership, Ownership newOwnership) { + return TestMCP.ofOneMCP(TEST_ENTITY_URN, oldOwnership, newOwnership, ENTITY_REGISTRY); + } +} diff --git a/gradle.properties b/gradle.properties index f410ff01bf397d..0b797d6a9ab209 100644 --- a/gradle.properties +++ b/gradle.properties @@ -3,8 +3,8 @@ org.gradle.configureondemand=true org.gradle.parallel=true org.gradle.caching=true -# Increase gradle JVM memory to 3GB to allow tests to run locally -org.gradle.jvmargs=-Xmx3000m +# Increase gradle JVM memory to 4GB to allow tests to run locally +org.gradle.jvmargs=-Xmx4000m # Increase retries to 5 (from default of 3) and increase interval from 125ms to 1s. # Based on this thread https://github.com/gradle/gradle/issues/4629, it's unclear # if we should be using systemProp or not. We're using both for now. diff --git a/ingestion-scheduler/src/main/java/com/datahub/metadata/ingestion/IngestionScheduler.java b/ingestion-scheduler/src/main/java/com/datahub/metadata/ingestion/IngestionScheduler.java index 3d3883b4af9df1..1472fa6b19f824 100644 --- a/ingestion-scheduler/src/main/java/com/datahub/metadata/ingestion/IngestionScheduler.java +++ b/ingestion-scheduler/src/main/java/com/datahub/metadata/ingestion/IngestionScheduler.java @@ -1,6 +1,5 @@ package com.datahub.metadata.ingestion; -import com.datahub.authentication.Authentication; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; @@ -24,6 +23,7 @@ import com.linkedin.metadata.utils.IngestionUtils; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; import java.time.ZonedDateTime; import java.util.ArrayList; import java.util.Collections; @@ -75,34 +75,34 @@ @RequiredArgsConstructor public class IngestionScheduler { - private final Authentication _systemAuthentication; - private final EntityClient _entityClient; + private final OperationContext systemOpContext; + private final EntityClient entityClient; // Maps a DataHubIngestionSource to a future representing the "next" scheduled execution of the // source // Visible for testing - final Map> _nextIngestionSourceExecutionCache = new HashMap<>(); + final Map> nextIngestionSourceExecutionCache = new HashMap<>(); // Shared executor service used for executing an ingestion source on a schedule - private final ScheduledExecutorService _sharedExecutorService = + private final ScheduledExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(1); - private final IngestionConfiguration _ingestionConfiguration; - private final int _batchGetDelayIntervalSeconds; - private final int _batchGetRefreshIntervalSeconds; + private final IngestionConfiguration ingestionConfiguration; + private final int batchGetDelayIntervalSeconds; + private final int batchGetRefreshIntervalSeconds; public void init() { final BatchRefreshSchedulesRunnable batchRefreshSchedulesRunnable = new BatchRefreshSchedulesRunnable( - _systemAuthentication, - _entityClient, + systemOpContext, + entityClient, this::scheduleNextIngestionSourceExecution, this::unscheduleAll); // Schedule a recurring batch-reload task. - _sharedExecutorService.scheduleAtFixedRate( + scheduledExecutorService.scheduleAtFixedRate( batchRefreshSchedulesRunnable, - _batchGetDelayIntervalSeconds, - _batchGetRefreshIntervalSeconds, + batchGetDelayIntervalSeconds, + batchGetRefreshIntervalSeconds, TimeUnit.SECONDS); } @@ -110,10 +110,10 @@ public void init() { public void unscheduleNextIngestionSourceExecution(final Urn ingestionSourceUrn) { log.info("Unscheduling ingestion source with urn {}", ingestionSourceUrn); // Deleting an ingestion source schedule. Un-schedule the next execution. - ScheduledFuture future = _nextIngestionSourceExecutionCache.get(ingestionSourceUrn); + ScheduledFuture future = nextIngestionSourceExecutionCache.get(ingestionSourceUrn); if (future != null) { future.cancel(false); // Do not interrupt running processes - _nextIngestionSourceExecutionCache.remove(ingestionSourceUrn); + nextIngestionSourceExecutionCache.remove(ingestionSourceUrn); } } @@ -125,7 +125,7 @@ public void unscheduleAll() { // Deleting an ingestion source schedule. Un-schedule the next execution. Set scheduledSources = new HashSet<>( - _nextIngestionSourceExecutionCache.keySet()); // Create copy to avoid concurrent mod. + nextIngestionSourceExecutionCache.keySet()); // Create copy to avoid concurrent mod. for (Urn urn : scheduledSources) { unscheduleNextIngestionSourceExecution(urn); } @@ -173,19 +173,19 @@ public void scheduleNextIngestionSourceExecution( // Schedule the ingestion source to run some time in the future. final ExecutionRequestRunnable executionRequestRunnable = new ExecutionRequestRunnable( - _systemAuthentication, - _entityClient, - _ingestionConfiguration, + systemOpContext, + entityClient, + ingestionConfiguration, ingestionSourceUrn, newInfo, - () -> _nextIngestionSourceExecutionCache.remove(ingestionSourceUrn), + () -> nextIngestionSourceExecutionCache.remove(ingestionSourceUrn), this::scheduleNextIngestionSourceExecution); // Schedule the next ingestion run final ScheduledFuture scheduledFuture = - _sharedExecutorService.schedule( + scheduledExecutorService.schedule( executionRequestRunnable, scheduleTime, TimeUnit.MILLISECONDS); - _nextIngestionSourceExecutionCache.put(ingestionSourceUrn, scheduledFuture); + nextIngestionSourceExecutionCache.put(ingestionSourceUrn, scheduledFuture); log.info( String.format( @@ -216,22 +216,22 @@ public void scheduleNextIngestionSourceExecution( @VisibleForTesting static class BatchRefreshSchedulesRunnable implements Runnable { - private final Authentication _systemAuthentication; - private final EntityClient _entityClient; - private final BiConsumer _scheduleNextIngestionSourceExecution; - private final Runnable _unscheduleAll; + private final OperationContext systemOpContext; + private final EntityClient entityClient; + private final BiConsumer scheduleNextIngestionSourceExecution; + private final Runnable unscheduleAll; public BatchRefreshSchedulesRunnable( - @Nonnull final Authentication systemAuthentication, + @Nonnull final OperationContext systemOpContext, @Nonnull final EntityClient entityClient, @Nonnull final BiConsumer scheduleNextIngestionSourceExecution, @Nonnull final Runnable unscheduleAll) { - _systemAuthentication = Objects.requireNonNull(systemAuthentication); - _entityClient = Objects.requireNonNull(entityClient); - _scheduleNextIngestionSourceExecution = + this.systemOpContext = systemOpContext; + this.entityClient = Objects.requireNonNull(entityClient); + this.scheduleNextIngestionSourceExecution = Objects.requireNonNull(scheduleNextIngestionSourceExecution); - _unscheduleAll = unscheduleAll; + this.unscheduleAll = unscheduleAll; } @Override @@ -239,7 +239,7 @@ public void run() { try { // First un-schedule all currently scheduled runs (to make sure consistency is maintained) - _unscheduleAll.run(); + unscheduleAll.run(); int start = 0; int count = 30; @@ -254,20 +254,20 @@ public void run() { // 1. List all ingestion source urns. final ListResult ingestionSourceUrns = - _entityClient.list( + entityClient.list( + systemOpContext, Constants.INGESTION_SOURCE_ENTITY_NAME, Collections.emptyMap(), start, - count, - _systemAuthentication); + count); // 2. Fetch all ingestion sources, specifically the "info" aspect. final Map ingestionSources = - _entityClient.batchGetV2( + entityClient.batchGetV2( Constants.INGESTION_SOURCE_ENTITY_NAME, new HashSet<>(ingestionSourceUrns.getEntities()), ImmutableSet.of(Constants.INGESTION_INFO_ASPECT_NAME), - _systemAuthentication); + systemOpContext.getSessionAuthentication()); // 3. Reschedule ingestion sources based on the fetched schedules (inside "info") log.debug( @@ -310,7 +310,7 @@ private void scheduleNextIngestionRuns( new DataHubIngestionSourceInfo(envelopedInfo.getValue().data()); // Invoke the "scheduleNextIngestionSourceExecution" (passed from parent) - _scheduleNextIngestionSourceExecution.accept(entityUrn, ingestionSourceInfo); + scheduleNextIngestionSourceExecution.accept(entityUrn, ingestionSourceInfo); } } } @@ -330,23 +330,23 @@ static class ExecutionRequestRunnable implements Runnable { private static final String VERSION_ARGUMENT_NAME = "version"; private static final String DEBUG_MODE_ARG_NAME = "debug_mode"; - private final Authentication _systemAuthentication; - private final EntityClient _entityClient; - private final IngestionConfiguration _ingestionConfiguration; + private final OperationContext systemOpContext; + private final EntityClient entityClient; + private final IngestionConfiguration ingestionConfiguration; // Information about the ingestion source being executed - private final Urn _ingestionSourceUrn; - private final DataHubIngestionSourceInfo _ingestionSourceInfo; + private final Urn ingestionSourceUrn; + private final DataHubIngestionSourceInfo ingestionSourceInfo; // Used for clearing the "next execution" cache once a corresponding execution request has been // created. - private final Runnable _deleteNextIngestionSourceExecution; + private final Runnable deleteNextIngestionSourceExecution; // Used for re-scheduling the ingestion source once it has executed! - private final BiConsumer _scheduleNextIngestionSourceExecution; + private final BiConsumer scheduleNextIngestionSourceExecution; public ExecutionRequestRunnable( - @Nonnull final Authentication systemAuthentication, + @Nonnull final OperationContext systemOpContext, @Nonnull final EntityClient entityClient, @Nonnull final IngestionConfiguration ingestionConfiguration, @Nonnull final Urn ingestionSourceUrn, @@ -355,14 +355,14 @@ public ExecutionRequestRunnable( @Nonnull final BiConsumer scheduleNextIngestionSourceExecution) { - _systemAuthentication = Objects.requireNonNull(systemAuthentication); - _entityClient = Objects.requireNonNull(entityClient); - _ingestionConfiguration = Objects.requireNonNull(ingestionConfiguration); - _ingestionSourceUrn = Objects.requireNonNull(ingestionSourceUrn); - _ingestionSourceInfo = Objects.requireNonNull(ingestionSourceInfo); - _deleteNextIngestionSourceExecution = + this.systemOpContext = systemOpContext; + this.entityClient = Objects.requireNonNull(entityClient); + this.ingestionConfiguration = Objects.requireNonNull(ingestionConfiguration); + this.ingestionSourceUrn = Objects.requireNonNull(ingestionSourceUrn); + this.ingestionSourceInfo = Objects.requireNonNull(ingestionSourceInfo); + this.deleteNextIngestionSourceExecution = Objects.requireNonNull(deleteNextIngestionSourceExecution); - _scheduleNextIngestionSourceExecution = + this.scheduleNextIngestionSourceExecution = Objects.requireNonNull(scheduleNextIngestionSourceExecution); } @@ -371,14 +371,14 @@ public void run() { // Remove the next ingestion execution as we are going to execute it now. (no retry logic // currently) - _deleteNextIngestionSourceExecution.run(); + deleteNextIngestionSourceExecution.run(); try { log.info( String.format( "Creating Execution Request for scheduled Ingestion Source with urn %s", - _ingestionSourceUrn)); + ingestionSourceUrn)); // Create a new Execution Request Proposal final MetadataChangeProposal proposal = new MetadataChangeProposal(); @@ -395,23 +395,23 @@ public void run() { input.setSource( new ExecutionRequestSource() .setType(EXECUTION_REQUEST_SOURCE_NAME) - .setIngestionSource(_ingestionSourceUrn)); - input.setExecutorId(_ingestionSourceInfo.getConfig().getExecutorId(), SetMode.IGNORE_NULL); + .setIngestionSource(ingestionSourceUrn)); + input.setExecutorId(ingestionSourceInfo.getConfig().getExecutorId(), SetMode.IGNORE_NULL); input.setRequestedAt(System.currentTimeMillis()); Map arguments = new HashMap<>(); String recipe = IngestionUtils.injectPipelineName( - _ingestionSourceInfo.getConfig().getRecipe(), _ingestionSourceUrn.toString()); + ingestionSourceInfo.getConfig().getRecipe(), ingestionSourceUrn.toString()); arguments.put(RECIPE_ARGUMENT_NAME, recipe); arguments.put( VERSION_ARGUMENT_NAME, - _ingestionSourceInfo.getConfig().hasVersion() - ? _ingestionSourceInfo.getConfig().getVersion() - : _ingestionConfiguration.getDefaultCliVersion()); + ingestionSourceInfo.getConfig().hasVersion() + ? ingestionSourceInfo.getConfig().getVersion() + : ingestionConfiguration.getDefaultCliVersion()); String debugMode = "false"; - if (_ingestionSourceInfo.getConfig().hasDebugMode()) { - debugMode = _ingestionSourceInfo.getConfig().isDebugMode() ? "true" : "false"; + if (ingestionSourceInfo.getConfig().hasDebugMode()) { + debugMode = ingestionSourceInfo.getConfig().isDebugMode() ? "true" : "false"; } arguments.put(DEBUG_MODE_ARG_NAME, debugMode); input.setArgs(new StringMap(arguments)); @@ -421,18 +421,18 @@ public void run() { proposal.setAspect(GenericRecordUtils.serializeAspect(input)); proposal.setChangeType(ChangeType.UPSERT); - _entityClient.ingestProposal(proposal, _systemAuthentication); + entityClient.ingestProposal(proposal, systemOpContext.getSystemAuthentication().get()); } catch (Exception e) { // TODO: This type of thing should likely be proactively reported. log.error( String.format( "Caught exception while attempting to create Execution Request for Ingestion Source with urn %s. Will retry on next scheduled attempt.", - _ingestionSourceUrn), + ingestionSourceUrn), e); } // 2. Re-Schedule the next execution request. - _scheduleNextIngestionSourceExecution.accept(_ingestionSourceUrn, _ingestionSourceInfo); + scheduleNextIngestionSourceExecution.accept(ingestionSourceUrn, ingestionSourceInfo); } } diff --git a/ingestion-scheduler/src/test/java/com/datahub/metadata/ingestion/IngestionSchedulerTest.java b/ingestion-scheduler/src/test/java/com/datahub/metadata/ingestion/IngestionSchedulerTest.java index 8174afc20765f0..f07a9df941a14a 100644 --- a/ingestion-scheduler/src/test/java/com/datahub/metadata/ingestion/IngestionSchedulerTest.java +++ b/ingestion-scheduler/src/test/java/com/datahub/metadata/ingestion/IngestionSchedulerTest.java @@ -2,7 +2,6 @@ import static org.testng.Assert.*; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.linkedin.common.UrnArray; @@ -18,6 +17,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.config.IngestionConfiguration; import com.linkedin.metadata.query.ListResult; +import io.datahubproject.metadata.context.OperationContext; import java.util.Collections; import java.util.concurrent.Future; import java.util.concurrent.ScheduledFuture; @@ -93,11 +93,11 @@ public void setupTest() throws Exception { // Set up mocks for ingestion source batch fetching Mockito.when( mockClient.list( + Mockito.any(), Mockito.eq(Constants.INGESTION_SOURCE_ENTITY_NAME), Mockito.eq(Collections.emptyMap()), Mockito.eq(0), - Mockito.eq(30), - Mockito.any())) + Mockito.eq(30))) .thenReturn( new ListResult() .setCount(30) @@ -117,7 +117,7 @@ public void setupTest() throws Exception { _ingestionScheduler = new IngestionScheduler( - Mockito.mock(Authentication.class), + Mockito.mock(OperationContext.class), mockClient, Mockito.mock(IngestionConfiguration.class), 1, @@ -128,11 +128,11 @@ public void setupTest() throws Exception { @Test public void testInvokeUpdateExistingSchedule() throws Exception { - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 1); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 1); Urn ingestionSourceUrn = Urn.createFromString("urn:li:dataHubIngestionSourceUrn:0"); Future beforeFuture = - _ingestionScheduler._nextIngestionSourceExecutionCache.get(ingestionSourceUrn); + _ingestionScheduler.nextIngestionSourceExecutionCache.get(ingestionSourceUrn); final DataHubIngestionSourceInfo newInfo = new DataHubIngestionSourceInfo(); newInfo.setSchedule( @@ -149,9 +149,9 @@ public void testInvokeUpdateExistingSchedule() throws Exception { // Assert that the new source has been scheduled successfully. _ingestionScheduler.scheduleNextIngestionSourceExecution(ingestionSourceUrn, newInfo); - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 1); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 1); Future newFuture = - _ingestionScheduler._nextIngestionSourceExecutionCache.get(ingestionSourceUrn); + _ingestionScheduler.nextIngestionSourceExecutionCache.get(ingestionSourceUrn); // Ensure that there is an overwritten future. Assert.assertNotSame(beforeFuture, newFuture); @@ -159,7 +159,7 @@ public void testInvokeUpdateExistingSchedule() throws Exception { @Test public void testInvokeNewSchedule() throws Exception { - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 1); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 1); final Urn urn = Urn.createFromString("urn:li:dataHubIngestionSourceUrn:2"); final DataHubIngestionSourceInfo newInfo = new DataHubIngestionSourceInfo(); @@ -177,12 +177,12 @@ public void testInvokeNewSchedule() throws Exception { // Assert that the new source has been scheduled successfully. _ingestionScheduler.scheduleNextIngestionSourceExecution(urn, newInfo); - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 2); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 2); } @Test public void testInvokeInvalidSchedule() throws Exception { - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 1); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 1); final Urn urn = Urn.createFromString("urn:li:dataHubIngestionSourceUrn:2"); final DataHubIngestionSourceInfo newInfo = new DataHubIngestionSourceInfo(); @@ -201,12 +201,12 @@ public void testInvokeInvalidSchedule() throws Exception { // Assert that no changes have been made to next execution cache. _ingestionScheduler.scheduleNextIngestionSourceExecution(urn, newInfo); - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 1); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 1); } @Test public void testInvokeMissingSchedule() throws Exception { - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 1); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 1); final Urn urn = Urn.createFromString("urn:li:dataHubIngestionSourceUrn:0"); final DataHubIngestionSourceInfo newInfo = new DataHubIngestionSourceInfo(); @@ -221,27 +221,27 @@ public void testInvokeMissingSchedule() throws Exception { // Assert that the schedule has been removed. _ingestionScheduler.scheduleNextIngestionSourceExecution(urn, newInfo); - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 0); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 0); } @Test public void testInvokeDelete() throws Exception { - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 1); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 1); // Attempt to delete an unscheduled urn final Urn urn1 = Urn.createFromString("urn:li:dataHubIngestionSource:not-scheduled"); _ingestionScheduler.unscheduleNextIngestionSourceExecution(urn1); - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 1); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 1); // Attempt to delete a scheduled urn final Urn urn2 = Urn.createFromString("urn:li:dataHubIngestionSourceUrn:0"); _ingestionScheduler.unscheduleNextIngestionSourceExecution(urn2); - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 0); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 0); } @Test public void testSchedule() throws Exception { - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 1); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 1); final Urn urn = Urn.createFromString("urn:li:dataHubIngestionSourceUrn:0"); final DataHubIngestionSourceInfo newInfo = new DataHubIngestionSourceInfo(); @@ -259,7 +259,7 @@ public void testSchedule() throws Exception { _ingestionScheduler.scheduleNextIngestionSourceExecution(urn, newInfo); - ScheduledFuture future = _ingestionScheduler._nextIngestionSourceExecutionCache.get(urn); + ScheduledFuture future = _ingestionScheduler.nextIngestionSourceExecutionCache.get(urn); Assert.assertTrue( future.getDelay(TimeUnit.SECONDS) < 60); // Next execution must always be less than a minute away. @@ -267,7 +267,7 @@ public void testSchedule() throws Exception { @Test public void testUnscheduleAll() throws Exception { - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 1); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 1); final Urn urn = Urn.createFromString("urn:li:dataHubIngestionSourceUrn:3"); final DataHubIngestionSourceInfo newInfo = new DataHubIngestionSourceInfo(); @@ -284,16 +284,16 @@ public void testUnscheduleAll() throws Exception { .setVersion("0.8.18")); _ingestionScheduler.scheduleNextIngestionSourceExecution(urn, newInfo); - assertEquals(_ingestionScheduler._nextIngestionSourceExecutionCache.size(), 2); + assertEquals(_ingestionScheduler.nextIngestionSourceExecutionCache.size(), 2); // Get reference to schedules futures - ScheduledFuture future = _ingestionScheduler._nextIngestionSourceExecutionCache.get(urn); + ScheduledFuture future = _ingestionScheduler.nextIngestionSourceExecutionCache.get(urn); // Unschedule all _ingestionScheduler.unscheduleAll(); // Ensure that the cache is empty - Assert.assertTrue(_ingestionScheduler._nextIngestionSourceExecutionCache.isEmpty()); + Assert.assertTrue(_ingestionScheduler.nextIngestionSourceExecutionCache.isEmpty()); // And that the future is cancelled Assert.assertTrue(future.isCancelled()); diff --git a/li-utils/src/main/java/com/linkedin/metadata/Constants.java b/li-utils/src/main/java/com/linkedin/metadata/Constants.java index 94be2f288521cf..6e089865956617 100644 --- a/li-utils/src/main/java/com/linkedin/metadata/Constants.java +++ b/li-utils/src/main/java/com/linkedin/metadata/Constants.java @@ -1,6 +1,7 @@ package com.linkedin.metadata; import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; /** Static class containing commonly-used constants across DataHub services. */ public class Constants { @@ -79,10 +80,13 @@ public class Constants { public static final String QUERY_ENTITY_NAME = "query"; public static final String DATA_PRODUCT_ENTITY_NAME = "dataProduct"; public static final String OWNERSHIP_TYPE_ENTITY_NAME = "ownershipType"; + public static final Urn DEFAULT_OWNERSHIP_TYPE_URN = + UrnUtils.getUrn("urn:li:ownershipType:__system__none"); public static final String STRUCTURED_PROPERTY_ENTITY_NAME = "structuredProperty"; public static final String DATA_TYPE_ENTITY_NAME = "dataType"; public static final String ENTITY_TYPE_ENTITY_NAME = "entityType"; public static final String FORM_ENTITY_NAME = "form"; + public static final String RESTRICTED_ENTITY_NAME = "restricted"; /** Aspects */ // Common diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ConjunctivePrivilegeGroup.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ConjunctivePrivilegeGroup.java index bc3a3c9f385a68..adbfdbe3236fc6 100644 --- a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ConjunctivePrivilegeGroup.java +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/ConjunctivePrivilegeGroup.java @@ -1,6 +1,6 @@ package com.datahub.authorization; -import java.util.List; +import java.util.Collection; /** * Represents a group of privileges that must ALL be required to authorize a request. @@ -8,13 +8,13 @@ *

That is, an AND of privileges. */ public class ConjunctivePrivilegeGroup { - private final List _requiredPrivileges; + private final Collection _requiredPrivileges; - public ConjunctivePrivilegeGroup(List requiredPrivileges) { + public ConjunctivePrivilegeGroup(Collection requiredPrivileges) { _requiredPrivileges = requiredPrivileges; } - public List getRequiredPrivileges() { + public Collection getRequiredPrivileges() { return _requiredPrivileges; } } diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/config/SearchAuthorizationConfiguration.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/config/SearchAuthorizationConfiguration.java new file mode 100644 index 00000000000000..cb176130d2e786 --- /dev/null +++ b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/config/SearchAuthorizationConfiguration.java @@ -0,0 +1,24 @@ +package com.datahub.authorization.config; + +import lombok.AccessLevel; +import lombok.AllArgsConstructor; +import lombok.Builder; +import lombok.Data; +import lombok.NoArgsConstructor; + +@Builder(toBuilder = true) +@Data +@AllArgsConstructor(access = AccessLevel.PACKAGE) +@NoArgsConstructor(access = AccessLevel.PACKAGE) +public class SearchAuthorizationConfiguration { + private boolean enabled; + private SearchAuthorizationRecommendationsConfiguration recommendations; + + @Builder(toBuilder = true) + @Data + @AllArgsConstructor(access = AccessLevel.PACKAGE) + @NoArgsConstructor(access = AccessLevel.PACKAGE) + public static class SearchAuthorizationRecommendationsConfiguration { + private boolean peerGroupEnabled; + } +} 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 a6baf0b5b282cf..3c113dcd2e0526 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 @@ -6,8 +6,13 @@ import com.datahub.authorization.AuthorizerContext; import com.datahub.authorization.EntitySpec; import com.datahub.plugins.Plugin; +import com.linkedin.common.urn.Urn; +import com.linkedin.policy.DataHubPolicyInfo; +import java.util.Collection; +import java.util.Collections; import java.util.Map; import java.util.Optional; +import java.util.Set; import javax.annotation.Nonnull; /** @@ -15,22 +20,53 @@ * privilege. */ public interface Authorizer extends Plugin { + Authorizer EMPTY = new Authorizer() {}; + /** * Initialize the Authorizer. Invoked once at boot time. * * @param authorizerConfig config provided to the authenticator derived from the Metadata Service * YAML config. This config comes from the "authorization.authorizers.config" configuration. */ - void init( - @Nonnull final Map authorizerConfig, @Nonnull final AuthorizerContext ctx); + default void init( + @Nonnull final Map authorizerConfig, @Nonnull final AuthorizerContext ctx) {} /** Authorizes an action based on the actor, the resource, and required privileges. */ - AuthorizationResult authorize(@Nonnull final AuthorizationRequest request); + default AuthorizationResult authorize(@Nonnull final AuthorizationRequest request) { + return new AuthorizationResult(request, AuthorizationResult.Type.DENY, "Not Implemented."); + } /** * Retrieves the current list of actors authorized to for a particular privilege against an * optional resource */ - AuthorizedActors authorizedActors( - final String privilege, final Optional resourceSpec); + default AuthorizedActors authorizedActors( + final String privilege, final Optional resourceSpec) { + return AuthorizedActors.builder() + .privilege(privilege) + .users(Collections.emptyList()) + .roles(Collections.emptyList()) + .groups(Collections.emptyList()) + .build(); + } + + /** + * Given the actor's urn retrieve the policies. + * + * @param actorUrn + * @return + */ + default Set getActorPolicies(@Nonnull Urn actorUrn) { + return Collections.emptySet(); + } + + /** Given the actor's urn retrieve the actor's groups */ + default Collection getActorGroups(@Nonnull Urn actorUrn) { + return Collections.emptyList(); + } + + /** Given an actor's urn retrieve the actor's peers */ + default Collection getActorPeers(@Nonnull Urn actorUrn) { + return Collections.emptyList(); + } } diff --git a/metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mae.json b/metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mae.json index bd4bd5878d3bb7..0f107f699428fa 100644 --- a/metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mae.json +++ b/metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mae.json @@ -16,6 +16,7 @@ "source": null } ], + "ownerTypes": null, "lastModified": { "time": 0, "actor": "urn:li:corpuser:foobar", diff --git a/metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mce.json b/metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mce.json index 8aba67eec4cee6..850660e59a60e5 100644 --- a/metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mce.json +++ b/metadata-events/mxe-utils-avro/src/test/resources/test-avro2pegasus-mce.json @@ -14,6 +14,7 @@ "source": null } ], + "ownerTypes": null, "lastModified": { "time": 0, "actor": "urn:li:corpuser:foobar", diff --git a/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py b/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py index cc8710325a8f18..7a2dfa7ae0705b 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py +++ b/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py @@ -226,7 +226,7 @@ def get_resource_owners_work_unit( if not current_ownership: # If we want to overwrite or there are no existing tags, create a new GlobalTags object - current_ownership = OwnershipClass(owners, get_audit_stamp()) + current_ownership = OwnershipClass(owners, lastModified=get_audit_stamp()) else: current_owner_urns: Set[str] = set( [owner.owner for owner in current_ownership.owners] diff --git a/metadata-ingestion/tests/unit/test_rest_sink.py b/metadata-ingestion/tests/unit/test_rest_sink.py index 82e02aced5a670..7bfa09a35951b4 100644 --- a/metadata-ingestion/tests/unit/test_rest_sink.py +++ b/metadata-ingestion/tests/unit/test_rest_sink.py @@ -235,7 +235,7 @@ "changeType": "UPSERT", "aspectName": "ownership", "aspect": { - "value": '{"owners": [{"owner": "urn:li:corpuser:fbar", "type": "DATAOWNER"}], "lastModified": {"time": 0, "actor": "urn:li:corpuser:fbar"}}', + "value": '{"owners": [{"owner": "urn:li:corpuser:fbar", "type": "DATAOWNER"}], "ownerTypes": {}, "lastModified": {"time": 0, "actor": "urn:li:corpuser:fbar"}}', "contentType": "application/json", }, } diff --git a/metadata-io/build.gradle b/metadata-io/build.gradle index f96517d93fca6a..e07a6ee55c21fb 100644 --- a/metadata-io/build.gradle +++ b/metadata-io/build.gradle @@ -9,6 +9,7 @@ configurations { dependencies { implementation project(':entity-registry') + implementation project(':metadata-service:auth-config') api project(':metadata-utils') api project(':metadata-events:mxe-avro') api project(':metadata-events:mxe-registration') @@ -17,6 +18,7 @@ dependencies { api project(':metadata-service:restli-client') api project(':metadata-service:configuration') api project(':metadata-service:services') + api project(':metadata-operation-context') implementation spec.product.pegasus.data implementation spec.product.pegasus.generator diff --git a/metadata-io/src/main/java/com/linkedin/metadata/client/EntityClientAspectRetriever.java b/metadata-io/src/main/java/com/linkedin/metadata/client/EntityClientAspectRetriever.java index 0fcb765b340cf2..6b3d1ad6e193c0 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/client/EntityClientAspectRetriever.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/client/EntityClientAspectRetriever.java @@ -17,11 +17,12 @@ import lombok.RequiredArgsConstructor; import org.springframework.stereotype.Component; +@Getter @Builder @Component @RequiredArgsConstructor public class EntityClientAspectRetriever implements CachingAspectRetriever { - @Getter private final EntityRegistry entityRegistry; + private final EntityRegistry entityRegistry; private final SystemEntityClient entityClient; /** diff --git a/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java b/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java index fed6379f921045..c261d7fefd4117 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java @@ -36,7 +36,6 @@ import com.linkedin.metadata.query.AutoCompleteResult; import com.linkedin.metadata.query.ListResult; import com.linkedin.metadata.query.ListUrnsResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.search.EntitySearchService; @@ -57,6 +56,7 @@ import com.linkedin.parseq.retry.backoff.BackoffPolicy; import com.linkedin.parseq.retry.backoff.ExponentialBackoff; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.time.Clock; @@ -84,6 +84,7 @@ public class JavaEntityClient implements EntityClient { private final Clock _clock = Clock.systemUTC(); + private final OperationContext opContext; private final EntityService entityService; private final DeleteEntityService deleteEntityService; private final EntitySearchService entitySearchService; @@ -165,15 +166,15 @@ public Map batchGet( */ @Nonnull public AutoCompleteResult autoComplete( + @Nonnull OperationContext opContext, @Nonnull String entityType, @Nonnull String query, @Nullable Filter requestFilters, @Nonnull int limit, - @Nullable String field, - @Nonnull final Authentication authentication) + @Nullable String field) throws RemoteInvocationException { return cachingEntitySearchService.autoComplete( - entityType, query, field, filterOrDefaultEmptyFilter(requestFilters), limit, null); + opContext, entityType, query, field, filterOrDefaultEmptyFilter(requestFilters), limit); } /** @@ -187,14 +188,14 @@ public AutoCompleteResult autoComplete( */ @Nonnull public AutoCompleteResult autoComplete( + @Nonnull OperationContext opContext, @Nonnull String entityType, @Nonnull String query, @Nullable Filter requestFilters, - @Nonnull int limit, - @Nonnull final Authentication authentication) + @Nonnull int limit) throws RemoteInvocationException { return cachingEntitySearchService.autoComplete( - entityType, query, "", filterOrDefaultEmptyFilter(requestFilters), limit, null); + opContext, entityType, query, "", filterOrDefaultEmptyFilter(requestFilters), limit); } /** @@ -207,18 +208,19 @@ public AutoCompleteResult autoComplete( * @param limit max number of datasets * @throws RemoteInvocationException */ + @Override @Nonnull public BrowseResult browse( + @Nonnull OperationContext opContext, @Nonnull String entityType, @Nonnull String path, @Nullable Map requestFilters, int start, - int limit, - @Nonnull final Authentication authentication) + int limit) throws RemoteInvocationException { return ValidationUtils.validateBrowseResult( cachingEntitySearchService.browse( - entityType, path, newFilter(requestFilters), start, limit, null), + opContext, entityType, path, newFilter(requestFilters), start, limit), entityService); } @@ -235,16 +237,15 @@ entityType, path, newFilter(requestFilters), start, limit, null), */ @Nonnull public BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nonnull Authentication authentication, - @Nullable SearchFlags searchFlags) { + int count) { // TODO: cache browseV2 results - return entitySearchService.browseV2(entityName, path, filter, input, start, count, searchFlags); + return entitySearchService.browseV2(opContext, entityName, path, filter, input, start, count); } /** @@ -260,17 +261,15 @@ public BrowseResultV2 browseV2( */ @Nonnull public BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nonnull Authentication authentication, - @Nullable SearchFlags searchFlags) { + int count) { // TODO: cache browseV2 results - return entitySearchService.browseV2( - entityNames, path, filter, input, start, count, searchFlags); + return entitySearchService.browseV2(opContext, entityNames, path, filter, input, start, count); } @SneakyThrows @@ -324,7 +323,6 @@ public void batchUpdate( * @param requestFilters search filters * @param start start offset for search results * @param count max number of search results requested - * @param searchFlags * @return a set of search results * @throws RemoteInvocationException */ @@ -332,18 +330,17 @@ public void batchUpdate( @WithSpan @Override public SearchResult search( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nonnull String input, @Nullable Map requestFilters, int start, - int count, - @Nonnull Authentication authentication, - @Nullable SearchFlags searchFlags) + int count) throws RemoteInvocationException { return ValidationUtils.validateSearchResult( entitySearchService.search( - List.of(entity), input, newFilter(requestFilters), null, start, count, searchFlags), + opContext, List.of(entity), input, newFilter(requestFilters), null, start, count), entityService); } @@ -358,18 +355,25 @@ public SearchResult search( * @return a set of list results * @throws RemoteInvocationException */ + @Override @Deprecated @Nonnull public ListResult list( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nullable Map requestFilters, int start, - int count, - @Nonnull final Authentication authentication) + int count) throws RemoteInvocationException { return ValidationUtils.validateListResult( toListResult( - entitySearchService.filter(entity, newFilter(requestFilters), null, start, count)), + entitySearchService.filter( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + entity, + newFilter(requestFilters), + null, + start, + count)), entityService); } @@ -387,34 +391,33 @@ public ListResult list( @Nonnull @Override public SearchResult search( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nonnull String input, @Nullable Filter filter, @Nullable SortCriterion sortCriterion, int start, - int count, - @Nonnull Authentication authentication, - @Nullable SearchFlags searchFlags) + int count) throws RemoteInvocationException { return ValidationUtils.validateSearchResult( entitySearchService.search( - List.of(entity), input, filter, sortCriterion, start, count, searchFlags), + opContext, List.of(entity), input, filter, sortCriterion, start, count), entityService); } + @Override @Nonnull public SearchResult searchAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter filter, int start, int count, - @Nullable SearchFlags searchFlags, - @Nullable SortCriterion sortCriterion, - @Nonnull final Authentication authentication) + @Nullable SortCriterion sortCriterion) throws RemoteInvocationException { return searchAcrossEntities( - entities, input, filter, start, count, searchFlags, sortCriterion, authentication, null); + opContext, entities, input, filter, start, count, sortCriterion, null); } /** @@ -430,49 +433,61 @@ public SearchResult searchAcrossEntities( * @return Snapshot key * @throws RemoteInvocationException */ + @Override @Nonnull public SearchResult searchAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter filter, int start, int count, - @Nullable SearchFlags searchFlags, @Nullable SortCriterion sortCriterion, - @Nonnull final Authentication authentication, @Nullable List facets) throws RemoteInvocationException { - final SearchFlags finalFlags = - searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true); + return ValidationUtils.validateSearchResult( searchService.searchAcrossEntities( - entities, input, filter, sortCriterion, start, count, finalFlags, facets), + opContext.withSearchFlags(flags -> flags.setFulltext(true)), + entities, + input, + filter, + sortCriterion, + start, + count, + facets), entityService); } @Nonnull @Override public ScrollResult scrollAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter filter, @Nullable String scrollId, @Nullable String keepAlive, - int count, - @Nullable SearchFlags searchFlags, - @Nonnull Authentication authentication) + int count) throws RemoteInvocationException { - final SearchFlags finalFlags = - searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true); + return ValidationUtils.validateScrollResult( searchService.scrollAcrossEntities( - entities, input, filter, null, scrollId, keepAlive, count, finalFlags), + opContext.withSearchFlags(flags -> flags.setFulltext(true)), + entities, + input, + filter, + null, + scrollId, + keepAlive, + count), entityService); } @Nonnull @Override public LineageSearchResult searchAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -481,12 +496,11 @@ public LineageSearchResult searchAcrossLineage( @Nullable Filter filter, @Nullable SortCriterion sortCriterion, int start, - int count, - @Nullable SearchFlags searchFlags, - @Nonnull final Authentication authentication) + int count) throws RemoteInvocationException { return ValidationUtils.validateLineageSearchResult( lineageSearchService.searchAcrossLineage( + opContext, sourceUrn, direction, entities, @@ -497,14 +511,14 @@ public LineageSearchResult searchAcrossLineage( start, count, null, - null, - searchFlags), + null), entityService); } @Nonnull @Override public LineageSearchResult searchAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -515,12 +529,11 @@ public LineageSearchResult searchAcrossLineage( int start, int count, @Nullable Long startTimeMillis, - @Nullable Long endTimeMillis, - @Nullable SearchFlags searchFlags, - @Nonnull final Authentication authentication) + @Nullable Long endTimeMillis) throws RemoteInvocationException { return ValidationUtils.validateLineageSearchResult( lineageSearchService.searchAcrossLineage( + opContext, sourceUrn, direction, entities, @@ -531,14 +544,14 @@ public LineageSearchResult searchAcrossLineage( start, count, startTimeMillis, - endTimeMillis, - searchFlags), + endTimeMillis), entityService); } @Nonnull @Override public LineageScrollResult scrollAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -550,14 +563,12 @@ public LineageScrollResult scrollAcrossLineage( @Nonnull String keepAlive, int count, @Nullable Long startTimeMillis, - @Nullable Long endTimeMillis, - @Nullable SearchFlags searchFlags, - @Nonnull final Authentication authentication) + @Nullable Long endTimeMillis) throws RemoteInvocationException { - final SearchFlags finalFlags = - searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true).setSkipCache(true); + return ValidationUtils.validateLineageScrollResult( lineageSearchService.scrollAcrossLineage( + opContext.withSearchFlags(flags -> flags.setFulltext(true).setSkipCache(true)), sourceUrn, direction, entities, @@ -569,8 +580,7 @@ public LineageScrollResult scrollAcrossLineage( keepAlive, count, startTimeMillis, - endTimeMillis, - finalFlags), + endTimeMillis), entityService); } @@ -581,35 +591,40 @@ public LineageScrollResult scrollAcrossLineage( * @return list of paths given urn * @throws RemoteInvocationException */ + @Override @Nonnull public StringArray getBrowsePaths(@Nonnull Urn urn, @Nonnull final Authentication authentication) throws RemoteInvocationException { return new StringArray(entitySearchService.getBrowsePaths(urn.getEntityType(), urn)); } + @Override public void setWritable(boolean canWrite, @Nonnull final Authentication authentication) throws RemoteInvocationException { entityService.setWritable(canWrite); } + @Override @Nonnull public Map batchGetTotalEntityCount( - @Nonnull List entityNames, @Nonnull final Authentication authentication) + @Nonnull OperationContext opContext, @Nonnull List entityNames) throws RemoteInvocationException { - return searchService.docCountPerEntity(entityNames); + return searchService.docCountPerEntity(opContext, entityNames); } /** List all urns existing for a particular Entity type. */ + @Override public ListUrnsResult listUrns( @Nonnull final String entityName, final int start, final int count, - @Nonnull final Authentication authentication) + @Nonnull Authentication authentication) throws RemoteInvocationException { return entityService.listUrns(entityName, start, count); } /** Hard delete an entity with a particular urn. */ + @Override public void deleteEntity(@Nonnull final Urn urn, @Nonnull final Authentication authentication) throws RemoteInvocationException { entityService.deleteUrn(urn); @@ -624,15 +639,22 @@ public void deleteEntityReferences(@Nonnull Urn urn, @Nonnull Authentication aut @Nonnull @Override public SearchResult filter( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nonnull Filter filter, @Nullable SortCriterion sortCriterion, int start, - int count, - @Nonnull final Authentication authentication) + int count) throws RemoteInvocationException { return ValidationUtils.validateSearchResult( - entitySearchService.filter(entity, filter, sortCriterion, start, count), entityService); + entitySearchService.filter( + opContext.withSearchFlags(flags -> flags.setFulltext(true)), + entity, + filter, + sortCriterion, + start, + count), + entityService); } @Override diff --git a/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java b/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java index fa020903c34f0a..1f0a43821a7f9b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java @@ -1,6 +1,5 @@ package com.linkedin.metadata.client; -import com.datahub.authentication.Authentication; import com.linkedin.entity.client.EntityClientCache; import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.metadata.config.cache.client.EntityClientCacheConfig; @@ -13,17 +12,21 @@ import com.linkedin.metadata.search.client.CachingEntitySearchService; import com.linkedin.metadata.service.RollbackService; import com.linkedin.metadata.timeseries.TimeseriesAspectService; -import javax.annotation.Nonnull; +import io.datahubproject.metadata.context.OperationContext; +import java.util.concurrent.ConcurrentHashMap; import lombok.Getter; +import org.springframework.beans.factory.annotation.Qualifier; /** Java backed SystemEntityClient */ @Getter public class SystemJavaEntityClient extends JavaEntityClient implements SystemEntityClient { private final EntityClientCache entityClientCache; - private final Authentication systemAuthentication; + private final OperationContext systemOperationContext; + private final ConcurrentHashMap operationContextMap; public SystemJavaEntityClient( + @Qualifier("systemOperationContext") OperationContext systemOperationContext, EntityService entityService, DeleteEntityService deleteEntityService, EntitySearchService entitySearchService, @@ -33,9 +36,9 @@ public SystemJavaEntityClient( TimeseriesAspectService timeseriesAspectService, RollbackService rollbackService, EventProducer eventProducer, - @Nonnull Authentication systemAuthentication, EntityClientCacheConfig cacheConfig) { super( + systemOperationContext, entityService, deleteEntityService, entitySearchService, @@ -45,8 +48,8 @@ public SystemJavaEntityClient( timeseriesAspectService, rollbackService, eventProducer); - this.systemAuthentication = systemAuthentication; - this.entityClientCache = - buildEntityClientCache(SystemJavaEntityClient.class, systemAuthentication, cacheConfig); + this.operationContextMap = new ConcurrentHashMap<>(); + this.systemOperationContext = systemOperationContext; + this.entityClientCache = buildEntityClientCache(SystemJavaEntityClient.class, cacheConfig); } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java index 3edb55f265dc13..aa60e7e5286739 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java @@ -63,6 +63,9 @@ public Pair>, List> toUpsertBatchItems( upsertItem = patchBatchItem.applyPatch(currentValue, aspectRetriever); } + // Populate old aspect for write hooks + upsertItem.setPreviousSystemAspect(latest); + return upsertItem; }) .collect(Collectors.toCollection(LinkedList::new)); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/MostPopularSource.java b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/MostPopularSource.java index f5c783014caa13..59109c8c3de643 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/MostPopularSource.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/MostPopularSource.java @@ -1,9 +1,9 @@ package com.linkedin.metadata.recommendation.candidatesource; import com.codahale.metrics.Timer; +import com.datahub.authorization.config.SearchAuthorizationConfiguration; import com.datahub.util.exception.ESQueryException; import com.google.common.collect.ImmutableSet; -import com.linkedin.common.urn.Urn; import com.linkedin.metadata.Constants; import com.linkedin.metadata.datahubusage.DataHubUsageEventConstants; import com.linkedin.metadata.datahubusage.DataHubUsageEventType; @@ -15,9 +15,11 @@ import com.linkedin.metadata.search.utils.ESUtils; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.metrics.MetricUtils; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.io.IOException; import java.util.List; +import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; import javax.annotation.Nonnull; @@ -29,6 +31,7 @@ import org.opensearch.client.RestHighLevelClient; import org.opensearch.client.indices.GetIndexRequest; import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.aggregations.AggregationBuilders; @@ -78,7 +81,7 @@ public RecommendationRenderType getRenderType() { @Override public boolean isEligible( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { boolean analyticsEnabled = false; try { analyticsEnabled = @@ -96,8 +99,8 @@ public boolean isEligible( @Override @WithSpan public List getRecommendations( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { - SearchRequest searchRequest = buildSearchRequest(userUrn); + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { + SearchRequest searchRequest = buildSearchRequest(opContext); try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getMostPopular").time()) { final SearchResponse searchResponse = _searchClient.search(searchRequest, RequestOptions.DEFAULT); @@ -121,11 +124,15 @@ public Set getSupportedEntityTypes() { return SUPPORTED_ENTITY_TYPES; } - private SearchRequest buildSearchRequest(@Nonnull Urn userUrn) { + private SearchRequest buildSearchRequest(@Nonnull OperationContext opContext) { // TODO: Proactively filter for entity types in the supported set. SearchRequest request = new SearchRequest(); SearchSourceBuilder source = new SearchSourceBuilder(); BoolQueryBuilder query = QueryBuilders.boolQuery(); + + // Potentially limit actors + restrictPeers(opContext).ifPresent(query::must); + // Filter for all entity view events query.must( QueryBuilders.termQuery( @@ -144,4 +151,23 @@ private SearchRequest buildSearchRequest(@Nonnull Urn userUrn) { request.indices(_indexConvention.getIndexName(DATAHUB_USAGE_INDEX)); return request; } + + // If search access controls enabled, restrict user activity to peers + private static Optional restrictPeers(@Nonnull OperationContext opContext) { + SearchAuthorizationConfiguration config = + opContext.getOperationContextConfig().getSearchAuthorizationConfiguration(); + + if (config.isEnabled() + && config.getRecommendations().isPeerGroupEnabled() + && !opContext.isSystemAuth()) { + return Optional.of( + QueryBuilders.termsQuery( + DataHubUsageEventConstants.ACTOR_URN + ".keyword", + opContext.getActorPeers().stream() + .map(Object::toString) + .collect(Collectors.toList()))); + } + + return Optional.empty(); + } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyEditedSource.java b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyEditedSource.java index 127b0f5c342c70..a2ba8cbcbcc61d 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyEditedSource.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyEditedSource.java @@ -15,6 +15,7 @@ import com.linkedin.metadata.search.utils.ESUtils; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.metrics.MetricUtils; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.io.IOException; import java.util.List; @@ -79,7 +80,7 @@ public RecommendationRenderType getRenderType() { @Override public boolean isEligible( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { boolean analyticsEnabled = false; try { analyticsEnabled = @@ -97,8 +98,8 @@ public boolean isEligible( @Override @WithSpan public List getRecommendations( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { - SearchRequest searchRequest = buildSearchRequest(userUrn); + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { + SearchRequest searchRequest = buildSearchRequest(opContext.getActorContext().getActorUrn()); try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getRecentlyEdited").time()) { final SearchResponse searchResponse = _searchClient.search(searchRequest, RequestOptions.DEFAULT); @@ -127,6 +128,11 @@ private SearchRequest buildSearchRequest(@Nonnull Urn userUrn) { SearchRequest request = new SearchRequest(); SearchSourceBuilder source = new SearchSourceBuilder(); BoolQueryBuilder query = QueryBuilders.boolQuery(); + // Filter for the entity edit events of the user requesting recommendation + query.must( + QueryBuilders.termQuery( + ESUtils.toKeywordField(DataHubUsageEventConstants.ACTOR_URN, false), + userUrn.toString())); // Filter for the entity action events query.must( QueryBuilders.termQuery( diff --git a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyViewedSource.java b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyViewedSource.java index 0ab5cf40cf4e5a..d6bf7d94aecb67 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyViewedSource.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyViewedSource.java @@ -15,6 +15,7 @@ import com.linkedin.metadata.search.utils.ESUtils; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.metrics.MetricUtils; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.io.IOException; import java.util.List; @@ -79,7 +80,7 @@ public RecommendationRenderType getRenderType() { @Override public boolean isEligible( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { boolean analyticsEnabled = false; try { analyticsEnabled = @@ -97,8 +98,8 @@ public boolean isEligible( @Override @WithSpan public List getRecommendations( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { - SearchRequest searchRequest = buildSearchRequest(userUrn); + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { + SearchRequest searchRequest = buildSearchRequest(opContext.getActorContext().getActorUrn()); try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getRecentlyViewed").time()) { final SearchResponse searchResponse = _searchClient.search(searchRequest, RequestOptions.DEFAULT); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/EntityLineageResultCacheKey.java b/metadata-io/src/main/java/com/linkedin/metadata/search/EntityLineageResultCacheKey.java index b862de320db363..3df3ad3945f59e 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/EntityLineageResultCacheKey.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/EntityLineageResultCacheKey.java @@ -4,10 +4,12 @@ import com.linkedin.metadata.graph.LineageDirection; import java.time.Instant; import java.time.temporal.TemporalUnit; +import javax.annotation.Nonnull; import lombok.Data; @Data public class EntityLineageResultCacheKey { + private final String contextId; private final Urn sourceUrn; private final LineageDirection direction; private final Long startTimeMillis; @@ -15,13 +17,14 @@ public class EntityLineageResultCacheKey { private final Integer maxHops; public EntityLineageResultCacheKey( + @Nonnull String contextId, Urn sourceUrn, LineageDirection direction, Long startTimeMillis, Long endTimeMillis, Integer maxHops, TemporalUnit resolution) { - + this.contextId = contextId; this.sourceUrn = sourceUrn; this.direction = direction; this.maxHops = maxHops; diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/LineageSearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/LineageSearchService.java index cf9279414a394d..622f92b0bd7e93 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/LineageSearchService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/LineageSearchService.java @@ -32,6 +32,7 @@ import com.linkedin.metadata.search.utils.FilterUtils; import com.linkedin.metadata.search.utils.QueryUtils; import com.linkedin.metadata.search.utils.SearchUtils; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.time.temporal.ChronoUnit; @@ -66,6 +67,7 @@ public class LineageSearchService { .setSkipCache(false) .setSkipAggregates(false) .setSkipHighlighting(true) + .setIncludeRestricted(false) .setGroupingSpec( new GroupingSpec() .setGroupingCriteria( @@ -79,7 +81,6 @@ public class LineageSearchService { @Nullable private final Cache cache; private final boolean cacheEnabled; private final SearchLineageCacheConfiguration cacheConfiguration; - private final ExecutorService cacheRefillExecutor = Executors.newFixedThreadPool(1); private static final String DEGREE_FILTER = "degree"; @@ -125,6 +126,7 @@ public class LineageSearchService { @Nonnull @WithSpan public LineageSearchResult searchAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -135,25 +137,30 @@ public LineageSearchResult searchAcrossLineage( int from, int size, @Nullable Long startTimeMillis, - @Nullable Long endTimeMillis, - @Nullable SearchFlags searchFlags) { - - final SearchFlags finalFlags = - applyDefaultSearchFlags(searchFlags, input, DEFAULT_SERVICE_SEARCH_FLAGS); + @Nullable Long endTimeMillis) { long startTime = System.nanoTime(); - log.debug("Cache enabled {}, Input :{}:", cacheEnabled, input); - if ((input == null) || (input.isEmpty())) { - input = "*"; - } + final String finalInput = input == null || input.isEmpty() ? "*" : input; + + log.debug("Cache enabled {}, Input :{}:", cacheEnabled, finalInput); if (maxHops == null) { maxHops = 1000; } + final OperationContext finalOpContext = + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, finalInput, DEFAULT_SERVICE_SEARCH_FLAGS)); + // Cache multihop result for faster performance final EntityLineageResultCacheKey cacheKey = new EntityLineageResultCacheKey( - sourceUrn, direction, startTimeMillis, endTimeMillis, maxHops, ChronoUnit.DAYS); + finalOpContext.getSearchContextId(), + sourceUrn, + direction, + startTimeMillis, + endTimeMillis, + maxHops, + ChronoUnit.DAYS); CachedEntityLineageResult cachedLineageResult = null; if (cacheEnabled) { @@ -166,7 +173,8 @@ public LineageSearchResult searchAcrossLineage( EntityLineageResult lineageResult; FreshnessStats freshnessStats = new FreshnessStats().setCached(Boolean.FALSE); - if (cachedLineageResult == null || finalFlags.isSkipCache()) { + if (cachedLineageResult == null + || finalOpContext.getSearchContext().getSearchFlags().isSkipCache()) { lineageResult = _graphService.getLineage( sourceUrn, direction, 0, MAX_RELATIONSHIPS, maxHops, startTimeMillis, endTimeMillis); @@ -218,7 +226,8 @@ public LineageSearchResult searchAcrossLineage( } } - if (SearchUtils.convertSchemaFieldToDataset(searchFlags)) { + if (SearchUtils.convertSchemaFieldToDataset( + finalOpContext.getSearchContext().getSearchFlags())) { // set schemaField relationship entity to be its reference urn LineageRelationshipArray updatedRelationships = convertSchemaFieldRelationships(lineageResult); @@ -243,7 +252,7 @@ public LineageSearchResult searchAcrossLineage( SearchUtils.removeCriteria( inputFilters, criterion -> criterion.getField().equals(DEGREE_FILTER_INPUT)); - if (canDoLightning(lineageRelationships, input, reducedFilters, sortCriterion)) { + if (canDoLightning(lineageRelationships, finalInput, reducedFilters, sortCriterion)) { codePath = "lightning"; // use lightning approach to return lineage search results LineageSearchResult lineageSearchResult = @@ -260,7 +269,13 @@ public LineageSearchResult searchAcrossLineage( codePath = "tortoise"; LineageSearchResult lineageSearchResult = getSearchResultInBatches( - lineageRelationships, input, reducedFilters, sortCriterion, from, size, finalFlags); + finalOpContext, + lineageRelationships, + finalInput, + reducedFilters, + sortCriterion, + from, + size); if (!lineageSearchResult.getEntities().isEmpty()) { log.debug( "Lineage entity results number -> {}; first -> {}", @@ -311,7 +326,7 @@ LineageSearchResult getLightningSearchResult( int size, Set entityNames) { - // Contruct result objects + // Construct result objects LineageSearchResult finalResult = new LineageSearchResult().setMetadata(new SearchResultMetadata()); LineageSearchEntityArray lineageSearchEntityArray = new LineageSearchEntityArray(); @@ -506,16 +521,13 @@ private Map generateUrnToRelationshipMap( // Search service can only take up to 50K term filter, so query search service in batches private LineageSearchResult getSearchResultInBatches( + @Nonnull OperationContext opContext, List lineageRelationships, @Nonnull String input, @Nullable Filter inputFilters, @Nullable SortCriterion sortCriterion, int from, - int size, - @Nonnull SearchFlags searchFlags) { - - final SearchFlags finalFlags = - applyDefaultSearchFlags(searchFlags, input, DEFAULT_SERVICE_SEARCH_FLAGS); + int size) { LineageSearchResult finalResult = new LineageSearchResult() @@ -540,13 +552,14 @@ private LineageSearchResult getSearchResultInBatches( LineageSearchResult resultForBatch = buildLineageSearchResult( _searchService.searchAcrossEntities( + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, input, DEFAULT_SERVICE_SEARCH_FLAGS)), entitiesToQuery, input, finalFilter, sortCriterion, queryFrom, - querySize, - finalFlags), + querySize), urnToRelationship); queryFrom = Math.max(0, from - resultForBatch.getNumEntities()); querySize = Math.max(0, size - resultForBatch.getEntities().size()); @@ -717,6 +730,7 @@ private LineageSearchEntity buildLineageSearchEntity( @Nonnull @WithSpan public LineageScrollResult scrollAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -728,12 +742,17 @@ public LineageScrollResult scrollAcrossLineage( @Nonnull String keepAlive, int size, @Nullable Long startTimeMillis, - @Nullable Long endTimeMillis, - @Nonnull SearchFlags searchFlags) { + @Nullable Long endTimeMillis) { // Cache multihop result for faster performance final EntityLineageResultCacheKey cacheKey = new EntityLineageResultCacheKey( - sourceUrn, direction, startTimeMillis, endTimeMillis, maxHops, ChronoUnit.DAYS); + opContext.getSearchContextId(), + sourceUrn, + direction, + startTimeMillis, + endTimeMillis, + maxHops, + ChronoUnit.DAYS); CachedEntityLineageResult cachedLineageResult = cacheEnabled ? cache.get(cacheKey, CachedEntityLineageResult.class) : null; EntityLineageResult lineageResult; @@ -767,28 +786,31 @@ public LineageScrollResult scrollAcrossLineage( SearchUtils.removeCriteria( inputFilters, criterion -> criterion.getField().equals(DEGREE_FILTER_INPUT)); return getScrollResultInBatches( + opContext, lineageRelationships, input != null ? input : "*", reducedFilters, sortCriterion, scrollId, keepAlive, - size, - searchFlags); + size); } // Search service can only take up to 50K term filter, so query search service in batches private LineageScrollResult getScrollResultInBatches( + @Nonnull OperationContext opContext, List lineageRelationships, @Nonnull String input, @Nullable Filter inputFilters, @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nonnull String keepAlive, - int size, - @Nonnull SearchFlags searchFlags) { - final SearchFlags finalFlags = - applyDefaultSearchFlags(searchFlags, input, DEFAULT_SERVICE_SEARCH_FLAGS); + int size) { + + OperationContext finalOpContext = + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, input, DEFAULT_SERVICE_SEARCH_FLAGS)); + LineageScrollResult finalResult = new LineageScrollResult() .setEntities(new LineageSearchEntityArray(Collections.emptyList())) @@ -810,14 +832,14 @@ private LineageScrollResult getScrollResultInBatches( LineageScrollResult resultForBatch = buildLineageScrollResult( _searchService.scrollAcrossEntities( + finalOpContext, entitiesToQuery, input, finalFilter, sortCriterion, scrollId, keepAlive, - querySize, - finalFlags), + querySize), urnToRelationship); querySize = Math.max(0, size - resultForBatch.getEntities().size()); finalResult = mergeScrollResult(finalResult, resultForBatch); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java index 3bcc163613c5ea..6e5bd631031901 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java @@ -4,7 +4,6 @@ import com.codahale.metrics.Timer; import com.linkedin.data.template.LongMap; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.search.cache.EntityDocCountCache; @@ -12,6 +11,7 @@ import com.linkedin.metadata.search.ranker.SearchRanker; import com.linkedin.metadata.utils.SearchUtil; import com.linkedin.metadata.utils.metrics.MetricUtils; +import io.datahubproject.metadata.context.OperationContext; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -38,14 +38,15 @@ public SearchService( _entityDocCountCache = entityDocCountCache; } - public Map docCountPerEntity(@Nonnull List entityNames) { + public Map docCountPerEntity( + @Nonnull OperationContext opContext, @Nonnull List entityNames) { return entityNames.stream() .collect( Collectors.toMap( Function.identity(), entityName -> _entityDocCountCache - .getEntityDocCount() + .getEntityDocCount(opContext) .getOrDefault(entityName.toLowerCase(), 0L))); } @@ -60,27 +61,26 @@ public Map docCountPerEntity(@Nonnull List entityNames) { * @param sortCriterion {@link SortCriterion} to be applied to search results * @param from index to start the search from * @param size the number of search hits to return - * @param searchFlags optional set of flags to control search behavior * @return a {@link SearchResult} that contains a list of matched documents and related search * result metadata */ @Nonnull public SearchResult search( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, - int size, - @Nullable SearchFlags searchFlags) { - List entitiesToSearch = getEntitiesToSearch(entityNames); + int size) { + List entitiesToSearch = getEntitiesToSearch(opContext, entityNames); if (entitiesToSearch.isEmpty()) { // Optimization: If the indices are all empty, return empty result return getEmptySearchResult(from, size); } SearchResult result = _cachingEntitySearchService.search( - entitiesToSearch, input, postFilters, sortCriterion, from, size, searchFlags, null); + opContext, entitiesToSearch, input, postFilters, sortCriterion, from, size, null); try { return result @@ -94,15 +94,15 @@ public SearchResult search( @Nonnull public SearchResult searchAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, - int size, - @Nullable SearchFlags searchFlags) { + int size) { return searchAcrossEntities( - entities, input, postFilters, sortCriterion, from, size, searchFlags, null); + opContext, entities, input, postFilters, sortCriterion, from, size, null); } /** @@ -116,20 +116,19 @@ public SearchResult searchAcrossEntities( * @param sortCriterion {@link SortCriterion} to be applied to search results * @param from index to start the search from * @param size the number of search hits to return - * @param searchFlags optional set of flags to control search behavior * @param facets list of facets we want aggregations for * @return a {@link SearchResult} that contains a list of matched documents and related search * result metadata */ @Nonnull public SearchResult searchAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, int size, - @Nullable SearchFlags searchFlags, @Nullable List facets) { log.debug( String.format( @@ -147,14 +146,14 @@ public SearchResult searchAcrossEntities( facets = new ArrayList<>(facets); facets.add(INDEX_VIRTUAL_FIELD); } - List nonEmptyEntities = getEntitiesToSearch(entities); + List nonEmptyEntities = getEntitiesToSearch(opContext, entities); if (nonEmptyEntities.isEmpty()) { // Optimization: If the indices are all empty, return empty result return getEmptySearchResult(from, size); } SearchResult result = _cachingEntitySearchService.search( - nonEmptyEntities, input, postFilters, sortCriterion, from, size, searchFlags, facets); + opContext, nonEmptyEntities, input, postFilters, sortCriterion, from, size, facets); if (facets == null || facets.contains("entity") || facets.contains("_entityType")) { Optional entityTypeAgg = result.getMetadata().getAggregations().stream() @@ -206,7 +205,8 @@ public SearchResult searchAcrossEntities( * @param inputEntities the requested entities * @return some entities to search */ - private List getEntitiesToSearch(@Nonnull List inputEntities) { + private List getEntitiesToSearch( + @Nonnull OperationContext opContext, @Nonnull List inputEntities) { List nonEmptyEntities; List lowercaseEntities = inputEntities.stream().map(String::toLowerCase).collect(Collectors.toList()); @@ -214,7 +214,7 @@ private List getEntitiesToSearch(@Nonnull List inputEntities) { if (lowercaseEntities.isEmpty()) { try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getNonEmptyEntities").time()) { - nonEmptyEntities = _entityDocCountCache.getNonEmptyEntities(); + nonEmptyEntities = _entityDocCountCache.getNonEmptyEntities(opContext); } } else { nonEmptyEntities = lowercaseEntities; @@ -234,38 +234,30 @@ private List getEntitiesToSearch(@Nonnull List inputEntities) { * @param sortCriterion {@link SortCriterion} to be applied to search results * @param scrollId opaque scroll identifier for passing to search backend * @param size the number of search hits to return - * @param searchFlags optional set of flags to control search behavior * @return a {@link ScrollResult} that contains a list of matched documents and related search * result metadata */ @Nonnull public ScrollResult scrollAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, - int size, - @Nullable SearchFlags searchFlags) { + int size) { log.debug( String.format( "Searching Search documents entities: %s, input: %s, postFilters: %s, sortCriterion: %s, from: %s, size: %s", entities, input, postFilters, sortCriterion, scrollId, size)); - List entitiesToSearch = getEntitiesToSearch(entities); + List entitiesToSearch = getEntitiesToSearch(opContext, entities); if (entitiesToSearch.isEmpty()) { // No indices with non-zero entries: skip querying and return empty result return getEmptyScrollResult(size); } return _cachingEntitySearchService.scroll( - entitiesToSearch, - input, - postFilters, - sortCriterion, - scrollId, - keepAlive, - size, - searchFlags); + opContext, entitiesToSearch, input, postFilters, sortCriterion, scrollId, keepAlive, size); } private static SearchResult getEmptySearchResult(int from, int size) { diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/CacheableSearcher.java b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/CacheableSearcher.java index 0ecdb83ed20ee1..28efa29c9fffa2 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/CacheableSearcher.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/CacheableSearcher.java @@ -3,17 +3,16 @@ import static com.datahub.util.RecordUtils.*; import com.codahale.metrics.Timer; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.utils.metrics.MetricUtils; +import io.datahubproject.metadata.context.OperationContext; import java.io.Serializable; import java.util.ArrayList; import java.util.List; import java.util.function.Function; import javax.annotation.Nonnull; -import javax.annotation.Nullable; import lombok.RequiredArgsConstructor; import lombok.Value; import org.springframework.cache.Cache; @@ -28,7 +27,6 @@ public class CacheableSearcher { private final Function searcher; // Function that generates the cache key given the query batch (from, size) private final Function cacheKeyGenerator; - @Nullable private final SearchFlags searchFlags; private final boolean enableCache; @Value @@ -43,7 +41,7 @@ public static class QueryPagination implements Serializable { * that return a variable number of results (we have no idea which batch the "from" "size" page * corresponds to) */ - public SearchResult getSearchResults(int from, int size) { + public SearchResult getSearchResults(@Nonnull OperationContext opContext, int from, int size) { try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getSearchResults").time()) { int resultsSoFar = 0; int batchId = 0; @@ -52,7 +50,7 @@ public SearchResult getSearchResults(int from, int size) { SearchResult batchedResult; // Use do-while to make sure we run at least one batch to fetch metadata do { - batchedResult = getBatch(batchId); + batchedResult = getBatch(opContext, batchId); int currentBatchSize = batchedResult.getEntities().size(); // If the number of results in this batch is 0, no need to continue if (currentBatchSize == 0) { @@ -85,13 +83,14 @@ private QueryPagination getBatchQuerySize(int batchId) { return new QueryPagination(batchId * batchSize, batchSize); } - private SearchResult getBatch(int batchId) { + private SearchResult getBatch(@Nonnull OperationContext opContext, int batchId) { try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getBatch").time()) { QueryPagination batch = getBatchQuerySize(batchId); SearchResult result; if (enableCache) { K cacheKey = cacheKeyGenerator.apply(batch); - if ((searchFlags == null || !searchFlags.isSkipCache())) { + if ((opContext.getSearchContext().getSearchFlags().isSkipCache() == null + || !opContext.getSearchContext().getSearchFlags().isSkipCache())) { try (Timer.Context ignored2 = MetricUtils.timer(this.getClass(), "getBatch_cache").time()) { Timer.Context cacheAccess = diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntityDocCountCache.java b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntityDocCountCache.java index 2c99c71acf7493..745ef6686d3204 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntityDocCountCache.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntityDocCountCache.java @@ -5,46 +5,56 @@ import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.utils.ConcurrencyUtils; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.util.List; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; import java.util.function.Function; import java.util.function.Supplier; import java.util.stream.Collectors; +import javax.annotation.Nonnull; public class EntityDocCountCache { - private final EntityRegistry _entityRegistry; - private final EntitySearchService _entitySearchService; - private final Supplier> entityDocCount; + private final EntityRegistry entityRegistry; + private final EntitySearchService entitySearchService; + private final EntityDocCountCacheConfiguration config; + private final Map>> entityDocCounts; public EntityDocCountCache( EntityRegistry entityRegistry, EntitySearchService entitySearchService, EntityDocCountCacheConfiguration config) { - _entityRegistry = entityRegistry; - _entitySearchService = entitySearchService; - entityDocCount = - Suppliers.memoizeWithExpiration( - this::fetchEntityDocCount, config.getTtlSeconds(), TimeUnit.SECONDS); + this.config = config; + this.entityRegistry = entityRegistry; + this.entitySearchService = entitySearchService; + this.entityDocCounts = new ConcurrentHashMap<>(); } - private Map fetchEntityDocCount() { + private Map fetchEntityDocCount(@Nonnull OperationContext opContext) { return ConcurrencyUtils.transformAndCollectAsync( - _entityRegistry.getEntitySpecs().keySet(), + entityRegistry.getEntitySpecs().keySet(), Function.identity(), - Collectors.toMap(Function.identity(), _entitySearchService::docCount)); + Collectors.toMap(Function.identity(), v -> entitySearchService.docCount(opContext, v))); } @WithSpan - public Map getEntityDocCount() { - return entityDocCount.get(); + public Map getEntityDocCount(@Nonnull OperationContext opContext) { + return entityDocCounts + .computeIfAbsent(opContext.getSearchContextId(), k -> buildSupplier(opContext)) + .get(); } - public List getNonEmptyEntities() { - return getEntityDocCount().entrySet().stream() + public List getNonEmptyEntities(@Nonnull OperationContext opContext) { + return getEntityDocCount(opContext).entrySet().stream() .filter(entry -> entry.getValue() > 0) .map(Map.Entry::getKey) .collect(Collectors.toList()); } + + private Supplier> buildSupplier(@Nonnull OperationContext opContext) { + return Suppliers.memoizeWithExpiration( + () -> fetchEntityDocCount(opContext), config.getTtlSeconds(), TimeUnit.SECONDS); + } } 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 eaeae0cfc15569..5db427fa901488 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 @@ -14,6 +14,7 @@ import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.search.cache.CacheableSearcher; import com.linkedin.metadata.utils.metrics.MetricUtils; +import io.datahubproject.metadata.context.OperationContext; import java.util.List; import java.util.Optional; import javax.annotation.Nonnull; @@ -42,52 +43,53 @@ public class CachingEntitySearchService { * Retrieves cached search results. If the query has been cached, this will return quickly. If * not, a full search request will be made. * - * @param entityName the name of the entity to search + * @param opContext the operation's context + * @param entityNames the names of the entity to search * @param query the search query * @param filters the filters to include * @param sortCriterion the sort criterion * @param from the start offset * @param size the count - * @param flags additional search flags * @param facets list of facets we want aggregations for * @return a {@link SearchResult} containing the requested batch of search results */ public SearchResult search( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String query, @Nullable Filter filters, @Nullable SortCriterion sortCriterion, int from, int size, - @Nullable SearchFlags flags, @Nullable List facets) { return getCachedSearchResults( - entityNames, query, filters, sortCriterion, from, size, flags, facets); + opContext, entityNames, query, filters, sortCriterion, from, size, facets); } /** * Retrieves cached auto complete results * + * @param opContext the operation's context * @param entityName the name of the entity to search * @param input the input query * @param filters the filters to include * @param limit the max number of results to return - * @param flags additional search flags * @return a {@link SearchResult} containing the requested batch of search results */ public AutoCompleteResult autoComplete( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String input, @Nullable String field, @Nullable Filter filters, - int limit, - @Nullable SearchFlags flags) { - return getCachedAutoCompleteResults(entityName, input, field, filters, limit, flags); + int limit) { + return getCachedAutoCompleteResults(opContext, entityName, input, field, filters, limit); } /** * Retrieves cached auto complete results * + * @param opContext the operation's context * @param entityName type of entity to query * @param path the path to be browsed * @param filters the request map with fields and values as filters @@ -96,19 +98,20 @@ public AutoCompleteResult autoComplete( * @return a {@link SearchResult} containing the requested batch of search results */ public BrowseResult browse( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filters, int from, - int size, - @Nullable SearchFlags flags) { - return getCachedBrowseResults(entityName, path, filters, from, size, flags); + int size) { + return getCachedBrowseResults(opContext, entityName, path, filters, from, size); } /** * Retrieves cached scroll results. If the query has been cached, this will return quickly. If * not, a full scroll request will be made. * + * @param opContext the operation's context * @param entities the names of the entities to search * @param query the search query * @param filters the filters to include @@ -116,20 +119,19 @@ public BrowseResult browse( * @param scrollId opaque scroll identifier for a scroll request * @param keepAlive the string representation of how long to keep point in time alive * @param size the count - * @param flags additional search flags * @return a {@link ScrollResult} containing the requested batch of scroll results */ public ScrollResult scroll( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String query, @Nullable Filter filters, @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, - int size, - @Nullable SearchFlags flags) { + int size) { return getCachedScrollResults( - entities, query, filters, sortCriterion, scrollId, keepAlive, size, flags); + opContext, entities, query, filters, sortCriterion, scrollId, keepAlive, size); } /** @@ -139,65 +141,64 @@ public ScrollResult scroll( * corresponds to) */ public SearchResult getCachedSearchResults( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String query, @Nullable Filter filters, @Nullable SortCriterion sortCriterion, int from, int size, - @Nullable SearchFlags flags, @Nullable List facets) { return new CacheableSearcher<>( cacheManager.getCache(ENTITY_SEARCH_SERVICE_SEARCH_CACHE_NAME), batchSize, querySize -> getRawSearchResults( + opContext, entityNames, query, filters, sortCriterion, querySize.getFrom(), querySize.getSize(), - flags, facets), querySize -> Septet.with( + opContext.getSearchContextId(), entityNames, query, filters != null ? toJsonString(filters) : null, sortCriterion != null ? toJsonString(sortCriterion) : null, - flags != null ? toJsonString(flags) : null, facets, querySize), - flags, enableCache) - .getSearchResults(from, size); + .getSearchResults(opContext, from, size); } /** Returns cached auto-complete results. */ public AutoCompleteResult getCachedAutoCompleteResults( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String input, @Nullable String field, @Nullable Filter filters, - int limit, - @Nullable SearchFlags flags) { + int limit) { try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getCachedAutoCompleteResults").time()) { Cache cache = cacheManager.getCache(ENTITY_SEARCH_SERVICE_AUTOCOMPLETE_CACHE_NAME); AutoCompleteResult result; - if (enableCache(flags)) { + if (enableCache(opContext.getSearchContext().getSearchFlags())) { try (Timer.Context ignored2 = MetricUtils.timer(this.getClass(), "getCachedAutoCompleteResults_cache").time()) { Timer.Context cacheAccess = MetricUtils.timer(this.getClass(), "autocomplete_cache_access").time(); Object cacheKey = Sextet.with( + opContext.getSearchContextId(), entityName, input, field, filters != null ? toJsonString(filters) : null, - flags != null ? toJsonString(flags) : null, limit); String json = cache.get(cacheKey, String.class); result = json != null ? toRecordTemplate(AutoCompleteResult.class, json) : null; @@ -205,14 +206,14 @@ public AutoCompleteResult getCachedAutoCompleteResults( if (result == null) { Timer.Context cacheMiss = MetricUtils.timer(this.getClass(), "autocomplete_cache_miss").time(); - result = getRawAutoCompleteResults(entityName, input, field, filters, limit); + result = getRawAutoCompleteResults(opContext, entityName, input, field, filters, limit); cache.put(cacheKey, toJsonString(result)); cacheMiss.stop(); MetricUtils.counter(this.getClass(), "autocomplete_cache_miss_count").inc(); } } } else { - result = getRawAutoCompleteResults(entityName, input, field, filters, limit); + result = getRawAutoCompleteResults(opContext, entityName, input, field, filters, limit); } return result; } @@ -220,27 +221,27 @@ public AutoCompleteResult getCachedAutoCompleteResults( /** Returns cached browse results. */ public BrowseResult getCachedBrowseResults( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filters, int from, - int size, - @Nullable SearchFlags flags) { + int size) { try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getCachedBrowseResults").time()) { Cache cache = cacheManager.getCache(ENTITY_SEARCH_SERVICE_BROWSE_CACHE_NAME); BrowseResult result; - if (enableCache(flags)) { + if (enableCache(opContext.getSearchContext().getSearchFlags())) { try (Timer.Context ignored2 = MetricUtils.timer(this.getClass(), "getCachedBrowseResults_cache").time()) { Timer.Context cacheAccess = MetricUtils.timer(this.getClass(), "browse_cache_access").time(); Object cacheKey = Sextet.with( + opContext.getSearchContextId(), entityName, path, filters != null ? toJsonString(filters) : null, - flags != null ? toJsonString(flags) : null, from, size); String json = cache.get(cacheKey, String.class); @@ -249,14 +250,14 @@ public BrowseResult getCachedBrowseResults( if (result == null) { Timer.Context cacheMiss = MetricUtils.timer(this.getClass(), "browse_cache_miss").time(); - result = getRawBrowseResults(entityName, path, filters, from, size); + result = getRawBrowseResults(opContext, entityName, path, filters, from, size); cache.put(cacheKey, toJsonString(result)); cacheMiss.stop(); MetricUtils.counter(this.getClass(), "browse_cache_miss_count").inc(); } } } else { - result = getRawBrowseResults(entityName, path, filters, from, size); + result = getRawBrowseResults(opContext, entityName, path, filters, from, size); } return result; } @@ -264,30 +265,33 @@ public BrowseResult getCachedBrowseResults( /** Returns cached scroll results. */ public ScrollResult getCachedScrollResults( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String query, @Nullable Filter filters, @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, - int size, - @Nullable SearchFlags flags) { + int size) { try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getCachedScrollResults").time()) { boolean isFullText = - Boolean.TRUE.equals(Optional.ofNullable(flags).orElse(new SearchFlags()).isFulltext()); + Boolean.TRUE.equals( + Optional.ofNullable(opContext.getSearchContext().getSearchFlags()) + .orElse(new SearchFlags()) + .isFulltext()); Cache cache = cacheManager.getCache(ENTITY_SEARCH_SERVICE_SCROLL_CACHE_NAME); ScrollResult result; - if (enableCache(flags)) { + if (enableCache(opContext.getSearchContext().getSearchFlags())) { Timer.Context cacheAccess = MetricUtils.timer(this.getClass(), "scroll_cache_access").time(); Object cacheKey = Septet.with( + opContext.getSearchContextId(), entities, query, filters != null ? toJsonString(filters) : null, sortCriterion != null ? toJsonString(sortCriterion) : null, - flags != null ? toJsonString(flags) : null, scrollId, size); String json = cache.get(cacheKey, String.class); @@ -297,6 +301,7 @@ public ScrollResult getCachedScrollResults( Timer.Context cacheMiss = MetricUtils.timer(this.getClass(), "scroll_cache_miss").time(); result = getRawScrollResults( + opContext, entities, query, filters, @@ -304,8 +309,7 @@ public ScrollResult getCachedScrollResults( scrollId, keepAlive, size, - isFullText, - flags); + isFullText); cache.put(cacheKey, toJsonString(result)); cacheMiss.stop(); MetricUtils.counter(this.getClass(), "scroll_cache_miss_count").inc(); @@ -313,6 +317,7 @@ public ScrollResult getCachedScrollResults( } else { result = getRawScrollResults( + opContext, entities, query, filters, @@ -320,8 +325,7 @@ public ScrollResult getCachedScrollResults( scrollId, keepAlive, size, - isFullText, - flags); + isFullText); } return result; } @@ -329,40 +333,43 @@ public ScrollResult getCachedScrollResults( /** Executes the expensive search query using the {@link EntitySearchService} */ private SearchResult getRawSearchResults( + @Nonnull OperationContext opContext, final List entityNames, final String input, final Filter filters, final SortCriterion sortCriterion, final int start, final int count, - @Nullable final SearchFlags searchFlags, @Nullable final List facets) { return entitySearchService.search( - entityNames, input, filters, sortCriterion, start, count, searchFlags, facets); + opContext, entityNames, input, filters, sortCriterion, start, count, facets); } /** Executes the expensive autocomplete query using the {@link EntitySearchService} */ private AutoCompleteResult getRawAutoCompleteResults( + @Nonnull OperationContext opContext, final String entityName, final String input, final String field, final Filter filters, final int limit) { - return entitySearchService.autoComplete(entityName, input, field, filters, limit); + return entitySearchService.autoComplete(opContext, entityName, input, field, filters, limit); } /** Executes the expensive autocomplete query using the {@link EntitySearchService} */ private BrowseResult getRawBrowseResults( + @Nonnull OperationContext opContext, final String entityName, final String input, final Filter filters, final int start, final int count) { - return entitySearchService.browse(entityName, input, filters, start, count); + return entitySearchService.browse(opContext, entityName, input, filters, start, count); } /** Executes the expensive search query using the {@link EntitySearchService} */ private ScrollResult getRawScrollResults( + @Nonnull OperationContext opContext, final List entities, final String input, final Filter filters, @@ -370,19 +377,18 @@ private ScrollResult getRawScrollResults( @Nullable final String scrollId, @Nullable final String keepAlive, final int count, - final boolean fulltext, - @Nullable final SearchFlags searchFlags) { + final boolean fulltext) { if (fulltext) { return entitySearchService.fullTextScroll( - entities, input, filters, sortCriterion, scrollId, keepAlive, count, searchFlags); + opContext, entities, input, filters, sortCriterion, scrollId, keepAlive, count); } else { return entitySearchService.structuredScroll( - entities, input, filters, sortCriterion, scrollId, keepAlive, count, searchFlags); + opContext, entities, input, filters, sortCriterion, scrollId, keepAlive, count); } } /** Returns true if the cache should be used or skipped when fetching search results */ - private boolean enableCache(final SearchFlags searchFlags) { + private boolean enableCache(@Nullable final SearchFlags searchFlags) { return enableCache && (searchFlags == null || !searchFlags.isSkipCache()); } } 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 936ecb6a8ead1c..0effed1d9a578c 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 @@ -1,5 +1,7 @@ package com.linkedin.metadata.search.elasticsearch; +import static com.linkedin.metadata.search.utils.SearchUtils.applyDefaultSearchFlags; + import com.linkedin.common.urn.Urn; import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.browse.BrowseResult; @@ -20,6 +22,7 @@ import com.linkedin.metadata.search.utils.SearchUtils; import com.linkedin.metadata.shared.ElasticSearchIndexed; import com.linkedin.structured.StructuredPropertyDefinition; +import io.datahubproject.metadata.context.OperationContext; import java.io.IOException; import java.util.Collection; import java.util.List; @@ -36,6 +39,16 @@ @RequiredArgsConstructor public class ElasticSearchService implements EntitySearchService, ElasticSearchIndexed { + public static final SearchFlags DEFAULT_SERVICE_SEARCH_FLAGS = + new SearchFlags() + .setFulltext(false) + .setMaxAggValues(20) + .setSkipCache(false) + .setSkipAggregates(false) + .setSkipHighlighting(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false); + private static final int MAX_RUN_IDS_INDEXED = 25; // Save the previous 25 run ids in the index. private final EntityIndexBuilders indexBuilders; private final ESSearchDAO esSearchDAO; @@ -76,8 +89,11 @@ public void clear() { } @Override - public long docCount(@Nonnull String entityName) { - return esSearchDAO.docCount(entityName); + public long docCount(@Nonnull OperationContext opContext, @Nonnull String entityName) { + return esSearchDAO.docCount( + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, null, DEFAULT_SERVICE_SEARCH_FLAGS)), + entityName); } @Override @@ -127,37 +143,47 @@ public void appendRunId(@Nonnull String entityName, @Nonnull Urn urn, @Nullable @Nonnull @Override public SearchResult search( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, - int size, - @Nullable SearchFlags searchFlags) { - return search(entityNames, input, postFilters, sortCriterion, from, size, searchFlags, null); + int size) { + return search(opContext, entityNames, input, postFilters, sortCriterion, from, size, null); } @Nonnull public SearchResult search( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, int size, - @Nullable SearchFlags searchFlags, @Nullable List facets) { log.debug( String.format( "Searching FullText Search documents entityName: %s, input: %s, postFilters: %s, sortCriterion: %s, from: %s, size: %s", entityNames, input, postFilters, sortCriterion, from, size)); + return esSearchDAO.search( - entityNames, input, postFilters, sortCriterion, from, size, searchFlags, facets); + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, input, DEFAULT_SERVICE_SEARCH_FLAGS)), + entityNames, + input, + postFilters, + sortCriterion, + from, + size, + facets); } @Nonnull @Override public SearchResult filter( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nullable Filter filters, @Nullable SortCriterion sortCriterion, @@ -167,12 +193,21 @@ public SearchResult filter( String.format( "Filtering Search documents entityName: %s, filters: %s, sortCriterion: %s, from: %s, size: %s", entityName, filters, sortCriterion, from, size)); - return esSearchDAO.filter(entityName, filters, sortCriterion, from, size); + + return esSearchDAO.filter( + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, null, DEFAULT_SERVICE_SEARCH_FLAGS)), + entityName, + filters, + sortCriterion, + from, + size); } @Nonnull @Override public AutoCompleteResult autoComplete( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String query, @Nullable String field, @@ -182,12 +217,21 @@ public AutoCompleteResult autoComplete( String.format( "Autocompleting query entityName: %s, query: %s, field: %s, requestParams: %s, limit: %s", entityName, query, field, requestParams, limit)); - return esSearchDAO.autoComplete(entityName, query, field, requestParams, limit); + + return esSearchDAO.autoComplete( + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, query, DEFAULT_SERVICE_SEARCH_FLAGS)), + entityName, + query, + field, + requestParams, + limit); } @Nonnull @Override public Map aggregateByValue( + @Nonnull OperationContext opContext, @Nullable List entityNames, @Nonnull String field, @Nullable Filter requestParams, @@ -198,12 +242,20 @@ public Map aggregateByValue( field, requestParams, limit); - return esSearchDAO.aggregateByValue(entityNames, field, requestParams, limit); + + return esSearchDAO.aggregateByValue( + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, null, DEFAULT_SERVICE_SEARCH_FLAGS)), + entityNames, + field, + requestParams, + limit); } @Nonnull @Override public BrowseResult browse( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filters, @@ -213,33 +265,58 @@ public BrowseResult browse( String.format( "Browsing entities entityName: %s, path: %s, filters: %s, from: %s, size: %s", entityName, path, filters, from, size)); - return esBrowseDAO.browse(entityName, path, filters, from, size); + return esBrowseDAO.browse( + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, null, DEFAULT_SERVICE_SEARCH_FLAGS)), + entityName, + path, + filters, + from, + size); } @Nonnull @Override public BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nullable SearchFlags searchFlags) { - return esBrowseDAO.browseV2(entityName, path, filter, input, start, count, searchFlags); + int count) { + + return esBrowseDAO.browseV2( + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, null, DEFAULT_SERVICE_SEARCH_FLAGS)), + entityName, + path, + filter, + input, + start, + count); } @Nonnull @Override public BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nullable SearchFlags searchFlags) { - return esBrowseDAO.browseV2(entityNames, path, filter, input, start, count, searchFlags); + int count) { + + return esBrowseDAO.browseV2( + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, input, DEFAULT_SERVICE_SEARCH_FLAGS)), + entityNames, + path, + filter, + input, + start, + count); } @Nonnull @@ -253,43 +330,61 @@ public List getBrowsePaths(@Nonnull String entityName, @Nonnull Urn urn) @Nonnull @Override public ScrollResult fullTextScroll( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, - int size, - @Nullable SearchFlags searchFlags) { + int size) { 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, flags); + opContext.withSearchFlags( + flags -> + applyDefaultSearchFlags(flags, input, DEFAULT_SERVICE_SEARCH_FLAGS) + .setFulltext(true)), + entities, + input, + postFilters, + sortCriterion, + scrollId, + keepAlive, + size); } @Nonnull @Override public ScrollResult structuredScroll( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, - int size, - @Nullable SearchFlags searchFlags) { + int size) { 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)); - SearchFlags flags = Optional.ofNullable(searchFlags).orElse(new SearchFlags()); - flags.setFulltext(false); + return esSearchDAO.scroll( - entities, input, postFilters, sortCriterion, scrollId, keepAlive, size, flags); + opContext.withSearchFlags( + flags -> + applyDefaultSearchFlags(flags, null, DEFAULT_SERVICE_SEARCH_FLAGS) + .setFulltext(false)), + entities, + input, + postFilters, + sortCriterion, + scrollId, + keepAlive, + size); } public Optional raw(@Nonnull String indexName, @Nullable String jsonQuery) { @@ -303,23 +398,25 @@ public int maxResultSize() { @Override public ExplainResponse explain( + @Nonnull OperationContext opContext, @Nonnull String query, @Nonnull String documentId, @Nonnull String entityName, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, - @Nullable SearchFlags searchFlags, @Nullable String scrollId, @Nullable String keepAlive, int size, @Nullable List facets) { + return esSearchDAO.explain( + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, null, DEFAULT_SERVICE_SEARCH_FLAGS)), query, documentId, entityName, postFilters, sortCriterion, - searchFlags, scrollId, keepAlive, size, 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 79f530f18a3451..1958bed33c92b7 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 @@ -3,6 +3,7 @@ import static com.linkedin.metadata.Constants.ENTITY_TYPE_URN_PREFIX; import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_MAPPING_FIELD; import static com.linkedin.metadata.models.StructuredPropertyUtils.sanitizeStructuredPropertyFQN; +import static com.linkedin.metadata.models.annotation.SearchableAnnotation.OBJECT_FIELD_TYPES; import static com.linkedin.metadata.search.elasticsearch.indexbuilder.SettingsBuilder.*; import com.google.common.collect.ImmutableMap; @@ -53,6 +54,7 @@ public static Map getPartialNgramConfigWithOverrides( public static final String PATH = "path"; public static final String PROPERTIES = "properties"; + public static final String DYNAMIC_TEMPLATES = "dynamic_templates"; private MappingsBuilder() {} @@ -100,6 +102,7 @@ public static Map getMappings( return merged.isEmpty() ? null : merged; }); } + return mappings; } @@ -221,7 +224,7 @@ private static Map getMappingsForField( mappingForField.put(TYPE, ESUtils.LONG_FIELD_TYPE); } else if (fieldType == FieldType.DATETIME) { mappingForField.put(TYPE, ESUtils.DATE_FIELD_TYPE); - } else if (fieldType == FieldType.OBJECT) { + } else if (OBJECT_FIELD_TYPES.contains(fieldType)) { mappingForField.put(TYPE, ESUtils.OBJECT_FIELD_TYPE); } else if (fieldType == FieldType.DOUBLE) { mappingForField.put(TYPE, ESUtils.DOUBLE_FIELD_TYPE); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ReindexConfig.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ReindexConfig.java index bb6905139f49d5..fbb7fcadba8bca 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ReindexConfig.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ReindexConfig.java @@ -1,6 +1,8 @@ package com.linkedin.metadata.search.elasticsearch.indexbuilder; import static com.linkedin.metadata.Constants.*; +import static com.linkedin.metadata.search.elasticsearch.indexbuilder.MappingsBuilder.PROPERTIES; +import static com.linkedin.metadata.search.elasticsearch.indexbuilder.SettingsBuilder.TYPE; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.StreamReadConstraints; @@ -8,6 +10,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.MapDifference; import com.google.common.collect.Maps; +import com.linkedin.metadata.search.utils.ESUtils; import java.util.List; import java.util.Map; import java.util.Objects; @@ -146,9 +149,10 @@ public ReindexConfig build() { if (super.exists) { /* Consider mapping changes */ MapDifference mappingsDiff = - Maps.difference( - getOrDefault(super.currentMappings, List.of("properties")), - getOrDefault(super.targetMappings, List.of("properties"))); + calculateMapDifference( + getOrDefault(super.currentMappings, List.of(PROPERTIES)), + getOrDefault(super.targetMappings, List.of(PROPERTIES))); + super.requiresApplyMappings = !mappingsDiff.entriesDiffering().isEmpty() || !mappingsDiff.entriesOnlyOnRight().isEmpty(); @@ -298,6 +302,47 @@ private boolean isSettingsReindexRequired() { (Map) indexSettings.get("analysis"), super.currentSettings.getByPrefix("index.analysis.")); } + + /** + * Dynamic fields should not be considered as part of the difference. This might need to be + * improved in the future for nested object fields. + * + * @param currentMappings current mappings + * @param targetMappings target mappings + * @return difference map + */ + private static MapDifference calculateMapDifference( + Map currentMappings, Map targetMappings) { + + // Identify dynamic object fields in target + Set targetObjectFields = + targetMappings.entrySet().stream() + .filter( + entry -> + ((Map) entry.getValue()).containsKey(TYPE) + && ((Map) entry.getValue()) + .get(TYPE) + .equals(ESUtils.OBJECT_FIELD_TYPE)) + .map(Map.Entry::getKey) + .collect(Collectors.toSet()); + + if (!targetObjectFields.isEmpty()) { + log.info("Object fields filtered from comparison: {}", targetObjectFields); + Map filteredCurrentMappings = + removeKeys(currentMappings, targetObjectFields); + Map filteredTargetMappings = removeKeys(targetMappings, targetObjectFields); + return Maps.difference(filteredCurrentMappings, filteredTargetMappings); + } + + return Maps.difference(currentMappings, targetMappings); + } + } + + private static Map removeKeys( + Map mapObject, Set keysToRemove) { + return mapObject.entrySet().stream() + .filter(entry -> !keysToRemove.contains(entry.getKey())) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); } private static boolean equalsGroup(Map newSettings, Settings oldSettings) { diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java index dd1c09853114d1..7209c1ce147beb 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java @@ -1,6 +1,7 @@ package com.linkedin.metadata.search.elasticsearch.query; -import static com.linkedin.metadata.utils.SearchUtil.filterSoftDeletedByDefault; +import static com.linkedin.metadata.search.utils.ESUtils.applyDefaultSearchFilters; +import static com.linkedin.metadata.search.utils.SearchUtils.applyDefaultSearchFlags; import com.codahale.metrics.Timer; import com.datahub.util.exception.ESQueryException; @@ -29,6 +30,7 @@ import com.linkedin.metadata.utils.SearchUtil; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.metrics.MetricUtils; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.ArrayList; import java.util.Arrays; @@ -86,6 +88,14 @@ public class ESBrowseDAO { // Set explicit max size for grouping private static final int AGGREGATION_MAX_SIZE = 2000; + private static final SearchFlags DEFAULT_BROWSE_SEARCH_FLAGS = + new SearchFlags() + .setFulltext(false) + .setSkipHighlighting(true) + .setGetSuggestions(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false); + @Value private class BrowseGroupsResult { List groups; @@ -112,6 +122,7 @@ private class BrowseGroupsResultV2 { */ @Nonnull public BrowseResult browse( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filters, @@ -119,6 +130,10 @@ public BrowseResult browse( int size) { final Map requestMap = SearchUtils.getRequestMap(filters); + final OperationContext finalOpContext = + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, path, DEFAULT_BROWSE_SEARCH_FLAGS)); + try { final String indexName = indexConvention.getIndexName( @@ -128,7 +143,8 @@ public BrowseResult browse( try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "esGroupSearch").time()) { groupsResponse = client.search( - constructGroupsSearchRequest(indexName, path, requestMap), RequestOptions.DEFAULT); + constructGroupsSearchRequest(finalOpContext, indexName, path, requestMap), + RequestOptions.DEFAULT); } final BrowseGroupsResult browseGroupsResult = extractGroupsResponse(groupsResponse, path, from, size); @@ -144,7 +160,8 @@ public BrowseResult browse( try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "esEntitiesSearch").time()) { entitiesResponse = client.search( - constructEntitiesSearchRequest(indexName, path, requestMap, entityFrom, entitySize), + constructEntitiesSearchRequest( + finalOpContext, indexName, path, requestMap, entityFrom, entitySize), RequestOptions.DEFAULT); } final int numEntities = (int) entitiesResponse.getHits().getTotalHits().value; @@ -194,11 +211,14 @@ private AggregationBuilder buildAggregations(@Nonnull String path) { */ @Nonnull protected SearchRequest constructGroupsSearchRequest( - @Nonnull String indexName, @Nonnull String path, @Nonnull Map requestMap) { + @Nonnull OperationContext opContext, + @Nonnull String indexName, + @Nonnull String path, + @Nonnull Map requestMap) { final SearchRequest searchRequest = new SearchRequest(indexName); final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchSourceBuilder.size(0); - searchSourceBuilder.query(buildQueryString(path, requestMap, true)); + searchSourceBuilder.query(buildQueryString(opContext, path, requestMap, true)); searchSourceBuilder.aggregation(buildAggregations(path)); searchRequest.source(searchSourceBuilder); return searchRequest; @@ -214,12 +234,15 @@ protected SearchRequest constructGroupsSearchRequest( */ @Nonnull private QueryBuilder buildQueryString( - @Nonnull String path, @Nonnull Map requestMap, boolean isGroupQuery) { + @Nonnull OperationContext opContext, + @Nonnull String path, + @Nonnull Map requestMap, + boolean isGroupQuery) { final int browseDepthVal = getPathDepth(path); final BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery(); - queryBuilder.mustNot(QueryBuilders.termQuery(REMOVED, "true")); + applyDefaultSearchFilters(opContext, null, queryBuilder); if (!path.isEmpty()) { queryBuilder.filter(QueryBuilders.termQuery(BROWSE_PATH, path)); @@ -247,6 +270,7 @@ private QueryBuilder buildQueryString( @VisibleForTesting @Nonnull SearchRequest constructEntitiesSearchRequest( + @Nonnull OperationContext opContext, @Nonnull String indexName, @Nonnull String path, @Nonnull Map requestMap, @@ -258,7 +282,7 @@ SearchRequest constructEntitiesSearchRequest( searchSourceBuilder.size(size); searchSourceBuilder.fetchSource(new String[] {BROWSE_PATH, URN}, null); searchSourceBuilder.sort(URN, SortOrder.ASC); - searchSourceBuilder.query(buildQueryString(path, requestMap, false)); + searchSourceBuilder.query(buildQueryString(opContext, path, requestMap, false)); searchRequest.source(searchSourceBuilder); return searchRequest; } @@ -276,6 +300,7 @@ SearchRequest constructEntitiesSearchRequest( @VisibleForTesting @Nonnull SearchRequest constructEntitiesSearchRequest( + @Nonnull OperationContext opContext, @Nonnull String indexName, @Nonnull String path, @Nonnull Map requestMap, @@ -291,7 +316,7 @@ SearchRequest constructEntitiesSearchRequest( searchSourceBuilder.size(size); searchSourceBuilder.fetchSource(new String[] {BROWSE_PATH, URN}, null); searchSourceBuilder.sort(URN, SortOrder.ASC); - searchSourceBuilder.query(buildQueryString(path, requestMap, false)); + searchSourceBuilder.query(buildQueryString(opContext, path, requestMap, false)); searchRequest.source(searchSourceBuilder); return searchRequest; } @@ -401,20 +426,25 @@ public List getBrowsePaths(@Nonnull String entityName, @Nonnull Urn urn) } public BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nullable SearchFlags searchFlags) { + int count) { try { final SearchResponse groupsResponse; + final OperationContext finalOpContext = + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, path, DEFAULT_BROWSE_SEARCH_FLAGS)); + try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "esGroupSearch").time()) { final String finalInput = input.isEmpty() ? "*" : input; groupsResponse = client.search( - constructGroupsSearchRequestV2(entityName, path, filter, finalInput, searchFlags), + constructGroupsSearchRequestV2( + finalOpContext, entityName, path, filter, finalInput), RequestOptions.DEFAULT); } @@ -438,22 +468,25 @@ public BrowseResultV2 browseV2( } public BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nullable SearchFlags searchFlags) { + int count) { try { final SearchResponse groupsResponse; + final OperationContext finalOpContext = + opContext.withSearchFlags( + flags -> applyDefaultSearchFlags(flags, path, DEFAULT_BROWSE_SEARCH_FLAGS)); try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "esGroupSearch").time()) { final String finalInput = input.isEmpty() ? "*" : input; groupsResponse = client.search( constructGroupsSearchRequestBrowseAcrossEntities( - entities, path, filter, finalInput, searchFlags), + finalOpContext, entities, path, filter, finalInput), RequestOptions.DEFAULT); } @@ -478,11 +511,11 @@ public BrowseResultV2 browseV2( @Nonnull private SearchRequest constructGroupsSearchRequestV2( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filter, - @Nonnull String input, - @Nullable SearchFlags searchFlags) { + @Nonnull String input) { final String indexName = indexConvention.getIndexName(aspectRetriever.getEntityRegistry().getEntitySpec(entityName)); final SearchRequest searchRequest = new SearchRequest(indexName); @@ -490,11 +523,11 @@ private SearchRequest constructGroupsSearchRequestV2( searchSourceBuilder.size(0); searchSourceBuilder.query( buildQueryStringV2( + opContext, entityName, path, SearchUtil.transformFilterForEntities(filter, indexConvention), - input, - searchFlags)); + input)); searchSourceBuilder.aggregation(buildAggregationsV2(path)); searchRequest.source(searchSourceBuilder); return searchRequest; @@ -502,11 +535,11 @@ private SearchRequest constructGroupsSearchRequestV2( @Nonnull private SearchRequest constructGroupsSearchRequestBrowseAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String path, @Nullable Filter filter, - @Nonnull String input, - @Nullable SearchFlags searchFlags) { + @Nonnull String input) { List entitySpecs = entities.stream() @@ -521,11 +554,11 @@ private SearchRequest constructGroupsSearchRequestBrowseAcrossEntities( searchSourceBuilder.size(0); searchSourceBuilder.query( buildQueryStringBrowseAcrossEntities( + opContext, entitySpecs, path, SearchUtil.transformFilterForEntities(filter, indexConvention), - input, - searchFlags)); + input)); searchSourceBuilder.aggregation(buildAggregationsV2(path)); searchRequest.source(searchSourceBuilder); return searchRequest; @@ -550,13 +583,16 @@ private static int getPathDepthV2(@Nonnull String path) { @Nonnull private QueryBuilder buildQueryStringV2( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filter, - @Nonnull String input, - @Nullable SearchFlags searchFlags) { - SearchFlags finalSearchFlags = - Optional.ofNullable(searchFlags).orElse(new SearchFlags().setFulltext(true)); + @Nonnull String input) { + + final OperationContext finalOpContext = + opContext.withSearchFlags( + flags -> Optional.ofNullable(flags).orElse(new SearchFlags().setFulltext(true))); + final int browseDepthVal = getPathDepthV2(path); final BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery(); @@ -565,11 +601,12 @@ private QueryBuilder buildQueryStringV2( QueryBuilder query = SearchRequestHandler.getBuilder( entitySpec, searchConfiguration, customSearchConfiguration, aspectRetriever) - .getQuery(input, Boolean.TRUE.equals(finalSearchFlags.isFulltext())); + .getQuery( + input, + Boolean.TRUE.equals( + finalOpContext.getSearchContext().getSearchFlags().isFulltext())); queryBuilder.must(query); - filterSoftDeletedByDefault(filter, queryBuilder); - if (!path.isEmpty()) { queryBuilder.filter(QueryBuilders.matchQuery(BROWSE_PATH_V2, path)); } @@ -578,20 +615,21 @@ private QueryBuilder buildQueryStringV2( queryBuilder.filter( SearchRequestHandler.getFilterQuery( - filter, entitySpec.getSearchableFieldTypes(), aspectRetriever)); + finalOpContext, filter, entitySpec.getSearchableFieldTypes(), aspectRetriever)); return queryBuilder; } @Nonnull private QueryBuilder buildQueryStringBrowseAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entitySpecs, @Nonnull String path, @Nullable Filter filter, - @Nonnull String input, - @Nullable SearchFlags searchFlags) { - SearchFlags finalSearchFlags = - Optional.ofNullable(searchFlags).orElse(new SearchFlags().setFulltext(true)); + @Nonnull String input) { + final OperationContext finalOpContext = + opContext.withSearchFlags( + flags -> Optional.ofNullable(flags).orElse(new SearchFlags().setFulltext(true))); final int browseDepthVal = getPathDepthV2(path); final BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery(); @@ -599,7 +637,10 @@ private QueryBuilder buildQueryStringBrowseAcrossEntities( QueryBuilder query = SearchRequestHandler.getBuilder( entitySpecs, searchConfiguration, customSearchConfiguration, aspectRetriever) - .getQuery(input, Boolean.TRUE.equals(finalSearchFlags.isFulltext())); + .getQuery( + input, + Boolean.TRUE.equals( + finalOpContext.getSearchContext().getSearchFlags().isFulltext())); queryBuilder.must(query); if (!path.isEmpty()) { @@ -620,7 +661,8 @@ private QueryBuilder buildQueryStringBrowseAcrossEntities( return set1; })); queryBuilder.filter( - SearchRequestHandler.getFilterQuery(filter, searchableFields, aspectRetriever)); + SearchRequestHandler.getFilterQuery( + finalOpContext, filter, searchableFields, aspectRetriever)); return queryBuilder; } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java index 4cfb54dacb5f03..3fd1062fb25c5b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java @@ -14,7 +14,6 @@ import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.search.AggregationMetadata; @@ -29,6 +28,7 @@ import com.linkedin.metadata.search.utils.QueryUtils; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.metrics.MetricUtils; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.io.IOException; import java.util.ArrayList; @@ -81,13 +81,13 @@ public class ESSearchDAO { @Nonnull private final SearchConfiguration searchConfiguration; @Nullable private final CustomSearchConfiguration customSearchConfiguration; - public long docCount(@Nonnull String entityName) { + public long docCount(@Nonnull OperationContext opContext, @Nonnull String entityName) { EntitySpec entitySpec = aspectRetriever.getEntityRegistry().getEntitySpec(entityName); CountRequest countRequest = new CountRequest(indexConvention.getIndexName(entitySpec)) .query( SearchRequestHandler.getFilterQuery( - null, entitySpec.getSearchableFieldTypes(), aspectRetriever)); + opContext, null, entitySpec.getSearchableFieldTypes(), aspectRetriever)); try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "docCount").time()) { return client.count(countRequest, RequestOptions.DEFAULT).getCount(); } catch (IOException e) { @@ -99,6 +99,7 @@ public long docCount(@Nonnull String entityName) { @Nonnull @WithSpan private SearchResult executeAndExtract( + @Nonnull OperationContext opContext, @Nonnull List entitySpec, @Nonnull SearchRequest searchRequest, @Nullable Filter filter, @@ -113,7 +114,7 @@ private SearchResult executeAndExtract( return transformIndexIntoEntityName( SearchRequestHandler.getBuilder( entitySpec, searchConfiguration, customSearchConfiguration, aspectRetriever) - .extractResult(searchResponse, filter, from, size)); + .extractResult(opContext, searchResponse, filter, from, size)); } catch (Exception e) { log.error("Search query failed", e); throw new ESQueryException("Search query failed:", e); @@ -184,6 +185,7 @@ private AggregationMetadataArray transformIndexIntoEntityName(AggregationMetadat @Nonnull @WithSpan private ScrollResult executeAndExtract( + @Nonnull OperationContext opContext, @Nonnull List entitySpecs, @Nonnull SearchRequest searchRequest, @Nullable Filter filter, @@ -198,7 +200,13 @@ private ScrollResult executeAndExtract( SearchRequestHandler.getBuilder( entitySpecs, searchConfiguration, customSearchConfiguration, aspectRetriever) .extractScrollResult( - searchResponse, filter, scrollId, keepAlive, size, supportsPointInTime())); + opContext, + searchResponse, + filter, + scrollId, + keepAlive, + size, + supportsPointInTime())); } catch (Exception e) { log.error("Search query failed: {}", searchRequest, e); throw new ESQueryException("Search query failed:", e); @@ -215,20 +223,19 @@ private ScrollResult executeAndExtract( * @param sortCriterion {@link SortCriterion} to be applied to search results * @param from index to start the search from * @param size the number of search hits to return - * @param searchFlags Structured or full text search modes, plus other misc options * @param facets list of facets we want aggregations for * @return a {@link SearchResult} that contains a list of matched documents and related search * result metadata */ @Nonnull public SearchResult search( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, int size, - @Nullable SearchFlags searchFlags, @Nullable List facets) { final String finalInput = input.isEmpty() ? "*" : input; Timer.Context searchRequestTimer = MetricUtils.timer(this.getClass(), "searchRequest").time(); @@ -242,12 +249,12 @@ public SearchResult search( SearchRequestHandler.getBuilder( entitySpecs, searchConfiguration, customSearchConfiguration, aspectRetriever) .getSearchRequest( - finalInput, transformedFilters, sortCriterion, from, size, searchFlags, facets); + opContext, finalInput, transformedFilters, sortCriterion, from, size, facets); searchRequest.indices( entityNames.stream().map(indexConvention::getEntityIndexName).toArray(String[]::new)); searchRequestTimer.stop(); // Step 2: execute the query and extract results, validated against document model as well - return executeAndExtract(entitySpecs, searchRequest, transformedFilters, from, size); + return executeAndExtract(opContext, entitySpecs, searchRequest, transformedFilters, from, size); } /** @@ -263,6 +270,7 @@ public SearchResult search( */ @Nonnull public SearchResult filter( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nullable Filter filters, @Nullable SortCriterion sortCriterion, @@ -273,10 +281,11 @@ public SearchResult filter( final SearchRequest searchRequest = SearchRequestHandler.getBuilder( entitySpec, searchConfiguration, customSearchConfiguration, aspectRetriever) - .getFilterRequest(transformedFilters, sortCriterion, from, size); + .getFilterRequest(opContext, transformedFilters, sortCriterion, from, size); searchRequest.indices(indexConvention.getIndexName(entitySpec)); - return executeAndExtract(List.of(entitySpec), searchRequest, transformedFilters, from, size); + return executeAndExtract( + opContext, List.of(entitySpec), searchRequest, transformedFilters, from, size); } /** @@ -293,18 +302,24 @@ public SearchResult filter( */ @Nonnull public AutoCompleteResult autoComplete( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String query, @Nullable String field, @Nullable Filter requestParams, int limit) { try { + EntitySpec entitySpec = aspectRetriever.getEntityRegistry().getEntitySpec(entityName); AutocompleteRequestHandler builder = AutocompleteRequestHandler.getBuilder(entitySpec, aspectRetriever); SearchRequest req = builder.getSearchRequest( - query, field, transformFilterForEntities(requestParams, indexConvention), limit); + opContext, + query, + field, + transformFilterForEntities(requestParams, indexConvention), + limit); req.indices(indexConvention.getIndexName(entitySpec)); SearchResponse searchResponse = client.search(req, RequestOptions.DEFAULT); return builder.extractResult(searchResponse, query); @@ -325,10 +340,12 @@ public AutoCompleteResult autoComplete( */ @Nonnull public Map aggregateByValue( + @Nonnull OperationContext opContext, @Nullable List entityNames, @Nonnull String field, @Nullable Filter requestParams, int limit) { + List entitySpecs; if (entityNames == null || entityNames.isEmpty()) { entitySpecs = QueryUtils.getQueryByDefaultEntitySpecs(aspectRetriever.getEntityRegistry()); @@ -342,7 +359,10 @@ public Map aggregateByValue( SearchRequestHandler.getBuilder( entitySpecs, searchConfiguration, customSearchConfiguration, aspectRetriever) .getAggregationRequest( - field, transformFilterForEntities(requestParams, indexConvention), limit); + opContext, + field, + transformFilterForEntities(requestParams, indexConvention), + limit); if (entityNames == null) { String indexName = indexConvention.getAllEntityIndicesPattern(); searchRequest.indices(indexName); @@ -381,14 +401,14 @@ public Map aggregateByValue( */ @Nonnull public ScrollResult scroll( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, - int size, - SearchFlags searchFlags) { + int size) { final String finalInput = input.isEmpty() ? "*" : input; String[] indexArray = entities.stream().map(indexConvention::getEntityIndexName).toArray(String[]::new); @@ -401,6 +421,7 @@ public ScrollResult scroll( // TODO: Align scroll and search using facets final SearchRequest searchRequest = getScrollRequest( + opContext, scrollId, keepAlive, indexArray, @@ -409,7 +430,6 @@ public ScrollResult scroll( entitySpecs, finalInput, sortCriterion, - searchFlags, null); // PIT specifies indices in creation so it doesn't support specifying indices on the request, so @@ -420,10 +440,11 @@ public ScrollResult scroll( scrollRequestTimer.stop(); return executeAndExtract( - entitySpecs, searchRequest, transformedFilters, scrollId, keepAlive, size); + opContext, entitySpecs, searchRequest, transformedFilters, scrollId, keepAlive, size); } private SearchRequest getScrollRequest( + @Nonnull OperationContext opContext, @Nullable String scrollId, @Nullable String keepAlive, String[] indexArray, @@ -432,7 +453,6 @@ private SearchRequest getScrollRequest( List entitySpecs, String finalInput, @Nullable SortCriterion sortCriterion, - @Nullable SearchFlags searchFlags, @Nullable List facets) { String pitId = null; Object[] sort = null; @@ -453,6 +473,7 @@ private SearchRequest getScrollRequest( return SearchRequestHandler.getBuilder( entitySpecs, searchConfiguration, customSearchConfiguration, aspectRetriever) .getSearchRequest( + opContext, finalInput, postFilters, sortCriterion, @@ -460,7 +481,6 @@ private SearchRequest getScrollRequest( pitId, keepAlive, size, - searchFlags, facets); } @@ -507,12 +527,12 @@ private String createPointInTime(String[] indexArray, String keepAlive) { } public ExplainResponse explain( + @Nonnull OperationContext opContext, @Nonnull String query, @Nonnull String documentId, @Nonnull String entityName, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, - @Nullable SearchFlags searchFlags, @Nullable String scrollId, @Nullable String keepAlive, int size, @@ -523,6 +543,7 @@ public ExplainResponse explain( final String finalQuery = query.isEmpty() ? "*" : query; final SearchRequest searchRequest = getScrollRequest( + opContext, scrollId, keepAlive, indexArray, @@ -531,7 +552,6 @@ public ExplainResponse explain( Collections.singletonList(entitySpec), finalQuery, sortCriterion, - searchFlags, facets); ; diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java index fb3b51930370c4..a32a8d1a54234d 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java @@ -149,7 +149,8 @@ private AggregationBuilder facetToAggregationBuilder(final String inputFacet) { case MISSING_SPECIAL_TYPE: aggBuilder = INDEX_VIRTUAL_FIELD.equalsIgnoreCase(specialTypeFields.get(1)) - ? AggregationBuilders.missing(inputFacet).field(getAggregationField("_index")) + ? AggregationBuilders.missing(inputFacet) + .field(getAggregationField(ES_INDEX_FIELD)) : AggregationBuilders.missing(inputFacet) .field(getAggregationField(specialTypeFields.get(1))); break; @@ -161,7 +162,7 @@ private AggregationBuilder facetToAggregationBuilder(final String inputFacet) { aggBuilder = facet.equalsIgnoreCase(INDEX_VIRTUAL_FIELD) ? AggregationBuilders.terms(inputFacet) - .field(getAggregationField("_index")) + .field(getAggregationField(ES_INDEX_FIELD)) .size(configs.getMaxTermBucketSize()) .minDocCount(0) : AggregationBuilders.terms(inputFacet) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java index de35d53bcde49b..3e6ce53b7af5c1 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java @@ -1,6 +1,7 @@ package com.linkedin.metadata.search.elasticsearch.query.request; import static com.linkedin.metadata.models.SearchableFieldSpecExtractor.PRIMARY_URN_SEARCH_PROPERTIES; +import static com.linkedin.metadata.search.utils.ESUtils.applyDefaultSearchFilters; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; @@ -14,6 +15,7 @@ import com.linkedin.metadata.query.AutoCompleteResult; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.utils.ESUtils; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.Collections; import java.util.HashSet; @@ -32,7 +34,6 @@ import org.opensearch.action.search.SearchResponse; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.MultiMatchQueryBuilder; -import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.search.SearchHit; import org.opensearch.search.builder.SearchSourceBuilder; @@ -84,11 +85,19 @@ public static AutocompleteRequestHandler getBuilder( } public SearchRequest getSearchRequest( - @Nonnull String input, @Nullable String field, @Nullable Filter filter, int limit) { + @Nonnull OperationContext opContext, + @Nonnull String input, + @Nullable String field, + @Nullable Filter filter, + int limit) { SearchRequest searchRequest = new SearchRequest(); SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchSourceBuilder.size(limit); - searchSourceBuilder.query(getQuery(input, field)); + // apply default filters + BoolQueryBuilder boolQueryBuilder = + applyDefaultSearchFilters(opContext, filter, getQuery(input, field)); + + searchSourceBuilder.query(boolQueryBuilder); searchSourceBuilder.postFilter( ESUtils.buildFilterQuery(filter, false, searchableFieldTypes, aspectRetriever)); searchSourceBuilder.highlighter(getHighlights(field)); @@ -96,11 +105,11 @@ public SearchRequest getSearchRequest( return searchRequest; } - private QueryBuilder getQuery(@Nonnull String query, @Nullable String field) { + private BoolQueryBuilder getQuery(@Nonnull String query, @Nullable String field) { return getQuery(getAutocompleteFields(field), query); } - public static QueryBuilder getQuery(List autocompleteFields, @Nonnull String query) { + public static BoolQueryBuilder getQuery(List autocompleteFields, @Nonnull String query) { BoolQueryBuilder finalQuery = QueryBuilders.boolQuery(); // Search for exact matches with higher boost and ngram matches MultiMatchQueryBuilder autocompleteQueryBuilder = @@ -126,8 +135,6 @@ public static QueryBuilder getQuery(List autocompleteFields, @Nonnull St }); finalQuery.should(autocompleteQueryBuilder); - - finalQuery.mustNot(QueryBuilders.matchQuery("removed", true)); return finalQuery; } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchFieldConfig.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchFieldConfig.java index 7709ff16f79409..d681df00546acb 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchFieldConfig.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchFieldConfig.java @@ -51,7 +51,8 @@ public class SearchFieldConfig { SearchableAnnotation.FieldType.BOOLEAN, SearchableAnnotation.FieldType.COUNT, SearchableAnnotation.FieldType.DATETIME, - SearchableAnnotation.FieldType.OBJECT); + SearchableAnnotation.FieldType.OBJECT, + SearchableAnnotation.FieldType.MAP_ARRAY); // NOT true for `urn` public static final Set TYPES_WITH_URN_TEXT = Set.of(SearchableAnnotation.FieldType.URN, SearchableAnnotation.FieldType.URN_PARTIAL); 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 db09c52d2099c8..0ae23445140e03 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 @@ -1,8 +1,7 @@ package com.linkedin.metadata.search.elasticsearch.query.request; import static com.linkedin.metadata.search.utils.ESUtils.NAME_SUGGESTION; -import static com.linkedin.metadata.search.utils.SearchUtils.applyDefaultSearchFlags; -import static com.linkedin.metadata.utils.SearchUtil.*; +import static com.linkedin.metadata.search.utils.ESUtils.applyDefaultSearchFilters; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; @@ -30,8 +29,10 @@ import com.linkedin.metadata.search.SearchSuggestion; import com.linkedin.metadata.search.SearchSuggestionArray; import com.linkedin.metadata.search.features.Features; +import com.linkedin.metadata.search.utils.ESAccessControlUtil; import com.linkedin.metadata.search.utils.ESUtils; import com.linkedin.util.Pair; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.util.ArrayList; @@ -64,13 +65,7 @@ @Slf4j public class SearchRequestHandler { - private static final SearchFlags DEFAULT_SERVICE_SEARCH_FLAGS = - new SearchFlags() - .setFulltext(false) - .setMaxAggValues(20) - .setSkipCache(false) - .setSkipAggregates(false) - .setSkipHighlighting(false); + private static final Map, SearchRequestHandler> REQUEST_HANDLER_BY_ENTITY_NAME = new ConcurrentHashMap<>(); private final List entitySpecs; @@ -170,18 +165,19 @@ private Set getDefaultQueryFieldNames(List annotat .collect(Collectors.toSet()); } - public BoolQueryBuilder getFilterQuery(@Nullable Filter filter) { - return getFilterQuery(filter, searchableFieldTypes, aspectRetriever); + public BoolQueryBuilder getFilterQuery( + @Nonnull OperationContext opContext, @Nullable Filter filter) { + return getFilterQuery(opContext, filter, searchableFieldTypes, aspectRetriever); } public static BoolQueryBuilder getFilterQuery( + @Nonnull OperationContext opContext, @Nullable Filter filter, Map> searchableFieldTypes, @Nonnull AspectRetriever aspectRetriever) { BoolQueryBuilder filterQuery = ESUtils.buildFilterQuery(filter, false, searchableFieldTypes, aspectRetriever); - - return filterSoftDeletedByDefault(filter, filterQuery); + return applyDefaultSearchFilters(opContext, filter, filterQuery); } /** @@ -194,23 +190,21 @@ public static BoolQueryBuilder getFilterQuery( * @param filter the search filter * @param from index to start the search from * @param size the number of search hits to return - * @param searchFlags Various flags controlling search query options * @param facets list of facets we want aggregations for * @return a valid search request */ @Nonnull @WithSpan public SearchRequest getSearchRequest( + @Nonnull OperationContext opContext, @Nonnull String input, @Nullable Filter filter, @Nullable SortCriterion sortCriterion, int from, int size, - @Nullable SearchFlags searchFlags, @Nullable List facets) { - SearchFlags finalSearchFlags = - applyDefaultSearchFlags(searchFlags, input, DEFAULT_SERVICE_SEARCH_FLAGS); + SearchFlags searchFlags = opContext.getSearchContext().getSearchFlags(); SearchRequest searchRequest = new SearchRequest(); SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); @@ -218,20 +212,20 @@ public SearchRequest getSearchRequest( searchSourceBuilder.size(size); searchSourceBuilder.fetchSource("urn", null); - BoolQueryBuilder filterQuery = getFilterQuery(filter); + BoolQueryBuilder filterQuery = getFilterQuery(opContext, filter); searchSourceBuilder.query( QueryBuilders.boolQuery() - .must(getQuery(input, Boolean.TRUE.equals(finalSearchFlags.isFulltext()))) + .must(getQuery(input, Boolean.TRUE.equals(searchFlags.isFulltext()))) .filter(filterQuery)); - if (Boolean.FALSE.equals(finalSearchFlags.isSkipAggregates())) { + if (Boolean.FALSE.equals(searchFlags.isSkipAggregates())) { aggregationQueryBuilder.getAggregations(facets).forEach(searchSourceBuilder::aggregation); } - if (Boolean.FALSE.equals(finalSearchFlags.isSkipHighlighting())) { + if (Boolean.FALSE.equals(searchFlags.isSkipHighlighting())) { searchSourceBuilder.highlighter(highlights); } ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, entitySpecs); - if (Boolean.TRUE.equals(finalSearchFlags.isGetSuggestions())) { + if (Boolean.TRUE.equals(searchFlags.isGetSuggestions())) { ESUtils.buildNameSuggestions(searchSourceBuilder, input); } @@ -256,6 +250,7 @@ public SearchRequest getSearchRequest( @Nonnull @WithSpan public SearchRequest getSearchRequest( + @Nonnull OperationContext opContext, @Nonnull String input, @Nullable Filter filter, @Nullable SortCriterion sortCriterion, @@ -263,11 +258,10 @@ public SearchRequest getSearchRequest( @Nullable String pitId, @Nullable String keepAlive, int size, - SearchFlags searchFlags, @Nullable List facets) { + SearchFlags searchFlags = opContext.getSearchContext().getSearchFlags(); SearchRequest searchRequest = new PITAwareSearchRequest(); - SearchFlags finalSearchFlags = - applyDefaultSearchFlags(searchFlags, input, DEFAULT_SERVICE_SEARCH_FLAGS); + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); ESUtils.setSearchAfter(searchSourceBuilder, sort, pitId, keepAlive); @@ -275,15 +269,15 @@ public SearchRequest getSearchRequest( searchSourceBuilder.size(size); searchSourceBuilder.fetchSource("urn", null); - BoolQueryBuilder filterQuery = getFilterQuery(filter); + BoolQueryBuilder filterQuery = getFilterQuery(opContext, filter); searchSourceBuilder.query( QueryBuilders.boolQuery() - .must(getQuery(input, Boolean.TRUE.equals(finalSearchFlags.isFulltext()))) + .must(getQuery(input, Boolean.TRUE.equals(searchFlags.isFulltext()))) .filter(filterQuery)); - if (Boolean.FALSE.equals(finalSearchFlags.isSkipAggregates())) { + if (Boolean.FALSE.equals(searchFlags.isSkipAggregates())) { aggregationQueryBuilder.getAggregations(facets).forEach(searchSourceBuilder::aggregation); } - if (Boolean.FALSE.equals(finalSearchFlags.isSkipHighlighting())) { + if (Boolean.FALSE.equals(searchFlags.isSkipHighlighting())) { searchSourceBuilder.highlighter(highlights); } ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, entitySpecs); @@ -306,10 +300,14 @@ public SearchRequest getSearchRequest( */ @Nonnull public SearchRequest getFilterRequest( - @Nullable Filter filters, @Nullable SortCriterion sortCriterion, int from, int size) { + @Nonnull OperationContext opContext, + @Nullable Filter filters, + @Nullable SortCriterion sortCriterion, + int from, + int size) { SearchRequest searchRequest = new SearchRequest(); - BoolQueryBuilder filterQuery = getFilterQuery(filters); + BoolQueryBuilder filterQuery = getFilterQuery(opContext, filters); final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchSourceBuilder.query(filterQuery); searchSourceBuilder.from(from).size(size); @@ -329,9 +327,13 @@ public SearchRequest getFilterRequest( */ @Nonnull public SearchRequest getAggregationRequest( - @Nonnull String field, @Nullable Filter filter, int limit) { + @Nonnull OperationContext opContext, + @Nonnull String field, + @Nullable Filter filter, + int limit) { + SearchRequest searchRequest = new SearchRequest(); - BoolQueryBuilder filterQuery = getFilterQuery(filter); + BoolQueryBuilder filterQuery = getFilterQuery(opContext, filter); final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchSourceBuilder.query(filterQuery); @@ -366,10 +368,15 @@ public HighlightBuilder getHighlights() { @WithSpan public SearchResult extractResult( - @Nonnull SearchResponse searchResponse, Filter filter, int from, int size) { + @Nonnull OperationContext opContext, + @Nonnull SearchResponse searchResponse, + Filter filter, + int from, + int size) { int totalCount = (int) searchResponse.getHits().getTotalHits().value; - List resultList = getResults(searchResponse); - SearchResultMetadata searchResultMetadata = extractSearchResultMetadata(searchResponse, filter); + List resultList = getResults(opContext, searchResponse); + SearchResultMetadata searchResultMetadata = + extractSearchResultMetadata(opContext, searchResponse, filter); return new SearchResult() .setEntities(new SearchEntityArray(resultList)) @@ -381,6 +388,7 @@ public SearchResult extractResult( @WithSpan public ScrollResult extractScrollResult( + @Nonnull OperationContext opContext, @Nonnull SearchResponse searchResponse, Filter filter, @Nullable String scrollId, @@ -388,8 +396,9 @@ public ScrollResult extractScrollResult( int size, boolean supportsPointInTime) { int totalCount = (int) searchResponse.getHits().getTotalHits().value; - List resultList = getResults(searchResponse); - SearchResultMetadata searchResultMetadata = extractSearchResultMetadata(searchResponse, filter); + List resultList = getResults(opContext, searchResponse); + SearchResultMetadata searchResultMetadata = + extractSearchResultMetadata(opContext, searchResponse, filter); SearchHit[] searchHits = searchResponse.getHits().getHits(); // Only return next scroll ID if there are more results, indicated by full size results String nextScrollId = null; @@ -484,10 +493,13 @@ private SearchEntity getResult(@Nonnull SearchHit hit) { * @return List of search entities */ @Nonnull - private List getResults(@Nonnull SearchResponse searchResponse) { - return Arrays.stream(searchResponse.getHits().getHits()) - .map(this::getResult) - .collect(Collectors.toList()); + private List getResults( + @Nonnull OperationContext opContext, @Nonnull SearchResponse searchResponse) { + return ESAccessControlUtil.restrictSearchResult( + opContext, + Arrays.stream(searchResponse.getHits().getHits()) + .map(this::getResult) + .collect(Collectors.toList())); } @Nonnull @@ -509,13 +521,18 @@ private Urn getUrnFromSearchHit(@Nonnull SearchHit hit) { */ @Nonnull private SearchResultMetadata extractSearchResultMetadata( - @Nonnull SearchResponse searchResponse, @Nullable Filter filter) { + @Nonnull OperationContext opContext, + @Nonnull SearchResponse searchResponse, + @Nullable Filter filter) { + final SearchFlags searchFlags = opContext.getSearchContext().getSearchFlags(); final SearchResultMetadata searchResultMetadata = new SearchResultMetadata().setAggregations(new AggregationMetadataArray()); - final List aggregationMetadataList = - aggregationQueryBuilder.extractAggregationMetadata(searchResponse, filter); - searchResultMetadata.setAggregations(new AggregationMetadataArray(aggregationMetadataList)); + if (Boolean.FALSE.equals(searchFlags.isSkipAggregates())) { + final List aggregationMetadataList = + aggregationQueryBuilder.extractAggregationMetadata(searchResponse, filter); + searchResultMetadata.setAggregations(new AggregationMetadataArray(aggregationMetadataList)); + } final List searchSuggestions = extractSearchSuggestions(searchResponse); searchResultMetadata.setSuggestions(new SearchSuggestionArray(searchSuggestions)); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java b/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java index 75c3d23d26c667..a291b27ebebeff 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java @@ -2,6 +2,7 @@ import static com.linkedin.metadata.Constants.*; import static com.linkedin.metadata.models.StructuredPropertyUtils.sanitizeStructuredPropertyFQN; +import static com.linkedin.metadata.models.annotation.SearchableAnnotation.OBJECT_FIELD_TYPES; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ArrayNode; @@ -26,6 +27,7 @@ import com.linkedin.structured.StructuredPropertyValueAssignment; import java.net.URISyntaxException; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.Optional; @@ -181,7 +183,7 @@ public void setSearchableValue( return; } - if (isArray || (valueType == DataSchema.Type.MAP && fieldType != FieldType.OBJECT)) { + if (isArray || (valueType == DataSchema.Type.MAP && !OBJECT_FIELD_TYPES.contains(fieldType))) { if (fieldType == FieldType.BROWSE_PATH_V2) { String browsePathV2Value = getBrowsePathV2Value(fieldValues); searchDocument.set(fieldName, JsonNodeFactory.instance.textNode(browsePathV2Value)); @@ -193,6 +195,25 @@ public void setSearchableValue( value -> getNodeForValue(valueType, value, fieldType).ifPresent(arrayNode::add)); searchDocument.set(fieldName, arrayNode); } + } else if (valueType == DataSchema.Type.MAP && FieldType.MAP_ARRAY.equals(fieldType)) { + ObjectNode dictDoc = JsonNodeFactory.instance.objectNode(); + fieldValues + .subList(0, Math.min(fieldValues.size(), maxObjectKeys)) + .forEach( + fieldValue -> { + String[] keyValues = fieldValue.toString().split("="); + String key = keyValues[0]; + ArrayNode values = JsonNodeFactory.instance.arrayNode(); + Arrays.stream(keyValues[1].substring(1, keyValues[1].length() - 1).split(", ")) + .forEach( + v -> { + if (!v.isEmpty()) { + values.add(v); + } + }); + dictDoc.set(key, values); + }); + searchDocument.set(fieldName, dictDoc); } else if (valueType == DataSchema.Type.MAP) { ObjectNode dictDoc = JsonNodeFactory.instance.objectNode(); fieldValues diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESAccessControlUtil.java b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESAccessControlUtil.java new file mode 100644 index 00000000000000..cf3eabb9c5a829 --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESAccessControlUtil.java @@ -0,0 +1,260 @@ +package com.linkedin.metadata.search.utils; + +import static com.linkedin.metadata.authorization.PoliciesConfig.VIEW_ENTITY_PRIVILEGES; +import static com.linkedin.metadata.utils.SearchUtil.ES_INDEX_FIELD; +import static com.linkedin.metadata.utils.SearchUtil.KEYWORD_SUFFIX; + +import com.datahub.authorization.AuthUtil; +import com.datahub.authorization.ConjunctivePrivilegeGroup; +import com.datahub.authorization.DisjunctivePrivilegeGroup; +import com.datahub.authorization.EntitySpec; +import com.google.common.collect.ImmutableList; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.StringArray; +import com.linkedin.metadata.aspect.hooks.OwnerTypeMap; +import com.linkedin.metadata.authorization.PoliciesConfig; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.search.SearchEntity; +import com.linkedin.metadata.search.SearchResult; +import com.linkedin.metadata.timeseries.elastic.indexbuilder.MappingsBuilder; +import com.linkedin.policy.DataHubActorFilter; +import com.linkedin.policy.DataHubPolicyInfo; +import com.linkedin.policy.PolicyMatchCriterion; +import com.linkedin.policy.PolicyMatchCriterionArray; +import io.datahubproject.metadata.context.ActorContext; +import io.datahubproject.metadata.context.OperationContext; +import java.util.Collection; +import java.util.List; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import javax.annotation.Nonnull; +import lombok.extern.slf4j.Slf4j; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.index.query.TermsQueryBuilder; + +@Slf4j +public class ESAccessControlUtil { + private ESAccessControlUtil() {} + + private static final String OWNER_TYPES_FIELD = "ownerTypes"; + private static final QueryBuilder MATCH_ALL = QueryBuilders.matchAllQuery(); + + /** + * Given the OperationContext produce a filter for search results + * + * @param opContext the OperationContext of the search + * @return + */ + public static Optional buildAccessControlFilters( + @Nonnull OperationContext opContext) { + Optional response = Optional.empty(); + + /* + If search authorization is enabled AND we're also not the system performing the query + */ + if (opContext.getOperationContextConfig().getSearchAuthorizationConfiguration().isEnabled() + && !opContext.isSystemAuth() + && !opContext.getSearchContext().isRestrictedSearch()) { + + BoolQueryBuilder builder = QueryBuilders.boolQuery(); + + // Apply access policies + streamViewQueries(opContext).distinct().forEach(builder::should); + + if (builder.should().isEmpty()) { + // default no filters + return Optional.of(builder.mustNot(MATCH_ALL)); + } else if (!builder.should().contains(MATCH_ALL)) { + // if MATCH_ALL is not present, apply filters requiring at least 1 + builder.minimumShouldMatch(1); + response = Optional.of(builder); + } + } + + // MATCH_ALL filter present or system user or disabled + return response; + } + + /** + * Given an OperationContext and SearchResult, mark the restricted entities. Currently, the entire + * entity is marked as restricted using the key aspect name + * + * @param searchResult restricted search result + */ + public static void restrictSearchResult( + @Nonnull OperationContext opContext, @Nonnull SearchResult searchResult) { + restrictSearchResult(opContext, searchResult.getEntities()); + } + + public static > T restrictSearchResult( + @Nonnull OperationContext opContext, T searchEntities) { + if (opContext.getOperationContextConfig().getSearchAuthorizationConfiguration().isEnabled() + && opContext.getSearchContext().isRestrictedSearch()) { + final EntityRegistry entityRegistry = opContext.getEntityRegistry(); + final String actorUrnStr = + opContext.getSessionActorContext().getAuthentication().getActor().toUrnStr(); + final DisjunctivePrivilegeGroup orGroup = + new DisjunctivePrivilegeGroup( + ImmutableList.of(new ConjunctivePrivilegeGroup(VIEW_ENTITY_PRIVILEGES))); + + for (SearchEntity searchEntity : searchEntities) { + final String entityType = searchEntity.getEntity().getEntityType(); + final Optional resourceSpec = + Optional.of(new EntitySpec(entityType, searchEntity.getEntity().toString())); + if (!AuthUtil.isAuthorized( + opContext.getAuthorizerContext().getAuthorizer(), actorUrnStr, resourceSpec, orGroup)) { + final String keyAspectName = + entityRegistry.getEntitySpecs().get(entityType.toLowerCase()).getKeyAspectName(); + searchEntity.setRestrictedAspects(new StringArray(List.of(keyAspectName))); + } + } + } + return searchEntities; + } + + private static final Function activeMetadataViewEntityPolicyFilter = + policy -> + policy.getPrivileges() != null + && PoliciesConfig.ACTIVE_POLICY_STATE.equals(policy.getState()) + && PoliciesConfig.METADATA_POLICY_TYPE.equals(policy.getType()) + && VIEW_ENTITY_PRIVILEGES.stream() + .anyMatch(priv -> policy.getPrivileges().contains(priv)); + + private static Stream streamViewQueries(OperationContext opContext) { + return opContext.getSessionActorContext().getPolicyInfoSet().stream() + .filter(activeMetadataViewEntityPolicyFilter::apply) + .map( + policy -> { + // Build actor query + QueryBuilder actorQuery = buildActorQuery(opContext, policy); + + if (!policy.hasResources()) { + // no resource restrictions + return actorQuery; + } else { + + // No filters or criteria + if (!policy.getResources().hasFilter() + || !policy.getResources().getFilter().hasCriteria()) { + return null; + } + + PolicyMatchCriterionArray criteriaArray = + policy.getResources().getFilter().getCriteria(); + // Cannot apply policy if we can't map every field + if (!criteriaArray.stream().allMatch(criteria -> toESField(criteria).isPresent())) { + return null; + } + + BoolQueryBuilder resourceQuery = QueryBuilders.boolQuery(); + // apply actor filter if present + if (!MATCH_ALL.equals(actorQuery)) { + resourceQuery.filter(actorQuery); + } + // add resource query + buildResourceQuery(opContext, criteriaArray).forEach(resourceQuery::filter); + return resourceQuery; + } + }) + .filter(Objects::nonNull); + } + + /** + * Build an entity index query for ownership policies. If no restrictions, returns MATCH_ALL query + * + * @param opContext context + * @param policy policy + * @return filter query + */ + private static QueryBuilder buildActorQuery( + OperationContext opContext, DataHubPolicyInfo policy) { + DataHubActorFilter actorFilter = policy.getActors(); + + if (!policy.hasActors() + || !(actorFilter.isResourceOwners() || actorFilter.hasResourceOwnersTypes())) { + // no owner restriction + return MATCH_ALL; + } + + ActorContext actorContext = opContext.getSessionActorContext(); + + // policy might apply to the actor via user or group + List actorAndGroupUrns = + Stream.concat( + Stream.of(actorContext.getAuthentication().getActor().toUrnStr()), + actorContext.getGroupMembership().stream().map(Urn::toString)) + .map(String::toLowerCase) + .distinct() + .collect(Collectors.toList()); + + if (!actorFilter.hasResourceOwnersTypes()) { + // owners without owner type restrictions + return QueryBuilders.termsQuery( + ESUtils.toKeywordField(MappingsBuilder.OWNERS_FIELD, false), actorAndGroupUrns); + } else { + // owners with type restrictions + BoolQueryBuilder orQuery = QueryBuilders.boolQuery(); + orQuery.minimumShouldMatch(1); + + Set typeFields = + actorFilter.getResourceOwnersTypes().stream() + .map( + typeUrn -> + String.format( + "%s.%s%s", + OWNER_TYPES_FIELD, + OwnerTypeMap.encodeFieldName(typeUrn.toString()), + KEYWORD_SUFFIX)) + .collect(Collectors.toSet()); + + typeFields.forEach( + field -> orQuery.should(QueryBuilders.termsQuery(field, actorAndGroupUrns))); + + return orQuery; + } + } + + private static Stream buildResourceQuery( + OperationContext opContext, PolicyMatchCriterionArray criteriaArray) { + return criteriaArray.stream() + .map( + criteria -> + QueryBuilders.termsQuery( + toESField(criteria).get(), toESValues(opContext, criteria))); + } + + private static Optional toESField(PolicyMatchCriterion criterion) { + switch (criterion.getField()) { + case "TYPE": + return Optional.of(ES_INDEX_FIELD); + case "URN": + return Optional.of(ESUtils.toKeywordField(MappingsBuilder.URN_FIELD, false)); + case "TAG": + return Optional.of(ESUtils.toKeywordField(MappingsBuilder.TAGS_FIELD, false)); + case "DOMAIN": + return Optional.of(ESUtils.toKeywordField(MappingsBuilder.DOMAINS_FIELD, false)); + default: + return Optional.empty(); + } + } + + private static Collection toESValues( + OperationContext opContext, PolicyMatchCriterion criterion) { + switch (criterion.getField()) { + case "TYPE": + return criterion.getValues().stream() + .map( + value -> + opContext.getSearchContext().getIndexConvention().getEntityIndexName(value)) + .collect(Collectors.toSet()); + default: + return criterion.getValues(); + } + } +} 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 5ca5087d5ac355..6c4507216482fd 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 @@ -1,6 +1,7 @@ package com.linkedin.metadata.search.utils; import static com.linkedin.metadata.Constants.*; +import static com.linkedin.metadata.models.annotation.SearchableAnnotation.OBJECT_FIELD_TYPES; import static com.linkedin.metadata.search.elasticsearch.query.request.SearchFieldConfig.KEYWORD_FIELDS; import static com.linkedin.metadata.search.elasticsearch.query.request.SearchFieldConfig.PATH_HIERARCHY_FIELDS; import static com.linkedin.metadata.search.utils.SearchUtils.isUrn; @@ -11,11 +12,13 @@ import com.linkedin.metadata.models.SearchableFieldSpec; import com.linkedin.metadata.models.StructuredPropertyUtils; import com.linkedin.metadata.models.annotation.SearchableAnnotation; +import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.Criterion; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; +import io.datahubproject.metadata.context.OperationContext; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; @@ -54,6 +57,7 @@ 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 = "|"; + private static final String REMOVED = "removed"; // Field types public static final String KEYWORD_FIELD_TYPE = "keyword"; @@ -149,6 +153,8 @@ public static BoolQueryBuilder buildFilterQuery( or -> finalQueryBuilder.should( ESUtils.buildConjunctiveFilterQuery(or, isTimeseries, searchableFieldTypes))); + // The default is not always 1 (ensure consistent default) + finalQueryBuilder.minimumShouldMatch(1); } else if (filter.getCriteria() != null) { // Otherwise, build boolean query from the deprecated "criteria" field. log.warn("Received query Filter with a deprecated field 'criteria'. Use 'or' instead."); @@ -165,6 +171,8 @@ public static BoolQueryBuilder buildFilterQuery( } }); finalQueryBuilder.should(andQueryBuilder); + // The default is not always 1 (ensure consistent default) + finalQueryBuilder.minimumShouldMatch(1); } return finalQueryBuilder; } @@ -263,7 +271,7 @@ public static String getElasticTypeForFieldType(SearchableAnnotation.FieldType f return LONG_FIELD_TYPE; } else if (fieldType == SearchableAnnotation.FieldType.DATETIME) { return DATE_FIELD_TYPE; - } else if (fieldType == SearchableAnnotation.FieldType.OBJECT) { + } else if (OBJECT_FIELD_TYPES.contains(fieldType)) { return OBJECT_FIELD_TYPE; } else if (fieldType == SearchableAnnotation.FieldType.DOUBLE) { return DOUBLE_FIELD_TYPE; @@ -666,4 +674,43 @@ private static QueryBuilder buildEqualsFromCriterionWithValue( .analyzer(KEYWORD_ANALYZER))); return filters; } + + @Nonnull + public static BoolQueryBuilder applyDefaultSearchFilters( + @Nonnull OperationContext opContext, + @Nullable Filter filter, + @Nonnull BoolQueryBuilder filterQuery) { + // filter soft deleted entities by default + filterSoftDeletedByDefault(filter, filterQuery, opContext.getSearchContext().getSearchFlags()); + // filter based on access controls + ESAccessControlUtil.buildAccessControlFilters(opContext).ifPresent(filterQuery::filter); + return filterQuery; + } + + /** + * Applies a default filter to remove entities that are soft deleted only if there isn't a filter + * for the REMOVED field already and soft delete entities are not being requested via search flags + */ + private static void filterSoftDeletedByDefault( + @Nullable Filter filter, + @Nonnull BoolQueryBuilder filterQuery, + @Nonnull SearchFlags searchFlags) { + if (Boolean.FALSE.equals(searchFlags.isIncludeSoftDeleted())) { + boolean removedInOrFilter = false; + if (filter != null) { + removedInOrFilter = + filter.getOr().stream() + .anyMatch( + or -> + or.getAnd().stream() + .anyMatch( + criterion -> + criterion.getField().equals(REMOVED) + || criterion.getField().equals(REMOVED + KEYWORD_SUFFIX))); + } + if (!removedInOrFilter) { + filterQuery.mustNot(QueryBuilders.matchQuery(REMOVED, true)); + } + } + } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/SearchUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/SearchUtils.java index 13ccfd7f972af3..3ddc004dd9fa9f 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/SearchUtils.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/SearchUtils.java @@ -198,6 +198,14 @@ public static SearchFlags applyDefaultSearchFlags( if (!finalSearchFlags.hasSkipCache() || finalSearchFlags.isSkipCache() == null) { finalSearchFlags.setSkipCache(defaultFlags.isSkipCache()); } + if (!finalSearchFlags.hasIncludeSoftDeleted() + || finalSearchFlags.isIncludeSoftDeleted() == null) { + finalSearchFlags.setIncludeSoftDeleted(defaultFlags.isIncludeSoftDeleted()); + } + if (!finalSearchFlags.hasIncludeRestricted() + || finalSearchFlags.isIncludeRestricted() == null) { + finalSearchFlags.setIncludeRestricted(defaultFlags.isIncludeRestricted()); + } if ((!finalSearchFlags.hasGroupingSpec() || finalSearchFlags.getGroupingSpec() == null) && (defaultFlags.getGroupingSpec() != null)) { finalSearchFlags.setGroupingSpec(defaultFlags.getGroupingSpec()); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/indexbuilder/MappingsBuilder.java b/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/indexbuilder/MappingsBuilder.java index 5bb523c8a8c1ee..b2f1571d99407f 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/indexbuilder/MappingsBuilder.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/indexbuilder/MappingsBuilder.java @@ -11,6 +11,9 @@ public class MappingsBuilder { public static final String URN_FIELD = "urn"; + public static final String TAGS_FIELD = "tags"; + public static final String OWNERS_FIELD = "owners"; + public static final String DOMAINS_FIELD = "domains"; public static final String MESSAGE_ID_FIELD = "messageId"; public static final String TIMESTAMP_FIELD = "@timestamp"; public static final String TIMESTAMP_MILLIS_FIELD = "timestampMillis"; diff --git a/metadata-io/src/test/java/com/linkedin/metadata/client/JavaEntityClientTest.java b/metadata-io/src/test/java/com/linkedin/metadata/client/JavaEntityClientTest.java index 5a4443904e2602..3ae2e52fe72089 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/client/JavaEntityClientTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/client/JavaEntityClientTest.java @@ -16,6 +16,7 @@ import com.linkedin.metadata.service.RollbackService; import com.linkedin.metadata.timeseries.TimeseriesAspectService; import com.linkedin.metadata.utils.metrics.MetricUtils; +import io.datahubproject.metadata.context.OperationContext; import java.util.function.Supplier; import org.mockito.MockedStatic; import org.testng.annotations.AfterMethod; @@ -35,6 +36,7 @@ public class JavaEntityClientTest { private MockedStatic _metricUtils; private RollbackService rollbackService; private Counter _counter; + private OperationContext opContext; @BeforeMethod public void setupTest() { @@ -50,6 +52,7 @@ public void setupTest() { _metricUtils = mockStatic(MetricUtils.class); _counter = mock(Counter.class); when(MetricUtils.counter(any(), any())).thenReturn(_counter); + opContext = mock(OperationContext.class); } @AfterMethod @@ -59,6 +62,7 @@ public void closeTest() { private JavaEntityClient getJavaEntityClient() { return new JavaEntityClient( + opContext, _entityService, _deleteEntityService, _entitySearchService, diff --git a/metadata-io/src/test/java/com/linkedin/metadata/graph/search/SearchGraphServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/graph/search/SearchGraphServiceTestBase.java index 7f0e4294e0cbf6..71f247ebfc29ae 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/graph/search/SearchGraphServiceTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/graph/search/SearchGraphServiceTestBase.java @@ -54,7 +54,7 @@ public abstract class SearchGraphServiceTestBase extends GraphServiceTestBase { @Nonnull protected abstract ESIndexBuilder getIndexBuilder(); - private final IndexConvention _indexConvention = new IndexConventionImpl(null); + private final IndexConvention _indexConvention = IndexConventionImpl.NO_PREFIX; private final String _indexName = _indexConvention.getIndexName(INDEX_NAME); private ElasticSearchGraphService _client; private boolean _enableMultiPathSearch = diff --git a/metadata-io/src/test/java/com/linkedin/metadata/recommendation/RecommendationsServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/recommendation/RecommendationsServiceTest.java index c0faf6fdfee6cd..f117c42572bd5b 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/recommendation/RecommendationsServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/recommendation/RecommendationsServiceTest.java @@ -1,14 +1,17 @@ package com.linkedin.metadata.recommendation; +import static org.mockito.Mockito.mock; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; import com.linkedin.metadata.TestEntityUtil; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.recommendation.candidatesource.TestSource; import com.linkedin.metadata.recommendation.ranker.RecommendationModuleRanker; import com.linkedin.metadata.recommendation.ranker.SimpleRecommendationRanker; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.net.URISyntaxException; import java.util.List; import java.util.stream.Collectors; @@ -79,7 +82,8 @@ public void testService() throws URISyntaxException { new RecommendationsService(ImmutableList.of(nonEligibleSource, emptySource), ranker); List result = service.listRecommendations( - Urn.createFromString("urn:li:corpuser:me"), + TestOperationContexts.userContextNoSearchAuthorization( + mock(EntityRegistry.class), Urn.createFromString("urn:li:corpuser:me")), new RecommendationRequestContext().setScenario(ScenarioType.HOME), 10); assertTrue(result.isEmpty()); @@ -90,7 +94,8 @@ public void testService() throws URISyntaxException { ImmutableList.of(nonEligibleSource, emptySource, valuesSource), ranker); result = service.listRecommendations( - Urn.createFromString("urn:li:corpuser:me"), + TestOperationContexts.userContextNoSearchAuthorization( + mock(EntityRegistry.class), Urn.createFromString("urn:li:corpuser:me")), new RecommendationRequestContext().setScenario(ScenarioType.HOME), 10); assertEquals(result.size(), 1); @@ -106,7 +111,8 @@ public void testService() throws URISyntaxException { ImmutableList.of(valuesSource, multiValuesSource, urnsSource, multiUrnsSource), ranker); result = service.listRecommendations( - Urn.createFromString("urn:li:corpuser:me"), + TestOperationContexts.userContextNoSearchAuthorization( + mock(EntityRegistry.class), Urn.createFromString("urn:li:corpuser:me")), new RecommendationRequestContext().setScenario(ScenarioType.HOME), 10); assertEquals(result.size(), 4); @@ -134,7 +140,8 @@ public void testService() throws URISyntaxException { // Test limit result = service.listRecommendations( - Urn.createFromString("urn:li:corpuser:me"), + TestOperationContexts.userContextNoSearchAuthorization( + mock(EntityRegistry.class), Urn.createFromString("urn:li:corpuser:me")), new RecommendationRequestContext().setScenario(ScenarioType.HOME), 2); assertEquals(result.size(), 2); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/EntitySearchAggregationCandidateSourceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/EntitySearchAggregationCandidateSourceTest.java index 2d60f3202b69f5..eb616ee15a292a 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/EntitySearchAggregationCandidateSourceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/EntitySearchAggregationCandidateSourceTest.java @@ -3,6 +3,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotNull; @@ -20,6 +21,8 @@ import com.linkedin.metadata.recommendation.RecommendationRequestContext; import com.linkedin.metadata.recommendation.ScenarioType; import com.linkedin.metadata.search.EntitySearchService; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.util.Collections; import java.util.List; import javax.annotation.Nonnull; @@ -30,10 +33,11 @@ import org.testng.annotations.Test; public class EntitySearchAggregationCandidateSourceTest { - private EntitySearchService _entitySearchService = Mockito.mock(EntitySearchService.class); - private EntityRegistry entityRegistry = Mockito.mock(EntityRegistry.class); + private EntitySearchService _entitySearchService = mock(EntitySearchService.class); + private EntityRegistry entityRegistry = mock(EntityRegistry.class); private EntitySearchAggregationSource _valueBasedCandidateSource; private EntitySearchAggregationSource _urnBasedCandidateSource; + private OperationContext opContext; private static final Urn USER = new CorpuserUrn("test"); private static final RecommendationRequestContext CONTEXT = @@ -41,6 +45,7 @@ public class EntitySearchAggregationCandidateSourceTest { @BeforeMethod public void setup() { + opContext = TestOperationContexts.userContextNoSearchAuthorization(entityRegistry, USER); Mockito.reset(_entitySearchService); _valueBasedCandidateSource = buildCandidateSource("testValue", false); _urnBasedCandidateSource = buildCandidateSource("testUrn", true); @@ -81,7 +86,8 @@ public RecommendationRenderType getRenderType() { @Override public boolean isEligible( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, + @Nonnull RecommendationRequestContext requestContext) { return true; } }; @@ -90,21 +96,24 @@ public boolean isEligible( @Test public void testWhenSearchServiceReturnsEmpty() { Mockito.when( - _entitySearchService.aggregateByValue(eq(null), eq("testValue"), eq(null), anyInt())) + _entitySearchService.aggregateByValue( + any(OperationContext.class), eq(null), eq("testValue"), eq(null), anyInt())) .thenReturn(Collections.emptyMap()); List candidates = - _valueBasedCandidateSource.getRecommendations(USER, CONTEXT); + _valueBasedCandidateSource.getRecommendations(opContext, CONTEXT); assertTrue(candidates.isEmpty()); - assertFalse(_valueBasedCandidateSource.getRecommendationModule(USER, CONTEXT).isPresent()); + assertFalse(_valueBasedCandidateSource.getRecommendationModule(opContext, CONTEXT).isPresent()); } @Test public void testWhenSearchServiceReturnsValueResults() { // One result - Mockito.when(_entitySearchService.aggregateByValue(any(), eq("testValue"), eq(null), anyInt())) + Mockito.when( + _entitySearchService.aggregateByValue( + any(OperationContext.class), any(), eq("testValue"), eq(null), anyInt())) .thenReturn(ImmutableMap.of("value1", 1L)); List candidates = - _valueBasedCandidateSource.getRecommendations(USER, CONTEXT); + _valueBasedCandidateSource.getRecommendations(opContext, CONTEXT); assertEquals(candidates.size(), 1); RecommendationContent content = candidates.get(0); assertEquals(content.getValue(), "value1"); @@ -119,12 +128,14 @@ public void testWhenSearchServiceReturnsValueResults() { new Criterion().setField("testValue").setValue("value1")); assertNotNull(params.getContentParams()); assertEquals(params.getContentParams().getCount().longValue(), 1L); - assertTrue(_valueBasedCandidateSource.getRecommendationModule(USER, CONTEXT).isPresent()); + assertTrue(_valueBasedCandidateSource.getRecommendationModule(opContext, CONTEXT).isPresent()); // Multiple result - Mockito.when(_entitySearchService.aggregateByValue(any(), eq("testValue"), eq(null), anyInt())) + Mockito.when( + _entitySearchService.aggregateByValue( + any(OperationContext.class), any(), eq("testValue"), eq(null), anyInt())) .thenReturn(ImmutableMap.of("value1", 1L, "value2", 2L, "value3", 3L)); - candidates = _valueBasedCandidateSource.getRecommendations(USER, CONTEXT); + candidates = _valueBasedCandidateSource.getRecommendations(opContext, CONTEXT); assertEquals(candidates.size(), 2); content = candidates.get(0); assertEquals(content.getValue(), "value3"); @@ -152,7 +163,7 @@ public void testWhenSearchServiceReturnsValueResults() { new Criterion().setField("testValue").setValue("value2")); assertNotNull(params.getContentParams()); assertEquals(params.getContentParams().getCount().longValue(), 2L); - assertTrue(_valueBasedCandidateSource.getRecommendationModule(USER, CONTEXT).isPresent()); + assertTrue(_valueBasedCandidateSource.getRecommendationModule(opContext, CONTEXT).isPresent()); } @Test @@ -161,10 +172,12 @@ public void testWhenSearchServiceReturnsUrnResults() { Urn testUrn1 = new TestEntityUrn("testUrn1", "testUrn1", "testUrn1"); Urn testUrn2 = new TestEntityUrn("testUrn2", "testUrn2", "testUrn2"); Urn testUrn3 = new TestEntityUrn("testUrn3", "testUrn3", "testUrn3"); - Mockito.when(_entitySearchService.aggregateByValue(any(), eq("testUrn"), eq(null), anyInt())) + Mockito.when( + _entitySearchService.aggregateByValue( + any(OperationContext.class), any(), eq("testUrn"), eq(null), anyInt())) .thenReturn(ImmutableMap.of(testUrn1.toString(), 1L)); List candidates = - _urnBasedCandidateSource.getRecommendations(USER, CONTEXT); + _urnBasedCandidateSource.getRecommendations(opContext, CONTEXT); assertEquals(candidates.size(), 1); RecommendationContent content = candidates.get(0); assertEquals(content.getValue(), testUrn1.toString()); @@ -179,14 +192,16 @@ public void testWhenSearchServiceReturnsUrnResults() { new Criterion().setField("testUrn").setValue(testUrn1.toString())); assertNotNull(params.getContentParams()); assertEquals(params.getContentParams().getCount().longValue(), 1L); - assertTrue(_urnBasedCandidateSource.getRecommendationModule(USER, CONTEXT).isPresent()); + assertTrue(_urnBasedCandidateSource.getRecommendationModule(opContext, CONTEXT).isPresent()); // Multiple result - Mockito.when(_entitySearchService.aggregateByValue(any(), eq("testUrn"), eq(null), anyInt())) + Mockito.when( + _entitySearchService.aggregateByValue( + any(OperationContext.class), any(), eq("testUrn"), eq(null), anyInt())) .thenReturn( ImmutableMap.of( testUrn1.toString(), 1L, testUrn2.toString(), 2L, testUrn3.toString(), 3L)); - candidates = _urnBasedCandidateSource.getRecommendations(USER, CONTEXT); + candidates = _urnBasedCandidateSource.getRecommendations(opContext, CONTEXT); assertEquals(candidates.size(), 2); content = candidates.get(0); assertEquals(content.getValue(), testUrn3.toString()); @@ -214,6 +229,6 @@ public void testWhenSearchServiceReturnsUrnResults() { new Criterion().setField("testUrn").setValue(testUrn2.toString())); assertNotNull(params.getContentParams()); assertEquals(params.getContentParams().getCount().longValue(), 2L); - assertTrue(_urnBasedCandidateSource.getRecommendationModule(USER, CONTEXT).isPresent()); + assertTrue(_urnBasedCandidateSource.getRecommendationModule(opContext, CONTEXT).isPresent()); } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationUtilsTest.java b/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationUtilsTest.java index 3998e45195b254..51b4ecf1410a1f 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationUtilsTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationUtilsTest.java @@ -1,9 +1,14 @@ package com.linkedin.metadata.recommendation.candidatesource; +import static org.mockito.Mockito.mock; + import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.metadata.Constants; +import com.linkedin.metadata.models.registry.EntityRegistry; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.util.Collections; import org.junit.Assert; import org.testng.annotations.Test; @@ -13,13 +18,17 @@ public class RecommendationUtilsTest { @Test private void testIsSupportedEntityType() { Urn testUrn = UrnUtils.getUrn("urn:li:corpuser:john"); + OperationContext opContext = + TestOperationContexts.userContextNoSearchAuthorization(mock(EntityRegistry.class), testUrn); + Assert.assertTrue( RecommendationUtils.isSupportedEntityType( - testUrn, + opContext, ImmutableSet.of(Constants.DATASET_ENTITY_NAME, Constants.CORP_USER_ENTITY_NAME))); Assert.assertFalse( RecommendationUtils.isSupportedEntityType( - testUrn, ImmutableSet.of(Constants.DATASET_ENTITY_NAME))); - Assert.assertFalse(RecommendationUtils.isSupportedEntityType(testUrn, Collections.emptySet())); + opContext, ImmutableSet.of(Constants.DATASET_ENTITY_NAME))); + Assert.assertFalse( + RecommendationUtils.isSupportedEntityType(opContext, Collections.emptySet())); } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/TestSource.java b/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/TestSource.java index 666deb2c419d73..4350cdd2662a89 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/TestSource.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/recommendation/candidatesource/TestSource.java @@ -1,9 +1,9 @@ package com.linkedin.metadata.recommendation.candidatesource; -import com.linkedin.common.urn.Urn; import com.linkedin.metadata.recommendation.RecommendationContent; import com.linkedin.metadata.recommendation.RecommendationRenderType; import com.linkedin.metadata.recommendation.RecommendationRequestContext; +import io.datahubproject.metadata.context.OperationContext; import java.util.List; import javax.annotation.Nonnull; import lombok.Getter; @@ -36,13 +36,13 @@ public RecommendationRenderType getRenderType() { @Override public boolean isEligible( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { return eligible; } @Override public List getRecommendations( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { return contents; } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchResultCacheKeyTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchResultCacheKeyTest.java index 57fa51ffbdd906..3c03f0b201aef0 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchResultCacheKeyTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchResultCacheKeyTest.java @@ -13,8 +13,8 @@ public class LineageSearchResultCacheKeyTest extends AbstractTestNGSpringContext public void testNulls() { // ensure no NPE assertEquals( - new EntityLineageResultCacheKey(null, null, null, null, null, ChronoUnit.DAYS), - new EntityLineageResultCacheKey(null, null, null, null, null, ChronoUnit.DAYS)); + new EntityLineageResultCacheKey("", null, null, null, null, null, ChronoUnit.DAYS), + new EntityLineageResultCacheKey("", null, null, null, null, null, ChronoUnit.DAYS)); } @Test @@ -22,13 +22,13 @@ public void testDateTruncation() { // expect start of day milli assertEquals( new EntityLineageResultCacheKey( - null, null, 1679529600000L, 1679615999999L, null, ChronoUnit.DAYS), + "", null, null, 1679529600000L, 1679615999999L, null, ChronoUnit.DAYS), new EntityLineageResultCacheKey( - null, null, 1679530293000L, 1679530293001L, null, ChronoUnit.DAYS)); + "", null, null, 1679530293000L, 1679530293001L, null, ChronoUnit.DAYS)); assertNotSame( new EntityLineageResultCacheKey( - null, null, 1679529600000L, 1679616000000L, null, ChronoUnit.DAYS), + "", null, null, 1679529600000L, 1679616000000L, null, ChronoUnit.DAYS), new EntityLineageResultCacheKey( - null, null, 1679530293000L, 1679530293001L, null, ChronoUnit.DAYS)); + "", null, null, 1679530293000L, 1679530293001L, null, ChronoUnit.DAYS)); } } 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 9588140bebd65f..52f91fb1b8c28a 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 @@ -17,6 +17,7 @@ import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; +import com.datahub.plugins.auth.authorization.Authorizer; import com.datahub.test.Snapshot; import com.fasterxml.jackson.databind.node.JsonNodeFactory; import com.fasterxml.jackson.databind.node.ObjectNode; @@ -42,7 +43,6 @@ import com.linkedin.metadata.graph.LineageRelationship; import com.linkedin.metadata.graph.LineageRelationshipArray; import com.linkedin.metadata.models.registry.SnapshotEntityRegistry; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; @@ -64,6 +64,8 @@ import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.net.URISyntaxException; import java.util.ArrayList; import java.util.Collections; @@ -74,6 +76,7 @@ import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; +import lombok.Getter; import org.junit.Assert; import org.mockito.ArgumentCaptor; import org.mockito.Mockito; @@ -112,6 +115,8 @@ public abstract class LineageServiceTestBase extends AbstractTestNGSpringContext private LineageSearchService lineageSearchService; private RestHighLevelClient searchClientSpy; + @Getter private OperationContext operationContext; + private static final String ENTITY_NAME = "testEntity"; private static final Urn TEST_URN = TestEntityUtil.getTestEntityUrn(); private static final String TEST = "test"; @@ -133,6 +138,10 @@ public void setup() throws RemoteInvocationException, URISyntaxException { .thenReturn(new SnapshotEntityRegistry(new Snapshot())); when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); indexConvention = new IndexConventionImpl("lineage_search_service_test"); + operationContext = + TestOperationContexts.systemContextNoSearchAuthorization( + aspectRetriever.getEntityRegistry(), indexConvention) + .asSession(Authorizer.EMPTY, TestOperationContexts.TEST_USER_AUTH); settingsBuilder = new SettingsBuilder(null); elasticSearchService = buildEntitySearchService(); elasticSearchService.configure(); @@ -368,6 +377,7 @@ public void testSearchService() throws Exception { searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(true)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -378,8 +388,7 @@ public void testSearchService() throws Exception { 0, 10, null, - null, - new SearchFlags().setSkipCache(false)); + null); assertEquals(searchResult.getNumEntities().intValue(), 1); Mockito.verify(graphService, times(1)) @@ -395,6 +404,7 @@ public void testSearchService() throws Exception { // Hit the cache on second attempt searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -405,8 +415,7 @@ public void testSearchService() throws Exception { 0, 10, null, - null, - new SearchFlags().setSkipCache(false)); + null); assertEquals(searchResult.getNumEntities().intValue(), 1); Mockito.verify(graphService, times(1)) .getLineage( @@ -434,6 +443,7 @@ public void testSearchService() throws Exception { searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), @@ -444,8 +454,7 @@ public void testSearchService() throws Exception { 0, 10, 0L, - 1L, - new SearchFlags().setSkipCache(false)); + 1L); assertEquals(searchResult.getNumEntities().intValue(), 1); Mockito.verify(graphService, times(1)) @@ -461,6 +470,7 @@ public void testSearchService() throws Exception { // Hit the cache on second attempt searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -471,8 +481,7 @@ public void testSearchService() throws Exception { 0, 10, 0L, - 1L, - new SearchFlags().setSkipCache(false)); + 1L); assertEquals(searchResult.getNumEntities().intValue(), 1); Mockito.verify(graphService, times(1)) .getLineage( @@ -764,6 +773,7 @@ public void testLightningSearchService() throws Exception { searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -774,8 +784,7 @@ public void testLightningSearchService() throws Exception { 0, 10, null, - null, - new SearchFlags().setSkipCache(false)); + null); assertEquals(searchResult.getNumEntities().intValue(), 1); verify(graphService, times(1)) @@ -793,6 +802,7 @@ public void testLightningSearchService() throws Exception { // Hit the cache on second attempt searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -803,8 +813,7 @@ public void testLightningSearchService() throws Exception { 0, 10, null, - null, - new SearchFlags().setSkipCache(false)); + null); assertEquals(searchResult.getNumEntities().intValue(), 1); verify(graphService, times(1)) .getLineage( @@ -834,6 +843,7 @@ public void testLightningSearchService() throws Exception { searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), @@ -844,8 +854,7 @@ public void testLightningSearchService() throws Exception { 0, 10, 0L, - 1L, - new SearchFlags().setSkipCache(false)); + 1L); assertEquals(searchResult.getNumEntities().intValue(), 1); verify(graphService, times(1)) @@ -863,6 +872,7 @@ public void testLightningSearchService() throws Exception { // Hit the cache on second attempt searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -873,8 +883,7 @@ public void testLightningSearchService() throws Exception { 0, 10, 0L, - 1L, - new SearchFlags().setSkipCache(false)); + 1L); assertEquals(searchResult.getNumEntities().intValue(), 1); verify(graphService, times(1)) .getLineage( @@ -896,6 +905,7 @@ public void testLightningSearchService() throws Exception { // Entity searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(DATASET_ENTITY_NAME), @@ -906,8 +916,7 @@ public void testLightningSearchService() throws Exception { 0, 10, null, - null, - new SearchFlags().setSkipCache(false)); + null); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(searchResult.getEntities().size(), 0); verify(lineageSearchService, times(1)) @@ -916,6 +925,7 @@ public void testLightningSearchService() throws Exception { // Cached searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(DATASET_ENTITY_NAME), @@ -926,8 +936,7 @@ public void testLightningSearchService() throws Exception { 0, 10, null, - null, - new SearchFlags().setSkipCache(false)); + null); Mockito.verify(graphService, times(1)) .getLineage( eq(TEST_URN), @@ -959,6 +968,7 @@ public void testLightningSearchService() throws Exception { searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -969,8 +979,7 @@ public void testLightningSearchService() throws Exception { 0, 10, null, - null, - new SearchFlags().setSkipCache(false)); + null); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(searchResult.getEntities().size(), 0); verify(lineageSearchService, times(3)) @@ -979,6 +988,7 @@ public void testLightningSearchService() throws Exception { // Cached searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -989,8 +999,7 @@ public void testLightningSearchService() throws Exception { 0, 10, null, - null, - new SearchFlags().setSkipCache(false)); + null); verify(graphService, times(1)) .getLineage( eq(TEST_URN), @@ -1009,6 +1018,7 @@ public void testLightningSearchService() throws Exception { Filter originFilter = QueryUtils.newFilter("origin", "PROD"); searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -1019,8 +1029,7 @@ public void testLightningSearchService() throws Exception { 0, 10, null, - null, - new SearchFlags().setSkipCache(false)); + null); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(searchResult.getEntities().size(), 0); verify(lineageSearchService, times(5)) @@ -1029,6 +1038,7 @@ public void testLightningSearchService() throws Exception { // Cached searchResult = lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(false)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -1039,8 +1049,7 @@ public void testLightningSearchService() throws Exception { 0, 10, null, - null, - new SearchFlags().setSkipCache(false)); + null); verify(graphService, times(1)) .getLineage( eq(TEST_URN), @@ -1284,6 +1293,7 @@ private LineageRelationship constructLineageRelationship(Urn urn) { // Convenience method to reduce spots where we're sending the same params private LineageSearchResult searchAcrossLineage(@Nullable Filter filter, @Nullable String input) { return lineageSearchService.searchAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(true)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), @@ -1294,13 +1304,13 @@ private LineageSearchResult searchAcrossLineage(@Nullable Filter filter, @Nullab 0, 10, null, - null, - new SearchFlags().setSkipCache(true)); + null); } private LineageScrollResult scrollAcrossLineage( @Nullable Filter filter, @Nullable String input, String scrollId, int size) { return lineageSearchService.scrollAcrossLineage( + getOperationContext().withSearchFlags(flags -> flags.setSkipCache(true)), TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), @@ -1312,8 +1322,7 @@ private LineageScrollResult scrollAcrossLineage( "5m", size, null, - null, - new SearchFlags().setSkipCache(true)); + null); } private LineageScrollResult scrollAcrossLineage(@Nullable Filter filter, @Nullable String input) { diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java index d860776a316815..3b233ed8ad7100 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java @@ -7,6 +7,7 @@ import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; +import com.datahub.plugins.auth.authorization.Authorizer; import com.datahub.test.Snapshot; import com.fasterxml.jackson.databind.node.JsonNodeFactory; import com.fasterxml.jackson.databind.node.ObjectNode; @@ -19,7 +20,6 @@ import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration; import com.linkedin.metadata.models.registry.SnapshotEntityRegistry; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; @@ -40,9 +40,12 @@ import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.net.URISyntaxException; import java.util.Map; import javax.annotation.Nonnull; +import lombok.Getter; import org.opensearch.client.RestHighLevelClient; import org.springframework.cache.CacheManager; import org.springframework.cache.concurrent.ConcurrentMapCacheManager; @@ -74,6 +77,7 @@ public abstract class SearchServiceTestBase extends AbstractTestNGSpringContextT private ElasticSearchService elasticSearchService; private CacheManager cacheManager; private SearchService searchService; + @Getter private OperationContext operationContext; private static final String ENTITY_NAME = "testEntity"; @@ -84,6 +88,10 @@ public void setup() throws RemoteInvocationException, URISyntaxException { .thenReturn(new SnapshotEntityRegistry(new Snapshot())); when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); indexConvention = new IndexConventionImpl("search_service_test"); + operationContext = + TestOperationContexts.systemContextNoSearchAuthorization( + aspectRetriever.getEntityRegistry(), indexConvention) + .asSession(Authorizer.EMPTY, TestOperationContexts.TEST_USER_AUTH); settingsBuilder = new SettingsBuilder(null); elasticSearchService = buildEntitySearchService(); elasticSearchService.configure(); @@ -156,17 +164,24 @@ private void clearCache() { public void testSearchService() throws Exception { SearchResult searchResult = searchService.searchAcrossEntities( + getOperationContext() + .withSearchFlags(flags -> flags.setFulltext(true).setSkipCache(true)), ImmutableList.of(ENTITY_NAME), "test", null, null, 0, - 10, - new SearchFlags().setFulltext(true).setSkipCache(true)); + 10); assertEquals(searchResult.getNumEntities().intValue(), 0); searchResult = searchService.searchAcrossEntities( - ImmutableList.of(), "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); + getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), + ImmutableList.of(), + "test", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 0); clearCache(); @@ -181,7 +196,13 @@ public void testSearchService() throws Exception { searchResult = searchService.searchAcrossEntities( - ImmutableList.of(), "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); + getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), + ImmutableList.of(), + "test", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); clearCache(); @@ -197,12 +218,20 @@ public void testSearchService() throws Exception { searchResult = searchService.searchAcrossEntities( - ImmutableList.of(), "'test2'", null, null, 0, 10, new SearchFlags().setFulltext(true)); + getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), + ImmutableList.of(), + "'test2'", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn2); clearCache(); - long docCount = elasticSearchService.docCount(ENTITY_NAME); + long docCount = + elasticSearchService.docCount( + getOperationContext().withSearchFlags(flags -> flags.setFulltext(false)), ENTITY_NAME); assertEquals(docCount, 2L); elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); @@ -210,7 +239,13 @@ public void testSearchService() throws Exception { syncAfterWrite(getBulkProcessor()); searchResult = searchService.searchAcrossEntities( - ImmutableList.of(), "'test2'", null, null, 0, 10, new SearchFlags().setFulltext(true)); + getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), + ImmutableList.of(), + "'test2'", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 0); } @@ -241,13 +276,13 @@ public void testAdvancedSearchOr() throws Exception { SearchResult searchResult = searchService.searchAcrossEntities( + getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), ImmutableList.of(ENTITY_NAME), "test", filterWithCondition, null, 0, - 10, - new SearchFlags().setFulltext(true)); + 10); assertEquals(searchResult.getNumEntities().intValue(), 0); clearCache(); @@ -286,13 +321,13 @@ public void testAdvancedSearchOr() throws Exception { searchResult = searchService.searchAcrossEntities( + getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), ImmutableList.of(), "test", filterWithCondition, null, 0, - 10, - new SearchFlags().setFulltext(true)); + 10); assertEquals(searchResult.getNumEntities().intValue(), 2); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); assertEquals(searchResult.getEntities().get(1).getEntity(), urn2); @@ -326,13 +361,13 @@ public void testAdvancedSearchSoftDelete() throws Exception { SearchResult searchResult = searchService.searchAcrossEntities( + getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), ImmutableList.of(ENTITY_NAME), "test", filterWithCondition, null, 0, - 10, - new SearchFlags().setFulltext(true)); + 10); assertEquals(searchResult.getNumEntities().intValue(), 0); clearCache(); @@ -374,13 +409,13 @@ public void testAdvancedSearchSoftDelete() throws Exception { searchResult = searchService.searchAcrossEntities( + getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), ImmutableList.of(), "test", filterWithCondition, null, 0, - 10, - new SearchFlags().setFulltext(true)); + 10); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); clearCache(); @@ -405,13 +440,13 @@ public void testAdvancedSearchNegated() throws Exception { SearchResult searchResult = searchService.searchAcrossEntities( + getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), ImmutableList.of(ENTITY_NAME), "test", filterWithCondition, null, 0, - 10, - new SearchFlags().setFulltext(true)); + 10); assertEquals(searchResult.getNumEntities().intValue(), 0); clearCache(); @@ -453,13 +488,13 @@ public void testAdvancedSearchNegated() throws Exception { searchResult = searchService.searchAcrossEntities( + getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)), ImmutableList.of(), "test", filterWithCondition, null, 0, - 10, - new SearchFlags().setFulltext(true)); + 10); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn3); clearCache(); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/TestEntityTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/TestEntityTestBase.java index 40ccc8dfb5047e..ae91f6a8876e9f 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/TestEntityTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/TestEntityTestBase.java @@ -14,7 +14,6 @@ import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.elasticsearch.ElasticSearchService; import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder; import com.linkedin.metadata.search.elasticsearch.indexbuilder.EntityIndexBuilders; @@ -25,6 +24,8 @@ import com.linkedin.metadata.search.elasticsearch.update.ESWriteDAO; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.util.List; import javax.annotation.Nonnull; import org.opensearch.client.RestHighLevelClient; @@ -60,12 +61,16 @@ public abstract class TestEntityTestBase extends AbstractTestNGSpringContextTest private IndexConvention indexConvention; private SettingsBuilder settingsBuilder; private ElasticSearchService elasticSearchService; + private OperationContext opContext; private static final String ENTITY_NAME = "testEntity"; @BeforeClass public void setup() { indexConvention = new IndexConventionImpl("es_service_test"); + opContext = + TestOperationContexts.systemContextNoSearchAuthorization( + aspectRetriever.getEntityRegistry(), indexConvention); settingsBuilder = new SettingsBuilder(null); elasticSearchService = buildService(); elasticSearchService.configure(); @@ -122,14 +127,35 @@ private ElasticSearchService buildService() { public void testElasticSearchServiceStructuredQuery() throws Exception { SearchResult searchResult = elasticSearchService.search( - List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(false)); + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + List.of(ENTITY_NAME), + "test", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 0); - BrowseResult browseResult = elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); + BrowseResult browseResult = + elasticSearchService.browse( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ENTITY_NAME, + "", + null, + 0, + 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 0); - assertEquals(elasticSearchService.docCount(ENTITY_NAME), 0); + assertEquals( + elasticSearchService.docCount( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), ENTITY_NAME), + 0); assertEquals( elasticSearchService - .aggregateByValue(ImmutableList.of(ENTITY_NAME), "textField", null, 10) + .aggregateByValue( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ImmutableList.of(ENTITY_NAME), + "textField", + null, + 10) .size(), 0); @@ -145,30 +171,57 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { searchResult = elasticSearchService.search( - List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(false)); + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + List.of(ENTITY_NAME), + "test", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); searchResult = elasticSearchService.search( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), List.of(ENTITY_NAME), "foreignKey:Node", null, null, 0, - 10, - new SearchFlags().setFulltext(false)); + 10); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); - browseResult = elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); + browseResult = + elasticSearchService.browse( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ENTITY_NAME, + "", + null, + 0, + 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 1); assertEquals(browseResult.getGroups().get(0).getName(), "a"); - browseResult = elasticSearchService.browse(ENTITY_NAME, "/a", null, 0, 10); + browseResult = + elasticSearchService.browse( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ENTITY_NAME, + "/a", + null, + 0, + 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 1); assertEquals(browseResult.getGroups().get(0).getName(), "b"); - assertEquals(elasticSearchService.docCount(ENTITY_NAME), 1); + assertEquals( + elasticSearchService.docCount( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), ENTITY_NAME), + 1); assertEquals( elasticSearchService.aggregateByValue( - ImmutableList.of(ENTITY_NAME), "textFieldOverride", null, 10), + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ImmutableList.of(ENTITY_NAME), + "textFieldOverride", + null, + 10), ImmutableMap.of("textFieldOverride", 1L)); Urn urn2 = new TestEntityUrn("test2", "urn2", "VALUE_2"); @@ -182,20 +235,47 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { searchResult = elasticSearchService.search( - List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(false)); + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + List.of(ENTITY_NAME), + "test2", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn2); - browseResult = elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); + browseResult = + elasticSearchService.browse( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ENTITY_NAME, + "", + null, + 0, + 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 2); assertEquals(browseResult.getGroups().get(0).getName(), "a"); assertEquals(browseResult.getGroups().get(1).getName(), "b"); - browseResult = elasticSearchService.browse(ENTITY_NAME, "/a", null, 0, 10); + browseResult = + elasticSearchService.browse( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ENTITY_NAME, + "/a", + null, + 0, + 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 1); assertEquals(browseResult.getGroups().get(0).getName(), "b"); - assertEquals(elasticSearchService.docCount(ENTITY_NAME), 2); + assertEquals( + elasticSearchService.docCount( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), ENTITY_NAME), + 2); assertEquals( elasticSearchService.aggregateByValue( - ImmutableList.of(ENTITY_NAME), "textFieldOverride", null, 10), + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ImmutableList.of(ENTITY_NAME), + "textFieldOverride", + null, + 10), ImmutableMap.of("textFieldOverride", 1L, "textFieldOverride2", 1L)); elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); @@ -203,14 +283,35 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { syncAfterWrite(getBulkProcessor()); searchResult = elasticSearchService.search( - List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(false)); + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + List.of(ENTITY_NAME), + "test2", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 0); - browseResult = elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); + browseResult = + elasticSearchService.browse( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ENTITY_NAME, + "", + null, + 0, + 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 0); - assertEquals(elasticSearchService.docCount(ENTITY_NAME), 0); + assertEquals( + elasticSearchService.docCount( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), ENTITY_NAME), + 0); assertEquals( elasticSearchService - .aggregateByValue(ImmutableList.of(ENTITY_NAME), "textField", null, 10) + .aggregateByValue( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ImmutableList.of(ENTITY_NAME), + "textField", + null, + 10) .size(), 0); } @@ -219,7 +320,13 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { public void testElasticSearchServiceFulltext() throws Exception { SearchResult searchResult = elasticSearchService.search( - List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); + opContext.withSearchFlags(flags -> flags.setFulltext(true)), + List.of(ENTITY_NAME), + "test", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 0); Urn urn = new TestEntityUrn("test", "urn1", "VALUE_1"); @@ -234,14 +341,27 @@ public void testElasticSearchServiceFulltext() throws Exception { searchResult = elasticSearchService.search( - List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); + opContext.withSearchFlags(flags -> flags.setFulltext(true)), + List.of(ENTITY_NAME), + "test", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); - assertEquals(elasticSearchService.docCount(ENTITY_NAME), 1); + assertEquals( + elasticSearchService.docCount( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), ENTITY_NAME), + 1); assertEquals( elasticSearchService.aggregateByValue( - ImmutableList.of(ENTITY_NAME), "textFieldOverride", null, 10), + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ImmutableList.of(ENTITY_NAME), + "textFieldOverride", + null, + 10), ImmutableMap.of("textFieldOverride", 1L)); Urn urn2 = new TestEntityUrn("test2", "urn2", "VALUE_2"); @@ -255,14 +375,27 @@ public void testElasticSearchServiceFulltext() throws Exception { searchResult = elasticSearchService.search( - List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(true)); + opContext.withSearchFlags(flags -> flags.setFulltext(true)), + List.of(ENTITY_NAME), + "test2", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn2); - assertEquals(elasticSearchService.docCount(ENTITY_NAME), 2); + assertEquals( + elasticSearchService.docCount( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), ENTITY_NAME), + 2); assertEquals( elasticSearchService.aggregateByValue( - ImmutableList.of(ENTITY_NAME), "textFieldOverride", null, 10), + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ImmutableList.of(ENTITY_NAME), + "textFieldOverride", + null, + 10), ImmutableMap.of("textFieldOverride", 1L, "textFieldOverride2", 1L)); elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); @@ -270,13 +403,27 @@ public void testElasticSearchServiceFulltext() throws Exception { syncAfterWrite(getBulkProcessor()); searchResult = elasticSearchService.search( - List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(true)); + opContext.withSearchFlags(flags -> flags.setFulltext(true)), + List.of(ENTITY_NAME), + "test2", + null, + null, + 0, + 10); assertEquals(searchResult.getNumEntities().intValue(), 0); - assertEquals(elasticSearchService.docCount(ENTITY_NAME), 0); + assertEquals( + elasticSearchService.docCount( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), ENTITY_NAME), + 0); assertEquals( elasticSearchService - .aggregateByValue(ImmutableList.of(ENTITY_NAME), "textField", null, 10) + .aggregateByValue( + opContext.withSearchFlags(flags -> flags.setFulltext(false)), + ImmutableList.of(ENTITY_NAME), + "textField", + null, + 10) .size(), 0); } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/cache/CacheableSearcherTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/cache/CacheableSearcherTest.java index 175c48e198185b..fc60119f775121 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/cache/CacheableSearcherTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/cache/CacheableSearcherTest.java @@ -1,17 +1,21 @@ package com.linkedin.metadata.search.cache; +import static org.mockito.Mockito.mock; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; import com.google.common.collect.Streams; import com.linkedin.common.urn.TestEntityUrn; import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.AggregationMetadataArray; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.search.SearchResultMetadata; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.util.List; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -26,42 +30,44 @@ public class CacheableSearcherTest { @Test public void testCacheableSearcherWhenEmpty() { + OperationContext opContext = + TestOperationContexts.systemContextNoSearchAuthorization(mock(EntityRegistry.class)); CacheableSearcher emptySearcher = new CacheableSearcher<>( cacheManager.getCache("emptySearcher"), 10, this::getEmptySearchResult, CacheableSearcher.QueryPagination::getFrom, - null, true); - assertTrue(emptySearcher.getSearchResults(0, 0).getEntities().isEmpty()); - assertTrue(emptySearcher.getSearchResults(0, 10).getEntities().isEmpty()); - assertTrue(emptySearcher.getSearchResults(5, 10).getEntities().isEmpty()); + assertTrue(emptySearcher.getSearchResults(opContext, 0, 0).getEntities().isEmpty()); + assertTrue(emptySearcher.getSearchResults(opContext, 0, 10).getEntities().isEmpty()); + assertTrue(emptySearcher.getSearchResults(opContext, 5, 10).getEntities().isEmpty()); } @Test public void testCacheableSearcherWithFixedNumResults() { + OperationContext opContext = + TestOperationContexts.systemContextNoSearchAuthorization(mock(EntityRegistry.class)); CacheableSearcher fixedBatchSearcher = new CacheableSearcher<>( cacheManager.getCache("fixedBatchSearcher"), 10, qs -> getSearchResult(qs, 10), CacheableSearcher.QueryPagination::getFrom, - null, true); - SearchResult result = fixedBatchSearcher.getSearchResults(0, 0); + SearchResult result = fixedBatchSearcher.getSearchResults(opContext, 0, 0); assertTrue(result.getEntities().isEmpty()); assertEquals(result.getNumEntities().intValue(), 1000); - result = fixedBatchSearcher.getSearchResults(0, 10); + result = fixedBatchSearcher.getSearchResults(opContext, 0, 10); assertEquals(result.getNumEntities().intValue(), 1000); assertEquals(result.getEntities().size(), 10); assertEquals( result.getEntities().stream().map(SearchEntity::getEntity).collect(Collectors.toList()), getUrns(0, 10)); - result = fixedBatchSearcher.getSearchResults(5, 10); + result = fixedBatchSearcher.getSearchResults(opContext, 5, 10); assertEquals(result.getNumEntities().intValue(), 1000); assertEquals(result.getEntities().size(), 10); assertEquals( @@ -72,27 +78,28 @@ public void testCacheableSearcherWithFixedNumResults() { @Test public void testCacheableSearcherWithVariableNumResults() { + OperationContext opContext = + TestOperationContexts.systemContextNoSearchAuthorization(mock(EntityRegistry.class)); CacheableSearcher variableBatchSearcher = new CacheableSearcher<>( cacheManager.getCache("variableBatchSearcher"), 10, qs -> getSearchResult(qs, qs.getFrom() + qs.getSize()), CacheableSearcher.QueryPagination::getFrom, - null, true); - SearchResult result = variableBatchSearcher.getSearchResults(0, 0); + SearchResult result = variableBatchSearcher.getSearchResults(opContext, 0, 0); assertTrue(result.getEntities().isEmpty()); assertEquals(result.getNumEntities().intValue(), 1000); - result = variableBatchSearcher.getSearchResults(0, 10); + result = variableBatchSearcher.getSearchResults(opContext, 0, 10); assertEquals(result.getNumEntities().intValue(), 1000); assertEquals(result.getEntities().size(), 10); assertEquals( result.getEntities().stream().map(SearchEntity::getEntity).collect(Collectors.toList()), getUrns(0, 10)); - result = variableBatchSearcher.getSearchResults(5, 10); + result = variableBatchSearcher.getSearchResults(opContext, 5, 10); assertEquals(result.getNumEntities().intValue(), 1000); assertEquals(result.getEntities().size(), 10); assertEquals( @@ -100,7 +107,7 @@ public void testCacheableSearcherWithVariableNumResults() { Streams.concat(getUrns(5, 10).stream(), getUrns(0, 5).stream()) .collect(Collectors.toList())); - result = variableBatchSearcher.getSearchResults(5, 100); + result = variableBatchSearcher.getSearchResults(opContext, 5, 100); assertEquals(result.getNumEntities().intValue(), 1000); assertEquals(result.getEntities().size(), 100); assertEquals( @@ -116,17 +123,18 @@ public void testCacheableSearcherWithVariableNumResults() { @Test public void testCacheableSearcherEnabled() { + OperationContext opContext = + TestOperationContexts.systemContextNoSearchAuthorization(mock(EntityRegistry.class)); // Verify cache is not interacted with when cache disabled - Cache mockCache = Mockito.mock(Cache.class); + Cache mockCache = mock(Cache.class); CacheableSearcher cacheDisabled = new CacheableSearcher<>( mockCache, 10, qs -> getSearchResult(qs, qs.getFrom() + qs.getSize()), CacheableSearcher.QueryPagination::getFrom, - null, false); - SearchResult result = cacheDisabled.getSearchResults(0, 10); + SearchResult result = cacheDisabled.getSearchResults(opContext, 0, 10); assertEquals(result.getNumEntities().intValue(), 1000); assertEquals(result.getEntities().size(), 10); assertEquals( @@ -142,9 +150,10 @@ public void testCacheableSearcherEnabled() { 10, qs -> getSearchResult(qs, qs.getFrom() + qs.getSize()), CacheableSearcher.QueryPagination::getFrom, - new SearchFlags().setSkipCache(true), true); - result = skipCache.getSearchResults(0, 10); + result = + skipCache.getSearchResults( + opContext.withSearchFlags(flags -> flags.setSkipCache(true)), 0, 10); assertEquals(result.getNumEntities().intValue(), 1000); assertEquals(result.getEntities().size(), 10); assertEquals( @@ -161,9 +170,9 @@ public void testCacheableSearcherEnabled() { 10, qs -> getSearchResult(qs, qs.getFrom() + qs.getSize()), CacheableSearcher.QueryPagination::getFrom, - null, true); - result = nullFlags.getSearchResults(0, 10); + result = + nullFlags.getSearchResults(opContext.withSearchFlags(flags -> new SearchFlags()), 0, 10); assertEquals(result.getNumEntities().intValue(), 1000); assertEquals(result.getEntities().size(), 10); assertEquals( @@ -180,9 +189,10 @@ public void testCacheableSearcherEnabled() { 10, qs -> getSearchResult(qs, qs.getFrom() + qs.getSize()), CacheableSearcher.QueryPagination::getFrom, - new SearchFlags().setSkipCache(false), true); - result = useCache.getSearchResults(0, 10); + result = + useCache.getSearchResults( + opContext.withSearchFlags(flags -> flags.setSkipCache(false)), 0, 10); assertEquals(result.getNumEntities().intValue(), 1000); assertEquals(result.getEntities().size(), 10); assertEquals( diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/LineageDataFixtureElasticSearchTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/LineageDataFixtureElasticSearchTest.java index 1fed3380a342d5..4c9c69a52debd5 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/LineageDataFixtureElasticSearchTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/LineageDataFixtureElasticSearchTest.java @@ -1,5 +1,6 @@ package com.linkedin.metadata.search.elasticsearch; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.LineageSearchService; import com.linkedin.metadata.search.SearchService; import com.linkedin.metadata.search.fixtures.LineageDataFixtureTestBase; @@ -27,6 +28,8 @@ public class LineageDataFixtureElasticSearchTest extends LineageDataFixtureTestB @Qualifier("searchLineageLineageSearchService") protected LineageSearchService lineageService; + @Autowired protected EntityRegistry entityRegistry; + @NotNull @Override protected LineageSearchService getLineageService() { @@ -43,4 +46,9 @@ protected SearchService getSearchService() { public void initTest() { AssertJUnit.assertNotNull(lineageService); } + + @Override + protected EntityRegistry getEntityRegistry() { + return entityRegistry; + } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/GoldenTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/GoldenTestBase.java index 4c125065deb4da..d384d8275bdd76 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/GoldenTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/GoldenTestBase.java @@ -1,11 +1,11 @@ package com.linkedin.metadata.search.fixtures; import static com.linkedin.metadata.Constants.*; -import static io.datahubproject.test.search.SearchTestUtils.searchAcrossCustomEntities; import static io.datahubproject.test.search.SearchTestUtils.searchAcrossEntities; import static org.testng.Assert.*; import static org.testng.AssertJUnit.assertNotNull; +import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; import com.linkedin.data.template.StringArray; @@ -21,6 +21,9 @@ import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.search.SearchService; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; +import io.datahubproject.test.search.SearchTestUtils; import java.util.Collections; import java.util.List; import java.util.stream.Collectors; @@ -48,6 +51,12 @@ public abstract class GoldenTestBase extends AbstractTestNGSpringContextTests { @Nonnull protected abstract SearchService getSearchService(); + @Nonnull + protected OperationContext getOperationContext() { + return TestOperationContexts.userContextNoSearchAuthorization( + getEntityRegistry(), Authorizer.EMPTY, TestOperationContexts.TEST_USER_AUTH); + } + @Test public void testNameMatchPetProfiles() { /* @@ -56,8 +65,11 @@ public void testNameMatchPetProfiles() { assertNotNull(getSearchService()); assertNotNull(getEntityRegistry()); SearchResult searchResult = - searchAcrossCustomEntities( - getSearchService(), "pet profiles", SEARCHABLE_LONGTAIL_ENTITIES); + searchAcrossEntities( + getOperationContext(), + getSearchService(), + SEARCHABLE_LONGTAIL_ENTITIES, + "pet profiles"); assertTrue(searchResult.getEntities().size() >= 2); Urn firstResultUrn = searchResult.getEntities().get(0).getEntity(); Urn secondResultUrn = searchResult.getEntities().get(1).getEntity(); @@ -73,7 +85,8 @@ public void testNameMatchPetProfile() { */ assertNotNull(getSearchService()); SearchResult searchResult = - searchAcrossEntities(getSearchService(), "pet profile", SEARCHABLE_LONGTAIL_ENTITIES); + searchAcrossEntities( + getOperationContext(), getSearchService(), SEARCHABLE_LONGTAIL_ENTITIES, "pet profile"); assertTrue(searchResult.getEntities().size() >= 2); Urn firstResultUrn = searchResult.getEntities().get(0).getEntity(); Urn secondResultUrn = searchResult.getEntities().get(1).getEntity(); @@ -90,7 +103,8 @@ public void testGlossaryTerms() { */ assertNotNull(getSearchService()); SearchResult searchResult = - searchAcrossEntities(getSearchService(), "ReturnRate", SEARCHABLE_LONGTAIL_ENTITIES); + searchAcrossEntities( + getOperationContext(), getSearchService(), SEARCHABLE_LONGTAIL_ENTITIES, "ReturnRate"); SearchEntityArray entities = searchResult.getEntities(); assertTrue(searchResult.getEntities().size() >= 4); MatchedFieldArray firstResultMatchedFields = entities.get(0).getMatchedFields(); @@ -113,7 +127,10 @@ public void testNameMatchPartiallyQualified() { assertNotNull(getSearchService()); SearchResult searchResult = searchAcrossEntities( - getSearchService(), "analytics.pet_details", SEARCHABLE_LONGTAIL_ENTITIES); + getOperationContext(), + getSearchService(), + SEARCHABLE_LONGTAIL_ENTITIES, + "analytics.pet_details"); assertTrue(searchResult.getEntities().size() >= 2); Urn firstResultUrn = searchResult.getEntities().get(0).getEntity(); Urn secondResultUrn = searchResult.getEntities().get(1).getEntity(); @@ -133,7 +150,10 @@ public void testNameMatchCollaborativeActionitems() { assertNotNull(getSearchService()); SearchResult searchResult = searchAcrossEntities( - getSearchService(), "collaborative actionitems", SEARCHABLE_LONGTAIL_ENTITIES); + getOperationContext(), + getSearchService(), + SEARCHABLE_LONGTAIL_ENTITIES, + "collaborative actionitems"); assertTrue(searchResult.getEntities().size() >= 2); Urn firstResultUrn = searchResult.getEntities().get(0).getEntity(); Urn secondResultUrn = searchResult.getEntities().get(1).getEntity(); @@ -158,7 +178,11 @@ public void testNameMatchCustomerOrders() { */ assertNotNull(getSearchService()); SearchResult searchResult = - searchAcrossEntities(getSearchService(), "customer orders", SEARCHABLE_LONGTAIL_ENTITIES); + searchAcrossEntities( + getOperationContext(), + getSearchService(), + SEARCHABLE_LONGTAIL_ENTITIES, + "customer orders"); assertTrue(searchResult.getEntities().size() >= 2); Urn firstResultUrn = searchResult.getEntities().get(0).getEntity(); @@ -194,12 +218,13 @@ public void testFilterOnCountField() { .setValue("") .setValues(new StringArray(ImmutableList.of("68")))))))); SearchResult searchResult = - searchAcrossEntities( + SearchTestUtils.facetAcrossEntities( + getOperationContext(), getSearchService(), - "*", SEARCHABLE_LONGTAIL_ENTITIES, - filter, - Collections.singletonList(DATASET_ENTITY_NAME)); + "*", + Collections.singletonList(DATASET_ENTITY_NAME), + filter); assertFalse(searchResult.getEntities().isEmpty()); Urn firstResultUrn = searchResult.getEntities().get(0).getEntity(); assertEquals( diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/LineageDataFixtureTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/LineageDataFixtureTestBase.java index 59942f76744dab..6950f62d45263e 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/LineageDataFixtureTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/fixtures/LineageDataFixtureTestBase.java @@ -5,12 +5,16 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; +import com.datahub.plugins.auth.authorization.Authorizer; import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.LineageSearchResult; import com.linkedin.metadata.search.LineageSearchService; import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.search.SearchService; import com.linkedin.util.Pair; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.net.URISyntaxException; import java.util.stream.Stream; import javax.annotation.Nonnull; @@ -25,15 +29,27 @@ public abstract class LineageDataFixtureTestBase extends AbstractTestNGSpringCon @Nonnull protected abstract SearchService getSearchService(); + @Nonnull + protected abstract EntityRegistry getEntityRegistry(); + + @Nonnull + protected OperationContext getOperationContext() { + return TestOperationContexts.userContextNoSearchAuthorization( + getEntityRegistry(), Authorizer.EMPTY, TestOperationContexts.TEST_USER_AUTH); + } + @Test public void testFixtureInitialization() { assertNotNull(getSearchService()); - SearchResult noResult = searchAcrossEntities(getSearchService(), "no results"); + SearchResult noResult = + searchAcrossEntities(getOperationContext(), getSearchService(), "no results"); assertEquals(noResult.getEntities().size(), 0); SearchResult result = searchAcrossEntities( - getSearchService(), "e3859789eed1cef55288b44f016ee08290d9fd08973e565c112d8"); + getOperationContext(), + getSearchService(), + "e3859789eed1cef55288b44f016ee08290d9fd08973e565c112d8"); assertEquals(result.getEntities().size(), 1); assertEquals( @@ -41,7 +57,8 @@ public void testFixtureInitialization() { "urn:li:dataset:(urn:li:dataPlatform:9cf8c96,e3859789eed1cef55288b44f016ee08290d9fd08973e565c112d8,PROD)"); LineageSearchResult lineageResult = - lineage(getLineageService(), result.getEntities().get(0).getEntity(), 1); + lineage( + getOperationContext(), getLineageService(), result.getEntities().get(0).getEntity(), 1); assertEquals(lineageResult.getEntities().size(), 10); } @@ -60,7 +77,11 @@ public void testDatasetLineage() throws URISyntaxException { hopsExpectedResultsStream.forEach( hopsExpectedResults -> { LineageSearchResult lineageResult = - lineage(getLineageService(), testUrn, hopsExpectedResults.getFirst()); + lineage( + getOperationContext(), + getLineageService(), + testUrn, + hopsExpectedResults.getFirst()); assertEquals(lineageResult.getEntities().size(), hopsExpectedResults.getSecond()); }); } 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 4742115b16e1bd..deda14c2216b1d 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 @@ -14,6 +14,7 @@ import com.datahub.authentication.Actor; import com.datahub.authentication.ActorType; import com.datahub.authentication.Authentication; +import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.linkedin.common.urn.Urn; @@ -28,7 +29,6 @@ import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.SearchableFieldSpec; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; @@ -45,6 +45,8 @@ import com.linkedin.metadata.search.elasticsearch.query.request.SearchFieldConfig; import com.linkedin.metadata.search.utils.ESUtils; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.io.IOException; import java.util.ArrayList; import java.util.Collections; @@ -86,6 +88,12 @@ public abstract class SampleDataFixtureTestBase extends AbstractTestNGSpringCont @Nonnull protected abstract RestHighLevelClient getSearchClient(); + @Nonnull + protected OperationContext getOperationContext() { + return TestOperationContexts.userContextNoSearchAuthorization( + getEntityRegistry(), Authorizer.EMPTY, AUTHENTICATION); + } + @Test public void testSearchFieldConfig() throws IOException { /* @@ -272,10 +280,12 @@ public void testDatasetHasTags() throws IOException { @Test public void testFixtureInitialization() { assertNotNull(getSearchService()); - SearchResult noResult = searchAcrossEntities(getSearchService(), "no results"); + SearchResult noResult = + searchAcrossEntities(getOperationContext(), getSearchService(), "no results"); assertEquals(0, noResult.getEntities().size()); - final SearchResult result = searchAcrossEntities(getSearchService(), "test"); + final SearchResult result = + searchAcrossEntities(getOperationContext(), getSearchService(), "test"); Map expectedTypes = Map.of( @@ -333,7 +343,8 @@ public void testDataPlatform() { expected.forEach( (key, value) -> { - SearchResult result = searchAcrossEntities(getSearchService(), key); + SearchResult result = + searchAcrossEntities(getOperationContext(), getSearchService(), key); assertEquals( result.getEntities().size(), value.intValue(), @@ -354,13 +365,17 @@ public void testUrn() { .forEach( query -> assertTrue( - searchAcrossEntities(getSearchService(), query).getEntities().size() >= 1, + searchAcrossEntities(getOperationContext(), getSearchService(), query) + .getEntities() + .size() + >= 1, String.format("Unexpected >1 urn result for `%s`", query))); } @Test public void testExactTable() { - SearchResult results = searchAcrossEntities(getSearchService(), "stg_customers"); + SearchResult results = + searchAcrossEntities(getOperationContext(), getSearchService(), "stg_customers"); assertEquals( results.getEntities().size(), 1, "Unexpected single urn result for `stg_customers`"); assertEquals( @@ -380,7 +395,8 @@ public void testStemming() { testSet -> { Integer expectedResults = null; for (String testQuery : testSet) { - SearchResult results = searchAcrossEntities(getSearchService(), testQuery); + SearchResult results = + searchAcrossEntities(getOperationContext(), getSearchService(), testQuery); assertTrue( results.hasEntities() && !results.getEntities().isEmpty(), @@ -402,7 +418,7 @@ public void testStemmingOverride() throws IOException { Set results = testSet.stream() - .map(test -> searchAcrossEntities(getSearchService(), test)) + .map(test -> searchAcrossEntities(getOperationContext(), getSearchService(), test)) .collect(Collectors.toSet()); results.forEach( @@ -456,7 +472,8 @@ public void testDelimitedSynonym() throws IOException { testSet.stream() .map( q -> { - SearchResult result = searchAcrossEntities(getSearchService(), q); + SearchResult result = + searchAcrossEntities(getOperationContext(), getSearchService(), q); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), "Expected search results for: " + q); @@ -592,7 +609,8 @@ public void testUrnSynonym() throws IOException { testSet.stream() .map( query -> { - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = + searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), "Expected search results for: " + query); @@ -824,7 +842,8 @@ public void testChartAutoComplete() throws InterruptedException, IOException { .forEach( query -> { try { - AutoCompleteResults result = autocomplete(new ChartType(getEntityClient()), query); + AutoCompleteResults result = + autocomplete(getOperationContext(), new ChartType(getEntityClient()), query); assertTrue( result.getEntities().size() == 2, String.format( @@ -853,7 +872,7 @@ public void testDatasetAutoComplete() { query -> { try { AutoCompleteResults result = - autocomplete(new DatasetType(getEntityClient()), query); + autocomplete(getOperationContext(), new DatasetType(getEntityClient()), query); assertTrue( result.getEntities().size() >= 1, String.format( @@ -879,7 +898,8 @@ public void testContainerAutoComplete() { query -> { try { AutoCompleteResults result = - autocomplete(new ContainerType(getEntityClient()), query); + autocomplete( + getOperationContext(), new ContainerType(getEntityClient()), query); assertTrue( result.getEntities().size() >= 1, String.format( @@ -898,7 +918,8 @@ public void testGroupAutoComplete() { query -> { try { AutoCompleteResults result = - autocomplete(new CorpGroupType(getEntityClient()), query); + autocomplete( + getOperationContext(), new CorpGroupType(getEntityClient()), query); assertTrue( result.getEntities().size() == 1, String.format( @@ -917,7 +938,8 @@ public void testUserAutoComplete() { query -> { try { AutoCompleteResults result = - autocomplete(new CorpUserType(getEntityClient(), null), query); + autocomplete( + getOperationContext(), new CorpUserType(getEntityClient(), null), query); assertTrue( result.getEntities().size() >= 1, String.format( @@ -953,7 +975,9 @@ public void testSmokeTestQueries() { .collect( Collectors.toMap( Map.Entry::getKey, - entry -> searchAcrossEntities(getSearchService(), entry.getKey()))); + entry -> + searchAcrossEntities( + getOperationContext(), getSearchService(), entry.getKey()))); results.forEach( (key, value) -> { @@ -978,7 +1002,9 @@ public void testSmokeTestQueries() { .collect( Collectors.toMap( Map.Entry::getKey, - entry -> searchStructured(getSearchService(), entry.getKey()))); + entry -> + searchStructured( + getOperationContext(), getSearchService(), entry.getKey()))); results.forEach( (key, value) -> { @@ -1033,7 +1059,8 @@ public void testUnderscore() throws IOException { @Test public void testFacets() { Set expectedFacets = Set.of("entity", "typeNames", "platform", "origin", "tags"); - SearchResult testResult = searchAcrossEntities(getSearchService(), "cypress"); + SearchResult testResult = + searchAcrossEntities(getOperationContext(), getSearchService(), "cypress"); expectedFacets.forEach( facet -> { assertTrue( @@ -1076,7 +1103,8 @@ public void testFacets() { public void testNestedAggregation() { Set expectedFacets = Set.of("platform"); SearchResult testResult = - searchAcrossEntities(getSearchService(), "cypress", List.copyOf(expectedFacets)); + facetAcrossEntities( + getOperationContext(), getSearchService(), "cypress", List.copyOf(expectedFacets)); assertEquals(testResult.getMetadata().getAggregations().size(), 1); expectedFacets.forEach( facet -> { @@ -1093,7 +1121,8 @@ public void testNestedAggregation() { expectedFacets = Set.of("platform", "typeNames", "_entityType", "entity"); SearchResult testResult2 = - searchAcrossEntities(getSearchService(), "cypress", List.copyOf(expectedFacets)); + facetAcrossEntities( + getOperationContext(), getSearchService(), "cypress", List.copyOf(expectedFacets)); assertEquals(testResult2.getMetadata().getAggregations().size(), 4); expectedFacets.forEach( facet -> { @@ -1140,7 +1169,8 @@ public void testNestedAggregation() { expectedFacets = Set.of("platform", "typeNames", "entity"); SearchResult testResult3 = - searchAcrossEntities(getSearchService(), "cypress", List.copyOf(expectedFacets)); + facetAcrossEntities( + getOperationContext(), getSearchService(), "cypress", List.copyOf(expectedFacets)); assertEquals(testResult3.getMetadata().getAggregations().size(), 4); expectedFacets.forEach( facet -> { @@ -1170,7 +1200,8 @@ public void testNestedAggregation() { String singleNestedFacet = String.format("_entityType%sowners", AGGREGATION_SEPARATOR_CHAR); expectedFacets = Set.of(singleNestedFacet); SearchResult testResultSingleNested = - searchAcrossEntities(getSearchService(), "cypress", List.copyOf(expectedFacets)); + facetAcrossEntities( + getOperationContext(), getSearchService(), "cypress", List.copyOf(expectedFacets)); assertEquals(testResultSingleNested.getMetadata().getAggregations().size(), 1); Map expectedNestedFacetCounts = new HashMap<>(); expectedNestedFacetCounts.put("datajob␞urn:li:corpuser:datahub", 2L); @@ -1192,7 +1223,8 @@ public void testNestedAggregation() { expectedFacets = Set.of("platform", singleNestedFacet, "typeNames", "origin"); SearchResult testResultNested = - searchAcrossEntities(getSearchService(), "cypress", List.copyOf(expectedFacets)); + facetAcrossEntities( + getOperationContext(), getSearchService(), "cypress", List.copyOf(expectedFacets)); assertEquals(testResultNested.getMetadata().getAggregations().size(), 4); expectedFacets.forEach( facet -> { @@ -1304,7 +1336,8 @@ public void testScrollAcrossEntities() throws IOException { int totalResults = 0; String scrollId = null; do { - ScrollResult result = scroll(getSearchService(), query, batchSize, scrollId); + ScrollResult result = + scroll(getOperationContext(), getSearchService(), query, batchSize, scrollId); int numResults = result.hasEntities() ? result.getEntities().size() : 0; assertTrue(numResults <= batchSize); totalResults += numResults; @@ -1317,13 +1350,19 @@ public void testScrollAcrossEntities() throws IOException { @Test public void testSearchAcrossMultipleEntities() { String query = "logging_events"; - SearchResult result = search(getSearchService(), query); + SearchResult result = search(getOperationContext(), getSearchService(), query); assertEquals((int) result.getNumEntities(), 8); - result = search(getSearchService(), List.of(DATASET_ENTITY_NAME, DATA_JOB_ENTITY_NAME), query); + result = + search( + getOperationContext(), + getSearchService(), + List.of(DATASET_ENTITY_NAME, DATA_JOB_ENTITY_NAME), + query); assertEquals((int) result.getNumEntities(), 8); - result = search(getSearchService(), List.of(DATASET_ENTITY_NAME), query); + result = search(getOperationContext(), getSearchService(), List.of(DATASET_ENTITY_NAME), query); assertEquals((int) result.getNumEntities(), 4); - result = search(getSearchService(), List.of(DATA_JOB_ENTITY_NAME), query); + result = + search(getOperationContext(), getSearchService(), List.of(DATA_JOB_ENTITY_NAME), query); assertEquals((int) result.getNumEntities(), 4); } @@ -1387,7 +1426,8 @@ public void testFragmentUrns() { testSet.forEach( query -> { - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = + searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1410,7 +1450,8 @@ public void testPlatformTest() { fieldName -> { final String query = String.format("%s:%s", fieldName, testPlatform.replaceAll(":", "\\\\:")); - SearchResult result = searchStructured(getSearchService(), query); + SearchResult result = + searchStructured(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1462,14 +1503,14 @@ public void testPlatformTest() { try { return getEntityClient() .search( + getOperationContext() + .withSearchFlags(flags -> flags.setFulltext(fulltextFlag)), "dataset", "*", filter, null, 0, - 100, - AUTHENTICATION, - new SearchFlags().setFulltext(fulltextFlag)); + 100); } catch (RemoteInvocationException e) { throw new RuntimeException(e); } @@ -1492,7 +1533,7 @@ public void testPlatformTest() { @Test public void testStructQueryFieldMatch() { String query = STRUCTURED_QUERY_PREFIX + "name: customers"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1507,7 +1548,7 @@ public void testStructQueryFieldMatch() { @Test public void testStructQueryFieldPrefixMatch() { String query = STRUCTURED_QUERY_PREFIX + "name: customers*"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1522,7 +1563,7 @@ public void testStructQueryFieldPrefixMatch() { @Test public void testStructQueryCustomPropertiesKeyPrefix() { String query = STRUCTURED_QUERY_PREFIX + "customProperties: node_type=*"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1537,7 +1578,7 @@ public void testStructQueryCustomPropertiesKeyPrefix() { @Test public void testStructQueryCustomPropertiesMatch() { String query = STRUCTURED_QUERY_PREFIX + "customProperties: node_type=model"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1563,7 +1604,9 @@ public void testCustomPropertiesQuoted() { .collect( Collectors.toMap( Map.Entry::getKey, - entry -> searchAcrossEntities(getSearchService(), entry.getKey()))); + entry -> + searchAcrossEntities( + getOperationContext(), getSearchService(), entry.getKey()))); results.forEach( (key, value) -> { @@ -1581,7 +1624,7 @@ public void testCustomPropertiesQuoted() { @Test public void testStructQueryFieldPaths() { String query = STRUCTURED_QUERY_PREFIX + "fieldPaths: customer_id"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1598,7 +1641,7 @@ public void testStructQueryBoolean() { String query = STRUCTURED_QUERY_PREFIX + "editedFieldTags:urn\\:li\\:tag\\:Legacy OR tags:urn\\:li\\:tag\\:testTag"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1610,7 +1653,7 @@ public void testStructQueryBoolean() { assertEquals(result.getEntities().size(), 2); query = STRUCTURED_QUERY_PREFIX + "editedFieldTags:urn\\:li\\:tag\\:Legacy"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1622,7 +1665,7 @@ public void testStructQueryBoolean() { assertEquals(result.getEntities().size(), 1); query = STRUCTURED_QUERY_PREFIX + "tags:urn\\:li\\:tag\\:testTag"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1637,7 +1680,7 @@ public void testStructQueryBoolean() { @Test public void testStructQueryBrowsePaths() { String query = STRUCTURED_QUERY_PREFIX + "browsePaths:*/dbt/*"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1652,7 +1695,7 @@ public void testStructQueryBrowsePaths() { @Test public void testOr() { String query = "stg_customers | logging_events"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1662,7 +1705,7 @@ public void testOr() { assertEquals(result.getEntities().size(), 9); query = "stg_customers"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1672,7 +1715,7 @@ public void testOr() { assertEquals(result.getEntities().size(), 1); query = "logging_events"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1685,7 +1728,7 @@ public void testOr() { @Test public void testNegate() { String query = "logging_events -bckp"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1695,7 +1738,7 @@ public void testNegate() { assertEquals(result.getEntities().size(), 7); query = "logging_events"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1708,7 +1751,7 @@ public void testNegate() { @Test public void testPrefix() { String query = "bigquery"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1718,7 +1761,7 @@ public void testPrefix() { assertEquals(result.getEntities().size(), 8); query = "big*"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1731,7 +1774,7 @@ public void testPrefix() { @Test public void testParens() { String query = "dbt | (bigquery + covid19)"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1741,7 +1784,7 @@ public void testParens() { assertEquals(result.getEntities().size(), 11); query = "dbt"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1751,7 +1794,7 @@ public void testParens() { assertEquals(result.getEntities().size(), 9); query = "bigquery + covid19"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1761,7 +1804,7 @@ public void testParens() { assertEquals(result.getEntities().size(), 2); query = "bigquery"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1771,7 +1814,7 @@ public void testParens() { assertEquals(result.getEntities().size(), 8); query = "covid19"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1784,7 +1827,7 @@ public void testParens() { @Test public void testGram() { String query = "jaffle shop customers"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1795,7 +1838,7 @@ public void testGram() { "Expected exact match in 1st position"); query = "shop customers source"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1806,7 +1849,7 @@ public void testGram() { "Expected ngram match in 1st position"); query = "jaffle shop stg customers"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1817,7 +1860,7 @@ public void testGram() { "Expected ngram match in 1st position"); query = "jaffle shop transformers customers"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1828,7 +1871,7 @@ public void testGram() { "Expected ngram match in 1st position"); query = "shop raw customers"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1842,7 +1885,7 @@ public void testGram() { @Test public void testPrefixVsExact() { String query = "\"customers\""; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1865,7 +1908,7 @@ public void testPrefixVsExactCaseSensitivity() { List insensitiveExactMatches = List.of("testExactMatchCase", "testexactmatchcase", "TESTEXACTMATCHCASE"); for (String query : insensitiveExactMatches) { - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), @@ -1885,7 +1928,7 @@ public void testPrefixVsExactCaseSensitivity() { @Test public void testColumnExactMatch() { String query = "unit_data"; - SearchResult result = searchAcrossEntities(getSearchService(), query); + SearchResult result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1902,7 +1945,7 @@ public void testColumnExactMatch() { "Expected table name exact match first"); query = "special_column_only_present_here_info"; - result = searchAcrossEntities(getSearchService(), query); + result = searchAcrossEntities(getOperationContext(), getSearchService(), query); assertTrue( result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1927,13 +1970,14 @@ public void testSortOrdering() { SearchResult result = getSearchService() .searchAcrossEntities( + getOperationContext() + .withSearchFlags(flags -> flags.setFulltext(true).setSkipCache(true)), SEARCHABLE_ENTITIES, query, null, criterion, 0, 100, - new SearchFlags().setFulltext(true).setSkipCache(true), null); assertTrue( result.getEntities().size() > 2, @@ -1957,11 +2001,11 @@ public void testFilterOnHasValuesField() { .setValues(new StringArray(ImmutableList.of("true")))))))); SearchResult searchResult = searchAcrossEntities( + getOperationContext(), getSearchService(), + Collections.singletonList(DATASET_ENTITY_NAME), "*", - SEARCHABLE_ENTITIES, - filter, - Collections.singletonList(DATASET_ENTITY_NAME)); + filter); assertEquals(searchResult.getEntities().size(), 8); } @@ -1982,11 +2026,11 @@ public void testFilterOnNumValuesField() { .setValues(new StringArray(ImmutableList.of("1")))))))); SearchResult searchResult = searchAcrossEntities( + getOperationContext(), getSearchService(), + Collections.singletonList(DATA_JOB_ENTITY_NAME), "*", - SEARCHABLE_ENTITIES, - filter, - Collections.singletonList(DATA_JOB_ENTITY_NAME)); + filter); assertEquals(searchResult.getEntities().size(), 4); } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/opensearch/LineageDataFixtureOpenSearchTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/opensearch/LineageDataFixtureOpenSearchTest.java index cc17e3287544ce..c67920a8ec1155 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/opensearch/LineageDataFixtureOpenSearchTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/opensearch/LineageDataFixtureOpenSearchTest.java @@ -1,5 +1,6 @@ package com.linkedin.metadata.search.opensearch; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.LineageSearchService; import com.linkedin.metadata.search.SearchService; import com.linkedin.metadata.search.fixtures.LineageDataFixtureTestBase; @@ -27,6 +28,10 @@ public class LineageDataFixtureOpenSearchTest extends LineageDataFixtureTestBase @Qualifier("searchLineageLineageSearchService") protected LineageSearchService lineageService; + @Autowired + @Qualifier("entityRegistry") + protected EntityRegistry entityRegistry; + @NotNull @Override protected LineageSearchService getLineageService() { @@ -43,4 +48,9 @@ protected SearchService getSearchService() { public void initTest() { AssertJUnit.assertNotNull(lineageService); } + + @Override + protected EntityRegistry getEntityRegistry() { + return entityRegistry; + } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java index b43d1556a68829..264aa280cac903 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java @@ -2,6 +2,7 @@ import static com.linkedin.metadata.Constants.*; import static com.linkedin.metadata.utils.SearchUtil.AGGREGATION_SEPARATOR_CHAR; +import static com.linkedin.metadata.utils.SearchUtil.ES_INDEX_FIELD; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -11,6 +12,7 @@ import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; +import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.linkedin.data.template.LongMap; import com.linkedin.data.template.StringArray; @@ -29,16 +31,20 @@ import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.search.SearchResultMetadata; +import com.linkedin.metadata.search.elasticsearch.ElasticSearchService; import com.linkedin.metadata.search.elasticsearch.query.ESSearchDAO; import com.linkedin.metadata.search.opensearch.SearchDAOOpenSearchTest; import com.linkedin.metadata.utils.SearchUtil; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.net.URISyntaxException; import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; +import lombok.Getter; import org.opensearch.action.explain.ExplainResponse; import org.opensearch.client.RestHighLevelClient; import org.springframework.test.context.testng.AbstractTestNGSpringContextTests; @@ -57,11 +63,18 @@ public abstract class SearchDAOTestBase extends AbstractTestNGSpringContextTests protected AspectRetriever aspectRetriever; + @Getter protected OperationContext operationContext; + @BeforeClass public void setup() throws RemoteInvocationException, URISyntaxException { aspectRetriever = mock(AspectRetriever.class); when(aspectRetriever.getEntityRegistry()).thenReturn(getEntityRegistry()); when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); + operationContext = + TestOperationContexts.userContextNoSearchAuthorization( + aspectRetriever.getEntityRegistry(), + Authorizer.EMPTY, + TestOperationContexts.TEST_USER_AUTH); } @Test @@ -124,7 +137,7 @@ public void testTransformFilterForEntitiesWithChanges() { .setValues(new StringArray(ImmutableList.of("smpldat_datasetindex_v2"))) .setNegated(false) .setCondition(Condition.EQUAL) - .setField("_index"); + .setField(ES_INDEX_FIELD); Filter expectedNewFilter = new Filter() @@ -168,7 +181,7 @@ public void testTransformFilterForEntitiesWithUnderscore() { .setValues(new StringArray(ImmutableList.of("smpldat_datajobindex_v2"))) .setNegated(false) .setCondition(Condition.EQUAL) - .setField("_index"); + .setField(ES_INDEX_FIELD); Filter expectedNewFilter = new Filter() @@ -220,7 +233,7 @@ public void testTransformFilterForEntitiesWithSomeChanges() { .setValues(new StringArray(ImmutableList.of("smpldat_datasetindex_v2"))) .setNegated(false) .setCondition(Condition.EQUAL) - .setField("_index"); + .setField(ES_INDEX_FIELD); Filter expectedNewFilter = new Filter() @@ -462,6 +475,8 @@ public void testExplain() { .setAspectRetriever(aspectRetriever); ExplainResponse explainResponse = searchDAO.explain( + getOperationContext() + .withSearchFlags(flags -> ElasticSearchService.DEFAULT_SERVICE_SEARCH_FLAGS), "*", "urn:li:dataset:(urn:li:dataPlatform:bigquery,bigquery-public-data.covid19_geotab_mobility_impact." + "ca_border_wait_times,PROD)", @@ -470,7 +485,6 @@ public void testExplain() { null, null, null, - null, 10, null); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java index bb37fb3f3b206a..6ec4f4e589f35b 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java @@ -6,7 +6,10 @@ import com.linkedin.metadata.TestEntitySpecBuilder; import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.elasticsearch.query.request.AutocompleteRequestHandler; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.util.List; import java.util.Map; import org.opensearch.action.search.SearchRequest; @@ -22,11 +25,14 @@ public class AutocompleteRequestHandlerTest { private AutocompleteRequestHandler handler = AutocompleteRequestHandler.getBuilder( TestEntitySpecBuilder.getSpec(), mock(AspectRetriever.class)); + private OperationContext mockOpContext = + TestOperationContexts.systemContextNoSearchAuthorization(mock(EntityRegistry.class)); @Test public void testDefaultAutocompleteRequest() { // When field is null - SearchRequest autocompleteRequest = handler.getSearchRequest("input", null, null, 10); + SearchRequest autocompleteRequest = + handler.getSearchRequest(mockOpContext, "input", null, null, 10); SearchSourceBuilder sourceBuilder = autocompleteRequest.source(); assertEquals(sourceBuilder.size(), 10); BoolQueryBuilder query = (BoolQueryBuilder) sourceBuilder.query(); @@ -64,7 +70,8 @@ public void testDefaultAutocompleteRequest() { @Test public void testAutocompleteRequestWithField() { // When field is null - SearchRequest autocompleteRequest = handler.getSearchRequest("input", "field", null, 10); + SearchRequest autocompleteRequest = + handler.getSearchRequest(mockOpContext, "input", "field", null, 10); SearchSourceBuilder sourceBuilder = autocompleteRequest.source(); assertEquals(sourceBuilder.size(), 10); BoolQueryBuilder query = (BoolQueryBuilder) sourceBuilder.query(); 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 14cc9e47913c1e..be128d7855f394 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 @@ -12,7 +12,6 @@ import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.config.search.WordGramConfiguration; import com.linkedin.metadata.models.EntitySpec; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; @@ -20,6 +19,7 @@ import com.linkedin.metadata.query.filter.CriterionArray; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.elasticsearch.query.request.SearchRequestHandler; +import io.datahubproject.metadata.context.OperationContext; import io.datahubproject.test.search.config.SearchCommonTestConfiguration; import java.util.ArrayList; import java.util.Collection; @@ -51,6 +51,8 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests { @Autowired private AspectRetriever aspectRetriever; + @Autowired private OperationContext operationContext; + public static SearchConfiguration testQueryConfig; static { @@ -106,12 +108,13 @@ public void testSearchRequestHandlerHighlightingTurnedOff() { TestEntitySpecBuilder.getSpec(), testQueryConfig, null, aspectRetriever); SearchRequest searchRequest = requestHandler.getSearchRequest( + operationContext.withSearchFlags( + flags -> flags.setFulltext(false).setSkipHighlighting(true)), "testQuery", null, null, 0, 10, - new SearchFlags().setFulltext(false).setSkipHighlighting(true), null); SearchSourceBuilder sourceBuilder = searchRequest.source(); assertEquals(sourceBuilder.from(), 0); @@ -146,7 +149,13 @@ public void testSearchRequestHandler() { TestEntitySpecBuilder.getSpec(), testQueryConfig, null, aspectRetriever); SearchRequest searchRequest = requestHandler.getSearchRequest( - "testQuery", null, null, 0, 10, new SearchFlags().setFulltext(false), null); + operationContext.withSearchFlags(flags -> flags.setFulltext(false)), + "testQuery", + null, + null, + 0, + 10, + null); SearchSourceBuilder sourceBuilder = searchRequest.source(); assertEquals(sourceBuilder.from(), 0); assertEquals(sourceBuilder.size(), 10); @@ -204,12 +213,12 @@ public void testAggregationsInSearch() { String.format("_entityType%stextFieldOverride", AGGREGATION_SEPARATOR_CHAR); SearchRequest searchRequest = requestHandler.getSearchRequest( + operationContext.withSearchFlags(flags -> flags.setFulltext(true)), "*", null, null, 0, 10, - new SearchFlags().setFulltext(true), List.of( "textFieldOverride", "_entityType", @@ -229,12 +238,12 @@ public void testAggregationsInSearch() { .size(testQueryConfig.getMaxTermBucketSize()); AggregationBuilder expectedEntityTypeAggregationBuilder = AggregationBuilders.terms("_entityType") - .field("_index") + .field(ES_INDEX_FIELD) .size(testQueryConfig.getMaxTermBucketSize()) .minDocCount(0); AggregationBuilder expectedNestedAggregationBuilder = AggregationBuilders.terms(nestedAggString) - .field("_index") + .field(ES_INDEX_FIELD) .size(testQueryConfig.getMaxTermBucketSize()) .minDocCount(0) .subAggregation( @@ -305,6 +314,7 @@ private BoolQueryBuilder constructFilterQuery( (BoolQueryBuilder) requestHandler .getSearchRequest( + operationContext.withSearchFlags(flags -> flags.setFulltext(false)), "testQuery", filterWithoutRemovedCondition, null, @@ -312,7 +322,6 @@ private BoolQueryBuilder constructFilterQuery( null, "5m", 10, - new SearchFlags().setFulltext(false), null) .source() .query(); @@ -321,12 +330,12 @@ private BoolQueryBuilder constructFilterQuery( (BoolQueryBuilder) requestHandler .getSearchRequest( + operationContext.withSearchFlags(flags -> flags.setFulltext(false)), "testQuery", filterWithoutRemovedCondition, null, 0, 10, - new SearchFlags().setFulltext(false), null) .source() .query(); @@ -380,6 +389,7 @@ private BoolQueryBuilder constructRemovedQuery( (BoolQueryBuilder) requestHandler .getSearchRequest( + operationContext.withSearchFlags(flags -> flags.setFulltext(false)), "testQuery", filterWithRemovedCondition, null, @@ -387,7 +397,6 @@ private BoolQueryBuilder constructRemovedQuery( null, "5m", 10, - new SearchFlags().setFulltext(false), null) .source() .query(); @@ -396,12 +405,12 @@ private BoolQueryBuilder constructRemovedQuery( (BoolQueryBuilder) requestHandler .getSearchRequest( + operationContext.withSearchFlags(flags -> flags.setFulltext(false)), "testQuery", filterWithRemovedCondition, null, 0, 10, - new SearchFlags().setFulltext(false), null) .source() .query(); @@ -621,7 +630,11 @@ public void testBrowsePathQueryFilter() { filter.setOr(conjunctiveCriterionArray); BoolQueryBuilder test = - SearchRequestHandler.getFilterQuery(filter, new HashMap<>(), aspectRetriever); + SearchRequestHandler.getFilterQuery( + operationContext.withSearchFlags(flags -> flags.setFulltext(false)), + filter, + new HashMap<>(), + aspectRetriever); assertEquals(test.should().size(), 1); @@ -649,7 +662,14 @@ private BoolQueryBuilder getQuery(final Criterion filterCriterion) { return (BoolQueryBuilder) requestHandler - .getSearchRequest("", filter, null, 0, 10, new SearchFlags().setFulltext(false), null) + .getSearchRequest( + operationContext.withSearchFlags(flags -> flags.setFulltext(false)), + "", + filter, + null, + 0, + 10, + null) .source() .query(); } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/utils/ESAccessControlUtilTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/utils/ESAccessControlUtilTest.java new file mode 100644 index 00000000000000..4d97bab1e4214c --- /dev/null +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/utils/ESAccessControlUtilTest.java @@ -0,0 +1,646 @@ +package com.linkedin.metadata.search.utils; + +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; + +import com.datahub.authentication.Actor; +import com.datahub.authentication.ActorType; +import com.datahub.authentication.Authentication; +import com.datahub.authorization.config.SearchAuthorizationConfiguration; +import com.datahub.plugins.auth.authorization.Authorizer; +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.metadata.aspect.hooks.OwnerTypeMap; +import com.linkedin.metadata.authorization.PoliciesConfig; +import com.linkedin.metadata.entity.TestEntityRegistry; +import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import com.linkedin.policy.DataHubActorFilter; +import com.linkedin.policy.DataHubPolicyInfo; +import com.linkedin.policy.DataHubResourceFilter; +import com.linkedin.policy.PolicyMatchCondition; +import com.linkedin.policy.PolicyMatchCriterion; +import com.linkedin.policy.PolicyMatchCriterionArray; +import com.linkedin.policy.PolicyMatchFilter; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.metadata.context.OperationContextConfig; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import org.opensearch.index.query.QueryBuilders; +import org.testng.annotations.Test; + +public class ESAccessControlUtilTest { + private static final Authentication SYSTEM_AUTH = + new Authentication(new Actor(ActorType.USER, "SYSTEM"), ""); + private static final Urn TEST_GROUP_A = UrnUtils.getUrn("urn:li:corpGroup:a"); + private static final Urn TEST_GROUP_B = UrnUtils.getUrn("urn:li:corpGroup:b"); + private static final Urn TEST_GROUP_C = UrnUtils.getUrn("urn:li:corpGroup:c"); + private static final Urn TEST_USER_A = UrnUtils.getUrn("urn:li:corpUser:a"); + private static final Urn TEST_USER_B = UrnUtils.getUrn("urn:li:corpUser:b"); + private static final Urn TECH_OWNER = + UrnUtils.getUrn("urn:li:ownershipType:__system__technical_owner"); + private static final Urn BUS_OWNER = + UrnUtils.getUrn("urn:li:ownershipType:__system__business_owner"); + private static final Authentication USER_AUTH = + new Authentication(new Actor(ActorType.USER, TEST_USER_A.getId()), ""); + private static final OperationContext ENABLED_CONTEXT = + OperationContext.asSystem( + OperationContextConfig.builder() + .allowSystemAuthentication(true) + .searchAuthorizationConfiguration( + SearchAuthorizationConfiguration.builder().enabled(true).build()) + .build(), + new TestEntityRegistry(), + SYSTEM_AUTH, + IndexConventionImpl.NO_PREFIX); + + @Test + public void testAllUserAllGroup() { + OperationContext allUsers = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllUsers(true)) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + OperationContext allGroups = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllGroups(true)) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(allUsers), + Optional.empty(), + "Expected no ES filters for all user access without resource restrictions"); + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(allGroups), + Optional.empty(), + "Expected no ES filters for all user access without resource restrictions"); + } + + @Test + public void testAllUserAllGroupEntityType() { + OperationContext resourceAllUsersPolicy = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllUsers(true)) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + new PolicyMatchCriterion() + .setField("TYPE") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues(new StringArray("dataset", "chart")))))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + OperationContext resourceAllGroupsPolicy = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllGroups(true)) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + new PolicyMatchCriterion() + .setField("TYPE") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues(new StringArray("dataset", "chart")))))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(resourceAllUsersPolicy), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .filter( + QueryBuilders.termsQuery( + "_index", List.of("datasetindex_v2", "chartindex_v2")))) + .minimumShouldMatch(1)), + "Expected index filter for each entity"); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(resourceAllGroupsPolicy), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .filter( + QueryBuilders.termsQuery( + "_index", List.of("datasetindex_v2", "chartindex_v2")))) + .minimumShouldMatch(1)), + "Expected index filter for each entity"); + } + + @Test + public void testAllUserAllGroupUrn() { + OperationContext resourceAllUsersPolicy = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllUsers(true)) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + new PolicyMatchCriterion() + .setField("URN") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues( + new StringArray( + "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.ShelterDogs,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.ecommerce.account,PROD)")))))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + OperationContext resourceAllGroupsPolicy = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllGroups(true)) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + new PolicyMatchCriterion() + .setField("URN") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues( + new StringArray( + "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.ShelterDogs,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.ecommerce.account,PROD)")))))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(resourceAllUsersPolicy), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .filter( + QueryBuilders.termsQuery( + "urn", + List.of( + "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.ShelterDogs,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.ecommerce.account,PROD)")))) + .minimumShouldMatch(1)), + "Expected filter for each urn"); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(resourceAllGroupsPolicy), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .filter( + QueryBuilders.termsQuery( + "urn", + List.of( + "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.ShelterDogs,PROD)", + "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.ecommerce.account,PROD)")))) + .minimumShouldMatch(1)), + "Expected filter for each urn"); + } + + @Test + public void testAllUserAllGroupTag() { + OperationContext resourceAllUsersPolicy = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllUsers(true)) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + new PolicyMatchCriterion() + .setField("TAG") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues( + new StringArray( + "urn:li:tag:pii", "urn:li:tag:prod")))))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + OperationContext resourceAllGroupsPolicy = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllGroups(true)) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + new PolicyMatchCriterion() + .setField("TAG") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues( + new StringArray( + "urn:li:tag:pii", "urn:li:tag:prod")))))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(resourceAllUsersPolicy), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .filter( + QueryBuilders.termsQuery( + "tags.keyword", List.of("urn:li:tag:pii", "urn:li:tag:prod")))) + .minimumShouldMatch(1)), + "Expected filter each tag"); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(resourceAllGroupsPolicy), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .filter( + QueryBuilders.termsQuery( + "tags.keyword", List.of("urn:li:tag:pii", "urn:li:tag:prod")))) + .minimumShouldMatch(1)), + "Expected filter each tag"); + } + + @Test + public void testAllUserAllGroupDomain() { + OperationContext resourceAllUsersPolicy = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllUsers(true)) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + new PolicyMatchCriterion() + .setField("DOMAIN") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues( + new StringArray( + "urn:li:domain:f9229a0b-c5ad-47e7-9ff3-f4248c5cb634", + "urn:li:domain:7d64d0fa-66c3-445c-83db-3a324723daf8")))))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + OperationContext resourceAllGroupsPolicy = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllGroups(true)) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + new PolicyMatchCriterion() + .setField("DOMAIN") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues( + new StringArray( + "urn:li:domain:f9229a0b-c5ad-47e7-9ff3-f4248c5cb634", + "urn:li:domain:7d64d0fa-66c3-445c-83db-3a324723daf8")))))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(resourceAllUsersPolicy), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .filter( + QueryBuilders.termsQuery( + "domains.keyword", + List.of( + "urn:li:domain:f9229a0b-c5ad-47e7-9ff3-f4248c5cb634", + "urn:li:domain:7d64d0fa-66c3-445c-83db-3a324723daf8")))) + .minimumShouldMatch(1)), + "Expected filter each domain"); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(resourceAllGroupsPolicy), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .filter( + QueryBuilders.termsQuery( + "domains.keyword", + List.of( + "urn:li:domain:f9229a0b-c5ad-47e7-9ff3-f4248c5cb634", + "urn:li:domain:7d64d0fa-66c3-445c-83db-3a324723daf8")))) + .minimumShouldMatch(1)), + "Expected filter each domain"); + } + + @Test + public void testAllUserAllGroupUnknownField() { + OperationContext resourceAllUsersPolicy = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllUsers(true)) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + new PolicyMatchCriterion() + .setField("UNKNOWN FIELD") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues(new StringArray("dataset", "chart")))))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + OperationContext resourceAllGroupsPolicy = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setAllGroups(true)) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + new PolicyMatchCriterion() + .setField("UNKNOWN FIELD") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues(new StringArray("dataset", "chart")))))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(resourceAllUsersPolicy), + Optional.of(QueryBuilders.boolQuery().mustNot(QueryBuilders.matchAllQuery())), + "Expected match-none query when an unknown field is encountered"); + + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(resourceAllGroupsPolicy), + Optional.of(QueryBuilders.boolQuery().mustNot(QueryBuilders.matchAllQuery())), + "Expected match-none query when an unknown field is encountered"); + } + + @Test + public void testUserGroupOwner() { + OperationContext ownerNoGroupsNoType = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors(new DataHubActorFilter().setResourceOwners(true)) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(ownerNoGroupsNoType), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.termsQuery( + "owners.keyword", + List.of( + TEST_USER_A.toString().toLowerCase(), + TEST_GROUP_A.toString().toLowerCase(), + TEST_GROUP_C.toString().toLowerCase()))) + .minimumShouldMatch(1)), + "Expected user filter for owners without group filter"); + + OperationContext ownerWithGroupsNoType = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors( + new DataHubActorFilter() + .setResourceOwners(true) + .setGroups(new UrnArray(TEST_GROUP_A))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(ownerWithGroupsNoType), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.termsQuery( + "owners.keyword", + List.of( + TEST_USER_A.toString().toLowerCase(), + TEST_GROUP_A.toString().toLowerCase(), + TEST_GROUP_C.toString().toLowerCase()))) + .minimumShouldMatch(1)), + "Expected user AND group filter for owners"); + } + + @Test + public void testUserGroupOwnerTypes() { + OperationContext ownerTypeBusinessNoUserNoGroup = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors( + new DataHubActorFilter().setResourceOwnersTypes(new UrnArray(BUS_OWNER))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(ownerTypeBusinessNoUserNoGroup), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .should( + QueryBuilders.termsQuery( + "ownerTypes." + + OwnerTypeMap.encodeFieldName(BUS_OWNER.toString()) + + ".keyword", + List.of( + TEST_USER_A.toString().toLowerCase(), + TEST_GROUP_A.toString().toLowerCase(), + TEST_GROUP_C.toString().toLowerCase()))) + .minimumShouldMatch(1)) + .minimumShouldMatch(1)), + "Expected user filter for business owner via user or group urn"); + + OperationContext ownerTypeBusinessMultiUserNoGroup = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors( + new DataHubActorFilter() + .setResourceOwnersTypes(new UrnArray(BUS_OWNER)) + .setUsers(new UrnArray(List.of(TEST_USER_A, TEST_USER_B)))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(ownerTypeBusinessMultiUserNoGroup), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .should( + QueryBuilders.termsQuery( + "ownerTypes." + + OwnerTypeMap.encodeFieldName(BUS_OWNER.toString()) + + ".keyword", + List.of( + TEST_USER_A.toString().toLowerCase(), + TEST_GROUP_A.toString().toLowerCase(), + TEST_GROUP_C.toString().toLowerCase()))) + .minimumShouldMatch(1)) + .minimumShouldMatch(1)), + "Expected user filter for `business owner` by owner user/group A urn (excluding other user/group B)"); + + OperationContext ownerWithGroupsBusTechMultiGroup = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors( + new DataHubActorFilter() + .setResourceOwnersTypes(new UrnArray(BUS_OWNER, TECH_OWNER)) + .setGroups(new UrnArray(TEST_GROUP_A, TEST_GROUP_B))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(ownerWithGroupsBusTechMultiGroup), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .should( + QueryBuilders.termsQuery( + "ownerTypes." + + OwnerTypeMap.encodeFieldName(BUS_OWNER.toString()) + + ".keyword", + List.of( + TEST_USER_A.toString().toLowerCase(), + TEST_GROUP_A.toString().toLowerCase(), + TEST_GROUP_C.toString().toLowerCase()))) + .should( + QueryBuilders.termsQuery( + "ownerTypes." + + OwnerTypeMap.encodeFieldName(TECH_OWNER.toString()) + + ".keyword", + List.of( + TEST_USER_A.toString().toLowerCase(), + TEST_GROUP_A.toString().toLowerCase(), + TEST_GROUP_C.toString().toLowerCase()))) + .minimumShouldMatch(1)) + .minimumShouldMatch(1)), + "Expected filter for business owner or technical owner by group A (excluding other group B and owner privilege)"); + + OperationContext ownerWithMultiUserMultiGroupsBusTech = + sessionWithPolicy( + Set.of( + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setType(PoliciesConfig.METADATA_POLICY_TYPE) + .setActors( + new DataHubActorFilter() + .setResourceOwnersTypes(new UrnArray(BUS_OWNER, TECH_OWNER)) + .setUsers(new UrnArray(List.of(TEST_USER_A, TEST_USER_B))) + .setGroups(new UrnArray(TEST_GROUP_A, TEST_GROUP_B))) + .setPrivileges( + new StringArray(List.of(PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType()))))); + assertEquals( + ESAccessControlUtil.buildAccessControlFilters(ownerWithMultiUserMultiGroupsBusTech), + Optional.of( + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .should( + QueryBuilders.termsQuery( + "ownerTypes." + + OwnerTypeMap.encodeFieldName(BUS_OWNER.toString()) + + ".keyword", + List.of( + TEST_USER_A.toString().toLowerCase(), + TEST_GROUP_A.toString().toLowerCase(), + TEST_GROUP_C.toString().toLowerCase()))) + .should( + QueryBuilders.termsQuery( + "ownerTypes." + + OwnerTypeMap.encodeFieldName(TECH_OWNER.toString()) + + ".keyword", + List.of( + TEST_USER_A.toString().toLowerCase(), + TEST_GROUP_A.toString().toLowerCase(), + TEST_GROUP_C.toString().toLowerCase()))) + .minimumShouldMatch(1)) + .minimumShouldMatch(1)), + "Expected filter for business owner or technical owner by user A and group A (excluding other group B and owner privilege)"); + } + + private static OperationContext sessionWithPolicy(Set policies) { + return sessionWithPolicy(policies, List.of(TEST_GROUP_A, TEST_GROUP_C)); + } + + private static OperationContext sessionWithPolicy( + Set policies, List groups) { + Authorizer mockAuthorizer = mock(Authorizer.class); + when(mockAuthorizer.getActorPolicies(eq(UrnUtils.getUrn(USER_AUTH.getActor().toUrnStr())))) + .thenReturn(policies); + when(mockAuthorizer.getActorGroups(eq(UrnUtils.getUrn(USER_AUTH.getActor().toUrnStr())))) + .thenReturn(groups); + + return ENABLED_CONTEXT.asSession(mockAuthorizer, USER_AUTH); + } +} diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/utils/SearchUtilsTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/utils/SearchUtilsTest.java index f4e82242545301..2b2d76fb5ec6f6 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/utils/SearchUtilsTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/utils/SearchUtilsTest.java @@ -20,7 +20,9 @@ private SearchFlags getDefaultSearchFlags() { .setSkipCache(true) .setSkipAggregates(true) .setMaxAggValues(1) - .setSkipHighlighting(true), + .setSkipHighlighting(true) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), true); } @@ -54,7 +56,9 @@ public void testApplyDefaultSearchFlags() { .setSkipCache(false) .setSkipAggregates(false) .setMaxAggValues(2) - .setSkipHighlighting(false), + .setSkipHighlighting(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), "not empty", defaultFlags), setConvertSchemaFieldsToDatasets( @@ -63,7 +67,9 @@ public void testApplyDefaultSearchFlags() { .setSkipAggregates(false) .setSkipCache(false) .setMaxAggValues(2) - .setSkipHighlighting(false), + .setSkipHighlighting(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), SearchUtils.convertSchemaFieldToDataset(defaultFlags)), "Expected no default values"); @@ -74,7 +80,9 @@ public void testApplyDefaultSearchFlags() { .setSkipCache(false) .setSkipAggregates(false) .setMaxAggValues(2) - .setSkipHighlighting(false), + .setSkipHighlighting(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), null, defaultFlags), setConvertSchemaFieldsToDatasets( @@ -83,7 +91,9 @@ public void testApplyDefaultSearchFlags() { .setSkipAggregates(false) .setSkipCache(false) .setMaxAggValues(2) - .setSkipHighlighting(true), + .setSkipHighlighting(true) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), SearchUtils.convertSchemaFieldToDataset(defaultFlags)), "Expected skip highlight due to query null query"); for (String query : Set.of("*", "")) { @@ -94,7 +104,9 @@ public void testApplyDefaultSearchFlags() { .setSkipCache(false) .setSkipAggregates(false) .setMaxAggValues(2) - .setSkipHighlighting(false), + .setSkipHighlighting(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), query, defaultFlags), setConvertSchemaFieldsToDatasets( @@ -103,7 +115,9 @@ public void testApplyDefaultSearchFlags() { .setSkipAggregates(false) .setSkipCache(false) .setMaxAggValues(2) - .setSkipHighlighting(true), + .setSkipHighlighting(true) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), SearchUtils.convertSchemaFieldToDataset(defaultFlags)), String.format("Expected skip highlight due to query string `%s`", query)); } @@ -117,7 +131,9 @@ public void testApplyDefaultSearchFlags() { .setSkipAggregates(true) .setSkipCache(true) .setMaxAggValues(1) - .setSkipHighlighting(true), + .setSkipHighlighting(true) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), SearchUtils.convertSchemaFieldToDataset(defaultFlags)), "Expected all default values except fulltext"); assertEquals( @@ -129,7 +145,9 @@ public void testApplyDefaultSearchFlags() { .setSkipAggregates(true) .setSkipCache(false) .setMaxAggValues(1) - .setSkipHighlighting(true), + .setSkipHighlighting(true) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), SearchUtils.convertSchemaFieldToDataset(defaultFlags)), "Expected all default values except skipCache"); assertEquals( @@ -141,7 +159,9 @@ public void testApplyDefaultSearchFlags() { .setSkipAggregates(false) .setSkipCache(true) .setMaxAggValues(1) - .setSkipHighlighting(true), + .setSkipHighlighting(true) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), SearchUtils.convertSchemaFieldToDataset(defaultFlags)), "Expected all default values except skipAggregates"); assertEquals( @@ -153,7 +173,9 @@ public void testApplyDefaultSearchFlags() { .setSkipAggregates(true) .setSkipCache(true) .setMaxAggValues(2) - .setSkipHighlighting(true), + .setSkipHighlighting(true) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), SearchUtils.convertSchemaFieldToDataset(defaultFlags)), "Expected all default values except maxAggValues"); assertEquals( @@ -165,7 +187,9 @@ public void testApplyDefaultSearchFlags() { .setSkipAggregates(true) .setSkipCache(true) .setMaxAggValues(1) - .setSkipHighlighting(false), + .setSkipHighlighting(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), SearchUtils.convertSchemaFieldToDataset(defaultFlags)), "Expected all default values except skipHighlighting"); } @@ -184,7 +208,9 @@ public void testImmutableDefaults() throws CloneNotSupportedException { .setSkipCache(false) .setSkipAggregates(false) .setMaxAggValues(2) - .setSkipHighlighting(false), + .setSkipHighlighting(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), SearchUtils.convertSchemaFieldToDataset(defaultFlags)), "not empty", defaultFlags), @@ -194,7 +220,9 @@ public void testImmutableDefaults() throws CloneNotSupportedException { .setSkipAggregates(false) .setSkipCache(false) .setMaxAggValues(2) - .setSkipHighlighting(false), + .setSkipHighlighting(false) + .setIncludeSoftDeleted(false) + .setIncludeRestricted(false), SearchUtils.convertSchemaFieldToDataset(defaultFlags)), "Expected no default values"); diff --git a/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SampleDataFixtureConfiguration.java b/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SampleDataFixtureConfiguration.java index 24acb7bbcb4a70..d627bd52cace9c 100644 --- a/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SampleDataFixtureConfiguration.java +++ b/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SampleDataFixtureConfiguration.java @@ -37,6 +37,7 @@ import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; import com.linkedin.metadata.version.GitVersion; +import io.datahubproject.metadata.context.OperationContext; import io.datahubproject.test.search.config.SearchCommonTestConfiguration; import java.io.IOException; import java.util.Map; @@ -72,6 +73,8 @@ public class SampleDataFixtureConfiguration { @Autowired private CustomSearchConfiguration _customSearchConfiguration; + @Autowired private OperationContext opContext; + @Bean(name = "sampleDataPrefix") protected String sampleDataPrefix() { return "smpldat"; @@ -297,6 +300,7 @@ private EntityClient entityClientHelper( PreProcessHooks preProcessHooks = new PreProcessHooks(); preProcessHooks.setUiEnabled(true); return new JavaEntityClient( + opContext, new EntityServiceImpl(mockAspectDao, null, entityRegistry, true, preProcessHooks, true), null, entitySearchService, diff --git a/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SearchLineageFixtureConfiguration.java b/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SearchLineageFixtureConfiguration.java index 1c43e623443c1e..74971c0b41eb1b 100644 --- a/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SearchLineageFixtureConfiguration.java +++ b/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SearchLineageFixtureConfiguration.java @@ -35,6 +35,7 @@ import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; import com.linkedin.metadata.version.GitVersion; +import io.datahubproject.metadata.context.OperationContext; import io.datahubproject.test.search.config.SearchCommonTestConfiguration; import io.datahubproject.test.search.config.SearchTestContainerConfiguration; import java.io.IOException; @@ -62,6 +63,8 @@ public class SearchLineageFixtureConfiguration { @Autowired private CustomSearchConfiguration customSearchConfiguration; + @Autowired private OperationContext opContext; + @Bean(name = "searchLineagePrefix") protected String indexPrefix() { return "srchlin"; @@ -125,6 +128,7 @@ protected ElasticSearchService entitySearchService( ESWriteDAO writeDAO = new ESWriteDAO( aspectRetriever.getEntityRegistry(), searchClient, indexConvention, bulkProcessor, 1); + return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO) .postConstruct(aspectRetriever); } @@ -232,6 +236,7 @@ protected EntityClient entityClient( PreProcessHooks preProcessHooks = new PreProcessHooks(); preProcessHooks.setUiEnabled(true); return new JavaEntityClient( + opContext, new EntityServiceImpl(null, null, entityRegistry, true, preProcessHooks, true), null, entitySearchService, diff --git a/metadata-io/src/test/java/io/datahubproject/test/search/SearchTestUtils.java b/metadata-io/src/test/java/io/datahubproject/test/search/SearchTestUtils.java index f3689f9b5d04ae..5dceed80f6542f 100644 --- a/metadata-io/src/test/java/io/datahubproject/test/search/SearchTestUtils.java +++ b/metadata-io/src/test/java/io/datahubproject/test/search/SearchTestUtils.java @@ -14,7 +14,6 @@ import com.linkedin.datahub.graphql.types.SearchableEntityType; import com.linkedin.datahub.graphql.types.entitytype.EntityTypeMapper; import com.linkedin.metadata.graph.LineageDirection; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.LineageSearchResult; import com.linkedin.metadata.search.LineageSearchService; @@ -22,6 +21,7 @@ import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.search.SearchService; import com.linkedin.metadata.search.elasticsearch.update.ESBulkProcessor; +import io.datahubproject.metadata.context.OperationContext; import java.util.List; import java.util.Optional; import java.util.stream.Collectors; @@ -54,94 +54,113 @@ public static void syncAfterWrite(ESBulkProcessor bulkProcessor) throws Interrup .collect(Collectors.toList()); } - public static SearchResult searchAcrossEntities(SearchService searchService, String query) { - return searchAcrossEntities(searchService, query, null); - } - - public static SearchResult searchAcrossEntities( - SearchService searchService, String query, @Nullable List facets) { - return searchService.searchAcrossEntities( - SEARCHABLE_ENTITIES, - query, - null, - null, - 0, - 100, - new SearchFlags().setFulltext(true).setSkipCache(true), - facets); + public static SearchResult facetAcrossEntities( + OperationContext opContext, + SearchService searchService, + String query, + @Nullable List facets) { + return facetAcrossEntities(opContext, searchService, SEARCHABLE_ENTITIES, query, facets, null); } - public static SearchResult searchAcrossEntities( + public static SearchResult facetAcrossEntities( + OperationContext opContext, SearchService searchService, + List entityNames, String query, @Nullable List facets, - Filter filter, - List entityNames) { + @Nullable Filter filter) { return searchService.searchAcrossEntities( + opContext.withSearchFlags(flags -> flags.setFulltext(true).setSkipCache(true)), entityNames, query, filter, null, 0, 100, - new SearchFlags().setFulltext(true).setSkipCache(true), facets); } - public static SearchResult searchAcrossCustomEntities( - SearchService searchService, String query, List searchableEntities) { + public static SearchResult searchAcrossEntities( + OperationContext opContext, SearchService searchService, String query) { + return searchAcrossEntities(opContext, searchService, SEARCHABLE_ENTITIES, query, null); + } + + public static SearchResult searchAcrossEntities( + OperationContext opContext, + SearchService searchService, + List entityNames, + String query) { + return searchAcrossEntities(opContext, searchService, entityNames, query, null); + } + + public static SearchResult searchAcrossEntities( + OperationContext opContext, + SearchService searchService, + List entityNames, + String query, + Filter filter) { return searchService.searchAcrossEntities( - searchableEntities, + opContext.withSearchFlags(flags -> flags.setFulltext(true).setSkipCache(true)), + entityNames, query, - null, + filter, null, 0, 100, - new SearchFlags().setFulltext(true).setSkipCache(true)); + null); } - public static SearchResult search(SearchService searchService, String query) { - return search(searchService, SEARCHABLE_ENTITIES, query); + public static SearchResult search( + OperationContext opContext, SearchService searchService, String query) { + return search(opContext, searchService, SEARCHABLE_ENTITIES, query); } public static SearchResult search( - SearchService searchService, List entities, String query) { + OperationContext opContext, + SearchService searchService, + List entities, + String query) { return searchService.search( + opContext.withSearchFlags(flags -> flags.setFulltext(true).setSkipCache(true)), entities, query, null, null, 0, - 100, - new SearchFlags().setFulltext(true).setSkipCache(true)); + 100); } public static ScrollResult scroll( - SearchService searchService, String query, int batchSize, @Nullable String scrollId) { + OperationContext opContext, + SearchService searchService, + String query, + int batchSize, + @Nullable String scrollId) { return searchService.scrollAcrossEntities( + opContext.withSearchFlags(flags -> flags.setFulltext(true).setSkipCache(true)), SEARCHABLE_ENTITIES, query, null, null, scrollId, "3m", - batchSize, - new SearchFlags().setFulltext(true).setSkipCache(true)); + batchSize); } - public static SearchResult searchStructured(SearchService searchService, String query) { + public static SearchResult searchStructured( + OperationContext opContext, SearchService searchService, String query) { return searchService.searchAcrossEntities( + opContext.withSearchFlags(flags -> flags.setFulltext(false).setSkipCache(true)), SEARCHABLE_ENTITIES, query, null, null, 0, - 100, - new SearchFlags().setFulltext(false).setSkipCache(true)); + 100); } public static LineageSearchResult lineage( - LineageSearchService lineageSearchService, Urn root, int hops) { + OperationContext opContext, LineageSearchService lineageSearchService, Urn root, int hops) { String degree = hops >= 3 ? "3+" : String.valueOf(hops); List filters = List.of( @@ -153,6 +172,7 @@ public static LineageSearchResult lineage( .build()); return lineageSearchService.searchAcrossLineage( + opContext.withSearchFlags(flags -> flags.setSkipCache(true)), root, LineageDirection.DOWNSTREAM, SEARCHABLE_ENTITY_TYPES.stream() @@ -165,12 +185,14 @@ public static LineageSearchResult lineage( 0, 100, null, - null, - new SearchFlags().setSkipCache(true)); + null); } public static AutoCompleteResults autocomplete( - SearchableEntityType searchableEntityType, String query) throws Exception { + OperationContext opContext, + SearchableEntityType searchableEntityType, + String query) + throws Exception { return searchableEntityType.autoComplete( query, null, @@ -191,6 +213,11 @@ public Authentication getAuthentication() { public Authorizer getAuthorizer() { return null; } + + @Override + public OperationContext getOperationContext() { + return opContext; + } }); } diff --git a/metadata-io/src/test/java/io/datahubproject/test/search/config/SearchCommonTestConfiguration.java b/metadata-io/src/test/java/io/datahubproject/test/search/config/SearchCommonTestConfiguration.java index ae81eaf1ef3884..8df86bef250d37 100644 --- a/metadata-io/src/test/java/io/datahubproject/test/search/config/SearchCommonTestConfiguration.java +++ b/metadata-io/src/test/java/io/datahubproject/test/search/config/SearchCommonTestConfiguration.java @@ -18,6 +18,8 @@ import com.linkedin.metadata.models.registry.EntityRegistryException; import com.linkedin.metadata.models.registry.SnapshotEntityRegistry; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.net.URISyntaxException; import java.util.Map; import org.springframework.boot.test.context.TestConfiguration; @@ -88,4 +90,9 @@ protected AspectRetriever snapshotRegistryAspectRetriever() when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); return aspectRetriever; } + + @Bean(name = "systemOperationContext") + public OperationContext systemOperationContext(final EntityRegistry entityRegistry) { + return TestOperationContexts.systemContextNoSearchAuthorization(entityRegistry); + } } diff --git a/metadata-jobs/mae-consumer-job/src/main/java/com/linkedin/metadata/kafka/MaeConsumerApplication.java b/metadata-jobs/mae-consumer-job/src/main/java/com/linkedin/metadata/kafka/MaeConsumerApplication.java index 018646fad07fcd..2b2643f088da4c 100644 --- a/metadata-jobs/mae-consumer-job/src/main/java/com/linkedin/metadata/kafka/MaeConsumerApplication.java +++ b/metadata-jobs/mae-consumer-job/src/main/java/com/linkedin/metadata/kafka/MaeConsumerApplication.java @@ -24,7 +24,8 @@ "com.linkedin.gms.factory.form", "com.linkedin.gms.factory.incident", "com.linkedin.gms.factory.timeline.eventgenerator", - "io.datahubproject.metadata.jobs.common.health.kafka" + "io.datahubproject.metadata.jobs.common.health.kafka", + "com.linkedin.gms.factory.context" }, excludeFilters = { @ComponentScan.Filter( diff --git a/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTest.java b/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTest.java index 22fbe7fc6b6ca8..68965e5bde2612 100644 --- a/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTest.java +++ b/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTest.java @@ -3,6 +3,7 @@ import static org.testng.AssertJUnit.*; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.service.FormService; import io.datahubproject.metadata.jobs.common.health.kafka.KafkaHealthIndicator; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.test.context.SpringBootTest; @@ -15,13 +16,16 @@ classes = {MaeConsumerApplication.class, MaeConsumerApplicationTestConfiguration.class}) public class MaeConsumerApplicationTest extends AbstractTestNGSpringContextTests { - @Autowired private EntityService _mockEntityService; + @Autowired private EntityService mockEntityService; @Autowired private KafkaHealthIndicator kafkaHealthIndicator; + @Autowired private FormService formService; + @Test public void testMaeConsumerAutoWiring() { - assertNotNull(_mockEntityService); + assertNotNull(mockEntityService); assertNotNull(kafkaHealthIndicator); + assertNotNull(formService); } } 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 7a1aaa7f6a0561..b212eb11e50c09 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 @@ -23,7 +23,6 @@ import com.linkedin.metadata.kafka.hook.MetadataChangeLogHook; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; @@ -38,6 +37,7 @@ import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.List; import java.util.stream.Collectors; @@ -71,20 +71,22 @@ public class SiblingAssociationHook implements MetadataChangeLogHook { public static final String SOURCE_SUBTYPE_V1 = "source"; public static final String SOURCE_SUBTYPE_V2 = "Source"; - private final EntityRegistry _entityRegistry; - private final SystemEntityClient _entityClient; - private final EntitySearchService _searchService; + private final EntityRegistry entityRegistry; + private final SystemEntityClient systemEntityClient; + private final EntitySearchService entitySearchService; private final boolean _isEnabled; + private final OperationContext opContext; @Autowired public SiblingAssociationHook( - @Nonnull final EntityRegistry entityRegistry, - @Nonnull final SystemEntityClient entityClient, + @Nonnull final OperationContext opContext, + @Nonnull final SystemEntityClient systemEntityClient, @Nonnull final EntitySearchService searchService, @Nonnull @Value("${siblings.enabled:true}") Boolean isEnabled) { - _entityRegistry = entityRegistry; - _entityClient = entityClient; - _searchService = searchService; + this.opContext = opContext; + this.entityRegistry = opContext.getEntityRegistryContext().getEntityRegistry(); + this.systemEntityClient = systemEntityClient; + entitySearchService = searchService; _isEnabled = isEnabled; } @@ -135,14 +137,16 @@ public void invoke(@Nonnull MetadataChangeLog event) { private void handleEntityKeyEvent(DatasetUrn datasetUrn) { Filter entitiesWithYouAsSiblingFilter = createFilterForEntitiesWithYouAsSibling(datasetUrn); final SearchResult searchResult = - _searchService.search( + entitySearchService.search( + opContext.withSearchFlags( + flags -> + flags.setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true)), List.of(DATASET_ENTITY_NAME), "*", entitiesWithYouAsSiblingFilter, null, 0, - 10, - new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true)); + 10); // we have a match of an entity with you as a sibling, associate yourself back searchResult @@ -265,7 +269,7 @@ private void setSiblingsAndSoftDeleteSibling(Urn dbtUrn, Urn sourceUrn) { dbtSiblingProposal.setEntityUrn(dbtUrn); try { - _entityClient.ingestProposal(dbtSiblingProposal, true); + systemEntityClient.ingestProposal(dbtSiblingProposal, true); } catch (RemoteInvocationException e) { log.error("Error while associating {} with {}: {}", dbtUrn, sourceUrn, e.toString()); throw new RuntimeException("Error ingesting sibling proposal. Skipping processing.", e); @@ -290,7 +294,7 @@ private void setSiblingsAndSoftDeleteSibling(Urn dbtUrn, Urn sourceUrn) { .filter( urn -> { try { - return _entityClient.exists(urn); + return systemEntityClient.exists(urn); } catch (RemoteInvocationException e) { log.error("Error while checking existence of {}: {}", urn, e.toString()); throw new RuntimeException("Error checking existence. Skipping processing.", e); @@ -312,7 +316,7 @@ private void setSiblingsAndSoftDeleteSibling(Urn dbtUrn, Urn sourceUrn) { sourceSiblingProposal.setEntityUrn(sourceUrn); try { - _entityClient.ingestProposal(sourceSiblingProposal, true); + systemEntityClient.ingestProposal(sourceSiblingProposal, true); } catch (RemoteInvocationException e) { log.error("Error while associating {} with {}: {}", dbtUrn, sourceUrn, e.toString()); throw new RuntimeException("Error ingesting sibling proposal. Skipping processing.", e); @@ -338,7 +342,7 @@ private boolean isEligibleForProcessing(final MetadataChangeLog event) { private Urn getUrnFromEvent(final MetadataChangeLog event) { EntitySpec entitySpec; try { - entitySpec = _entityRegistry.getEntitySpec(event.getEntityType()); + entitySpec = entityRegistry.getEntitySpec(event.getEntityType()); } catch (IllegalArgumentException e) { log.error("Error while processing entity type {}: {}", event.getEntityType(), e.toString()); throw new RuntimeException( @@ -359,7 +363,7 @@ private UpstreamLineage getUpstreamLineageFromEvent(final MetadataChangeLog even } try { - entitySpec = _entityRegistry.getEntitySpec(event.getEntityType()); + entitySpec = entityRegistry.getEntitySpec(event.getEntityType()); } catch (IllegalArgumentException e) { log.error("Error while processing entity type {}: {}", event.getEntityType(), e.toString()); throw new RuntimeException( @@ -383,7 +387,7 @@ private SubTypes getSubtypesFromEvent(final MetadataChangeLog event) { } try { - entitySpec = _entityRegistry.getEntitySpec(event.getEntityType()); + entitySpec = entityRegistry.getEntitySpec(event.getEntityType()); } catch (IllegalArgumentException e) { log.error("Error while processing entity type {}: {}", event.getEntityType(), e.toString()); throw new RuntimeException( @@ -427,7 +431,7 @@ private Filter createFilterForEntitiesWithYouAsSibling(final Urn entityUrn) { private SubTypes getSubtypesFromEntityClient(final Urn urn) { try { EntityResponse entityResponse = - _entityClient.getV2(urn, ImmutableSet.of(SUB_TYPES_ASPECT_NAME)); + systemEntityClient.getV2(urn, ImmutableSet.of(SUB_TYPES_ASPECT_NAME)); if (entityResponse != null && entityResponse.hasAspects() @@ -445,7 +449,7 @@ private SubTypes getSubtypesFromEntityClient(final Urn urn) { private UpstreamLineage getUpstreamLineageFromEntityClient(final Urn urn) { try { EntityResponse entityResponse = - _entityClient.getV2(urn, ImmutableSet.of(UPSTREAM_LINEAGE_ASPECT_NAME)); + systemEntityClient.getV2(urn, ImmutableSet.of(UPSTREAM_LINEAGE_ASPECT_NAME)); if (entityResponse != null && entityResponse.hasAspects() @@ -467,7 +471,7 @@ private UpstreamLineage getUpstreamLineageFromEntityClient(final Urn urn) { private Siblings getSiblingsFromEntityClient(final Urn urn) { try { EntityResponse entityResponse = - _entityClient.getV2(urn, ImmutableSet.of(SIBLINGS_ASPECT_NAME)); + systemEntityClient.getV2(urn, ImmutableSet.of(SIBLINGS_ASPECT_NAME)); if (entityResponse != null && entityResponse.hasAspects() 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 3823668adeace9..2a83a2310518f3 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 @@ -2,6 +2,7 @@ import static com.linkedin.metadata.Constants.*; import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.when; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; @@ -26,7 +27,6 @@ import com.linkedin.metadata.key.DatasetKey; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; @@ -34,6 +34,8 @@ import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.MetadataChangeProposal; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.net.URISyntaxException; import org.mockito.Mockito; import org.testng.annotations.BeforeMethod; @@ -43,6 +45,7 @@ public class SiblingAssociationHookTest { private SiblingAssociationHook _siblingAssociationHook; SystemEntityClient _mockEntityClient; EntitySearchService _mockSearchService; + OperationContext opContext; @BeforeMethod public void setupTest() { @@ -53,8 +56,9 @@ public void setupTest() { .getResourceAsStream("test-entity-registry-siblings.yml")); _mockEntityClient = Mockito.mock(SystemEntityClient.class); _mockSearchService = Mockito.mock(EntitySearchService.class); + opContext = TestOperationContexts.systemContextNoSearchAuthorization(registry); _siblingAssociationHook = - new SiblingAssociationHook(registry, _mockEntityClient, _mockSearchService, true); + new SiblingAssociationHook(opContext, _mockEntityClient, _mockSearchService, true); _siblingAssociationHook.setEnabled(true); } @@ -69,13 +73,12 @@ public void testInvokeWhenThereIsAPairWithDbtSourceNode() throws Exception { EntityResponse mockResponse = new EntityResponse(); mockResponse.setAspects(mockResponseMap); - Mockito.when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); + when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); - Mockito.when( - _mockEntityClient.getV2( - Urn.createFromString( - "urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)"), - ImmutableSet.of(SUB_TYPES_ASPECT_NAME))) + when(_mockEntityClient.getV2( + Urn.createFromString( + "urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)"), + ImmutableSet.of(SUB_TYPES_ASPECT_NAME))) .thenReturn(mockResponse); MetadataChangeLog event = @@ -145,7 +148,7 @@ public void testInvokeWhenThereIsNoPairWithDbtModel() throws Exception { SubTypes mockSourceSubtypesAspect = new SubTypes(); mockSourceSubtypesAspect.setTypeNames(new StringArray(ImmutableList.of("model"))); - Mockito.when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); + when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); EnvelopedAspectMap mockResponseMap = new EnvelopedAspectMap(); mockResponseMap.put( @@ -154,13 +157,12 @@ public void testInvokeWhenThereIsNoPairWithDbtModel() throws Exception { EntityResponse mockResponse = new EntityResponse(); mockResponse.setAspects(mockResponseMap); - Mockito.when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); + when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); - Mockito.when( - _mockEntityClient.getV2( - Urn.createFromString( - "urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)"), - ImmutableSet.of(SUB_TYPES_ASPECT_NAME))) + when(_mockEntityClient.getV2( + Urn.createFromString( + "urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)"), + ImmutableSet.of(SUB_TYPES_ASPECT_NAME))) .thenReturn(mockResponse); MetadataChangeLog event = @@ -206,7 +208,7 @@ public void testInvokeWhenThereIsNoPairWithDbtModel() throws Exception { @Test public void testInvokeWhenThereIsAPairWithBigqueryDownstreamNode() throws Exception { - Mockito.when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); + when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); MetadataChangeLog event = createEvent(DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, ChangeType.UPSERT); @@ -271,7 +273,7 @@ public void testInvokeWhenThereIsAPairWithBigqueryDownstreamNode() throws Except @Test public void testInvokeWhenThereIsAKeyBeingReingested() throws Exception { - Mockito.when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); + when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); SearchResult returnSearchResult = new SearchResult(); SearchEntityArray returnEntityArray = new SearchEntityArray(); @@ -283,19 +285,8 @@ public void testInvokeWhenThereIsAKeyBeingReingested() throws Exception { returnSearchResult.setEntities(returnEntityArray); - Mockito.when( - _mockSearchService.search( - any(), - anyString(), - any(), - any(), - anyInt(), - anyInt(), - eq( - new SearchFlags() - .setFulltext(false) - .setSkipAggregates(true) - .setSkipHighlighting(true)))) + when(_mockSearchService.search( + any(OperationContext.class), any(), anyString(), any(), any(), anyInt(), anyInt())) .thenReturn(returnSearchResult); MetadataChangeLog event = diff --git a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java index a80017a0956b22..789918a4b164c4 100644 --- a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java +++ b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java @@ -19,6 +19,10 @@ import com.linkedin.metadata.service.FormService; import com.linkedin.metadata.systemmetadata.SystemMetadataService; import com.linkedin.metadata.timeseries.TimeseriesAspectService; +import com.linkedin.metadata.utils.elasticsearch.IndexConvention; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.metadata.context.OperationContextConfig; +import io.datahubproject.test.metadata.context.TestOperationContexts; import org.apache.avro.generic.GenericRecord; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.boot.test.mock.mockito.MockBean; @@ -80,4 +84,18 @@ public SystemEntityClient systemEntityClient( @MockBean public SchemaRegistryService schemaRegistryService; @MockBean public EntityIndexBuilders entityIndexBuilders; + + @Bean(name = "systemOperationContext") + public OperationContext operationContext( + final EntityRegistry entityRegistry, + @Qualifier("systemAuthentication") final Authentication systemAuthentication, + final IndexConvention indexConvention) { + when(systemAuthentication.getActor()) + .thenReturn(TestOperationContexts.TEST_SYSTEM_AUTH.getActor()); + return OperationContext.asSystem( + OperationContextConfig.builder().build(), + entityRegistry, + systemAuthentication, + indexConvention); + } } diff --git a/metadata-jobs/mce-consumer-job/src/main/java/com/linkedin/metadata/kafka/MceConsumerApplication.java b/metadata-jobs/mce-consumer-job/src/main/java/com/linkedin/metadata/kafka/MceConsumerApplication.java index 1210bf37059b43..0b0c8f622efd2b 100644 --- a/metadata-jobs/mce-consumer-job/src/main/java/com/linkedin/metadata/kafka/MceConsumerApplication.java +++ b/metadata-jobs/mce-consumer-job/src/main/java/com/linkedin/metadata/kafka/MceConsumerApplication.java @@ -32,7 +32,8 @@ "com.linkedin.metadata.dao.producer", "com.linkedin.gms.factory.form", "com.linkedin.metadata.dao.producer", - "io.datahubproject.metadata.jobs.common.health.kafka" + "io.datahubproject.metadata.jobs.common.health.kafka", + "com.linkedin.gms.factory.context" }, excludeFilters = { @ComponentScan.Filter( diff --git a/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java b/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java index 93a6ae8fb47970..ce093e3115f8e1 100644 --- a/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java +++ b/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java @@ -1,10 +1,10 @@ package com.linkedin.metadata.kafka; -import com.datahub.authentication.Authentication; import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.gms.factory.auth.SystemAuthenticationFactory; import com.linkedin.gms.factory.config.ConfigurationProvider; +import com.linkedin.gms.factory.context.SystemOperationContextFactory; import com.linkedin.metadata.dao.producer.KafkaHealthChecker; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.graph.SiblingGraphService; @@ -15,6 +15,7 @@ import com.linkedin.metadata.timeseries.TimeseriesAspectService; import com.linkedin.parseq.retry.backoff.ExponentialBackoff; import com.linkedin.restli.client.Client; +import io.datahubproject.metadata.context.OperationContext; import io.ebean.Database; import java.net.URI; import org.springframework.beans.factory.annotation.Autowired; @@ -27,7 +28,7 @@ import org.springframework.context.annotation.Primary; @TestConfiguration -@Import(value = {SystemAuthenticationFactory.class}) +@Import(value = {SystemAuthenticationFactory.class, SystemOperationContextFactory.class}) public class MceConsumerApplicationTestConfiguration { @Autowired private TestRestTemplate restTemplate; @@ -39,15 +40,15 @@ public class MceConsumerApplicationTestConfiguration { @Bean @Primary public SystemEntityClient systemEntityClient( - @Qualifier("configurationProvider") final ConfigurationProvider configurationProvider, - @Qualifier("systemAuthentication") final Authentication systemAuthentication) { + @Qualifier("systemOperationContext") final OperationContext systemOperationContext, + @Qualifier("configurationProvider") final ConfigurationProvider configurationProvider) { String selfUri = restTemplate.getRootUri(); final Client restClient = DefaultRestliClientFactory.getRestLiClient(URI.create(selfUri), null); return new SystemRestliEntityClient( + systemOperationContext, restClient, new ExponentialBackoff(1), 1, - systemAuthentication, configurationProvider.getCache().getClient().getEntityClient()); } diff --git a/metadata-models/src/main/pegasus/com/linkedin/common/Ownership.pdl b/metadata-models/src/main/pegasus/com/linkedin/common/Ownership.pdl index cee882762814e2..d76b60c061b4bd 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/common/Ownership.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/common/Ownership.pdl @@ -13,6 +13,17 @@ record Ownership { */ owners: array[Owner] + /** + * Ownership type to Owners map, populated via mutation hook. + */ + @Searchable = { + "/*": { + "fieldType": "MAP_ARRAY", + "queryByDefault": false + } + } + ownerTypes: optional map[string, array[Urn]] = {} + /** * Audit stamp containing who last modified the record and when. A value of 0 in the time field indicates missing data. */ diff --git a/metadata-models/src/main/pegasus/com/linkedin/metadata/query/SearchFlags.pdl b/metadata-models/src/main/pegasus/com/linkedin/metadata/query/SearchFlags.pdl index 67f41ea175b51f..355a8bb7a5cb34 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/metadata/query/SearchFlags.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/metadata/query/SearchFlags.pdl @@ -38,4 +38,14 @@ record SearchFlags { * Instructions for grouping results before returning */ groupingSpec: optional GroupingSpec + + /** + * include soft deleted entities in results + */ + includeSoftDeleted:optional boolean = false + + /** + * include restricted entities in results (default is to filter) + */ + includeRestricted:optional boolean = false } diff --git a/metadata-models/src/main/pegasus/com/linkedin/metadata/search/SearchEntity.pdl b/metadata-models/src/main/pegasus/com/linkedin/metadata/search/SearchEntity.pdl index 1010e3e2330cc0..df457c1ba26f99 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/metadata/search/SearchEntity.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/metadata/search/SearchEntity.pdl @@ -29,4 +29,10 @@ record SearchEntity { features: optional map[string, double] score: optional double + + /** + * A list of the the restricted aspects on the entity. + * If the key aspect is present, assume ALL aspects should be restricted including the entity's Urn. + */ + restrictedAspects: optional array[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 cf8eb738e2443f..22fe6551eb5285 100644 --- a/metadata-models/src/main/resources/entity-registry.yml +++ b/metadata-models/src/main/resources/entity-registry.yml @@ -574,4 +574,13 @@ plugins: enabled: true supportedEntityAspectNames: - entityName: '*' - aspectName: structuredProperties \ No newline at end of file + aspectName: structuredProperties + - className: 'com.linkedin.metadata.aspect.hooks.OwnerTypeMap' + enabled: true + supportedOperations: + - UPSERT + - CREATE + - RESTATE + supportedEntityAspectNames: + - entityName: '*' + aspectName: ownership \ No newline at end of file diff --git a/metadata-operation-context/build.gradle b/metadata-operation-context/build.gradle new file mode 100644 index 00000000000000..1be98cb0140f3a --- /dev/null +++ b/metadata-operation-context/build.gradle @@ -0,0 +1,15 @@ +plugins { + id 'java-library' +} + +dependencies { + api project(':metadata-utils') + api project(':metadata-auth:auth-api') + + implementation externalDependency.slf4jApi + compileOnly externalDependency.lombok + + annotationProcessor externalDependency.lombok + testImplementation externalDependency.testng + testImplementation externalDependency.mockito +} \ No newline at end of file diff --git a/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/ActorContext.java b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/ActorContext.java new file mode 100644 index 00000000000000..6c8077923e67f8 --- /dev/null +++ b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/ActorContext.java @@ -0,0 +1,71 @@ +package io.datahubproject.metadata.context; + +import com.datahub.authentication.Authentication; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.metadata.authorization.PoliciesConfig; +import com.linkedin.policy.DataHubPolicyInfo; +import java.util.Collection; +import java.util.Collections; +import java.util.Optional; +import java.util.Set; +import lombok.Builder; +import lombok.Getter; + +@Builder +@Getter +public class ActorContext implements ContextInterface { + + public static ActorContext asSystem(Authentication systemAuthentication) { + return ActorContext.builder().systemAuth(true).authentication(systemAuthentication).build(); + } + + public static ActorContext asSessionRestricted( + Authentication authentication, + Set dataHubPolicySet, + Collection groupMembership) { + return ActorContext.builder() + .systemAuth(false) + .authentication(authentication) + .policyInfoSet(dataHubPolicySet) + .groupMembership(groupMembership) + .build(); + } + + private final Authentication authentication; + @Builder.Default private final Set policyInfoSet = Collections.emptySet(); + @Builder.Default private final Collection groupMembership = Collections.emptyList(); + private final boolean systemAuth; + + public Urn getActorUrn() { + return UrnUtils.getUrn(authentication.getActor().toUrnStr()); + } + + /** + * The current implementation creates a cache entry unique for the set of policies. + * + *

We are relying on the consistent hash code implementation of String and the consistent + * conversion of the policy into a String + * + * @return + */ + @Override + public Optional getCacheKeyComponent() { + return Optional.of( + policyInfoSet.stream() + .filter(policy -> PoliciesConfig.ACTIVE_POLICY_STATE.equals(policy.getState())) + .mapToInt( + policy -> { + if (policy.getActors().hasResourceOwners() + || policy.getActors().hasResourceOwnersTypes()) { + // results are based on actor, distinct() added to remove duplicate sums of + // multiple owner policies + return authentication.getActor().toUrnStr().hashCode(); + } else { + return policy.toString().hashCode(); + } + }) + .distinct() + .sum()); + } +} diff --git a/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/AuthorizerContext.java b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/AuthorizerContext.java new file mode 100644 index 00000000000000..fdd84f6d64557d --- /dev/null +++ b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/AuthorizerContext.java @@ -0,0 +1,28 @@ +package io.datahubproject.metadata.context; + +import com.datahub.plugins.auth.authorization.Authorizer; +import java.util.Optional; +import javax.annotation.Nonnull; +import lombok.Builder; +import lombok.Getter; + +@Builder +@Getter +public class AuthorizerContext implements ContextInterface { + + public static final AuthorizerContext EMPTY = + AuthorizerContext.builder().authorizer(Authorizer.EMPTY).build(); + + @Nonnull private final Authorizer authorizer; + + /** + * No need to consider the authorizer in the cache context since it is ultimately determined by + * the underlying search context + * + * @return + */ + @Override + public Optional getCacheKeyComponent() { + return Optional.empty(); + } +} diff --git a/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/ContextInterface.java b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/ContextInterface.java new file mode 100644 index 00000000000000..6e02f273cbb216 --- /dev/null +++ b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/ContextInterface.java @@ -0,0 +1,21 @@ +package io.datahubproject.metadata.context; + +import java.util.Optional; + +public interface ContextInterface { + /** + * Caching layers must take into account the operation's context to avoid returning incorrect or + * restricted results. + * + *

A consistent hash must be produced in a distributed cache so that multiple jvms produce the + * same keys for the same objects within the same context. This generally rules out hashCode() + * since those are not guaranteed properties. + * + *

We are however leveraging the special case of hashCode() for String which should be + * consistent across jvms and executions. + * + *

The overall context must produce a unique id to be included with cache keys. Each component + * of the OperationContext must produce a unique identifier to be used for this purpose. + */ + Optional getCacheKeyComponent(); +} diff --git a/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/EntityRegistryContext.java b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/EntityRegistryContext.java new file mode 100644 index 00000000000000..c026a453604eeb --- /dev/null +++ b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/EntityRegistryContext.java @@ -0,0 +1,22 @@ +package io.datahubproject.metadata.context; + +import com.linkedin.metadata.models.registry.EntityRegistry; +import java.util.Optional; +import javax.annotation.Nullable; +import lombok.Builder; +import lombok.Getter; + +@Builder +@Getter +public class EntityRegistryContext implements ContextInterface { + public static EntityRegistryContext EMPTY = EntityRegistryContext.builder().build(); + + @Nullable private final EntityRegistry entityRegistry; + + @Override + public Optional getCacheKeyComponent() { + return entityRegistry == null + ? Optional.empty() + : Optional.ofNullable(entityRegistry.getIdentifier()).map(String::hashCode); + } +} diff --git a/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/OperationContext.java b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/OperationContext.java new file mode 100644 index 00000000000000..d2c038c26e3251 --- /dev/null +++ b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/OperationContext.java @@ -0,0 +1,285 @@ +package io.datahubproject.metadata.context; + +import com.datahub.authentication.Authentication; +import com.datahub.plugins.auth.authorization.Authorizer; +import com.google.common.collect.ImmutableSet; +import com.linkedin.common.AuditStamp; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.query.SearchFlags; +import com.linkedin.metadata.utils.AuditStampUtils; +import com.linkedin.metadata.utils.elasticsearch.IndexConvention; +import java.util.Collection; +import java.util.Objects; +import java.util.Optional; +import java.util.function.Function; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import lombok.Builder; +import lombok.Getter; + +/** + * These contexts define a read/write context which allows more flexibility when reading and writing + * to various data stores. This context can be considered per **operation** and allows for + * supporting database read replicas, mirroring or sharding across multiple databases/elasticsearch + * instances, and separation of data at the storage level. + * + *

Different operations might also include different EntityRegistries + * + *

An integral part of the operation's context is additionally the user's identity and this + * context encompasses the `Authentication` context. + */ +@Builder(toBuilder = true) +@Getter +public class OperationContext { + + /** + * This should be the primary entry point when a request is made to Rest.li, OpenAPI, Graphql or + * other service layers. + * + *

Copy the context from a system level context to a specific request/user context. Inheriting + * all other contexts except for the sessionActor. Consider this a down leveling of the access. + * + *

This allows the context to contain system context such as elasticsearch and database + * contexts which are inherited from the system. + * + * @param systemOperationContext the base operation context + * @param sessionAuthentication the lower level authentication + * @param allowSystemAuthentication whether the context is allowed to escalate as needed + * @return the new context + */ + public static OperationContext asSession( + OperationContext systemOperationContext, + @Nonnull Authorizer authorizer, + @Nonnull Authentication sessionAuthentication, + boolean allowSystemAuthentication) { + return systemOperationContext.toBuilder() + .operationContextConfig( + // update allowed system authentication + systemOperationContext.getOperationContextConfig().toBuilder() + .allowSystemAuthentication(allowSystemAuthentication) + .build()) + .authorizerContext(AuthorizerContext.builder().authorizer(authorizer).build()) + .build(sessionAuthentication); + } + + /** + * Apply a set of default flags on top of any existing search flags + * + * @param opContext + * @param flagDefaults + * @return + */ + public static OperationContext withSearchFlags( + OperationContext opContext, Function flagDefaults) { + + return opContext.toBuilder() + // update search flags for the request's session + .searchContext(opContext.getSearchContext().withFlagDefaults(flagDefaults)) + .build(opContext.getSessionAuthentication()); + } + + /** + * Set the system authentication object AND allow escalation of privilege for the session. This + * OperationContext typically serves the default. + * + *

If you'd like to set the system authentication but not allow escalation, use the + * systemActorContext() directly which does not reconfigure the escalation configuration. + * + * @param systemAuthentication the system authentication + * @return builder + */ + public static OperationContext asSystem( + @Nonnull OperationContextConfig config, + @Nonnull EntityRegistry entityRegistry, + @Nonnull Authentication systemAuthentication, + @Nonnull IndexConvention indexConvention) { + + ActorContext systemActorContext = + ActorContext.builder().systemAuth(true).authentication(systemAuthentication).build(); + OperationContextConfig systemConfig = + config.toBuilder().allowSystemAuthentication(true).build(); + SearchContext systemSearchContext = + SearchContext.builder().indexConvention(indexConvention).build(); + + return OperationContext.builder() + .operationContextConfig(systemConfig) + .systemActorContext(systemActorContext) + .searchContext(systemSearchContext) + .entityRegistryContext( + EntityRegistryContext.builder().entityRegistry(entityRegistry).build()) + // Authorizer.EMPTY doesn't actually apply to system auth + .authorizerContext(AuthorizerContext.builder().authorizer(Authorizer.EMPTY).build()) + .build(systemAuthentication); + } + + @Nonnull private final OperationContextConfig operationContextConfig; + @Nonnull private final ActorContext sessionActorContext; + @Nullable private final ActorContext systemActorContext; + @Nonnull private final SearchContext searchContext; + @Nonnull private final AuthorizerContext authorizerContext; + @Nonnull private final EntityRegistryContext entityRegistryContext; + + public OperationContext withSearchFlags( + @Nonnull Function flagDefaults) { + return OperationContext.withSearchFlags(this, flagDefaults); + } + + public OperationContext asSession( + @Nonnull Authorizer authorizer, @Nonnull Authentication sessionAuthentication) { + return OperationContext.asSession( + this, + authorizer, + sessionAuthentication, + getOperationContextConfig().isAllowSystemAuthentication()); + } + + @Nonnull + public EntityRegistry getEntityRegistry() { + return getEntityRegistryContext().getEntityRegistry(); + } + + /** + * Requests for a generic authentication should return the system first if allowed. + * + * @return an entity client + */ + @Nonnull + public ActorContext getActorContext() { + if (operationContextConfig.isAllowSystemAuthentication() && systemActorContext != null) { + return systemActorContext; + } else { + return sessionActorContext; + } + } + + /** + * Other users within the same group as the actor + * + * @return + */ + public Collection getActorPeers() { + return authorizerContext.getAuthorizer().getActorPeers(sessionActorContext.getActorUrn()); + } + + /** + * Whether default authentication is system level + * + * @return + */ + public boolean isSystemAuth() { + return operationContextConfig.isAllowSystemAuthentication() + && sessionActorContext.isSystemAuth(); + } + + /** + * Requests for a generic authentication should return the system first if allowed. + * + * @return an entity client + */ + public Authentication getAuthentication() { + return getActorContext().getAuthentication(); + } + + public Authentication getSessionAuthentication() { + return sessionActorContext.getAuthentication(); + } + + public Optional getSystemAuthentication() { + return Optional.ofNullable(systemActorContext).map(ActorContext::getAuthentication); + } + + /** AuditStamp prefer session authentication */ + public AuditStamp getAuditStamp(@Nullable Long currentTimeMs) { + return AuditStampUtils.getAuditStamp( + UrnUtils.getUrn(sessionActorContext.getAuthentication().getActor().toUrnStr()), + currentTimeMs); + } + + public AuditStamp getAuditStamp() { + return getAuditStamp(null); + } + + /** + * Return a unique id for this context. Typically useful for building cache keys. We combine the + * different context components to create a single string representation of the hashcode across + * the contexts. + * + *

The overall context id can be comprised of one or more other contexts depending on the + * requirements. + * + * @return id representing this context instance's unique identifier + */ + public String getGlobalContextId() { + return String.valueOf( + ImmutableSet.builder() + .add(getOperationContextConfig()) + .add(getAuthorizerContext()) + .add(getActorContext()) + .add(getSearchContext()) + .add(getEntityRegistryContext()) + .build() + .stream() + .map(ContextInterface::getCacheKeyComponent) + .filter(Optional::isPresent) + .mapToInt(Optional::get) + .sum()); + } + + // Context id specific to contexts which impact search responses + public String getSearchContextId() { + return String.valueOf( + ImmutableSet.builder() + .add(getOperationContextConfig()) + .add(getActorContext()) + .add(getSearchContext()) + .add(getEntityRegistryContext()) + .build() + .stream() + .map(ContextInterface::getCacheKeyComponent) + .filter(Optional::isPresent) + .mapToInt(Optional::get) + .sum()); + } + + // Context id specific to entity lookups (not search) + public String getEntityContextId() { + return String.valueOf( + ImmutableSet.builder() + .add(getOperationContextConfig()) + .add(getActorContext()) + .add(getEntityRegistryContext()) + .build() + .stream() + .map(ContextInterface::getCacheKeyComponent) + .filter(Optional::isPresent) + .mapToInt(Optional::get) + .sum()); + } + + public static class OperationContextBuilder { + + public OperationContext build(@Nonnull Authentication sessionAuthentication) { + final Urn actorUrn = UrnUtils.getUrn(sessionAuthentication.getActor().toUrnStr()); + return new OperationContext( + this.operationContextConfig, + ActorContext.builder() + .authentication(sessionAuthentication) + .systemAuth( + this.systemActorContext != null + && this.systemActorContext.getAuthentication().equals(sessionAuthentication)) + .policyInfoSet(this.authorizerContext.getAuthorizer().getActorPolicies(actorUrn)) + .groupMembership(this.authorizerContext.getAuthorizer().getActorGroups(actorUrn)) + .build(), + this.systemActorContext, + Objects.requireNonNull(this.searchContext), + Objects.requireNonNull(this.authorizerContext), + Objects.requireNonNull(this.entityRegistryContext)); + } + + private OperationContext build() { + return null; + } + } +} diff --git a/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/OperationContextConfig.java b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/OperationContextConfig.java new file mode 100644 index 00000000000000..f0e12f5a0ce2b2 --- /dev/null +++ b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/OperationContextConfig.java @@ -0,0 +1,24 @@ +package io.datahubproject.metadata.context; + +import com.datahub.authorization.config.SearchAuthorizationConfiguration; +import java.util.Optional; +import lombok.Builder; +import lombok.Getter; + +@Builder(toBuilder = true) +@Getter +public class OperationContextConfig implements ContextInterface { + /** + * Whether the given session authentication is allowed to assume the system authentication as + * needed + */ + private final boolean allowSystemAuthentication; + + /** Configuration for search authorization */ + private final SearchAuthorizationConfiguration searchAuthorizationConfiguration; + + @Override + public Optional getCacheKeyComponent() { + return Optional.of(searchAuthorizationConfiguration.hashCode()); + } +} diff --git a/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/SearchContext.java b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/SearchContext.java new file mode 100644 index 00000000000000..d4e3712309d6c0 --- /dev/null +++ b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/SearchContext.java @@ -0,0 +1,86 @@ +package io.datahubproject.metadata.context; + +import com.linkedin.metadata.query.SearchFlags; +import com.linkedin.metadata.utils.elasticsearch.IndexConvention; +import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import java.util.Optional; +import java.util.function.Function; +import java.util.stream.Stream; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import lombok.Builder; +import lombok.Getter; + +@Builder(toBuilder = true) +@Getter +public class SearchContext implements ContextInterface { + + public static SearchContext EMPTY = + SearchContext.builder().indexConvention(IndexConventionImpl.NO_PREFIX).build(); + + public static SearchContext withFlagDefaults( + @Nonnull SearchContext searchContext, + @Nonnull Function flagDefaults) { + return searchContext.toBuilder() + // update search flags + .searchFlags(flagDefaults.apply(searchContext.getSearchFlags())) + .build(); + } + + @Nonnull private final IndexConvention indexConvention; + @Nonnull private final SearchFlags searchFlags; + + public boolean isRestrictedSearch() { + return Optional.ofNullable(searchFlags.isIncludeRestricted()).orElse(false); + } + + public SearchContext withFlagDefaults(Function flagDefaults) { + return SearchContext.withFlagDefaults(this, flagDefaults); + } + + /** + * Currently relying on the consistent hashing of String + * + * @return + */ + @Override + public Optional getCacheKeyComponent() { + return Optional.of( + Stream.of(indexConvention.getPrefix().orElse(""), keySearchFlags().toString()) + .mapToInt(String::hashCode) + .sum()); + } + + /** + * Only certain flags change the cache key + * + * @return + */ + private SearchFlags keySearchFlags() { + try { + // whether cache is enabled or not does not impact the result + return searchFlags.clone().setSkipCache(false); + } catch (CloneNotSupportedException e) { + throw new RuntimeException(e); + } + } + + public static class SearchContextBuilder { + + public SearchContextBuilder searchFlags(@Nullable SearchFlags searchFlags) { + this.searchFlags = searchFlags != null ? searchFlags : buildDefaultSearchFlags(); + return this; + } + + public SearchContext build() { + if (this.searchFlags == null) { + searchFlags(buildDefaultSearchFlags()); + } + return new SearchContext(this.indexConvention, this.searchFlags); + } + } + + private static SearchFlags buildDefaultSearchFlags() { + return new SearchFlags().setSkipCache(false); + } +} diff --git a/metadata-operation-context/src/main/java/io/datahubproject/test/metadata/context/TestOperationContexts.java b/metadata-operation-context/src/main/java/io/datahubproject/test/metadata/context/TestOperationContexts.java new file mode 100644 index 00000000000000..88e2f7f04ca5a6 --- /dev/null +++ b/metadata-operation-context/src/main/java/io/datahubproject/test/metadata/context/TestOperationContexts.java @@ -0,0 +1,69 @@ +package io.datahubproject.test.metadata.context; + +import com.datahub.authentication.Actor; +import com.datahub.authentication.ActorType; +import com.datahub.authentication.Authentication; +import com.datahub.authorization.config.SearchAuthorizationConfiguration; +import com.datahub.plugins.auth.authorization.Authorizer; +import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.utils.elasticsearch.IndexConvention; +import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.metadata.context.OperationContextConfig; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +/** + * Useful for testing. If the defaults are not sufficient, try using the .toBuilder() and replacing + * the parts that you are interested in customizing. + */ +public class TestOperationContexts { + public static final Authentication TEST_SYSTEM_AUTH = + new Authentication(new Actor(ActorType.USER, "testSystemUser"), ""); + public static final Authentication TEST_USER_AUTH = + new Authentication(new Actor(ActorType.USER, "datahub"), ""); + public static final IndexConvention TEST_EMPTY_INDEX_CONVENTION = IndexConventionImpl.NO_PREFIX; + + public static OperationContext systemContextNoSearchAuthorization( + @Nonnull EntityRegistry entityRegistry) { + return systemContextNoSearchAuthorization(entityRegistry, null); + } + + public static OperationContext systemContextNoSearchAuthorization( + @Nonnull EntityRegistry entityRegistry, @Nullable IndexConvention indexConvention) { + return OperationContext.asSystem( + OperationContextConfig.builder() + .searchAuthorizationConfiguration( + SearchAuthorizationConfiguration.builder().enabled(false).build()) + .build(), + entityRegistry, + TEST_SYSTEM_AUTH, + indexConvention != null ? indexConvention : TEST_EMPTY_INDEX_CONVENTION); + } + + public static OperationContext userContextNoSearchAuthorization( + @Nonnull EntityRegistry entityRegistry, @Nonnull Urn userUrn) { + return userContextNoSearchAuthorization(entityRegistry, Authorizer.EMPTY, userUrn); + } + + public static OperationContext userContextNoSearchAuthorization( + @Nonnull EntityRegistry entityRegistry, + @Nonnull Authorizer authorizer, + @Nonnull Urn userUrn) { + return userContextNoSearchAuthorization( + entityRegistry, + authorizer, + new Authentication(new Actor(ActorType.USER, userUrn.getId()), "")); + } + + public static OperationContext userContextNoSearchAuthorization( + @Nonnull EntityRegistry entityRegistry, + @Nonnull Authorizer authorizer, + @Nonnull Authentication sessionAuthorization) { + return systemContextNoSearchAuthorization(entityRegistry) + .asSession(authorizer, sessionAuthorization); + } + + private TestOperationContexts() {} +} diff --git a/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/ActorContextTest.java b/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/ActorContextTest.java new file mode 100644 index 00000000000000..15fe2bc277b9b9 --- /dev/null +++ b/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/ActorContextTest.java @@ -0,0 +1,129 @@ +package io.datahubproject.metadata.context; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotEquals; + +import com.datahub.authentication.Actor; +import com.datahub.authentication.ActorType; +import com.datahub.authentication.Authentication; +import com.linkedin.common.UrnArray; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.data.template.StringArray; +import com.linkedin.metadata.authorization.PoliciesConfig; +import com.linkedin.policy.DataHubActorFilter; +import com.linkedin.policy.DataHubPolicyInfo; +import com.linkedin.policy.DataHubResourceFilter; +import com.linkedin.policy.PolicyMatchCondition; +import com.linkedin.policy.PolicyMatchCriterion; +import com.linkedin.policy.PolicyMatchCriterionArray; +import com.linkedin.policy.PolicyMatchFilter; +import java.util.List; +import java.util.Set; +import org.testng.annotations.Test; + +public class ActorContextTest { + + private static final DataHubPolicyInfo POLICY_ABC = + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setActors( + new DataHubActorFilter() + .setUsers( + new UrnArray( + UrnUtils.getUrn("urn:li:corpUser:userA"), + UrnUtils.getUrn("urn:li:corpUser:userB")))) + .setPrivileges(new StringArray(List.of("a", "b", "c"))); + + private static final DataHubPolicyInfo POLICY_D = + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setActors( + new DataHubActorFilter() + .setUsers( + new UrnArray( + UrnUtils.getUrn("urn:li:corpUser:userA"), + UrnUtils.getUrn("urn:li:corpUser:userB")))) + .setPrivileges(new StringArray(List.of("d"))); + + private static final DataHubPolicyInfo POLICY_ABC_RESOURCE = + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setActors( + new DataHubActorFilter() + .setUsers( + new UrnArray( + UrnUtils.getUrn("urn:li:corpUser:userA"), + UrnUtils.getUrn("urn:li:corpUser:userB")))) + .setResources( + new DataHubResourceFilter() + .setFilter( + new PolicyMatchFilter() + .setCriteria( + new PolicyMatchCriterionArray( + List.of( + new PolicyMatchCriterion() + .setField("tag") + .setCondition(PolicyMatchCondition.EQUALS) + .setValues(new StringArray("urn:li:tag:test"))))))) + .setPrivileges(new StringArray(List.of("a", "b", "c"))); + + private static final DataHubPolicyInfo POLICY_D_OWNER = + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setActors(new DataHubActorFilter().setResourceOwners(true)) + .setPrivileges(new StringArray(List.of("d"))); + + private static final DataHubPolicyInfo POLICY_D_OWNER_TYPE = + new DataHubPolicyInfo() + .setState(PoliciesConfig.ACTIVE_POLICY_STATE) + .setActors( + new DataHubActorFilter() + .setResourceOwnersTypes( + new UrnArray(UrnUtils.getUrn("urn:li:ownershipType:test")))) + .setPrivileges(new StringArray(List.of("d"))); + + @Test + public void actorContextId() { + Authentication userAuth = new Authentication(new Actor(ActorType.USER, "USER"), ""); + + assertEquals( + ActorContext.asSessionRestricted(userAuth, Set.of(), Set.of()).getCacheKeyComponent(), + ActorContext.asSessionRestricted(userAuth, Set.of(), Set.of()).getCacheKeyComponent(), + "Expected equality across instances"); + + assertEquals( + ActorContext.asSessionRestricted(userAuth, Set.of(), Set.of()).getCacheKeyComponent(), + ActorContext.asSessionRestricted( + userAuth, Set.of(), Set.of(UrnUtils.getUrn("urn:li:corpGroup:group1"))) + .getCacheKeyComponent(), + "Expected no impact to cache context from group membership"); + + assertEquals( + ActorContext.asSessionRestricted(userAuth, Set.of(POLICY_ABC, POLICY_D), Set.of()) + .getCacheKeyComponent(), + ActorContext.asSessionRestricted(userAuth, Set.of(POLICY_ABC, POLICY_D), Set.of()) + .getCacheKeyComponent(), + "Expected equality when non-ownership policies are identical"); + + assertNotEquals( + ActorContext.asSessionRestricted(userAuth, Set.of(POLICY_ABC_RESOURCE, POLICY_D), Set.of()) + .getCacheKeyComponent(), + ActorContext.asSessionRestricted(userAuth, Set.of(POLICY_ABC, POLICY_D), Set.of()) + .getCacheKeyComponent(), + "Expected differences with non-identical resource policy"); + + assertNotEquals( + ActorContext.asSessionRestricted(userAuth, Set.of(POLICY_D_OWNER), Set.of()) + .getCacheKeyComponent(), + ActorContext.asSessionRestricted(userAuth, Set.of(POLICY_D), Set.of()) + .getCacheKeyComponent(), + "Expected differences with ownership policy"); + + assertNotEquals( + ActorContext.asSessionRestricted(userAuth, Set.of(POLICY_D_OWNER_TYPE), Set.of()) + .getCacheKeyComponent(), + ActorContext.asSessionRestricted(userAuth, Set.of(POLICY_D), Set.of()) + .getCacheKeyComponent(), + "Expected differences with ownership type policy"); + } +} diff --git a/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/OperationContextTest.java b/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/OperationContextTest.java new file mode 100644 index 00000000000000..81583deba0e6ca --- /dev/null +++ b/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/OperationContextTest.java @@ -0,0 +1,67 @@ +package io.datahubproject.metadata.context; + +import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertEquals; + +import com.datahub.authentication.Actor; +import com.datahub.authentication.ActorType; +import com.datahub.authentication.Authentication; +import com.datahub.plugins.auth.authorization.Authorizer; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import org.testng.annotations.Test; + +public class OperationContextTest { + + @Test + public void testSystemPrivilegeEscalation() { + Authentication systemAuth = new Authentication(new Actor(ActorType.USER, "SYSTEM"), ""); + Authentication userAuth = new Authentication(new Actor(ActorType.USER, "USER"), ""); + + // Allows system authentication + OperationContext systemOpContext = + OperationContext.asSystem( + OperationContextConfig.builder().build(), + mock(EntityRegistry.class), + systemAuth, + IndexConventionImpl.NO_PREFIX); + + OperationContext opContext = systemOpContext.asSession(Authorizer.EMPTY, userAuth); + + assertEquals( + opContext.getAuthentication(), systemAuth, "Expected system authentication when allowed"); + assertEquals( + opContext.getAuditStamp().getActor().getId(), + "USER", + "Audit stamp expected to match the user's identity"); + assertEquals(opContext.getSessionAuthentication(), userAuth); + assertEquals(opContext.getSessionActorContext().getAuthentication(), userAuth); + assertEquals(opContext.getActorContext().getAuthentication(), systemAuth); + assertEquals(opContext.getSystemActorContext().getAuthentication(), systemAuth); + assertEquals(opContext.getSystemAuthentication().get(), systemAuth); + + // Do not allow system auth + OperationContext opContextNoSystem = + systemOpContext.toBuilder() + .operationContextConfig( + systemOpContext.getOperationContextConfig().toBuilder() + .allowSystemAuthentication(false) + .build()) + .build(userAuth); + + assertEquals( + opContextNoSystem.getAuthentication(), + userAuth, + "Expect user authentication when system authentication is not allowed"); + assertEquals( + opContextNoSystem.getAuditStamp().getActor().getId(), + "USER", + "Audit stamp expected to match the user's identity"); + assertEquals(opContextNoSystem.getSessionActorContext().getAuthentication(), userAuth); + assertEquals(opContextNoSystem.getActorContext().getAuthentication(), userAuth); + assertEquals(opContextNoSystem.getSystemActorContext().getAuthentication(), systemAuth); + assertEquals(opContextNoSystem.getSystemAuthentication().get(), systemAuth); + assertEquals(opContextNoSystem.getSystemActorContext().getAuthentication(), systemAuth); + assertEquals(opContextNoSystem.getSessionAuthentication(), userAuth); + } +} diff --git a/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/SearchContextTest.java b/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/SearchContextTest.java new file mode 100644 index 00000000000000..26365c283fc57d --- /dev/null +++ b/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/SearchContextTest.java @@ -0,0 +1,79 @@ +package io.datahubproject.metadata.context; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotEquals; + +import com.linkedin.metadata.query.SearchFlags; +import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import org.testng.annotations.Test; + +public class SearchContextTest { + + @Test + public void searchContextId() { + SearchContext testNoFlags = + SearchContext.builder().indexConvention(IndexConventionImpl.NO_PREFIX).build(); + + assertEquals( + testNoFlags.getCacheKeyComponent(), + SearchContext.builder() + .indexConvention(IndexConventionImpl.NO_PREFIX) + .build() + .getCacheKeyComponent(), + "Expected consistent context ids across instances"); + + SearchContext testWithFlags = + SearchContext.builder() + .indexConvention(IndexConventionImpl.NO_PREFIX) + .searchFlags(new SearchFlags().setFulltext(true)) + .build(); + + assertEquals( + testWithFlags.getCacheKeyComponent(), + SearchContext.builder() + .indexConvention(IndexConventionImpl.NO_PREFIX) + .searchFlags(new SearchFlags().setFulltext(true)) + .build() + .getCacheKeyComponent(), + "Expected consistent context ids across instances"); + + assertNotEquals( + testNoFlags.getCacheKeyComponent(), + testWithFlags.getCacheKeyComponent(), + "Expected differences in search flags to result in different caches"); + assertNotEquals( + testWithFlags.getCacheKeyComponent(), + SearchContext.builder() + .indexConvention(IndexConventionImpl.NO_PREFIX) + .searchFlags(new SearchFlags().setFulltext(true).setIncludeRestricted(true)) + .build() + .getCacheKeyComponent(), + "Expected differences in search flags to result in different caches"); + + assertNotEquals( + testNoFlags.getCacheKeyComponent(), + SearchContext.builder() + .indexConvention(new IndexConventionImpl("Some Prefix")) + .searchFlags(null) + .build() + .getCacheKeyComponent(), + "Expected differences in index convention to result in different caches"); + + assertNotEquals( + SearchContext.builder() + .indexConvention(IndexConventionImpl.NO_PREFIX) + .searchFlags( + new SearchFlags() + .setFulltext(false) + .setIncludeRestricted(true) + .setSkipAggregates(true)) + .build() + .getCacheKeyComponent(), + SearchContext.builder() + .indexConvention(IndexConventionImpl.NO_PREFIX) + .searchFlags(new SearchFlags().setFulltext(true).setIncludeRestricted(true)) + .build() + .getCacheKeyComponent(), + "Expected differences in search flags to result in different caches"); + } +} diff --git a/metadata-service/auth-config/src/main/java/com/datahub/authorization/AuthorizationConfiguration.java b/metadata-service/auth-config/src/main/java/com/datahub/authorization/AuthorizationConfiguration.java index 5ed69d3e2ff8c4..aefa4d17b42c92 100644 --- a/metadata-service/auth-config/src/main/java/com/datahub/authorization/AuthorizationConfiguration.java +++ b/metadata-service/auth-config/src/main/java/com/datahub/authorization/AuthorizationConfiguration.java @@ -1,5 +1,6 @@ package com.datahub.authorization; +import com.datahub.authorization.config.SearchAuthorizationConfiguration; import com.datahub.plugins.auth.authorization.Authorizer; import java.util.List; import lombok.Data; @@ -12,4 +13,6 @@ public class AuthorizationConfiguration { /** List of configurations for {@link Authorizer}s to be registered */ private List authorizers; + + private SearchAuthorizationConfiguration search; } diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authentication/invite/InviteTokenService.java b/metadata-service/auth-impl/src/main/java/com/datahub/authentication/invite/InviteTokenService.java index 73add48958f609..731cf081853848 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authentication/invite/InviteTokenService.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authentication/invite/InviteTokenService.java @@ -20,6 +20,7 @@ import com.linkedin.metadata.secret.SecretService; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.Collections; import javax.annotation.Nonnull; @@ -58,34 +59,31 @@ public Urn getInviteTokenRole( @Nonnull public String getInviteToken( - @Nullable final String roleUrnStr, - boolean regenerate, - @Nonnull final Authentication authentication) + @Nonnull OperationContext opContext, @Nullable final String roleUrnStr, boolean regenerate) throws Exception { final Filter inviteTokenFilter = roleUrnStr == null ? createInviteTokenFilter() : createInviteTokenFilter(roleUrnStr); final SearchResult searchResult = - _entityClient.filter( - INVITE_TOKEN_ENTITY_NAME, inviteTokenFilter, null, 0, 10, authentication); + _entityClient.filter(opContext, INVITE_TOKEN_ENTITY_NAME, inviteTokenFilter, null, 0, 10); final int numEntities = searchResult.getEntities().size(); // If there is more than one invite token, wipe all of them and generate a fresh one if (numEntities > 1) { - deleteExistingInviteTokens(searchResult, authentication); - return createInviteToken(roleUrnStr, authentication); + deleteExistingInviteTokens(searchResult, opContext.getAuthentication()); + return createInviteToken(roleUrnStr, opContext.getAuthentication()); } // If we want to regenerate, or there are no entities in the result, create a new invite token. if (regenerate || numEntities == 0) { - return createInviteToken(roleUrnStr, authentication); + return createInviteToken(roleUrnStr, opContext.getAuthentication()); } final SearchEntity searchEntity = searchResult.getEntities().get(0); final Urn inviteTokenUrn = searchEntity.getEntity(); com.linkedin.identity.InviteToken inviteToken = - getInviteTokenEntity(inviteTokenUrn, authentication); + getInviteTokenEntity(inviteTokenUrn, opContext.getAuthentication()); return _secretService.decrypt(inviteToken.getToken()); } 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 9e8c1928c9de09..5663ffffdb3d60 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 @@ -2,7 +2,9 @@ import com.datahub.plugins.auth.authorization.Authorizer; import com.linkedin.common.urn.Urn; +import com.linkedin.policy.DataHubPolicyInfo; import java.util.ArrayList; +import java.util.Collection; import java.util.Collections; import java.util.HashSet; import java.util.List; @@ -10,6 +12,7 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; @@ -154,4 +157,26 @@ private AuthorizedActors mergeAuthorizedActors( public DataHubAuthorizer getDefaultAuthorizer() { return (DataHubAuthorizer) defaultAuthorizer; } + + @Override + public Set getActorPolicies(@Nonnull Urn actorUrn) { + return authorizers.stream() + .flatMap(authorizer -> authorizer.getActorPolicies(actorUrn).stream()) + .collect(Collectors.toSet()); + } + + @Override + public Collection getActorGroups(@Nonnull Urn actorUrn) { + return authorizers.stream() + .flatMap(authorizer -> authorizer.getActorGroups(actorUrn).stream()) + .collect(Collectors.toList()); + } + + @Override + public Collection getActorPeers(@Nonnull Urn actorUrn) { + return authorizers.stream() + .flatMap(authorizer -> authorizer.getActorPeers(actorUrn).stream()) + .distinct() + .collect(Collectors.toList()); + } } 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 350d57aae37834..b5c6910776e52b 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,19 +8,23 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.policy.DataHubPolicyInfo; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.ArrayList; +import java.util.Collection; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.Set; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.stream.Collectors; import javax.annotation.Nonnull; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -46,9 +50,6 @@ public enum AuthorizationMode { ALLOW_ALL } - // Credentials used to make / authorize requests as the internal system actor. - private final Authentication _systemAuthentication; - // Maps privilege name to the associated set of policies for fast access. // Not concurrent data structure because writes are always against the entire thing. private final Map> _policyCache = @@ -62,22 +63,24 @@ public enum AuthorizationMode { private final PolicyEngine _policyEngine; private EntitySpecResolver _entitySpecResolver; private AuthorizationMode _mode; + private final OperationContext systemOpContext; public static final String ALL = "ALL"; public DataHubAuthorizer( - final Authentication systemAuthentication, + @Nonnull final OperationContext systemOpContext, final EntityClient entityClient, final int delayIntervalSeconds, final int refreshIntervalSeconds, final AuthorizationMode mode, final int policyFetchSize) { - _systemAuthentication = Objects.requireNonNull(systemAuthentication); + this.systemOpContext = systemOpContext; _mode = Objects.requireNonNull(mode); - _policyEngine = new PolicyEngine(systemAuthentication, Objects.requireNonNull(entityClient)); + _policyEngine = + new PolicyEngine(systemOpContext.getAuthentication(), Objects.requireNonNull(entityClient)); _policyRefreshRunnable = new PolicyRefreshRunnable( - systemAuthentication, + systemOpContext, new PolicyFetcher(entityClient), _policyCache, readWriteLock.writeLock(), @@ -95,7 +98,7 @@ public void init(@Nonnull Map authorizerConfig, @Nonnull Authori public AuthorizationResult authorize(@Nonnull final AuthorizationRequest request) { // 0. Short circuit: If the action is being performed by the system (root), always allow it. - if (isSystemRequest(request, this._systemAuthentication)) { + if (isSystemRequest(request, systemOpContext.getAuthentication())) { return new AuthorizationResult(request, AuthorizationResult.Type.ALLOW, null); } @@ -135,6 +138,45 @@ public List getGrantedPrivileges( policiesToEvaluate, resolvedActorSpec, resolvedResourceSpec); } + @Override + public Set getActorPolicies(@Nonnull Urn actorUrn) { + // 1. Fetch all policies + final List policiesToEvaluate = getOrDefault(ALL, new ArrayList<>()); + + // 2. Actor identity + final ResolvedEntitySpec resolvedActorSpec = + _entitySpecResolver.resolve(new EntitySpec(actorUrn.getEntityType(), actorUrn.toString())); + + return policiesToEvaluate.stream() + .filter(policy -> PoliciesConfig.ACTIVE_POLICY_STATE.equals(policy.getState())) + .filter( + policy -> + policy.getActors().isResourceOwners() + || _policyEngine.isActorMatch( + resolvedActorSpec, + policy.getActors(), + Optional.empty(), + new PolicyEngine.PolicyEvaluationContext())) + .collect(Collectors.toSet()); + } + + @Override + public Collection getActorGroups(@Nonnull Urn actorUrn) { + // 1. Actor identity + final ResolvedEntitySpec resolvedActorSpec = + _entitySpecResolver.resolve(new EntitySpec(actorUrn.getEntityType(), actorUrn.toString())); + + return resolvedActorSpec.getGroupMembership().stream() + .map(UrnUtils::getUrn) + .collect(Collectors.toList()); + } + + @Override + public Collection getActorPeers(@Nonnull Urn actorUrn) { + // TODO: Fetch users from groups the actor is a member of + return List.of(actorUrn); + } + /** * Retrieves the current list of actors authorized to for a particular privilege against an * optional resource @@ -261,7 +303,7 @@ private List getOrDefault(String key, List @RequiredArgsConstructor static class PolicyRefreshRunnable implements Runnable { - private final Authentication _systemAuthentication; + private final OperationContext systemOpContext; private final PolicyFetcher _policyFetcher; private final Map> _policyCache; private final Lock writeLock; @@ -278,7 +320,7 @@ public void run() { while (total == null || scrollId != null) { try { final PolicyFetcher.PolicyFetchResult policyFetchResult = - _policyFetcher.fetchPolicies(count, scrollId, null, _systemAuthentication); + _policyFetcher.fetchPolicies(systemOpContext, count, scrollId, null); addPoliciesToCache(newCache, policyFetchResult.getPolicies()); 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 f078d2d316cae2..13c50059031b63 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 @@ -256,11 +256,10 @@ private boolean checkCondition( } /** - * 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. + * Returns true if the actor portion of a DataHub policy matches the actor being evaluated, false + * otherwise. */ - private boolean isActorMatch( + boolean isActorMatch( final ResolvedEntitySpec resolvedActorSpec, final DataHubActorFilter actorFilter, final Optional resourceSpec, diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyFetcher.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyFetcher.java index 0485e3000ad173..cbd80c7755adb9 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyFetcher.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyFetcher.java @@ -3,12 +3,10 @@ import static com.linkedin.metadata.Constants.DATAHUB_POLICY_INFO_ASPECT_NAME; import static com.linkedin.metadata.Constants.POLICY_ENTITY_NAME; -import com.datahub.authentication.Authentication; import com.linkedin.common.urn.Urn; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.entity.client.EntityClient; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.query.filter.SortOrder; @@ -16,6 +14,7 @@ import com.linkedin.metadata.search.SearchEntity; import com.linkedin.policy.DataHubPolicyInfo; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.Collections; import java.util.HashSet; @@ -33,7 +32,7 @@ @Slf4j @RequiredArgsConstructor public class PolicyFetcher { - private final EntityClient _entityClient; + private final EntityClient entityClient; private static final SortCriterion POLICY_SORT_CRITERION = new SortCriterion().setField("lastUpdatedTimestamp").setOrder(SortOrder.DESCENDING); @@ -46,7 +45,7 @@ public class PolicyFetcher { */ @Deprecated public CompletableFuture fetchPolicies( - int start, String query, int count, Filter filter, Authentication authentication) { + OperationContext opContext, int start, String query, int count, Filter filter) { return CompletableFuture.supplyAsync( () -> { try { @@ -57,7 +56,7 @@ public CompletableFuture fetchPolicies( while (PolicyFetchResult.EMPTY.equals(result) && scrollId != null) { PolicyFetchResult tmpResult = fetchPolicies( - query, count, scrollId.isEmpty() ? null : scrollId, filter, authentication); + opContext, query, count, scrollId.isEmpty() ? null : scrollId, filter); fetchedResults += tmpResult.getPolicies().size(); scrollId = tmpResult.getScrollId(); if (fetchedResults > start) { @@ -73,35 +72,32 @@ public CompletableFuture fetchPolicies( } public PolicyFetchResult fetchPolicies( - int count, @Nullable String scrollId, Filter filter, Authentication authentication) + OperationContext opContext, int count, @Nullable String scrollId, Filter filter) throws RemoteInvocationException, URISyntaxException { - return fetchPolicies("", count, scrollId, filter, authentication); + return fetchPolicies(opContext, "", count, scrollId, filter); } public PolicyFetchResult fetchPolicies( - String query, - int count, - @Nullable String scrollId, - Filter filter, - Authentication authentication) + OperationContext opContext, String query, int count, @Nullable String scrollId, Filter filter) throws RemoteInvocationException, URISyntaxException { log.debug(String.format("Batch fetching policies. count: %s, scroll: %s", count, scrollId)); // First fetch all policy urns ScrollResult result = - _entityClient.scrollAcrossEntities( + entityClient.scrollAcrossEntities( + opContext.withSearchFlags( + flags -> + flags + .setSkipCache(true) + .setSkipAggregates(true) + .setSkipHighlighting(true) + .setFulltext(true)), List.of(POLICY_ENTITY_NAME), query, filter, scrollId, null, - count, - new SearchFlags() - .setSkipCache(true) - .setSkipAggregates(true) - .setSkipHighlighting(true) - .setFulltext(true), - authentication); + count); List policyUrns = result.getEntities().stream().map(SearchEntity::getEntity).collect(Collectors.toList()); @@ -111,8 +107,8 @@ public PolicyFetchResult fetchPolicies( // Fetch DataHubPolicyInfo aspects for each urn final Map policyEntities = - _entityClient.batchGetV2( - POLICY_ENTITY_NAME, new HashSet<>(policyUrns), null, authentication); + entityClient.batchGetV2( + POLICY_ENTITY_NAME, new HashSet<>(policyUrns), null, opContext.getAuthentication()); return new PolicyFetchResult( policyUrns.stream() .map(policyEntities::get) diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/invite/InviteTokenServiceTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/invite/InviteTokenServiceTest.java index cd9d5972103c12..775039766a2c90 100644 --- a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/invite/InviteTokenServiceTest.java +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/invite/InviteTokenServiceTest.java @@ -18,6 +18,7 @@ import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.secret.SecretService; +import io.datahubproject.metadata.context.OperationContext; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -35,6 +36,7 @@ public class InviteTokenServiceTest { private EntityClient _entityClient; private SecretService _secretService; private InviteTokenService _inviteTokenService; + private OperationContext opContext; @BeforeMethod public void setupTest() throws Exception { @@ -42,7 +44,7 @@ public void setupTest() throws Exception { roleUrn = Urn.createFromString(ROLE_URN_STRING); _entityClient = mock(EntityClient.class); _secretService = mock(SecretService.class); - + opContext = mock(OperationContext.class); _inviteTokenService = new InviteTokenService(_entityClient, _secretService); } @@ -129,27 +131,22 @@ public void testGetInviteTokenRole() throws Exception { public void getInviteTokenRoleUrnDoesNotExist() throws Exception { when(_entityClient.exists(eq(roleUrn), eq(SYSTEM_AUTHENTICATION))).thenReturn(false); - assertThrows( - () -> _inviteTokenService.getInviteToken(roleUrn.toString(), false, SYSTEM_AUTHENTICATION)); + assertThrows(() -> _inviteTokenService.getInviteToken(opContext, roleUrn.toString(), false)); } @Test public void getInviteTokenRegenerate() throws Exception { final SearchResult searchResult = new SearchResult(); searchResult.setEntities(new SearchEntityArray()); + when(opContext.getAuthentication()).thenReturn(SYSTEM_AUTHENTICATION); when(_entityClient.filter( - eq(INVITE_TOKEN_ENTITY_NAME), - any(), - any(), - anyInt(), - anyInt(), - eq(SYSTEM_AUTHENTICATION))) + eq(opContext), eq(INVITE_TOKEN_ENTITY_NAME), any(), any(), anyInt(), anyInt())) .thenReturn(searchResult); when(_secretService.generateUrlSafeToken(anyInt())).thenReturn(INVITE_TOKEN_STRING); when(_secretService.hashString(anyString())).thenReturn(HASHED_INVITE_TOKEN_STRING); when(_secretService.encrypt(anyString())).thenReturn(ENCRYPTED_INVITE_TOKEN_STRING); - _inviteTokenService.getInviteToken(null, true, SYSTEM_AUTHENTICATION); + _inviteTokenService.getInviteToken(opContext, null, true); verify(_entityClient, times(1)).ingestProposal(any(), eq(SYSTEM_AUTHENTICATION)); } @@ -157,19 +154,15 @@ public void getInviteTokenRegenerate() throws Exception { public void getInviteTokenEmptySearchResult() throws Exception { final SearchResult searchResult = new SearchResult(); searchResult.setEntities(new SearchEntityArray()); + when(opContext.getAuthentication()).thenReturn(SYSTEM_AUTHENTICATION); when(_entityClient.filter( - eq(INVITE_TOKEN_ENTITY_NAME), - any(), - any(), - anyInt(), - anyInt(), - eq(SYSTEM_AUTHENTICATION))) + eq(opContext), eq(INVITE_TOKEN_ENTITY_NAME), any(), any(), anyInt(), anyInt())) .thenReturn(searchResult); when(_secretService.generateUrlSafeToken(anyInt())).thenReturn(INVITE_TOKEN_STRING); when(_secretService.hashString(anyString())).thenReturn(HASHED_INVITE_TOKEN_STRING); when(_secretService.encrypt(anyString())).thenReturn(ENCRYPTED_INVITE_TOKEN_STRING); - _inviteTokenService.getInviteToken(null, false, SYSTEM_AUTHENTICATION); + _inviteTokenService.getInviteToken(opContext, null, false); verify(_entityClient, times(1)).ingestProposal(any(), eq(SYSTEM_AUTHENTICATION)); } @@ -180,19 +173,15 @@ public void getInviteTokenNullEntity() throws Exception { final SearchEntity searchEntity = new SearchEntity().setEntity(inviteTokenUrn); searchEntityArray.add(searchEntity); searchResult.setEntities(searchEntityArray); + when(opContext.getAuthentication()).thenReturn(SYSTEM_AUTHENTICATION); when(_entityClient.filter( - eq(INVITE_TOKEN_ENTITY_NAME), - any(), - any(), - anyInt(), - anyInt(), - eq(SYSTEM_AUTHENTICATION))) + eq(opContext), eq(INVITE_TOKEN_ENTITY_NAME), any(), any(), anyInt(), anyInt())) .thenReturn(searchResult); when(_entityClient.getV2( eq(INVITE_TOKEN_ENTITY_NAME), eq(inviteTokenUrn), any(), eq(SYSTEM_AUTHENTICATION))) .thenReturn(null); - assertThrows(() -> _inviteTokenService.getInviteToken(null, false, SYSTEM_AUTHENTICATION)); + assertThrows(() -> _inviteTokenService.getInviteToken(opContext, null, false)); } @Test @@ -203,12 +192,7 @@ public void getInviteTokenNoInviteTokenAspect() throws Exception { searchEntityArray.add(searchEntity); searchResult.setEntities(searchEntityArray); when(_entityClient.filter( - eq(INVITE_TOKEN_ENTITY_NAME), - any(), - any(), - anyInt(), - anyInt(), - eq(SYSTEM_AUTHENTICATION))) + eq(opContext), eq(INVITE_TOKEN_ENTITY_NAME), any(), any(), anyInt(), anyInt())) .thenReturn(searchResult); final EntityResponse entityResponse = new EntityResponse().setAspects(new EnvelopedAspectMap()); @@ -218,7 +202,7 @@ public void getInviteTokenNoInviteTokenAspect() throws Exception { when(_secretService.encrypt(anyString())).thenReturn(ENCRYPTED_INVITE_TOKEN_STRING); - assertThrows(() -> _inviteTokenService.getInviteToken(null, false, SYSTEM_AUTHENTICATION)); + assertThrows(() -> _inviteTokenService.getInviteToken(opContext, null, false)); } @Test @@ -228,13 +212,9 @@ public void getInviteToken() throws Exception { final SearchEntity searchEntity = new SearchEntity().setEntity(inviteTokenUrn); searchEntityArray.add(searchEntity); searchResult.setEntities(searchEntityArray); + when(opContext.getAuthentication()).thenReturn(SYSTEM_AUTHENTICATION); when(_entityClient.filter( - eq(INVITE_TOKEN_ENTITY_NAME), - any(), - any(), - anyInt(), - anyInt(), - eq(SYSTEM_AUTHENTICATION))) + eq(opContext), eq(INVITE_TOKEN_ENTITY_NAME), any(), any(), anyInt(), anyInt())) .thenReturn(searchResult); final EntityResponse entityResponse = new EntityResponse(); @@ -251,8 +231,6 @@ public void getInviteToken() throws Exception { when(_secretService.decrypt(eq(ENCRYPTED_INVITE_TOKEN_STRING))).thenReturn(INVITE_TOKEN_STRING); - assertEquals( - _inviteTokenService.getInviteToken(null, false, SYSTEM_AUTHENTICATION), - INVITE_TOKEN_STRING); + assertEquals(_inviteTokenService.getInviteToken(opContext, null, false), INVITE_TOKEN_STRING); } } 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 588cdf57269ef3..c37dc70ef0649c 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 @@ -39,14 +39,17 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.identity.GroupMembership; import com.linkedin.identity.RoleMembership; -import com.linkedin.metadata.query.SearchFlags; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.ScrollResult; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; +import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.policy.DataHubActorFilter; import com.linkedin.policy.DataHubPolicyInfo; import com.linkedin.policy.DataHubResourceFilter; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.metadata.context.OperationContextConfig; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -69,6 +72,7 @@ public class DataHubAuthorizerTest { private EntityClient _entityClient; private DataHubAuthorizer _dataHubAuthorizer; + private OperationContext systemOpContext; @BeforeMethod public void setupTest() throws Exception { @@ -158,19 +162,13 @@ public void setupTest() throws Exception { ImmutableList.of(new SearchEntity().setEntity(adminPolicyUrn)))); when(_entityClient.scrollAcrossEntities( + any(OperationContext.class), eq(List.of("dataHubPolicy")), eq(""), isNull(), any(), isNull(), - anyInt(), - eq( - new SearchFlags() - .setFulltext(true) - .setSkipAggregates(true) - .setSkipHighlighting(true) - .setSkipCache(true)), - any())) + anyInt())) .thenReturn(policySearchResult1) .thenReturn(policySearchResult2) .thenReturn(policySearchResult3) @@ -268,10 +266,16 @@ public void setupTest() throws Exception { final Authentication systemAuthentication = new Authentication(new Actor(ActorType.USER, DATAHUB_SYSTEM_CLIENT_ID), ""); + systemOpContext = + OperationContext.asSystem( + OperationContextConfig.builder().build(), + mock(EntityRegistry.class), + systemAuthentication, + mock(IndexConvention.class)); _dataHubAuthorizer = new DataHubAuthorizer( - systemAuthentication, + systemOpContext, _entityClient, 10, 10, @@ -358,14 +362,13 @@ public void testInvalidateCache() throws Exception { emptyResult.setEntities(new SearchEntityArray()); when(_entityClient.search( + any(OperationContext.class), eq("dataHubPolicy"), eq(""), isNull(), any(), anyInt(), - anyInt(), - any(), - eq(new SearchFlags().setFulltext(true)))) + anyInt())) .thenReturn(emptyResult); when(_entityClient.batchGetV2( eq(POLICY_ENTITY_NAME), eq(Collections.emptySet()), eq(null), any())) diff --git a/metadata-service/configuration/src/main/resources/application.yml b/metadata-service/configuration/src/main/resources/application.yml index 494c18b75f0234..9e824303788274 100644 --- a/metadata-service/configuration/src/main/resources/application.yml +++ b/metadata-service/configuration/src/main/resources/application.yml @@ -45,6 +45,11 @@ authorization: cachePolicyFetchSize: ${POLICY_CACHE_FETCH_SIZE:1000} # Enables authorization of reads, writes, and deletes on REST APIs. Defaults to false for backwards compatibility, but should become true down the road restApiAuthorization: ${REST_API_AUTHORIZATION_ENABLED:false} + search: + enabled: ${SEARCH_AUTHORIZATION_ENABLED:false} + recommendations: + # Currently limited to the actor only, see TODO: DataHubAuthorizer + peerGroupEnabled: ${SEARCH_AUTHORIZATION_RECOMMENDATIONS_PEER_GROUP_ENABLED:true} ingestion: # The value of cliMajorVersion is substituted in by the processResources Gradle task. @@ -327,6 +332,11 @@ systemUpdate: batchSize: ${BOOTSTRAP_SYSTEM_UPDATE_POLICY_FIELDS_BATCH_SIZE:5000} reprocess: enabled: ${REPROCESS_DEFAULT_POLICY_FIELDS:false} + ownershipTypes: + enabled: ${BOOTSTRAP_SYSTEM_UPDATE_OWNERSHIP_TYPES_ENABLED:true} + batchSize: ${BOOTSTRAP_SYSTEM_UPDATE_OWNERSHIP_TYPES_BATCH_SIZE:1000} + reprocess: + enabled: ${BOOTSTRAP_SYSTEM_UPDATE_OWNERSHIP_TYPES_REPROCESS:false} structuredProperties: enabled: ${ENABLE_STRUCTURED_PROPERTIES_HOOK:true} # applies structured properties mappings diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubAuthorizerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubAuthorizerFactory.java index 0935e8ad0e7d4d..d29b770cbf0eec 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubAuthorizerFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubAuthorizerFactory.java @@ -3,7 +3,9 @@ import com.datahub.authorization.DataHubAuthorizer; import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.metadata.spring.YamlPropertySourceFactory; +import io.datahubproject.metadata.context.OperationContext; import javax.annotation.Nonnull; +import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.beans.factory.annotation.Value; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -26,7 +28,9 @@ public class DataHubAuthorizerFactory { @Bean(name = "dataHubAuthorizer") @Scope("singleton") @Nonnull - protected DataHubAuthorizer dataHubAuthorizer(final SystemEntityClient systemEntityClient) { + protected DataHubAuthorizer dataHubAuthorizer( + @Qualifier("systemOperationContext") final OperationContext systemOpContext, + final SystemEntityClient systemEntityClient) { final DataHubAuthorizer.AuthorizationMode mode = policiesEnabled @@ -34,7 +38,7 @@ protected DataHubAuthorizer dataHubAuthorizer(final SystemEntityClient systemEnt : DataHubAuthorizer.AuthorizationMode.ALLOW_ALL; return new DataHubAuthorizer( - systemEntityClient.getSystemAuthentication(), + systemOpContext, systemEntityClient, 10, policyCacheRefreshIntervalSeconds, diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/RestrictedServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/RestrictedServiceFactory.java new file mode 100644 index 00000000000000..de161023faed99 --- /dev/null +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/RestrictedServiceFactory.java @@ -0,0 +1,28 @@ +package com.linkedin.gms.factory.auth; + +import com.linkedin.metadata.secret.SecretService; +import com.linkedin.metadata.service.RestrictedService; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; +import javax.annotation.Nonnull; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Qualifier; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.context.annotation.PropertySource; +import org.springframework.context.annotation.Scope; + +@Configuration +@PropertySource(value = "classpath:/application.yml", factory = YamlPropertySourceFactory.class) +public class RestrictedServiceFactory { + + @Autowired + @Qualifier("dataHubSecretService") + private SecretService _secretService; + + @Bean(name = "restrictedService") + @Scope("singleton") + @Nonnull + protected RestrictedService getInstance() throws Exception { + return new RestrictedService(_secretService); + } +} diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/context/SystemOperationContextFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/context/SystemOperationContextFactory.java new file mode 100644 index 00000000000000..0dfdee5fcbbbc6 --- /dev/null +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/context/SystemOperationContextFactory.java @@ -0,0 +1,44 @@ +package com.linkedin.gms.factory.context; + +import com.datahub.authentication.Authentication; +import com.linkedin.gms.factory.config.ConfigurationProvider; +import com.linkedin.gms.factory.search.BaseElasticSearchComponentsFactory; +import com.linkedin.metadata.models.registry.EntityRegistry; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.metadata.context.OperationContextConfig; +import javax.annotation.Nonnull; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Qualifier; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +@Configuration +public class SystemOperationContextFactory { + + @Autowired + @Qualifier("baseElasticSearchComponents") + private BaseElasticSearchComponentsFactory.BaseElasticSearchComponents components; + + @Bean(name = "systemOperationContext") + @Nonnull + protected OperationContext systemOperationContext( + @Nonnull final EntityRegistry entityRegistry, + @Nonnull @Qualifier("systemAuthentication") final Authentication systemAuthentication, + @Nonnull final OperationContextConfig operationContextConfig) { + + return OperationContext.asSystem( + operationContextConfig, + entityRegistry, + systemAuthentication, + components.getIndexConvention()); + } + + @Bean + @Nonnull + protected OperationContextConfig operationContextConfig( + final ConfigurationProvider configurationProvider) { + return OperationContextConfig.builder() + .searchAuthorizationConfiguration(configurationProvider.getAuthorization().getSearch()) + .build(); + } +} diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/JavaEntityClientFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/JavaEntityClientFactory.java index 530136e32662f8..9c9c2b7a70a653 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/JavaEntityClientFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/JavaEntityClientFactory.java @@ -1,6 +1,5 @@ package com.linkedin.gms.factory.entityclient; -import com.datahub.authentication.Authentication; import com.linkedin.entity.client.EntityClient; import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.metadata.client.JavaEntityClient; @@ -16,6 +15,7 @@ import com.linkedin.metadata.service.RollbackService; import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.timeseries.TimeseriesAspectService; +import io.datahubproject.metadata.context.OperationContext; import javax.inject.Singleton; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; @@ -32,6 +32,7 @@ public class JavaEntityClientFactory { @Bean("entityClient") @Singleton public EntityClient entityClient( + final OperationContext opContext, final @Qualifier("entityService") EntityService _entityService, final @Qualifier("deleteEntityService") DeleteEntityService _deleteEntityService, final @Qualifier("searchService") SearchService _searchService, @@ -43,6 +44,7 @@ public EntityClient entityClient( final @Qualifier("kafkaEventProducer") EventProducer _eventProducer, final RollbackService rollbackService) { return new JavaEntityClient( + opContext, _entityService, _deleteEntityService, _entitySearchService, @@ -57,6 +59,7 @@ public EntityClient entityClient( @Bean("systemEntityClient") @Singleton public SystemEntityClient systemEntityClient( + final @Qualifier("systemOperationContext") OperationContext systemOperationContext, final @Qualifier("entityService") EntityService _entityService, final @Qualifier("deleteEntityService") DeleteEntityService _deleteEntityService, final @Qualifier("searchService") SearchService _searchService, @@ -67,9 +70,9 @@ public SystemEntityClient systemEntityClient( final @Qualifier("relationshipSearchService") LineageSearchService _lineageSearchService, final @Qualifier("kafkaEventProducer") EventProducer _eventProducer, final RollbackService rollbackService, - final EntityClientCacheConfig entityClientCacheConfig, - @Qualifier("systemAuthentication") final Authentication systemAuthentication) { + final EntityClientCacheConfig entityClientCacheConfig) { return new SystemJavaEntityClient( + systemOperationContext, _entityService, _deleteEntityService, _entitySearchService, @@ -79,7 +82,6 @@ public SystemEntityClient systemEntityClient( _timeseriesAspectService, rollbackService, _eventProducer, - systemAuthentication, entityClientCacheConfig); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/RestliEntityClientFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/RestliEntityClientFactory.java index 88989b1833e78d..f2c970df681e6b 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/RestliEntityClientFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/RestliEntityClientFactory.java @@ -1,6 +1,5 @@ package com.linkedin.gms.factory.entityclient; -import com.datahub.authentication.Authentication; import com.linkedin.entity.client.EntityClient; import com.linkedin.entity.client.RestliEntityClient; import com.linkedin.entity.client.SystemEntityClient; @@ -10,6 +9,7 @@ import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.parseq.retry.backoff.ExponentialBackoff; import com.linkedin.restli.client.Client; +import io.datahubproject.metadata.context.OperationContext; import java.net.URI; import javax.inject.Singleton; import org.springframework.beans.factory.annotation.Qualifier; @@ -48,6 +48,7 @@ public EntityClient entityClient( @Bean("systemEntityClient") @Singleton public SystemEntityClient systemEntityClient( + @Qualifier("systemOperationContext") final OperationContext systemOperationContext, @Value("${datahub.gms.host}") String gmsHost, @Value("${datahub.gms.port}") int gmsPort, @Value("${datahub.gms.useSSL}") boolean gmsUseSSL, @@ -55,8 +56,7 @@ public SystemEntityClient systemEntityClient( @Value("${datahub.gms.sslContext.protocol}") String gmsSslProtocol, @Value("${entityClient.retryInterval:2}") int retryInterval, @Value("${entityClient.numRetries:3}") int numRetries, - final EntityClientCacheConfig entityClientCacheConfig, - @Qualifier("systemAuthentication") final Authentication systemAuthentication) { + final EntityClientCacheConfig entityClientCacheConfig) { final Client restClient; if (gmsUri != null) { @@ -66,10 +66,10 @@ public SystemEntityClient systemEntityClient( DefaultRestliClientFactory.getRestLiClient(gmsHost, gmsPort, gmsUseSSL, gmsSslProtocol); } return new SystemRestliEntityClient( + systemOperationContext, restClient, new ExponentialBackoff(retryInterval), numRetries, - systemAuthentication, entityClientCacheConfig); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/form/FormServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/form/FormServiceFactory.java index 73be819028f57c..16631c52e9103b 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/form/FormServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/form/FormServiceFactory.java @@ -3,7 +3,9 @@ import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.metadata.service.FormService; import com.linkedin.metadata.spring.YamlPropertySourceFactory; +import io.datahubproject.metadata.context.OperationContext; import javax.annotation.Nonnull; +import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.context.annotation.PropertySource; @@ -15,7 +17,10 @@ public class FormServiceFactory { @Bean(name = "formService") @Scope("singleton") @Nonnull - protected FormService getInstance(final SystemEntityClient entityClient) throws Exception { - return new FormService(entityClient, entityClient.getSystemAuthentication()); + protected FormService getInstance( + @Qualifier("systemOperationContext") OperationContext systemOpContext, + final SystemEntityClient entityClient) + throws Exception { + return new FormService(systemOpContext, entityClient); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java index 15bf674581b6a8..73fb026398d2d2 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java @@ -32,6 +32,7 @@ import com.linkedin.metadata.service.LineageService; import com.linkedin.metadata.service.OwnershipTypeService; import com.linkedin.metadata.service.QueryService; +import com.linkedin.metadata.service.RestrictedService; import com.linkedin.metadata.service.SettingsService; import com.linkedin.metadata.service.ViewService; import com.linkedin.metadata.timeline.TimelineService; @@ -167,6 +168,10 @@ public class GraphQLEngineFactory { @Qualifier("formService") private FormService formService; + @Autowired + @Qualifier("restrictedService") + private RestrictedService restrictedService; + @Value("${platformAnalytics.enabled}") // TODO: Migrate to DATAHUB_ANALYTICS_ENABLED private Boolean isAnalyticsEnabled; @@ -213,6 +218,7 @@ protected GraphQLEngine graphQLEngine( args.setQueryService(queryService); args.setFeatureFlags(configProvider.getFeatureFlags()); args.setFormService(formService); + args.setRestrictedService(restrictedService); args.setDataProductService(dataProductService); args.setGraphQLQueryComplexityLimit( configProvider.getGraphQL().getQuery().getComplexityLimit()); diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ingestion/IngestionSchedulerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ingestion/IngestionSchedulerFactory.java index 0ba953d66730c2..10833109e6a59b 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ingestion/IngestionSchedulerFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ingestion/IngestionSchedulerFactory.java @@ -5,6 +5,7 @@ import com.linkedin.gms.factory.auth.SystemAuthenticationFactory; import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.metadata.spring.YamlPropertySourceFactory; +import io.datahubproject.metadata.context.OperationContext; import javax.annotation.Nonnull; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; @@ -33,9 +34,11 @@ public class IngestionSchedulerFactory { @Bean(name = "ingestionScheduler") @Scope("singleton") @Nonnull - protected IngestionScheduler getInstance(final SystemEntityClient entityClient) { + protected IngestionScheduler getInstance( + @Qualifier("systemOperationContext") final OperationContext systemOpContext, + final SystemEntityClient entityClient) { return new IngestionScheduler( - entityClient.getSystemAuthentication(), + systemOpContext, entityClient, _configProvider.getIngestion(), _delayIntervalSeconds, diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/DomainsCandidateSourceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/DomainsCandidateSourceFactory.java index a7c2dde8b7d25e..1df73399185fd4 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/DomainsCandidateSourceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/DomainsCandidateSourceFactory.java @@ -4,6 +4,7 @@ import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.recommendation.candidatesource.DomainsCandidateSource; import com.linkedin.metadata.search.EntitySearchService; +import io.datahubproject.metadata.context.OperationContext; import javax.annotation.Nonnull; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; @@ -21,7 +22,8 @@ public class DomainsCandidateSourceFactory { @Bean(name = "domainsCandidateSource") @Nonnull - protected DomainsCandidateSource getInstance(final EntityRegistry entityRegistry) { + protected DomainsCandidateSource getInstance( + final OperationContext opContext, final EntityRegistry entityRegistry) { return new DomainsCandidateSource(entitySearchService, entityRegistry); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/usage/UsageClientFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/usage/UsageClientFactory.java index 03e066a912e446..d0fe095ddfd91f 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/usage/UsageClientFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/usage/UsageClientFactory.java @@ -1,6 +1,5 @@ package com.linkedin.gms.factory.usage; -import com.datahub.authentication.Authentication; import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.metadata.restli.DefaultRestliClientFactory; import com.linkedin.metadata.spring.YamlPropertySourceFactory; @@ -8,6 +7,7 @@ import com.linkedin.r2.transport.http.client.HttpClientFactory; import com.linkedin.restli.client.Client; import com.linkedin.usage.UsageClient; +import io.datahubproject.metadata.context.OperationContext; import java.util.HashMap; import java.util.Map; import org.springframework.beans.factory.annotation.Autowired; @@ -48,7 +48,7 @@ public class UsageClientFactory { @Bean("usageClient") public UsageClient getUsageClient( - @Qualifier("systemAuthentication") final Authentication systemAuthentication) { + @Qualifier("systemOperationContext") final OperationContext systemOperationContext) { Map params = new HashMap<>(); params.put(HttpClientFactory.HTTP_REQUEST_TIMEOUT, String.valueOf(timeoutMs)); @@ -56,10 +56,10 @@ public UsageClient getUsageClient( DefaultRestliClientFactory.getRestLiClient( gmsHost, gmsPort, gmsUseSSL, gmsSslProtocol, params); return new UsageClient( + systemOperationContext, restClient, new ExponentialBackoff(retryInterval), numRetries, - systemAuthentication, configurationProvider.getCache().getClient().getUsageClient()); } } 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 b808c3da5d8d0c..4bf4313f4ad7e0 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 @@ -33,6 +33,7 @@ import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.SearchService; import com.linkedin.metadata.search.transformer.SearchDocumentTransformer; +import io.datahubproject.metadata.context.OperationContext; import java.util.ArrayList; import java.util.List; import javax.annotation.Nonnull; @@ -103,10 +104,12 @@ public class BootstrapManagerFactory { @Bean(name = "bootstrapManager") @Scope("singleton") @Nonnull - protected BootstrapManager createInstance() { + protected BootstrapManager createInstance( + @Qualifier("systemOperationContext") final OperationContext systemOpContext) { final IngestRootUserStep ingestRootUserStep = new IngestRootUserStep(_entityService); final IngestPoliciesStep ingestPoliciesStep = new IngestPoliciesStep( + systemOpContext, _entityRegistry, _entityService, _entitySearchService, @@ -118,7 +121,8 @@ protected BootstrapManager createInstance() { final IngestDataPlatformInstancesStep ingestDataPlatformInstancesStep = new IngestDataPlatformInstancesStep(_entityService, _migrationsDao); final RestoreGlossaryIndices restoreGlossaryIndicesStep = - new RestoreGlossaryIndices(_entityService, _entitySearchService, _entityRegistry); + new RestoreGlossaryIndices( + systemOpContext, _entityService, _entitySearchService, _entityRegistry); final IndexDataPlatformsStep indexDataPlatformsStep = new IndexDataPlatformsStep(_entityService, _entitySearchService, _entityRegistry); final RestoreDbtSiblingsIndices restoreDbtSiblingsIndices = @@ -161,7 +165,8 @@ protected BootstrapManager createInstance() { } if (_backfillBrowsePathsV2Enabled) { - finalSteps.add(new BackfillBrowsePathsV2Step(_entityService, _searchService)); + finalSteps.add( + new BackfillBrowsePathsV2Step(systemOpContext, _entityService, _searchService)); } return new BootstrapManager(finalSteps); diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java index 49a86406c1ecd5..30cfc304926c12 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java @@ -24,6 +24,7 @@ import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; +import io.datahubproject.metadata.context.OperationContext; import java.util.Set; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -46,11 +47,14 @@ public class BackfillBrowsePathsV2Step extends UpgradeStep { private static final String UPGRADE_ID = "backfill-default-browse-paths-v2-step"; private static final Integer BATCH_SIZE = 5000; - private final SearchService _searchService; + private final SearchService searchService; + private final OperationContext opContext; - public BackfillBrowsePathsV2Step(EntityService entityService, SearchService searchService) { + public BackfillBrowsePathsV2Step( + OperationContext opContext, EntityService entityService, SearchService searchService) { super(entityService, VERSION, UPGRADE_ID); - _searchService = searchService; + this.searchService = searchService; + this.opContext = opContext; } @Nonnull @@ -106,8 +110,8 @@ private String backfillBrowsePathsV2(String entityType, AuditStamp auditStamp, S filter.setOr(conjunctiveCriterionArray); final ScrollResult scrollResult = - _searchService.scrollAcrossEntities( - ImmutableList.of(entityType), "*", filter, null, scrollId, "5m", BATCH_SIZE, null); + searchService.scrollAcrossEntities( + opContext, ImmutableList.of(entityType), "*", filter, null, scrollId, "5m", BATCH_SIZE); if (scrollResult.getNumEntities() == 0 || scrollResult.getEntities().size() == 0) { return null; } 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 f925c96e333fd5..3a6b704613a569 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 @@ -27,6 +27,7 @@ import com.linkedin.mxe.GenericAspect; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.policy.DataHubPolicyInfo; +import io.datahubproject.metadata.context.OperationContext; import java.io.IOException; import java.net.URISyntaxException; import java.util.Collections; @@ -45,6 +46,7 @@ public class IngestPoliciesStep implements BootstrapStep { private static final String POLICY_ENTITY_NAME = "dataHubPolicy"; private static final String POLICY_INFO_ASPECT_NAME = "dataHubPolicyInfo"; + private final OperationContext systemOpContext; private final EntityRegistry _entityRegistry; private final EntityService _entityService; private final EntitySearchService _entitySearchService; @@ -113,7 +115,7 @@ public void execute() throws IOException, URISyntaxException { // If search index for policies is empty, update the policy index with the ingested policies // from previous step. // Directly update the ES index, does not produce MCLs - if (_entitySearchService.docCount(Constants.POLICY_ENTITY_NAME) == 0) { + if (_entitySearchService.docCount(systemOpContext, Constants.POLICY_ENTITY_NAME) == 0) { updatePolicyIndex(); } log.info("Successfully ingested default access policies."); diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java index 5c2b2c28e6dcf3..e0b912f483520b 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java @@ -12,10 +12,10 @@ import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchResult; +import io.datahubproject.metadata.context.OperationContext; import java.util.Collections; import java.util.HashSet; import java.util.LinkedList; @@ -34,26 +34,29 @@ public class RestoreGlossaryIndices extends UpgradeStep { private static final String UPGRADE_ID = "restore-glossary-indices-ui"; private static final Integer BATCH_SIZE = 1000; - private final EntitySearchService _entitySearchService; - private final EntityRegistry _entityRegistry; + private final OperationContext opContext; + private final EntitySearchService entitySearchService; + private final EntityRegistry entityRegistry; public RestoreGlossaryIndices( + OperationContext opContext, EntityService entityService, EntitySearchService entitySearchService, EntityRegistry entityRegistry) { super(entityService, VERSION, UPGRADE_ID); - _entitySearchService = entitySearchService; - _entityRegistry = entityRegistry; + this.opContext = opContext; + this.entitySearchService = entitySearchService; + this.entityRegistry = entityRegistry; } @Override public void upgrade() throws Exception { final AspectSpec termAspectSpec = - _entityRegistry + entityRegistry .getEntitySpec(Constants.GLOSSARY_TERM_ENTITY_NAME) .getAspectSpec(Constants.GLOSSARY_TERM_INFO_ASPECT_NAME); final AspectSpec nodeAspectSpec = - _entityRegistry + entityRegistry .getEntitySpec(Constants.GLOSSARY_NODE_ENTITY_NAME) .getAspectSpec(Constants.GLOSSARY_NODE_INFO_ASPECT_NAME); final AuditStamp auditStamp = @@ -85,14 +88,16 @@ public ExecutionMode getExecutionMode() { private int getAndRestoreTermAspectIndices( int start, AuditStamp auditStamp, AspectSpec termAspectSpec) throws Exception { SearchResult termsResult = - _entitySearchService.search( + entitySearchService.search( + opContext.withSearchFlags( + flags -> + flags.setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true)), List.of(Constants.GLOSSARY_TERM_ENTITY_NAME), "", null, null, start, - BATCH_SIZE, - new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true)); + BATCH_SIZE); List termUrns = termsResult.getEntities().stream() .map(SearchEntity::getEntity) @@ -153,14 +158,16 @@ private int getAndRestoreTermAspectIndices( private int getAndRestoreNodeAspectIndices( int start, AuditStamp auditStamp, AspectSpec nodeAspectSpec) throws Exception { SearchResult nodesResult = - _entitySearchService.search( + entitySearchService.search( + opContext.withSearchFlags( + flags -> + flags.setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true)), List.of(Constants.GLOSSARY_NODE_ENTITY_NAME), "", null, null, start, - BATCH_SIZE, - new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true)); + BATCH_SIZE); List nodeUrns = nodesResult.getEntities().stream() .map(SearchEntity::getEntity) diff --git a/metadata-service/factories/src/test/java/com/linkedin/gms/factory/search/CacheTest.java b/metadata-service/factories/src/test/java/com/linkedin/gms/factory/search/CacheTest.java index 6cc1d293e24e67..d28bb291ad2363 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/gms/factory/search/CacheTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/gms/factory/search/CacheTest.java @@ -2,6 +2,7 @@ import static com.datahub.util.RecordUtils.*; import static com.linkedin.metadata.search.client.CachingEntitySearchService.*; +import static org.mockito.Mockito.mock; import com.google.common.collect.ImmutableList; import com.hazelcast.config.Config; @@ -14,6 +15,7 @@ import com.linkedin.metadata.graph.LineageDirection; import com.linkedin.metadata.graph.LineageRelationship; import com.linkedin.metadata.graph.LineageRelationshipArray; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; @@ -29,6 +31,8 @@ import com.linkedin.metadata.search.SearchResultMetadata; import com.linkedin.metadata.search.cache.CacheableSearcher; import com.linkedin.metadata.search.cache.CachedEntityLineageResult; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.time.temporal.ChronoUnit; import java.util.List; import org.javatuples.Quintet; @@ -56,6 +60,8 @@ public CacheTest() { @Test public void hazelcastTest() { + OperationContext systemOpContext = + TestOperationContexts.systemContextNoSearchAuthorization(mock(EntityRegistry.class)); CorpuserUrn corpuserUrn = new CorpuserUrn("user"); SearchEntity searchEntity = new SearchEntity().setEntity(corpuserUrn); SearchResult searchResult = @@ -83,7 +89,6 @@ public void hazelcastTest() { 10, querySize -> searchResult, querySize -> quintet, - null, true); CacheableSearcher< @@ -94,18 +99,18 @@ public void hazelcastTest() { 10, querySize -> searchResult, querySize -> quintet, - null, true); // Cache result - SearchResult result = cacheableSearcher1.getSearchResults(0, 1); + SearchResult result = cacheableSearcher1.getSearchResults(systemOpContext, 0, 1); Assert.assertNotEquals(result, null); Assert.assertEquals( instance1.getMap("test").get(quintet), instance2.getMap("test").get(quintet)); - Assert.assertEquals(cacheableSearcher1.getSearchResults(0, 1), searchResult); + Assert.assertEquals(cacheableSearcher1.getSearchResults(systemOpContext, 0, 1), searchResult); Assert.assertEquals( - cacheableSearcher1.getSearchResults(0, 1), cacheableSearcher2.getSearchResults(0, 1)); + cacheableSearcher1.getSearchResults(systemOpContext, 0, 1), + cacheableSearcher2.getSearchResults(systemOpContext, 0, 1)); } @Test @@ -181,7 +186,7 @@ public void testLineageCaching() { EntityLineageResultCacheKey key = new EntityLineageResultCacheKey( - corpuserUrn, LineageDirection.DOWNSTREAM, 0L, 1L, 1, ChronoUnit.DAYS); + "", corpuserUrn, LineageDirection.DOWNSTREAM, 0L, 1L, 1, ChronoUnit.DAYS); cache1.put(key, cachedEntityLineageResult); diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java index 0858736e39021a..86707ac0c6db55 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java @@ -2,6 +2,7 @@ import static com.linkedin.metadata.Constants.CONTAINER_ASPECT_NAME; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; import com.google.common.collect.ImmutableList; import com.linkedin.common.AuditStamp; @@ -20,6 +21,7 @@ import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchService; import com.linkedin.mxe.MetadataChangeProposal; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.Collections; import java.util.HashMap; @@ -86,19 +88,19 @@ public void testExecuteNoExistingBrowsePaths() throws Exception { .thenReturn(null); BackfillBrowsePathsV2Step backfillBrowsePathsV2Step = - new BackfillBrowsePathsV2Step(mockService, mockSearchService); + new BackfillBrowsePathsV2Step(mock(OperationContext.class), mockService, mockSearchService); backfillBrowsePathsV2Step.execute(); Mockito.verify(mockSearchService, Mockito.times(9)) .scrollAcrossEntities( + any(OperationContext.class), any(), Mockito.eq("*"), any(Filter.class), Mockito.eq(null), Mockito.eq(null), Mockito.eq("5m"), - Mockito.eq(5000), - Mockito.eq(null)); + Mockito.eq(5000)); // Verify that 11 aspects are ingested, 2 for the upgrade request / result, 9 for ingesting 1 of // each entity type Mockito.verify(mockService, Mockito.times(11)) @@ -107,7 +109,7 @@ public void testExecuteNoExistingBrowsePaths() throws Exception { @Test public void testDoesNotRunWhenAlreadyExecuted() throws Exception { - final EntityService mockService = Mockito.mock(EntityService.class); + final EntityService mockService = mock(EntityService.class); final SearchService mockSearchService = initMockSearchService(); final Urn upgradeEntityUrn = Urn.createFromString(UPGRADE_URN); @@ -127,7 +129,7 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { .thenReturn(response); BackfillBrowsePathsV2Step backfillBrowsePathsV2Step = - new BackfillBrowsePathsV2Step(mockService, mockSearchService); + new BackfillBrowsePathsV2Step(mock(OperationContext.class), mockService, mockSearchService); backfillBrowsePathsV2Step.execute(); Mockito.verify(mockService, Mockito.times(0)) @@ -136,7 +138,7 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { } private EntityService initMockService() throws URISyntaxException { - final EntityService mockService = Mockito.mock(EntityService.class); + final EntityService mockService = mock(EntityService.class); final EntityRegistry registry = new UpgradeDefaultBrowsePathsStepTest.TestEntityRegistry(); Mockito.when(mockService.getEntityRegistry()).thenReturn(registry); @@ -153,19 +155,19 @@ private EntityService initMockService() throws URISyntaxException { } private SearchService initMockSearchService() { - final SearchService mockSearchService = Mockito.mock(SearchService.class); + final SearchService mockSearchService = mock(SearchService.class); for (int i = 0; i < ENTITY_TYPES.size(); i++) { Mockito.when( mockSearchService.scrollAcrossEntities( + Mockito.any(OperationContext.class), Mockito.eq(ImmutableList.of(ENTITY_TYPES.get(i))), Mockito.eq("*"), any(Filter.class), Mockito.eq(null), Mockito.eq(null), Mockito.eq("5m"), - Mockito.eq(5000), - Mockito.eq(null))) + Mockito.eq(5000))) .thenReturn( new ScrollResult() .setNumEntities(1) diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java index 4a4532763f02bc..426e52d3636f7c 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java @@ -15,13 +15,13 @@ import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; import com.linkedin.metadata.search.SearchResult; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.util.Pair; +import io.datahubproject.metadata.context.OperationContext; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -54,16 +54,13 @@ private void mockGetTermInfo( .setAspects(new EnvelopedAspectMap(termInfoAspects))); Mockito.when( mockSearchService.search( - List.of(Constants.GLOSSARY_TERM_ENTITY_NAME), - "", - null, - null, - 0, - 1000, - new SearchFlags() - .setFulltext(false) - .setSkipAggregates(true) - .setSkipHighlighting(true))) + Mockito.any(), + Mockito.eq(List.of(Constants.GLOSSARY_TERM_ENTITY_NAME)), + Mockito.eq(""), + Mockito.any(), + Mockito.any(), + Mockito.eq(0), + Mockito.eq(1000))) .thenReturn( new SearchResult() .setNumEntities(1) @@ -93,16 +90,13 @@ private void mockGetNodeInfo( .setAspects(new EnvelopedAspectMap(nodeInfoAspects))); Mockito.when( mockSearchService.search( - List.of(Constants.GLOSSARY_NODE_ENTITY_NAME), - "", - null, - null, - 0, - 1000, - new SearchFlags() - .setFulltext(false) - .setSkipAggregates(true) - .setSkipHighlighting(true))) + Mockito.any(), + Mockito.eq(List.of(Constants.GLOSSARY_NODE_ENTITY_NAME)), + Mockito.eq(""), + Mockito.any(), + Mockito.any(), + Mockito.eq(0), + Mockito.eq(1000))) .thenReturn( new SearchResult() .setNumEntities(1) @@ -171,7 +165,8 @@ public void testExecuteFirstTime() throws Exception { AspectSpec aspectSpec = mockGlossaryAspectSpecs(mockRegistry); RestoreGlossaryIndices restoreIndicesStep = - new RestoreGlossaryIndices(mockService, mockSearchService, mockRegistry); + new RestoreGlossaryIndices( + Mockito.mock(OperationContext.class), mockService, mockSearchService, mockRegistry); restoreIndicesStep.execute(); Mockito.verify(mockRegistry, Mockito.times(1)) @@ -254,7 +249,8 @@ public void testExecutesWithNewVersion() throws Exception { AspectSpec aspectSpec = mockGlossaryAspectSpecs(mockRegistry); RestoreGlossaryIndices restoreIndicesStep = - new RestoreGlossaryIndices(mockService, mockSearchService, mockRegistry); + new RestoreGlossaryIndices( + Mockito.mock(OperationContext.class), mockService, mockSearchService, mockRegistry); restoreIndicesStep.execute(); Mockito.verify(mockRegistry, Mockito.times(1)) @@ -319,7 +315,8 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { .thenReturn(response); RestoreGlossaryIndices restoreIndicesStep = - new RestoreGlossaryIndices(mockService, mockSearchService, mockRegistry); + new RestoreGlossaryIndices( + Mockito.mock(OperationContext.class), mockService, mockSearchService, mockRegistry); restoreIndicesStep.execute(); Mockito.verify(mockRegistry, Mockito.times(0)) @@ -328,22 +325,22 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { .getEntitySpec(Constants.GLOSSARY_NODE_ENTITY_NAME); Mockito.verify(mockSearchService, Mockito.times(0)) .search( - List.of(Constants.GLOSSARY_TERM_ENTITY_NAME), - "", - null, - null, - 0, - 1000, - new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true)); + Mockito.any(), + Mockito.eq(List.of(Constants.GLOSSARY_TERM_ENTITY_NAME)), + Mockito.eq(""), + Mockito.any(), + Mockito.any(), + Mockito.eq(0), + Mockito.eq(1000)); Mockito.verify(mockSearchService, Mockito.times(0)) .search( - List.of(Constants.GLOSSARY_NODE_ENTITY_NAME), - "", - null, - null, - 0, - 1000, - new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true)); + Mockito.any(), + Mockito.eq(List.of(Constants.GLOSSARY_NODE_ENTITY_NAME)), + Mockito.eq(""), + Mockito.any(), + Mockito.any(), + Mockito.eq(0), + Mockito.eq(1000)); Mockito.verify(mockService, Mockito.times(0)) .ingestProposal( Mockito.any(MetadataChangeProposal.class), diff --git a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java index da4fac1451e469..1c4f2824c6357e 100644 --- a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java +++ b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java @@ -11,10 +11,12 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.inject.name.Named; import com.linkedin.datahub.graphql.GraphQLEngine; import com.linkedin.datahub.graphql.exception.DataHubGraphQLError; import com.linkedin.metadata.utils.metrics.MetricUtils; import graphql.ExecutionResult; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.api.trace.Span; import jakarta.inject.Inject; import jakarta.servlet.http.HttpServletRequest; @@ -25,6 +27,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; +import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; import org.springframework.http.HttpEntity; import org.springframework.http.HttpStatus; @@ -47,6 +50,11 @@ public GraphQLController() { @Inject AuthorizerChain _authorizerChain; + @Nonnull + @Inject + @Named("systemOperationContext") + private OperationContext systemOperationContext; + @PostMapping(value = "/graphql", produces = "application/json;charset=utf-8") CompletableFuture> postGraphQL(HttpEntity httpEntity) { @@ -95,7 +103,8 @@ CompletableFuture> postGraphQL(HttpEntity httpEnt * Init QueryContext */ Authentication authentication = AuthenticationContext.getAuthentication(); - SpringQueryContext context = new SpringQueryContext(true, authentication, _authorizerChain); + SpringQueryContext context = + new SpringQueryContext(true, authentication, _authorizerChain, systemOperationContext); Span.current().setAttribute("actor.urn", context.getActorUrn()); return CompletableFuture.supplyAsync( diff --git a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/SpringQueryContext.java b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/SpringQueryContext.java index 379521eda0c1a3..20e06945e1d6ba 100644 --- a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/SpringQueryContext.java +++ b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/SpringQueryContext.java @@ -3,34 +3,27 @@ import com.datahub.authentication.Authentication; import com.datahub.plugins.auth.authorization.Authorizer; import com.linkedin.datahub.graphql.QueryContext; +import io.datahubproject.metadata.context.OperationContext; +import javax.annotation.Nonnull; +import lombok.Getter; +@Getter public class SpringQueryContext implements QueryContext { private final boolean isAuthenticated; private final Authentication authentication; private final Authorizer authorizer; + @Nonnull private final OperationContext operationContext; public SpringQueryContext( final boolean isAuthenticated, final Authentication authentication, - final Authorizer authorizer) { + final Authorizer authorizer, + @Nonnull final OperationContext systemOperationContext) { this.isAuthenticated = isAuthenticated; this.authentication = authentication; this.authorizer = authorizer; - } - - @Override - public boolean isAuthenticated() { - return this.isAuthenticated; - } - - @Override - public Authentication getAuthentication() { - return this.authentication; - } - - @Override - public Authorizer getAuthorizer() { - return this.authorizer; + this.operationContext = + OperationContext.asSession(systemOperationContext, authorizer, authentication, true); } } diff --git a/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImpl.java b/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImpl.java index 1e375f90fc38a9..20fdb0db0bd090 100644 --- a/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImpl.java +++ b/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImpl.java @@ -19,6 +19,7 @@ import com.linkedin.metadata.search.ScrollResult; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchService; +import io.datahubproject.metadata.context.OperationContext; import io.datahubproject.openapi.dto.UpsertAspectRequest; import io.datahubproject.openapi.dto.UrnResponseMap; import io.datahubproject.openapi.entities.EntitiesController; @@ -71,6 +72,7 @@ import org.springframework.http.ResponseEntity; public class EntityApiDelegateImpl { + private final OperationContext systemOperationContext; private final EntityRegistry _entityRegistry; private final EntityService _entityService; private final SearchService _searchService; @@ -85,6 +87,7 @@ public class EntityApiDelegateImpl { private final StackWalker walker = StackWalker.getInstance(); public EntityApiDelegateImpl( + OperationContext systemOperationContext, EntityService entityService, SearchService searchService, EntitiesController entitiesController, @@ -93,6 +96,7 @@ public EntityApiDelegateImpl( Class reqClazz, Class respClazz, Class scrollRespClazz) { + this.systemOperationContext = systemOperationContext; this._entityService = entityService; this._searchService = searchService; this._entityRegistry = entityService.getEntityRegistry(); @@ -458,7 +462,13 @@ public ResponseEntity scroll( @Valid SortOrder sortOrder, @Valid String query) { + SearchFlags searchFlags = + new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true); + Authentication authentication = AuthenticationContext.getAuthentication(); + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, _authorizationChain, authentication, true); com.linkedin.metadata.models.EntitySpec entitySpec = OpenApiEntitiesUtil.responseClassToEntitySpec(_entityRegistry, _respClazz); checkScrollAuthorized(authentication, entitySpec); @@ -470,19 +480,16 @@ public ResponseEntity scroll( com.linkedin.metadata.query.filter.SortOrder.valueOf( Optional.ofNullable(sortOrder).map(Enum::name).orElse("ASCENDING"))); - SearchFlags searchFlags = - new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true); - ScrollResult result = _searchService.scrollAcrossEntities( + opContext.withSearchFlags(flags -> searchFlags), List.of(entitySpec.getName()), query, null, sortCriterion, scrollId, null, - count, - searchFlags); + count); String[] urns = result.getEntities().stream() diff --git a/metadata-service/openapi-entity-servlet/src/main/resources/JavaSpring/apiController.mustache b/metadata-service/openapi-entity-servlet/src/main/resources/JavaSpring/apiController.mustache index 7ac087f220561f..f9717b8cb16fb1 100644 --- a/metadata-service/openapi-entity-servlet/src/main/resources/JavaSpring/apiController.mustache +++ b/metadata-service/openapi-entity-servlet/src/main/resources/JavaSpring/apiController.mustache @@ -1,6 +1,7 @@ package {{package}}; import io.datahubproject.openapi.v2.delegates.EntityApiDelegateImpl; +import io.datahubproject.metadata.context.OperationContext; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.search.SearchService; import io.datahubproject.openapi.entities.EntitiesController; @@ -91,12 +92,14 @@ public class {{classname}}Controller implements {{classname}} { private final EntityApiDelegateImpl<{{requestClass}}, {{responseClass}}, {{scrollResponseClass}}> delegate; @org.springframework.beans.factory.annotation.Autowired - public {{classname}}Controller(ObjectMapper objectMapper, HttpServletRequest request, EntityService entityService, + public {{classname}}Controller(ObjectMapper objectMapper, HttpServletRequest request, + @org.springframework.beans.factory.annotation.Qualifier("systemOperationContext") OperationContext systemOperationContext, + EntityService entityService, SearchService searchService, EntitiesController v1Controller, AuthorizerChain authorizationChain, @Value("${authorization.restApiAuthorization:false}") boolean restApiAuthorizationEnabled) { this.objectMapper = objectMapper; this.request = request; - this.delegate = new EntityApiDelegateImpl<{{requestClass}}, {{responseClass}}, {{scrollResponseClass}}>(entityService, searchService, v1Controller, + this.delegate = new EntityApiDelegateImpl<{{requestClass}}, {{responseClass}}, {{scrollResponseClass}}>(systemOperationContext, entityService, searchService, v1Controller, restApiAuthorizationEnabled, authorizationChain, {{requestClass}}.class, {{responseClass}}.class, {{scrollResponseClass}}.class); } {{#isJava8or11}} diff --git a/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/config/OpenAPIEntityTestConfiguration.java b/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/config/OpenAPIEntityTestConfiguration.java index 920a13d998985d..4ff40b4beed5a7 100644 --- a/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/config/OpenAPIEntityTestConfiguration.java +++ b/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/config/OpenAPIEntityTestConfiguration.java @@ -27,11 +27,13 @@ import com.linkedin.metadata.search.SearchService; import com.linkedin.metadata.systemmetadata.SystemMetadataService; import com.linkedin.metadata.timeline.TimelineService; +import io.datahubproject.metadata.context.OperationContext; import io.datahubproject.openapi.dto.UrnResponseMap; import io.datahubproject.openapi.entities.EntitiesController; import io.datahubproject.openapi.generated.EntityResponse; import io.datahubproject.openapi.relationships.RelationshipsController; import io.datahubproject.openapi.timeline.TimelineController; +import io.datahubproject.test.metadata.context.TestOperationContexts; import java.util.Arrays; import java.util.Map; import java.util.stream.Collectors; @@ -62,7 +64,7 @@ public EntityService entityService(final EntityRegistry mockRegistry) { public SearchService searchService() { SearchService searchService = mock(SearchService.class); when(searchService.scrollAcrossEntities( - anyList(), any(), any(), any(), any(), any(), anyInt(), any())) + any(OperationContext.class), anyList(), any(), any(), any(), any(), any(), anyInt())) .thenReturn(new ScrollResult().setEntities(new SearchEntityArray())); return searchService; @@ -134,4 +136,9 @@ public EntitiesController entitiesController() { @MockBean public TimelineController timelineController; @MockBean public RelationshipsController relationshipsController; + + @Bean(name = "systemOperationContext") + public OperationContext operationContext(final EntityRegistry entityRegistry) { + return TestOperationContexts.systemContextNoSearchAuthorization(entityRegistry); + } } diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/elastic/OperationsController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/elastic/OperationsController.java index 1d6d3067d23f73..f3617801e6b552 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/elastic/OperationsController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/elastic/OperationsController.java @@ -3,9 +3,9 @@ import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; import com.datahub.authorization.AuthUtil; -import com.datahub.authorization.AuthorizerChain; import com.datahub.authorization.ConjunctivePrivilegeGroup; import com.datahub.authorization.DisjunctivePrivilegeGroup; +import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.query.SearchFlags; @@ -15,6 +15,7 @@ import com.linkedin.metadata.systemmetadata.SystemMetadataService; import com.linkedin.metadata.timeseries.TimeseriesAspectService; import com.linkedin.timeseries.TimeseriesIndexSizeResult; +import io.datahubproject.metadata.context.OperationContext; import io.datahubproject.openapi.util.ElasticsearchUtils; import io.swagger.v3.oas.annotations.Operation; import io.swagger.v3.oas.annotations.Parameter; @@ -46,7 +47,8 @@ name = "ElasticSearchOperations", description = "An API for managing your elasticsearch instance") public class OperationsController { - private final AuthorizerChain authorizerChain; + private final Authorizer authorizerChain; + private final OperationContext systemOperationContext; @Value("${authorization.restApiAuthorization:false}") private boolean restApiAuthorizationEnabled; @@ -57,11 +59,12 @@ public class OperationsController { private final EntitySearchService searchService; public OperationsController( - AuthorizerChain authorizerChain, + OperationContext systemOperationContext, SystemMetadataService systemMetadataService, TimeseriesAspectService timeseriesAspectService, EntitySearchService searchService) { - this.authorizerChain = authorizerChain; + this.systemOperationContext = systemOperationContext; + this.authorizerChain = systemOperationContext.getAuthorizerContext().getAuthorizer(); this.systemMetadataService = systemMetadataService; this.timeseriesAspectService = timeseriesAspectService; this.searchService = searchService; @@ -231,14 +234,19 @@ public ResponseEntity explainSearchQuery( log.error("{} is not authorized to get timeseries index sizes", actorUrnStr); return ResponseEntity.status(HttpStatus.FORBIDDEN).body(null); } + OperationContext opContext = + systemOperationContext + .asSession(authorizerChain, authentication) + .withSearchFlags(flags -> searchFlags); + ExplainResponse response = searchService.explain( + opContext, query, documentId, entityName, filters, sortCriterion, - searchFlags, scrollId, keepAlive, size, diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java index 656d6542483cf3..7a11d60a567f95 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java @@ -37,6 +37,7 @@ import com.linkedin.metadata.utils.SearchUtil; import com.linkedin.mxe.SystemMetadata; import com.linkedin.util.Pair; +import io.datahubproject.metadata.context.OperationContext; import io.datahubproject.openapi.v2.models.GenericEntity; import io.datahubproject.openapi.v2.models.GenericScrollResult; import io.swagger.v3.oas.annotations.Operation; @@ -55,6 +56,7 @@ import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.http.MediaType; import org.springframework.http.ResponseEntity; import org.springframework.web.bind.annotation.DeleteMapping; @@ -82,6 +84,10 @@ public class EntityController { @Autowired private boolean restApiAuthorizationEnabled; @Autowired private ObjectMapper objectMapper; + @Qualifier("systemOperationContext") + @Autowired + private OperationContext systemOperationContext; + @Tag(name = "Generic Entities", description = "API for interacting with generic entities.") @GetMapping(value = "/{entityName}", produces = MediaType.APPLICATION_JSON_VALUE) @Operation(summary = "Scroll entities") @@ -100,28 +106,31 @@ public ResponseEntity> getEntities( EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + Authentication authentication = AuthenticationContext.getAuthentication(); if (restApiAuthorizationEnabled) { - Authentication authentication = AuthenticationContext.getAuthentication(); checkAuthorized( authorizationChain, authentication.getActor(), entitySpec, ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE.getType())); } + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, authorizationChain, authentication, true); // TODO: support additional and multiple sort params SortCriterion sortCriterion = SearchUtil.sortBy(sortField, SortOrder.valueOf(sortOrder)); ScrollResult result = searchService.scrollAcrossEntities( + opContext.withSearchFlags(flags -> DEFAULT_SEARCH_FLAGS), List.of(entitySpec.getName()), query, null, sortCriterion, scrollId, null, - count, - DEFAULT_SEARCH_FLAGS); + count); return ResponseEntity.ok( GenericScrollResult.builder() diff --git a/metadata-service/restli-api/src/main/idl/com.linkedin.entity.entities.restspec.json b/metadata-service/restli-api/src/main/idl/com.linkedin.entity.entities.restspec.json index eac1cc690a60d1..c5fc14c45c522d 100644 --- a/metadata-service/restli-api/src/main/idl/com.linkedin.entity.entities.restspec.json +++ b/metadata-service/restli-api/src/main/idl/com.linkedin.entity.entities.restspec.json @@ -74,6 +74,10 @@ }, { "name" : "limit", "type" : "int" + }, { + "name" : "searchFlags", + "type" : "com.linkedin.metadata.query.SearchFlags", + "optional" : true } ], "returns" : "com.linkedin.metadata.query.AutoCompleteResult" }, { @@ -114,6 +118,10 @@ }, { "name" : "limit", "type" : "int" + }, { + "name" : "searchFlags", + "type" : "com.linkedin.metadata.query.SearchFlags", + "optional" : true } ], "returns" : "com.linkedin.metadata.browse.BrowseResult" }, { diff --git a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.aspects.snapshot.json b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.aspects.snapshot.json index fe16d24e3475ac..4375dfa587e7d0 100644 --- a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.aspects.snapshot.json +++ b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.aspects.snapshot.json @@ -1222,6 +1222,24 @@ "items" : "Owner" }, "doc" : "List of owners of the entity." + }, { + "name" : "ownerTypes", + "type" : { + "type" : "map", + "values" : { + "type" : "array", + "items" : "Urn" + } + }, + "doc" : "Owners to ownership type map, populated with mutation hook.", + "default" : { }, + "optional" : true, + "Searchable" : { + "/*" : { + "fieldType" : "OBJECT", + "queryByDefault" : false + } + } }, { "name" : "lastModified", "type" : "AuditStamp", diff --git a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json index 55fed125936eb0..f6ccf6a9c63260 100644 --- a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json +++ b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json @@ -1258,6 +1258,24 @@ "items" : "Owner" }, "doc" : "List of owners of the entity." + }, { + "name" : "ownerTypes", + "type" : { + "type" : "map", + "values" : { + "type" : "array", + "items" : "Urn" + } + }, + "doc" : "Owners to ownership type map, populated with mutation hook.", + "default" : { }, + "optional" : true, + "Searchable" : { + "/*" : { + "fieldType" : "OBJECT", + "queryByDefault" : false + } + } }, { "name" : "lastModified", "type" : "AuditStamp", @@ -5785,6 +5803,18 @@ "type" : "GroupingSpec", "doc" : "Instructions for grouping results before returning", "optional" : true + }, { + "name" : "includeSoftDeleted", + "type" : "boolean", + "doc" : "include soft deleted entities in results", + "default" : false, + "optional" : true + }, { + "name" : "includeRestricted", + "type" : "boolean", + "doc" : "include restricted entities in results (default is to filter)", + "default" : false, + "optional" : true } ] }, { "type" : "enum", @@ -6111,6 +6141,14 @@ "name" : "score", "type" : "double", "optional" : true + }, { + "name" : "restrictedAspects", + "type" : { + "type" : "array", + "items" : "string" + }, + "doc" : "A list of the the restricted aspects on the entity.\nIf the key aspect is present, assume ALL aspects should be restricted including the entity's Urn.", + "optional" : true } ] } ], "fields" : [ { @@ -6420,6 +6458,10 @@ }, { "name" : "limit", "type" : "int" + }, { + "name" : "searchFlags", + "type" : "com.linkedin.metadata.query.SearchFlags", + "optional" : true } ], "returns" : "com.linkedin.metadata.query.AutoCompleteResult" }, { @@ -6460,6 +6502,10 @@ }, { "name" : "limit", "type" : "int" + }, { + "name" : "searchFlags", + "type" : "com.linkedin.metadata.query.SearchFlags", + "optional" : true } ], "returns" : "com.linkedin.metadata.browse.BrowseResult" }, { diff --git a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.runs.snapshot.json b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.runs.snapshot.json index f9f1999923ec0b..2168ee950957a0 100644 --- a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.runs.snapshot.json +++ b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.runs.snapshot.json @@ -964,6 +964,24 @@ "items" : "Owner" }, "doc" : "List of owners of the entity." + }, { + "name" : "ownerTypes", + "type" : { + "type" : "map", + "values" : { + "type" : "array", + "items" : "Urn" + } + }, + "doc" : "Owners to ownership type map, populated with mutation hook.", + "default" : { }, + "optional" : true, + "Searchable" : { + "/*" : { + "fieldType" : "OBJECT", + "queryByDefault" : false + } + } }, { "name" : "lastModified", "type" : "AuditStamp", diff --git a/metadata-service/restli-api/src/main/snapshot/com.linkedin.operations.operations.snapshot.json b/metadata-service/restli-api/src/main/snapshot/com.linkedin.operations.operations.snapshot.json index 88dad7e49152aa..2c093f753e3a6b 100644 --- a/metadata-service/restli-api/src/main/snapshot/com.linkedin.operations.operations.snapshot.json +++ b/metadata-service/restli-api/src/main/snapshot/com.linkedin.operations.operations.snapshot.json @@ -964,6 +964,24 @@ "items" : "Owner" }, "doc" : "List of owners of the entity." + }, { + "name" : "ownerTypes", + "type" : { + "type" : "map", + "values" : { + "type" : "array", + "items" : "Urn" + } + }, + "doc" : "Owners to ownership type map, populated with mutation hook.", + "default" : { }, + "optional" : true, + "Searchable" : { + "/*" : { + "fieldType" : "OBJECT", + "queryByDefault" : false + } + } }, { "name" : "lastModified", "type" : "AuditStamp", diff --git a/metadata-service/restli-api/src/main/snapshot/com.linkedin.platform.platform.snapshot.json b/metadata-service/restli-api/src/main/snapshot/com.linkedin.platform.platform.snapshot.json index 4d34126cd59fcd..13172d3e4d5eca 100644 --- a/metadata-service/restli-api/src/main/snapshot/com.linkedin.platform.platform.snapshot.json +++ b/metadata-service/restli-api/src/main/snapshot/com.linkedin.platform.platform.snapshot.json @@ -1258,6 +1258,24 @@ "items" : "Owner" }, "doc" : "List of owners of the entity." + }, { + "name" : "ownerTypes", + "type" : { + "type" : "map", + "values" : { + "type" : "array", + "items" : "Urn" + } + }, + "doc" : "Owners to ownership type map, populated with mutation hook.", + "default" : { }, + "optional" : true, + "Searchable" : { + "/*" : { + "fieldType" : "OBJECT", + "queryByDefault" : false + } + } }, { "name" : "lastModified", "type" : "AuditStamp", diff --git a/metadata-service/restli-client/build.gradle b/metadata-service/restli-client/build.gradle index 86336755dc0954..9bee54da9ff6e6 100644 --- a/metadata-service/restli-client/build.gradle +++ b/metadata-service/restli-client/build.gradle @@ -9,6 +9,7 @@ dependencies { api project(path: ':metadata-service:restli-api', configuration: 'restClient') api project(':metadata-events:mxe-schemas') api project(':metadata-utils') + api project(':metadata-operation-context') implementation project(':metadata-service:configuration') implementation externalDependency.caffeine diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java index 65169344776b7b..07ecdf2408d9d1 100644 --- a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java +++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java @@ -21,7 +21,6 @@ import com.linkedin.metadata.query.AutoCompleteResult; import com.linkedin.metadata.query.ListResult; import com.linkedin.metadata.query.ListUrnsResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.search.LineageScrollResult; @@ -32,6 +31,7 @@ import com.linkedin.mxe.PlatformEvent; import com.linkedin.mxe.SystemMetadata; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.Collection; import java.util.List; @@ -49,7 +49,7 @@ public interface EntityClient { default void postConstruct(AspectRetriever aspectRetriever) {} @Nullable - public EntityResponse getV2( + EntityResponse getV2( @Nonnull String entityName, @Nonnull final Urn urn, @Nullable final Set aspectNames, @@ -58,11 +58,11 @@ public EntityResponse getV2( @Nonnull @Deprecated - public Entity get(@Nonnull final Urn urn, @Nonnull final Authentication authentication) + Entity get(@Nonnull final Urn urn, @Nonnull final Authentication authentication) throws RemoteInvocationException; @Nonnull - public Map batchGetV2( + Map batchGetV2( @Nonnull String entityName, @Nonnull final Set urns, @Nullable final Set aspectNames, @@ -79,7 +79,7 @@ Map batchGetVersionedV2( @Nonnull @Deprecated - public Map batchGet( + Map batchGet( @Nonnull final Set urns, @Nonnull final Authentication authentication) throws RemoteInvocationException; @@ -93,13 +93,13 @@ public Map batchGet( * @throws RemoteInvocationException */ @Nonnull - public AutoCompleteResult autoComplete( + AutoCompleteResult autoComplete( + @Nonnull OperationContext opContext, @Nonnull String entityType, @Nonnull String query, @Nullable Filter requestFilters, @Nonnull int limit, - @Nullable String field, - @Nonnull Authentication authentication) + @Nullable String field) throws RemoteInvocationException; /** @@ -111,12 +111,12 @@ public AutoCompleteResult autoComplete( * @throws RemoteInvocationException */ @Nonnull - public AutoCompleteResult autoComplete( + AutoCompleteResult autoComplete( + @Nonnull OperationContext opContext, @Nonnull String entityType, @Nonnull String query, @Nullable Filter requestFilters, - @Nonnull int limit, - @Nonnull Authentication authentication) + @Nonnull int limit) throws RemoteInvocationException; /** @@ -130,13 +130,13 @@ public AutoCompleteResult autoComplete( * @throws RemoteInvocationException */ @Nonnull - public BrowseResult browse( + BrowseResult browse( + @Nonnull OperationContext opContext, @Nonnull String entityType, @Nonnull String path, @Nullable Map requestFilters, int start, - int limit, - @Nonnull Authentication authentication) + int limit) throws RemoteInvocationException; /** @@ -151,15 +151,14 @@ public BrowseResult browse( * @throws RemoteInvocationException */ @Nonnull - public BrowseResultV2 browseV2( + BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nonnull Authentication authentication, - @Nullable SearchFlags searchFlags) + int count) throws RemoteInvocationException; /** @@ -174,30 +173,29 @@ public BrowseResultV2 browseV2( * @throws RemoteInvocationException */ @Nonnull - public BrowseResultV2 browseV2( + BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nonnull Authentication authentication, - @Nullable SearchFlags searchFlags) + int count) throws RemoteInvocationException; @Deprecated - public void update(@Nonnull final Entity entity, @Nonnull final Authentication authentication) + void update(@Nonnull final Entity entity, @Nonnull final Authentication authentication) throws RemoteInvocationException; @Deprecated - public void updateWithSystemMetadata( + void updateWithSystemMetadata( @Nonnull final Entity entity, @Nullable final SystemMetadata systemMetadata, @Nonnull final Authentication authentication) throws RemoteInvocationException; @Deprecated - public void batchUpdate( + void batchUpdate( @Nonnull final Set entities, @Nonnull final Authentication authentication) throws RemoteInvocationException; @@ -208,19 +206,17 @@ public void batchUpdate( * @param requestFilters search filters * @param start start offset for search results * @param count max number of search results requested - * @param searchFlags configuration flags for the search request * @return a set of search results * @throws RemoteInvocationException */ @Nonnull - public SearchResult search( + SearchResult search( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nonnull String input, @Nullable Map requestFilters, int start, - int count, - @Nonnull Authentication authentication, - @Nullable SearchFlags searchFlags) + int count) throws RemoteInvocationException; /** @@ -235,12 +231,12 @@ public SearchResult search( * @throws RemoteInvocationException */ @Nonnull - public ListResult list( + ListResult list( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nullable Map requestFilters, int start, - int count, - @Nonnull Authentication authentication) + int count) throws RemoteInvocationException; /** @@ -251,20 +247,18 @@ public ListResult list( * @param sortCriterion sort criterion * @param start start offset for search results * @param count max number of search results requested - * @param searchFlags configuration flags for the search request * @return Snapshot key * @throws RemoteInvocationException */ @Nonnull - public SearchResult search( + SearchResult search( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nonnull String input, @Nullable Filter filter, SortCriterion sortCriterion, int start, - int count, - @Nonnull Authentication authentication, - @Nullable SearchFlags searchFlags) + int count) throws RemoteInvocationException; /** @@ -275,20 +269,18 @@ public SearchResult search( * @param filter search filters * @param start start offset for search results * @param count max number of search results requested - * @param searchFlags configuration flags for the search request * @return Snapshot key * @throws RemoteInvocationException */ @Nonnull - public SearchResult searchAcrossEntities( + SearchResult searchAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter filter, int start, int count, - @Nullable SearchFlags searchFlags, - @Nullable SortCriterion sortCriterion, - @Nonnull Authentication authentication) + @Nullable SortCriterion sortCriterion) throws RemoteInvocationException; /** @@ -299,21 +291,19 @@ public SearchResult searchAcrossEntities( * @param filter search filters * @param start start offset for search results * @param count max number of search results requested - * @param searchFlags configuration flags for the search request * @param facets list of facets we want aggregations for * @return Snapshot key * @throws RemoteInvocationException */ @Nonnull - public SearchResult searchAcrossEntities( + SearchResult searchAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter filter, int start, int count, - @Nullable SearchFlags searchFlags, @Nullable SortCriterion sortCriterion, - @Nonnull Authentication authentication, List facets) throws RemoteInvocationException; @@ -331,14 +321,13 @@ public SearchResult searchAcrossEntities( */ @Nonnull ScrollResult scrollAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter filter, @Nullable String scrollId, @Nullable String keepAlive, - int count, - @Nullable SearchFlags searchFlags, - @Nonnull Authentication authentication) + int count) throws RemoteInvocationException; /** @@ -353,12 +342,12 @@ ScrollResult scrollAcrossEntities( * @param sortCriterion {@link SortCriterion} to be applied to search results * @param start index to start the search from * @param count the number of search hits to return - * @param searchFlags configuration flags for the search request * @return a {@link SearchResult} that contains a list of matched documents and related search * result metadata */ @Nonnull - public LineageSearchResult searchAcrossLineage( + LineageSearchResult searchAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -367,9 +356,7 @@ public LineageSearchResult searchAcrossLineage( @Nullable Filter filter, @Nullable SortCriterion sortCriterion, int start, - int count, - @Nullable SearchFlags searchFlags, - @Nonnull final Authentication authentication) + int count) throws RemoteInvocationException; /** @@ -386,13 +373,12 @@ public LineageSearchResult searchAcrossLineage( * @param count the number of search hits to return * @param endTimeMillis end time to filter to * @param startTimeMillis start time to filter from - * @param searchFlags configuration flags for the search request - * @param authentication a reference to an authentication * @return a {@link SearchResult} that contains a list of matched documents and related search * result metadata */ @Nonnull - public LineageSearchResult searchAcrossLineage( + LineageSearchResult searchAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -403,9 +389,7 @@ public LineageSearchResult searchAcrossLineage( int start, int count, @Nullable final Long startTimeMillis, - @Nullable final Long endTimeMillis, - @Nullable SearchFlags searchFlags, - @Nonnull final Authentication authentication) + @Nullable final Long endTimeMillis) throws RemoteInvocationException; /** @@ -428,6 +412,7 @@ public LineageSearchResult searchAcrossLineage( */ @Nonnull LineageScrollResult scrollAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -439,9 +424,7 @@ LineageScrollResult scrollAcrossLineage( @Nonnull String keepAlive, int count, @Nullable final Long startTimeMillis, - @Nullable final Long endTimeMillis, - @Nullable SearchFlags searchFlags, - @Nonnull final Authentication authentication) + @Nullable final Long endTimeMillis) throws RemoteInvocationException; /** @@ -452,19 +435,19 @@ LineageScrollResult scrollAcrossLineage( * @throws RemoteInvocationException */ @Nonnull - public StringArray getBrowsePaths(@Nonnull Urn urn, @Nonnull Authentication authentication) + StringArray getBrowsePaths(@Nonnull Urn urn, @Nonnull Authentication authentication) throws RemoteInvocationException; - public void setWritable(boolean canWrite, @Nonnull Authentication authentication) + void setWritable(boolean canWrite, @Nonnull Authentication authentication) throws RemoteInvocationException; @Nonnull - public Map batchGetTotalEntityCount( - @Nonnull List entityName, @Nonnull Authentication authentication) + Map batchGetTotalEntityCount( + @Nonnull OperationContext opContext, @Nonnull List entityName) throws RemoteInvocationException; /** List all urns existing for a particular Entity type. */ - public ListUrnsResult listUrns( + ListUrnsResult listUrns( @Nonnull final String entityName, final int start, final int count, @@ -472,12 +455,11 @@ public ListUrnsResult listUrns( throws RemoteInvocationException; /** Hard delete an entity with a particular urn. */ - public void deleteEntity(@Nonnull final Urn urn, @Nonnull final Authentication authentication) + void deleteEntity(@Nonnull final Urn urn, @Nonnull final Authentication authentication) throws RemoteInvocationException; /** Delete all references to an entity with a particular urn. */ - public void deleteEntityReferences( - @Nonnull final Urn urn, @Nonnull final Authentication authentication) + void deleteEntityReferences(@Nonnull final Urn urn, @Nonnull final Authentication authentication) throws RemoteInvocationException; /** @@ -492,13 +474,13 @@ public void deleteEntityReferences( * @throws RemoteInvocationException */ @Nonnull - public SearchResult filter( + SearchResult filter( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nonnull Filter filter, @Nullable SortCriterion sortCriterion, int start, - int count, - @Nonnull Authentication authentication) + int count) throws RemoteInvocationException; /** @@ -510,12 +492,12 @@ public SearchResult filter( * @throws RemoteInvocationException */ @Nonnull - public boolean exists(@Nonnull Urn urn, @Nonnull Authentication authentication) + boolean exists(@Nonnull Urn urn, @Nonnull Authentication authentication) throws RemoteInvocationException; @Nullable @Deprecated - public VersionedAspect getAspect( + VersionedAspect getAspect( @Nonnull String urn, @Nonnull String aspect, @Nonnull Long version, @@ -524,7 +506,7 @@ public VersionedAspect getAspect( @Nullable @Deprecated - public VersionedAspect getAspectOrNull( + VersionedAspect getAspectOrNull( @Nonnull String urn, @Nonnull String aspect, @Nonnull Long version, @@ -545,7 +527,7 @@ default List getTimeseriesAspectValues( urn, entity, aspect, startTimeMillis, endTimeMillis, limit, filter, null, authentication); } - public List getTimeseriesAspectValues( + List getTimeseriesAspectValues( @Nonnull String urn, @Nonnull String entity, @Nonnull String aspect, @@ -609,7 +591,7 @@ default List batchIngestProposals( @Nonnull @Deprecated - public Optional getVersionedAspect( + Optional getVersionedAspect( @Nonnull String urn, @Nonnull String aspect, @Nonnull Long version, @@ -618,21 +600,21 @@ public Optional getVersionedAspect( throws RemoteInvocationException; @Deprecated - public DataMap getRawAspect( + DataMap getRawAspect( @Nonnull String urn, @Nonnull String aspect, @Nonnull Long version, @Nonnull Authentication authentication) throws RemoteInvocationException; - public void producePlatformEvent( + void producePlatformEvent( @Nonnull String name, @Nullable String key, @Nonnull PlatformEvent event, @Nonnull Authentication authentication) throws Exception; - public void rollbackIngestion( + void rollbackIngestion( @Nonnull String runId, @Nonnull Authorizer authorizer, @Nonnull Authentication authentication) throws Exception; diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClientCache.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClientCache.java index b5c6beaf4aa5b0..2fcb02946ca466 100644 --- a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClientCache.java +++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClientCache.java @@ -11,6 +11,7 @@ import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.metadata.config.cache.client.EntityClientCacheConfig; import com.linkedin.util.Pair; +import io.datahubproject.metadata.context.OperationContext; import java.util.Collection; import java.util.Map; import java.util.Optional; @@ -28,14 +29,19 @@ public class EntityClientCache { @NonNull private EntityClientCacheConfig config; @NonNull private final ClientCache cache; - @NonNull private BiFunction, Set, Map> loadFunction; + @NonNull private Function> loadFunction; - public EntityResponse getV2(@Nonnull final Urn urn, @Nonnull final Set aspectNames) { - return batchGetV2(Set.of(urn), aspectNames).get(urn); + public EntityResponse getV2( + @Nonnull OperationContext opContext, + @Nonnull final Urn urn, + @Nonnull final Set aspectNames) { + return batchGetV2(opContext, Set.of(urn), aspectNames).get(urn); } public Map batchGetV2( - @Nonnull final Set urns, @Nonnull final Set aspectNames) { + @Nonnull OperationContext opContext, + @Nonnull final Set urns, + @Nonnull final Set aspectNames) { final Map response; if (config.isEnabled()) { @@ -43,7 +49,14 @@ public Map batchGetV2( urns.stream() .flatMap( urn -> - aspectNames.stream().map(a -> Key.builder().urn(urn).aspectName(a).build())) + aspectNames.stream() + .map( + a -> + Key.builder() + .contextId(opContext.getEntityContextId()) + .urn(urn) + .aspectName(a) + .build())) .collect(Collectors.toSet()); Map envelopedAspects = cache.getAll(keys); @@ -61,7 +74,13 @@ public Map batchGetV2( response = responses.stream().collect(Collectors.toMap(EntityResponse::getUrn, Function.identity())); } else { - response = loadFunction.apply(urns, aspectNames); + response = + loadFunction.apply( + CollectionKey.builder() + .contextId(opContext.getEntityContextId()) + .urns(urns) + .aspectNames(aspectNames) + .build()); } return response; @@ -93,39 +112,16 @@ public EntityClientCache build(Class metricClazz) { // batch loads data from entity client (restli or java) Function, Map> loader = (Iterable keys) -> { - Map> keysByEntity = - StreamSupport.stream(keys.spliterator(), false) - .collect(Collectors.groupingBy(Key::getEntityName, Collectors.toSet())); - - Map results = - keysByEntity.entrySet().stream() - .flatMap( - entry -> { - Set urns = - entry.getValue().stream() - .map(Key::getUrn) - .collect(Collectors.toSet()); - Set aspects = - entry.getValue().stream() - .map(Key::getAspectName) - .collect(Collectors.toSet()); - return loadFunction.apply(urns, aspects).entrySet().stream(); - }) - .flatMap( - resp -> - resp.getValue().getAspects().values().stream() - .map( - envAspect -> { - Key key = - Key.builder() - .urn(resp.getKey()) - .aspectName(envAspect.getName()) - .build(); - return Map.entry(key, envAspect); - })) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - - return results; + Map>> keysByContextEntity = groupByContextEntity(keys); + + // load responses by context and combine + return keysByContextEntity.entrySet().stream() + .flatMap( + entry -> + loadByEntity(entry.getKey(), entry.getValue(), loadFunction) + .entrySet() + .stream()) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); }; // ideally the cache time comes from caching headers from service, but configuration driven @@ -149,9 +145,64 @@ public EntityClientCache build(Class metricClazz) { } } + private static Map>> groupByContextEntity( + Iterable keys) { + // group by context + Map> byContext = + StreamSupport.stream(keys.spliterator(), false) + .collect(Collectors.groupingBy(Key::getContextId, Collectors.toSet())); + + // then by entity + return byContext.entrySet().stream() + .map( + contextSet -> + Pair.of( + contextSet.getKey(), + contextSet.getValue().stream() + .collect(Collectors.groupingBy(Key::getEntityName, Collectors.toSet())))) + .collect(Collectors.toMap(Pair::getKey, Pair::getValue)); + } + + private static Map loadByEntity( + String contextId, + Map> keysByEntity, + Function> loadFunction) { + return keysByEntity.entrySet().stream() + .flatMap( + entry -> { + Set urns = + entry.getValue().stream().map(Key::getUrn).collect(Collectors.toSet()); + Set aspects = + entry.getValue().stream().map(Key::getAspectName).collect(Collectors.toSet()); + return loadFunction + .apply( + CollectionKey.builder() + .contextId(contextId) + .urns(urns) + .aspectNames(aspects) + .build()) + .entrySet() + .stream(); + }) + .flatMap( + resp -> + resp.getValue().getAspects().values().stream() + .map( + envAspect -> { + Key key = + Key.builder() + .urn(resp.getKey()) + .aspectName(envAspect.getName()) + .build(); + return Map.entry(key, envAspect); + })) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + } + @Data @Builder protected static class Key { + private final String contextId; private final Urn urn; private final String aspectName; @@ -159,4 +210,12 @@ public String getEntityName() { return urn.getEntityType(); } } + + @Data + @Builder + public static class CollectionKey { + private final String contextId; + private final Set urns; + private final Set aspectNames; + } } diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java index 653ef046ffc021..4dfe36b49bf11e 100644 --- a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java +++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java @@ -75,6 +75,7 @@ import com.linkedin.restli.client.Client; import com.linkedin.restli.client.RestLiResponseException; import com.linkedin.restli.common.HttpStatus; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.Collection; import java.util.HashMap; @@ -114,6 +115,7 @@ public RestliEntityClient( super(restliClient, backoffPolicy, retryCount); } + @Override @Nullable public EntityResponse getV2( @Nonnull String entityName, @@ -126,6 +128,7 @@ public EntityResponse getV2( return sendClientRequest(requestBuilder, authentication).getEntity(); } + @Override @Nonnull public Entity get(@Nonnull final Urn urn, @Nonnull final Authentication authentication) throws RemoteInvocationException { @@ -143,6 +146,7 @@ public Entity get(@Nonnull final Urn urn, @Nonnull final Authentication authenti * @param authentication the authentication to include in the request to the Metadata Service * @throws RemoteInvocationException */ + @Override @Nonnull public Map batchGet( @Nonnull final Set urns, @Nonnull final Authentication authentication) @@ -195,6 +199,7 @@ public Map batchGet( * @param authentication the authentication to include in the request to the Metadata Service * @throws RemoteInvocationException */ + @Override @Nonnull public Map batchGetV2( @Nonnull String entityName, @@ -237,6 +242,7 @@ public Map batchGetV2( * @param authentication the authentication to include in the request to the Metadata Service * @throws RemoteInvocationException */ + @Override @Nonnull public Map batchGetVersionedV2( @Nonnull String entityName, @@ -280,14 +286,15 @@ public Map batchGetVersionedV2( * @param field the field to autocomplete against, e.g. 'name' * @throws RemoteInvocationException */ + @Override @Nonnull public AutoCompleteResult autoComplete( + @Nonnull OperationContext opContext, @Nonnull String entityType, @Nonnull String query, @Nullable Filter requestFilters, @Nonnull int limit, - @Nullable String field, - @Nonnull final Authentication authentication) + @Nullable String field) throws RemoteInvocationException { EntitiesDoAutocompleteRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS @@ -297,7 +304,7 @@ public AutoCompleteResult autoComplete( .fieldParam(field) .filterParam(filterOrDefaultEmptyFilter(requestFilters)) .limitParam(limit); - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } /** @@ -309,13 +316,14 @@ public AutoCompleteResult autoComplete( * @param limit max number of autocomplete results * @throws RemoteInvocationException */ + @Override @Nonnull public AutoCompleteResult autoComplete( + @Nonnull OperationContext opContext, @Nonnull String entityType, @Nonnull String query, @Nullable Filter requestFilters, - @Nonnull int limit, - @Nonnull final Authentication authentication) + @Nonnull int limit) throws RemoteInvocationException { EntitiesDoAutocompleteRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS @@ -324,7 +332,7 @@ public AutoCompleteResult autoComplete( .queryParam(query) .filterParam(filterOrDefaultEmptyFilter(requestFilters)) .limitParam(limit); - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } /** @@ -337,14 +345,15 @@ public AutoCompleteResult autoComplete( * @param limit max number of datasets * @throws RemoteInvocationException */ + @Override @Nonnull public BrowseResult browse( + @Nonnull OperationContext opContext, @Nonnull String entityType, @Nonnull String path, @Nullable Map requestFilters, int start, - int limit, - @Nonnull final Authentication authentication) + int limit) throws RemoteInvocationException { EntitiesDoBrowseRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS @@ -356,7 +365,7 @@ public BrowseResult browse( if (requestFilters != null) { requestBuilder.filterParam(newFilter(requestFilters)); } - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } /** @@ -370,34 +379,34 @@ public BrowseResult browse( * @param count max number of results requested * @throws RemoteInvocationException */ + @Override @Nonnull public BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nonnull Authentication authentication, - @Nullable SearchFlags searchFlags) { + int count) { throw new NotImplementedException("BrowseV2 is not implemented in Restli yet"); } @Nonnull @Override public BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nonnull Authentication authentication, - @Nullable SearchFlags searchFlags) + int count) throws RemoteInvocationException { throw new NotImplementedException("BrowseV2 is not implemented in Restli yet"); } + @Override public void update(@Nonnull final Entity entity, @Nonnull final Authentication authentication) throws RemoteInvocationException { EntitiesDoIngestRequestBuilder requestBuilder = @@ -405,6 +414,7 @@ public void update(@Nonnull final Entity entity, @Nonnull final Authentication a sendClientRequest(requestBuilder, authentication); } + @Override public void updateWithSystemMetadata( @Nonnull final Entity entity, @Nullable final SystemMetadata systemMetadata, @@ -424,6 +434,7 @@ public void updateWithSystemMetadata( sendClientRequest(requestBuilder, authentication); } + @Override public void batchUpdate( @Nonnull final Set entities, @Nonnull final Authentication authentication) throws RemoteInvocationException { @@ -440,22 +451,22 @@ public void batchUpdate( * @param requestFilters search filters * @param start start offset for search results * @param count max number of search results requested - * @param searchFlags configuration flags for the search request * @return a set of search results * @throws RemoteInvocationException */ @Nonnull @Override public SearchResult search( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nonnull String input, @Nullable Map requestFilters, int start, - int count, - @Nonnull final Authentication authentication, - @Nullable SearchFlags searchFlags) + int count) throws RemoteInvocationException { + SearchFlags searchFlags = opContext.getSearchContext().getSearchFlags(); + final EntitiesDoSearchRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS .actionSearch() @@ -465,11 +476,9 @@ public SearchResult search( .startParam(start) .fulltextParam(searchFlags != null ? searchFlags.isFulltext() : null) .countParam(count); - if (searchFlags != null) { - requestBuilder.searchFlagsParam(searchFlags); - } + requestBuilder.searchFlagsParam(opContext.getSearchContext().getSearchFlags()); - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } /** @@ -481,13 +490,14 @@ public SearchResult search( * @return a set of list results * @throws RemoteInvocationException */ + @Override @Nonnull public ListResult list( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nullable Map requestFilters, int start, - int count, - @Nonnull final Authentication authentication) + int count) throws RemoteInvocationException { final EntitiesDoListRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS @@ -497,7 +507,7 @@ public ListResult list( .startParam(start) .countParam(count); - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } /** @@ -514,16 +524,16 @@ public ListResult list( @Nonnull @Override public SearchResult search( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nonnull String input, @Nullable Filter filter, SortCriterion sortCriterion, int start, - int count, - @Nonnull final Authentication authentication, - @Nullable SearchFlags searchFlags) + int count) throws RemoteInvocationException { + SearchFlags searchFlags = opContext.getSearchContext().getSearchFlags(); final EntitiesDoSearchRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS .actionSearch() @@ -547,22 +557,22 @@ public SearchResult search( } } - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } + @Override @Nonnull public SearchResult searchAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter filter, int start, int count, - @Nullable SearchFlags searchFlags, - @Nullable SortCriterion sortCriterion, - @Nonnull final Authentication authentication) + @Nullable SortCriterion sortCriterion) throws RemoteInvocationException { return searchAcrossEntities( - entities, input, filter, start, count, searchFlags, sortCriterion, authentication, null); + opContext, entities, input, filter, start, count, sortCriterion, null); } /** @@ -577,19 +587,20 @@ public SearchResult searchAcrossEntities( * @return Snapshot key * @throws RemoteInvocationException */ + @Override @Nonnull public SearchResult searchAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter filter, int start, int count, - @Nullable SearchFlags searchFlags, @Nullable SortCriterion sortCriterion, - @Nonnull final Authentication authentication, @Nullable List facets) throws RemoteInvocationException { + SearchFlags searchFlags = opContext.getSearchContext().getSearchFlags(); final EntitiesDoSearchAcrossEntitiesRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS .actionSearchAcrossEntities() @@ -611,21 +622,21 @@ public SearchResult searchAcrossEntities( requestBuilder.sortParam(sortCriterion); } - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } @Nonnull @Override public ScrollResult scrollAcrossEntities( + @Nonnull OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter filter, @Nullable String scrollId, @Nullable String keepAlive, - int count, - @Nullable SearchFlags searchFlags, - @Nonnull Authentication authentication) + int count) throws RemoteInvocationException { + final SearchFlags searchFlags = opContext.getSearchContext().getSearchFlags(); final EntitiesDoScrollAcrossEntitiesRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS.actionScrollAcrossEntities().inputParam(input).countParam(count); @@ -645,12 +656,13 @@ public ScrollResult scrollAcrossEntities( requestBuilder.keepAliveParam(keepAlive); } - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } @Nonnull @Override public LineageSearchResult searchAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -659,9 +671,7 @@ public LineageSearchResult searchAcrossLineage( @Nullable Filter filter, @Nullable SortCriterion sortCriterion, int start, - int count, - @Nullable SearchFlags searchFlags, - @Nonnull final Authentication authentication) + int count) throws RemoteInvocationException { final EntitiesDoSearchAcrossLineageRequestBuilder requestBuilder = @@ -679,16 +689,15 @@ public LineageSearchResult searchAcrossLineage( if (filter != null) { requestBuilder.filterParam(filter); } - if (searchFlags != null) { - requestBuilder.searchFlagsParam(searchFlags); - } + requestBuilder.searchFlagsParam(opContext.getSearchContext().getSearchFlags()); - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } @Nonnull @Override public LineageSearchResult searchAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -699,9 +708,7 @@ public LineageSearchResult searchAcrossLineage( int start, int count, @Nullable final Long startTimeMillis, - @Nullable final Long endTimeMillis, - @Nullable SearchFlags searchFlags, - @Nonnull final Authentication authentication) + @Nullable final Long endTimeMillis) throws RemoteInvocationException { final EntitiesDoSearchAcrossLineageRequestBuilder requestBuilder = @@ -725,15 +732,14 @@ public LineageSearchResult searchAcrossLineage( if (endTimeMillis != null) { requestBuilder.endTimeMillisParam(endTimeMillis); } - if (searchFlags != null) { - requestBuilder.searchFlagsParam(searchFlags); - } + requestBuilder.searchFlagsParam(opContext.getSearchContext().getSearchFlags()); - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } @Override public LineageScrollResult scrollAcrossLineage( + @Nonnull OperationContext opContext, @Nonnull Urn sourceUrn, @Nonnull LineageDirection direction, @Nonnull List entities, @@ -745,9 +751,7 @@ public LineageScrollResult scrollAcrossLineage( @Nonnull String keepAlive, int count, @Nullable final Long startTimeMillis, - @Nullable final Long endTimeMillis, - @Nullable final SearchFlags searchFlags, - @Nonnull final Authentication authentication) + @Nullable final Long endTimeMillis) throws RemoteInvocationException { final EntitiesDoScrollAcrossLineageRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS @@ -773,11 +777,9 @@ public LineageScrollResult scrollAcrossLineage( if (endTimeMillis != null) { requestBuilder.endTimeMillisParam(endTimeMillis); } - if (searchFlags != null) { - requestBuilder.searchFlagsParam(searchFlags); - } + requestBuilder.searchFlagsParam(opContext.getSearchContext().getSearchFlags()); - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } /** @@ -802,18 +804,20 @@ public void setWritable(boolean canWrite, @Nonnull final Authentication authenti sendClientRequest(requestBuilder, authentication); } + @Override @Nonnull public Map batchGetTotalEntityCount( - @Nonnull List entityName, @Nonnull final Authentication authentication) + @Nonnull OperationContext opContext, @Nonnull List entityName) throws RemoteInvocationException { EntitiesDoBatchGetTotalEntityCountRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS .actionBatchGetTotalEntityCount() .entitiesParam(new StringArray(entityName)); - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } /** List all urns existing for a particular Entity type. */ + @Override public ListUrnsResult listUrns( @Nonnull final String entityName, final int start, @@ -849,12 +853,12 @@ public void deleteEntityReferences(@Nonnull Urn urn, @Nonnull Authentication aut @Nonnull @Override public SearchResult filter( + @Nonnull OperationContext opContext, @Nonnull String entity, @Nonnull Filter filter, @Nullable SortCriterion sortCriterion, int start, - int count, - @Nonnull final Authentication authentication) + int count) throws RemoteInvocationException { EntitiesDoFilterRequestBuilder requestBuilder = ENTITIES_REQUEST_BUILDERS @@ -866,7 +870,7 @@ public SearchResult filter( if (sortCriterion != null) { requestBuilder.sortParam(sortCriterion); } - return sendClientRequest(requestBuilder, authentication).getEntity(); + return sendClientRequest(requestBuilder, opContext.getAuthentication()).getEntity(); } @Nonnull @@ -885,6 +889,7 @@ public boolean exists(@Nonnull Urn urn, @Nonnull final Authentication authentica * @return list of paths given urn * @throws RemoteInvocationException on remote request error. */ + @Override @Nonnull public VersionedAspect getAspect( @Nonnull String urn, @@ -906,6 +911,7 @@ public VersionedAspect getAspect( * @return list of paths given urn * @throws RemoteInvocationException on remote request error. */ + @Override @Nullable public VersionedAspect getAspectOrNull( @Nonnull String urn, @@ -940,6 +946,7 @@ public VersionedAspect getAspectOrNull( * @return the list of EnvelopedAspect values satisfying the input parameters. * @throws RemoteInvocationException on remote request error. */ + @Override @Nonnull public List getTimeseriesAspectValues( @Nonnull String urn, @@ -1002,6 +1009,7 @@ public String ingestProposal( return sendClientRequest(requestBuilder, authentication).getEntity(); } + @Override public Optional getVersionedAspect( @Nonnull String urn, @Nonnull String aspect, diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java index 243e8a40bf4b75..72af56b8c7fc74 100644 --- a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java +++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java @@ -5,16 +5,14 @@ import com.linkedin.entity.Aspect; import com.linkedin.entity.EntityResponse; import com.linkedin.metadata.config.cache.client.EntityClientCacheConfig; -import com.linkedin.metadata.query.SearchFlags; -import com.linkedin.metadata.query.filter.Filter; -import com.linkedin.metadata.search.ScrollResult; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.PlatformEvent; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; -import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -23,65 +21,43 @@ public interface SystemEntityClient extends EntityClient { EntityClientCache getEntityClientCache(); - Authentication getSystemAuthentication(); + @Nonnull + ConcurrentHashMap getOperationContextMap(); - /** - * Searches for entities matching to a given query and filters across multiple entity types - * - * @param entities entity types to search (if empty, searches all entities) - * @param input search query - * @param filter search filters - * @param scrollId opaque scroll ID indicating offset - * @param keepAlive string representation of time to keep point in time alive, ex: 5m - * @param count max number of search results requested - * @return Snapshot key - * @throws RemoteInvocationException - */ @Nonnull - default ScrollResult scrollAcrossEntities( - @Nonnull List entities, - @Nonnull String input, - @Nullable Filter filter, - @Nullable String scrollId, - @Nullable String keepAlive, - int count, - @Nullable SearchFlags searchFlags) - throws RemoteInvocationException { - return scrollAcrossEntities( - entities, - input, - filter, - scrollId, - keepAlive, - count, - searchFlags, - getSystemAuthentication()); + OperationContext getSystemOperationContext(); + + default Authentication getSystemAuthentication() { + return getSystemOperationContext().getAuthentication(); } /** * Builds the cache * - * @param systemAuthentication system authentication * @param cacheConfig cache configuration * @return the cache */ default EntityClientCache buildEntityClientCache( - Class metricClazz, - Authentication systemAuthentication, - EntityClientCacheConfig cacheConfig) { + Class metricClazz, EntityClientCacheConfig cacheConfig) { return EntityClientCache.builder() .config(cacheConfig) .loadFunction( - (Set urns, Set aspectNames) -> { + (EntityClientCache.CollectionKey collectionKey) -> { try { - String entityName = urns.stream().findFirst().map(Urn::getEntityType).get(); + String entityName = + collectionKey.getUrns().stream().findFirst().map(Urn::getEntityType).get(); - if (urns.stream().anyMatch(urn -> !urn.getEntityType().equals(entityName))) { + if (collectionKey.getUrns().stream() + .anyMatch(urn -> !urn.getEntityType().equals(entityName))) { throw new IllegalArgumentException( "Urns must be of the same entity type. RestliEntityClient API limitation."); } - return batchGetV2(entityName, urns, aspectNames, systemAuthentication); + return batchGetV2( + entityName, + collectionKey.getUrns(), + collectionKey.getAspectNames(), + getSystemOperationContext().getAuthentication()); } catch (RemoteInvocationException | URISyntaxException e) { throw new RuntimeException(e); } @@ -101,7 +77,7 @@ default EntityClientCache buildEntityClientCache( @Nullable default EntityResponse getV2(@Nonnull Urn urn, @Nonnull Set aspectNames) throws RemoteInvocationException, URISyntaxException { - return getEntityClientCache().getV2(urn, aspectNames); + return getEntityClientCache().getV2(getSystemOperationContext(), urn, aspectNames); } /** @@ -115,7 +91,7 @@ default EntityResponse getV2(@Nonnull Urn urn, @Nonnull Set aspectNames) default Map batchGetV2( @Nonnull Set urns, @Nonnull Set aspectNames) throws RemoteInvocationException, URISyntaxException { - return getEntityClientCache().batchGetV2(urns, aspectNames); + return getEntityClientCache().batchGetV2(getSystemOperationContext(), urns, aspectNames); } default void producePlatformEvent( diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemRestliEntityClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemRestliEntityClient.java index 0f179c4da7b74c..7aad31b5beeba2 100644 --- a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemRestliEntityClient.java +++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemRestliEntityClient.java @@ -1,9 +1,10 @@ package com.linkedin.entity.client; -import com.datahub.authentication.Authentication; import com.linkedin.metadata.config.cache.client.EntityClientCacheConfig; import com.linkedin.parseq.retry.backoff.BackoffPolicy; import com.linkedin.restli.client.Client; +import io.datahubproject.metadata.context.OperationContext; +import java.util.concurrent.ConcurrentHashMap; import javax.annotation.Nonnull; import lombok.Getter; @@ -11,17 +12,18 @@ @Getter public class SystemRestliEntityClient extends RestliEntityClient implements SystemEntityClient { private final EntityClientCache entityClientCache; - private final Authentication systemAuthentication; + private final OperationContext systemOperationContext; + private final ConcurrentHashMap operationContextMap; public SystemRestliEntityClient( + @Nonnull OperationContext systemOperationContext, @Nonnull final Client restliClient, @Nonnull final BackoffPolicy backoffPolicy, int retryCount, - @Nonnull Authentication systemAuthentication, EntityClientCacheConfig cacheConfig) { super(restliClient, backoffPolicy, retryCount); - this.systemAuthentication = systemAuthentication; - this.entityClientCache = - buildEntityClientCache(SystemRestliEntityClient.class, systemAuthentication, cacheConfig); + this.operationContextMap = new ConcurrentHashMap<>(); + this.systemOperationContext = systemOperationContext; + this.entityClientCache = buildEntityClientCache(SystemRestliEntityClient.class, cacheConfig); } } diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClient.java index 747e1e0e1a2887..461c2e50fac548 100644 --- a/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClient.java +++ b/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClient.java @@ -8,6 +8,7 @@ import com.linkedin.parseq.retry.backoff.BackoffPolicy; import com.linkedin.r2.RemoteInvocationException; import com.linkedin.restli.client.Client; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import javax.annotation.Nonnull; @@ -16,22 +17,27 @@ public class UsageClient extends BaseClient { private static final UsageStatsRequestBuilders USAGE_STATS_REQUEST_BUILDERS = new UsageStatsRequestBuilders(); + private final OperationContext systemOperationContext; private final UsageClientCache usageClientCache; public UsageClient( + @Nonnull OperationContext systemOperationContext, @Nonnull final Client restliClient, @Nonnull final BackoffPolicy backoffPolicy, int retryCount, - Authentication systemAuthentication, UsageClientCacheConfig cacheConfig) { super(restliClient, backoffPolicy, retryCount); + this.systemOperationContext = systemOperationContext; this.usageClientCache = UsageClientCache.builder() .config(cacheConfig) .loadFunction( - (String resource, UsageTimeRange range) -> { + (UsageClientCache.Key cacheKey) -> { try { - return getUsageStats(resource, range, systemAuthentication); + return getUsageStats( + cacheKey.getResource(), + cacheKey.getRange(), + systemOperationContext.getAuthentication()); } catch (RemoteInvocationException | URISyntaxException e) { throw new RuntimeException(e); } @@ -45,7 +51,7 @@ public UsageClient( */ @Nonnull public UsageQueryResult getUsageStats(@Nonnull String resource, @Nonnull UsageTimeRange range) { - return usageClientCache.getUsageStats(resource, range); + return usageClientCache.getUsageStats(systemOperationContext, resource, range); } /** Gets a specific version of downstream {@link EntityRelationships} for the given dataset. */ diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClientCache.java b/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClientCache.java index b56b59d0feec0d..1d12f0395e3f6d 100644 --- a/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClientCache.java +++ b/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClientCache.java @@ -4,6 +4,7 @@ import com.github.benmanes.caffeine.cache.Weigher; import com.linkedin.common.client.ClientCache; import com.linkedin.metadata.config.cache.client.UsageClientCacheConfig; +import io.datahubproject.metadata.context.OperationContext; import java.util.Map; import java.util.function.BiFunction; import java.util.function.Function; @@ -18,13 +19,22 @@ public class UsageClientCache { @NonNull private UsageClientCacheConfig config; @NonNull private final ClientCache cache; - @NonNull private BiFunction loadFunction; + @NonNull private Function loadFunction; - public UsageQueryResult getUsageStats(@Nonnull String resource, @Nonnull UsageTimeRange range) { + public UsageQueryResult getUsageStats( + @Nonnull OperationContext opContext, + @Nonnull String resource, + @Nonnull UsageTimeRange range) { + Key cacheKey = + Key.builder() + .contextId(opContext.getEntityContextId()) + .resource(resource) + .range(range) + .build(); if (config.isEnabled()) { - return cache.get(Key.builder().resource(resource).range(range).build()); + return cache.get(cacheKey); } else { - return loadFunction.apply(resource, range); + return loadFunction.apply(cacheKey); } } @@ -43,7 +53,7 @@ public UsageClientCache build() { Function, Map> loader = (Iterable keys) -> StreamSupport.stream(keys.spliterator(), false) - .map(k -> Map.entry(k, loadFunction.apply(k.getResource(), k.getRange()))) + .map(k -> Map.entry(k, loadFunction.apply(k))) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); // default ttl only @@ -64,7 +74,8 @@ public UsageClientCache build() { @Data @Builder - protected static class Key { + public static class Key { + private final String contextId; private final String resource; private final UsageTimeRange range; } 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 afdaf06802a11d..27620e0b49d0cf 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 @@ -12,6 +12,8 @@ import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; import com.datahub.authorization.EntitySpec; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.metadata.context.OperationContextConfig; import com.datahub.plugins.auth.authorization.Authorizer; import com.google.common.collect.ImmutableList; import com.linkedin.common.AuditStamp; @@ -159,6 +161,10 @@ public class EntityResource extends CollectionResourceTaskTemplate search( throw new RestLiServiceException( HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } + + OperationContext opContext = OperationContext.asSession(systemOperationContext, _authorizer, auth, true); + log.info("GET SEARCH RESULTS for {} with query {}", entityName, input); // TODO - change it to use _searchService once we are confident on it's latency return RestliUtil.toTask( @@ -376,8 +385,8 @@ public Task search( final SearchResult result; // This API is not used by the frontend for search bars so we default to structured result = - _entitySearchService.search( - List.of(entityName), input, filter, sortCriterion, start, count, searchFlags); + _entitySearchService.search(opContext, + List.of(entityName), input, filter, sortCriterion, start, count); return validateSearchResult(result, _entityService); }, MetricRegistry.name(this.getClass(), "search")); @@ -404,15 +413,17 @@ public Task searchAcrossEntities( throw new RestLiServiceException( HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } + OperationContext opContext = OperationContext.asSession( + systemOperationContext, _authorizer, auth, true) + .withSearchFlags(flags -> searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true)); + List entityList = entities == null ? Collections.emptyList() : Arrays.asList(entities); log.info("GET SEARCH RESULTS ACROSS ENTITIES for {} with query {}", entityList, input); - final SearchFlags finalFlags = - searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true); return RestliUtil.toTask( () -> validateSearchResult( - _searchService.searchAcrossEntities( - entityList, input, filter, sortCriterion, start, count, finalFlags), + _searchService.searchAcrossEntities(opContext, + entityList, input, filter, sortCriterion, start, count), _entityService), "searchAcrossEntities"); } @@ -429,26 +440,30 @@ public Task scrollAcrossEntities( @ActionParam(PARAM_KEEP_ALIVE) String keepAlive, @ActionParam(PARAM_COUNT) int count, @ActionParam(PARAM_SEARCH_FLAGS) @Optional SearchFlags searchFlags) { + Authentication auth = AuthenticationContext.getAuthentication(); + OperationContext opContext = OperationContext.asSession( + systemOperationContext, _authorizer, auth, true) + .withSearchFlags(flags -> searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true)); + List entityList = entities == null ? Collections.emptyList() : Arrays.asList(entities); log.info( "GET SCROLL RESULTS ACROSS ENTITIES for {} with query {} and scroll ID: {}", entityList, input, scrollId); - final SearchFlags finalFlags = - searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true); + return RestliUtil.toTask( () -> validateScrollResult( _searchService.scrollAcrossEntities( + opContext, entityList, input, filter, sortCriterion, scrollId, keepAlive, - count, - finalFlags), + count), _entityService), "scrollAcrossEntities"); } @@ -470,7 +485,12 @@ public Task searchAcrossLineage( @ActionParam(PARAM_END_TIME_MILLIS) @Optional @Nullable Long endTimeMillis, @Optional @Nullable @ActionParam(PARAM_SEARCH_FLAGS) SearchFlags searchFlags) throws URISyntaxException { + Authentication auth = AuthenticationContext.getAuthentication(); + OperationContext opContext = OperationContext.asSession( + systemOperationContext, _authorizer, auth, true) + .withSearchFlags(flags -> searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true)); + if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) && !isAuthorized( auth, @@ -492,6 +512,7 @@ public Task searchAcrossLineage( () -> validateLineageSearchResult( _lineageSearchService.searchAcrossLineage( + opContext, urn, LineageDirection.valueOf(direction), entityList, @@ -502,8 +523,7 @@ public Task searchAcrossLineage( start, count, startTimeMillis, - endTimeMillis, - searchFlags), + endTimeMillis), _entityService), "searchAcrossRelationships"); } @@ -526,6 +546,12 @@ public Task scrollAcrossLineage( @ActionParam(PARAM_END_TIME_MILLIS) @Optional @Nullable Long endTimeMillis, @ActionParam(PARAM_SEARCH_FLAGS) @Optional @Nullable SearchFlags searchFlags) throws URISyntaxException { + + Authentication auth = AuthenticationContext.getAuthentication(); + OperationContext opContext = OperationContext.asSession( + systemOperationContext, _authorizer, auth, true) + .withSearchFlags(flags -> searchFlags != null ? searchFlags : new SearchFlags().setSkipCache(true)); + Urn urn = Urn.createFromString(urnStr); List entityList = entities == null ? Collections.emptyList() : Arrays.asList(entities); log.info( @@ -534,12 +560,12 @@ public Task scrollAcrossLineage( direction, entityList, input); - final SearchFlags finalFlags = - searchFlags != null ? searchFlags : new SearchFlags().setSkipCache(true); + return RestliUtil.toTask( () -> validateLineageScrollResult( _lineageSearchService.scrollAcrossLineage( + opContext, urn, LineageDirection.valueOf(direction), entityList, @@ -551,8 +577,7 @@ public Task scrollAcrossLineage( keepAlive, count, startTimeMillis, - endTimeMillis, - finalFlags), + endTimeMillis), _entityService), "scrollAcrossLineage"); } @@ -577,12 +602,16 @@ public Task list( throw new RestLiServiceException( HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } + OperationContext opContext = OperationContext.asSession( + systemOperationContext, _authorizer, auth, true) + .withSearchFlags(flags -> new SearchFlags().setFulltext(false)); + log.info("GET LIST RESULTS for {} with filter {}", entityName, filter); return RestliUtil.toTask( () -> validateListResult( toListResult( - _entitySearchService.filter(entityName, filter, sortCriterion, start, count)), + _entitySearchService.filter(opContext, entityName, filter, sortCriterion, start, count)), _entityService), MetricRegistry.name(this.getClass(), "filter")); } @@ -595,7 +624,8 @@ public Task autocomplete( @ActionParam(PARAM_QUERY) @Nonnull String query, @ActionParam(PARAM_FIELD) @Optional @Nullable String field, @ActionParam(PARAM_FILTER) @Optional @Nullable Filter filter, - @ActionParam(PARAM_LIMIT) int limit) { + @ActionParam(PARAM_LIMIT) int limit, + @ActionParam(PARAM_SEARCH_FLAGS) @Optional @Nullable SearchFlags searchFlags) { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) @@ -607,8 +637,12 @@ public Task autocomplete( throw new RestLiServiceException( HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } + OperationContext opContext = OperationContext.asSession( + systemOperationContext, _authorizer, auth, true) + .withSearchFlags(flags -> searchFlags != null ? searchFlags : flags); + return RestliUtil.toTask( - () -> _entitySearchService.autoComplete(entityName, query, field, filter, limit), + () -> _entitySearchService.autoComplete(opContext, entityName, query, field, filter, limit), MetricRegistry.name(this.getClass(), "autocomplete")); } @@ -620,7 +654,8 @@ public Task browse( @ActionParam(PARAM_PATH) @Nonnull String path, @ActionParam(PARAM_FILTER) @Optional @Nullable Filter filter, @ActionParam(PARAM_START) int start, - @ActionParam(PARAM_LIMIT) int limit) { + @ActionParam(PARAM_LIMIT) int limit, + @ActionParam(PARAM_SEARCH_FLAGS) @Optional @Nullable SearchFlags searchFlags) { Authentication auth = AuthenticationContext.getAuthentication(); if (Boolean.parseBoolean(System.getenv(REST_API_AUTHORIZATION_ENABLED_ENV)) @@ -632,11 +667,15 @@ public Task browse( throw new RestLiServiceException( HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } + OperationContext opContext = OperationContext.asSession( + systemOperationContext, _authorizer, auth, true) + .withSearchFlags(flags -> searchFlags != null ? searchFlags : flags); + log.info("GET BROWSE RESULTS for {} at path {}", entityName, path); return RestliUtil.toTask( () -> validateBrowseResult( - _entitySearchService.browse(entityName, path, filter, start, limit), + _entitySearchService.browse(opContext, entityName, path, filter, start, limit), _entityService), MetricRegistry.name(this.getClass(), "browse")); } @@ -937,7 +976,9 @@ public Task getTotalEntityCount(@ActionParam(PARAM_ENTITY) @Nonnull String throw new RestLiServiceException( HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get entity counts."); } - return RestliUtil.toTask(() -> _entitySearchService.docCount(entityName)); + OperationContext opContext = OperationContext.asSession( + systemOperationContext, _authorizer, auth, true); + return RestliUtil.toTask(() -> _entitySearchService.docCount(opContext, entityName)); } @Action(name = "batchGetTotalEntityCount") @@ -955,8 +996,10 @@ public Task batchGetTotalEntityCount( throw new RestLiServiceException( HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to get entity counts."); } + OperationContext opContext = OperationContext.asSession( + systemOperationContext, _authorizer, auth, true); return RestliUtil.toTask( - () -> new LongMap(_searchService.docCountPerEntity(Arrays.asList(entityNames)))); + () -> new LongMap(_searchService.docCountPerEntity(opContext, Arrays.asList(entityNames)))); } @Action(name = ACTION_LIST_URNS) @@ -1030,11 +1073,13 @@ public Task filter( throw new RestLiServiceException( HttpStatus.S_401_UNAUTHORIZED, "User is unauthorized to search."); } + OperationContext opContext = OperationContext.asSession( + systemOperationContext, _authorizer, auth, true); log.info("FILTER RESULTS for {} with filter {}", entityName, filter); return RestliUtil.toTask( () -> validateSearchResult( - _entitySearchService.filter(entityName, filter, sortCriterion, start, count), + _entitySearchService.filter(opContext.withSearchFlags(flags -> flags.setFulltext(true)), entityName, filter, sortCriterion, start, count), _entityService), MetricRegistry.name(this.getClass(), "search")); } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/AspectUtils.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/AspectUtils.java index 55373730e7b673..915756fc8da35e 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/AspectUtils.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/AspectUtils.java @@ -2,7 +2,6 @@ import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableSet; -import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; import com.linkedin.data.template.RecordTemplate; import com.linkedin.entity.Aspect; @@ -19,7 +18,6 @@ import java.util.Set; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; -import org.joda.time.DateTimeUtils; @Slf4j public class AspectUtils { @@ -70,13 +68,6 @@ public static MetadataChangeProposal buildMetadataChangeProposal( return proposal; } - public static AuditStamp getAuditStamp(Urn actor) { - AuditStamp auditStamp = new AuditStamp(); - auditStamp.setTime(DateTimeUtils.currentTimeMillis()); - auditStamp.setActor(actor); - return auditStamp; - } - public static AspectSpec validateAspect(MetadataChangeLog mcl, EntitySpec entitySpec) { if (!mcl.hasAspectName() || (!ChangeType.DELETE.equals(mcl.getChangeType()) && !mcl.hasAspect())) { diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/RecommendationsService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/RecommendationsService.java index 5676dc9ebac54e..fcea114446c498 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/RecommendationsService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/RecommendationsService.java @@ -1,9 +1,9 @@ package com.linkedin.metadata.recommendation; -import com.linkedin.common.urn.Urn; import com.linkedin.metadata.recommendation.candidatesource.RecommendationSource; import com.linkedin.metadata.recommendation.ranker.RecommendationModuleRanker; import com.linkedin.metadata.utils.ConcurrencyUtils; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.util.List; import java.util.Map; @@ -47,7 +47,7 @@ private void validateRecommendationSources(final List cand /** * Return the list of recommendation modules given input context * - * @param userUrn User requesting recommendations + * @param opContext User's context requesting recommendations * @param requestContext Context of where the recommendations are being requested * @param limit Max number of modules to return * @return List of recommendation modules @@ -55,14 +55,17 @@ private void validateRecommendationSources(final List cand @Nonnull @WithSpan public List listRecommendations( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext, int limit) { + @Nonnull OperationContext opContext, + @Nonnull RecommendationRequestContext requestContext, + int limit) { + // Get recommendation candidates from sources which are eligible, in parallel final List candidateModules = ConcurrencyUtils.transformAndCollectAsync( _candidateSources.stream() - .filter(source -> source.isEligible(userUrn, requestContext)) + .filter(source -> source.isEligible(opContext, requestContext)) .collect(Collectors.toList()), - source -> source.getRecommendationModule(userUrn, requestContext), + source -> source.getRecommendationModule(opContext, requestContext), (source, exception) -> { log.error( "Error while fetching candidate modules from source {}", source, exception); @@ -74,6 +77,6 @@ public List listRecommendations( .collect(Collectors.toList()); // Rank recommendation modules, which determines their ordering during rendering - return _moduleRanker.rank(candidateModules, userUrn, requestContext, limit); + return _moduleRanker.rank(opContext, requestContext, candidateModules, limit); } } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/DomainsCandidateSource.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/DomainsCandidateSource.java index e34fa8ff1bde57..dd2bdfa57dfbf9 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/DomainsCandidateSource.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/DomainsCandidateSource.java @@ -1,11 +1,11 @@ package com.linkedin.metadata.recommendation.candidatesource; -import com.linkedin.common.urn.Urn; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.recommendation.RecommendationRenderType; import com.linkedin.metadata.recommendation.RecommendationRequestContext; import com.linkedin.metadata.recommendation.ScenarioType; import com.linkedin.metadata.search.EntitySearchService; +import io.datahubproject.metadata.context.OperationContext; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -36,7 +36,7 @@ public RecommendationRenderType getRenderType() { @Override public boolean isEligible( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { return requestContext.getScenario() == ScenarioType.HOME; } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/EntitySearchAggregationSource.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/EntitySearchAggregationSource.java index 8d6ccb22660fb2..ece3a4e177b3cc 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/EntitySearchAggregationSource.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/EntitySearchAggregationSource.java @@ -13,6 +13,7 @@ import com.linkedin.metadata.recommendation.SearchParams; import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.utils.QueryUtils; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.util.Collections; @@ -71,10 +72,10 @@ protected boolean isValidCandidate(T candidate) { @Override @WithSpan public List getRecommendations( - @Nonnull Urn userUrn, @Nullable RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nullable RecommendationRequestContext requestContext) { Map aggregationResult = entitySearchService.aggregateByValue( - getEntityNames(entityRegistry), getSearchFieldName(), null, getMaxContent()); + opContext, getEntityNames(entityRegistry), getSearchFieldName(), null, getMaxContent()); if (aggregationResult.isEmpty()) { return Collections.emptyList(); diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlySearchedSource.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlySearchedSource.java index e133e3dc75ff39..96b266f88406e6 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlySearchedSource.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlySearchedSource.java @@ -13,6 +13,7 @@ import com.linkedin.metadata.recommendation.SearchParams; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.metrics.MetricUtils; +import io.datahubproject.metadata.context.OperationContext; import java.io.IOException; import java.util.List; import java.util.Optional; @@ -60,7 +61,7 @@ public RecommendationRenderType getRenderType() { @Override public boolean isEligible( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { boolean analyticsEnabled = false; try { analyticsEnabled = @@ -77,8 +78,8 @@ public boolean isEligible( @Override public List getRecommendations( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { - SearchRequest searchRequest = buildSearchRequest(userUrn); + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { + SearchRequest searchRequest = buildSearchRequest(opContext.getActorContext().getActorUrn()); try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getRecentlySearched").time()) { final SearchResponse searchResponse = _searchClient.search(searchRequest, RequestOptions.DEFAULT); diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationSource.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationSource.java index 788ef728e294fb..95c5df64ed2f29 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationSource.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationSource.java @@ -1,11 +1,11 @@ package com.linkedin.metadata.recommendation.candidatesource; -import com.linkedin.common.urn.Urn; import com.linkedin.metadata.recommendation.RecommendationContent; import com.linkedin.metadata.recommendation.RecommendationContentArray; import com.linkedin.metadata.recommendation.RecommendationModule; import com.linkedin.metadata.recommendation.RecommendationRenderType; import com.linkedin.metadata.recommendation.RecommendationRequestContext; +import io.datahubproject.metadata.context.OperationContext; import io.opentelemetry.extension.annotations.WithSpan; import java.util.List; import java.util.Optional; @@ -26,37 +26,38 @@ public interface RecommendationSource { /** * Whether or not this module is eligible for resolution given the context * - * @param userUrn User requesting recommendations + * @param opContext User's context requesting recommendations * @param requestContext Context of where the recommendations are being requested * @return whether this source is eligible */ - boolean isEligible(@Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext); + boolean isEligible( + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext); /** * Get recommended items (candidates / content) provided the context * - * @param userUrn User requesting recommendations + * @param opContext User's context requesting recommendations * @param requestContext Context of where the recommendations are being requested * @return list of recommendation candidates */ @WithSpan List getRecommendations( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext); + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext); /** * Get the full recommendations module itself provided the request context. * - * @param userUrn User requesting recommendations + * @param opContext User's context requesting recommendations * @param requestContext Context of where the recommendations are being requested * @return list of recommendation candidates */ default Optional getRecommendationModule( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { - if (!isEligible(userUrn, requestContext)) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { + if (!isEligible(opContext, requestContext)) { return Optional.empty(); } - List recommendations = getRecommendations(userUrn, requestContext); + List recommendations = getRecommendations(opContext, requestContext); if (recommendations.isEmpty()) { return Optional.empty(); } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationUtils.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationUtils.java index 1fa47d1a13645f..2c1b183b9344b7 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationUtils.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecommendationUtils.java @@ -1,6 +1,6 @@ package com.linkedin.metadata.recommendation.candidatesource; -import com.linkedin.common.urn.Urn; +import io.datahubproject.metadata.context.OperationContext; import java.util.Set; import javax.annotation.Nonnull; @@ -14,8 +14,8 @@ public class RecommendationUtils { * @return true if the type of the urn is in the set of valid entity types, false otherwise. */ public static boolean isSupportedEntityType( - @Nonnull final Urn urn, @Nonnull final Set entityTypes) { - final String entityType = urn.getEntityType(); + @Nonnull OperationContext opContext, @Nonnull final Set entityTypes) { + final String entityType = opContext.getActorContext().getActorUrn().getEntityType(); return entityTypes.contains(entityType); } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopPlatformsSource.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopPlatformsSource.java index 1a5f1ff4b2ca46..f43f5894a09f75 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopPlatformsSource.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopPlatformsSource.java @@ -11,6 +11,7 @@ import com.linkedin.metadata.recommendation.RecommendationRequestContext; import com.linkedin.metadata.recommendation.ScenarioType; import com.linkedin.metadata.search.EntitySearchService; +import io.datahubproject.metadata.context.OperationContext; import java.util.List; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -64,7 +65,7 @@ public RecommendationRenderType getRenderType() { @Override public boolean isEligible( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { return requestContext.getScenario() == ScenarioType.HOME; } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopTagsSource.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopTagsSource.java index 0897d441335fac..1eb00fa8eae303 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopTagsSource.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopTagsSource.java @@ -1,11 +1,11 @@ package com.linkedin.metadata.recommendation.candidatesource; -import com.linkedin.common.urn.Urn; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.recommendation.RecommendationRenderType; import com.linkedin.metadata.recommendation.RecommendationRequestContext; import com.linkedin.metadata.recommendation.ScenarioType; import com.linkedin.metadata.search.EntitySearchService; +import io.datahubproject.metadata.context.OperationContext; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -35,7 +35,7 @@ public RecommendationRenderType getRenderType() { @Override public boolean isEligible( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { return requestContext.getScenario() == ScenarioType.HOME || requestContext.getScenario() == ScenarioType.SEARCH_RESULTS; } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopTermsSource.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopTermsSource.java index 0fab9a28b51ea4..2909a56d263689 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopTermsSource.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/candidatesource/TopTermsSource.java @@ -1,11 +1,11 @@ package com.linkedin.metadata.recommendation.candidatesource; -import com.linkedin.common.urn.Urn; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.recommendation.RecommendationRenderType; import com.linkedin.metadata.recommendation.RecommendationRequestContext; import com.linkedin.metadata.recommendation.ScenarioType; import com.linkedin.metadata.search.EntitySearchService; +import io.datahubproject.metadata.context.OperationContext; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -35,7 +35,7 @@ public RecommendationRenderType getRenderType() { @Override public boolean isEligible( - @Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) { + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) { return requestContext.getScenario() == ScenarioType.HOME || requestContext.getScenario() == ScenarioType.SEARCH_RESULTS; } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/ranker/RecommendationModuleRanker.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/ranker/RecommendationModuleRanker.java index f09f83fd6ec259..ab736cf4b25217 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/ranker/RecommendationModuleRanker.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/ranker/RecommendationModuleRanker.java @@ -1,8 +1,8 @@ package com.linkedin.metadata.recommendation.ranker; -import com.linkedin.common.urn.Urn; import com.linkedin.metadata.recommendation.RecommendationModule; import com.linkedin.metadata.recommendation.RecommendationRequestContext; +import io.datahubproject.metadata.context.OperationContext; import java.util.List; import javax.annotation.Nonnull; @@ -10,15 +10,15 @@ public interface RecommendationModuleRanker { /** * Rank and return the final list of modules * + * @param opContext the user's context * @param candidates Candidate modules to rank - * @param userUrn User requesting recommendations * @param requestContext Context of where the recommendations are being requested * @param limit Max number of modules to return * @return ranked list of modules */ List rank( - @Nonnull List candidates, - @Nonnull Urn userUrn, + @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext, + @Nonnull List candidates, int limit); } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/ranker/SimpleRecommendationRanker.java b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/ranker/SimpleRecommendationRanker.java index 13bc5af91c9e9c..4599b42d88e4b7 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/ranker/SimpleRecommendationRanker.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/recommendation/ranker/SimpleRecommendationRanker.java @@ -1,8 +1,8 @@ package com.linkedin.metadata.recommendation.ranker; -import com.linkedin.common.urn.Urn; import com.linkedin.metadata.recommendation.RecommendationModule; import com.linkedin.metadata.recommendation.RecommendationRequestContext; +import io.datahubproject.metadata.context.OperationContext; import java.util.List; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -10,9 +10,9 @@ public class SimpleRecommendationRanker implements RecommendationModuleRanker { @Override public List rank( - @Nonnull List candidates, - @Nonnull Urn userUrn, + @Nonnull OperationContext opContext, @Nullable RecommendationRequestContext requestContext, + @Nonnull List candidates, int limit) { return candidates.subList(0, Math.min(candidates.size(), limit)); } 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 03b5c7f5547e7b..01b109e7cd924f 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 @@ -5,9 +5,9 @@ import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.browse.BrowseResultV2; import com.linkedin.metadata.query.AutoCompleteResult; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; +import io.datahubproject.metadata.context.OperationContext; import java.util.List; import java.util.Map; import javax.annotation.Nonnull; @@ -33,7 +33,7 @@ public interface EntitySearchService { * * @param entityName name of the entity */ - long docCount(@Nonnull String entityName); + long docCount(@Nonnull OperationContext opContext, @Nonnull String entityName); /** * Updates or inserts the given search document. @@ -76,19 +76,18 @@ public interface EntitySearchService { * @param sortCriterion {@link SortCriterion} to be applied to search results * @param from index to start the search from * @param size the number of search hits to return - * @param searchFlags flags controlling search options * @return a {@link SearchResult} that contains a list of matched documents and related search * result metadata */ @Nonnull SearchResult search( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, - int size, - @Nullable SearchFlags searchFlags); + int size); /** * Gets a list of documents that match given search request. The results are aggregated and @@ -105,20 +104,19 @@ SearchResult search( * @param sortCriterion {@link SortCriterion} to be applied to search results * @param from index to start the search from * @param size the number of search hits to return - * @param searchFlags flags controlling search options * @param facets list of facets we want aggregations for * @return a {@link SearchResult} that contains a list of matched documents and related search * result metadata */ @Nonnull SearchResult search( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, int size, - @Nullable SearchFlags searchFlags, @Nullable List facets); /** @@ -135,6 +133,7 @@ SearchResult search( */ @Nonnull SearchResult filter( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nullable Filter filters, @Nullable SortCriterion sortCriterion, @@ -156,6 +155,7 @@ SearchResult filter( */ @Nonnull AutoCompleteResult autoComplete( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String query, @Nullable String field, @@ -174,6 +174,7 @@ AutoCompleteResult autoComplete( */ @Nonnull Map aggregateByValue( + @Nonnull OperationContext opContext, @Nullable List entityNames, @Nonnull String field, @Nullable Filter requestParams, @@ -191,6 +192,7 @@ Map aggregateByValue( */ @Nonnull BrowseResult browse( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter requestParams, @@ -206,17 +208,16 @@ BrowseResult browse( * @param input search query * @param start start offset of first group * @param count max number of results requested - * @param searchFlags configuration options for search */ @Nonnull - public BrowseResultV2 browseV2( + BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull String entityName, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nullable SearchFlags searchFlags); + int count); /** * Gets browse snapshot of a given path @@ -227,17 +228,16 @@ public BrowseResultV2 browseV2( * @param input search query * @param start start offset of first group * @param count max number of results requested - * @param searchFlags configuration options for search */ @Nonnull - public BrowseResultV2 browseV2( + BrowseResultV2 browseV2( + @Nonnull OperationContext opContext, @Nonnull List entityNames, @Nonnull String path, @Nullable Filter filter, @Nonnull String input, int start, - int count, - @Nullable SearchFlags searchFlags); + int count); /** * Gets a list of paths for a given urn. @@ -260,20 +260,19 @@ public BrowseResultV2 browseV2( * @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 OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, - int size, - @Nullable SearchFlags searchFlags); + int size); /** * Gets a list of documents that match given search request. The results are aggregated and @@ -286,31 +285,30 @@ ScrollResult fullTextScroll( * @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 OperationContext opContext, @Nonnull List entities, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, @Nullable String scrollId, @Nullable String keepAlive, - int size, - @Nullable SearchFlags searchFlags); + int size); /** Max result size returned by the underlying search backend */ int maxResultSize(); ExplainResponse explain( + @Nonnull OperationContext opContext, @Nonnull String query, @Nonnull String documentId, @Nonnull String entityName, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, - @Nullable SearchFlags searchFlags, @Nullable String scrollId, @Nullable String keepAlive, int size, diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/service/FormService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/service/FormService.java index 59d40b29e7383d..f8b9aa2ea6d64e 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/service/FormService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/service/FormService.java @@ -44,6 +44,7 @@ import com.linkedin.structured.StructuredProperties; import com.linkedin.structured.StructuredPropertyValueAssignment; import com.linkedin.structured.StructuredPropertyValueAssignmentArray; +import io.datahubproject.metadata.context.OperationContext; import java.net.URISyntaxException; import java.util.ArrayList; import java.util.List; @@ -67,10 +68,12 @@ public class FormService extends BaseService { private static final int BATCH_FORM_ENTITY_COUNT = 500; + private final OperationContext systemOpContext; + public FormService( - @Nonnull final EntityClient entityClient, - @Nonnull final Authentication systemAuthentication) { - super(entityClient, systemAuthentication); + @Nonnull OperationContext systemOpContext, @Nonnull final EntityClient entityClient) { + super(entityClient, systemOpContext.getAuthentication()); + this.systemOpContext = systemOpContext; } /** Batch associated a form to a given set of entities by urn. */ @@ -162,7 +165,7 @@ public void upsertFormAssignmentRunner( @Nonnull final Urn formUrn, @Nonnull final DynamicFormAssignment formFilters) { try { SearchBasedFormAssignmentRunner.assign( - formFilters, formUrn, BATCH_FORM_ENTITY_COUNT, entityClient, systemAuthentication); + systemOpContext, formFilters, formUrn, BATCH_FORM_ENTITY_COUNT, entityClient); } catch (Exception e) { throw new RuntimeException( String.format("Failed to dynamically assign form with urn: %s", formUrn), e); diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/service/LineageService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/service/LineageService.java index cd5202ce75b648..7cbdde5959b36f 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/service/LineageService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/service/LineageService.java @@ -1,6 +1,7 @@ package com.linkedin.metadata.service; import static com.linkedin.metadata.entity.AspectUtils.*; +import static com.linkedin.metadata.utils.AuditStampUtils.getAuditStamp; import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableSet; diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/service/RestrictedService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/service/RestrictedService.java new file mode 100644 index 00000000000000..d1aa8e9f5dbb5c --- /dev/null +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/service/RestrictedService.java @@ -0,0 +1,29 @@ +package com.linkedin.metadata.service; + +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.metadata.secret.SecretService; +import javax.annotation.Nonnull; + +public class RestrictedService { + + private final SecretService _secretService; + + public RestrictedService(@Nonnull SecretService secretService) { + this._secretService = secretService; + } + + public Urn encryptRestrictedUrn(@Nonnull final Urn entityUrn) { + final String encryptedEntityUrn = this._secretService.encrypt(entityUrn.toString()); + try { + return new Urn("restricted", encryptedEntityUrn); + } catch (Exception e) { + throw new RuntimeException("Error when creating restricted entity urn", e); + } + } + + public Urn decryptRestrictedUrn(@Nonnull final Urn restrictedUrn) { + final String encryptedUrn = restrictedUrn.getId(); + return UrnUtils.getUrn(this._secretService.decrypt(encryptedUrn)); + } +} diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/service/util/SearchBasedFormAssignmentManager.java b/metadata-service/services/src/main/java/com/linkedin/metadata/service/util/SearchBasedFormAssignmentManager.java index 73e3bc130ac9d4..9c29e2a41a633f 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/service/util/SearchBasedFormAssignmentManager.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/service/util/SearchBasedFormAssignmentManager.java @@ -1,6 +1,5 @@ package com.linkedin.metadata.service.util; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; import com.linkedin.entity.client.EntityClient; @@ -10,6 +9,7 @@ import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.service.FormService; import com.linkedin.r2.RemoteInvocationException; +import io.datahubproject.metadata.context.OperationContext; import java.util.List; import java.util.stream.Collectors; import lombok.extern.slf4j.Slf4j; @@ -21,31 +21,30 @@ public class SearchBasedFormAssignmentManager { ImmutableList.of(Constants.DATASET_ENTITY_NAME); public static void apply( + OperationContext opContext, DynamicFormAssignment formFilters, Urn formUrn, int batchFormEntityCount, - EntityClient entityClient, - Authentication authentication) + EntityClient entityClient) throws Exception { try { int totalResults = 0; int numResults = 0; String scrollId = null; - FormService formService = new FormService(entityClient, authentication); + FormService formService = new FormService(opContext, entityClient); do { ScrollResult results = entityClient.scrollAcrossEntities( + opContext, ENTITY_TYPES, "*", formFilters.getFilter(), scrollId, "5m", - batchFormEntityCount, - null, - authentication); + batchFormEntityCount); if (!results.hasEntities() || results.getNumEntities() == 0 diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/service/util/SearchBasedFormAssignmentRunner.java b/metadata-service/services/src/main/java/com/linkedin/metadata/service/util/SearchBasedFormAssignmentRunner.java index a20f71f550c65d..f24307c30793de 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/service/util/SearchBasedFormAssignmentRunner.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/service/util/SearchBasedFormAssignmentRunner.java @@ -1,27 +1,27 @@ package com.linkedin.metadata.service.util; -import com.datahub.authentication.Authentication; import com.linkedin.common.urn.Urn; import com.linkedin.entity.client.EntityClient; import com.linkedin.form.DynamicFormAssignment; +import io.datahubproject.metadata.context.OperationContext; import lombok.extern.slf4j.Slf4j; @Slf4j public class SearchBasedFormAssignmentRunner { public static void assign( + OperationContext opContext, DynamicFormAssignment formFilters, Urn formUrn, int batchFormEntityCount, - EntityClient entityClient, - Authentication authentication) { + EntityClient entityClient) { Runnable runnable = new Runnable() { @Override public void run() { try { SearchBasedFormAssignmentManager.apply( - formFilters, formUrn, batchFormEntityCount, entityClient, authentication); + opContext, formFilters, formUrn, batchFormEntityCount, entityClient); } catch (Exception e) { log.error( "SearchBasedFormAssignmentRunner failed to run. " diff --git a/metadata-service/services/src/test/java/com/linkedin/metadata/service/RestrictedServiceTest.java b/metadata-service/services/src/test/java/com/linkedin/metadata/service/RestrictedServiceTest.java new file mode 100644 index 00000000000000..239e56cfd97a8b --- /dev/null +++ b/metadata-service/services/src/test/java/com/linkedin/metadata/service/RestrictedServiceTest.java @@ -0,0 +1,37 @@ +package com.linkedin.metadata.service; + +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.metadata.secret.SecretService; +import org.mockito.Mockito; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class RestrictedServiceTest { + + private static final Urn TEST_DATASET_URN = + UrnUtils.getUrn("urn:li:dataset:(urn:li:dataPlatform:foo,bar,baz1)"); + private static final String ENCRYPED_DATASET_URN = "12d3as456tgs"; + private static final Urn TEST_RESTRICTED_URN = + UrnUtils.getUrn(String.format("urn:li:restricted:%s", ENCRYPED_DATASET_URN)); + + @Test + private void testEncryptRestrictedUrn() throws Exception { + SecretService mockSecretService = Mockito.mock(SecretService.class); + Mockito.when(mockSecretService.encrypt(TEST_DATASET_URN.toString())) + .thenReturn(ENCRYPED_DATASET_URN); + final RestrictedService service = new RestrictedService(mockSecretService); + + Assert.assertEquals(service.encryptRestrictedUrn(TEST_DATASET_URN), TEST_RESTRICTED_URN); + } + + @Test + private void testDecryptRestrictedUrn() throws Exception { + SecretService mockSecretService = Mockito.mock(SecretService.class); + Mockito.when(mockSecretService.decrypt(ENCRYPED_DATASET_URN)) + .thenReturn(TEST_DATASET_URN.toString()); + final RestrictedService service = new RestrictedService(mockSecretService); + + Assert.assertEquals(service.decryptRestrictedUrn(TEST_RESTRICTED_URN), TEST_DATASET_URN); + } +} diff --git a/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java b/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java index afaeb9c81039bf..ffa0d600a2351d 100644 --- a/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java +++ b/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java @@ -10,8 +10,8 @@ import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.elasticsearch.query.request.SearchRequestHandler; +import io.datahubproject.metadata.context.OperationContext; import jakarta.servlet.http.HttpServlet; import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletResponse; @@ -43,6 +43,10 @@ private AspectRetriever getAspectRetriever(WebApplicationContext ctx) { return (AspectRetriever) ctx.getBean("aspectRetriever"); } + private OperationContext getOperationContext(WebApplicationContext ctx) { + return (OperationContext) ctx.getBean("systemOperationContext"); + } + private void writeSearchCsv(WebApplicationContext ctx, PrintWriter pw) { SearchConfiguration searchConfiguration = getConfigProvider(ctx).getElasticSearch().getSearch(); AspectRetriever aspectRetriever = getAspectRetriever(ctx); @@ -84,15 +88,18 @@ private void writeSearchCsv(WebApplicationContext ctx, PrintWriter pw) { SearchRequestHandler.getBuilder( entitySpec, searchConfiguration, null, aspectRetriever) .getSearchRequest( + getOperationContext(ctx) + .withSearchFlags( + flags -> + flags + .setFulltext(true) + .setSkipHighlighting(true) + .setSkipAggregates(true)), "*", null, null, 0, 0, - new SearchFlags() - .setFulltext(true) - .setSkipHighlighting(true) - .setSkipAggregates(true), null); FunctionScoreQueryBuilder rankingQuery = diff --git a/metadata-service/war/src/main/resources/boot/policies.json b/metadata-service/war/src/main/resources/boot/policies.json index 66a5de48790c27..b89ee970c875f3 100644 --- a/metadata-service/war/src/main/resources/boot/policies.json +++ b/metadata-service/war/src/main/resources/boot/policies.json @@ -55,6 +55,7 @@ "privileges":[ "EDIT_ENTITY", "VIEW_ENTITY_PAGE", + "VIEW_ENTITY", "EDIT_LINEAGE", "EDIT_ENTITY_ASSERTIONS", "SEARCH_PRIVILEGE", @@ -106,6 +107,7 @@ }, "privileges":[ "VIEW_ENTITY_PAGE", + "VIEW_ENTITY", "SEARCH_PRIVILEGE", "GET_COUNTS_PRIVILEGE", "GET_TIMESERIES_ASPECT_PRIVILEGE", @@ -199,6 +201,7 @@ }, "privileges":[ "VIEW_ENTITY_PAGE", + "VIEW_ENTITY", "EDIT_ENTITY_TAGS", "EDIT_ENTITY_GLOSSARY_TERMS", "EDIT_ENTITY_OWNERS", @@ -281,6 +284,7 @@ }, "privileges":[ "VIEW_ENTITY_PAGE", + "VIEW_ENTITY", "EDIT_ENTITY_TAGS", "EDIT_ENTITY_GLOSSARY_TERMS", "EDIT_ENTITY_DOCS", @@ -399,6 +403,7 @@ }, "privileges":[ "VIEW_ENTITY_PAGE", + "VIEW_ENTITY", "VIEW_DATASET_USAGE", "VIEW_DATASET_PROFILE", "SEARCH_PRIVILEGE", @@ -427,6 +432,7 @@ }, "privileges":[ "VIEW_ENTITY_PAGE", + "VIEW_ENTITY", "EDIT_ENTITY_TAGS", "EDIT_ENTITY_GLOSSARY_TERMS", "EDIT_ENTITY_OWNERS", diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/authorization/PoliciesConfig.java b/metadata-utils/src/main/java/com/linkedin/metadata/authorization/PoliciesConfig.java index 6889d56d4ebc27..d6d2d24109874d 100644 --- a/metadata-utils/src/main/java/com/linkedin/metadata/authorization/PoliciesConfig.java +++ b/metadata-utils/src/main/java/com/linkedin/metadata/authorization/PoliciesConfig.java @@ -3,6 +3,7 @@ import com.google.common.collect.ImmutableList; import java.util.Collection; import java.util.List; +import java.util.Set; import java.util.stream.Collectors; import java.util.stream.Stream; import lombok.AllArgsConstructor; @@ -139,6 +140,19 @@ public class PoliciesConfig { public static final Privilege VIEW_ENTITY_PAGE_PRIVILEGE = Privilege.of("VIEW_ENTITY_PAGE", "View Entity Page", "The ability to view the entity page."); + public static final Privilege VIEW_ENTITY_PRIVILEGE = + Privilege.of( + "VIEW_ENTITY", "View Entity", "The ability to view the entity in search results."); + + /* + These two privileges are logically the same for search for now. + In the future, we might allow search but not the entity page view. + */ + public static final Set VIEW_ENTITY_PRIVILEGES = + Set.of( + PoliciesConfig.VIEW_ENTITY_PRIVILEGE.getType(), + PoliciesConfig.VIEW_ENTITY_PAGE_PRIVILEGE.getType()); + public static final Privilege EDIT_ENTITY_TAGS_PRIVILEGE = Privilege.of( "EDIT_ENTITY_TAGS", "Edit Tags", "The ability to add and remove tags to an asset."); @@ -242,6 +256,7 @@ public class PoliciesConfig { EDIT_ENTITY_DEPRECATION_PRIVILEGE, EDIT_ENTITY_PRIVILEGE, DELETE_ENTITY_PRIVILEGE, + VIEW_ENTITY_PRIVILEGE, EDIT_ENTITY_INCIDENTS_PRIVILEGE); // Dataset Privileges diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/utils/AuditStampUtils.java b/metadata-utils/src/main/java/com/linkedin/metadata/utils/AuditStampUtils.java index 6ba311cf166d4e..404dd001353a69 100644 --- a/metadata-utils/src/main/java/com/linkedin/metadata/utils/AuditStampUtils.java +++ b/metadata-utils/src/main/java/com/linkedin/metadata/utils/AuditStampUtils.java @@ -5,25 +5,31 @@ import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; -import java.net.URISyntaxException; -import java.time.Clock; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; +import org.joda.time.DateTimeUtils; @Slf4j public class AuditStampUtils { private AuditStampUtils() {} public static AuditStamp createDefaultAuditStamp() { - return new AuditStamp() - .setActor(UrnUtils.getUrn(SYSTEM_ACTOR)) - .setTime(Clock.systemUTC().millis()); + return getAuditStamp(UrnUtils.getUrn(SYSTEM_ACTOR)); } - public static AuditStamp createAuditStamp(@Nonnull String actorUrn) throws URISyntaxException { + public static AuditStamp createAuditStamp(@Nonnull String actorUrn) { + return getAuditStamp(UrnUtils.getUrn(actorUrn)); + } + + public static AuditStamp getAuditStamp(Urn actor) { + return getAuditStamp(actor, null); + } + + public static AuditStamp getAuditStamp(@Nonnull Urn actor, @Nullable Long currentTimeMs) { AuditStamp auditStamp = new AuditStamp(); - auditStamp.setActor(Urn.createFromString(actorUrn)); - auditStamp.setTime(Clock.systemUTC().millis()); + auditStamp.setTime(currentTimeMs != null ? currentTimeMs : DateTimeUtils.currentTimeMillis()); + auditStamp.setActor(actor); return auditStamp; } } diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/utils/SearchUtil.java b/metadata-utils/src/main/java/com/linkedin/metadata/utils/SearchUtil.java index 9df708c6e9fdcd..c3c9cac6280ed2 100644 --- a/metadata-utils/src/main/java/com/linkedin/metadata/utils/SearchUtil.java +++ b/metadata-utils/src/main/java/com/linkedin/metadata/utils/SearchUtil.java @@ -23,8 +23,6 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; -import org.opensearch.index.query.BoolQueryBuilder; -import org.opensearch.index.query.QueryBuilders; @Slf4j public class SearchUtil { @@ -33,9 +31,9 @@ public class SearchUtil { public static final String AGGREGATION_SPECIAL_TYPE_DELIMITER = "␝"; public static final String MISSING_SPECIAL_TYPE = "missing"; public static final String INDEX_VIRTUAL_FIELD = "_entityType"; + public static final String ES_INDEX_FIELD = "_index"; public static final String KEYWORD_SUFFIX = ".keyword"; private static final String URN_PREFIX = "urn:"; - private static final String REMOVED = "removed"; private SearchUtil() {} @@ -73,7 +71,7 @@ public static FilterValue createFilterValue(String value, Long facetCount, Boole private static Criterion transformEntityTypeCriterion( Criterion criterion, IndexConvention indexConvention) { return criterion - .setField("_index") + .setField(ES_INDEX_FIELD) .setValues( new StringArray( criterion.getValues().stream() @@ -124,30 +122,6 @@ public static Filter transformFilterForEntities( return filter; } - /** - * Applies a default filter to remove entities that are soft deleted only if there isn't a filter - * for the REMOVED field already - */ - public static BoolQueryBuilder filterSoftDeletedByDefault( - @Nullable Filter filter, @Nullable BoolQueryBuilder filterQuery) { - boolean removedInOrFilter = false; - if (filter != null) { - removedInOrFilter = - filter.getOr().stream() - .anyMatch( - or -> - or.getAnd().stream() - .anyMatch( - criterion -> - criterion.getField().equals(REMOVED) - || criterion.getField().equals(REMOVED + KEYWORD_SUFFIX))); - } - if (!removedInOrFilter) { - filterQuery.mustNot(QueryBuilders.matchQuery(REMOVED, true)); - } - return filterQuery; - } - public static SortCriterion sortBy(@Nonnull String field, @Nullable SortOrder direction) { SortCriterion sortCriterion = new SortCriterion(); sortCriterion.setField(field); diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/utils/elasticsearch/IndexConventionImpl.java b/metadata-utils/src/main/java/com/linkedin/metadata/utils/elasticsearch/IndexConventionImpl.java index 764630eb739733..47801cd2054fa4 100644 --- a/metadata-utils/src/main/java/com/linkedin/metadata/utils/elasticsearch/IndexConventionImpl.java +++ b/metadata-utils/src/main/java/com/linkedin/metadata/utils/elasticsearch/IndexConventionImpl.java @@ -12,6 +12,8 @@ // Default implementation of search index naming convention public class IndexConventionImpl implements IndexConvention { + public static final IndexConvention NO_PREFIX = new IndexConventionImpl(null); + // Map from Entity name -> Index name private final Map indexNameMapping = new ConcurrentHashMap<>(); private final Optional _prefix; diff --git a/metadata-utils/src/test/java/com/linkedin/metadata/utils/elasticsearch/IndexConventionImplTest.java b/metadata-utils/src/test/java/com/linkedin/metadata/utils/elasticsearch/IndexConventionImplTest.java index f3e52c99897751..8074f344cd2441 100644 --- a/metadata-utils/src/test/java/com/linkedin/metadata/utils/elasticsearch/IndexConventionImplTest.java +++ b/metadata-utils/src/test/java/com/linkedin/metadata/utils/elasticsearch/IndexConventionImplTest.java @@ -10,7 +10,7 @@ public class IndexConventionImplTest { @Test public void testIndexConventionNoPrefix() { - IndexConvention indexConventionNoPrefix = new IndexConventionImpl(null); + IndexConvention indexConventionNoPrefix = IndexConventionImpl.NO_PREFIX; String entityName = "dataset"; String expectedIndexName = "datasetindex_v2"; assertEquals(indexConventionNoPrefix.getEntityIndexName(entityName), expectedIndexName); @@ -42,7 +42,7 @@ public void testIndexConventionPrefix() { @Test public void testTimeseriesIndexConventionNoPrefix() { - IndexConvention indexConventionNoPrefix = new IndexConventionImpl(null); + IndexConvention indexConventionNoPrefix = IndexConventionImpl.NO_PREFIX; String entityName = "dataset"; String aspectName = "datasetusagestatistics"; String expectedIndexName = "dataset_datasetusagestatisticsaspect_v1"; diff --git a/settings.gradle b/settings.gradle index 4614b6ed4ccaf0..e58c1c851c8f12 100644 --- a/settings.gradle +++ b/settings.gradle @@ -64,3 +64,4 @@ include 'mock-entity-registry' include 'metadata-service:services' include 'metadata-service:configuration' include ':metadata-jobs:common' +include ':metadata-operation-context' diff --git a/smoke-test/cypress-dev.sh b/smoke-test/cypress-dev.sh index b1c6571e1a0658..0359c6eb547e62 100755 --- a/smoke-test/cypress-dev.sh +++ b/smoke-test/cypress-dev.sh @@ -10,6 +10,8 @@ fi source venv/bin/activate +export KAFKA_BROKER_CONTAINER="datahub-kafka-broker-1" +export KAFKA_BOOTSTRAP_SERVER="broker:9092" python -c 'from tests.cypress.integration_test import ingest_data; ingest_data()' cd tests/cypress 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 feabcc5f9d6558..d5e1ade663dffd 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 @@ -106,6 +106,11 @@ def test_group_upsert(wait_for_healthchecks: Any) -> None: "owners": [ {"owner": "urn:li:corpuser:user1", "type": "TECHNICAL_OWNER"} ], + "ownerTypes": { + "urn:li:ownershipType:__system__none": [ + "urn:li:corpuser:user1", + ], + }, }, "status": {"removed": False}, } diff --git a/smoke-test/tests/consistency_utils.py b/smoke-test/tests/consistency_utils.py index 5ffc642d494697..4335e2a874c1e7 100644 --- a/smoke-test/tests/consistency_utils.py +++ b/smoke-test/tests/consistency_utils.py @@ -8,6 +8,10 @@ ELASTICSEARCH_REFRESH_INTERVAL_SECONDS: int = int( os.getenv("ELASTICSEARCH_REFRESH_INTERVAL_SECONDS", 5) ) +KAFKA_BROKER_CONTAINER: str = str( + os.getenv("KAFKA_BROKER_CONTAINER", "datahub-broker-1") +) +KAFKA_BOOTSTRAP_SERVER: str = str(os.getenv("KAFKA_BOOTSTRAP_SERVER", "broker:29092")) logger = logging.getLogger(__name__) @@ -21,19 +25,26 @@ def wait_for_writes_to_sync(max_timeout_in_sec: int = 120) -> None: lag_zero = False while not lag_zero and (time.time() - start_time) < max_timeout_in_sec: time.sleep(1) # micro-sleep - completed_process = subprocess.run( - "docker exec datahub-broker-1 /bin/kafka-consumer-groups --bootstrap-server broker:29092 --group generic-mae-consumer-job-client --describe | grep -v LAG | awk '{print $6}'", - capture_output=True, - shell=True, - text=True, - ) - result = str(completed_process.stdout) - lines = result.splitlines() - lag_values = [int(line) for line in lines if line != ""] - maximum_lag = max(lag_values) if lag_values else 0 - if maximum_lag == 0: - lag_zero = True + cmd = ( + f"docker exec {KAFKA_BROKER_CONTAINER} /bin/kafka-consumer-groups --bootstrap-server {KAFKA_BOOTSTRAP_SERVER} --group generic-mae-consumer-job-client --describe | grep -v LAG " + + "| awk '{print $6}'" + ) + try: + completed_process = subprocess.run( + cmd, + capture_output=True, + shell=True, + text=True, + ) + result = str(completed_process.stdout) + lines = result.splitlines() + lag_values = [int(line) for line in lines if line != ""] + maximum_lag = max(lag_values) + if maximum_lag == 0: + lag_zero = True + except ValueError: + logger.warning(f"Error reading kafka lag using command: {cmd}") if not lag_zero: logger.warning( From 29274c1be43f64439b7aa1ced856bd8e79765eaf Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Thu, 29 Feb 2024 12:11:57 -0800 Subject: [PATCH 09/13] feat(ingest/sql-parser): add alias for mariadb (#9956) --- metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py b/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py index 710fffc4afbd30..0dffc4291132b4 100644 --- a/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py +++ b/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py @@ -16,12 +16,13 @@ def _get_dialect_str(platform: str) -> str: return "tsql" elif platform == "athena": return "trino" - elif platform == "mysql": + elif platform in {"mysql", "mariadb"}: # In sqlglot v20+, MySQL is now case-sensitive by default, which is the # default behavior on Linux. However, MySQL's default case sensitivity # actually depends on the underlying OS. # For us, it's simpler to just assume that it's case-insensitive, and # let the fuzzy resolution logic handle it. + # MariaDB is a fork of MySQL, so we reuse the same dialect. return "mysql, normalization_strategy = lowercase" else: return platform From 0c7ca4cd3e27ee12823082465a737b82346983f2 Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Thu, 29 Feb 2024 12:39:32 -0800 Subject: [PATCH 10/13] docs(ingest/lookml): update known discrepancy list (#9941) Co-authored-by: Tamas Nemeth --- docs-website/build.gradle | 24 +++++++++++-------- .../airflow-plugin/build.gradle | 2 +- metadata-ingestion/build.gradle | 4 +++- .../docs/sources/looker/lookml_post.md | 2 +- 4 files changed, 19 insertions(+), 13 deletions(-) diff --git a/docs-website/build.gradle b/docs-website/build.gradle index 2644491a2a5f80..a096d9cf4971b8 100644 --- a/docs-website/build.gradle +++ b/docs-website/build.gradle @@ -45,7 +45,7 @@ def projectMdFiles = project.fileTree("${project.projectDir}") { include '**/*.ts' exclude 'node_modules' exclude '**/dist/**' - } +} // Combine GraphQL schemas for documentation. task generateGraphQLSchema(type: Exec) { @@ -68,6 +68,16 @@ task yarnInstall(type: YarnTask) { } else { args = ['install'] } + + // The node_modules directory can contain built artifacts, so + // it's not really safe to cache it. + outputs.cacheIf { false } + + inputs.files( + file('yarn.lock'), + file('package.json'), + ) + outputs.dir('node_modules') } task yarnGenerate(type: YarnTask, dependsOn: [yarnInstall, @@ -94,17 +104,11 @@ task fastReload(type: YarnTask) { task yarnLint(type: YarnTask, dependsOn: [yarnInstall, yarnGenerate]) { inputs.files(projectMdFiles) args = ['run', 'lint-check'] - outputs.dir("dist") - // tell gradle to apply the build cache - outputs.cacheIf { true } } task yarnLintFix(type: YarnTask, dependsOn: [yarnInstall]) { inputs.files(projectMdFiles) args = ['run', 'lint-fix'] - outputs.dir("dist") - // tell gradle to apply the build cache - outputs.cacheIf { true } } task serve(type: YarnTask, dependsOn: [yarnInstall] ) { @@ -123,11 +127,11 @@ task yarnBuild(type: YarnTask, dependsOn: [yarnLint, yarnGenerate, downloadHisto outputs.cacheIf { true } // See https://stackoverflow.com/questions/53230823/fatal-error-ineffective-mark-compacts-near-heap-limit-allocation-failed-java // and https://github.com/facebook/docusaurus/issues/8329. - // TODO: As suggested in https://github.com/facebook/docusaurus/issues/4765, try switching to swc-loader. + // TODO: As suggested in https://github.com/facebook/docusaurus/issues/4765, try switching to swc-loader or esbuild minification. if (project.hasProperty('useSystemNode') && project.getProperty('useSystemNode').toBoolean()) { - environment = ['NODE_OPTIONS': '--max-old-space-size=10248'] + environment = ['NODE_OPTIONS': '--max-old-space-size=14336'] } else { - environment = ['NODE_OPTIONS': '--max-old-space-size=10248 --openssl-legacy-provider'] + environment = ['NODE_OPTIONS': '--max-old-space-size=14336 --openssl-legacy-provider'] } args = ['run', 'build'] diff --git a/metadata-ingestion-modules/airflow-plugin/build.gradle b/metadata-ingestion-modules/airflow-plugin/build.gradle index 554da819af41d4..350c43b2568084 100644 --- a/metadata-ingestion-modules/airflow-plugin/build.gradle +++ b/metadata-ingestion-modules/airflow-plugin/build.gradle @@ -109,7 +109,7 @@ task cleanPythonCache(type: Exec) { task buildWheel(type: Exec, dependsOn: [install, cleanPythonCache]) { commandLine 'bash', '-c', "source ${venv_name}/bin/activate && set -x && " + - 'pip install build && RELEASE_VERSION="\${RELEASE_VERSION:-0.0.0.dev1}" RELEASE_SKIP_TEST=1 RELEASE_SKIP_UPLOAD=1 ./scripts/release.sh' + 'uv pip install build && RELEASE_VERSION="\${RELEASE_VERSION:-0.0.0.dev1}" RELEASE_SKIP_TEST=1 RELEASE_SKIP_INSTALL=1 RELEASE_SKIP_UPLOAD=1 ./scripts/release.sh' } build.dependsOn install diff --git a/metadata-ingestion/build.gradle b/metadata-ingestion/build.gradle index 269424d3de72b2..81031575e9bbc6 100644 --- a/metadata-ingestion/build.gradle +++ b/metadata-ingestion/build.gradle @@ -212,7 +212,9 @@ task cleanPythonCache(type: Exec) { "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 && " + 'uv pip install build && RELEASE_VERSION="\${RELEASE_VERSION:-0.0.0.dev1}" RELEASE_SKIP_TEST=1 RELEASE_SKIP_UPLOAD=1 ./scripts/release.sh' + commandLine 'bash', '-c', + "source ${venv_name}/bin/activate && " + + 'uv pip install build && RELEASE_VERSION="\${RELEASE_VERSION:-0.0.0.dev1}" RELEASE_SKIP_TEST=1 RELEASE_SKIP_INSTALL=1 RELEASE_SKIP_UPLOAD=1 ./scripts/release.sh' } build.dependsOn install diff --git a/metadata-ingestion/docs/sources/looker/lookml_post.md b/metadata-ingestion/docs/sources/looker/lookml_post.md index 8ebbab4b9ed48d..773a917202f247 100644 --- a/metadata-ingestion/docs/sources/looker/lookml_post.md +++ b/metadata-ingestion/docs/sources/looker/lookml_post.md @@ -84,7 +84,7 @@ If you see messages like `my_file.view.lkml': "failed to load view file: Unable The first thing to check is that the Looker IDE can validate the file without issues. You can check this by clicking this "Validate LookML" button in the IDE when in development mode. If that's not the issue, it might be because DataHub's parser, which is based on the [joshtemple/lkml](https://github.com/joshtemple/lkml) library, is slightly more strict than the official Looker parser. -Note that there's currently only one known discrepancy between the two parsers, and it's related to using [multiple colon characters](https://github.com/joshtemple/lkml/issues/82) when defining parameters. +Note that there's currently only one known discrepancy between the two parsers, and it's related to using [leading colons in blocks](https://github.com/joshtemple/lkml/issues/90). To check if DataHub can parse your LookML file syntax, you can use the `lkml` CLI tool. If this raises an exception, DataHub will fail to parse the file. From 9a7c0976f2e44d93c1a566a39340bb0879c09737 Mon Sep 17 00:00:00 2001 From: pankajmahato-visa <154867659+pankajmahato-visa@users.noreply.github.com> Date: Fri, 1 Mar 2024 03:07:21 +0530 Subject: [PATCH 11/13] chore(vulnerability): Bumped up versions for vulnerability fix (#9929) --- build.gradle | 29 ++++++++++++++++------------- buildSrc/build.gradle | 2 +- 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/build.gradle b/build.gradle index 228a8a9f5ff0a0..cbee21fe1d4c1b 100644 --- a/build.gradle +++ b/build.gradle @@ -47,7 +47,7 @@ buildscript { ext.log4jVersion = '2.19.0' ext.slf4jVersion = '1.7.36' ext.logbackClassic = '1.4.14' - ext.hadoop3Version = '3.3.5' + ext.hadoop3Version = '3.3.6' ext.kafkaVersion = '2.3.0' ext.hazelcastVersion = '5.3.6' ext.ebeanVersion = '12.16.1' @@ -99,8 +99,8 @@ project.ext.spec = [ project.ext.externalDependency = [ 'akkaHttp': 'com.typesafe.akka:akka-http-core_2.12:10.2.10', - 'antlr4Runtime': 'org.antlr:antlr4-runtime:4.7.2', - 'antlr4': 'org.antlr:antlr4:4.7.2', + 'antlr4Runtime': 'org.antlr:antlr4-runtime:4.9.3', + 'antlr4': 'org.antlr:antlr4:4.9.3', 'assertJ': 'org.assertj:assertj-core:3.11.1', 'avro': 'org.apache.avro:avro:1.11.3', 'avroCompiler': 'org.apache.avro:avro-compiler:1.11.3', @@ -111,7 +111,7 @@ project.ext.externalDependency = [ 'awsRds':'software.amazon.awssdk:rds:2.18.24', 'cacheApi': 'javax.cache:cache-api:1.1.0', 'commonsCli': 'commons-cli:commons-cli:1.5.0', - 'commonsIo': 'commons-io:commons-io:2.4', + 'commonsIo': 'commons-io:commons-io:2.14.0', 'commonsLang': 'commons-lang:commons-lang:2.6', 'commonsText': 'org.apache.commons:commons-text:1.10.0', 'commonsCollections': 'commons-collections:commons-collections:3.2.2', @@ -155,7 +155,7 @@ project.ext.externalDependency = [ 'javatuples': 'org.javatuples:javatuples:1.2', 'javaxInject' : 'javax.inject:javax.inject:1', 'javaxValidation' : 'javax.validation:validation-api:2.0.1.Final', - 'jerseyCore': 'org.glassfish.jersey.core:jersey-client:2.25.1', + 'jerseyCore': 'org.glassfish.jersey.core:jersey-client:2.39.1', 'jerseyGuava': 'org.glassfish.jersey.bundles.repackaged:jersey-guava:2.25.1', 'jettyJaas': "org.eclipse.jetty:jetty-jaas:$jettyVersion", 'jettyClient': "org.eclipse.jetty:jetty-client:$jettyVersion", @@ -171,10 +171,10 @@ project.ext.externalDependency = [ 'junitJupiterParams': "org.junit.jupiter:junit-jupiter-params:$junitJupiterVersion", 'junitJupiterEngine': "org.junit.jupiter:junit-jupiter-engine:$junitJupiterVersion", // avro-serde includes dependencies for `kafka-avro-serializer` `kafka-schema-registry-client` and `avro` - 'kafkaAvroSerde': 'io.confluent:kafka-streams-avro-serde:5.5.1', - 'kafkaAvroSerializer': 'io.confluent:kafka-avro-serializer:5.1.4', + 'kafkaAvroSerde': 'io.confluent:kafka-streams-avro-serde:5.5.14', + 'kafkaAvroSerializer': 'io.confluent:kafka-avro-serializer:5.5.14', 'kafkaClients': "org.apache.kafka:kafka-clients:$kafkaVersion", - 'snappy': 'org.xerial.snappy:snappy-java:1.1.10.4', + 'snappy': 'org.xerial.snappy:snappy-java:1.1.10.5', 'logbackClassic': "ch.qos.logback:logback-classic:$logbackClassic", 'slf4jApi': "org.slf4j:slf4j-api:$slf4jVersion", 'log4jCore': "org.apache.logging.log4j:log4j-core:$log4jVersion", @@ -187,8 +187,8 @@ project.ext.externalDependency = [ 'mixpanel': 'com.mixpanel:mixpanel-java:1.4.4', 'mockito': 'org.mockito:mockito-core:4.11.0', 'mockitoInline': 'org.mockito:mockito-inline:4.11.0', - 'mockServer': 'org.mock-server:mockserver-netty:5.11.2', - 'mockServerClient': 'org.mock-server:mockserver-client-java:5.11.2', + 'mockServer': 'org.mock-server:mockserver-netty:5.13.0', + 'mockServerClient': 'org.mock-server:mockserver-client-java:5.13.0', 'mysqlConnector': 'mysql:mysql-connector-java:8.0.20', 'neo4jHarness': 'org.neo4j.test:neo4j-harness:' + neo4jTestVersion, 'neo4jJavaDriver': 'org.neo4j.driver:neo4j-java-driver:' + neo4jVersion, @@ -213,7 +213,7 @@ project.ext.externalDependency = [ 'playFilters': "com.typesafe.play:filters-helpers_2.12:$playVersion", 'pac4j': 'org.pac4j:pac4j-oidc:4.5.7', 'playPac4j': 'org.pac4j:play-pac4j_2.12:9.0.2', - 'postgresql': 'org.postgresql:postgresql:42.3.8', + 'postgresql': 'org.postgresql:postgresql:42.7.2', 'protobuf': 'com.google.protobuf:protobuf-java:3.19.6', 'grpcProtobuf': 'io.grpc:grpc-protobuf:1.53.0', 'rangerCommons': 'org.apache.ranger:ranger-plugins-common:2.3.0', @@ -255,9 +255,9 @@ project.ext.externalDependency = [ 'typesafeConfig':'com.typesafe:config:1.4.1', 'wiremock':'com.github.tomakehurst:wiremock:2.10.0', 'zookeeper': 'org.apache.zookeeper:zookeeper:3.7.2', - 'wire': 'com.squareup.wire:wire-compiler:3.7.1', + 'wire': 'com.squareup.wire:wire-compiler:4.9.1', 'charle': 'com.charleskorn.kaml:kaml:0.53.0', - 'common': 'commons-io:commons-io:2.7', + 'common': 'commons-io:commons-io:2.14.0', 'jline':'jline:jline:1.4.1', 'jetbrains':' org.jetbrains.kotlin:kotlin-stdlib:1.6.0', 'annotationApi': 'javax.annotation:javax.annotation-api:1.3.2', @@ -347,6 +347,9 @@ configure(subprojects.findAll {! it.name.startsWith('spark-lineage')}) { exclude group: "org.slf4j", module: "slf4j-log4j12" exclude group: "org.slf4j", module: "slf4j-nop" exclude group: "org.slf4j", module: "slf4j-ext" + exclude group: "commons-httpclient", module: "commons-httpclient" + exclude group: "org.codehaus.jackson", module: "jackson-mapper-asl" + exclude group: "software.amazon.ion", module: "ion-java" } } diff --git a/buildSrc/build.gradle b/buildSrc/build.gradle index 88900e06d48451..e628a6a173bccd 100644 --- a/buildSrc/build.gradle +++ b/buildSrc/build.gradle @@ -21,7 +21,7 @@ dependencies { implementation 'com.google.guava:guava:32.1.2-jre' implementation 'com.fasterxml.jackson.core:jackson-databind:2.13.5' implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.13.5' - implementation 'commons-io:commons-io:2.11.0' + implementation 'commons-io:commons-io:2.14.0' compileOnly 'org.projectlombok:lombok:1.18.30' annotationProcessor 'org.projectlombok:lombok:1.18.30' From 8b6790e936d1201d87244bffe616cc41cad2c9c0 Mon Sep 17 00:00:00 2001 From: RyanHolstien Date: Thu, 29 Feb 2024 16:11:04 -0600 Subject: [PATCH 12/13] =?UTF-8?q?Revert=20"chore(vulnerability):=20Bumped?= =?UTF-8?q?=20up=20versions=20for=20vulnerability=20fi=E2=80=A6=20(#9961)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- build.gradle | 29 +++++++++++++---------------- buildSrc/build.gradle | 2 +- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/build.gradle b/build.gradle index cbee21fe1d4c1b..228a8a9f5ff0a0 100644 --- a/build.gradle +++ b/build.gradle @@ -47,7 +47,7 @@ buildscript { ext.log4jVersion = '2.19.0' ext.slf4jVersion = '1.7.36' ext.logbackClassic = '1.4.14' - ext.hadoop3Version = '3.3.6' + ext.hadoop3Version = '3.3.5' ext.kafkaVersion = '2.3.0' ext.hazelcastVersion = '5.3.6' ext.ebeanVersion = '12.16.1' @@ -99,8 +99,8 @@ project.ext.spec = [ project.ext.externalDependency = [ 'akkaHttp': 'com.typesafe.akka:akka-http-core_2.12:10.2.10', - 'antlr4Runtime': 'org.antlr:antlr4-runtime:4.9.3', - 'antlr4': 'org.antlr:antlr4:4.9.3', + 'antlr4Runtime': 'org.antlr:antlr4-runtime:4.7.2', + 'antlr4': 'org.antlr:antlr4:4.7.2', 'assertJ': 'org.assertj:assertj-core:3.11.1', 'avro': 'org.apache.avro:avro:1.11.3', 'avroCompiler': 'org.apache.avro:avro-compiler:1.11.3', @@ -111,7 +111,7 @@ project.ext.externalDependency = [ 'awsRds':'software.amazon.awssdk:rds:2.18.24', 'cacheApi': 'javax.cache:cache-api:1.1.0', 'commonsCli': 'commons-cli:commons-cli:1.5.0', - 'commonsIo': 'commons-io:commons-io:2.14.0', + 'commonsIo': 'commons-io:commons-io:2.4', 'commonsLang': 'commons-lang:commons-lang:2.6', 'commonsText': 'org.apache.commons:commons-text:1.10.0', 'commonsCollections': 'commons-collections:commons-collections:3.2.2', @@ -155,7 +155,7 @@ project.ext.externalDependency = [ 'javatuples': 'org.javatuples:javatuples:1.2', 'javaxInject' : 'javax.inject:javax.inject:1', 'javaxValidation' : 'javax.validation:validation-api:2.0.1.Final', - 'jerseyCore': 'org.glassfish.jersey.core:jersey-client:2.39.1', + 'jerseyCore': 'org.glassfish.jersey.core:jersey-client:2.25.1', 'jerseyGuava': 'org.glassfish.jersey.bundles.repackaged:jersey-guava:2.25.1', 'jettyJaas': "org.eclipse.jetty:jetty-jaas:$jettyVersion", 'jettyClient': "org.eclipse.jetty:jetty-client:$jettyVersion", @@ -171,10 +171,10 @@ project.ext.externalDependency = [ 'junitJupiterParams': "org.junit.jupiter:junit-jupiter-params:$junitJupiterVersion", 'junitJupiterEngine': "org.junit.jupiter:junit-jupiter-engine:$junitJupiterVersion", // avro-serde includes dependencies for `kafka-avro-serializer` `kafka-schema-registry-client` and `avro` - 'kafkaAvroSerde': 'io.confluent:kafka-streams-avro-serde:5.5.14', - 'kafkaAvroSerializer': 'io.confluent:kafka-avro-serializer:5.5.14', + 'kafkaAvroSerde': 'io.confluent:kafka-streams-avro-serde:5.5.1', + 'kafkaAvroSerializer': 'io.confluent:kafka-avro-serializer:5.1.4', 'kafkaClients': "org.apache.kafka:kafka-clients:$kafkaVersion", - 'snappy': 'org.xerial.snappy:snappy-java:1.1.10.5', + 'snappy': 'org.xerial.snappy:snappy-java:1.1.10.4', 'logbackClassic': "ch.qos.logback:logback-classic:$logbackClassic", 'slf4jApi': "org.slf4j:slf4j-api:$slf4jVersion", 'log4jCore': "org.apache.logging.log4j:log4j-core:$log4jVersion", @@ -187,8 +187,8 @@ project.ext.externalDependency = [ 'mixpanel': 'com.mixpanel:mixpanel-java:1.4.4', 'mockito': 'org.mockito:mockito-core:4.11.0', 'mockitoInline': 'org.mockito:mockito-inline:4.11.0', - 'mockServer': 'org.mock-server:mockserver-netty:5.13.0', - 'mockServerClient': 'org.mock-server:mockserver-client-java:5.13.0', + 'mockServer': 'org.mock-server:mockserver-netty:5.11.2', + 'mockServerClient': 'org.mock-server:mockserver-client-java:5.11.2', 'mysqlConnector': 'mysql:mysql-connector-java:8.0.20', 'neo4jHarness': 'org.neo4j.test:neo4j-harness:' + neo4jTestVersion, 'neo4jJavaDriver': 'org.neo4j.driver:neo4j-java-driver:' + neo4jVersion, @@ -213,7 +213,7 @@ project.ext.externalDependency = [ 'playFilters': "com.typesafe.play:filters-helpers_2.12:$playVersion", 'pac4j': 'org.pac4j:pac4j-oidc:4.5.7', 'playPac4j': 'org.pac4j:play-pac4j_2.12:9.0.2', - 'postgresql': 'org.postgresql:postgresql:42.7.2', + 'postgresql': 'org.postgresql:postgresql:42.3.8', 'protobuf': 'com.google.protobuf:protobuf-java:3.19.6', 'grpcProtobuf': 'io.grpc:grpc-protobuf:1.53.0', 'rangerCommons': 'org.apache.ranger:ranger-plugins-common:2.3.0', @@ -255,9 +255,9 @@ project.ext.externalDependency = [ 'typesafeConfig':'com.typesafe:config:1.4.1', 'wiremock':'com.github.tomakehurst:wiremock:2.10.0', 'zookeeper': 'org.apache.zookeeper:zookeeper:3.7.2', - 'wire': 'com.squareup.wire:wire-compiler:4.9.1', + 'wire': 'com.squareup.wire:wire-compiler:3.7.1', 'charle': 'com.charleskorn.kaml:kaml:0.53.0', - 'common': 'commons-io:commons-io:2.14.0', + 'common': 'commons-io:commons-io:2.7', 'jline':'jline:jline:1.4.1', 'jetbrains':' org.jetbrains.kotlin:kotlin-stdlib:1.6.0', 'annotationApi': 'javax.annotation:javax.annotation-api:1.3.2', @@ -347,9 +347,6 @@ configure(subprojects.findAll {! it.name.startsWith('spark-lineage')}) { exclude group: "org.slf4j", module: "slf4j-log4j12" exclude group: "org.slf4j", module: "slf4j-nop" exclude group: "org.slf4j", module: "slf4j-ext" - exclude group: "commons-httpclient", module: "commons-httpclient" - exclude group: "org.codehaus.jackson", module: "jackson-mapper-asl" - exclude group: "software.amazon.ion", module: "ion-java" } } diff --git a/buildSrc/build.gradle b/buildSrc/build.gradle index e628a6a173bccd..88900e06d48451 100644 --- a/buildSrc/build.gradle +++ b/buildSrc/build.gradle @@ -21,7 +21,7 @@ dependencies { implementation 'com.google.guava:guava:32.1.2-jre' implementation 'com.fasterxml.jackson.core:jackson-databind:2.13.5' implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.13.5' - implementation 'commons-io:commons-io:2.14.0' + implementation 'commons-io:commons-io:2.11.0' compileOnly 'org.projectlombok:lombok:1.18.30' annotationProcessor 'org.projectlombok:lombok:1.18.30' From 3284235850afbfcb88e8eb4ba4b462b9062196f9 Mon Sep 17 00:00:00 2001 From: david-leifker <114954101+david-leifker@users.noreply.github.com> Date: Thu, 29 Feb 2024 17:42:07 -0600 Subject: [PATCH 13/13] bump(kafka-setup): client version bump (#9962) --- docker/kafka-setup/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/kafka-setup/Dockerfile b/docker/kafka-setup/Dockerfile index 53353863b6e5f6..6f8b52c5e0bb6e 100644 --- a/docker/kafka-setup/Dockerfile +++ b/docker/kafka-setup/Dockerfile @@ -22,7 +22,7 @@ ARG ALPINE_REPO_URL ARG APACHE_DOWNLOAD_URL ARG GITHUB_REPO_URL -ENV KAFKA_VERSION 3.4.1 +ENV KAFKA_VERSION 3.5.2 ENV SCALA_VERSION 2.13 LABEL name="kafka" version=${KAFKA_VERSION}