alamb commented on code in PR #5015:
URL: https://github.com/apache/arrow-datafusion/pull/5015#discussion_r1083296581
##########
datafusion/core/tests/sqllogictests/postgres/postgres_create_table.sql:
##########
@@ -1,21 +0,0 @@
-CREATE TABLE aggregate_test_100_by_sql
Review Comment:
these statements are now run in the individual setup files
##########
datafusion/core/tests/sqllogictests/src/engines/postgres/mod.rs:
##########
@@ -29,69 +30,157 @@ use postgres_types::Type;
use rust_decimal::Decimal;
use tokio_postgres::{Column, Row};
-pub mod image;
+// default connect string, can be overridden by environment PG_DSN
+const PG_DSN: &str = "postgresql://[email protected]/test";
+
+/// DataFusion sql-logicaltest error
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+ #[error("Postgres error: {0}")]
+ Postgres(#[from] tokio_postgres::error::Error),
+ #[error("Error handling copy command: {0}")]
+ Copy(String),
+}
+
+pub type Result<T, E = Error> = std::result::Result<T, E>;
pub struct Postgres {
- client: Arc<tokio_postgres::Client>,
+ client: tokio_postgres::Client,
join_handle: JoinHandle<()>,
+ /// Filename, for display purposes
file_name: String,
}
impl Postgres {
- pub async fn connect_with_retry(
- file_name: String,
- host: &str,
- port: u16,
- db: &str,
- user: &str,
- pass: &str,
- ) -> Result<Self, tokio_postgres::error::Error> {
+ /// Creates a runner for executiong queries against an existing
+ /// posgres connection. `file_name` is used for display output
+ ///
+ /// The database connection details can be overridden by the
+ /// `PG_DSN` environment variable.
+ ///
+ /// This defaults to
+ ///
+ /// ```text
+ /// PG_DSN="postgresql://[email protected]/test"
+ /// ```
+ ///
+ /// See
https://docs.rs/tokio-postgres/latest/tokio_postgres/config/struct.Config.html#url
for format
+ pub async fn connect(file_name: impl Into<String>) -> Result<Self> {
+ let file_name = file_name.into();
+
+ let dsn = if let Ok(val) = std::env::var("PG_DSN") {
+ val
+ } else {
+ PG_DSN.to_string()
+ };
+
+ debug!("Using posgres dsn: {dsn}");
+
+ let config = tokio_postgres::Config::from_str(&dsn)?;
+
let mut retry = 0;
+
loop {
- let connection_result =
- Postgres::connect(file_name.clone(), host, port, db, user,
pass).await;
+ let connection_result =
config.connect(tokio_postgres::NoTls).await;
+
match connection_result {
Err(e) if retry <= 3 => {
debug!("Retrying connection error '{:?}'", e);
retry += 1;
tokio::time::sleep(Duration::from_secs(1)).await;
}
- result => break result,
+ Err(e) => return Err(Error::from(e)),
+ Ok((client, connection)) => {
+ let join_handle = tokio::spawn(async move {
+ if let Err(e) = connection.await {
+ log::error!("Postgres connection error: {:?}", e);
+ }
+ });
+
+ return Ok(Self {
+ client,
+ join_handle,
+ file_name,
+ });
+ }
}
}
}
- async fn connect(
- file_name: String,
- host: &str,
- port: u16,
- db: &str,
- user: &str,
- pass: &str,
- ) -> Result<Self, tokio_postgres::error::Error> {
- let (client, connection) = tokio_postgres::Config::new()
- .host(host)
- .port(port)
- .dbname(db)
- .user(user)
- .password(pass)
- .connect(tokio_postgres::NoTls)
- .await?;
-
- let join_handle = tokio::spawn(async move {
- if let Err(e) = connection.await {
- log::error!("Postgres connection error: {:?}", e);
+ /// Special COPY command support. "COPY 'filename'" requires the
+ /// server to read the file which may not be possible (maybe it is
+ /// remote or running in some other docker container).
+ ///
+ /// Thus, we rewrite sql statements like
+ ///
+ /// ```sql
+ /// COPY ... FROM 'filename' ...
+ /// ```
+ ///
+ /// Into
+ ///
+ /// ```sql
+ /// COPY ... FROM STDIN ...
+ /// ```
+ ///
+ /// And read the file locally.
+ async fn run_copy_command(&mut self, sql: &str) -> Result<DBOutput> {
Review Comment:
I was quite pleased with this :bowtie: -- it allows the tests to run `COPY
FROM 'file'`
##########
.github/workflows/rust.yml:
##########
@@ -232,9 +232,21 @@ jobs:
name: "Run sqllogictest with Postgres runner"
needs: [linux-build-lib]
runs-on: ubuntu-latest
+ services:
+ postgres:
Review Comment:
this starts the container using the code from `integration-tests` (above in
this file) rather than running docker from within the harness
##########
datafusion/core/tests/sqllogictests/test_files/pg_compat_simple.slt:
##########
@@ -15,6 +15,67 @@
# specific language governing permissions and limitations
# under the License.
+###
+## Setup test for postgres
+###
+
+onlyif postgres
Review Comment:
I used the `onlyif` statement to run different setups for DataFusion and
posgres
##########
datafusion/core/tests/sqllogictests/test_files/aggregate.slt:
##########
@@ -15,6 +15,28 @@
# specific language governing permissions and limitations
# under the License.
+
+statement ok
Review Comment:
drive by clean up -- this can be run directly in the .slt file
--
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]
