kosiew commented on code in PR #22839: URL: https://github.com/apache/datafusion/pull/22839#discussion_r3385197198
########## datafusion-cli/src/object_storage/stdin.rs: ########## @@ -0,0 +1,270 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Exposes the process's standard input as a `stdin://` object store so that +//! piped data (e.g. `cat data.csv | datafusion-cli`) can be queried via +//! `CREATE EXTERNAL TABLE ... LOCATION '/dev/stdin'`. + +use std::io::Read; +use std::sync::Arc; + +use datafusion::common::exec_datafusion_err; +use datafusion::config::ConfigFileType; +use datafusion::error::Result; +use datafusion::execution::context::SessionState; + +use object_store::memory::InMemory; +use object_store::path::Path as ObjectStorePath; +use object_store::{ObjectStore, ObjectStoreExt}; +use url::Url; + +/// Utilities for exposing the process's standard input as an object store. +/// +/// stdin is surfaced as a `stdin://` object store and dispatched alongside the +/// other schemes (`s3`, `gs`, `http`, ...) so that reading piped data flows +/// through the normal object-store/listing code path, conceptually similar to +/// DuckDB's `PipeFileSystem`. +pub(crate) struct StdinUtils; + +impl StdinUtils { + /// The URL scheme used to expose stdin as an object store, mirroring how + /// `s3`, `gs`, `http`, etc. are addressed. + pub(crate) const SCHEME: &'static str = "stdin"; + + /// Filesystem paths that refer to the process's standard input. + /// + /// These are intentionally limited to the well known pseudo-files exposed + /// by the operating system so that ordinary files are never accidentally + /// treated as stdin. + const LOCATIONS: [&'static str; 3] = ["/dev/stdin", "/dev/fd/0", "/proc/self/fd/0"]; + + /// Returns `true` if `location` refers to the process's standard input. + fn is_stdin_location(location: &str) -> bool { + Self::LOCATIONS.contains(&location) + } + + /// Rewrites the well known stdin pseudo-paths (e.g. `/dev/stdin`) to a + /// canonical `stdin://` URL so that reading from standard input flows + /// through the same object-store/listing code path as any other scheme. + /// Non-stdin locations are returned unchanged. + /// + /// The listing layer filters candidate files by extension, so the canonical + /// object is named with the extension matching the declared `STORED AS` + /// format. + pub(crate) fn rewrite_location( + location: &str, + format: Option<&ConfigFileType>, + ) -> String { + if !Self::is_stdin_location(location) { + return location.to_string(); + } + + let object_name = match format { + Some(ConfigFileType::CSV) => "stdin.csv", + Some(ConfigFileType::JSON) => "stdin.json", + Some(ConfigFileType::PARQUET) => "stdin.parquet", + _ => "stdin", + }; + format!("{}:///{object_name}", Self::SCHEME) + } + + /// Returns the object store backing the `stdin://` scheme, reading and + /// buffering standard input on first use and reusing that buffer for any + /// subsequent `stdin://` table created in the same session. + /// + /// stdin is a one-shot stream: it can only be read once. The object store + /// registry keys by scheme/authority, so every `stdin://` URL maps to the + /// same store. Without this guard, a second `CREATE EXTERNAL TABLE ... + /// LOCATION '/dev/stdin'` would re-read (now-EOF) stdin, build an empty + /// store, and overwrite the populated one, silently emptying the earlier + /// table. Reusing the already-registered store avoids that. + pub(crate) async fn get_or_create( + state: &SessionState, + url: &Url, + ) -> Result<Arc<dyn ObjectStore>> { + if let Ok(existing) = state.runtime_env().object_store_registry.get_store(url) { Review Comment: Nice catch reusing the existing registered store, but I think there's still a corner case here. `rewrite_location` produces format-specific paths such as `stdin:///stdin.csv`, `stdin:///stdin.json`, and `stdin:///stdin.parquet`. The object-store registry lookup is keyed by scheme/authority rather than path, so after the first table registers `stdin:///stdin.csv`, a later `/dev/stdin` table rewritten to `stdin:///stdin.json` will hit this branch and reuse the existing store. The problem is that the store only contains `/stdin.csv`. The second table then attempts to read `/stdin.json` from that store and fails (or appears missing) rather than reusing the buffered stdin bytes or producing a clear error. So this fixes the reported CSV/CSV case, but not the broader invariant that subsequent `/dev/stdin` registrations should remain safe and predictable. Could we either make the stored object path canonical across formats, populate the newly requested path from the buffered stdin data, or explicitly reject a second `/dev/stdin` registration when the rewritten path differs and return a clear error? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
