acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/sql/mysql.py

@@ -1,12 +1,18 @@
 # This import verifies that the dependencies are available.
+import logging
+from typing import TYPE_CHECKING, Any, List, Optional
 
 import pymysql  # noqa: F401
 from pydantic.fields import Field
-from sqlalchemy import util
+from sqlalchemy import create_engine, event, inspect, util
 from sqlalchemy.dialects.mysql import BIT, base
 from sqlalchemy.dialects.mysql.enumerated import SET
 from sqlalchemy.engine.reflection import Inspector
 
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Engine
+
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.ingestion.api.decorators import (
     SourceCapability,
     SupportStatus,
@@ -15,16 +21,27 @@ from datahub.ingestion.api.decorators import (
     platform_name,
     support_status,
 )
+from datahub.ingestion.source.aws.aws_common import (
+    AwsConnectionConfig,
+    RDSIAMTokenManager,
+)
 from datahub.ingestion.source.sql.sql_common import (
     make_sqlalchemy_type,
     register_custom_type,
 )
 from datahub.ingestion.source.sql.sql_config import SQLAlchemyConnectionConfig
+from datahub.ingestion.source.sql.sqlalchemy_uri import parse_host_port
+from datahub.ingestion.source.sql.stored_procedures.base import (
+    BaseProcedure,
+)
 from datahub.ingestion.source.sql.two_tier_sql_source import (
     TwoTierSQLAlchemyConfig,
     TwoTierSQLAlchemySource,
 )
 from datahub.metadata.schema_classes import BytesTypeClass
+from datahub.utilities.str_enum import StrEnum
+
+logger = logging.getLogger(__name__)
 
 SET.__repr__ = util.generic_repr  # type:ignore
 
@@ -48,16 +65,49 @@ base.ischema_names["polygon"] = POLYGON
 base.ischema_names["decimal128"] = DECIMAL128
 
 
+class MySQLAuthMode(StrEnum):
+    """Authentication mode for MySQL connection."""
+
+    PASSWORD = "PASSWORD"
+    AWS_IAM = "AWS_IAM"
+
+
 class MySQLConnectionConfig(SQLAlchemyConnectionConfig):
     # defaults
     host_port: str = Field(default="localhost:3306", description="MySQL host URL.")
-    scheme: str = "mysql+pymysql"
+    scheme: HiddenFromDocs[str] = "mysql+pymysql"
+
+    # Authentication configuration
+    auth_mode: MySQLAuthMode = Field(
+        default=MySQLAuthMode.PASSWORD,
+        description="Authentication mode to use for the MySQL connection. "
+        "Options are 'PASSWORD' (default) for standard username/password authentication, "
+        "or 'AWS_IAM' for AWS RDS IAM authentication.",
+    )
+    aws_config: AwsConnectionConfig = Field(
+        default_factory=AwsConnectionConfig,
+        description="AWS configuration for RDS IAM authentication (only used when auth_mode is AWS_IAM). "
+        "Provides full control over AWS credentials, region, profiles, role assumption, retry logic, and proxy settings. "
+        "If not explicitly configured, boto3 will automatically use the default credential chain and region from "
+        "environment variables (AWS_DEFAULT_REGION, AWS_REGION), AWS config files (~/.aws/config), or IAM role metadata.",
+    )
 
 
 class MySQLConfig(MySQLConnectionConfig, TwoTierSQLAlchemyConfig):
     def get_identifier(self, *, schema: str, table: str) -> str:
         return f"{schema}.{table}"
 
+    include_stored_procedures: bool = Field(
+        default=True,
+        description="Include ingest of stored procedures.",
+    )
+
+    procedure_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex patterns for stored procedures to filter in ingestion."
+        "Specify regex to match the entire procedure name in database.schema.procedure_name format. e.g. to match all procedures starting with customer in Customer database and public schema, use the regex 'Customer.public.customer.*'",
+    )
+
 
 @platform_name("MySQL")
 @config_class(MySQLConfig)
@@ -65,7 +115,6 @@ class MySQLConfig(MySQLConnectionConfig, TwoTierSQLAlchemyConfig):
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
 class MySQLSource(TwoTierSQLAlchemySource):
     """
     This plugin extracts the following:
@@ -75,9 +124,27 @@ class MySQLSource(TwoTierSQLAlchemySource):
     Table, row, and column statistics via optional SQL profiling
     """
 
-    def __init__(self, config, ctx):
+    config: MySQLConfig
+
+    def __init__(self, config: MySQLConfig, ctx: Any):
         super().__init__(config, ctx, self.get_platform())
 
+        self._rds_iam_token_manager: Optional[RDSIAMTokenManager] = None
+        if config.auth_mode == MySQLAuthMode.AWS_IAM:
+            hostname, port = parse_host_port(config.host_port, default_port=3306)
+            if port is None:
+                raise ValueError("Port must be specified for RDS IAM authentication")
+
+            if not config.username:
+                raise ValueError("username is required for RDS IAM authentication")
+
+            self._rds_iam_token_manager = RDSIAMTokenManager(
+                endpoint=hostname,
+                username=config.username,
+                port=port,
+                aws_config=config.aws_config,
+            )
+
     def get_platform(self):
         return "mysql"
 
@@ -86,6 +153,52 @@ class MySQLSource(TwoTierSQLAlchemySource):
         config = MySQLConfig.parse_obj(config_dict)
         return cls(config, ctx)
 
+    def _setup_rds_iam_event_listener(
+        self, engine: "Engine", database_name: Optional[str] = None
+    ) -> None:
+        """Setup SQLAlchemy event listener to inject RDS IAM tokens."""
+        if not (
+            self.config.auth_mode == MySQLAuthMode.AWS_IAM
+            and self._rds_iam_token_manager
+        ):
+            return
+
+        def do_connect_listener(_dialect, _conn_rec, _cargs, cparams):
+            if not self._rds_iam_token_manager:
+                raise RuntimeError("RDS IAM Token Manager is not initialized")
+            cparams["password"] = self._rds_iam_token_manager.get_token()
+            # PyMySQL requires SSL to be enabled for RDS IAM authentication.
+            # Preserve any existing SSL configuration, otherwise enable with default settings.
+            # The {"ssl": True} dict is a workaround to make PyMySQL recognize that SSL
+            # should be enabled, since the library requires a truthy value in the ssl parameter.
+            # See https://pymysql.readthedocs.io/en/latest/modules/connections.html#pymysql.connections.Connection
+            cparams["ssl"] = cparams.get("ssl") or {"ssl": True}
+
+        event.listen(engine, "do_connect", do_connect_listener)  # type: ignore[misc]
+
+    def get_inspectors(self):
+        url = self.config.get_sql_alchemy_url()
+        logger.debug(f"sql_alchemy_url={url}")
+
+        engine = create_engine(url, **self.config.options)
+        self._setup_rds_iam_event_listener(engine)
+
+        with engine.connect() as conn:
+            inspector = inspect(conn)
+            if self.config.database and self.config.database != "":
+                databases = [self.config.database]
+            else:
+                databases = inspector.get_schema_names()
+            for db in databases:
+                if self.config.database_pattern.allowed(db):
+                    url = self.config.get_sql_alchemy_url(current_db=db)
+                    db_engine = create_engine(url, **self.config.options)
+                    self._setup_rds_iam_event_listener(db_engine, database_name=db)
+
+                    with db_engine.connect() as conn:
+                        inspector = inspect(conn)
+                        yield inspector
+
     def add_profile_metadata(self, inspector: Inspector) -> None:
         if not self.config.is_profiling_enabled():
             return
@@ -96,3 +209,40 @@ class MySQLSource(TwoTierSQLAlchemySource):
                 self.profile_metadata_info.dataset_name_to_storage_bytes[
                     f"{row.TABLE_SCHEMA}.{row.TABLE_NAME}"
                 ] = row.DATA_LENGTH
+
+    def get_procedures_for_schema(
+        self, inspector: Inspector, schema: str, db_name: str
+    ) -> List[BaseProcedure]:
+        """
+        Get stored procedures for a specific schema.
+        """
+        base_procedures = []
+        with inspector.engine.connect() as conn:
+            procedures = conn.execute(
+                """
+                SELECT ROUTINE_NAME AS name, 
+                    ROUTINE_DEFINITION AS definition, 
+                    EXTERNAL_LANGUAGE AS language
+                FROM information_schema.ROUTINES
+                WHERE ROUTINE_TYPE = 'PROCEDURE'
+                AND ROUTINE_SCHEMA = %s
+                """,
+                (schema,),
+            )
+
+            procedure_rows = list(procedures)
+            for row in procedure_rows:
+                base_procedures.append(
+                    BaseProcedure(
+                        name=row.name,
+                        language=row.language,
+                        argument_signature=None,
+                        return_type=None,
+                        procedure_definition=row.definition,
+                        created=None,
+                        last_altered=None,
+                        extra_properties=None,
+                        comment=None,
+                    )
+                )
+            return base_procedures

datahub/ingestion/source/sql/oracle.py

@@ -37,7 +37,7 @@ from datahub.ingestion.source.sql.sql_config import (
 
 logger = logging.getLogger(__name__)
 
-oracledb.version = "8.3.0"
+oracledb.version = "8.3.0"  # type: ignore[assignment]
 sys.modules["cx_Oracle"] = oracledb
 
 extra_oracle_types = {
@@ -110,10 +110,10 @@ class OracleConfig(BasicSQLAlchemyConfig):
         return v
 
     @pydantic.validator("data_dictionary_mode")
-    def check_data_dictionary_mode(cls, values):
-        if values not in ("ALL", "DBA"):
+    def check_data_dictionary_mode(cls, value):
+        if value not in ("ALL", "DBA"):
             raise ValueError("Specify one of data dictionary views mode: 'ALL', 'DBA'.")
-        return values
+        return value
 
     @pydantic.validator("thick_mode_lib_dir", always=True)
     def check_thick_mode_lib_dir(cls, v, values):
@@ -441,7 +441,7 @@ class OracleInspectorObjectWrapper:
             "\nac.constraint_name,"
             "\nac.constraint_type,"
             "\nacc.column_name AS local_column,"
-            "\nac.r_table_name AS remote_table,"
+            "\nac.table_name AS remote_table,"
             "\nrcc.column_name AS remote_column,"
             "\nac.r_owner AS remote_owner,"
             "\nacc.position AS loc_pos,"

datahub/ingestion/source/sql/postgres.py

@@ -1,6 +1,6 @@
 import logging
 from collections import defaultdict
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union
 
 # This import verifies that the dependencies are available.
 import psycopg2  # noqa: F401
@@ -14,9 +14,12 @@ import sqlalchemy.dialects.postgresql as custom_types
 from geoalchemy2 import Geometry  # noqa: F401
 from pydantic import BaseModel
 from pydantic.fields import Field
-from sqlalchemy import create_engine, inspect
+from sqlalchemy import create_engine, event, inspect
 from sqlalchemy.engine.reflection import Inspector
 
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Engine
+
 from datahub.configuration.common import AllowDenyPattern
 from datahub.emitter import mce_builder
 from datahub.emitter.mcp_builder import mcps_from_mce
@@ -30,17 +33,26 @@ from datahub.ingestion.api.decorators import (
     support_status,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.aws.aws_common import (
+    AwsConnectionConfig,
+    RDSIAMTokenManager,
+)
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     SqlWorkUnit,
     register_custom_type,
 )
 from datahub.ingestion.source.sql.sql_config import BasicSQLAlchemyConfig
+from datahub.ingestion.source.sql.sqlalchemy_uri import parse_host_port
+from datahub.ingestion.source.sql.stored_procedures.base import (
+    BaseProcedure,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     ArrayTypeClass,
     BytesTypeClass,
     MapTypeClass,
 )
+from datahub.utilities.str_enum import StrEnum
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -97,12 +109,34 @@ class ViewLineageEntry(BaseModel):
     dependent_schema: str
 
 
+class PostgresAuthMode(StrEnum):
+    """Authentication mode for PostgreSQL connection."""
+
+    PASSWORD = "PASSWORD"
+    AWS_IAM = "AWS_IAM"
+
+
 class BasePostgresConfig(BasicSQLAlchemyConfig):
     scheme: str = Field(default="postgresql+psycopg2", description="database scheme")
     schema_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern(deny=["information_schema"])
     )
 
+    # Authentication configuration
+    auth_mode: PostgresAuthMode = Field(
+        default=PostgresAuthMode.PASSWORD,
+        description="Authentication mode to use for the PostgreSQL connection. "
+        "Options are 'PASSWORD' (default) for standard username/password authentication, "
+        "or 'AWS_IAM' for AWS RDS IAM authentication.",
+    )
+    aws_config: AwsConnectionConfig = Field(
+        default_factory=AwsConnectionConfig,
+        description="AWS configuration for RDS IAM authentication (only used when auth_mode is AWS_IAM). "
+        "Provides full control over AWS credentials, region, profiles, role assumption, retry logic, and proxy settings. "
+        "If not explicitly configured, boto3 will automatically use the default credential chain and region from "
+        "environment variables (AWS_DEFAULT_REGION, AWS_REGION), AWS config files (~/.aws/config), or IAM role metadata.",
+    )
+
 
 class PostgresConfig(BasePostgresConfig):
     database_pattern: AllowDenyPattern = Field(
@@ -124,6 +158,17 @@ class PostgresConfig(BasePostgresConfig):
         ),
     )
 
+    include_stored_procedures: bool = Field(
+        default=True,
+        description="Include ingest of stored procedures.",
+    )
+
+    procedure_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex patterns for stored procedures to filter in ingestion."
+        "Specify regex to match the entire procedure name in database.schema.procedure_name format. e.g. to match all procedures starting with customer in Customer database and public schema, use the regex 'Customer.public.customer.*'",
+    )
+
 
 @platform_name("Postgres")
 @config_class(PostgresConfig)
@@ -131,12 +176,11 @@ class PostgresConfig(BasePostgresConfig):
 @capability(SourceCapability.DOMAINS, "Enabled by default")
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
-@capability(SourceCapability.LINEAGE_COARSE, "Optionally enabled via configuration")
 class PostgresSource(SQLAlchemySource):
     """
     This plugin extracts the following:
 
-    - Metadata for databases, schemas, views, and tables
+    - Metadata for databases, schemas, views, tables, and stored procedures
     - Column types associated with each table
     - Also supports PostGIS extensions
     - Table, row, and column statistics via optional SQL profiling
@@ -147,6 +191,22 @@ class PostgresSource(SQLAlchemySource):
     def __init__(self, config: PostgresConfig, ctx: PipelineContext):
         super().__init__(config, ctx, self.get_platform())
 
+        self._rds_iam_token_manager: Optional[RDSIAMTokenManager] = None
+        if config.auth_mode == PostgresAuthMode.AWS_IAM:
+            hostname, port = parse_host_port(config.host_port, default_port=5432)
+            if port is None:
+                raise ValueError("Port must be specified for RDS IAM authentication")
+
+            if not config.username:
+                raise ValueError("username is required for RDS IAM authentication")
+
+            self._rds_iam_token_manager = RDSIAMTokenManager(
+                endpoint=hostname,
+                username=config.username,
+                port=port,
+                aws_config=config.aws_config,
+            )
+
     def get_platform(self):
         return "postgres"
 
@@ -155,13 +215,36 @@ class PostgresSource(SQLAlchemySource):
         config = PostgresConfig.parse_obj(config_dict)
         return cls(config, ctx)
 
+    def _setup_rds_iam_event_listener(
+        self, engine: "Engine", database_name: Optional[str] = None
+    ) -> None:
+        """Setup SQLAlchemy event listener to inject RDS IAM tokens."""
+        if not (
+            self.config.auth_mode == PostgresAuthMode.AWS_IAM
+            and self._rds_iam_token_manager
+        ):
+            return
+
+        def do_connect_listener(_dialect, _conn_rec, _cargs, cparams):
+            if not self._rds_iam_token_manager:
+                raise RuntimeError("RDS IAM Token Manager is not initialized")
+            cparams["password"] = self._rds_iam_token_manager.get_token()
+            if cparams.get("sslmode") not in ("require", "verify-ca", "verify-full"):
+                cparams["sslmode"] = "require"
+
+        event.listen(engine, "do_connect", do_connect_listener)  # type: ignore[misc]
+
     def get_inspectors(self) -> Iterable[Inspector]:
         # Note: get_sql_alchemy_url will choose `sqlalchemy_uri` over the passed in database
         url = self.config.get_sql_alchemy_url(
             database=self.config.database or self.config.initial_database
         )
+
         logger.debug(f"sql_alchemy_url={url}")
+
         engine = create_engine(url, **self.config.options)
+        self._setup_rds_iam_event_listener(engine)
+
         with engine.connect() as conn:
             if self.config.database or self.config.sqlalchemy_uri:
                 inspector = inspect(conn)
@@ -169,14 +252,21 @@ class PostgresSource(SQLAlchemySource):
             else:
                 # pg_database catalog -  https://www.postgresql.org/docs/current/catalog-pg-database.html
                 # exclude template databases - https://www.postgresql.org/docs/current/manage-ag-templatedbs.html
+                # exclude rdsadmin - AWS RDS administrative database
                 databases = conn.execute(
-                    "SELECT datname from pg_database where datname not in ('template0', 'template1')"
+                    "SELECT datname from pg_database where datname not in ('template0', 'template1', 'rdsadmin')"
                 )
                 for db in databases:
                     if not self.config.database_pattern.allowed(db["datname"]):
                         continue
+
                     url = self.config.get_sql_alchemy_url(database=db["datname"])
-                    with create_engine(url, **self.config.options).connect() as conn:
+                    db_engine = create_engine(url, **self.config.options)
+                    self._setup_rds_iam_event_listener(
+                        db_engine, database_name=db["datname"]
+                    )
+
+                    with db_engine.connect() as conn:
                         inspector = inspect(conn)
                         yield inspector
 
@@ -292,3 +382,49 @@ class PostgresSource(SQLAlchemySource):
                     ] = row.table_size
         except Exception as e:
             logger.error(f"failed to fetch profile metadata: {e}")
+
+    def get_procedures_for_schema(
+        self, inspector: Inspector, schema: str, db_name: str
+    ) -> List[BaseProcedure]:
+        """
+        Get stored procedures for a specific schema.
+        """
+        base_procedures = []
+        with inspector.engine.connect() as conn:
+            procedures = conn.execute(
+                """
+                    SELECT
+                        p.proname AS name,
+                        l.lanname AS language,
+                        pg_get_function_arguments(p.oid) AS arguments,
+                        pg_get_functiondef(p.oid) AS definition,
+                        obj_description(p.oid, 'pg_proc') AS comment
+                    FROM
+                        pg_proc p
+                    JOIN
+                        pg_namespace n ON n.oid = p.pronamespace
+                    JOIN
+                        pg_language l ON l.oid = p.prolang
+                    WHERE
+                        p.prokind = 'p'
+                        AND n.nspname = %s;
+                """,
+                (schema,),
+            )
+
+            procedure_rows = list(procedures)
+            for row in procedure_rows:
+                base_procedures.append(
+                    BaseProcedure(
+                        name=row.name,
+                        language=row.language,
+                        argument_signature=row.arguments,
+                        return_type=None,
+                        procedure_definition=row.definition,
+                        created=None,
+                        last_altered=None,
+                        comment=row.comment,
+                        extra_properties=None,
+                    )
+                )
+            return base_procedures

datahub/ingestion/source/sql/presto.py

@@ -8,6 +8,7 @@ from sqlalchemy import exc, sql
 from sqlalchemy.engine import reflection
 from sqlalchemy.engine.base import Engine
 
+from datahub.configuration.common import HiddenFromDocs
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -87,7 +88,7 @@ PrestoDialect._get_full_table = _get_full_table
 
 class PrestoConfig(TrinoConfig):
     # defaults
-    scheme: str = Field(default="presto", description="", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="presto")
 
 
 @platform_name("Presto", doc_order=1)