ray-project--ray/python/ray/data/tests/test_catalog.py

"""Tests for the Catalog connector API (ray.data.catalog)."""

import contextlib
import os
import pickle
from unittest import mock

import pyarrow as pa
import pyarrow.fs as pafs
import pyarrow.parquet as pq
import pytest

import ray
from ray.data.catalog import (
    Catalog,
    DatabricksUnityCatalog,
    ReaderFormat,
    ResolvedSource,
)

# conftest provides ray_start_regular_shared
from ray.data.tests.conftest import *  # noqa: F401,F403

pytest.importorskip("databricks.sdk")
from databricks.sdk.service.catalog import (  # noqa: E402
    AwsCredentials,
    AzureUserDelegationSas,
    DataSourceFormat,
    GcpOauthToken,
    GenerateTemporaryTableCredentialResponse,
    TableInfo,
)

AWS_RESP = GenerateTemporaryTableCredentialResponse(
    url="s3://bucket/path",
    aws_temp_credentials=AwsCredentials(
        access_key_id="AKIA",
        secret_access_key="secret",
        session_token="token",
    ),
)


def _mock_uc_sdk(*, data_source_format="DELTA", storage_location=None, creds=None):
    """Patch DatabricksUnityCatalog._workspace_client to return canned SDK responses.

    Replaces the catalog's two SDK calls (``tables.get`` /
    ``generate_temporary_table_credentials``) so no Databricks workspace is hit.
    Returns an already-started patcher; the caller is responsible for ``stop()``.
    """
    creds = creds if creds is not None else AWS_RESP
    table_info = TableInfo(
        table_id="tid-123",
        data_source_format=(
            DataSourceFormat(data_source_format) if data_source_format else None
        ),
        storage_location=storage_location,
    )
    client = mock.MagicMock()
    client.tables.get.return_value = table_info
    gen = client.temporary_table_credentials.generate_temporary_table_credentials
    gen.return_value = creds

    patcher = mock.patch.object(
        DatabricksUnityCatalog, "_workspace_client", return_value=client
    )
    patcher.start()
    return patcher


@pytest.fixture
def uc_catalog():
    return DatabricksUnityCatalog(
        url="https://dbc-test.cloud.databricks.com",
        token="dapi-test",
        region="us-west-2",
    )


@pytest.fixture
def isolated_env(monkeypatch):
    # resolve() exports vended creds to os.environ and may call ray.init.
    # Isolate the env mutation and skip the real ray.init for unit tests.
    monkeypatch.setattr("ray.is_initialized", lambda: True)
    monkeypatch.setattr(os, "environ", dict(os.environ))
    return os.environ


# ---------------------------------------------------------------------------
# DatabricksUnityCatalog.resolve
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("reader", [ReaderFormat.DELTA, ReaderFormat.PARQUET])
def test_resolve_storage_aws(uc_catalog, isolated_env, reader):
    patcher = _mock_uc_sdk()
    try:
        resolved = uc_catalog.resolve("main.sales.txns", reader=reader)
    finally:
        patcher.stop()

    assert resolved.path == "s3://bucket/path"
    # AWS creds are exported to the environment for both readers.
    assert isolated_env["AWS_ACCESS_KEY_ID"] == "AKIA"
    assert isolated_env["AWS_SESSION_TOKEN"] == "token"
    # Delta additionally gets an explicit S3FileSystem for the data scan;
    # Parquet reads via the environment-configured filesystem.
    if reader is ReaderFormat.DELTA:
        assert isinstance(resolved.filesystem, pafs.S3FileSystem)
    else:
        assert resolved.filesystem is None
    assert resolved.data_format is ReaderFormat.DELTA


def test_resolve_aws_requires_region(isolated_env):
    catalog = DatabricksUnityCatalog(
        url="https://h.databricks.com", token="t"
    )  # no region
    patcher = _mock_uc_sdk()
    try:
        with pytest.raises(ValueError, match="region"):
            catalog.resolve("main.sales.txns", reader=ReaderFormat.DELTA)
    finally:
        patcher.stop()


def test_resolve_initializes_ray_with_runtime_env(uc_catalog, monkeypatch):
    # When Ray isn't running, vended creds are propagated via runtime_env.
    monkeypatch.setattr(os, "environ", dict(os.environ))
    monkeypatch.setattr("ray.is_initialized", lambda: False)
    init_kwargs = {}
    monkeypatch.setattr("ray.init", lambda **kw: init_kwargs.update(kw))

    patcher = _mock_uc_sdk()
    try:
        uc_catalog.resolve("main.sales.txns", reader=ReaderFormat.PARQUET)
    finally:
        patcher.stop()

    env_vars = init_kwargs["runtime_env"]["env_vars"]
    assert env_vars["AWS_ACCESS_KEY_ID"] == "AKIA"
    assert env_vars["AWS_SESSION_TOKEN"] == "token"


def test_resolve_retries_once_on_401(uc_catalog, isolated_env):
    # On a 401, the provider is invalidated and the SDK call is retried once
    # (mirrors the old request_with_401_retry behavior).
    from databricks.sdk.errors import Unauthenticated

    client = mock.MagicMock()
    client.tables.get.side_effect = [
        Unauthenticated("expired"),
        TableInfo(
            table_id="tid-123",
            data_source_format=DataSourceFormat("DELTA"),
            storage_location="s3://bucket/path",
        ),
    ]
    gen = client.temporary_table_credentials.generate_temporary_table_credentials
    gen.return_value = AWS_RESP

    with mock.patch.object(
        DatabricksUnityCatalog, "_workspace_client", return_value=client
    ), mock.patch.object(uc_catalog._provider, "invalidate") as invalidate:
        resolved = uc_catalog.resolve("main.sales.txns", reader=ReaderFormat.PARQUET)

    invalidate.assert_called_once()
    assert client.tables.get.call_count == 2
    assert resolved.path == "s3://bucket/path"


def test_resolve_iceberg(uc_catalog):
    # Iceberg resolution does not hit the credential-vending REST endpoints.
    resolved = uc_catalog.resolve("main.sales.txns", reader=ReaderFormat.ICEBERG)

    assert resolved.path is None
    assert resolved.filesystem is None
    assert resolved.data_format is ReaderFormat.ICEBERG
    ckw = resolved.catalog_kwargs
    assert ckw["type"] == "rest"
    assert ckw["uri"] == (
        "https://dbc-test.cloud.databricks.com/api/2.1/unity-catalog/iceberg-rest"
    )
    assert ckw["token"] == "dapi-test"
    assert ckw["header.X-Iceberg-Access-Delegation"] == "vended-credentials"


def test_resolve_azure_sets_env(isolated_env):
    catalog = DatabricksUnityCatalog(url="https://h.databricks.com", token="t")
    azure_resp = GenerateTemporaryTableCredentialResponse(
        url="abfss://c@acct.dfs.core.windows.net/path",
        azure_user_delegation_sas=AzureUserDelegationSas(sas_token="sv=2021&sig=abc"),
    )
    patcher = _mock_uc_sdk(creds=azure_resp)
    try:
        resolved = catalog.resolve("main.sales.txns", reader=ReaderFormat.DELTA)
    finally:
        patcher.stop()

    # Azure creds flow via the environment (read by both pyarrow and the
    # deltalake object_store log reader); no filesystem/storage_options.
    assert resolved.filesystem is None
    assert resolved.storage_options is None
    assert isolated_env["AZURE_STORAGE_SAS_TOKEN"] == "sv=2021&sig=abc"


def test_resolve_azure_strips_leading_question_mark(isolated_env):
    # UC may return the SAS as a full query string ("?sv=..."); the leading "?"
    # must be stripped for AZURE_STORAGE_SAS_TOKEN.
    catalog = DatabricksUnityCatalog(url="https://h.databricks.com", token="t")
    azure_resp = GenerateTemporaryTableCredentialResponse(
        url="abfss://c@acct.dfs.core.windows.net/path",
        azure_user_delegation_sas=AzureUserDelegationSas(sas_token="?sv=2021&sig=abc"),
    )
    patcher = _mock_uc_sdk(creds=azure_resp)
    try:
        catalog.resolve("main.sales.txns", reader=ReaderFormat.DELTA)
    finally:
        patcher.stop()

    assert isolated_env["AZURE_STORAGE_SAS_TOKEN"] == "sv=2021&sig=abc"


def _gcp_resp():
    return GenerateTemporaryTableCredentialResponse(
        url="gs://bucket/path",
        gcp_oauth_token=GcpOauthToken(oauth_token="ya29.tok"),
        expiration_time=4102444800000,  # far-future epoch ms
    )


def test_resolve_gcp_parquet_builds_filesystem(uc_catalog, isolated_env):
    # GCP vends an OAuth token (not a service-account JSON); for Parquet it rides
    # on an explicit GcsFileSystem (the data scan), never an env var.
    patcher = _mock_uc_sdk(creds=_gcp_resp(), storage_location="gs://bucket/path")
    try:
        resolved = uc_catalog.resolve("main.sales.txns", reader=ReaderFormat.PARQUET)
    finally:
        patcher.stop()

    assert isinstance(resolved.filesystem, pafs.GcsFileSystem)
    assert "GOOGLE_APPLICATION_CREDENTIALS" not in isolated_env


def test_resolve_gcp_delta_raises(uc_catalog, isolated_env):
    # deltalake's object_store can't use a GCS OAuth token, so GCP + Delta is
    # rejected up front with an actionable error instead of failing deep in the
    # log read against the GCE metadata server.
    patcher = _mock_uc_sdk(creds=_gcp_resp(), storage_location="gs://bucket/path")
    try:
        with pytest.raises(RuntimeError, match="GCP-backed Delta"):
            uc_catalog.resolve("main.sales.txns", reader=ReaderFormat.DELTA)
    finally:
        patcher.stop()


# ---------------------------------------------------------------------------
# Reader integration via a fake catalog (no network)
# ---------------------------------------------------------------------------


class _FakeCatalog(Catalog):
    """Returns a pre-baked ResolvedSource; records the reader it was asked for."""

    def __init__(self, resolved):
        self._resolved = resolved
        self.calls = []

    def resolve(self, table, *, reader):
        self.calls.append((table, reader))
        return self._resolved


def test_read_parquet_with_catalog(ray_start_regular_shared, tmp_path):
    path = str(tmp_path / "data.parquet")
    pq.write_table(pa.table({"id": [1, 2, 3]}), path)

    catalog = _FakeCatalog(ResolvedSource(path=path))
    ds = ray.data.read_parquet("main.db.tbl", catalog=catalog)

    assert sorted(r["id"] for r in ds.take_all()) == [1, 2, 3]
    assert catalog.calls == [("main.db.tbl", ReaderFormat.PARQUET)]


@pytest.mark.parametrize("reader", ["parquet", "delta"])
def test_catalog_filesystem_overrides_with_warning(reader):
    # The catalog-resolved filesystem overrides a user-supplied one, but warns.
    # The warning fires at the top of the reader body; suppress any downstream
    # failure from the (intentionally unreachable) s3 path.
    if reader == "delta":
        pytest.importorskip("deltalake")
    fs = pafs.S3FileSystem(
        access_key="AKIA", secret_key="secret", session_token="t", region="us-west-2"
    )
    catalog = _FakeCatalog(ResolvedSource(path="s3://b/p", filesystem=fs))
    read_fn = ray.data.read_parquet if reader == "parquet" else ray.data.read_delta

    with mock.patch.object(ray.data.read_api.logger, "warning") as warn:
        with contextlib.suppress(Exception):
            read_fn("main.db.tbl", catalog=catalog, filesystem=pafs.LocalFileSystem())

    assert any(
        "Overriding the provided `filesystem`" in str(c) for c in warn.call_args_list
    )


def test_read_delta_with_catalog(ray_start_regular_shared, tmp_path):
    deltalake = pytest.importorskip("deltalake")
    path = str(tmp_path / "delta-table")
    deltalake.write_deltalake(path, pa.table({"id": [1, 2, 3]}))

    catalog = _FakeCatalog(ResolvedSource(path=path))
    ds = ray.data.read_delta("main.db.tbl", catalog=catalog)

    assert sorted(r["id"] for r in ds.take_all()) == [1, 2, 3]
    assert catalog.calls == [("main.db.tbl", ReaderFormat.DELTA)]


def test_read_iceberg_uses_catalog_resolved_kwargs():
    catalog = _FakeCatalog(
        ResolvedSource(catalog_kwargs={"type": "rest", "uri": "u", "token": "tk"})
    )
    with mock.patch(
        "ray.data._internal.datasource.iceberg_datasource.IcebergDatasource"
    ) as ds_cls, mock.patch("ray.data.read_api.read_datasource"):
        ray.data.read_iceberg(table_identifier="main.db.tbl", catalog=catalog)

    _, kwargs = ds_cls.call_args
    assert kwargs["catalog_kwargs"] == {"type": "rest", "uri": "u", "token": "tk"}
    assert catalog.calls == [("main.db.tbl", ReaderFormat.ICEBERG)]


def test_read_iceberg_explicit_catalog_kwargs_take_precedence():
    # When both catalog and catalog_kwargs are given, catalog is ignored.
    catalog = _FakeCatalog(ResolvedSource(catalog_kwargs={"type": "rest", "uri": "u"}))
    with mock.patch(
        "ray.data._internal.datasource.iceberg_datasource.IcebergDatasource"
    ) as ds_cls, mock.patch("ray.data.read_api.read_datasource"):
        ray.data.read_iceberg(
            table_identifier="main.db.tbl",
            catalog=catalog,
            catalog_kwargs={"type": "sql", "uri": "explicit"},
        )

    _, kwargs = ds_cls.call_args
    assert kwargs["catalog_kwargs"] == {"type": "sql", "uri": "explicit"}
    assert catalog.calls == []  # catalog was not consulted


# ---------------------------------------------------------------------------
# Serialization
# ---------------------------------------------------------------------------


def test_unity_catalog_is_picklable(uc_catalog):
    restored = pickle.loads(pickle.dumps(uc_catalog))
    assert isinstance(restored, DatabricksUnityCatalog)
    assert restored.region == "us-west-2"


def test_resolved_source_with_filesystem_is_picklable():
    fs = pafs.S3FileSystem(
        access_key="AKIA", secret_key="secret", session_token="t", region="us-west-2"
    )
    src = ResolvedSource(path="s3://b/p", filesystem=fs, data_format=ReaderFormat.DELTA)
    restored = pickle.loads(pickle.dumps(src))
    assert restored.path == "s3://b/p"
    assert isinstance(restored.filesystem, pafs.S3FileSystem)
    assert restored.data_format is ReaderFormat.DELTA


# ---------------------------------------------------------------------------
# Deprecated read_unity_catalog shim
# ---------------------------------------------------------------------------


def test_read_unity_catalog_deprecation_delegates():
    with mock.patch("ray.data.read_api.read_delta") as read_delta:
        with pytest.warns(DeprecationWarning, match="read_unity_catalog"):
            ray.data.read_unity_catalog(
                table="main.db.tbl",
                url="https://h.databricks.com",
                token="t",
                data_format="delta",
            )

    read_delta.assert_called_once()
    _, kwargs = read_delta.call_args
    assert isinstance(kwargs["catalog"], DatabricksUnityCatalog)


def test_read_unity_catalog_infers_format_from_cred_url():
    # Metadata omits both data_source_format and storage_location; the vended
    # credential URL extension must still identify the format.
    creds = GenerateTemporaryTableCredentialResponse(url="s3://bucket/data.parquet")
    patcher = _mock_uc_sdk(data_source_format=None, storage_location=None, creds=creds)
    try:
        with mock.patch("ray.data.read_api.read_parquet") as read_parquet, pytest.warns(
            DeprecationWarning
        ):
            ray.data.read_unity_catalog(
                table="main.db.tbl", url="https://h.databricks.com", token="t"
            )
    finally:
        patcher.stop()

    read_parquet.assert_called_once()


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))