psqlpy-python · chandr-andr · May 19, 2025 · May 17, 2025
diff --git a/src/connection/impls.rs b/src/connection/impls.rs
@@ -1,15 +1,22 @@
+use std::sync::{Arc, RwLock};
+
 use bytes::Buf;
 use pyo3::{PyAny, Python};
-use tokio_postgres::{CopyInSink, Row, Statement, ToStatement};
+use tokio_postgres::{CopyInSink, Portal as tp_Portal, Row, Statement, ToStatement};
 
 use crate::{
+    driver::portal::Portal,
     exceptions::rust_errors::{PSQLPyResult, RustPSQLDriverError},
     options::{IsolationLevel, ReadVariant},
     query_result::{PSQLDriverPyQueryResult, PSQLDriverSinglePyQueryResult},
     statement::{statement::PsqlpyStatement, statement_builder::StatementBuilder},
+    transaction::structs::PSQLPyTransaction,
     value_converter::to_python::postgres_to_py,
 };
 
+use deadpool_postgres::Transaction as dp_Transaction;
+use tokio_postgres::Transaction as tp_Transaction;
+
 use super::{
     structs::{PSQLPyConnection, PoolConnection, SingleConnection},
     traits::{CloseTransaction, Connection, Cursor, StartTransaction, Transaction},
@@ -516,4 +523,42 @@ impl PSQLPyConnection {
             }
         }
     }
+
+    pub async fn transaction(&mut self) -> PSQLPyResult<PSQLPyTransaction> {
+        match self {
+            PSQLPyConnection::PoolConn(conn) => {
+                let transaction = unsafe {
+                    std::mem::transmute::<dp_Transaction<'_>, dp_Transaction<'static>>(
+                        conn.connection.transaction().await?,
+                    )
+                };
+                Ok(PSQLPyTransaction::PoolTransaction(transaction))
+            }
+            PSQLPyConnection::SingleConnection(conn) => {
+                let transaction = unsafe {
+                    std::mem::transmute::<tp_Transaction<'_>, tp_Transaction<'static>>(
+                        conn.connection.transaction().await?,
+                    )
+                };
+                Ok(PSQLPyTransaction::SingleTransaction(transaction))
+            }
+        }
+    }
+
+    pub async fn portal(
+        &mut self,
+        querystring: String,
+        parameters: Option<pyo3::Py<PyAny>>,
+    ) -> PSQLPyResult<(PSQLPyTransaction, tp_Portal)> {
+        let statement = StatementBuilder::new(querystring, parameters, self, Some(false))
+            .build()
+            .await?;
+
+        let transaction = self.transaction().await?;
+        let inner_portal = transaction
+            .portal(statement.raw_query(), &statement.params())
+            .await?;
+
+        Ok((transaction, inner_portal))
+    }
 }
diff --git a/src/driver/connection.rs b/src/driver/connection.rs
@@ -18,7 +18,9 @@ use crate::{
     runtime::tokio_runtime,
 };
 
-use super::{connection_pool::connect_pool, cursor::Cursor, transaction::Transaction};
+use super::{
+    connection_pool::connect_pool, cursor::Cursor, portal::Portal, transaction::Transaction,
+};
 
 /// Make new connection pool.
 ///
@@ -396,17 +398,16 @@ impl Connection {
         read_variant: Option<ReadVariant>,
         deferrable: Option<bool>,
     ) -> PSQLPyResult<Transaction> {
-        if let Some(db_client) = &self.db_client {
-            return Ok(Transaction::new(
-                Some(db_client.clone()),
-                self.pg_config.clone(),
-                isolation_level,
-                read_variant,
-                deferrable,
-            ));
-        }
-
-        Err(RustPSQLDriverError::ConnectionClosedError)
+        let Some(conn) = &self.db_client else {
+            return Err(RustPSQLDriverError::ConnectionClosedError);
+        };
+        Ok(Transaction::new(
+            Some(conn.clone()),
+            self.pg_config.clone(),
+            isolation_level,
+            read_variant,
+            deferrable,
+        ))
     }
 
     /// Create new cursor object.
@@ -428,19 +429,39 @@ impl Connection {
         scroll: Option<bool>,
         prepared: Option<bool>,
     ) -> PSQLPyResult<Cursor> {
-        if let Some(db_client) = &self.db_client {
-            return Ok(Cursor::new(
-                db_client.clone(),
-                self.pg_config.clone(),
-                querystring,
-                parameters,
-                fetch_number.unwrap_or(10),
-                scroll,
-                prepared,
-            ));
-        }
+        let Some(conn) = &self.db_client else {
+            return Err(RustPSQLDriverError::ConnectionClosedError);
+        };
+
+        Ok(Cursor::new(
+            conn.clone(),
+            self.pg_config.clone(),
+            querystring,
+            parameters,
+            fetch_number.unwrap_or(10),
+            scroll,
+            prepared,
+        ))
+    }
 
-        Err(RustPSQLDriverError::ConnectionClosedError)
+    #[pyo3(signature = (
+        querystring,
+        parameters=None,
+        fetch_number=None,
+    ))]
+    pub fn portal(
+        &self,
+        querystring: String,
+        parameters: Option<Py<PyAny>>,
+        fetch_number: Option<i32>,
+    ) -> PSQLPyResult<Portal> {
+        println!("{:?}", fetch_number);
+        Ok(Portal::new(
+            self.db_client.clone(),
+            querystring,
+            parameters,
+            fetch_number,
+        ))
     }
 
     #[allow(clippy::needless_pass_by_value)]