diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/mappers/InputFieldsMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/mappers/InputFieldsMapper.java index 49c2d17ce09585..a4e40750f0d659 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/mappers/InputFieldsMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/mappers/InputFieldsMapper.java @@ -5,10 +5,14 @@ import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.InputField; import com.linkedin.datahub.graphql.types.dataset.mappers.SchemaFieldMapper; +import java.net.URISyntaxException; import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; +import lombok.extern.slf4j.Slf4j; + +@Slf4j public class InputFieldsMapper { public static final InputFieldsMapper INSTANCE = new InputFieldsMapper(); @@ -31,13 +35,19 @@ public com.linkedin.datahub.graphql.generated.InputFields apply( .map( field -> { InputField fieldResult = new InputField(); + Urn parentUrn = entityUrn; - if (field.hasSchemaField()) { - fieldResult.setSchemaField( - SchemaFieldMapper.map(context, field.getSchemaField(), entityUrn)); - } if (field.hasSchemaFieldUrn()) { fieldResult.setSchemaFieldUrn(field.getSchemaFieldUrn().toString()); + try { + parentUrn = Urn.createFromString(field.getSchemaFieldUrn().getEntityKey().get(0)); + } catch (URISyntaxException e) { + log.error("Field urn resolution: failed to extract parentUrn successfully from {}. Falling back to {}", field.getSchemaFieldUrn(), entityUrn, e); + } + } + if (field.hasSchemaField()) { + fieldResult.setSchemaField( + SchemaFieldMapper.map(context, field.getSchemaField(), parentUrn)); } return fieldResult; }) diff --git a/datahub-web-react/src/app/settings/SettingsPage.tsx b/datahub-web-react/src/app/settings/SettingsPage.tsx index e0a15c73a626d0..24bcd17ca7f9c0 100644 --- a/datahub-web-react/src/app/settings/SettingsPage.tsx +++ b/datahub-web-react/src/app/settings/SettingsPage.tsx @@ -8,6 +8,7 @@ import { FilterOutlined, TeamOutlined, PushpinOutlined, + ControlOutlined, } from '@ant-design/icons'; import { Redirect, Route, useHistory, useLocation, useRouteMatch, Switch } from 'react-router'; import styled from 'styled-components'; @@ -17,11 +18,17 @@ import { ManagePermissions } from '../permissions/ManagePermissions'; import { useAppConfig } from '../useAppConfig'; import { AccessTokens } from './AccessTokens'; import { Preferences } from './Preferences'; +import { Features } from './features/Features'; import { ManageViews } from '../entity/view/ManageViews'; import { useUserContext } from '../context/useUserContext'; import { ManageOwnership } from '../entity/ownership/ManageOwnership'; import ManagePosts from './posts/ManagePosts'; +const MenuItem = styled(Menu.Item)` + display: flex; + align-items: center; +`; + const PageContainer = styled.div` display: flex; overflow: auto; @@ -59,6 +66,17 @@ const ItemTitle = styled.span` const menuStyle = { width: 256, 'margin-top': 8, overflow: 'hidden auto' }; +const NewTag = styled.span` + padding: 4px 8px; + margin-left: 8px; + + border-radius: 24px; + background: #f1fbfe; + + color: #09739a; + font-size: 12px; +`; + /** * URL Paths for each settings page. */ @@ -70,6 +88,7 @@ const PATHS = [ { path: 'views', content: }, { path: 'ownership', content: }, { path: 'posts', content: }, + { path: 'features', content: }, ]; /** @@ -80,6 +99,7 @@ const DEFAULT_PATH = PATHS[0]; export const SettingsPage = () => { const { path, url } = useRouteMatch(); const { pathname } = useLocation(); + const history = useHistory(); const subRoutes = PATHS.map((p) => p.path.replace('/', '')); const currPathName = pathname.replace(path, ''); @@ -101,6 +121,7 @@ export const SettingsPage = () => { const showViews = isViewsEnabled || false; const showOwnershipTypes = me && me?.platformPrivileges?.manageOwnershipTypes; const showHomePagePosts = me && me?.platformPrivileges?.manageGlobalAnnouncements && !readOnlyModeEnabled; + const showFeatures = true; // TODO: Add feature flag for this return ( @@ -143,6 +164,13 @@ export const SettingsPage = () => { )} {(showViews || showOwnershipTypes || showHomePagePosts) && ( + {showFeatures && ( + + + Features + New! + + )} {showViews && ( My Views diff --git a/datahub-web-react/src/app/settings/features/Feature.tsx b/datahub-web-react/src/app/settings/features/Feature.tsx new file mode 100644 index 00000000000000..2c090aae696f88 --- /dev/null +++ b/datahub-web-react/src/app/settings/features/Feature.tsx @@ -0,0 +1,179 @@ +import React from 'react'; + +import styled from 'styled-components'; + +import { Divider, Typography, Switch, Card, Button, Tooltip } from 'antd'; +import { ArrowRightOutlined } from '@ant-design/icons'; +import { ANTD_GRAY } from '../../entity/shared/constants'; + +const Title = styled(Typography.Title)` + && { + margin-bottom: 8px; + } +`; + +const FeatureRow = styled.div` + display: flex; + align-items: flex-start; + justify-content: space-between; +`; + +const FeatureOptionRow = styled.div` + display: flex; + justify-content: space-between; + + &:not(:last-child) { + margin-bottom: 8px; + } +`; + +const SettingsOptionRow = styled.div` + display: flex; + justify-content: space-between; + align-items: center; + padding: 0 16px; + + &:not(:last-child) { + margin-bottom: 8px; + } +`; + +const DescriptionText = styled(Typography.Text)` + color: ${ANTD_GRAY[7]}; + font-size: 11px; +`; + +const SettingTitle = styled.div` + display: flex; + align-items: center; + gap: 8px; + font-size: 14px; + margin-bottom: 4px; +`; + +const OptionTitle = styled(Typography.Text)` + display: flex; + align-items: center; + gap: 8px; + font-size: 12px; +`; + +const learnMoreLinkStyle = { + flex: 1, + display: 'flex', + alignItems: 'center', + gap: '8px', + color: '#1890FF', + fontSize: '12px', + cursor: 'pointer', +}; + +const NewTag = styled.div` + padding: 4px 8px; + + border-radius: 24px; + background: #f1fbfe; + + color: #09739a; + font-size: 12px; +`; + +const DataHubOnlyTag = styled.div` + padding: 2px 8px; + + border-radius: 24px; + background: #c9fff2; + + color: #50a494; + font-size: 12px; +`; + +export interface FeatureType { + key: string; + title: string; + description: string; + settings: Array<{ + key: string; + title: string; + isAvailable: boolean; + buttonText: string; + onClick?: () => void; + }>; + options: Array<{ + key: string; + title: string; + description: string; + isAvailable: boolean; + checked: boolean; + onChange?: (checked: boolean) => void; + }>; + isNew: boolean; + learnMoreLink?: string; +} + +export const Feature = ({ key, title, description, settings, options, isNew, learnMoreLink }: FeatureType) => ( + + +
+ + + {title} + + {isNew && New!} + +
+ {description} +
+
+
+ {learnMoreLink && ( + + Learn more + + )} +
+ + + {settings.map((option) => ( + <> + + + + {option.title} + + + + + + + + ))} + + {options.map((option, index) => ( + <> + + + + {option.title} + {!option.isAvailable && ( + Only available on DataHub Cloud + )} + +
+ {option.description} +
+ + (option.onChange ? option.onChange(checked) : null)} + disabled={!option.isAvailable} + /> + + {index !== options.length - 1 && } + + ))} + + +); diff --git a/datahub-web-react/src/app/settings/features/Features.tsx b/datahub-web-react/src/app/settings/features/Features.tsx new file mode 100644 index 00000000000000..ee8d7c628c1eff --- /dev/null +++ b/datahub-web-react/src/app/settings/features/Features.tsx @@ -0,0 +1,110 @@ +import React from 'react'; + +import styled from 'styled-components'; + +import { Divider, Typography } from 'antd'; +import { v4 as uuidv4 } from 'uuid'; + +import { Feature, FeatureType } from './Feature'; + +import { useGetDocPropagationSettings, useUpdateDocPropagationSettings } from './useDocPropagationSettings'; + +const Page = styled.div` + width: 100%; + display: flex; + justify-content: center; +`; + +const SourceContainer = styled.div` + width: 80%; + padding-top: 20px; + padding-right: 40px; + padding-left: 40px; +`; +const Container = styled.div` + padding-top: 0px; +`; + +const Title = styled(Typography.Title)` + && { + margin-bottom: 8px; + } +`; + +export const Features = () => { + /* + * Note: When adding new features, make sure to update the features array below + * and create a hook file for the new feature in the same directory + */ + + // Hooks to get and update the document propagation settings + const { isColPropagateChecked, setIsColPropagateChecked } = useGetDocPropagationSettings(); + const { updateDocPropagation } = useUpdateDocPropagationSettings(); + + // Features to display + const features: FeatureType[] = [ + { + key: uuidv4(), + title: 'Documentation Propagation', + description: 'Automatically propagate documentation from upstream to downstream columns and assets.', + settings: [ + { + key: uuidv4(), + title: 'Rollback Propagation Changes', + isAvailable: false, + buttonText: 'Rollback', + }, + { + key: uuidv4(), + title: 'Backfill existing documentation from upstream to downstream columns/assets', + isAvailable: false, + buttonText: 'Initialize', + }, + ], + options: [ + { + key: uuidv4(), + title: 'Column Level Propagation', + description: + 'Propagate new documentation from upstream to downstream columns based on column-level lineage relationships.', + isAvailable: true, + checked: isColPropagateChecked, + onChange: (checked: boolean) => { + setIsColPropagateChecked(checked); + updateDocPropagation(checked); + }, + }, + { + key: uuidv4(), + title: 'Asset Level Propagation', + description: + 'Propagate new documentation from upstream to downstream assets based on data lineage relationships.', + isAvailable: false, + checked: false, + }, + ], + isNew: true, + learnMoreLink: 'https://datahubproject.io/docs/automations/doc-propagation', + }, + ]; + + // Render + return ( + + + +
+ Features + + Explore and configure specific features + +
+ + + {features.map((feature) => ( + + ))} + + + ); +}; diff --git a/datahub-web-react/src/app/settings/features/useDocPropagationSettings.ts b/datahub-web-react/src/app/settings/features/useDocPropagationSettings.ts new file mode 100644 index 00000000000000..c93b610cff9d1b --- /dev/null +++ b/datahub-web-react/src/app/settings/features/useDocPropagationSettings.ts @@ -0,0 +1,50 @@ +import { useEffect, useState } from 'react'; + +import { message } from 'antd'; + +import { + useGetDocPropagationSettingsQuery, + useUpdateDocPropagationSettingsMutation, +} from '../../../graphql/app.generated'; + +// Hook to get the document propagation settings & manage state +export const useGetDocPropagationSettings = () => { + const { data, refetch } = useGetDocPropagationSettingsQuery(); + const [isColPropagateChecked, setIsColPropagateChecked] = useState(false); + + useEffect(() => { + const docPropSetting = data?.docPropagationSettings?.docColumnPropagation; + if (docPropSetting !== undefined) setIsColPropagateChecked(!!docPropSetting); + }, [data]); + + return { + isColPropagateChecked, + setIsColPropagateChecked, + refetch, + }; +}; + +// Hook to update the document propagation settings +export const useUpdateDocPropagationSettings = () => { + const [updateDocPropagationSettings] = useUpdateDocPropagationSettingsMutation(); + const { refetch } = useGetDocPropagationSettingsQuery(); + + const updateDocPropagation = async (checked: boolean) => { + try { + await updateDocPropagationSettings({ + variables: { + input: { + docColumnPropagation: checked, + }, + }, + }); + refetch(); + message.success('Successfully updated documentation propagation settings'); + } catch (e) { + message.error('Failed to update documentation propagation settings'); + refetch(); + } + }; + + return { updateDocPropagation }; +}; diff --git a/datahub-web-react/src/graphql/app.graphql b/datahub-web-react/src/graphql/app.graphql index bfca27a4ad106a..e058a6fbb58e00 100644 --- a/datahub-web-react/src/graphql/app.graphql +++ b/datahub-web-react/src/graphql/app.graphql @@ -89,6 +89,16 @@ query getGlobalViewsSettings { } } +query getDocPropagationSettings { + docPropagationSettings { + docColumnPropagation + } +} + mutation updateGlobalViewsSettings($input: UpdateGlobalViewsSettingsInput!) { updateGlobalViewsSettings(input: $input) } + +mutation updateDocPropagationSettings($input: UpdateDocPropagationSettingsInput!) { + updateDocPropagationSettings(input: $input) +} diff --git a/metadata-models/src/main/pegasus/com/linkedin/dataset/PartitionSummary.pdl b/metadata-models/src/main/pegasus/com/linkedin/dataset/PartitionSummary.pdl new file mode 100644 index 00000000000000..3984277a314171 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/dataset/PartitionSummary.pdl @@ -0,0 +1,24 @@ +namespace com.linkedin.dataset + +import com.linkedin.common.AuditStamp + +/** + * Defines how the data is partitioned + */ +record PartitionSummary { + /** + * A unique id / value for the partition for which statistics were collected, + * generated by applying the key definition to a given row. + */ + partition: string + + /** + * The created time for a given partition. + */ + created: optional AuditStamp + + /** + * The last modified / touched time for a given partition. + */ + lastModified: optional AuditStamp +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/dataset/PartitionsSummary.pdl b/metadata-models/src/main/pegasus/com/linkedin/dataset/PartitionsSummary.pdl new file mode 100644 index 00000000000000..34e696890d64ff --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/dataset/PartitionsSummary.pdl @@ -0,0 +1,19 @@ +namespace com.linkedin.dataset + +/** + * Defines how the data is partitioned for Data Lake tables (e.g. Hive, S3, Iceberg, Delta, Hudi, etc). + */ +@Aspect = { + "name": "partitionsSummary" +} +record PartitionsSummary { + /** + * The minimum partition as ordered + */ + minPartition: optional PartitionSummary + + /** + * The maximum partition as ordered + */ + maxPartition: optional PartitionSummary +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/metadata/key/DataHubActionKey.pdl b/metadata-models/src/main/pegasus/com/linkedin/metadata/key/DataHubActionKey.pdl new file mode 100644 index 00000000000000..8205ecbb807168 --- /dev/null +++ b/metadata-models/src/main/pegasus/com/linkedin/metadata/key/DataHubActionKey.pdl @@ -0,0 +1,14 @@ +namespace com.linkedin.metadata.key + +/** + * Key for a DataHub Action Pipeline + */ +@Aspect = { + "name": "dataHubActionKey" +} +record DataHubActionKey { + /** + * A unique id for the Action, either generated or provided + */ + id: string +} \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/schema/SchemaField.pdl b/metadata-models/src/main/pegasus/com/linkedin/schema/SchemaField.pdl index afb0263057b6d9..f91e2004401cf9 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/schema/SchemaField.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/schema/SchemaField.pdl @@ -150,6 +150,8 @@ record SchemaField { /** * For Datasets which are partitioned, this determines the partitioning key. + * Note that multiple columns can be part of a partitioning key, but currently we do not support + * rendering the ordered partitioning key. */ isPartitioningKey: optional boolean diff --git a/metadata-models/src/main/pegasus/com/linkedin/timeseries/PartitionSpec.pdl b/metadata-models/src/main/pegasus/com/linkedin/timeseries/PartitionSpec.pdl index 084af1513ec807..146a285e24dbab 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/timeseries/PartitionSpec.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/timeseries/PartitionSpec.pdl @@ -1,24 +1,28 @@ namespace com.linkedin.timeseries /** - * Defines how the data is partitioned + * A reference to a specific partition in a dataset. */ record PartitionSpec { - - type: enum PartitionType { - FULL_TABLE, - QUERY, - PARTITION - } = "PARTITION" - /** - * String representation of the partition + * A unique id / value for the partition for which statistics were collected, + * generated by applying the key definition to a given row. */ @TimeseriesField = {} partition: string /** - * Time window of the partition if applicable + * Time window of the partition, if we are able to extract it from the partition key. */ timePartition: optional TimeWindow + + /** + * Unused! + */ + @deprecated + type: enum PartitionType { + FULL_TABLE, + QUERY, + PARTITION + } = "PARTITION" } \ 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 3af4af5e4767e2..f8520990a09841 100644 --- a/metadata-models/src/main/resources/entity-registry.yml +++ b/metadata-models/src/main/resources/entity-registry.yml @@ -45,6 +45,7 @@ entities: - access - structuredProperties - forms + - partitionsSummary - name: dataHubPolicy doc: DataHub Policies represent access policies granted to users or groups on metadata operations like edit, view etc. category: internal @@ -74,6 +75,7 @@ entities: - forms - subTypes - incidentsSummary + - testResults - name: dataFlow category: core keyAspect: dataFlowKey @@ -95,12 +97,14 @@ entities: - incidentsSummary - forms - subTypes + - testResults - name: dataProcess keyAspect: dataProcessKey aspects: - dataProcessInfo - ownership - status + - testResults - name: dataProcessInstance doc: DataProcessInstance represents an instance of a datajob/jobflow run keyAspect: dataProcessInstanceKey @@ -111,6 +115,7 @@ entities: - dataProcessInstanceRelationships - dataProcessInstanceRunEvent - status + - testResults - name: chart category: core keyAspect: chartKey @@ -136,6 +141,7 @@ entities: - structuredProperties - incidentsSummary - forms + - testResults - name: dashboard keyAspect: dashboardKey aspects: @@ -159,6 +165,7 @@ entities: - structuredProperties - incidentsSummary - forms + - testResults - name: notebook doc: Notebook represents a combination of query, text, chart and etc. This is in BETA version keyAspect: notebookKey @@ -176,6 +183,7 @@ entities: - subTypes - dataPlatformInstance - browsePathsV2 + - testResults - name: corpuser doc: CorpUser represents an identity of a person (or an account) in the enterprise. keyAspect: corpUserKey @@ -193,6 +201,7 @@ entities: - roleMembership - structuredProperties - forms + - testResults - name: corpGroup doc: CorpGroup represents an identity of a group of users in the enterprise. keyAspect: corpGroupKey @@ -206,6 +215,7 @@ entities: - roleMembership - structuredProperties - forms + - testResults - name: domain doc: A data domain within an organization. category: core @@ -216,6 +226,7 @@ entities: - ownership - structuredProperties - forms + - testResults - name: container doc: A container of related data assets. category: core @@ -236,6 +247,7 @@ entities: - browsePathsV2 - structuredProperties - forms + - testResults - name: tag category: core keyAspect: tagKey @@ -244,6 +256,7 @@ entities: - ownership - deprecation - status + - testResults - name: glossaryTerm category: core keyAspect: glossaryTermKey @@ -259,6 +272,7 @@ entities: - browsePaths - structuredProperties - forms + - testResults - name: glossaryNode category: core keyAspect: glossaryNodeKey @@ -269,6 +283,7 @@ entities: - status - structuredProperties - forms + - testResults - name: dataHubIngestionSource category: internal keyAspect: dataHubIngestionSourceKey @@ -340,6 +355,7 @@ entities: - browsePathsV2 - structuredProperties - forms + - testResults - name: mlModelGroup category: core keyAspect: mlModelGroupKey @@ -357,6 +373,7 @@ entities: - browsePathsV2 - structuredProperties - forms + - testResults - name: mlModelDeployment category: core keyAspect: mlModelDeploymentKey @@ -367,6 +384,7 @@ entities: - deprecation - globalTags - dataPlatformInstance + - testResults - name: mlFeatureTable category: core keyAspect: mlFeatureTableKey @@ -385,6 +403,7 @@ entities: - browsePathsV2 - structuredProperties - forms + - testResults - name: mlFeature category: core keyAspect: mlFeatureKey @@ -403,6 +422,7 @@ entities: - browsePathsV2 - structuredProperties - forms + - testResults - name: mlPrimaryKey category: core keyAspect: mlPrimaryKeyKey @@ -419,6 +439,7 @@ entities: - dataPlatformInstance - structuredProperties - forms + - testResults - name: telemetry category: internal keyAspect: telemetryKey @@ -455,6 +476,7 @@ entities: - forms - businessAttributes - documentation + - testResults - name: globalSettings doc: Global settings for an the platform category: internal @@ -522,6 +544,7 @@ entities: - status - structuredProperties - forms + - testResults - name: ownershipType doc: Ownership Type represents a user-created ownership category for a person or group who is responsible for an asset. category: core @@ -549,6 +572,10 @@ entities: keyAspect: dataHubPersonaKey aspects: - dataHubPersonaInfo + - name: dataHubAction + category: internal + keyAspect: dataHubActionKey + aspects: [] - name: entityType doc: A type of entity in the DataHub Metadata Model. category: core diff --git a/smoke-test/tests/schema_fields/__init__.py b/smoke-test/tests/schema_fields/__init__.py new file mode 100644 index 00000000000000..e69de29bb2d1d6 diff --git a/smoke-test/tests/schema_fields/queries/get_chart_field.gql b/smoke-test/tests/schema_fields/queries/get_chart_field.gql new file mode 100644 index 00000000000000..424e5ad686ab9c --- /dev/null +++ b/smoke-test/tests/schema_fields/queries/get_chart_field.gql @@ -0,0 +1,20 @@ +query($urn:String!) { + chart(urn: $urn) { + inputFields { + fields { + schemaFieldUrn + schemaField { + schemaFieldEntity { + urn + fieldPath + documentation { + documentations { + documentation + } + } + } + } + } + } + } +} \ No newline at end of file diff --git a/smoke-test/tests/schema_fields/test_schemafields.py b/smoke-test/tests/schema_fields/test_schemafields.py new file mode 100644 index 00000000000000..31f237308e2a81 --- /dev/null +++ b/smoke-test/tests/schema_fields/test_schemafields.py @@ -0,0 +1,177 @@ +import logging +import os +import tempfile +import time +from random import randint + +import datahub.metadata.schema_classes as models +import pytest +from datahub.emitter.mce_builder import ( + make_dataset_urn, + make_schema_field_urn, +) +from datahub.emitter.mcp import MetadataChangeProposalWrapper +from datahub.ingestion.api.common import PipelineContext, RecordEnvelope +from datahub.ingestion.api.sink import NoopWriteCallback +from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph +from datahub.ingestion.sink.file import FileSink, FileSinkConfig + +from tests.utils import ( + delete_urns_from_file, + get_gms_url, + get_sleep_info, + ingest_file_via_rest, + wait_for_writes_to_sync, +) + +logger = logging.getLogger(__name__) + + +start_index = randint(10, 10000) +dataset_urns = [ + make_dataset_urn("snowflake", f"table_foo_{i}") + for i in range(start_index, start_index + 10) +] + + +class FileEmitter: + def __init__(self, filename: str) -> None: + self.sink: FileSink = FileSink( + ctx=PipelineContext(run_id="create_test_data"), + config=FileSinkConfig(filename=filename), + ) + + def emit(self, event): + self.sink.write_record_async( + record_envelope=RecordEnvelope(record=event, metadata={}), + write_callback=NoopWriteCallback(), + ) + + def close(self): + self.sink.close() + + +@pytest.fixture(scope="module") +def chart_urn(): + return "urn:li:chart:(looker,chart_foo)" + + +@pytest.fixture(scope="module") +def upstream_schema_field_urn(): + return make_schema_field_urn(make_dataset_urn("snowflake", "table_bar"), "field1") + + +def create_test_data(filename: str, chart_urn: str, upstream_schema_field_urn: str): + documentation_mcp = MetadataChangeProposalWrapper( + entityUrn=upstream_schema_field_urn, + aspect=models.DocumentationClass( + documentations=[ + models.DocumentationAssociationClass( + documentation="test documentation", + attribution=models.MetadataAttributionClass( + time=int(time.time() * 1000), + actor="urn:li:corpuser:datahub", + source="urn:li:dataHubAction:documentation_propagation", + ), + ) + ] + ), + ) + + input_fields_mcp = MetadataChangeProposalWrapper( + entityUrn=chart_urn, + aspect=models.InputFieldsClass( + fields=[ + models.InputFieldClass( + schemaFieldUrn=upstream_schema_field_urn, + schemaField=models.SchemaFieldClass( + fieldPath="field1", + type=models.SchemaFieldDataTypeClass(models.StringTypeClass()), + nativeDataType="STRING", + ), + ) + ] + ), + ) + + file_emitter = FileEmitter(filename) + for mcps in [documentation_mcp, input_fields_mcp]: + file_emitter.emit(mcps) + + file_emitter.close() + + +sleep_sec, sleep_times = get_sleep_info() + + +@pytest.fixture(scope="module", autouse=False) +def ingest_cleanup_data(request, chart_urn, upstream_schema_field_urn): + new_file, filename = tempfile.mkstemp(suffix=".json") + try: + create_test_data(filename, chart_urn, upstream_schema_field_urn) + print("ingesting schema fields test data") + ingest_file_via_rest(filename) + yield + print("removing schema fields test data") + delete_urns_from_file(filename) + wait_for_writes_to_sync() + finally: + os.remove(filename) + + +@pytest.mark.dependency() +def test_healthchecks(wait_for_healthchecks): + # Call to wait_for_healthchecks fixture will do the actual functionality. + pass + + +def get_gql_query(filename: str) -> str: + with open(filename) as fp: + return fp.read() + + +def validate_schema_field_urn_for_chart( + graph: DataHubGraph, chart_urn: str, upstream_schema_field_urn: str +) -> None: + # Validate listing + result = graph.execute_graphql( + get_gql_query("tests/schema_fields/queries/get_chart_field.gql"), + {"urn": chart_urn}, + ) + assert "chart" in result + assert "inputFields" in result["chart"] + assert len(result["chart"]["inputFields"]["fields"]) == 1 + assert ( + result["chart"]["inputFields"]["fields"][0]["schemaField"]["schemaFieldEntity"][ + "urn" + ] + == upstream_schema_field_urn + ) + assert ( + result["chart"]["inputFields"]["fields"][0]["schemaField"]["schemaFieldEntity"][ + "fieldPath" + ] + == "field1" + ) + assert ( + result["chart"]["inputFields"]["fields"][0]["schemaFieldUrn"] + == upstream_schema_field_urn + ) + assert ( + result["chart"]["inputFields"]["fields"][0]["schemaField"]["schemaFieldEntity"][ + "documentation" + ]["documentations"][0]["documentation"] + == "test documentation" + ) + + +# @tenacity.retry( +# stop=tenacity.stop_after_attempt(sleep_times), wait=tenacity.wait_fixed(sleep_sec) +# ) +@pytest.mark.dependency(depends=["test_healthchecks"]) +def test_schema_field_gql_mapper_for_charts( + ingest_cleanup_data, chart_urn, upstream_schema_field_urn +): + graph: DataHubGraph = DataHubGraph(config=DatahubClientConfig(server=get_gms_url())) + + validate_schema_field_urn_for_chart(graph, chart_urn, upstream_schema_field_urn)