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

datahub/ingestion/source/slack/slack.py

@@ -1,14 +1,16 @@
+import json
 import logging
 import textwrap
 from dataclasses import dataclass
-from typing import Iterable, List, Optional, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Tuple
 
-from pydantic import Field, SecretStr
+from pydantic import BaseModel, Field, SecretStr
 from slack_sdk import WebClient
 from tenacity import retry, wait_exponential
 from tenacity.before_sleep import before_sleep_log
 
 import datahub.emitter.mce_builder as builder
+from datahub.emitter.mce_builder import datahub_guid, make_dataplatform_instance_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
@@ -32,16 +34,153 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
 )
 from datahub.metadata.schema_classes import (
     CorpUserEditableInfoClass,
+    CorpUserSettingsClass,
+    DataPlatformInstanceClass,
+    DataPlatformInstancePropertiesClass,
     DatasetPropertiesClass,
     DeprecationClass,
+    NotificationSettingsClass,
+    PlatformResourceInfoClass,
+    SerializedValueClass,
+    SerializedValueContentTypeClass,
+    SerializedValueSchemaTypeClass,
+    SlackNotificationSettingsClass,
+    SlackUserInfoClass as SlackUserInfo,
+    StatusClass,
     SubTypesClass,
+    _Aspect,
 )
 from datahub.utilities.ratelimiter import RateLimiter
+from datahub.utilities.str_enum import StrEnum
 from datahub.utilities.urns.urn import Urn
 
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+# TODO: Relocate this function to a utility module
+def is_picture_default_or_missing(picture_link: Optional[str]) -> bool:
+    if not picture_link:
+        return True
+    return picture_link.endswith("default_avatar.png")
+
+
+def is_slack_image(picture_link: Optional[str]) -> bool:
+    """
+    Guesses if the picture link is a slack image.
+    """
+    if not picture_link:
+        return False
+    return "slack-edge.com" in picture_link
+
+
+class ResourceType(StrEnum):
+    USER_INFO = "user-info"
+    CHANNEL_INFO = "channel-info"
+
+
+class SlackInstance(BaseModel):
+    id: str
+    name: Optional[str] = None
+    description: Optional[str] = None
+    external_url: Optional[str] = None
+    custom_properties: Optional[Dict[str, str]] = None
+
+    def to_platform_instance_urn(self) -> str:
+        return make_dataplatform_instance_urn(
+            platform=DATA_PLATFORM_SLACK_URN, instance=self.id
+        )
+
+    def with_slack_team_info(self, team_info: dict) -> "SlackInstance":
+        """
+        team_info looks like this
+        {'id': 'T22BUCL1LKW', 'name': 'DataHub', 'url': 'https://datahubspace.slack.com/', 'domain': 'datahub', 'email_domain': '', 'icon': {'image_default': False, 'image_34': 'https://avatars.slack-edge.com/2021-07-05/2228585180071_63e6f300a919abc488bb_34.png', 'image_44': 'https://avatars.slack-edge.com/2021-07-05/2228585180071_63e6f300a919abc488bb_44.png', 'image_68': 'https://avatars.slack-edge.com/2021-07-05/2228585180071_63e6f300a919abc488bb_68.png', 'image_88': 'https://avatars.slack-edge.com/2021-07-05/2228585180071_63e6f300a919abc488bb_88.png', 'image_102': 'https://avatars.slack-edge.com/2021-07-05/2228585180071_63e6f300a919abc488bb_102.png', 'image_230': 'https://avatars.slack-edge.com/2021-07-05/2228585180071_63e6f300a919abc488bb_230.png', 'image_132': 'https://avatars.slack-edge.com/2021-07-05/2228585180071_63e6f300a919abc488bb_132.png'}, 'avatar_base_url': 'https://ca.slack-edge.com/', 'is_verified': False, 'external_org_migrations': {'date_updated': 1722672564, 'current': []}, 'discoverable': 'closed', 'enterprise_id': 'E06TPM5T1G9', 'enterprise_name': 'DataHub', 'enterprise_domain': 'datahubspace', 'lob_sales_home_enabled': False}
+        """
+        self.name = team_info.get("name")
+        self.description = team_info.get("name")
+        self.external_url = team_info.get("url")
+        self.custom_properties = {
+            k: v
+            for k, v in {
+                "domain": team_info.get("domain"),
+                "enterprise_id": team_info.get("enterprise_id"),
+                "enterprise_name": team_info.get("enterprise_name"),
+                "enterprise_domain": team_info.get("enterprise_domain"),
+                "icon": team_info.get("icon", {}).get("image_102"),
+            }.items()
+            if v is not None
+        }
+        return self
+
+    def to_mcps(self) -> Iterable[MetadataChangeProposalWrapper]:
+        return [
+            MetadataChangeProposalWrapper(
+                entityUrn=self.to_platform_instance_urn(),
+                aspect=DataPlatformInstancePropertiesClass(
+                    name=self.name or self.id,
+                    description=self.description,
+                    externalUrl=self.external_url or None,
+                    customProperties=self.custom_properties or {},
+                ),
+            )
+        ]
+
+
+def to_serialized_value(value: _Aspect) -> SerializedValueClass:
+    # HACK: we remove the .pegasus2avro from the schema type since we want to refer to
+    # the original pdl type
+    schema_type = value.RECORD_SCHEMA.fullname.replace(".pegasus2avro", "")
+    serialized_value = SerializedValueClass(
+        blob=json.dumps(value.to_obj()).encode("utf-8"),
+        contentType=SerializedValueContentTypeClass.JSON,
+        schemaType=SerializedValueSchemaTypeClass.PEGASUS,
+        schemaRef=schema_type,
+    )
+    return serialized_value
+
+
+class SlackUserDetails:
+    def __init__(self, slack_user_info: SlackUserInfo):
+        self.slack_user_info = slack_user_info
+
+    def to_guid(self) -> str:
+        """
+        A slack user is uniquely identified by the combination of their id and teamId.
+        """
+        return datahub_guid(
+            {"id": self.slack_user_info.id, "dpi": self.slack_user_info.teamId}
+        )
+
+    def get_resource_urn(self) -> str:
+        return f"urn:li:platformResource:{self.to_guid()}"
+
+    def to_mcps(self) -> Iterable[MetadataChangeProposalWrapper]:
+        resource_urn = self.get_resource_urn()
+
+        dpi = DataPlatformInstanceClass(
+            platform=DATA_PLATFORM_SLACK_URN,
+            instance=self.slack_user_info.slackInstance,
+        )
+
+        secondary_keys = []
+        if self.slack_user_info.email:
+            secondary_keys.append(self.slack_user_info.email)
+
+        resource_info = PlatformResourceInfoClass(
+            resourceType=ResourceType.USER_INFO.value,
+            value=to_serialized_value(self.slack_user_info),
+            primaryKey=self.slack_user_info.id,
+            secondaryKeys=secondary_keys,
+        )
+
+        status = StatusClass(
+            removed=self.slack_user_info.isDeleted,
+        )
+
+        yield from MetadataChangeProposalWrapper.construct_many(
+            resource_urn, aspects=[dpi, resource_info, status]
+        )
+
+
 @dataclass
 class CorpUser:
     urn: Optional[str] = None
@@ -52,18 +191,26 @@ class CorpUser:
     phone: Optional[str] = None
     real_name: Optional[str] = None
     slack_display_name: Optional[str] = None
+    team_id: Optional[str] = None
+    team_domain: Optional[str] = None
+    is_team_enterprise: Optional[bool] = None
 
 
 class SlackSourceConfig(
     StatefulIngestionConfigBase,
 ):
     bot_token: SecretStr = Field(
-        description="Bot token for the Slack workspace. Needs `users:read`, `users:read.email` and `users.profile:read` scopes.",
+        description="Bot token for the Slack workspace. Needs `users:read`, `users:read.email`, `users.profile:read`, and `team:read` scopes.",
     )
     enrich_user_metadata: bool = Field(
         type=bool,
         default=True,
-        description="Whether to enrich user metadata.",
+        description="When enabled, will enrich provisioned DataHub users' metadata with information from Slack.",
+    )
+    ingest_users: bool = Field(
+        type=bool,
+        default=True,
+        description="Whether to ingest users. When set to true, will ingest all users in the Slack workspace (as platform resources) to simplify user enrichment after they are provisioned on DataHub.",
     )
     api_requests_per_min: int = Field(
         type=int,
@@ -96,9 +243,11 @@ class SlackSourceConfig(
 class SlackSourceReport(StaleEntityRemovalSourceReport):
     channels_reported: int = 0
     archived_channels_reported: int = 0
+    users_reported: int = 0
 
 
 PLATFORM_NAME = "slack"
+DATA_PLATFORM_SLACK_URN: str = builder.make_data_platform_urn(PLATFORM_NAME)
 
 
 @platform_name("Slack")
@@ -124,6 +273,38 @@ class SlackSource(StatefulIngestionSourceBase):
     def get_slack_client(self) -> WebClient:
         return WebClient(token=self.config.bot_token.get_secret_value())
 
+    @staticmethod
+    def populate_slack_member_from_response(
+        user: Dict[str, Any], slack_instance: SlackInstance
+    ) -> SlackUserDetails:
+        profile = user.get("profile", {})
+
+        user_info = SlackUserInfo(
+            slackInstance=slack_instance.to_platform_instance_urn(),
+            id=user["id"],
+            name=user["name"],
+            realName=user.get("real_name", ""),
+            displayName=profile.get("display_name", ""),
+            email=profile.get("email"),
+            teamId=user["team_id"],
+            isDeleted=user.get("deleted", False),
+            isAdmin=user.get("is_admin", False),
+            isOwner=user.get("is_owner", False),
+            isPrimaryOwner=user.get("is_primary_owner", False),
+            isBot=user.get("is_bot", False),
+            timezone=user.get("tz"),
+            timezoneOffset=user.get("tz_offset"),
+            title=profile.get("title"),
+            phone=profile.get("phone"),
+            profilePictureUrl=profile.get(
+                "image_192"
+            ),  # Using 192px image as an example
+            statusText=profile.get("status_text"),
+            statusEmoji=profile.get("status_emoji"),
+            lastUpdatedSeconds=user.get("updated"),
+        )
+        return SlackUserDetails(slack_user_info=user_info)
+
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
         return [
             *super().get_workunit_processors(),
@@ -143,46 +324,108 @@ class SlackSource(StatefulIngestionSourceBase):
         logger.info(auth_resp.data)
         if self.config.ingest_public_channels:
             yield from self.get_public_channels()
-        if self.config.enrich_user_metadata:
+        if self.config.enrich_user_metadata or self.config.ingest_users:
             yield from self.get_user_info()
 
+    def _get_datahub_user_info(
+        self,
+    ) -> Dict[str, Tuple[CorpUser, Optional[CorpUserEditableInfoClass]]]:
+        # get_user_to_be_updated ensures that the email field is not None
+        users = {
+            user_obj.email: (user_obj, editable_properties)
+            for user_obj, editable_properties in self.get_user_to_be_updated()
+            if user_obj.email
+        }
+        return users
+
     def get_user_info(self) -> Iterable[MetadataWorkUnit]:
-        assert self.ctx.graph is not None
-        for user_obj in self.get_user_to_be_updated():
-            self.populate_slack_id_from_email(user_obj)
-            if user_obj.slack_id is None:
-                continue
-            self.populate_user_profile(user_obj)
-            if user_obj.urn is None:
-                continue
-            logger.info(f"User: {user_obj}")
-            corpuser_editable_info = (
-                self.ctx.graph.get_aspect(
-                    entity_urn=user_obj.urn, aspect_type=CorpUserEditableInfoClass
-                )
-                or CorpUserEditableInfoClass()
-            )
-            corpuser_editable_info.email = user_obj.email
-            corpuser_editable_info.slack = user_obj.slack_id
-            corpuser_editable_info.title = user_obj.title
-            if user_obj.image_url:
-                corpuser_editable_info.pictureLink = user_obj.image_url
-            if user_obj.phone:
-                corpuser_editable_info.phone = user_obj.phone
-            if (
-                not corpuser_editable_info.displayName
-                or corpuser_editable_info.displayName == corpuser_editable_info.email
-            ):
-                # let's fill out a real name
-                corpuser_editable_info.displayName = user_obj.real_name
-            yield MetadataWorkUnit(
-                id=f"{user_obj.urn}",
-                mcp=MetadataChangeProposalWrapper(
-                    entityUrn=user_obj.urn,
-                    aspect=corpuser_editable_info,
-                ),
+        # Get team information to populate for all users
+        slack_instance: Optional[SlackInstance] = None
+        with self.rate_limiter:
+            team_response = self.get_slack_client().team_info()
+            if team_response and "team" in team_response:
+                team_info = team_response["team"]
+                slack_instance = SlackInstance(id=team_info.get("id"))
+                slack_instance = slack_instance.with_slack_team_info(team_info)
+
+        if slack_instance:
+            for mcp in slack_instance.to_mcps():
+                yield mcp.as_workunit()
+        else:
+            logger.error("Failed to fetch team information")
+            self.report.report_failure(
+                "team_info", "Failed to fetch team information for users"
             )
 
+        assert slack_instance
+
+        # Fetch all DataHub users that need to be updated
+        if self.config.enrich_user_metadata:
+            datahub_users = self._get_datahub_user_info()
+        else:
+            datahub_users = {}
+        cursor = None
+        while True:
+            with self.rate_limiter:
+                response = self.get_slack_client().users_list(cursor=cursor)
+            assert isinstance(response.data, dict)
+            if not response.data["ok"]:
+                self.report.report_failure("users", "Failed to fetch users")
+                return
+
+            assert self.ctx.graph is not None
+            for user in response.data["members"]:
+                # Query all slack users and ingest them into the generic
+                # slackMember aspect
+                slack_user_details: SlackUserDetails = (
+                    self.populate_slack_member_from_response(user, slack_instance)
+                )
+                if self.config.ingest_users:
+                    for mcp in slack_user_details.to_mcps():
+                        yield mcp.as_workunit()
+
+                platform_resource_urn = slack_user_details.get_resource_urn()
+                # If user is in DataHub, compute and emit CorpUserEditableInfo
+                # aspect. This code will be removed once we have server side
+                # processing of raw slackMember aspects. This code path can also
+                # be turned off by setting enrich_user_metadata to False.
+                user_obj_props_tuple = datahub_users.get(user["profile"].get("email"))
+                if user_obj_props_tuple is None:
+                    # User is not in DataHub or enrichment is disabled
+                    continue
+                user_obj, editable_properties = user_obj_props_tuple
+                slack_user_profile = user.get("profile", {})
+                user_obj.slack_id = user.get("id")
+                user_obj.title = slack_user_profile.get("title")
+                user_obj.image_url = slack_user_profile.get("image_192")
+                user_obj.phone = slack_user_profile.get("phone")
+                user_obj.real_name = slack_user_profile.get("real_name")
+                user_obj.slack_display_name = slack_user_profile.get("display_name")
+                corpuser_editable_info = editable_properties or (
+                    CorpUserEditableInfoClass()
+                )
+                emittable_corpuser_editable_info = self.populate_corpuser_editable_info(
+                    corpuser_editable_info,
+                    user_obj,
+                    platform_resource_urn=platform_resource_urn,
+                    slack_instance=slack_instance,
+                )
+                if emittable_corpuser_editable_info:
+                    yield MetadataChangeProposalWrapper(
+                        entityUrn=user_obj.urn, aspect=emittable_corpuser_editable_info
+                    ).as_workunit()
+                    # if we update corpusereditable info, we also update
+                    # slackuserinfo. This will be removed once we have server
+                    # side processing of raw slackMember aspects.
+                    yield MetadataChangeProposalWrapper(
+                        entityUrn=user_obj.urn,
+                        aspect=slack_user_details.slack_user_info,
+                    ).as_workunit()
+                yield from self.emit_corp_user_slack_settings(user_obj)
+            cursor = str(response.data["response_metadata"]["next_cursor"])
+            if not cursor:
+                break
+
     def _get_channel_info(
         self, cursor: Optional[str]
     ) -> Tuple[List[MetadataWorkUnit], Optional[str]]:
@@ -259,6 +502,58 @@ class SlackSource(StatefulIngestionSourceBase):
         cursor = str(response.data["response_metadata"]["next_cursor"])
         return result_channels, cursor
 
+    def populate_corpuser_editable_info(
+        self,
+        corpuser_editable_info: CorpUserEditableInfoClass,
+        user_obj: CorpUser,
+        platform_resource_urn: str,
+        slack_instance: SlackInstance,
+    ) -> Optional[CorpUserEditableInfoClass]:
+        """
+        Populate CorpUserEditableInfo aspect with user information from Slack.
+        If changes are not required, None is returned.
+        If changes are required, the updated aspect is returned.
+        """
+        mutation_required = False
+        if not corpuser_editable_info.email and user_obj.email:
+            mutation_required = True
+            corpuser_editable_info.email = user_obj.email
+        if not corpuser_editable_info.slack and user_obj.slack_id:
+            mutation_required = True
+            corpuser_editable_info.slack = user_obj.slack_id
+        if not corpuser_editable_info.title and user_obj.title:
+            mutation_required = True
+            corpuser_editable_info.title = user_obj.title
+        if user_obj.image_url and (
+            is_picture_default_or_missing(corpuser_editable_info.pictureLink)
+            or (
+                is_slack_image(corpuser_editable_info.pictureLink)
+                and user_obj.image_url != corpuser_editable_info.pictureLink
+            )
+        ):
+            mutation_required = True
+            corpuser_editable_info.pictureLink = user_obj.image_url
+        if user_obj.phone and not corpuser_editable_info.phone:
+            mutation_required = True
+            corpuser_editable_info.phone = user_obj.phone
+        if (
+            not corpuser_editable_info.displayName
+            or corpuser_editable_info.displayName == corpuser_editable_info.email
+        ) and user_obj.real_name:
+            mutation_required = True
+            corpuser_editable_info.displayName = user_obj.real_name
+        if mutation_required:
+            # update informationSources
+            corpuser_editable_info.informationSources = (
+                []
+                if not corpuser_editable_info.informationSources
+                else corpuser_editable_info.informationSources
+            )
+            if platform_resource_urn not in corpuser_editable_info.informationSources:
+                corpuser_editable_info.informationSources.append(platform_resource_urn)
+            return corpuser_editable_info
+        return None
+
     def get_public_channels(self) -> Iterable[MetadataWorkUnit]:
         cursor = None
         while True:
@@ -270,58 +565,80 @@ class SlackSource(StatefulIngestionSourceBase):
             if not cursor:
                 break
 
-    def populate_user_profile(self, user_obj: CorpUser) -> None:
-        if not user_obj.slack_id:
+    def emit_slack_member_aspect(
+        self, user: SlackUserInfo
+    ) -> Iterable[MetadataWorkUnit]:
+        slack_user = SlackUserDetails(slack_user_info=user)
+        for mcp in slack_user.to_mcps():
+            yield mcp.as_workunit()
+
+    def emit_corp_user_slack_settings(
+        self, user_obj: CorpUser
+    ) -> Iterable[MetadataWorkUnit]:
+        assert self.ctx.graph is not None
+
+        if not user_obj.urn:
             return
-        try:
-            # https://api.slack.com/methods/users.profile.get
-            with self.rate_limiter:
-                if self._use_users_info:
-                    user_profile_res = self.get_slack_client().users_info(
-                        user=user_obj.slack_id
-                    )
-                    user_profile_res = user_profile_res.get("user", {})
-                else:
-                    user_profile_res = self.get_slack_client().users_profile_get(
-                        user=user_obj.slack_id
-                    )
-            logger.debug(f"User profile: {user_profile_res}")
-            user_profile = user_profile_res.get("profile", {})
-            user_obj.title = user_profile.get("title")
-            user_obj.image_url = user_profile.get("image_192")
-            user_obj.phone = user_profile.get("phone")
-            user_obj.real_name = user_profile.get("real_name")
-            user_obj.slack_display_name = user_profile.get("display_name")
-
-        except Exception as e:
-            if "missing_scope" in str(e):
-                if self._use_users_info:
-                    raise e
-                self._use_users_info = True
-                self.populate_user_profile(user_obj)
+
+        corp_user_settings = self.ctx.graph.get_aspect(
+            user_obj.urn, CorpUserSettingsClass
+        )
+        if not corp_user_settings:
             return
 
-    def populate_slack_id_from_email(self, user_obj: CorpUser) -> None:
-        if user_obj.email is None:
+        notification_settings = corp_user_settings.notificationSettings
+
+        if not notification_settings:
+            corp_user_settings.notificationSettings = NotificationSettingsClass(
+                sinkTypes=[],
+                slackSettings=SlackNotificationSettingsClass(
+                    userHandle=user_obj.slack_id
+                ),
+            )
+        elif (
+            not notification_settings.slackSettings
+            or not notification_settings.slackSettings.userHandle
+        ):
+            notification_settings.slackSettings = SlackNotificationSettingsClass(
+                userHandle=user_obj.slack_id
+            )
+        else:
             return
-        try:
-            # https://api.slack.com/methods/users.lookupByEmail
-            with self.rate_limiter:
-                user_info_res = self.get_slack_client().users_lookupByEmail(
-                    email=user_obj.email
-                )
-            user_info = user_info_res.get("user", {})
-            user_obj.slack_id = user_info.get("id")
-        except Exception as e:
-            if "users_not_found" in str(e):
-                return
-            raise e
+
+        yield MetadataWorkUnit(
+            id=f"{user_obj.urn}",
+            mcp=MetadataChangeProposalWrapper(
+                entityUrn=user_obj.urn,
+                aspect=corp_user_settings,
+            ),
+        )
+
+    def get_user_to_be_updated(
+        self,
+    ) -> Iterable[Tuple[CorpUser, Optional[CorpUserEditableInfoClass]]]:
+        assert self.ctx.graph is not None
+        for urn in self.ctx.graph.get_urns_by_filter(
+            entity_types=["corpuser"], query="*"
+        ):
+            user_obj = CorpUser()
+            user_obj.urn = urn
+            editable_properties = self.ctx.graph.get_aspect(
+                urn, CorpUserEditableInfoClass
+            )
+            if editable_properties and editable_properties.email:
+                user_obj.email = editable_properties.email
+            else:
+                urn_id = Urn.from_string(user_obj.urn).get_entity_id_as_string()
+                if "@" in urn_id:
+                    user_obj.email = urn_id
+            if user_obj.email is not None:
+                yield (user_obj, editable_properties)
 
     @retry(
         wait=wait_exponential(multiplier=2, min=4, max=60),
         before_sleep=before_sleep_log(logger, logging.ERROR, True),
     )
-    def get_user_to_be_updated(self) -> Iterable[CorpUser]:
+    def get_user_to_be_updated_oss(self) -> Iterable[CorpUser]:
         graphql_query = textwrap.dedent(
             """
             query listUsers($input: ListUsersInput!) {

datahub/ingestion/source/snowflake/constants.py

@@ -54,6 +54,7 @@ class SnowflakeObjectDomain(StrEnum):
     COLUMN = "column"
     ICEBERG_TABLE = "iceberg table"
     STREAM = "stream"
+    PROCEDURE = "procedure"
 
 
 GENERIC_PERMISSION_ERROR_KEY = "permission-error"

datahub/ingestion/source/snowflake/snowflake_config.py

@@ -100,7 +100,15 @@ class SnowflakeFilterConfig(SQLFilterConfig):
 
     stream_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern.allow_all(),
-        description="Regex patterns for streams to filter in ingestion. Note: Defaults to table_pattern if not specified. Specify regex to match the entire view name in database.schema.view format. e.g. to match all views starting with customer in Customer database and public schema, use the regex 'Customer.public.customer.*'",
+        description="Regex patterns for streams to filter in ingestion. Specify regex to match the entire view name in database.schema.view format. e.g. to match all views starting with customer in Customer database and public schema, use the regex 'Customer.public.customer.*'",
+    )
+
+    procedure_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex patterns for procedures to filter in ingestion. "
+        "Specify regex to match the entire procedure name in database.schema.procedure format. "
+        "e.g. to match all procedures starting with customer in Customer database and public schema,"
+        " use the regex 'Customer.public.customer.*'",
     )
 
     match_fully_qualified_names: bool = Field(
@@ -284,6 +292,11 @@ class SnowflakeV2Config(
         description="If enabled, streams will be ingested as separate entities from tables/views.",
     )
 
+    include_procedures: bool = Field(
+        default=True,
+        description="If enabled, procedures will be ingested as pipelines/tasks.",
+    )
+
     structured_property_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern.allow_all(),
         description=(

datahub/ingestion/source/snowflake/snowflake_query.py

@@ -164,6 +164,23 @@ class SnowflakeQuery:
         and table_type in ('BASE TABLE', 'EXTERNAL TABLE')
         order by table_schema, table_name"""
 
+    @staticmethod
+    def procedures_for_database(db_name: Optional[str]) -> str:
+        db_clause = f'"{db_name}".' if db_name is not None else ""
+        return f"""
+        SELECT procedure_catalog AS "PROCEDURE_CATALOG",
+        procedure_schema AS "PROCEDURE_SCHEMA",
+        procedure_name AS "PROCEDURE_NAME",
+        procedure_language AS "PROCEDURE_LANGUAGE",
+        argument_signature AS "ARGUMENT_SIGNATURE",
+        data_type AS "PROCEDURE_RETURN_TYPE",
+        procedure_definition AS "PROCEDURE_DEFINITION",
+        created AS "CREATED",
+        last_altered AS "LAST_ALTERED",
+        comment AS "COMMENT"
+        FROM {db_clause}information_schema.procedures
+        order by procedure_schema, procedure_name"""
+
     @staticmethod
     def get_all_tags():
         return """

datahub/ingestion/source/snowflake/snowflake_report.py

@@ -105,6 +105,7 @@ class SnowflakeV2Report(
     databases_scanned: int = 0
     tags_scanned: int = 0
     streams_scanned: int = 0
+    procedures_scanned: int = 0
 
     include_usage_stats: bool = False
     include_operational_stats: bool = False
@@ -163,6 +164,8 @@ class SnowflakeV2Report(
             self.tags_scanned += 1
         elif ent_type == "stream":
             self.streams_scanned += 1
+        elif ent_type == "procedure":
+            self.procedures_scanned += 1
         else:
             raise KeyError(f"Unknown entity {ent_type}.")