picopyn

11class Client:
12    """
13    Async client for managing connections to a picodata cluster using a connection pool.
14
15    This client handles connection pooling, automatic node discovery (if enabled),
16    and supports load balancing strategies for query distribution.
17
18    :param dsn (str): The data source name (e.g., "postgresql://user:pass@host:port") for the cluster.
19    :param balance_strategy (callable, optional): A custom strategy function to select a connection
20        from the pool. If None, round-robin strategy is used.
21    :param connect_kwargs: Additional keyword arguments passed to each connection.
22
23    Example:
24        >>> def random_strategy(connections):
25        ...     import random
26        ...     return random.choice(connections)
27
28        >>> client = Client(
29        ...     dsn="postgresql://admin:pass@localhost:5432",
30        ...     balance_strategy=random_strategy
31        ... )
32    """
33
34    def __init__(
35        self,
36        dsn: str,
37        pool_size: int | None = None,
38        balance_strategy: Callable[[list[Connection]], Connection] | None = None,
39        **connect_kwargs: Any,
40    ) -> None:
41        self._pool = Pool(
42            dsn=dsn,
43            max_size=pool_size or 10,
44            enable_discovery=True,
45            balance_strategy=balance_strategy,
46            **connect_kwargs,
47        )
48
49    async def connect(self) -> None:
50        """
51        Prepares the client by connection connection pool.
52
53        This should be called before using the client to ensure connections are available.
54        """
55        await self._pool.connect()
56
57    async def execute(self, query: str, *args: Any) -> str:
58        """
59        Executes a query that does not return rows (e.g. INSERT, UPDATE, DELETE).
60
61        :param query: The SQL query string.
62        :param args: Optional parameters for the SQL query.
63        :return: The result of the query execution.
64        """
65        return await self._pool.execute(query, *args)
66
67    async def fetch(self, query: str, *args: Any) -> list[asyncpg.Record]:
68        """
69        Executes a query and fetches all resulting rows.
70
71        :param query: The SQL query string.
72        :param args: Optional parameters for the SQL query.
73        :return: A list of rows returned by the query.
74        """
75        return await self._pool.fetch(query, *args)
76
77    async def fetchrow(self, query: str, *args: Any) -> asyncpg.Record | None:
78        """
79        Executes a query and fetches a single row (first row).
80
81        :param query: The SQL query string.
82        :param args: Optional parameters for the SQL query.
83        :return: A single row returned by the query.
84        """
85        return await self._pool.fetchrow(query, *args)
86
87    async def close(self) -> None:
88        """
89        Closes all connections in the pool.
90
91        This should be called during application shutdown to clean up resources.
92        """
93        await self._pool.close()

 7class Connection:
 8    """
 9    A representation of a database session.
10
11    :param dsn (str):  The data source name (e.g., "postgresql://user:pass@host:port") for the picodata node.
12    """
13
14    def __init__(self, dsn: str) -> None:
15        if dsn is None:
16            raise ValueError("dsn can not be None")
17        self.dsn = dsn
18        self.conn = None
19
20    async def connect(self) -> None:
21        """
22        Create new connection to Picodata
23        """
24
25        try:
26            self.conn = await asyncpg.connect(self.dsn)
27        except Exception as e:
28            raise RuntimeError(
29                f"Failed to connect to picodata instance using DSN {self.dsn}: {e}"
30            ) from e
31
32    async def execute(self, *args: Any, **kwargs: Any) -> str:
33        """
34        Execute an SQL command
35        """
36
37        if not self.conn:
38            raise OSError("No active connection. Try to call .connect() before.")
39
40        try:
41            return await self.conn.execute(*args, **kwargs)
42        except Exception as e:
43            raise RuntimeError(f"Failed to execute SQL query: {e}. Query: {args}") from e
44
45    async def fetchrow(self, *args: Any, **kwargs: Any) -> asyncpg.Record | None:
46        """
47        Run a query and return the first row.
48        """
49
50        if not self.conn:
51            raise OSError("No active connection. Try to call .connect() before")
52
53        try:
54            return await self.conn.fetchrow(*args, **kwargs)
55        except Exception as e:
56            raise RuntimeError(
57                f"Failed to execute SQL query and fetch row: {e}. Query: {args}"
58            ) from e
59
60    async def fetch(self, *args: Any, **kwargs: Any) -> list[asyncpg.Record]:
61        """
62        Run a query and return the results as a list.
63        """
64
65        if not self.conn:
66            raise OSError("No active connection. Try to call .connect() before")
67
68        try:
69            return await self.conn.fetch(*args, **kwargs)
70        except Exception as e:
71            raise RuntimeError(
72                f"Failed to execute SQL query and fetch result: {e}. Query: {args}"
73            ) from e
74
75    async def close(self, *args: Any, **kwargs: Any) -> None:
76        """
77        Close the connection gracefully.
78        """
79        if self.conn:
80            try:
81                return await self.conn.close(*args, **kwargs)
82            except Exception as e:
83                raise RuntimeError(
84                    f"Failed to disconnect from picodata instance {self.dsn}: {e}"
85                ) from e

 16class Pool:
 17    """A connection pool.
 18
 19    Connection pool can be used to manage a set of connections to the database.
 20    Connections are first acquired from the pool, then used, and then released
 21    back to the pool
 22
 23    :param dsn (str):  The data source name (e.g., "postgresql://user:pass@host:port") for the cluster.
 24    :param balance_strategy (callable, optional): A custom strategy function to select a connection
 25        from the pool. If None, round-robin strategy is used.
 26    :param max_size (int): Maximum number of connections in the pool. Must be at least 1.
 27    :param enable_discovery (bool): If True, the pool will automatically discover available
 28        picodata instances. If False, only the given `dsn` will be used.
 29    :param balance_strategy (callable, optional): A function that selects a connection from the pool.
 30        If None, a default round-robin strategy will be used.
 31    """
 32
 33    def __init__(
 34        self,
 35        dsn: str,
 36        max_size: int = 10,
 37        enable_discovery: bool = False,
 38        balance_strategy: Callable[[list[Connection]], Connection] | None = None,
 39        **connect_kwargs: Any,
 40    ) -> None:
 41        if max_size < 1:
 42            raise ValueError("max_size must be at least 1")
 43
 44        self._dsn = dsn
 45        self._connect_kwargs = connect_kwargs
 46        self._max_size = max_size
 47        self._pool: deque[Connection] = deque()
 48        self._used: set[Connection] = set()
 49        self._lock: asyncio.Lock = asyncio.Lock()
 50        self._default_acquire_timeout_sec = 5
 51        # node discovery mode
 52        # if disabled, pool will be filled with given address connections
 53        # if enabled, pool will be filled with available picodata instances
 54        self.enable_discovery = enable_discovery
 55        # load balancing strategy:
 56        # if None, a simple round-robin strategy will be used.
 57        # otherwise, the provided callable will be used to select connections.
 58        if balance_strategy is not None and not callable(balance_strategy):
 59            raise ValueError("balance_strategy must be callable or None")
 60        self._balance_strategy = balance_strategy
 61
 62    async def connect(self) -> None:
 63        """
 64        Prepares the pool by opening up to `max_size` connections.
 65
 66        This should be called before using the pool to ensure connections are available.
 67        """
 68        async with self._lock:
 69            if len(self._pool) == self._max_size:
 70                return
 71
 72            # if node discovery is enabled, then connect to all alive picodata instances
 73            # (if they fit within the max_size limit)
 74            if self.enable_discovery:
 75                try:
 76                    instance_addrs = await self._discover_instances()
 77                except Exception as e:
 78                    raise RuntimeError(
 79                        f"Failed to discover instances using DSN {self._dsn}: {e}"
 80                    ) from e
 81
 82                parsed_url = urlparse(self._dsn)
 83
 84                addr_index = 0
 85                # fill the connection pool with connections to all available nodes, up to the max_size.
 86                # this ensures the pool is evenly populated across all nodes.
 87                # if a node fails to connect, it will be skipped and removed from the list.
 88                # the loop will exit early if no nodes remain to avoid an infinite loop.
 89                while len(self._pool) < self._max_size and instance_addrs:
 90                    address = instance_addrs[addr_index % len(instance_addrs)]
 91                    dsn = f"{parsed_url.scheme}://{parsed_url.username}:{parsed_url.password}@{address}"
 92
 93                    try:
 94                        conn = Connection(dsn, **self._connect_kwargs)
 95                        await conn.connect()
 96                        self._pool.append(conn)
 97                    except Exception as e:
 98                        print(f"Could not connect to node {address} for pool: {e}")
 99                        instance_addrs.remove(address)
100                        if not instance_addrs:
101                            break
102                        continue
103
104                    addr_index += 1
105
106            # then fill the connection pool up to max_size with main mode connections
107            while len(self._pool) < self._max_size:
108                main_node_conn = Connection(self._dsn, **self._connect_kwargs)
109                try:
110                    await main_node_conn.connect()
111                    self._pool.append(main_node_conn)
112                except Exception as e:
113                    raise RuntimeError(
114                        f"Could not connect to main node {self._dsn} for pool: {e}"
115                    ) from e
116
117            # rotate the pool to randomize the order of connections.
118            # this helps to distribute the initial load more evenly across nodes
119            # when using round-robin or when multiple clients start simultaneously.
120            shift = random.randint(0, len(self._pool) - 1)
121            self._pool.rotate(shift)
122
123            return
124
125    async def _discover_instances(self) -> list[str]:
126        # make temporary connection
127        temp_conn = Connection(self._dsn, **self._connect_kwargs)
128
129        try:
130            await temp_conn.connect()
131
132            # all instance addresses excluding connected node
133            alive_instances_info = await temp_conn.fetch(
134                """
135                WITH my_uuid AS (SELECT instance_uuid() AS uuid)
136                SELECT i.name, i.raft_id, i.current_state, p.address
137                FROM _pico_instance i
138                JOIN _pico_peer_address p ON i.raft_id = p.raft_id
139                JOIN my_uuid u ON 1 = 1
140                WHERE p.connection_type = 'pgproto' AND i.uuid != u.uuid;
141            """
142            )
143
144            online_addresses = []
145            # place connected node as first node to be sure that
146            # it will be in the pool independ on pool size
147            parsed_url = urlparse(self._dsn)
148            online_addresses.append(f"{parsed_url.hostname}:{parsed_url.port}")
149            for r in alive_instances_info:
150                if not r.get("current_state", None):
151                    continue
152
153                try:
154                    current_state = json.loads(r.get("current_state", None))
155                except json.JSONDecodeError:
156                    print(
157                        f"Failed to decode current state of picodata instance {r.get('current_state', None)}"
158                    )
159                    continue
160
161                if "Online" in current_state:
162                    online_addresses.append(r["address"])
163
164            return online_addresses
165        finally:
166            await temp_conn.close()
167
168    async def acquire(self, timeout: float | None = None) -> Connection:
169        """
170        Acquire a connection from the pool.
171
172        If no connections are available, this method will wait until one is released.
173
174        :return: A database connection.
175        """
176        start_time = time.monotonic()
177        effective_timeout = timeout if timeout is not None else self._default_acquire_timeout_sec
178
179        while True:
180            async with self._lock:
181                # сheck if there are any available connections in the pool
182                if self._pool:
183                    # round-robin strategy
184                    if self._balance_strategy is None:
185                        conn = self._pool.popleft()
186                    # custom strategy
187                    else:
188                        try:
189                            conn = self._balance_strategy(list(self._pool))
190                        except Exception as e:
191                            raise RuntimeError(f"balance_strategy raised an exception: {e}") from e
192
193                        if conn not in self._pool:
194                            raise RuntimeError("balance_strategy returned a connection not in pool")
195                        self._pool.remove(conn)
196
197                    # mark it as currently in use
198                    self._used.add(conn)
199                    return conn
200
201            if (time.monotonic() - start_time) >= effective_timeout:
202                raise TimeoutError("Timed out waiting for a free connection in the pool")
203
204            # if no connections are available, wait briefly before retrying
205            # this gives other coroutines (like `release`) a chance to return a connection to the pool
206            await asyncio.sleep(0.1)
207
208    async def release(self, conn: Connection) -> None:
209        """
210        Release a previously acquired connection back to the pool.
211
212        :param conn: The connection to release.
213        """
214        async with self._lock:
215            if conn in self._used:
216                self._used.remove(conn)
217                self._pool.append(conn)
218
219    async def close(self) -> None:
220        """
221        Closes all connections in the pool.
222
223        This should be called during application shutdown to clean up resources.
224        """
225        async with self._lock:
226            while self._pool:
227                conn = self._pool.popleft()
228                await conn.close()
229            for conn in self._used:
230                await conn.close()
231            self._used.clear()
232
233    async def execute(self, query: str, *args: Any) -> str:
234        """
235        Executes a query that does not return rows (e.g. INSERT, UPDATE, DELETE).
236
237        :param query: The SQL query string.
238        :param args: Optional parameters for the SQL query.
239        :return: The result of the query execution.
240        """
241        conn = await self.acquire()
242        try:
243            return await conn.execute(query, *args)
244        finally:
245            await self.release(conn)
246
247    async def fetch(self, query: str, *args: Any) -> list[asyncpg.Record]:
248        """
249        Executes a query and fetches all resulting rows.
250
251        :param query: The SQL query string.
252        :param args: Optional parameters for the SQL query.
253        :return: A list of rows returned by the query.
254        """
255        conn = await self.acquire()
256        try:
257            return await conn.fetch(query, *args)
258        finally:
259            await self.release(conn)
260
261    async def fetchrow(self, query: str, *args: Any) -> asyncpg.Record | None:
262        """
263        Executes a query and fetches a single row (first row).
264
265        :param query: The SQL query string.
266        :param args: Optional parameters for the SQL query.
267        :return: A single row returned by the query.
268        """
269        conn = await self.acquire()
270        try:
271            return await conn.fetchrow(query, *args)
272        finally:
273            await self.release(conn)