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

datahub/ingestion/source/bigquery_v2/bigquery_config.py

@@ -1,12 +1,13 @@
 import logging
-import os
 import re
+from copy import deepcopy
 from datetime import timedelta
 from typing import Dict, List, Optional, Union
 
 from pydantic import Field, PositiveInt, PrivateAttr, root_validator, validator
 
-from datahub.configuration.common import AllowDenyPattern, ConfigModel
+from datahub.configuration.common import AllowDenyPattern, ConfigModel, HiddenFromDocs
+from datahub.configuration.env_vars import get_bigquery_schema_parallelism
 from datahub.configuration.source_common import (
     EnvConfigMixin,
     LowerCaseDatasetUrnConfigMixin,
@@ -24,15 +25,14 @@ from datahub.ingestion.source.sql.sql_config import SQLCommonConfig, SQLFilterCo
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulLineageConfigMixin,
     StatefulProfilingConfigMixin,
+    StatefulTimeWindowConfigMixin,
     StatefulUsageConfigMixin,
 )
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_BQ_SCHEMA_PARALLELISM = int(
-    os.getenv("DATAHUB_BIGQUERY_SCHEMA_PARALLELISM", 20)
-)
+DEFAULT_BQ_SCHEMA_PARALLELISM = get_bigquery_schema_parallelism()
 
 # Regexp for sharded tables.
 # A sharded table is a table that has a suffix of the form _yyyymmdd or yyyymmdd, where yyyymmdd is a date.
@@ -73,8 +73,10 @@ class BigQueryBaseConfig(ConfigModel):
             ) from e
         return v
 
-    @root_validator(pre=True, skip_on_failure=True)
+    @root_validator(pre=True)
     def project_id_backward_compatibility_configs_set(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
         project_id = values.pop("project_id", None)
         project_ids = values.get("project_ids")
 
@@ -182,13 +184,14 @@ class BigQueryFilterConfig(SQLFilterConfig):
     )
 
     # NOTE: `schema_pattern` is added here only to hide it from docs.
-    schema_pattern: AllowDenyPattern = Field(
+    schema_pattern: HiddenFromDocs[AllowDenyPattern] = Field(
         default=AllowDenyPattern.allow_all(),
-        hidden_from_docs=True,
     )
 
     @root_validator(pre=False, skip_on_failure=True)
     def backward_compatibility_configs_set(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
         dataset_pattern: Optional[AllowDenyPattern] = values.get("dataset_pattern")
         schema_pattern = values.get("schema_pattern")
         if (
@@ -269,6 +272,7 @@ class BigQueryV2Config(
     SQLCommonConfig,
     StatefulUsageConfigMixin,
     StatefulLineageConfigMixin,
+    StatefulTimeWindowConfigMixin,
     StatefulProfilingConfigMixin,
     ClassificationSourceConfigMixin,
 ):
@@ -320,8 +324,7 @@ class BigQueryV2Config(
         description="Include full payload into events. It is only for debugging and internal use.",
     )
 
-    number_of_datasets_process_in_batch: int = Field(
-        hidden_from_docs=True,
+    number_of_datasets_process_in_batch: HiddenFromDocs[int] = Field(
         default=10000,
         description="Number of table queried in batch when getting metadata. This is a low level config property "
         "which should be touched with care.",
@@ -342,7 +345,7 @@ class BigQueryV2Config(
     )
 
     use_queries_v2: bool = Field(
-        default=False,
+        default=True,
         description="If enabled, uses the new queries extractor to extract queries from bigquery.",
     )
     include_queries: bool = Field(
@@ -436,17 +439,15 @@ class BigQueryV2Config(
 
     upstream_lineage_in_report: bool = Field(
         default=False,
-        description="Useful for debugging lineage information. Set to True to see the raw lineage created internally.",
+        description="Useful for debugging lineage information. Set to True to see the raw lineage created internally. Only works with legacy approach (`use_queries_v2: False`).",
     )
 
-    run_optimized_column_query: bool = Field(
-        hidden_from_docs=True,
+    run_optimized_column_query: HiddenFromDocs[bool] = Field(
         default=False,
         description="Run optimized column query to get column information. This is an experimental feature and may not work for all cases.",
     )
 
-    file_backed_cache_size: int = Field(
-        hidden_from_docs=True,
+    file_backed_cache_size: HiddenFromDocs[int] = Field(
         default=2000,
         description="Maximum number of entries for the in-memory caches of FileBacked data structures.",
     )
@@ -456,10 +457,9 @@ class BigQueryV2Config(
         description="Option to exclude empty projects from being ingested.",
     )
 
-    schema_resolution_batch_size: int = Field(
+    schema_resolution_batch_size: HiddenFromDocs[int] = Field(
         default=100,
         description="The number of tables to process in a batch when resolving schema from DataHub.",
-        hidden_from_schema=True,
     )
 
     max_threads_dataset_parallelism: int = Field(
@@ -480,6 +480,8 @@ class BigQueryV2Config(
 
     @root_validator(pre=True)
     def set_include_schema_metadata(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
         # Historically this is used to disable schema ingestion
         if (
             "include_tables" in values
@@ -498,6 +500,8 @@ class BigQueryV2Config(
 
     @root_validator(skip_on_failure=True)
     def profile_default_settings(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
         # Extra default SQLAlchemy option for better connection pooling and threading.
         # https://docs.sqlalchemy.org/en/14/core/pooling.html#sqlalchemy.pool.QueuePool.params.max_overflow
         values["options"].setdefault("max_overflow", -1)
@@ -515,9 +519,33 @@ class BigQueryV2Config(
 
         return v
 
+    @validator("upstream_lineage_in_report")
+    def validate_upstream_lineage_in_report(cls, v: bool, values: Dict) -> bool:
+        if v and values.get("use_queries_v2", True):
+            logging.warning(
+                "`upstream_lineage_in_report` is enabled but will be ignored because `use_queries_v2` is enabled."
+                "This debugging feature only works with the legacy lineage approach (`use_queries_v2: false`)."
+            )
+
+        return v
+
+    @root_validator(pre=False, skip_on_failure=True)
+    def validate_queries_v2_stateful_ingestion(cls, values: Dict) -> Dict:
+        if values.get("use_queries_v2"):
+            if values.get("enable_stateful_lineage_ingestion") or values.get(
+                "enable_stateful_usage_ingestion"
+            ):
+                logger.warning(
+                    "enable_stateful_lineage_ingestion and enable_stateful_usage_ingestion are deprecated "
+                    "when using use_queries_v2=True. These configs only work with the legacy (non-queries v2) extraction path. "
+                    "For queries v2, use enable_stateful_time_window instead to enable stateful ingestion "
+                    "for the unified time window extraction (lineage + usage + operations + queries)."
+                )
+        return values
+
     def get_table_pattern(self, pattern: List[str]) -> str:
         return "|".join(pattern) if pattern else ""
 
-    platform_instance_not_supported_for_bigquery = pydantic_removed_field(
+    _platform_instance_not_supported_for_bigquery = pydantic_removed_field(
         "platform_instance"
     )

datahub/ingestion/source/bigquery_v2/bigquery_connection.py

@@ -2,16 +2,23 @@ import logging
 import os
 from typing import Any, Dict, Optional
 
+from google.api_core.client_info import ClientInfo
 from google.cloud import bigquery, datacatalog_v1, resourcemanager_v3
 from google.cloud.logging_v2.client import Client as GCPLoggingClient
 from pydantic import Field, PrivateAttr
 
+from datahub._version import __version__
 from datahub.configuration.common import ConfigModel
 from datahub.ingestion.source.common.gcp_credentials_config import GCPCredential
 
 logger = logging.getLogger(__name__)
 
 
+def _get_bigquery_client_info() -> ClientInfo:
+    """Get ClientInfo with DataHub user-agent for BigQuery client identification"""
+    return ClientInfo(user_agent=f"datahub/{__version__}")
+
+
 class BigQueryConnectionConfig(ConfigModel):
     credential: Optional[GCPCredential] = Field(
         default=None, description="BigQuery credential informations"
@@ -41,7 +48,11 @@ class BigQueryConnectionConfig(ConfigModel):
 
     def get_bigquery_client(self) -> bigquery.Client:
         client_options = self.extra_client_options
-        return bigquery.Client(self.project_on_behalf, **client_options)
+        return bigquery.Client(
+            self.project_on_behalf,
+            client_info=_get_bigquery_client_info(),
+            **client_options,
+        )
 
     def get_projects_client(self) -> resourcemanager_v3.ProjectsClient:
         return resourcemanager_v3.ProjectsClient()

datahub/ingestion/source/bigquery_v2/bigquery_queries.py

@@ -7,6 +7,7 @@ from typing_extensions import Self
 
 from datahub.configuration.time_window_config import BaseTimeWindowConfig
 from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.decorators import SupportStatus, support_status
 from datahub.ingestion.api.source import Source, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.bigquery_v2.bigquery_config import (
@@ -50,6 +51,7 @@ class BigQueryQueriesSourceConfig(
     )
 
 
+@support_status(SupportStatus.CERTIFIED)
 class BigQueryQueriesSource(Source):
     def __init__(self, ctx: PipelineContext, config: BigQueryQueriesSourceConfig):
         self.ctx = ctx
@@ -94,3 +96,4 @@ class BigQueryQueriesSource(Source):
     def close(self) -> None:
         self.queries_extractor.close()
         self.connection.close()
+        super().close()

datahub/ingestion/source/bigquery_v2/bigquery_report.py

@@ -9,7 +9,6 @@ import pydantic
 from datahub.ingestion.api.report import Report
 from datahub.ingestion.glossary.classification_mixin import ClassificationReportMixin
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
-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, LossyList, LossySet
@@ -78,7 +77,6 @@ class BigQueryQueriesExtractorReport(Report):
 @dataclass
 class BigQueryV2Report(
     SQLSourceReport,
-    IngestionStageReport,
     BaseTimeWindowReport,
     ClassificationReportMixin,
 ):

datahub/ingestion/source/bigquery_v2/bigquery_schema.py

@@ -283,23 +283,30 @@ class BigQuerySchemaApi:
         with self.report.list_datasets_timer:
             self.report.num_list_datasets_api_requests += 1
             datasets = self.bq_client.list_datasets(project_id, max_results=maxResults)
-            return [
-                BigqueryDataset(
-                    name=d.dataset_id,
-                    labels=d.labels,
-                    location=(
-                        d._properties.get("location")
-                        if hasattr(d, "_properties") and isinstance(d._properties, dict)
-                        else None
-                    ),
-                    # TODO: Fetch dataset description individually impacts overall performance if the number of datasets is high (hundreds); instead we should fetch in batch for all datasets.
-                    # TODO: Given we are calling get_dataset for each dataset, we may consume and publish other fields too, such as created, modified, etc...
-                    # https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_get_dataset
-                    # https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.dataset.Dataset
-                    comment=self.bq_client.get_dataset(d.reference).description,
+            result = []
+            for d in datasets:
+                # TODO: Fetch dataset description individually impacts overall performance if the number of datasets is high (hundreds); instead we should fetch in batch for all datasets.
+                # https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_get_dataset
+                # https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.dataset.Dataset
+                dataset = self.bq_client.get_dataset(d.reference)
+
+                location = (
+                    d._properties.get("location")
+                    if hasattr(d, "_properties") and isinstance(d._properties, dict)
+                    else None
+                )
+
+                result.append(
+                    BigqueryDataset(
+                        name=d.dataset_id,
+                        labels=d.labels,
+                        location=location,
+                        comment=dataset.description,
+                        created=dataset.created,
+                        last_altered=dataset.modified,
+                    )
                 )
-                for d in datasets
-            ]
+            return result
 
     # This is not used anywhere
     def get_datasets_for_project_id_with_information_schema(

datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py

@@ -12,6 +12,7 @@ from datahub.emitter.mce_builder import (
     make_dataset_urn_with_platform_instance,
     make_schema_field_urn,
     make_tag_urn,
+    make_ts_millis,
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import BigQueryDatasetKey, ContainerKey, ProjectIdKey
@@ -65,7 +66,7 @@ from datahub.ingestion.source.sql.sql_utils import (
 )
 from datahub.ingestion.source_report.ingestion_stage import (
     METADATA_EXTRACTION,
-    PROFILING,
+    IngestionHighStage,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import (
     Status,
@@ -286,6 +287,7 @@ class BigQuerySchemaGenerator:
         yield from gen_database_container(
             database=database,
             name=database,
+            qualified_name=database,
             sub_types=[DatasetContainerSubTypes.BIGQUERY_PROJECT],
             domain_registry=self.domain_registry,
             domain_config=self.config.domain,
@@ -299,6 +301,8 @@ class BigQuerySchemaGenerator:
         description: Optional[str] = None,
         tags: Optional[Dict[str, str]] = None,
         extra_properties: Optional[Dict[str, str]] = None,
+        created: Optional[int] = None,
+        last_modified: Optional[int] = None,
     ) -> Iterable[MetadataWorkUnit]:
         schema_container_key = self.gen_dataset_key(project_id, dataset)
 
@@ -332,6 +336,7 @@ class BigQuerySchemaGenerator:
         yield from gen_schema_container(
             database=project_id,
             schema=dataset,
+            qualified_name=f"{project_id}.{dataset}",
             sub_types=[DatasetContainerSubTypes.BIGQUERY_DATASET],
             domain_registry=self.domain_registry,
             domain_config=self.config.domain,
@@ -347,6 +352,8 @@ class BigQuerySchemaGenerator:
             ),
             tags=tags_joined,
             extra_properties=extra_properties,
+            created=created,
+            last_modified=last_modified,
         )
 
     def _process_project(
@@ -409,7 +416,7 @@ class BigQuerySchemaGenerator:
 
         if self.config.is_profiling_enabled():
             logger.info(f"Starting profiling project {project_id}")
-            with self.report.new_stage(f"{project_id}: {PROFILING}"):
+            with self.report.new_high_stage(IngestionHighStage.PROFILING):
                 yield from self.profiler.get_workunits(
                     project_id=project_id,
                     tables=db_tables,
@@ -442,10 +449,12 @@ class BigQuerySchemaGenerator:
                 ):
                     yield wu
             except Exception as e:
-                if self.config.is_profiling_enabled():
-                    action_mesage = "Does your service account have bigquery.tables.list, bigquery.routines.get, bigquery.routines.list permission, bigquery.tables.getData permission?"
+                # If configuration indicates we need table data access (for profiling or use_tables_list_query_v2),
+                # include bigquery.tables.getData in the error message since that's likely the missing permission
+                if self.config.have_table_data_read_permission:
+                    action_mesage = "Does your service account have bigquery.tables.list, bigquery.routines.get, bigquery.routines.list, bigquery.tables.getData permissions?"
                 else:
-                    action_mesage = "Does your service account have bigquery.tables.list, bigquery.routines.get, bigquery.routines.list permission?"
+                    action_mesage = "Does your service account have bigquery.tables.list, bigquery.routines.get, bigquery.routines.list permissions?"
 
                 self.report.failure(
                     title="Unable to get tables for dataset",
@@ -482,6 +491,12 @@ class BigQuerySchemaGenerator:
                     else None
                 ),
                 description=bigquery_dataset.comment,
+                created=make_ts_millis(bigquery_dataset.created)
+                if bigquery_dataset.created
+                else None,
+                last_modified=make_ts_millis(bigquery_dataset.last_altered)
+                if bigquery_dataset.last_altered
+                else None,
             )
 
         columns = None

datahub/ingestion/source/bigquery_v2/common.py

@@ -63,7 +63,7 @@ class BigQueryIdentifierBuilder:
         )
 
     def gen_user_urn(self, user_email: str) -> str:
-        return make_user_urn(user_email.split("@")[0])
+        return make_user_urn(user_email)
 
     def make_data_platform_urn(self) -> str:
         return make_data_platform_urn(self.platform)

datahub/ingestion/source/bigquery_v2/profiler.py

@@ -189,6 +189,7 @@ WHERE
 
         if len(profile_requests) == 0:
             return
+
         yield from self.generate_profile_workunits(
             profile_requests,
             max_workers=self.config.profiling.max_workers,
@@ -226,10 +227,11 @@ WHERE
             db_name, schema_name, bq_table, self.config.profiling.partition_datetime
         )
 
-        if partition is None and bq_table.partition_info:
+        # For partitioned tables, if it has a row count but not a valid partition, that means something went wrong with the partition detection.
+        if partition is None and bq_table.partition_info and bq_table.rows_count:
             self.report.report_warning(
                 title="Profile skipped for partitioned table",
-                message="profile skipped as partitioned table is empty or partition id or type was invalid",
+                message="profile skipped as partition id or type was invalid",
                 context=profile_request.pretty_name,
             )
             return None

datahub/ingestion/source/bigquery_v2/queries.py

@@ -45,12 +45,12 @@ SELECT
   tos.OPTION_VALUE as comment,
   t.is_insertable_into,
   t.ddl,
-  ts.row_count,
+  ts.row_count as row_count,
   ts.size_bytes as bytes,
   p.num_partitions,
   p.max_partition_id,
-  p.active_billable_bytes,
-  p.long_term_billable_bytes,
+  p.active_billable_bytes as active_billable_bytes,
+  IFNULL(p.long_term_billable_bytes, 0) as long_term_billable_bytes,
   REGEXP_EXTRACT(t.table_name, r"(?:(?:.+\\D)[_$]?)(\\d\\d\\d\\d(?:0[1-9]|1[012])(?:0[1-9]|[12][0-9]|3[01]))$") as table_suffix,
   REGEXP_REPLACE(t.table_name, r"(?:[_$]?)(\\d\\d\\d\\d(?:0[1-9]|1[012])(?:0[1-9]|[12][0-9]|3[01]))$", "") as table_base
 

datahub/ingestion/source/bigquery_v2/queries_extractor.py

@@ -8,7 +8,7 @@ from typing import Collection, Dict, Iterable, List, Optional, TypedDict
 from google.cloud.bigquery import Client
 from pydantic import Field, PositiveInt
 
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.configuration.time_window_config import (
     BaseTimeWindowConfig,
     get_time_bucket,
@@ -36,6 +36,9 @@ from datahub.ingestion.source.bigquery_v2.common import (
     BigQueryFilter,
     BigQueryIdentifierBuilder,
 )
+from datahub.ingestion.source.state.redundant_run_skip_handler import (
+    RedundantQueriesRunSkipHandler,
+)
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
 from datahub.metadata.urns import CorpUserUrn
 from datahub.sql_parsing.schema_resolver import SchemaResolver
@@ -86,12 +89,11 @@ class BigQueryQueriesExtractorConfig(BigQueryBaseConfig):
     # TODO: Support stateful ingestion for the time windows.
     window: BaseTimeWindowConfig = BaseTimeWindowConfig()
 
-    local_temp_path: Optional[pathlib.Path] = Field(
-        default=None,
-        description="Local path to store the audit log.",
+    local_temp_path: HiddenFromDocs[Optional[pathlib.Path]] = Field(
         # TODO: For now, this is simply an advanced config to make local testing easier.
         # Eventually, we will want to store date-specific files in the directory and use it as a cache.
-        hidden_from_docs=True,
+        default=None,
+        description="Local path to store the audit log.",
     )
 
     user_email_pattern: AllowDenyPattern = Field(
@@ -136,6 +138,7 @@ class BigQueryQueriesExtractor(Closeable):
         structured_report: SourceReport,
         filters: BigQueryFilter,
         identifiers: BigQueryIdentifierBuilder,
+        redundant_run_skip_handler: Optional[RedundantQueriesRunSkipHandler] = None,
         graph: Optional[DataHubGraph] = None,
         schema_resolver: Optional[SchemaResolver] = None,
         discovered_tables: Optional[Collection[str]] = None,
@@ -159,6 +162,9 @@ class BigQueryQueriesExtractor(Closeable):
         )
 
         self.structured_report = structured_report
+        self.redundant_run_skip_handler = redundant_run_skip_handler
+
+        self.start_time, self.end_time = self._get_time_window()
 
         self.aggregator = SqlParsingAggregator(
             platform=self.identifiers.platform,
@@ -173,8 +179,8 @@ class BigQueryQueriesExtractor(Closeable):
             generate_query_usage_statistics=self.config.include_query_usage_statistics,
             usage_config=BaseUsageConfig(
                 bucket_duration=self.config.window.bucket_duration,
-                start_time=self.config.window.start_time,
-                end_time=self.config.window.end_time,
+                start_time=self.start_time,
+                end_time=self.end_time,
                 user_email_pattern=self.config.user_email_pattern,
                 top_n_queries=self.config.top_n_queries,
             ),
@@ -200,6 +206,34 @@ class BigQueryQueriesExtractor(Closeable):
         logger.info(f"Using local temp path: {path}")
         return path
 
+    def _get_time_window(self) -> tuple[datetime, datetime]:
+        if self.redundant_run_skip_handler:
+            start_time, end_time = (
+                self.redundant_run_skip_handler.suggest_run_time_window(
+                    self.config.window.start_time,
+                    self.config.window.end_time,
+                )
+            )
+        else:
+            start_time = self.config.window.start_time
+            end_time = self.config.window.end_time
+
+        # Usage statistics are aggregated per bucket (typically per day).
+        # To ensure accurate aggregated metrics, we need to align the start_time
+        # to the beginning of a bucket so that we include complete bucket periods.
+        if self.config.include_usage_statistics:
+            start_time = get_time_bucket(start_time, self.config.window.bucket_duration)
+
+        return start_time, end_time
+
+    def _update_state(self) -> None:
+        if self.redundant_run_skip_handler:
+            self.redundant_run_skip_handler.update_state(
+                self.config.window.start_time,
+                self.config.window.end_time,
+                self.config.window.bucket_duration,
+            )
+
     def is_temp_table(self, name: str) -> bool:
         try:
             table = BigqueryTableIdentifier.from_string_name(name)
@@ -300,6 +334,8 @@ class BigQueryQueriesExtractor(Closeable):
             shared_connection.close()
             audit_log_file.unlink(missing_ok=True)
 
+        self._update_state()
+
     def deduplicate_queries(
         self, queries: FileBackedList[ObservedQuery]
     ) -> FileBackedDict[Dict[int, ObservedQuery]]:
@@ -356,8 +392,8 @@ class BigQueryQueriesExtractor(Closeable):
         query_log_query = _build_enriched_query_log_query(
             project_id=project.id,
             region=region,
-            start_time=self.config.window.start_time,
-            end_time=self.config.window.end_time,
+            start_time=self.start_time,
+            end_time=self.end_time,
         )
 
         logger.info(f"Fetching query log from BQ Project {project.id} for {region}")

datahub/ingestion/source/cassandra/cassandra.py

@@ -80,7 +80,7 @@ class KeyspaceKey(ContainerKey):
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
 class CassandraSource(StatefulIngestionSourceBase):
@@ -296,13 +296,11 @@ class CassandraSource(StatefulIngestionSourceBase):
             qualified_name=dataset_name,
             description=view.comment,
             custom_properties=self._get_dataset_custom_props(view),
-            extra_aspects=[
-                ViewPropertiesClass(
-                    materialized=True,
-                    viewLogic=view.where_clause,  # Use the WHERE clause as view logic
-                    viewLanguage="CQL",  # Use "CQL" as the language
-                ),
-            ],
+            view_definition=ViewPropertiesClass(
+                materialized=True,
+                viewLogic=view.where_clause,  # Use the WHERE clause as view logic
+                viewLanguage="CQL",  # Use "CQL" as the language
+            ),
         )
 
         # Construct and emit lineage off of 'base_table_name'

datahub/ingestion/source/cassandra/cassandra_api.py

@@ -132,7 +132,23 @@ class CassandraAPI:
 
             ssl_context = None
             if self.config.ssl_ca_certs:
-                ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+                # Map SSL version string to ssl module constant
+                ssl_version_map = {
+                    "TLS_CLIENT": ssl.PROTOCOL_TLS_CLIENT,
+                    "TLSv1": ssl.PROTOCOL_TLSv1,
+                    "TLSv1_1": ssl.PROTOCOL_TLSv1_1,
+                    "TLSv1_2": ssl.PROTOCOL_TLSv1_2,
+                    "TLSv1_3": ssl.PROTOCOL_TLSv1_2,  # Python's ssl module uses TLSv1_2 for TLS 1.3
+                }
+
+                ssl_protocol = (
+                    ssl_version_map.get(
+                        self.config.ssl_version, ssl.PROTOCOL_TLS_CLIENT
+                    )
+                    if self.config.ssl_version
+                    else ssl.PROTOCOL_TLS_CLIENT
+                )
+                ssl_context = ssl.SSLContext(ssl_protocol)
                 ssl_context.load_verify_locations(self.config.ssl_ca_certs)
                 if self.config.ssl_certfile and self.config.ssl_keyfile:
                     ssl_context.load_cert_chain(

datahub/ingestion/source/cassandra/cassandra_config.py

@@ -94,6 +94,11 @@ class CassandraSourceConfig(
         description="Path to the SSL key file for SSL connections.",
     )
 
+    ssl_version: Optional[str] = Field(
+        default="TLS_CLIENT",
+        description="SSL protocol version to use for connections. Options: TLS_CLIENT, TLSv1, TLSv1_1, TLSv1_2, TLSv1_3. Defaults to TLS_CLIENT.",
+    )
+
     keyspace_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern.allow_all(),
         description="Regex patterns to filter keyspaces for ingestion.",

datahub/ingestion/source/cassandra/cassandra_profiling.py

@@ -18,7 +18,7 @@ from datahub.ingestion.source.cassandra.cassandra_api import (
 )
 from datahub.ingestion.source.cassandra.cassandra_config import CassandraSourceConfig
 from datahub.ingestion.source.cassandra.cassandra_utils import CassandraSourceReport
-from datahub.ingestion.source_report.ingestion_stage import PROFILING
+from datahub.ingestion.source_report.ingestion_stage import IngestionHighStage
 from datahub.metadata.schema_classes import (
     DatasetFieldProfileClass,
     DatasetProfileClass,
@@ -70,11 +70,12 @@ class CassandraProfiler:
     ) -> Iterable[MetadataWorkUnit]:
         for keyspace_name in cassandra_data.keyspaces:
             tables = cassandra_data.tables.get(keyspace_name, [])
-            with self.report.new_stage(
-                f"{keyspace_name}: {PROFILING}"
-            ), ThreadPoolExecutor(
-                max_workers=self.config.profiling.max_workers
-            ) as executor:
+            with (
+                self.report.new_high_stage(IngestionHighStage.PROFILING),
+                ThreadPoolExecutor(
+                    max_workers=self.config.profiling.max_workers
+                ) as executor,
+            ):
                 future_to_dataset = {
                     executor.submit(
                         self.generate_profile,

datahub/ingestion/source/cassandra/cassandra_utils.py

@@ -6,7 +6,6 @@ from datahub.ingestion.source.cassandra.cassandra_api import CassandraColumn
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalSourceReport,
 )
-from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     SchemaField,
     SchemaFieldDataType,
@@ -35,7 +34,7 @@ SYSTEM_KEYSPACE_LIST = set(
 
 
 @dataclass
-class CassandraSourceReport(StaleEntityRemovalSourceReport, IngestionStageReport):
+class CassandraSourceReport(StaleEntityRemovalSourceReport):
     num_tables_failed: int = 0
     num_views_failed: int = 0
     tables_scanned: int = 0