Handle target_session_attrs

Author: sfackler

tokio-postgres/src/connect.rs

@@ -2,7 +2,9 @@ use crate::config::{Host, TargetSessionAttrs};
 use crate::connect_raw::connect_raw;
 use crate::connect_socket::connect_socket;
 use crate::tls::{MakeTlsConnect, TlsConnect};
-use crate::{Client, Config, Connection, Error, Socket};
+use crate::{Client, Config, Connection, Error, SimpleQueryMessage, Socket};
+use futures::TryStreamExt;
+use std::io;
 
 pub async fn connect<T>(
     mut tls: T,
@@ -50,10 +52,27 @@ where
     T: TlsConnect<Socket>,
 {
     let socket = connect_socket(idx, config).await?;
-    let (client, connection) = connect_raw(socket, tls, config, Some(idx)).await?;
+    let (mut client, connection) = connect_raw(socket, tls, config, Some(idx)).await?;
 
     if let TargetSessionAttrs::ReadWrite = config.target_session_attrs {
-        unimplemented!()
+        let mut rows = client.simple_query("SHOW transaction_read_only");
+
+        loop {
+            match rows.try_next().await? {
+                Some(SimpleQueryMessage::Row(row)) => {
+                    if row.try_get(0)? == Some("on") {
+                        return Err(Error::connect(io::Error::new(
+                            io::ErrorKind::PermissionDenied,
+                            "database does not allow writes",
+                        )));
+                    } else {
+                        break;
+                    }
+                }
+                Some(_) => {}
+                None => return Err(Error::unexpected_message()),
+            }
+        }
     }
 
     Ok((client, connection))

tokio-postgres/tests/test/runtime.rs

@@ -49,7 +49,6 @@ async fn wrong_port_count() {
         .unwrap();
 }
 
-/*
 #[tokio::test]
 async fn target_session_attrs_ok() {
     tokio_postgres::connect(
@@ -67,9 +66,13 @@ async fn target_session_attrs_err() {
         "host=localhost port=5433 user=postgres target_session_attrs=read-write
          options='-c default_transaction_read_only=on'",
         NoTls,
-    ).await.err().unwrap();
+    )
+    .await
+    .err()
+    .unwrap();
 }
 
+/*
 #[test]
 fn cancel_query() {
     let mut runtime = Runtime::new().unwrap();