acryl-datahub 1.0.0.1rc2__py3-none-any.whl → 1.0.0.1rc3__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_schema_gen.py

@@ -41,6 +41,7 @@ from datahub.ingestion.source.snowflake.snowflake_query import SnowflakeQuery
 from datahub.ingestion.source.snowflake.snowflake_report import SnowflakeV2Report
 from datahub.ingestion.source.snowflake.snowflake_schema import (
     SCHEMA_PARALLELISM,
+    BaseProcedure,
     SnowflakeColumn,
     SnowflakeDatabase,
     SnowflakeDataDictionary,
@@ -63,12 +64,14 @@ from datahub.ingestion.source.snowflake.snowflake_utils import (
 from datahub.ingestion.source.sql.sql_utils import (
     add_table_to_schema_container,
     gen_database_container,
-    gen_database_key,
     gen_schema_container,
-    gen_schema_key,
     get_dataplatform_instance_aspect,
     get_domain_wu,
 )
+from datahub.ingestion.source.sql.stored_procedures.base import (
+    generate_procedure_container_workunits,
+    generate_procedure_workunits,
+)
 from datahub.ingestion.source_report.ingestion_stage import (
     EXTERNAL_TABLE_DDL_LINEAGE,
     LINEAGE_EXTRACTION,
@@ -448,10 +451,15 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         if self.config.include_streams:
             self.report.num_get_streams_for_schema_queries += 1
             streams = self.fetch_streams_for_schema(
-                snowflake_schema, db_name, schema_name
+                snowflake_schema,
+                db_name,
             )
             yield from self._process_streams(streams, snowflake_schema, db_name)
 
+        if self.config.include_procedures:
+            procedures = self.fetch_procedures_for_schema(snowflake_schema, db_name)
+            yield from self._process_procedures(procedures, snowflake_schema, db_name)
+
         if self.config.include_technical_schema and snowflake_schema.tags:
             yield from self._process_tags_in_schema(snowflake_schema)
 
@@ -536,6 +544,26 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         for stream in streams:
             yield from self._process_stream(stream, snowflake_schema, db_name)
 
+    def _process_procedures(
+        self,
+        procedures: List[BaseProcedure],
+        snowflake_schema: SnowflakeSchema,
+        db_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        if self.config.include_technical_schema:
+            if procedures:
+                yield from generate_procedure_container_workunits(
+                    self.identifiers.gen_database_key(
+                        db_name,
+                    ),
+                    self.identifiers.gen_schema_key(
+                        db_name=db_name,
+                        schema_name=snowflake_schema.name,
+                    ),
+                )
+            for procedure in procedures:
+                yield from self._process_procedure(procedure, snowflake_schema, db_name)
+
     def _process_tags_in_schema(
         self, snowflake_schema: SnowflakeSchema
     ) -> Iterable[MetadataWorkUnit]:
@@ -819,13 +847,7 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
             entityUrn=dataset_urn, aspect=dataset_properties
         ).as_workunit()
 
-        schema_container_key = gen_schema_key(
-            db_name=self.snowflake_identifier(db_name),
-            schema=self.snowflake_identifier(schema_name),
-            platform=self.platform,
-            platform_instance=self.config.platform_instance,
-            env=self.config.env,
-        )
+        schema_container_key = self.identifiers.gen_schema_key(db_name, schema_name)
 
         if self.config.extract_tags_as_structured_properties:
             yield from self.gen_column_tags_as_structured_properties(dataset_urn, table)
@@ -1094,11 +1116,8 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
     def gen_database_containers(
         self, database: SnowflakeDatabase
     ) -> Iterable[MetadataWorkUnit]:
-        database_container_key = gen_database_key(
-            self.snowflake_identifier(database.name),
-            platform=self.platform,
-            platform_instance=self.config.platform_instance,
-            env=self.config.env,
+        database_container_key = self.identifiers.gen_database_key(
+            database.name,
         )
 
         yield from gen_database_container(
@@ -1147,21 +1166,9 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
     def gen_schema_containers(
         self, schema: SnowflakeSchema, db_name: str
     ) -> Iterable[MetadataWorkUnit]:
-        schema_name = self.snowflake_identifier(schema.name)
-        database_container_key = gen_database_key(
-            database=self.snowflake_identifier(db_name),
-            platform=self.platform,
-            platform_instance=self.config.platform_instance,
-            env=self.config.env,
-        )
+        database_container_key = self.identifiers.gen_database_key(db_name)
 
-        schema_container_key = gen_schema_key(
-            db_name=self.snowflake_identifier(db_name),
-            schema=schema_name,
-            platform=self.platform,
-            platform_instance=self.config.platform_instance,
-            env=self.config.env,
-        )
+        schema_container_key = self.identifiers.gen_schema_key(db_name, schema.name)
 
         yield from gen_schema_container(
             name=schema.name,
@@ -1290,13 +1297,13 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
             )
 
     def fetch_streams_for_schema(
-        self, snowflake_schema: SnowflakeSchema, db_name: str, schema_name: str
+        self, snowflake_schema: SnowflakeSchema, db_name: str
     ) -> List[SnowflakeStream]:
         try:
             streams: List[SnowflakeStream] = []
-            for stream in self.get_streams_for_schema(schema_name, db_name):
+            for stream in self.get_streams_for_schema(snowflake_schema.name, db_name):
                 stream_identifier = self.identifiers.get_dataset_identifier(
-                    stream.name, schema_name, db_name
+                    stream.name, snowflake_schema.name, db_name
                 )
 
                 self.report.report_entity_scanned(stream_identifier, "stream")
@@ -1310,16 +1317,15 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
             snowflake_schema.streams = [stream.name for stream in streams]
             return streams
         except Exception as e:
-            if isinstance(e, SnowflakePermissionError):
-                error_msg = f"Failed to get streams for schema {db_name}.{schema_name}. Please check permissions."
-                raise SnowflakePermissionError(error_msg) from e.__cause__
-            else:
-                self.structured_reporter.warning(
-                    "Failed to get streams for schema",
-                    f"{db_name}.{schema_name}",
-                    exc=e,
-                )
-                return []
+            self.structured_reporter.warning(
+                title="Failed to get streams for schema",
+                message="Please check permissions"
+                if isinstance(e, SnowflakePermissionError)
+                else "",
+                context=f"{db_name}.{snowflake_schema.name}",
+                exc=e,
+            )
+            return []
 
     def get_streams_for_schema(
         self, schema_name: str, db_name: str
@@ -1328,6 +1334,42 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
 
         return streams.get(schema_name, [])
 
+    def fetch_procedures_for_schema(
+        self, snowflake_schema: SnowflakeSchema, db_name: str
+    ) -> List[BaseProcedure]:
+        try:
+            procedures: List[BaseProcedure] = []
+            for procedure in self.get_procedures_for_schema(snowflake_schema, db_name):
+                procedure_qualified_name = self.identifiers.get_dataset_identifier(
+                    procedure.name, snowflake_schema.name, db_name
+                )
+                self.report.report_entity_scanned(procedure_qualified_name, "procedure")
+
+                if self.filters.is_procedure_allowed(procedure_qualified_name):
+                    procedures.append(procedure)
+                else:
+                    self.report.report_dropped(procedure_qualified_name)
+            return procedures
+        except Exception as e:
+            self.structured_reporter.warning(
+                title="Failed to get procedures for schema",
+                message="Please check permissions"
+                if isinstance(e, SnowflakePermissionError)
+                else "",
+                context=f"{db_name}.{snowflake_schema.name}",
+                exc=e,
+            )
+            return []
+
+    def get_procedures_for_schema(
+        self,
+        snowflake_schema: SnowflakeSchema,
+        db_name: str,
+    ) -> List[BaseProcedure]:
+        procedures = self.data_dictionary.get_procedures_for_database(db_name)
+
+        return procedures.get(snowflake_schema.name, [])
+
     def _process_stream(
         self,
         stream: SnowflakeStream,
@@ -1350,6 +1392,34 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
                 "Failed to get columns for stream:", stream.name, exc=e
             )
 
+    def _process_procedure(
+        self,
+        procedure: BaseProcedure,
+        snowflake_schema: SnowflakeSchema,
+        db_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        try:
+            # TODO: For CLL, we should process procedures after all tables are processed
+            yield from generate_procedure_workunits(
+                procedure,
+                database_key=self.identifiers.gen_database_key(
+                    db_name,
+                ),
+                schema_key=self.identifiers.gen_schema_key(
+                    db_name, snowflake_schema.name
+                ),
+                schema_resolver=(
+                    self.aggregator._schema_resolver if self.aggregator else None
+                ),
+            )
+        except Exception as e:
+            self.structured_reporter.warning(
+                title="Failed to ingest stored procedure",
+                message="",
+                context=procedure.name,
+                exc=e,
+            )
+
     def get_columns_for_stream(
         self,
         source_object: str,  # Qualified name of source table/view

datahub/ingestion/source/snowflake/snowflake_utils.py

@@ -3,7 +3,10 @@ from functools import cached_property
 from typing import ClassVar, List, Literal, Optional, Tuple
 
 from datahub.configuration.pattern_utils import is_schema_allowed
-from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance
+from datahub.emitter.mce_builder import (
+    make_dataset_urn_with_platform_instance,
+)
+from datahub.emitter.mcp_builder import DatabaseKey, SchemaKey
 from datahub.ingestion.api.source import SourceReport
 from datahub.ingestion.source.snowflake.constants import (
     SNOWFLAKE_REGION_CLOUD_REGION_MAPPING,
@@ -16,6 +19,7 @@ from datahub.ingestion.source.snowflake.snowflake_config import (
     SnowflakeV2Config,
 )
 from datahub.ingestion.source.snowflake.snowflake_report import SnowflakeV2Report
+from datahub.ingestion.source.sql.sql_utils import gen_database_key, gen_schema_key
 
 
 class SnowflakeStructuredReportMixin(abc.ABC):
@@ -180,6 +184,9 @@ class SnowflakeFilter:
 
         return True
 
+    def is_procedure_allowed(self, procedure_name: str) -> bool:
+        return self.filter_config.procedure_pattern.allowed(procedure_name)
+
 
 def _combine_identifier_parts(
     *, table_name: str, schema_name: str, db_name: str
@@ -330,6 +337,23 @@ class SnowflakeIdentifierBuilder:
             else user_name
         )
 
+    def gen_schema_key(self, db_name: str, schema_name: str) -> SchemaKey:
+        return gen_schema_key(
+            db_name=self.snowflake_identifier(db_name),
+            schema=self.snowflake_identifier(schema_name),
+            platform=self.platform,
+            platform_instance=self.identifier_config.platform_instance,
+            env=self.identifier_config.env,
+        )
+
+    def gen_database_key(self, db_name: str) -> DatabaseKey:
+        return gen_database_key(
+            database=self.snowflake_identifier(db_name),
+            platform=self.platform,
+            platform_instance=self.identifier_config.platform_instance,
+            env=self.identifier_config.env,
+        )
+
 
 class SnowflakeCommonMixin(SnowflakeStructuredReportMixin):
     platform = "snowflake"

datahub/ingestion/source/sql/mssql/job_models.py

@@ -15,6 +15,7 @@ from datahub.ingestion.source.common.subtypes import (
     FlowContainerSubTypes,
     JobContainerSubTypes,
 )
+from datahub.ingestion.source.sql.stored_procedures.base import BaseProcedure
 from datahub.metadata.schema_classes import (
     ContainerClass,
     DataFlowInfoClass,
@@ -135,6 +136,19 @@ class StoredProcedure:
     def escape_full_name(self) -> str:
         return f"[{self.db}].[{self.schema}].[{self.formatted_name}]"
 
+    def to_base_procedure(self) -> BaseProcedure:
+        return BaseProcedure(
+            name=self.formatted_name,
+            procedure_definition=self.code,
+            created=None,
+            last_altered=None,
+            comment=None,
+            argument_signature=None,
+            return_type=None,
+            language="SQL",
+            extra_properties=None,
+        )
+
 
 @dataclass
 class JobStep:
@@ -222,7 +236,7 @@ class MSSQLDataJob:
         type = (
             JobContainerSubTypes.MSSQL_JOBSTEP
             if isinstance(self.entity, JobStep)
-            else JobContainerSubTypes.MSSQL_STORED_PROCEDURE
+            else JobContainerSubTypes.STORED_PROCEDURE
         )
         return SubTypesClass(
             typeNames=[type],

datahub/ingestion/source/sql/mssql/source.py

@@ -37,9 +37,6 @@ from datahub.ingestion.source.sql.mssql.job_models import (
     ProcedureParameter,
     StoredProcedure,
 )
-from datahub.ingestion.source.sql.mssql.stored_procedure_lineage import (
-    generate_procedure_lineage,
-)
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     SqlWorkUnit,
@@ -50,6 +47,9 @@ from datahub.ingestion.source.sql.sql_config import (
     make_sqlalchemy_uri,
 )
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
+from datahub.ingestion.source.sql.stored_procedures.base import (
+    generate_procedure_lineage,
+)
 from datahub.utilities.file_backed_collections import FileBackedList
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -65,6 +65,8 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
     # defaults
     host_port: str = Field(default="localhost:1433", description="MSSQL host URL.")
     scheme: str = Field(default="mssql+pytds", description="", hidden_from_docs=True)
+
+    # TODO: rename to include_procedures ?
     include_stored_procedures: bool = Field(
         default=True,
         description="Include ingest of stored procedures. Requires access to the 'sys' schema.",
@@ -763,9 +765,11 @@ class SQLServerSource(SQLAlchemySource):
                 yield from auto_workunit(
                     generate_procedure_lineage(
                         schema_resolver=self.get_schema_resolver(),
-                        procedure=procedure,
+                        procedure=procedure.to_base_procedure(),
                         procedure_job_urn=MSSQLDataJob(entity=procedure).urn,
                         is_temp_table=self.is_temp_table,
+                        default_db=procedure.db,
+                        default_schema=procedure.schema,
                     )
                 )
 

datahub/ingestion/source/sql/stored_procedures/base.py

@@ -0,0 +1,242 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Callable, Dict, Iterable, Optional
+
+from datahub.emitter.mce_builder import (
+    DEFAULT_ENV,
+    datahub_guid,
+    make_data_flow_urn,
+    make_data_job_urn,
+    make_data_platform_urn,
+    make_dataplatform_instance_urn,
+)
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.mcp_builder import DatabaseKey, SchemaKey
+from datahub.ingestion.api.source_helpers import auto_workunit
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import (
+    FlowContainerSubTypes,
+    JobContainerSubTypes,
+)
+from datahub.ingestion.source.sql.stored_procedures.lineage import parse_procedure_code
+from datahub.metadata.schema_classes import (
+    ContainerClass,
+    DataFlowInfoClass,
+    DataJobInfoClass,
+    DataPlatformInstanceClass,
+    DataTransformClass,
+    DataTransformLogicClass,
+    QueryStatementClass,
+    SubTypesClass,
+)
+from datahub.sql_parsing.schema_resolver import SchemaResolver
+
+
+@dataclass
+class BaseProcedure:
+    name: str
+    procedure_definition: Optional[str]
+    created: Optional[datetime]
+    last_altered: Optional[datetime]
+    comment: Optional[str]
+    argument_signature: Optional[str]
+    return_type: Optional[str]
+    language: str
+    extra_properties: Optional[Dict[str, str]]
+
+    def get_procedure_identifier(
+        self,
+    ) -> str:
+        if self.argument_signature:
+            argument_signature_hash = datahub_guid(
+                dict(argument_signature=self.argument_signature)
+            )
+            return f"{self.name}_{argument_signature_hash}"
+
+        return self.name
+
+    def to_urn(self, database_key: DatabaseKey, schema_key: Optional[SchemaKey]) -> str:
+        return make_data_job_urn(
+            orchestrator=database_key.platform,
+            flow_id=_get_procedure_flow_name(database_key, schema_key),
+            job_id=self.get_procedure_identifier(),
+            cluster=database_key.env or DEFAULT_ENV,
+            platform_instance=database_key.instance,
+        )
+
+
+def _generate_flow_workunits(
+    database_key: DatabaseKey, schema_key: Optional[SchemaKey]
+) -> Iterable[MetadataWorkUnit]:
+    """Generate flow workunits for database and schema"""
+
+    procedure_flow_name = _get_procedure_flow_name(database_key, schema_key)
+
+    flow_urn = make_data_flow_urn(
+        orchestrator=database_key.platform,
+        flow_id=procedure_flow_name,
+        cluster=database_key.env or DEFAULT_ENV,
+        platform_instance=database_key.instance,
+    )
+
+    yield MetadataChangeProposalWrapper(
+        entityUrn=flow_urn,
+        aspect=DataFlowInfoClass(
+            name=procedure_flow_name,
+        ),
+    ).as_workunit()
+
+    yield MetadataChangeProposalWrapper(
+        entityUrn=flow_urn,
+        aspect=SubTypesClass(
+            typeNames=[FlowContainerSubTypes.MSSQL_PROCEDURE_CONTAINER],
+        ),
+    ).as_workunit()
+
+    if database_key.instance:
+        yield MetadataChangeProposalWrapper(
+            entityUrn=flow_urn,
+            aspect=DataPlatformInstanceClass(
+                platform=make_data_platform_urn(database_key.platform),
+                instance=make_dataplatform_instance_urn(
+                    platform=database_key.platform,
+                    instance=database_key.instance,
+                ),
+            ),
+        ).as_workunit()
+
+    yield MetadataChangeProposalWrapper(
+        entityUrn=flow_urn,
+        aspect=ContainerClass(container=database_key.as_urn()),
+    ).as_workunit()
+
+
+def _get_procedure_flow_name(
+    database_key: DatabaseKey, schema_key: Optional[SchemaKey]
+) -> str:
+    if schema_key:
+        procedure_flow_name = (
+            f"{schema_key.database}.{schema_key.db_schema}.stored_procedures"
+        )
+    else:
+        procedure_flow_name = f"{database_key.database}.stored_procedures"
+    return procedure_flow_name
+
+
+def _generate_job_workunits(
+    database_key: DatabaseKey,
+    schema_key: Optional[SchemaKey],
+    procedure: BaseProcedure,
+) -> Iterable[MetadataWorkUnit]:
+    """Generate job workunits for database, schema and procedure"""
+
+    job_urn = procedure.to_urn(database_key, schema_key)
+
+    yield MetadataChangeProposalWrapper(
+        entityUrn=job_urn,
+        aspect=DataJobInfoClass(
+            name=procedure.name,
+            type=JobContainerSubTypes.STORED_PROCEDURE,
+            description=procedure.comment,
+            customProperties=procedure.extra_properties,
+        ),
+    ).as_workunit()
+
+    yield MetadataChangeProposalWrapper(
+        entityUrn=job_urn,
+        aspect=SubTypesClass(
+            typeNames=[JobContainerSubTypes.STORED_PROCEDURE],
+        ),
+    ).as_workunit()
+
+    if database_key.instance:
+        yield MetadataChangeProposalWrapper(
+            entityUrn=job_urn,
+            aspect=DataPlatformInstanceClass(
+                platform=make_data_platform_urn(database_key.platform),
+                instance=make_dataplatform_instance_urn(
+                    platform=database_key.platform,
+                    instance=database_key.instance,
+                ),
+            ),
+        ).as_workunit()
+
+    container_key = schema_key or database_key  # database_key for 2-tier
+    yield MetadataChangeProposalWrapper(
+        entityUrn=job_urn,
+        aspect=ContainerClass(container=container_key.as_urn()),
+    ).as_workunit()
+
+    # TODO: Config whether to ingest procedure code
+    if procedure.procedure_definition:
+        yield MetadataChangeProposalWrapper(
+            entityUrn=job_urn,
+            aspect=DataTransformLogicClass(
+                transforms=[
+                    DataTransformClass(
+                        queryStatement=QueryStatementClass(
+                            value=procedure.procedure_definition,
+                            language=procedure.language,
+                        ),
+                    )
+                ]
+            ),
+        ).as_workunit()
+
+
+def generate_procedure_lineage(
+    *,
+    schema_resolver: SchemaResolver,
+    procedure: BaseProcedure,
+    procedure_job_urn: str,
+    default_db: Optional[str] = None,
+    default_schema: Optional[str] = None,
+    is_temp_table: Callable[[str], bool] = lambda _: False,
+    raise_: bool = False,
+) -> Iterable[MetadataChangeProposalWrapper]:
+    if procedure.procedure_definition and procedure.language == "SQL":
+        datajob_input_output = parse_procedure_code(
+            schema_resolver=schema_resolver,
+            default_db=default_db,
+            default_schema=default_schema,
+            code=procedure.procedure_definition,
+            is_temp_table=is_temp_table,
+            raise_=raise_,
+        )
+
+        if datajob_input_output:
+            yield MetadataChangeProposalWrapper(
+                entityUrn=procedure_job_urn,
+                aspect=datajob_input_output,
+            )
+
+
+def generate_procedure_container_workunits(
+    database_key: DatabaseKey,
+    schema_key: Optional[SchemaKey],
+) -> Iterable[MetadataWorkUnit]:
+    """Generate container workunits for database and schema"""
+
+    yield from _generate_flow_workunits(database_key, schema_key)
+
+
+def generate_procedure_workunits(
+    procedure: BaseProcedure,
+    database_key: DatabaseKey,
+    schema_key: Optional[SchemaKey],
+    schema_resolver: Optional[SchemaResolver],
+) -> Iterable[MetadataWorkUnit]:
+    yield from _generate_job_workunits(database_key, schema_key, procedure)
+
+    if schema_resolver:
+        job_urn = procedure.to_urn(database_key, schema_key)
+
+        yield from auto_workunit(
+            generate_procedure_lineage(
+                schema_resolver=schema_resolver,
+                procedure=procedure,
+                procedure_job_urn=job_urn,
+                default_db=database_key.database,
+                default_schema=schema_key.db_schema if schema_key else None,
+            )
+        )

datahub/ingestion/source/sql/{mssql/stored_procedure_lineage.py → stored_procedures/lineage.py}

@@ -1,8 +1,6 @@
 import logging
-from typing import Callable, Iterable, Optional
+from typing import Callable, Optional
 
-from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.ingestion.source.sql.mssql.job_models import StoredProcedure
 from datahub.metadata.schema_classes import DataJobInputOutputClass
 from datahub.sql_parsing.datajob import to_datajob_input_output
 from datahub.sql_parsing.schema_resolver import SchemaResolver
@@ -56,29 +54,3 @@ def parse_procedure_code(
         mcps=mcps,
         ignore_extra_mcps=True,
     )
-
-
-# Is procedure handling generic enough to be added to SqlParsingAggregator?
-def generate_procedure_lineage(
-    *,
-    schema_resolver: SchemaResolver,
-    procedure: StoredProcedure,
-    procedure_job_urn: str,
-    is_temp_table: Callable[[str], bool] = lambda _: False,
-    raise_: bool = False,
-) -> Iterable[MetadataChangeProposalWrapper]:
-    if procedure.code:
-        datajob_input_output = parse_procedure_code(
-            schema_resolver=schema_resolver,
-            default_db=procedure.db,
-            default_schema=procedure.schema,
-            code=procedure.code,
-            is_temp_table=is_temp_table,
-            raise_=raise_,
-        )
-
-        if datajob_input_output:
-            yield MetadataChangeProposalWrapper(
-                entityUrn=procedure_job_urn,
-                aspect=datajob_input_output,
-            )

datahub/ingestion/source/vertexai/vertexai.py

@@ -107,7 +107,6 @@ class ContainerKeyWithId(ContainerKey):
     SourceCapability.DESCRIPTIONS,
     "Extract descriptions for Vertex AI Registered Models and Model Versions",
 )
-@capability(SourceCapability.TAGS, "Extract tags for Vertex AI Registered Model Stages")
 class VertexAISource(Source):
     platform: str = "vertexai"
 
@@ -602,6 +601,7 @@ class VertexAISource(Source):
                         else None
                     ),
                     customProperties=None,
+                    externalUrl=self._make_model_external_url(model),
                 ),
                 SubTypesClass(typeNames=[MLAssetSubTypes.VERTEX_MODEL_GROUP]),
                 ContainerClass(container=self._get_project_container().as_urn()),

datahub/metadata/schema.avsc

@@ -9453,7 +9453,9 @@
                       "fields": [
                         {
                           "Searchable": {
+                            "addToFilters": true,
                             "fieldType": "BOOLEAN",
+                            "filterNameOverride": "Deprecated",
                             "weightsPerFieldValue": {
                               "true": 0.5
                             }