picopyn

 12class Client:
 13    """
 14    Async client for managing connections to a picodata cluster using a connection pool.
 15
 16    This client handles connection pooling, automatic node discovery (if enabled),
 17    and supports load balancing strategies for query distribution.
 18
 19    :param dsn (str): The data source name (e.g., "postgresql://user:pass@host:port") for the cluster.
 20    :param pool_size (int, optional): Maximum number of connections in the pool. Must be at least 1. Default value is 10
 21    :param balance_strategy (callable, optional): A custom strategy function to select a connection
 22        from the pool. If None, round-robin strategy is used.
 23    :param forbidden_tiers (str, optional): A comma-separated list of Picodata node tiers for which connection is forbidden.
 24        (e.g., "arbiter,readonly"). If None, connections to all tiers are allowed.
 25    :param connect_kwargs: Additional keyword arguments passed to each connection.
 26
 27    Example:
 28
 29    Client with random balance strategy
 30    >>> def random_strategy(connections):
 31    ...     import random
 32    ...     return random.choice(connections)
 33    >>> client = Client(
 34    ...     dsn="postgresql://admin:pass@localhost:5432",
 35    ...     balance_strategy=random_strategy
 36    ... )
 37
 38    Client with strategy always choose first connection
 39    >>> def custom_strategy(pool_conns):
 40    ...     return pool_conns[0]
 41    >>> client = Client(
 42    ...     dsn="postgresql://admin:pass@localhost:5432",
 43    ...     balance_strategy=custom_strategy
 44    ... )
 45
 46    Client with multi-host DSN string:
 47    >>> client = Client(
 48    ...     dsn="postgresql://admin:pass@host1:5432,host2:5432"
 49    ... )
 50    """
 51
 52    def __init__(
 53        self,
 54        dsn: str,
 55        pool_size: int | None = None,
 56        balance_strategy: Callable[[list[Connection]], Connection] | None = None,
 57        forbidden_tiers: str | None = None,
 58        **connect_kwargs: Any,
 59    ) -> None:
 60        self._pool = Pool(
 61            dsn=dsn,
 62            max_size=pool_size or 10,
 63            enable_discovery=True,
 64            balance_strategy=balance_strategy,
 65            forbidden_tiers=forbidden_tiers,
 66            **connect_kwargs,
 67        )
 68
 69    async def connect(self) -> None:
 70        """
 71        Prepares the client by connection connection pool.
 72
 73        This should be called before using the client to ensure connections are available.
 74
 75        Example:
 76
 77        >>> client = Client(dsn="postgresql://admin:pass@localhost:5432")
 78        >>> await client.connect()
 79        """
 80        await self._pool.connect()
 81
 82    async def execute(self, query: str, *args: Any) -> str:
 83        """
 84        Executes a query that does not return rows (e.g. INSERT, UPDATE, DELETE).
 85
 86        :param query: The SQL query string.
 87        :param args: Optional parameters for the SQL query.
 88        :return: The result of the query execution.
 89
 90        Example:
 91
 92        >>> ddl = 'CREATE TABLE "warehouse" (id INTEGER NOT NULL, item TEXT NOT NULL, PRIMARY KEY (id)) USING memtx DISTRIBUTED BY (id);'
 93        >>> await client.execute(ddl)
 94        'CREATE TABLE'
 95        >>> dml = 'INSERT INTO "warehouse" VALUES ($1::int, $2::varchar)'
 96        >>> await client.execute(dml, 1, "test")
 97        'INSERT 0 1'
 98        """
 99        return await self._pool.execute(query, *args)
100
101    async def fetch(self, query: str, *args: Any) -> list[asyncpg.Record]:
102        """
103        Executes a query and fetches all resulting rows.
104
105        :param query: The SQL query string.
106        :param args: Optional parameters for the SQL query.
107        :return: A list of rows returned by the query.
108
109        Example:
110
111        >>> dql = 'SELECT * FROM "warehouse";'
112        >>> await client.fetch(dql)
113        [<Record id=1 item='test'>]
114        """
115        return await self._pool.fetch(query, *args)
116
117    async def fetchrow(self, query: str, *args: Any) -> asyncpg.Record | None:
118        """
119        Executes a query and fetches a single row (first row).
120
121        :param query: The SQL query string.
122        :param args: Optional parameters for the SQL query.
123        :return: A single row returned by the query.
124
125        Example:
126
127        >>> dql = 'SELECT * FROM "warehouse";'
128        >>> await client.fetchrow(dql)
129        <Record id=1 item='test'>
130        """
131        return await self._pool.fetchrow(query, *args)
132
133    async def explain(
134        self, query: str, *args: Any, raw: bool = False
135    ) -> ExplainPlan | ExplainRawPlan:
136        """
137        Executes EXPLAIN for a query and returns a structured plan.
138
139        :param query: The SQL query string without EXPLAIN prefix.
140        :param args: Optional parameters for the SQL query.
141        :param raw: If True, uses EXPLAIN (RAW).
142        :return: ExplainPlan for plain EXPLAIN or ExplainRawPlan for RAW mode.
143
144        Example:
145
146        >>> plan = await client.explain("SELECT * FROM warehouse")
147        >>> raw_plan = await client.explain("SELECT * FROM warehouse", raw=True)
148        """
149        return await self._pool.explain(query, *args, raw=raw)
150
151    async def close(self) -> None:
152        """
153        Closes all connections in the pool.
154
155        This should be called during application shutdown to clean up resources.
156        """
157        await self._pool.close()

 17class Connection:
 18    """
 19    A representation of a database session.
 20
 21    :param dsn (str):  The data source name (e.g., "postgresql://user:pass@host:port" or "postgresql://user:pass@host1:port1,host2:port2") for the picodata node.
 22    :param kwargs (Any): Additional keyword arguments to pass to asyncpg.connect() (e.g., ssl).
 23
 24    SSL should be configured using the appropriate parameters in single source by DSN or by providing ssl parameter in kwargs.
 25    """
 26
 27    def __init__(self, dsn: str, **_connect_kwargs: Any) -> None:
 28        if dsn is None:
 29            raise ValueError("dsn can not be None")
 30
 31        validate_ssl_source(dsn, _connect_kwargs)
 32
 33        self.dsn = dsn
 34        self._connect_kwargs = _connect_kwargs
 35        self.conn = None
 36
 37    async def connect(self) -> None:
 38        """
 39        Create new connection to Picodata
 40        """
 41
 42        try:
 43            self.conn = await asyncpg.connect(self.dsn, **self._connect_kwargs)
 44        except Exception as e:
 45            raise RuntimeError(
 46                f"Failed to connect to picodata instance using DSN {self.dsn}: {e}"
 47            ) from e
 48
 49    async def execute(self, *args: Any, **kwargs: Any) -> str:
 50        """
 51        Execute an SQL command
 52        """
 53
 54        if not self.conn:
 55            raise OSError("No active connection. Try to call .connect() before.")
 56
 57        try:
 58            return await self.conn.execute(*args, **kwargs)
 59        except Exception as e:
 60            raise RuntimeError(f"Failed to execute SQL query: {e}. Query: {args}") from e
 61
 62    async def fetchrow(self, *args: Any, **kwargs: Any) -> asyncpg.Record | None:
 63        """
 64        Run a query and return the first row.
 65        """
 66
 67        if not self.conn:
 68            raise OSError("No active connection. Try to call .connect() before")
 69
 70        try:
 71            return await self.conn.fetchrow(*args, **kwargs)
 72        except Exception as e:
 73            raise RuntimeError(
 74                f"Failed to execute SQL query and fetch row: {e}. Query: {args}"
 75            ) from e
 76
 77    async def fetch(self, *args: Any, **kwargs: Any) -> list[asyncpg.Record]:
 78        """
 79        Run a query and return the results as a list.
 80        """
 81
 82        if not self.conn:
 83            raise OSError("No active connection. Try to call .connect() before")
 84
 85        try:
 86            return await self.conn.fetch(*args, **kwargs)
 87        except Exception as e:
 88            raise RuntimeError(
 89                f"Failed to execute SQL query and fetch result: {e}. Query: {args}"
 90            ) from e
 91
 92    async def explain(
 93        self, query: str, *args: Any, raw: bool = False
 94    ) -> ExplainPlan | ExplainRawPlan:
 95        """
 96        Run EXPLAIN for a query and return a structured plan.
 97
 98        :param query: The SQL query string without EXPLAIN prefix.
 99        :param args: Optional parameters for the SQL query.
100        :param raw: If True, uses EXPLAIN (RAW).
101        :return: ExplainPlan for plain EXPLAIN or ExplainRawPlan for RAW mode.
102        """
103
104        explain_query = build_explain_query(query=query, raw=raw)
105        try:
106            rows = await self.fetch(explain_query, *args)
107            lines = rows_to_lines(rows)
108            if raw:
109                return parse_raw_explain_lines(lines)
110            return parse_plain_explain_lines(lines)
111        except ValueError as e:
112            raise RuntimeError(
113                f"Failed to parse EXPLAIN output: {e}. Query: {explain_query}"
114            ) from e
115
116    async def close(self, *args: Any, **kwargs: Any) -> None:
117        """
118        Close the connection gracefully.
119        """
120        if self.conn:
121            try:
122                return await self.conn.close(*args, **kwargs)
123            except Exception as e:
124                raise RuntimeError(
125                    f"Failed to disconnect from picodata instance {self.dsn}: {e}"
126                ) from e

25@dataclass(slots=True)
26class ExplainPlan:
27    root: PlanNode
28    execution_options: dict[str, int]
29    buckets: str

47@dataclass(slots=True)
48class ExplainRawPlan:
49    sections: list[RawSection]

123class Pool:
124    """A connection pool.
125
126    Connection pool can be used to manage a set of connections to the database.
127    Connections are first acquired from the pool, then used, and then released
128    back to the pool
129
130    :param dsn (str):  The data source name (e.g., "postgresql://user:pass@host:port") for the cluster.
131    :param balance_strategy (callable, optional): A custom strategy function to select a connection
132        from the pool. If None, round-robin strategy is used.
133    :param max_size (int, optional): Maximum number of connections in the pool. Must be at least 1. Default value is 10
134    :param enable_discovery (bool, optional): If True, the pool will automatically discover available
135        picodata instances. If False, only the given `dsn` will be used.
136    :param forbidden_tiers (str, optional): A comma-separated list of Picodata node tiers for which connection is forbidden.
137        (e.g., "arbiter,readonly"). If None, connections to all tiers are allowed.
138    :param connect_kwargs (Any): Additional keyword arguments to pass to asyncpg.connect() (e.g., ssl).
139
140    Example:
141    >>> def random_strategy(connections):
142    ...     import random
143    ...     return random.choice(connections)
144
145    >>> pool = Pool(
146    ...     dsn="postgresql://admin:pass@localhost:5432",
147    ...     balance_strategy=random_strategy,
148    ...     max_size=5
149    ... )
150
151    Also you can specify multiple hosts in the DSN string, separated by commas.
152    Allows connecting to multiple nodes.
153    If the enable_discovery flag is enabled, the driver can also fall back to a fallback node.
154    The nodes are iterated in a round-robin fashion until a connection to one of them succeeds,
155    allowing the driver to perform cluster discovery.
156
157    Example:
158    >>> pool = Pool(
159    ...     dsn="postgresql://admin:pass@host1:5432,host2:5432",
160    ...     max_size=10
161    ... )
162
163    When discovery is enabled, connection pool addresses are taken as-is from
164    the `_pico_peer_address` system table; DSN hosts are used only for initial
165    topology discovery. Address configuration (e.g. instance.pg.listen) must be
166    handled on the cluster side: https://docs.picodata.io/picodata/25.5/reference/config/
167
168    SSL example (kwargs-based):
169    >>> import ssl
170    >>> ctx = ssl.create_default_context(cafile="/path/to/ca.crt")
171    >>> ctx.load_cert_chain(certfile="/path/to/client.crt", keyfile="/path/to/client.key")
172    >>> pool = Pool(
173    ...     dsn="postgresql://admin:pass@host1:5432,host2:5432",
174    ...     ssl=ctx,
175    ... )
176
177    SSL example (DSN query params):
178    >>> pool = Pool(
179    ...     dsn=(
180    ...         "postgresql://admin:pass@host1:5432,host2:5432/db"
181    ...         "?sslmode=verify-ca"
182    ...         "&sslrootcert=/path/to/ca.crt"
183    ...         "&sslcert=/path/to/client.crt"
184    ...         "&sslkey=/path/to/client.key"
185    ...     ),
186    ... )
187    """
188
189    def __init__(
190        self,
191        dsn: str,
192        max_size: int = 10,
193        enable_discovery: bool = False,
194        balance_strategy: Callable[[list[Connection]], Connection] | None = None,
195        forbidden_tiers: str | None = None,
196        **connect_kwargs: Any,
197    ) -> None:
198        if max_size < 1:
199            raise ValueError("max_size must be at least 1")
200
201        self._dsns: list[str] = _parse_multihost_dsn(dsn)
202        self._raw_dsn = dsn
203        self._connect_kwargs = connect_kwargs
204        self._max_size = max_size
205        self._pool: deque[Connection] = deque()
206        self._used: set[Connection] = set()
207        self._forbidden_tiers = (
208            set(t.strip() for t in forbidden_tiers.split(",") if t.strip())
209            if forbidden_tiers
210            else set()
211        )
212        self._lock: asyncio.Lock = asyncio.Lock()
213        self._default_acquire_timeout_sec = 5
214        # node discovery mode
215        # if disabled, pool will be filled with given address connections
216        # if enabled, pool will be filled with available picodata instances
217        self._enable_discovery = enable_discovery
218        # load balancing strategy:
219        # if None, a simple round-robin strategy will be used.
220        # otherwise, the provided callable will be used to select connections.
221        if balance_strategy is not None and not callable(balance_strategy):
222            raise ValueError("balance_strategy must be callable or None")
223        self._balance_strategy = balance_strategy
224
225    async def connect(self) -> None:
226        """
227        Prepares the pool by opening up to `max_size` connections.
228
229        This should be called before using the pool to ensure connections are available.
230        """
231        async with self._lock:
232            if len(self._pool) == self._max_size:
233                return
234
235            # if node discovery is enabled, then connect to all alive picodata instances
236            # (if they fit within the max_size limit)
237            if self._enable_discovery:
238                await self._fill_pool_from_discovery()
239            else:
240                await self._fill_pool_from_bootstrap_dsns()
241
242            if len(self._pool) < self._max_size:
243                raise RuntimeError(
244                    f"Failed to initialize connection pool: only {len(self._pool)} "
245                    f"out of {self._max_size} connections established for DSN {self._raw_dsn}"
246                )
247
248            # rotate the pool to randomize the order of connections.
249            # this helps to distribute the initial load more evenly across nodes
250            # when using round-robin or when multiple clients start simultaneously.
251            shift = random.randint(0, len(self._pool) - 1)
252            self._pool.rotate(shift)
253
254            return
255
256    async def _fill_pool_from_discovery(self) -> None:
257        try:
258            instance_addrs = await self._discover_instances()
259        except Exception as e:
260            raise RuntimeError(
261                f"Failed to discover instances using DSN {self._raw_dsn}: {e}"
262            ) from e
263
264        addr_index = 0
265        # fill the connection pool with connections to all available nodes, up to the max_size.
266        # this ensures the pool is evenly populated across all nodes.
267        # if a node fails to connect, it will be skipped and removed from the list.
268        # the loop will exit early if no nodes remain to avoid an infinite loop.
269        while len(self._pool) < self._max_size and instance_addrs:
270            address = instance_addrs[addr_index % len(instance_addrs)]
271            dsn = _replace_dsn_host(self._dsns[0], address)
272
273            try:
274                conn = Connection(dsn, **self._connect_kwargs)
275                await conn.connect()
276                self._pool.append(conn)
277            except Exception as e:
278                print(f"Could not connect to node {address} for pool: {e}")
279                instance_addrs.remove(address)
280                if not instance_addrs:
281                    break
282                continue
283
284            addr_index += 1
285
286    async def _fill_pool_from_bootstrap_dsns(self) -> None:
287        """
288        Fill the pool with nodes from DSNs if they pass filters.
289
290        Current filter is `forbidden_tiers`.
291        """
292        available = list(self._dsns)
293        idx = 0
294        while len(self._pool) < self._max_size and available:
295            candidate = available[idx % len(available)]
296            try:
297                conn = Connection(candidate, **self._connect_kwargs)
298                await conn.connect()
299
300                if self._forbidden_tiers:
301                    tier = await self._fetch_tier(conn)
302                    if tier and tier in self._forbidden_tiers:
303                        await conn.close()
304                        available.remove(candidate)
305                        continue
306                self._pool.append(conn)
307                idx += 1
308            except Exception:
309                available.remove(candidate)
310
311    async def _fetch_tier(self, conn: Connection) -> str | None:
312        row = await conn.fetchrow(
313            """
314            WITH my_uuid AS (SELECT instance_uuid() AS uuid)
315            SELECT i.tier
316            FROM _pico_instance i
317            JOIN my_uuid u ON i.uuid = u.uuid
318            """
319        )
320        return row["tier"] if row else None
321
322    async def _discover_instances(self) -> list[str]:
323        # try each bootstrap DSN until one succeeds
324        temp_conn: Connection | None = None
325        last_error: Exception | None = None
326        for dsn in self._dsns:
327            try:
328                candidate = Connection(dsn, **self._connect_kwargs)
329                await candidate.connect()
330                temp_conn = candidate
331                break
332            except Exception as e:
333                last_error = e
334
335        if temp_conn is None:
336            raise RuntimeError(
337                f"Could not connect to any bootstrap node {self._raw_dsn}: {last_error}"
338            ) from last_error
339
340        try:
341            alive_instances_info = await self._fetch_discovery_rows(temp_conn)
342            online_addresses = self._extract_online_addresses(alive_instances_info)
343
344            if not online_addresses:
345                if self._forbidden_tiers:
346                    raise ValueError(
347                        f"No online nodes available after applying forbidden_tiers filter: {self._forbidden_tiers}"
348                    )
349                raise ValueError("No online nodes discovered")
350
351            return online_addresses
352        finally:
353            await temp_conn.close()
354
355    async def _fetch_discovery_rows(self, conn: Connection) -> list[asyncpg.Record]:
356        # all instance addresses excluding forbidden tiers
357        if self._forbidden_tiers:
358            placeholders = ", ".join(f"${i + 1}" for i in range(len(self._forbidden_tiers)))
359            query = f"""
360            SELECT i.current_state, p.address
361            FROM _pico_instance i
362            JOIN _pico_peer_address p ON i.raft_id = p.raft_id
363            WHERE p.connection_type = 'pgproto'
364            AND i.tier NOT IN ({placeholders});
365            """
366            return await conn.fetch(query, *list(self._forbidden_tiers))
367
368        # all instance addresses
369        return await conn.fetch(
370            """
371            SELECT i.current_state, p.address
372            FROM _pico_instance i
373            JOIN _pico_peer_address p ON i.raft_id = p.raft_id
374            WHERE p.connection_type = 'pgproto';
375            """
376        )
377
378    def _extract_online_addresses(self, alive_instances_info: list[asyncpg.Record]) -> list[str]:
379        online_addresses: list[str] = []
380        for r in alive_instances_info:
381            if not r.get("current_state"):
382                continue
383
384            try:
385                current_state = json.loads(r["current_state"])
386            except json.JSONDecodeError:
387                print(
388                    f"Failed to decode current state of picodata instance {r.get('current_state')}"
389                )
390                continue
391
392            if "Online" in current_state:
393                online_addresses.append(r["address"])
394
395        return online_addresses
396
397    async def _acquire_raw(self, timeout: float | None = None) -> Connection:
398        """
399        Acquire a connection from the pool.
400
401        If no connections are available, this method will wait until one is released.
402        
403        :param timeout: Maximum time to wait for a connection if the pool is exhausted. If None, a default timeout is used.
404
405        :return: A database connection.
406        """
407        start_time = time.monotonic()
408        effective_timeout = timeout if timeout is not None else self._default_acquire_timeout_sec
409
410        while True:
411            async with self._lock:
412                # сheck if there are any available connections in the pool
413                if self._pool:
414                    # round-robin strategy
415                    if self._balance_strategy is None:
416                        conn = self._pool.popleft()
417                    # custom strategy
418                    else:
419                        try:
420                            conn = self._balance_strategy(list(self._pool))
421                        except Exception as e:
422                            raise RuntimeError(f"balance_strategy raised an exception: {e}") from e
423
424                        if conn not in self._pool:
425                            raise RuntimeError("balance_strategy returned a connection not in pool")
426                        self._pool.remove(conn)
427
428                    # mark it as currently in use
429                    self._used.add(conn)
430                    return conn
431
432            if (time.monotonic() - start_time) >= effective_timeout:
433                raise TimeoutError("Timed out waiting for a free connection in the pool")
434
435            # if no connections are available, wait briefly before retrying
436            # this gives other coroutines (like `release`) a chance to return a connection to the pool
437            await asyncio.sleep(0.1)
438
439    def acquire(self, timeout: float | None = None) -> _PoolAcquireContext:
440        """
441        Acquire a connection from the pool.
442
443        If no connections are available, this method will wait until one is released.
444
445        :param timeout: Maximum time to wait for a connection if the pool is exhausted. If None, a default timeout is used.
446        :return: _PoolAcquireContext: a context manager that can be used with "async with" or "await" to acquire a connection.
447
448        Example (context manager, recommended, safely returns the connection):
449        >>> async with pool.acquire() as conn:
450        ...     await conn.execute("UPDATE ...")
451
452        Example (explicit acquisition, e.g., for compatibility with prepared statements):
453        >>> ctx = pool.acquire()
454        ... conn = await ctx
455        ... try:
456        ...     await conn.execute(...)
457        ... finally:
458        ...     await pool.release(conn)
459
460        Important: if you use acquire/release manually, prepared statements may be
461        dropped when the connection is returned to the pool. To preserve prepared
462        statements, use the connection within a single context block.
463
464        """
465        return _PoolAcquireContext(self, timeout)
466
467    async def release(self, conn: Connection) -> None:
468        """
469        Release a previously acquired connection back to the pool.
470
471        :param conn: The connection to release.
472        """
473        async with self._lock:
474            if conn in self._used:
475                self._used.remove(conn)
476                self._pool.append(conn)
477
478    async def close(self) -> None:
479        """
480        Closes all connections in the pool.
481
482        This should be called during application shutdown to clean up resources.
483        """
484        async with self._lock:
485            while self._pool:
486                conn = self._pool.popleft()
487                await conn.close()
488            for conn in self._used:
489                await conn.close()
490            self._used.clear()
491
492    async def execute(self, query: str, *args: Any) -> str:
493        """
494        Executes a query that does not return rows (e.g. INSERT, UPDATE, DELETE).
495
496        :param query: The SQL query string.
497        :param args: Optional parameters for the SQL query.
498        :return: The result of the query execution.
499        """
500        async with self.acquire() as conn:
501            return await conn.execute(query, *args)
502
503    async def fetch(self, query: str, *args: Any) -> list[asyncpg.Record]:
504        """
505        Executes a query and fetches all resulting rows.
506
507        :param query: The SQL query string.
508        :param args: Optional parameters for the SQL query.
509        :return: A list of rows returned by the query.
510        """
511        async with self.acquire() as conn:
512            return await conn.fetch(query, *args)
513
514    async def fetchrow(self, query: str, *args: Any) -> asyncpg.Record | None:
515        """
516        Executes a query and fetches a single row (first row).
517
518        :param query: The SQL query string.
519        :param args: Optional parameters for the SQL query.
520        :return: A single row returned by the query.
521        """
522        async with self.acquire() as conn:
523            return await conn.fetchrow(query, *args)
524
525    async def explain(
526        self, query: str, *args: Any, raw: bool = False
527    ) -> ExplainPlan | ExplainRawPlan:
528        """
529        Executes EXPLAIN for a query and returns a structured plan.
530
531        :param query: The SQL query string without EXPLAIN prefix.
532        :param args: Optional parameters for the SQL query.
533        :param raw: If True, uses EXPLAIN (RAW).
534        :return: ExplainPlan for plain EXPLAIN or ExplainRawPlan for RAW mode.
535        """
536        async with self.acquire() as conn:
537            return await conn.explain(query, *args, raw=raw)