acryl-datahub 0.15.0.4rc3__py3-none-any.whl → 0.15.0.5__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_config.py

@@ -249,6 +249,12 @@ class SnowflakeV2Config(
         description="If enabled along with `extract_tags`, extracts snowflake's key-value tags as DataHub structured properties instead of DataHub tags.",
     )
 
+    structured_properties_template_cache_invalidation_interval: int = Field(
+        hidden_from_docs=True,
+        default=60,
+        description="Interval in seconds to invalidate the structured properties template cache.",
+    )
+
     include_external_url: bool = Field(
         default=True,
         description="Whether to populate Snowsight url for Snowflake Objects",
@@ -302,6 +308,13 @@ class SnowflakeV2Config(
         " assertions CLI in snowflake",
     )
 
+    pushdown_deny_usernames: List[str] = Field(
+        default=[],
+        description="List of snowflake usernames which will not be considered for lineage/usage/queries extraction. "
+        "This is primarily useful for improving performance by filtering out users with extremely high query volumes. "
+        "Only applicable if `use_queries_v2` is enabled.",
+    )
+
     @validator("convert_urns_to_lowercase")
     def validate_convert_urns_to_lowercase(cls, v):
         if not v:

datahub/ingestion/source/snowflake/snowflake_query.py

@@ -159,6 +159,17 @@ class SnowflakeQuery:
         and table_type in ('BASE TABLE', 'EXTERNAL TABLE', 'HYBRID TABLE')
         order by table_schema, table_name"""
 
+    @staticmethod
+    def get_all_tags():
+        return """
+        SELECT tag_database as "TAG_DATABASE",
+        tag_schema AS "TAG_SCHEMA",
+        tag_name AS "TAG_NAME",
+        FROM snowflake.account_usage.tag_references
+        GROUP BY TAG_DATABASE , TAG_SCHEMA, tag_name
+        ORDER BY TAG_DATABASE, TAG_SCHEMA, TAG_NAME  ASC;
+        """
+
     @staticmethod
     def get_all_tags_on_object_with_propagation(
         db_name: str, quoted_identifier: str, domain: str

datahub/ingestion/source/snowflake/snowflake_report.py

@@ -12,6 +12,7 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
 from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.ingestion.source_report.time_window import BaseTimeWindowReport
 from datahub.sql_parsing.sql_parsing_aggregator import SqlAggregatorReport
+from datahub.utilities.lossy_collections import LossyDict
 from datahub.utilities.perf_timer import PerfTimer
 
 if TYPE_CHECKING:
@@ -66,7 +67,7 @@ class SnowflakeReport(SQLSourceReport, BaseTimeWindowReport):
     num_external_table_edges_scanned: int = 0
     ignore_start_time_lineage: Optional[bool] = None
     upstream_lineage_in_report: Optional[bool] = None
-    upstream_lineage: Dict[str, List[str]] = field(default_factory=dict)
+    upstream_lineage: LossyDict[str, List[str]] = field(default_factory=LossyDict)
 
     lineage_start_time: Optional[datetime] = None
     lineage_end_time: Optional[datetime] = None
@@ -114,6 +115,7 @@ class SnowflakeV2Report(
     num_tables_with_known_upstreams: int = 0
     num_upstream_lineage_edge_parsing_failed: int = 0
     num_secure_views_missing_definition: int = 0
+    num_structured_property_templates_created: int = 0
 
     data_dictionary_cache: Optional["SnowflakeDataDictionary"] = None
 

datahub/ingestion/source/snowflake/snowflake_schema.py

@@ -285,6 +285,23 @@ class SnowflakeDataDictionary(SupportsAsObj):
 
         return secure_view_definitions
 
+    def get_all_tags(self) -> List[SnowflakeTag]:
+        cur = self.connection.query(
+            SnowflakeQuery.get_all_tags(),
+        )
+
+        tags = [
+            SnowflakeTag(
+                database=tag["TAG_DATABASE"],
+                schema=tag["TAG_SCHEMA"],
+                name=tag["TAG_NAME"],
+                value="",
+            )
+            for tag in cur
+        ]
+
+        return tags
+
     @serialized_lru_cache(maxsize=1)
     def get_tables_for_database(
         self, db_name: str

datahub/ingestion/source/snowflake/snowflake_schema_gen.py

@@ -1,10 +1,10 @@
 import itertools
 import logging
+import time
 from typing import Dict, Iterable, List, Optional, Union
 
 from datahub.configuration.pattern_utils import is_schema_allowed
 from datahub.emitter.mce_builder import (
-    get_sys_time,
     make_data_platform_urn,
     make_dataset_urn_with_platform_instance,
     make_schema_field_urn,
@@ -74,7 +74,6 @@ from datahub.ingestion.source_report.ingestion_stage import (
     PROFILING,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import (
-    AuditStamp,
     GlobalTags,
     Status,
     SubTypes,
@@ -101,15 +100,8 @@ from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     StringType,
     TimeType,
 )
-from datahub.metadata.com.linkedin.pegasus2avro.structured import (
-    StructuredPropertyDefinition,
-)
 from datahub.metadata.com.linkedin.pegasus2avro.tag import TagProperties
 from datahub.metadata.urns import (
-    ContainerUrn,
-    DatasetUrn,
-    DataTypeUrn,
-    EntityTypeUrn,
     SchemaFieldUrn,
     StructuredPropertyUrn,
 )
@@ -191,7 +183,7 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         self.domain_registry: Optional[DomainRegistry] = domain_registry
         self.classification_handler = ClassificationHandler(self.config, self.report)
         self.tag_extractor = SnowflakeTagExtractor(
-            config, self.data_dictionary, self.report
+            config, self.data_dictionary, self.report, identifiers
         )
         self.profiler: Optional[SnowflakeProfiler] = profiler
         self.snowsight_url_builder: Optional[SnowsightUrlBuilder] = (
@@ -217,6 +209,16 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         return self.identifiers.snowflake_identifier(identifier)
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+        if self.config.extract_tags_as_structured_properties:
+            logger.info("Creating structured property templates for tags")
+            yield from self.tag_extractor.create_structured_property_templates()
+            # We have to wait until cache invalidates to make sure the structured property template is available
+            logger.info(
+                f"Waiting for {self.config.structured_properties_template_cache_invalidation_interval} seconds for structured properties cache to invalidate"
+            )
+            time.sleep(
+                self.config.structured_properties_template_cache_invalidation_interval
+            )
         self.databases = []
         for database in self.get_databases() or []:
             self.report.report_entity_scanned(database.name, "database")
@@ -491,15 +493,25 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         try:
             view_definitions = self.data_dictionary.get_secure_view_definitions()
             return view_definitions[db_name][schema_name][table_name]
+        except KeyError:
+            # Received secure view definitions but the view is not present in results
+            self.structured_reporter.info(
+                title="Secure view definition not found",
+                message="Lineage will be missing for the view.",
+                context=f"{db_name}.{schema_name}.{table_name}",
+            )
+            return None
         except Exception as e:
-            if isinstance(e, SnowflakePermissionError):
-                error_msg = (
-                    "Failed to get secure views definitions. Please check permissions."
-                )
-            else:
-                error_msg = "Failed to get secure views definitions"
+            action_msg = (
+                "Please check permissions."
+                if isinstance(e, SnowflakePermissionError)
+                else ""
+            )
+
             self.structured_reporter.warning(
-                error_msg,
+                title="Failed to get secure views definitions",
+                message=f"Lineage will be missing for the view. {action_msg}",
+                context=f"{db_name}.{schema_name}.{table_name}",
                 exc=e,
             )
             return None
@@ -688,6 +700,7 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
 
     def _process_tag(self, tag: SnowflakeTag) -> Iterable[MetadataWorkUnit]:
         use_sp = self.config.extract_tags_as_structured_properties
+
         identifier = (
             self.snowflake_identifier(tag.structured_property_identifier())
             if use_sp
@@ -698,10 +711,11 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
             return
 
         self.report.report_tag_processed(identifier)
+
         if use_sp:
-            yield from self.gen_tag_as_structured_property_workunits(tag)
-        else:
-            yield from self.gen_tag_workunits(tag)
+            return
+
+        yield from self.gen_tag_workunits(tag)
 
     def _format_tags_as_structured_properties(
         self, tags: List[SnowflakeTag]
@@ -722,6 +736,7 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         if table.tags:
             for tag in table.tags:
                 yield from self._process_tag(tag)
+
         for column_name in table.column_tags:
             for tag in table.column_tags[column_name]:
                 yield from self._process_tag(tag)
@@ -893,29 +908,6 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
             entityUrn=tag_urn, aspect=tag_properties_aspect
         ).as_workunit()
 
-    def gen_tag_as_structured_property_workunits(
-        self, tag: SnowflakeTag
-    ) -> Iterable[MetadataWorkUnit]:
-        identifier = self.snowflake_identifier(tag.structured_property_identifier())
-        urn = StructuredPropertyUrn(identifier).urn()
-        aspect = StructuredPropertyDefinition(
-            qualifiedName=identifier,
-            displayName=tag.name,
-            valueType=DataTypeUrn("datahub.string").urn(),
-            entityTypes=[
-                EntityTypeUrn(f"datahub.{ContainerUrn.ENTITY_TYPE}").urn(),
-                EntityTypeUrn(f"datahub.{DatasetUrn.ENTITY_TYPE}").urn(),
-                EntityTypeUrn(f"datahub.{SchemaFieldUrn.ENTITY_TYPE}").urn(),
-            ],
-            lastModified=AuditStamp(
-                time=get_sys_time(), actor="urn:li:corpuser:datahub"
-            ),
-        )
-        yield MetadataChangeProposalWrapper(
-            entityUrn=urn,
-            aspect=aspect,
-        ).as_workunit()
-
     def gen_column_tags_as_structured_properties(
         self, dataset_urn: str, table: Union[SnowflakeTable, SnowflakeView]
     ) -> Iterable[MetadataWorkUnit]:

datahub/ingestion/source/snowflake/snowflake_tag.py

@@ -1,6 +1,9 @@
 import logging
-from typing import Dict, List, Optional
+from typing import Dict, Iterable, List, Optional
 
+from datahub.emitter.mce_builder import get_sys_time
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.snowflake.constants import SnowflakeObjectDomain
 from datahub.ingestion.source.snowflake.snowflake_config import (
     SnowflakeV2Config,
@@ -12,7 +15,22 @@ from datahub.ingestion.source.snowflake.snowflake_schema import (
     SnowflakeTag,
     _SnowflakeTagCache,
 )
-from datahub.ingestion.source.snowflake.snowflake_utils import SnowflakeCommonMixin
+from datahub.ingestion.source.snowflake.snowflake_utils import (
+    SnowflakeCommonMixin,
+    SnowflakeIdentifierBuilder,
+)
+from datahub.metadata.com.linkedin.pegasus2avro.common import AuditStamp
+from datahub.metadata.com.linkedin.pegasus2avro.structured import (
+    StructuredPropertyDefinition,
+)
+from datahub.metadata.urns import (
+    ContainerUrn,
+    DatasetUrn,
+    DataTypeUrn,
+    EntityTypeUrn,
+    SchemaFieldUrn,
+    StructuredPropertyUrn,
+)
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -23,11 +41,12 @@ class SnowflakeTagExtractor(SnowflakeCommonMixin):
         config: SnowflakeV2Config,
         data_dictionary: SnowflakeDataDictionary,
         report: SnowflakeV2Report,
+        snowflake_identifiers: SnowflakeIdentifierBuilder,
     ) -> None:
         self.config = config
         self.data_dictionary = data_dictionary
         self.report = report
-
+        self.snowflake_identifiers = snowflake_identifiers
         self.tag_cache: Dict[str, _SnowflakeTagCache] = {}
 
     def _get_tags_on_object_without_propagation(
@@ -59,6 +78,41 @@ class SnowflakeTagExtractor(SnowflakeCommonMixin):
             raise ValueError(f"Unknown domain {domain}")
         return tags
 
+    def create_structured_property_templates(self) -> Iterable[MetadataWorkUnit]:
+        for tag in self.data_dictionary.get_all_tags():
+            if not self.config.structured_property_pattern.allowed(
+                tag.tag_identifier()
+            ):
+                continue
+            if self.config.extract_tags_as_structured_properties:
+                self.report.num_structured_property_templates_created += 1
+                yield from self.gen_tag_as_structured_property_workunits(tag)
+
+    def gen_tag_as_structured_property_workunits(
+        self, tag: SnowflakeTag
+    ) -> Iterable[MetadataWorkUnit]:
+        identifier = self.snowflake_identifiers.snowflake_identifier(
+            tag.structured_property_identifier()
+        )
+        urn = StructuredPropertyUrn(identifier).urn()
+        aspect = StructuredPropertyDefinition(
+            qualifiedName=identifier,
+            displayName=tag.name,
+            valueType=DataTypeUrn("datahub.string").urn(),
+            entityTypes=[
+                EntityTypeUrn(f"datahub.{ContainerUrn.ENTITY_TYPE}").urn(),
+                EntityTypeUrn(f"datahub.{DatasetUrn.ENTITY_TYPE}").urn(),
+                EntityTypeUrn(f"datahub.{SchemaFieldUrn.ENTITY_TYPE}").urn(),
+            ],
+            lastModified=AuditStamp(
+                time=get_sys_time(), actor="urn:li:corpuser:datahub"
+            ),
+        )
+        yield MetadataChangeProposalWrapper(
+            entityUrn=urn,
+            aspect=aspect,
+        ).as_workunit()
+
     def _get_tags_on_object_with_propagation(
         self,
         domain: str,

datahub/ingestion/source/snowflake/snowflake_v2.py

@@ -5,6 +5,7 @@ import logging
 import os
 import os.path
 import platform
+import re
 from dataclasses import dataclass
 from typing import Dict, Iterable, List, Optional, Union
 
@@ -33,6 +34,7 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.snowflake.constants import (
     GENERIC_PERMISSION_ERROR_KEY,
     SnowflakeEdition,
+    SnowflakeObjectDomain,
 )
 from datahub.ingestion.source.snowflake.snowflake_assertion import (
     SnowflakeAssertionsHandler,
@@ -162,6 +164,8 @@ class SnowflakeV2Source(
         self.data_dictionary = SnowflakeDataDictionary(connection=self.connection)
         self.lineage_extractor: Optional[SnowflakeLineageExtractor] = None
 
+        self.discovered_datasets: Optional[List[str]] = None
+
         self.aggregator: SqlParsingAggregator = self._exit_stack.enter_context(
             SqlParsingAggregator(
                 platform=self.identifiers.platform,
@@ -182,6 +186,8 @@ class SnowflakeV2Source(
                 generate_usage_statistics=False,
                 generate_operations=False,
                 format_queries=self.config.format_sql_queries,
+                is_temp_table=self._is_temp_table,
+                is_allowed_table=self._is_allowed_table,
             )
         )
         self.report.sql_aggregator = self.aggregator.report
@@ -444,6 +450,34 @@ class SnowflakeV2Source(
 
         return _report
 
+    def _is_temp_table(self, name: str) -> bool:
+        if any(
+            re.match(pattern, name, flags=re.IGNORECASE)
+            for pattern in self.config.temporary_tables_pattern
+        ):
+            return True
+
+        # This is also a temp table if
+        #   1. this name would be allowed by the dataset patterns, and
+        #   2. we have a list of discovered tables, and
+        #   3. it's not in the discovered tables list
+        if (
+            self.filters.is_dataset_pattern_allowed(name, SnowflakeObjectDomain.TABLE)
+            and self.discovered_datasets
+            and name not in self.discovered_datasets
+        ):
+            return True
+
+        return False
+
+    def _is_allowed_table(self, name: str) -> bool:
+        if self.discovered_datasets and name not in self.discovered_datasets:
+            return False
+
+        return self.filters.is_dataset_pattern_allowed(
+            name, SnowflakeObjectDomain.TABLE
+        )
+
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
         return [
             *super().get_workunit_processors(),
@@ -513,7 +547,7 @@ class SnowflakeV2Source(
             )
             return
 
-        discovered_datasets = discovered_tables + discovered_views
+        self.discovered_datasets = discovered_tables + discovered_views
 
         if self.config.use_queries_v2:
             with self.report.new_stage(f"*: {VIEW_PARSING}"):
@@ -533,18 +567,20 @@ class SnowflakeV2Source(
                         include_queries=self.config.include_queries,
                         include_query_usage_statistics=self.config.include_query_usage_statistics,
                         user_email_pattern=self.config.user_email_pattern,
+                        pushdown_deny_usernames=self.config.pushdown_deny_usernames,
                     ),
                     structured_report=self.report,
                     filters=self.filters,
                     identifiers=self.identifiers,
                     schema_resolver=schema_resolver,
-                    discovered_tables=discovered_datasets,
+                    discovered_tables=self.discovered_datasets,
                     graph=self.ctx.graph,
                 )
 
                 # TODO: This is slightly suboptimal because we create two SqlParsingAggregator instances with different configs
                 # but a shared schema resolver. That's fine for now though - once we remove the old lineage/usage extractors,
                 # it should be pretty straightforward to refactor this and only initialize the aggregator once.
+                # This also applies for the _is_temp_table and _is_allowed_table methods above, duplicated from SnowflakeQueriesExtractor.
                 self.report.queries_extractor = queries_extractor.report
                 yield from queries_extractor.get_workunits_internal()
                 queries_extractor.close()
@@ -568,12 +604,14 @@ class SnowflakeV2Source(
             if (
                 self.config.include_usage_stats or self.config.include_operational_stats
             ) and self.usage_extractor:
-                yield from self.usage_extractor.get_usage_workunits(discovered_datasets)
+                yield from self.usage_extractor.get_usage_workunits(
+                    self.discovered_datasets
+                )
 
         if self.config.include_assertion_results:
             yield from SnowflakeAssertionsHandler(
                 self.config, self.report, self.connection, self.identifiers
-            ).get_assertion_workunits(discovered_datasets)
+            ).get_assertion_workunits(self.discovered_datasets)
 
         self.connection.close()
 

datahub/ingestion/source/sql/clickhouse.py

@@ -53,7 +53,6 @@ from datahub.metadata.com.linkedin.pegasus2avro.schema import (
 )
 from datahub.metadata.schema_classes import (
     DatasetLineageTypeClass,
-    DatasetPropertiesClass,
     DatasetSnapshotClass,
     UpstreamClass,
 )
@@ -418,41 +417,11 @@ class ClickHouseSource(TwoTierSQLAlchemySource):
                 dataset_snapshot: DatasetSnapshotClass = wu.metadata.proposedSnapshot
                 assert dataset_snapshot
 
-                lineage_mcp, lineage_properties_aspect = self.get_lineage_mcp(
-                    wu.metadata.proposedSnapshot.urn
-                )
+                lineage_mcp = self.get_lineage_mcp(wu.metadata.proposedSnapshot.urn)
 
                 if lineage_mcp is not None:
                     yield lineage_mcp.as_workunit()
 
-                if lineage_properties_aspect:
-                    aspects = dataset_snapshot.aspects
-                    if aspects is None:
-                        aspects = []
-
-                    dataset_properties_aspect: Optional[DatasetPropertiesClass] = None
-
-                    for aspect in aspects:
-                        if isinstance(aspect, DatasetPropertiesClass):
-                            dataset_properties_aspect = aspect
-
-                    if dataset_properties_aspect is None:
-                        dataset_properties_aspect = DatasetPropertiesClass()
-                        aspects.append(dataset_properties_aspect)
-
-                    custom_properties = (
-                        {
-                            **dataset_properties_aspect.customProperties,
-                            **lineage_properties_aspect.customProperties,
-                        }
-                        if dataset_properties_aspect.customProperties
-                        else lineage_properties_aspect.customProperties
-                    )
-                    dataset_properties_aspect.customProperties = custom_properties
-                    dataset_snapshot.aspects = aspects
-
-                    dataset_snapshot.aspects.append(dataset_properties_aspect)
-
             # Emit the work unit from super.
             yield wu
 
@@ -656,19 +625,16 @@ class ClickHouseSource(TwoTierSQLAlchemySource):
 
     def get_lineage_mcp(
         self, dataset_urn: str
-    ) -> Tuple[
-        Optional[MetadataChangeProposalWrapper], Optional[DatasetPropertiesClass]
-    ]:
+    ) -> Optional[MetadataChangeProposalWrapper]:
         dataset_key = mce_builder.dataset_urn_to_key(dataset_urn)
         if dataset_key is None:
-            return None, None
+            return None
 
         if not self._lineage_map:
             self._populate_lineage()
         assert self._lineage_map is not None
 
         upstream_lineage: List[UpstreamClass] = []
-        custom_properties: Dict[str, str] = {}
 
         if dataset_key.name in self._lineage_map:
             item = self._lineage_map[dataset_key.name]
@@ -684,16 +650,12 @@ class ClickHouseSource(TwoTierSQLAlchemySource):
                 )
                 upstream_lineage.append(upstream_table)
 
-        properties = None
-        if custom_properties:
-            properties = DatasetPropertiesClass(customProperties=custom_properties)
-
         if not upstream_lineage:
-            return None, properties
+            return None
 
         mcp = MetadataChangeProposalWrapper(
             entityUrn=dataset_urn,
             aspect=UpstreamLineage(upstreams=upstream_lineage),
         )
 
-        return mcp, properties
+        return mcp

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

@@ -7,7 +7,12 @@ from datahub.emitter.mce_builder import (
     make_data_platform_urn,
     make_dataplatform_instance_urn,
 )
+from datahub.emitter.mcp_builder import (
+    DatabaseKey,
+    SchemaKey,
+)
 from datahub.metadata.schema_classes import (
+    ContainerClass,
     DataFlowInfoClass,
     DataJobInfoClass,
     DataJobInputOutputClass,
@@ -171,11 +176,7 @@ class MSSQLDataJob:
             flow_id=self.entity.flow.formatted_name,
             job_id=self.entity.formatted_name,
             cluster=self.entity.flow.cluster,
-            platform_instance=(
-                self.entity.flow.platform_instance
-                if self.entity.flow.platform_instance
-                else None
-            ),
+            platform_instance=self.entity.flow.platform_instance,
         )
 
     def add_property(
@@ -222,6 +223,26 @@ class MSSQLDataJob:
             )
         return None
 
+    @property
+    def as_container_aspect(self) -> ContainerClass:
+        key_args = dict(
+            platform=self.entity.flow.orchestrator,
+            instance=self.entity.flow.platform_instance,
+            env=self.entity.flow.env,
+            database=self.entity.flow.db,
+        )
+        container_key = (
+            SchemaKey(
+                schema=self.entity.schema,
+                **key_args,
+            )
+            if isinstance(self.entity, StoredProcedure)
+            else DatabaseKey(
+                **key_args,
+            )
+        )
+        return ContainerClass(container=container_key.as_urn())
+
 
 @dataclass
 class MSSQLDataFlow:
@@ -244,9 +265,7 @@ class MSSQLDataFlow:
             orchestrator=self.entity.orchestrator,
             flow_id=self.entity.formatted_name,
             cluster=self.entity.cluster,
-            platform_instance=(
-                self.entity.platform_instance if self.entity.platform_instance else None
-            ),
+            platform_instance=self.entity.platform_instance,
         )
 
     @property
@@ -267,3 +286,13 @@ class MSSQLDataFlow:
                 ),
             )
         return None
+
+    @property
+    def as_container_aspect(self) -> ContainerClass:
+        databaseKey = DatabaseKey(
+            platform=self.entity.orchestrator,
+            instance=self.entity.platform_instance,
+            env=self.entity.env,
+            database=self.entity.db,
+        )
+        return ContainerClass(container=databaseKey.as_urn())

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

@@ -108,6 +108,10 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
         default=True,
         description="Enable lineage extraction for stored procedures",
     )
+    include_containers_for_pipelines: bool = Field(
+        default=False,
+        description="Enable the container aspects ingestion for both pipelines and tasks. Note that this feature requires the corresponding model support in the backend, which was introduced in version 0.15.0.1.",
+    )
 
     @pydantic.validator("uri_args")
     def passwords_match(cls, v, values, **kwargs):
@@ -641,6 +645,12 @@ class SQLServerSource(SQLAlchemySource):
                 aspect=data_platform_instance_aspect,
             ).as_workunit()
 
+        if self.config.include_containers_for_pipelines:
+            yield MetadataChangeProposalWrapper(
+                entityUrn=data_job.urn,
+                aspect=data_job.as_container_aspect,
+            ).as_workunit()
+
         if include_lineage:
             yield MetadataChangeProposalWrapper(
                 entityUrn=data_job.urn,
@@ -683,6 +693,13 @@ class SQLServerSource(SQLAlchemySource):
                 entityUrn=data_flow.urn,
                 aspect=data_platform_instance_aspect,
             ).as_workunit()
+
+        if self.config.include_containers_for_pipelines:
+            yield MetadataChangeProposalWrapper(
+                entityUrn=data_flow.urn,
+                aspect=data_flow.as_container_aspect,
+            ).as_workunit()
+
         # TODO: Add SubType when it appear
 
     def get_inspectors(self) -> Iterable[Inspector]: