kosiew commented on code in PR #22839: URL: https://github.com/apache/datafusion/pull/22839#discussion_r3379639116
########## datafusion-cli/src/object_storage/stdin.rs: ########## @@ -0,0 +1,230 @@ +// 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 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) + } + + /// Builds the object store backing the `stdin://` scheme by reading all of + /// standard input into memory. + /// + /// A pipe (e.g. `cat data.csv | datafusion-cli`) is not seekable and reports + /// a size of `0`, so it cannot be read directly by the file based formats + /// (CSV requires seeking, Parquet needs the footer at the end of the file). + /// Buffering the whole input up front sidesteps these limitations and lets + /// the data be read like any other object, including being scanned more than + /// once. + pub(crate) async fn object_store(url: &Url) -> Result<Arc<dyn ObjectStore>> { Review Comment: I think there's a correctness issue here. `object_store` re-reads process stdin every time a `stdin://` store is registered. After the first `CREATE EXTERNAL TABLE ... LOCATION '/dev/stdin'`, the stdin contents have already been consumed and buffered. If a second stdin-backed table is created in the same session, this path reads EOF and replaces the registered `stdin://` store with an empty one. As a result, the first table silently changes from containing data to appearing empty. Repro: ```sql CREATE EXTERNAL TABLE t ... LOCATION '/dev/stdin'; SELECT count(*) FROM t; -- 2 CREATE EXTERNAL TABLE t2 ... LOCATION '/dev/stdin'; SELECT count(*) FROM t; -- 0 ``` Could we enforce the one-shot stdin invariant when registering the store? Two approaches that seem reasonable are: * Reject subsequent `/dev/stdin` registrations with a clear error before replacing the existing store. * Cache and reuse the first buffered store for the same `stdin://` URL. The important part is avoiding replacement of the backing object after tables have already been registered. -- 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]
