Per-source client_config registry on Catalog

src/datachain/catalog/catalog.py

@@ -506,6 +506,24 @@ def clone_catalog_with_cache(catalog: "Catalog", cache: "Cache") -> "Catalog":
     return clone
 
 
+def _copy_client_config(value):
+    """Recursively copy the mapping/sequence structure of a client config so
+    later caller-side mutation (including nested ``client_kwargs`` etc.)
+    cannot change a registered configuration. Leaf objects (credential
+    providers, SSL contexts, ...) are kept by reference."""
+    if isinstance(value, dict):
+        return {k: _copy_client_config(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple)):
+        return type(value)(_copy_client_config(v) for v in value)
+    return value
+
+
+# Auto-detected public-bucket access (see ``read_storage``). Unlike a
+# user-supplied config it is a derived guess, so an explicit config may
+# overwrite it in the registry.
+AUTO_ANON_CLIENT_CONFIG = {"anon": True}
+
+
 class Catalog:
     def __init__(
         self,
@@ -514,6 +532,7 @@ def __init__(
         cache_dir=None,
         tmp_dir=None,
         client_config: dict[str, Any] | None = None,
+        source_client_configs: dict[str, dict[str, Any]] | None = None,
         in_memory: bool = False,
     ):
         datachain_dir = DataChainDir(cache=cache_dir, tmp=tmp_dir)
@@ -522,6 +541,11 @@ def __init__(
         self.warehouse = warehouse
         self.cache = Cache(datachain_dir.cache, datachain_dir.tmp)
         self.client_config = client_config if client_config is not None else {}
+        # Per-source client configs: {File.source -> config}, registered by
+        # read_storage. An entry wins over `client_config` for its source.
+        self.source_client_configs: dict[str, dict[str, Any]] = dict(
+            source_client_configs or {}
+        )
         self._init_params = {
             "cache_dir": cache_dir,
             "tmp_dir": tmp_dir,
@@ -539,6 +563,7 @@ def get_init_params(self) -> dict[str, Any]:
         return {
             **self._init_params,
             "client_config": self.client_config,
+            "source_client_configs": self.source_client_configs,
         }
 
     def copy(self, cache=True, db=True):
@@ -575,11 +600,49 @@ def __exit__(self, exc_type, exc_value, traceback) -> None:
     def generate_query_dataset_name(cls) -> str:
         return f"{QUERY_DATASET_PREFIX}_{uuid4().hex}"
 
+    def register_client_config(
+        self, source: "str | os.PathLike[str]", config: dict[str, Any]
+    ) -> None:
+        """Register `config` for a storage source (a ``File.source`` value:
+        the bucket for cloud storage, the directory for local paths).
+
+        Re-registering the same config is a no-op. An auto-detected
+        ``{"anon": True}`` entry is a derived guess: an explicit config may
+        replace it, and it never replaces or conflicts with an existing
+        entry. Any other mismatch raises, because one catalog cannot hold
+        two configurations for the same source.
+        """
+        key = str(source).rstrip("/")
+        existing = self.source_client_configs.get(key)
+        if existing is not None and config == AUTO_ANON_CLIENT_CONFIG:
+            # A derived guess never overrides or conflicts with an existing
+            # entry.
+            return
+        if existing is not None and existing not in (config, AUTO_ANON_CLIENT_CONFIG):
+            raise ValueError(
+                f"{key} was already accessed with a different client_config "
+                "in this session; pass an explicit Session(client_config=...) "
+                "to isolate it"
+            )
+        self.source_client_configs[key] = _copy_client_config(config)
+
+    def client_config_for(self, source: "str | os.PathLike[str]") -> dict[str, Any]:
+        """The config registered for `source` (a ``File.source`` value), else
+        the catalog-wide default."""
+        config = self.source_client_configs.get(str(source).rstrip("/"))
+        if config is None:
+            return self.client_config
+        if config == AUTO_ANON_CLIENT_CONFIG:
+            # Auto-detected anon refines the default config (which may carry
+            # e.g. an endpoint URL); an explicit config replaces it.
+            return {**self.client_config, **config}
+        return config
+
     def get_client(self, uri: str, **config: Any) -> Client:
         """
         Return the client corresponding to the given source `uri`.
         """
-        config = config or self.client_config
+        config = config or self.client_config_for(uri)
         cls = Client.get_implementation(uri)
         return cls.from_source(StorageURI(uri), self.cache, **config)
 

src/datachain/lib/dc/storage.py

@@ -3,6 +3,7 @@
 from functools import reduce
 from typing import TYPE_CHECKING
 
+from datachain.catalog.catalog import AUTO_ANON_CLIENT_CONFIG
 from datachain.client import Client
 from datachain.lib.dc.storage_pattern import (
     apply_glob_filter,
@@ -180,24 +181,31 @@ def read_storage(
         session.catalog.client_config if session is not None else None
     )
 
-    if (
-        anon is None
-        and not _backends_have_credentials(uris, probe_config)
-        and _all_buckets_anonymous(uris, probe_config)
+    if anon is not None:
+        # Normalize to a plain bool: callers pass truthy strings ("True"),
+        # and the value participates in registry equality checks.
+        anon = bool(anon)
+    elif not _backends_have_credentials(uris, probe_config) and _all_buckets_anonymous(
+        uris, probe_config
     ):
         anon = True
 
     if anon is not None:
         client_config = (client_config or {}) | {"anon": anon}
-    session = Session.get(session, client_config=client_config, in_memory=in_memory)
+    session = Session.get(session, in_memory=in_memory)
     catalog = session.catalog
     cache = catalog.cache
-    client_config = session.catalog.client_config
-    if anon is not None:
-        # Session.get discards our client_config when an existing session is
-        # passed. Re-apply anon locally for the listing path without mutating
-        # the caller's session.
-        client_config = client_config | {"anon": anon}
+    # The per-call config (including auto-detected anon) is registered for
+    # each listed source below, so every later access to it — listing, file
+    # reads inside UDFs (including worker processes), exports — resolves the
+    # same config through the catalog instead of a session-wide default.
+    per_source_config = client_config or None
+    # A bare {"anon": True} refines the catalog default at use time (see
+    # Catalog.client_config_for); give get_listing the same effective view so
+    # its file/dir probe talks to the right endpoint.
+    listing_config = per_source_config
+    if listing_config == AUTO_ANON_CLIENT_CONFIG:
+        listing_config = catalog.client_config | listing_config
     listing_namespace_name = catalog.metastore.system_namespace_name
     listing_project_name = catalog.metastore.listing_project_name
 
@@ -229,9 +237,18 @@ def read_storage(
             update_single_uri = True
 
         list_ds_name, list_uri, list_path, _ = get_listing(
-            list_uri_to_use, session, update=update_single_uri
+            list_uri_to_use,
+            session,
+            update=update_single_uri,
+            client_config=listing_config,
         )
 
+        # `list_uri` parses to the same source the listed files carry.
+        source, _ = Client.parse_url(list_uri)
+        if per_source_config:
+            catalog.register_client_config(source, per_source_config)
+        client_config = catalog.client_config_for(source)
+
         # list_ds_name is None if object is a file, we don't want to use cache
         # or do listing in that case - just read that single object
         if not list_ds_name:
@@ -256,10 +273,10 @@ def read_storage(
         dc._query.update = update
         dc.signals_schema = dc.signals_schema.mutate({f"{column}": file_type})
 
-        def lst_fn(ds_name, lst_uri):
+        def lst_fn(ds_name, lst_uri, lst_config):
             # Seed for .gen() iteration. content_hash=None because hash_callable
             # doesn't capture list_func's closure (which holds `lst_uri`, `cache`,
-            # `client_config`) -- auto-hashing the seed would let the UDF
+            # `lst_config`) -- auto-hashing the seed would let the UDF
             # checkpoint cache return a stale listing across URIs/runs.
             (
                 create_records_dataset(
@@ -276,16 +293,19 @@ def lst_fn(ds_name, lst_uri):
                     project=listing_project_name,
                 )
                 .gen(
-                    list_bucket(lst_uri, cache, client_config=client_config),
+                    list_bucket(lst_uri, cache, client_config=lst_config),
                     output={f"{column}": file_type},
                 )
                 # for internal listing datasets, we always bump major version
                 .save(ds_name, listing=True, update_version="major")
             )
 
         # Always attach listing_fn so resolve_listing can refresh stale listings.
+        # Bind loop variables via defaults: `client_config` differs per URI.
         dc._query.set_listing_fn(
-            lambda ds_name=list_ds_name, lst_uri=list_uri: lst_fn(ds_name, lst_uri)
+            lambda ds_name=list_ds_name, lst_uri=list_uri, cfg=client_config: lst_fn(
+                ds_name, lst_uri, cfg
+            )
         )
 
         # If a glob pattern was detected, use it for filtering

src/datachain/lib/file.py

@@ -571,7 +571,7 @@ def open(
         if self._catalog is None:
             raise RuntimeError("Cannot open file: catalog is not set")
 
-        base_cfg = getattr(self._catalog, "client_config", {}) or {}
+        base_cfg = self._catalog.client_config_for(self.source)
         merged_cfg = {**base_cfg, **(client_config or {})}
         client: Client = self._catalog.get_client(self.source, **merged_cfg)
 
@@ -741,7 +741,11 @@ def save(
     def _resolve_destination(
         self, destination: str, client_config: dict | None = None
     ) -> "tuple[Client, str]":
-        """Return (client rooted at the storage, rel_path) for *destination*."""
+        """Return (client rooted at the storage, rel_path) for *destination*.
+
+        With no explicit `client_config`, the client resolves the config
+        registered for the destination's source, if any (via `get_client`).
+        """
         from datachain.client.fsspec import Client as FSClient
 
         uri = path_to_fsspec_uri(destination)
@@ -831,7 +835,7 @@ def export(
 
         suffix = self._get_destination_suffix(placement)
         output_str = stringify_path(output)
-        client = self._catalog.get_client(output_str, **(client_config or {}))
+        client, _ = self._resolve_destination(output_str, client_config)
 
         # Normalization and traversal safety: for local exports, resolve to absolute
         # and validate the suffix. Cloud exports skip this — the cloud client already

src/datachain/lib/listing.py

@@ -256,7 +256,10 @@ def _reraise_as_client_error() -> Iterator[None]:
 
 
 def get_listing(
-    uri: str | os.PathLike[str], session: "Session", update: bool = False
+    uri: str | os.PathLike[str],
+    session: "Session",
+    update: bool = False,
+    client_config: dict | None = None,
 ) -> tuple[str | None, str, str, bool]:
     """Returns correct listing dataset name that must be used for saving listing
     operation. It takes into account existing listings and reusability of those.
@@ -270,12 +273,13 @@ def get_listing(
 
     catalog = session.catalog
     cache = catalog.cache
-    client_config = catalog.client_config
+    if not isinstance(uri, str):
+        uri = os.fspath(uri)
+    if client_config is None:
+        client_config = catalog.client_config_for(Client.parse_url(uri)[0])
 
     client = Client.get_client(uri, cache, **client_config)
     telemetry.log_param("client", client.PREFIX)
-    if not isinstance(uri, str):
-        uri = os.fspath(uri)
 
     # we don't want to use cached dataset (e.g. for a single file listing)
     isfile = _reraise_as_client_error()(fsutils.isfile)

src/datachain/lib/zarr.py

@@ -79,7 +79,8 @@ def _open(self, mode: Literal["r", "r+", "a", "w", "w-"] = "r") -> Any:
         storage_options = None
         if f.source and not f.source.startswith("file://"):
             catalog = getattr(f, "_catalog", None)
-            storage_options = getattr(catalog, "client_config", None) or None
+            if catalog is not None:
+                storage_options = catalog.client_config_for(f.source) or None
         if storage_options:
             return zarr.open(url, mode=mode, storage_options=storage_options)
         return zarr.open(url, mode=mode)

src/datachain/query/session.py

@@ -298,12 +298,18 @@ def get(
         client_config: dict | None = None,
         in_memory: bool = False,
     ) -> "Session":
-        """Creates a Session() object from a catalog.
+        """Resolve the session to use: an explicit `session`, else the active
+        context, else the process-global session (created on first use).
 
         Parameters:
-            session (Session): Optional Session(). If not provided a new session will
-                    be created. It's needed mostly for simple API purposes.
-            catalog (Catalog): Optional catalog. By default, a new catalog is created.
+            session (Session): Optional Session(). If not provided the ambient
+                    session is resolved as described above.
+            catalog (Catalog): Optional catalog; used only when this call
+                    creates the global session.
+            client_config (dict): Optional storage client config; used only
+                    when this call creates the global session. Per-source
+                    configs never fork a session — they are registered on the
+                    catalog (see `Catalog.register_client_config`).
         """
         if session:
             return session
@@ -327,15 +333,6 @@ def get(
         else:
             session = cls.GLOBAL_SESSION_CTX
 
-        if client_config and session.catalog.client_config != client_config:
-            session = Session(
-                "session" + uuid4().hex[:4],
-                catalog,
-                client_config=client_config,
-                in_memory=in_memory,
-            )
-            session.__enter__()
-
         return session
 
     @staticmethod