acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

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):
@@ -123,16 +123,7 @@ class CassandraSource(StatefulIngestionSourceBase):
             ).workunit_processor,
         ]
 
-    def get_workunits_internal(
-        self,
-    ) -> Iterable[MetadataWorkUnit]:
-        for metadata in self._get_metadata():
-            if isinstance(metadata, MetadataWorkUnit):
-                yield metadata
-            else:
-                yield from metadata.as_workunits()
-
-    def _get_metadata(self) -> Iterable[Union[MetadataWorkUnit, Entity]]:
+    def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, Entity]]:
         if not self.cassandra_api.authenticate():
             return
         keyspaces: List[CassandraKeyspace] = self.cassandra_api.get_keyspaces()
@@ -305,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

@@ -1,3 +1,4 @@
+import ssl
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional
 
@@ -128,6 +129,39 @@ class CassandraAPI:
 
                 self._cassandra_session = cluster.connect()
                 return True
+
+            ssl_context = None
+            if self.config.ssl_ca_certs:
+                # 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(
+                        certfile=self.config.ssl_certfile,
+                        keyfile=self.config.ssl_keyfile,
+                    )
+                elif self.config.ssl_certfile or self.config.ssl_keyfile:
+                    # If one is provided, the other must be too.
+                    # This is a simplification; real-world scenarios might allow one without the other depending on setup.
+                    raise ValueError(
+                        "Both ssl_certfile and ssl_keyfile must be provided if one is specified."
+                    )
+
             if self.config.username and self.config.password:
                 auth_provider = PlainTextAuthProvider(
                     username=self.config.username, password=self.config.password
@@ -136,12 +170,14 @@ class CassandraAPI:
                     [self.config.contact_point],
                     port=self.config.port,
                     auth_provider=auth_provider,
+                    ssl_context=ssl_context,
                     load_balancing_policy=None,
                 )
             else:
                 cluster = Cluster(
                     [self.config.contact_point],
                     port=self.config.port,
+                    ssl_context=ssl_context,
                     load_balancing_policy=None,
                 )
 

datahub/ingestion/source/cassandra/cassandra_config.py

@@ -79,6 +79,26 @@ class CassandraSourceConfig(
         description="Configuration for cloud-based Cassandra, such as DataStax Astra DB.",
     )
 
+    ssl_ca_certs: Optional[str] = Field(
+        default=None,
+        description="Path to the CA certificate file for SSL connections.",
+    )
+
+    ssl_certfile: Optional[str] = Field(
+        default=None,
+        description="Path to the SSL certificate file for SSL connections.",
+    )
+
+    ssl_keyfile: Optional[str] = Field(
+        default=None,
+        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,30 +70,32 @@ 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}"):
-                with ThreadPoolExecutor(
+            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,
-                            keyspace_name,
-                            table_name,
-                            cassandra_data.columns.get(table_name, []),
-                        ): table_name
-                        for table_name in tables
-                    }
-                    for future in as_completed(future_to_dataset):
-                        table_name = future_to_dataset[future]
-                        try:
-                            yield from future.result()
-                        except Exception as exc:
-                            self.report.profiling_skipped_other[table_name] += 1
-                            self.report.failure(
-                                message="Failed to profile for table",
-                                context=f"{keyspace_name}.{table_name}",
-                                exc=exc,
-                            )
+                ) as executor,
+            ):
+                future_to_dataset = {
+                    executor.submit(
+                        self.generate_profile,
+                        keyspace_name,
+                        table_name,
+                        cassandra_data.columns.get(table_name, []),
+                    ): table_name
+                    for table_name in tables
+                }
+                for future in as_completed(future_to_dataset):
+                    table_name = future_to_dataset[future]
+                    try:
+                        yield from future.result()
+                    except Exception as exc:
+                        self.report.profiling_skipped_other[table_name] += 1
+                        self.report.failure(
+                            message="Failed to profile for table",
+                            context=f"{keyspace_name}.{table_name}",
+                            exc=exc,
+                        )
 
     def generate_profile(
         self,

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

datahub/ingestion/source/common/data_platforms.py

@@ -0,0 +1,23 @@
+# This is a pretty limited list, and is not really complete yet. Right now it's only used to allow
+# automatic platform mapping when generating lineage and we have a manual override, so
+# it being incomplete is ok. This should not be used for urn validation.
+KNOWN_VALID_PLATFORM_NAMES = [
+    "bigquery",
+    "cassandra",
+    "databricks",
+    "delta-lake",
+    "dbt",
+    "feast",
+    "file",
+    "gcs",
+    "hdfs",
+    "hive",
+    "mssql",
+    "mysql",
+    "oracle",
+    "postgres",
+    "redshift",
+    "s3",
+    "sagemaker",
+    "snowflake",
+]

datahub/ingestion/source/common/gcp_credentials_config.py

@@ -9,7 +9,9 @@ from datahub.configuration.validate_multiline_string import pydantic_multiline_s
 
 
 class GCPCredential(ConfigModel):
-    project_id: Optional[str] = Field(description="Project id to set the credentials")
+    project_id: Optional[str] = Field(
+        None, description="Project id to set the credentials"
+    )
     private_key_id: str = Field(description="Private key id")
     private_key: str = Field(
         description="Private key in a form of '-----BEGIN PRIVATE KEY-----\\nprivate-key\\n-----END PRIVATE KEY-----\\n'"
@@ -51,3 +53,9 @@ class GCPCredential(ConfigModel):
             cred_json = json.dumps(configs, indent=4, separators=(",", ": "))
             fp.write(cred_json.encode())
             return fp.name
+
+    def to_dict(self, project_id: Optional[str] = None) -> Dict[str, str]:
+        configs = self.dict()
+        if project_id:
+            configs["project_id"] = project_id
+        return configs

datahub/ingestion/source/common/subtypes.py

@@ -1,5 +1,10 @@
+import logging
+from typing import Any, Dict
+
 from datahub.utilities.str_enum import StrEnum
 
+logger = logging.getLogger(__name__)
+
 
 class DatasetSubTypes(StrEnum):
     # Generic SubTypes
@@ -25,6 +30,12 @@ class DatasetSubTypes(StrEnum):
     NEO4J_NODE = "Neo4j Node"
     NEO4J_RELATIONSHIP = "Neo4j Relationship"
     SNOWFLAKE_STREAM = "Snowflake Stream"
+    DYNAMIC_TABLE = "Dynamic Table"
+    API_ENDPOINT = "API Endpoint"
+    SLACK_CHANNEL = "Slack Channel"
+    PROJECTIONS = "Projections"
+    GOOGLE_SHEETS = "Google Sheets"
+    GOOGLE_SHEETS_NAMED_RANGE = "Google Sheets Named Range"
 
     # TODO: Create separate entity...
     NOTEBOOK = "Notebook"
@@ -44,13 +55,19 @@ class DatasetContainerSubTypes(StrEnum):
     GCS_BUCKET = "GCS bucket"
     ABS_CONTAINER = "ABS container"
     KEYSPACE = "Keyspace"  # Cassandra
+    NAMESPACE = "Namespace"  # Iceberg
+    DREMIO_SPACE = "Dremio Space"
+    DREMIO_SOURCE = "Dremio Source"
 
 
 class BIContainerSubTypes(StrEnum):
     LOOKER_FOLDER = "Folder"
     LOOKML_PROJECT = "LookML Project"
     LOOKML_MODEL = "LookML Model"
+    TABLEAU_SITE = "Site"
+    TABLEAU_PROJECT = "Project"
     TABLEAU_WORKBOOK = "Workbook"
+    POWERBI_WORKSPACE = "Workspace"
     POWERBI_DATASET = "Semantic Model"
     POWERBI_DATASET_TABLE = "Table"
     QLIK_SPACE = "Qlik Space"
@@ -58,6 +75,8 @@ class BIContainerSubTypes(StrEnum):
     SIGMA_WORKSPACE = "Sigma Workspace"
     SIGMA_WORKBOOK = "Sigma Workbook"
     MODE_COLLECTION = "Collection"
+    GRAFANA_FOLDER = "Folder"
+    GRAFANA_DASHBOARD = "Dashboard"
 
 
 class FlowContainerSubTypes(StrEnum):
@@ -68,10 +87,13 @@ class FlowContainerSubTypes(StrEnum):
 class JobContainerSubTypes(StrEnum):
     NIFI_PROCESS_GROUP = "Process Group"
     MSSQL_JOBSTEP = "Job Step"
-    MSSQL_STORED_PROCEDURE = "Stored Procedure"
+    STORED_PROCEDURE = "Stored Procedure"
 
 
 class BIAssetSubTypes(StrEnum):
+    DASHBOARD = "Dashboard"
+    CHART = "Chart"
+
     # Generic SubTypes
     REPORT = "Report"
 
@@ -93,7 +115,57 @@ class BIAssetSubTypes(StrEnum):
     SAC_STORY = "Story"
     SAC_APPLICATION = "Application"
 
+    # Hex
+    HEX_PROJECT = "Project"
+    HEX_COMPONENT = "Component"
+
 
 class MLAssetSubTypes(StrEnum):
     MLFLOW_TRAINING_RUN = "ML Training Run"
     MLFLOW_EXPERIMENT = "ML Experiment"
+    VERTEX_EXPERIMENT = "Experiment"
+    VERTEX_EXPERIMENT_RUN = "Experiment Run"
+    VERTEX_EXECUTION = "Execution"
+
+    VERTEX_MODEL = "ML Model"
+    VERTEX_MODEL_GROUP = "ML Model Group"
+    VERTEX_TRAINING_JOB = "Training Job"
+    VERTEX_ENDPOINT = "Endpoint"
+    VERTEX_DATASET = "Dataset"
+    VERTEX_PROJECT = "Project"
+    VERTEX_PIPELINE = "Pipeline Job"
+    VERTEX_PIPELINE_TASK = "Pipeline Task"
+    VERTEX_PIPELINE_TASK_RUN = "Pipeline Task Run"
+
+
+def create_source_capability_modifier_enum():
+    all_values: Dict[str, Any] = {}
+    source_enums = [
+        DatasetSubTypes,
+        DatasetContainerSubTypes,
+        BIContainerSubTypes,
+        FlowContainerSubTypes,
+        JobContainerSubTypes,
+        BIAssetSubTypes,
+        MLAssetSubTypes,
+    ]
+
+    for enum_class in source_enums:
+        for member in enum_class:  # type: ignore[var-annotated]
+            if member.name in all_values:
+                logger.debug(
+                    f"Warning: {member.name} already exists with value {all_values[member.name]}, skipping {member.value}"
+                )
+                continue
+            all_values[member.name] = member.value
+
+    enum_code = "class SourceCapabilityModifier(StrEnum):\n"
+    for name, value in all_values.items():
+        enum_code += f'    {name} = "{value}"\n'
+
+    exec(enum_code, globals())
+    return globals()["SourceCapabilityModifier"]
+
+
+# This will have all values from the enums above
+SourceCapabilityModifier = create_source_capability_modifier_enum()

datahub/ingestion/source/data_lake_common/data_lake_utils.py

@@ -11,23 +11,30 @@ from datahub.emitter.mcp_builder import (
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.aws.s3_util import (
-    get_bucket_name,
     get_bucket_relative_path,
     get_s3_prefix,
     is_s3_uri,
 )
 from datahub.ingestion.source.azure.abs_utils import (
     get_abs_prefix,
-    get_container_name,
     get_container_relative_path,
     is_abs_uri,
 )
 from datahub.ingestion.source.common.subtypes import DatasetContainerSubTypes
+from datahub.ingestion.source.data_lake_common.object_store import (
+    get_object_store_bucket_name,
+    get_object_store_for_uri,
+)
+from datahub.ingestion.source.data_lake_common.path_spec import PathSpec
 from datahub.ingestion.source.gcs.gcs_utils import (
-    get_gcs_bucket_name,
     get_gcs_prefix,
     is_gcs_uri,
 )
+from datahub.metadata.schema_classes import (
+    SchemaFieldClass,
+    SchemaFieldDataTypeClass,
+    StringTypeClass,
+)
 
 # hide annoying debug errors from py4j
 logging.getLogger("py4j").setLevel(logging.ERROR)
@@ -38,6 +45,37 @@ PLATFORM_GCS = "gcs"
 PLATFORM_ABS = "abs"
 
 
+def add_partition_columns_to_schema(
+    path_spec: PathSpec, full_path: str, fields: List[SchemaFieldClass]
+) -> None:
+    # Check if using fieldPath v2 format
+    is_fieldpath_v2 = any(
+        field.fieldPath.startswith("[version=2.0]") for field in fields
+    )
+
+    # Extract partition information from path
+    partition_keys = path_spec.get_partition_from_path(full_path)
+    if not partition_keys:
+        return
+
+    # Add partition fields to schema
+    for partition_key in partition_keys:
+        fields.append(
+            SchemaFieldClass(
+                fieldPath=(
+                    f"{partition_key[0]}"
+                    if not is_fieldpath_v2
+                    else f"[version=2.0].[type=string].{partition_key[0]}"
+                ),
+                nativeDataType="string",
+                type=SchemaFieldDataTypeClass(StringTypeClass()),
+                isPartitioningKey=True,
+                nullable=False,
+                recursive=False,
+            )
+        )
+
+
 class ContainerWUCreator:
     processed_containers: List[str]
 
@@ -87,6 +125,13 @@ class ContainerWUCreator:
 
     @staticmethod
     def get_protocol(path: str) -> str:
+        object_store = get_object_store_for_uri(path)
+        if object_store:
+            prefix = object_store.get_prefix(path)
+            if prefix:
+                return prefix
+
+        # Legacy fallback
         protocol: Optional[str] = None
         if is_s3_uri(path):
             protocol = get_s3_prefix(path)
@@ -104,13 +149,12 @@ class ContainerWUCreator:
 
     @staticmethod
     def get_bucket_name(path: str) -> str:
-        if is_s3_uri(path):
-            return get_bucket_name(path)
-        elif is_gcs_uri(path):
-            return get_gcs_bucket_name(path)
-        elif is_abs_uri(path):
-            return get_container_name(path)
-        raise ValueError(f"Unable to get bucket name from path: {path}")
+        """
+        Get the bucket/container name from any supported object store URI.
+
+        Delegates to the abstract get_object_store_bucket_name function.
+        """
+        return get_object_store_bucket_name(path)
 
     def get_sub_types(self) -> str:
         if self.platform == PLATFORM_S3:
@@ -122,6 +166,11 @@ class ContainerWUCreator:
         raise ValueError(f"Unable to sub type for platform: {self.platform}")
 
     def get_base_full_path(self, path: str) -> str:
+        object_store = get_object_store_for_uri(path)
+        if object_store:
+            return object_store.get_object_key(path)
+
+        # Legacy fallback
         if self.platform == "s3" or self.platform == "gcs":
             return get_bucket_relative_path(path)
         elif self.platform == "abs":