Merge branch 'main' into develop

chandr-andr · chandr-andr · commit 94f80651421a · 2024-08-09T17:11:15.000+02:00
diff --git a/.github/workflows/release_docs.yaml b/.github/workflows/release_docs.yaml
@@ -30,7 +30,7 @@ jobs:
         uses: peaceiris/actions-gh-pages@v3
         with:
           personal_token: ${{ secrets.PERSONAL_TOKEN }}
-          external_repository: qaspen-python/qaspen-python.github.io
+          external_repository: psqlpy-python/psqlpy-python.github.io
           publish_dir: ./docs_dist
           user_name: "github-actions[bot]"
           user_email: "github-actions[bot]@users.noreply.github.com"
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "psqlpy"
-version = "0.7.3"
+version = "0.7.5"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/README.md b/README.md
@@ -8,7 +8,7 @@ Driver for PostgreSQL written fully in Rust and exposed to Python.
 Main goals of the library is speed and type safety.
 
 ## Documentation
-You can find full documentation here - [PSQLPy documentation](https://qaspen-python.github.io/)
+You can find full documentation here - [PSQLPy documentation](https://psqlpy-python.github.io/)
 
 ## Installation
 
diff --git a/docs/.vuepress/config.ts b/docs/.vuepress/config.ts
@@ -12,15 +12,15 @@ export default defineUserConfig({
   bundler: viteBundler(),
 
   theme: hopeTheme({
-    repo: "qaspen-python/psqlpy",
+    repo: "psqlpy-python/psqlpy",
 
     repoLabel: "GitHub",
 
     repoDisplay: true,
 
     sidebar,
 
-    hostname: "https://qaspen-python.github.io/",
+    hostname: "https://psqlpy-python.github.io/",
 
     plugins: {
       readingTime: false,
diff --git a/docs/.vuepress/theme.ts b/docs/.vuepress/theme.ts
@@ -2,13 +2,13 @@ import { hopeTheme } from "vuepress-theme-hope";
 import sidebar from "./sidebar.js";
 
 export default hopeTheme({
-  hostname: "https://github.com/qaspen-python/psqlpy",
+  hostname: "https://github.com/psqlpy-python/psqlpy",
 
   iconAssets: "fontawesome-with-brands",
 
   logo: "./logo.png",
 
-  repo: "qaspen-python/psqlpy",
+  repo: "psqlpy-python/psqlpy",
 
   docsDir: "src",
 
diff --git a/docs/README.md b/docs/README.md
@@ -52,7 +52,7 @@ poetry add psqlpy
 @tab git
 
 ```bash
-pip install git+https://github.com/qaspen-python/psqlpy
+pip install git+https://github.com/psqlpy-python/psqlpy
 ```
 
 :::
diff --git a/docs/components/connection_pool.md b/docs/components/connection_pool.md
@@ -13,8 +13,6 @@ You cannot set the minimum size for the connection pool, by it is 0.
 So, if you set `max_db_pool_size` to 100, pool will create new connection every time there aren't enough connections to handle the load.
 :::
 
-## Connection pool methods
-
 ### All available ConnectionPool parameters
 
 - `dsn`: Full dsn connection string.
@@ -105,9 +103,6 @@ db_pool: Final = ConnectionPool(
     db_name="postgres",
     max_db_pool_size=10,
 )
-
-async def main() -> None:
-
 ```
 
 ### Initialize Connection Pool with DSN
@@ -124,9 +119,6 @@ db_pool: Final = ConnectionPool(
     dsn="postgres://postgres:postgres@localhost:5432/postgres",
     max_db_pool_size=10,
 )
-
-async def main() -> None:
-
 ```
 
 ### Create Connection Pool with one function
@@ -143,6 +135,26 @@ db_pool: Final = connect(
 ```
 `connect` function has the same parameters as `ConnectionPool`.
 
+### Use Connection Pool as context manager
+```py
+from typing import Final
+
+from psqlpy import ConnectionPool
+
+
+async def main() -> None:
+    with ConnectionPool(
+        dsn="postgres://postgres:postgres@localhost:5432/postgres",
+        max_db_pool_size=10,
+    ) as db_pool:
+        # ConnectionPool is opened
+        await db_pool.execute("SOME_SQL")
+        # ConnectionPool is opened
+    # ConnectionPool is closed
+```
+
+## Connection pool methods
+
 ### Resize
 Resize connection pool capacity.
 
diff --git a/docs/contribute.md b/docs/contribute.md
@@ -8,7 +8,7 @@ There are several rules for contributors:
 - Please do not add malware.
 - Please make sure that your request solves the problem.
 
-If you struggle with something or feel frustrated, you either create an issue, create a [discussions](https://github.com/qaspen-python/psqlpy/discussions). page or publish a draft PR and ask your question in the description.
+If you struggle with something or feel frustrated, you either create an issue, create a [discussions](https://github.com/psqlpy-python/psqlpy/discussions). page or publish a draft PR and ask your question in the description.
 
 We have lots of tests in CI. But since CI runs from first-time contributors should be approved, you better test locally. It just takes less time to prepare PR for merging.
 
diff --git a/docs/introduction/lets_start.md b/docs/introduction/lets_start.md
@@ -22,7 +22,7 @@ poetry add psqlpy
 @tab git
 
 ```bash
-pip install git+https://github.com/qaspen-python/psqlpy
+pip install git+https://github.com/psqlpy-python/psqlpy
 ```
 
 :::
diff --git a/docs/usage/frameworks/fastapi.md b/docs/usage/frameworks/fastapi.md
@@ -31,7 +31,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
     )
     app.state.db_pool = db_pool
     yield
-    await db_pool.close()
+    db_pool.close()
 
 
 app = FastAPI(lifespan=lifespan)
diff --git a/docs/usage/frameworks/litestar.md b/docs/usage/frameworks/litestar.md
@@ -35,7 +35,7 @@ def start_db_pool(app: Litestar) -> ConnectionPool:
 
 async def stop_db_pool(app: Litestar) -> None:
     """Close database connection pool."""
-    if getattr(app.state, "engine", None):
+    if getattr(app.state, "db_pool", None):
         db_pool = cast(ConnectionPool, app.state.db_pool)
         db_pool.close()
 
diff --git a/docs/usage/frameworks/robyn.md b/docs/usage/frameworks/robyn.md
@@ -37,7 +37,7 @@ async def pg_pool_example(request: Request) -> list[dict[Any, Any]]:
 
 async def main() -> None:
     try:
-        app.start(host="127.0.0.1", port=8000)
+        await app.start(host="127.0.0.1", port=8000)
     finally:
         db_pool.close()
 
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "psqlpy-docs",
-  "repository": "git@github.com:qaspen-python/psqlpy.git",
+  "repository": "git@github.com:psqlpy-python/psqlpy.git",
   "author": "Kiselev Aleksandr <askiselev00@gmail.com>",
   "license": "MIT",
   "private": false,
diff --git a/psqlpy-stress/poetry.lock b/psqlpy-stress/poetry.lock
diff --git a/psqlpy-stress/pyproject.toml b/psqlpy-stress/pyproject.toml
@@ -21,7 +21,7 @@ pytz = "^2024.1"
 psycopg2-binary = "^2.9.9"
 gunicorn = "^21.2.0"
 uvicorn = "^0.29.0"
-psqlpy-piccolo = { git = "https://github.com/qaspen-python/psqlpy-piccolo" }
+psqlpy-piccolo = { git = "https://github.com/psqlpy-python/psqlpy-piccolo" }
 piccolo = "^1.12.0"
 
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -44,9 +44,9 @@ module-name = "psqlpy._internal"
 features = ["pyo3/extension-module"]
 
 [project.urls]
-homepage = "https://github.com/qaspen-python/psqlpy"
-repository = "https://github.com/qaspen-python/psqlpy"
-documentation = "https://qaspen-python.github.io/"
+homepage = "https://github.com/psqlpy-python/psqlpy"
+repository = "https://github.com/psqlpy-python/psqlpy"
+documentation = "https://psqlpy-python.github.io/"
 
 [tool.isort]
 profile = "black"
diff --git a/python/psqlpy/_internal/__init__.pyi b/python/psqlpy/_internal/__init__.pyi
@@ -1215,6 +1215,14 @@ class ConnectionPool:
         - `ca_file`: Loads trusted root certificates from a file.
             The file should contain a sequence of PEM-formatted CA certificates.
         """
+    def __iter__(self: Self) -> Self: ...
+    def __enter__(self: Self) -> Self: ...
+    def __exit__(
+        self: Self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: types.TracebackType | None,
+    ) -> None: ...
     def status(self: Self) -> ConnectionPoolStatus:
         """Return information about connection pool.
 
diff --git a/python/tests/test_connection_pool.py b/python/tests/test_connection_pool.py
@@ -175,3 +175,15 @@ async def test_close_connection_pool() -> None:
 
     with pytest.raises(expected_exception=RustPSQLDriverPyBaseError):
         await pg_pool.execute("SELECT 1")
+
+
+async def test_connection_pool_as_context_manager() -> None:
+    """Test connection pool as context manager."""
+    with ConnectionPool(
+        dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test",
+    ) as pg_pool:
+        res = await pg_pool.execute("SELECT 1")
+        assert res.result()
+
+    with pytest.raises(expected_exception=RustPSQLDriverPyBaseError):
+        await pg_pool.execute("SELECT 1")
diff --git a/python/tests/test_ssl_mode.py b/python/tests/test_ssl_mode.py
@@ -74,3 +74,25 @@ async def test_ssl_mode_require_pool_builder(
     pool = builder.build()
 
     await pool.execute("SELECT 1")
+
+
+async def test_ssl_mode_require_without_ca_file(
+    postgres_host: str,
+    postgres_user: str,
+    postgres_password: str,
+    postgres_port: int,
+    postgres_dbname: str,
+) -> None:
+    builder = (
+        ConnectionPoolBuilder()
+        .max_pool_size(10)
+        .host(postgres_host)
+        .port(postgres_port)
+        .user(postgres_user)
+        .password(postgres_password)
+        .dbname(postgres_dbname)
+        .ssl_mode(SslMode.Require)
+    )
+    pool = builder.build()
+
+    await pool.execute("SELECT 1")
diff --git a/src/Untitled-1.py b/src/Untitled-1.py
diff --git a/src/driver/common_options.rs b/src/driver/common_options.rs
@@ -64,7 +64,7 @@ impl TargetSessionAttrs {
 }
 
 #[pyclass]
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, PartialEq)]
 pub enum SslMode {
     /// Do not use TLS.
     Disable,
diff --git a/src/driver/connection_pool.rs b/src/driver/connection_pool.rs
@@ -1,8 +1,8 @@
 use crate::runtime::tokio_runtime;
 use deadpool_postgres::{Manager, ManagerConfig, Object, Pool, RecyclingMethod};
-use openssl::ssl::{SslConnector, SslMethod};
+use openssl::ssl::{SslConnector, SslMethod, SslVerifyMode};
 use postgres_openssl::MakeTlsConnector;
-use pyo3::{pyclass, pyfunction, pymethods, PyAny};
+use pyo3::{pyclass, pyfunction, pymethods, Py, PyAny};
 use std::{sync::Arc, vec};
 use tokio_postgres::NoTls;
 
@@ -13,7 +13,7 @@ use crate::{
 };
 
 use super::{
-    common_options::{ConnRecyclingMethod, LoadBalanceHosts, SslMode, TargetSessionAttrs},
+    common_options::{self, ConnRecyclingMethod, LoadBalanceHosts, SslMode, TargetSessionAttrs},
     connection::Connection,
     utils::build_connection_config,
 };
@@ -104,6 +104,15 @@ pub fn connect(
         builder.set_ca_file(ca_file)?;
         let tls_connector = MakeTlsConnector::new(builder.build());
         mgr = Manager::from_config(pg_config, tls_connector, mgr_config);
+    } else if let Some(ssl_mode) = ssl_mode {
+        if ssl_mode == common_options::SslMode::Require {
+            let mut builder = SslConnector::builder(SslMethod::tls())?;
+            builder.set_verify(SslVerifyMode::NONE);
+            let tls_connector = MakeTlsConnector::new(builder.build());
+            mgr = Manager::from_config(pg_config, tls_connector, mgr_config);
+        } else {
+            mgr = Manager::from_config(pg_config, NoTls, mgr_config);
+        }
     } else {
         mgr = Manager::from_config(pg_config, NoTls, mgr_config);
     }
@@ -244,6 +253,28 @@ impl ConnectionPool {
         )
     }
 
+    #[must_use]
+    pub fn __iter__(self_: Py<Self>) -> Py<Self> {
+        self_
+    }
+
+    #[allow(clippy::needless_pass_by_value)]
+    fn __enter__(self_: Py<Self>) -> Py<Self> {
+        self_
+    }
+
+    #[allow(clippy::needless_pass_by_value)]
+    fn __exit__(
+        self_: Py<Self>,
+        _exception_type: Py<PyAny>,
+        _exception: Py<PyAny>,
+        _traceback: Py<PyAny>,
+    ) {
+        pyo3::Python::with_gil(|gil| {
+            self_.borrow(gil).close();
+        });
+    }
+
     #[must_use]
     pub fn status(&self) -> ConnectionPoolStatus {
         let inner_status = self.0.status();
@@ -413,7 +444,7 @@ impl ConnectionPool {
         Ok(Connection::new(Some(Arc::new(db_connection)), None))
     }
 
-    /// Return new single connection.
+    /// Close connection pool.
     ///
     /// # Errors
     /// May return Err Result if cannot get new connection from the pool.
diff --git a/src/driver/connection_pool_builder.rs b/src/driver/connection_pool_builder.rs

Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:`
`31`	`31`	`)`
`32`	`32`	`app.state.db_pool = db_pool`
`33`	`33`	`yield`
`34`		`- await db_pool.close()`
	`34`	`+ db_pool.close()`
`35`	`35`
`36`	`36`
`37`	`37`	`app = FastAPI(lifespan=lifespan)`
Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "psqlpy-docs",`
`3`		`- "repository": "git@github.com:qaspen-python/psqlpy.git",`
	`3`	`+ "repository": "git@github.com:psqlpy-python/psqlpy.git",`
`4`	`4`	`"author": "Kiselev Aleksandr <askiselev00@gmail.com>",`
`5`	`5`	`"license": "MIT",`
`6`	`6`	`"private": false,`
Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,7 @@ impl TargetSessionAttrs {`
`64`	`64`	`}`
`65`	`65`
`66`	`66`	`#[pyclass]`
`67`		`-#[derive(Clone, Copy)]`
	`67`	`+#[derive(Clone, Copy, PartialEq)]`
`68`	`68`	`pub enum SslMode {`
`69`	`69`	`/// Do not use TLS.`
`70`	`70`	`Disable,`