Skip to content

Commit 84845c6

Browse files
sfc-gh-fpawlowskisfc-gh-fpawlowski
committed
NO-SNOW: cleaned up code
1 parent 56ec7fa commit 84845c6

File tree

2 files changed

+177
-218
lines changed

2 files changed

+177
-218
lines changed

‎ASYNC_CONNECT_IMPLEMENTATION.md‎

Lines changed: 55 additions & 205 deletions
Original file line numberDiff line numberDiff line change
@@ -2,229 +2,95 @@
22

33
## Overview
44

5-
This document describes the hybrid wrapper pattern used to implement `snowflake.connector.aio.connect()` that preserves metadata from `SnowflakeConnection.__init__` while supporting both simple await and async context manager usage patterns, with full coroutine protocol support.
5+
This document describes the hybrid wrapper pattern used to implement `snowflake.connector.aio.connect()` that preserves metadata from `SnowflakeConnection.__init__` while supporting both simple await and async context manager usage patterns with full coroutine protocol support.
66

77
## Problem Statement
88

9-
The synchronous `snowflake.connector.connect()` uses:
10-
```python
11-
@wraps(SnowflakeConnection.__init__)
12-
def Connect(**kwargs) -> SnowflakeConnection:
13-
return SnowflakeConnection(**kwargs)
14-
```
15-
16-
The async version cannot be decorated with `@wraps` on a raw async function. We needed a solution that:
17-
1. Preserves metadata for IDE introspection and tooling
18-
2. Supports `conn = await aio.connect(...)`
19-
3. Supports `async with aio.connect(...) as conn:`
20-
4. Implements the full coroutine protocol (following aiohttp's pattern)
9+
The async version of `connect()` must:
10+
1. Preserve metadata for IDE introspection and tooling (like the synchronous version)
11+
2. Support `conn = await aio.connect(...)` pattern
12+
3. Support `async with aio.connect(...) as conn:` pattern
13+
4. Implement the full coroutine protocol for ecosystem compatibility (following aiohttp's pattern)
2114

22-
## Implementation
23-
24-
### Architecture
15+
## Architecture
2516

2617
```
2718
connect = _AsyncConnectWrapper()
2819
↓ (calls __call__)
29-
_AsyncConnectContextManager (coroutine wrapper)
30-
├─ Coroutine Protocol: send(), throw(), close()
31-
├─ __await__(), __iter__()
32-
└─ __aenter__(), __aexit__() (async context manager)
20+
_AsyncConnectContextManager (implements HybridCoroutineContextManager protocol)
21+
├─ Coroutine Protocol: send(), throw(), close(), __await__(), __iter__()
22+
└─ Async Context Manager: __aenter__(), __aexit__()
3323
```
3424

35-
### Class: _AsyncConnectContextManager
25+
## Key Components
3626

37-
Makes a coroutine both awaitable and an async context manager, while implementing the full coroutine protocol.
27+
### HybridCoroutineContextManager Protocol
3828

39-
```python
40-
class _AsyncConnectContextManager:
41-
"""Hybrid wrapper that enables both awaiting and async context manager usage.
29+
Combines PEP 492 (coroutine) and PEP 343 (async context manager) protocols. Allows the same object to be managed by external code expecting either interface (e.g., timeout handlers, async schedulers).
4230

43-
Implements the full coroutine protocol for maximum compatibility.
44-
"""
31+
### _AsyncConnectContextManager
4532

46-
__slots__ = ("_coro", "_conn")
33+
Implements `HybridCoroutineContextManager[SnowflakeConnection]`:
34+
- **`send()`, `throw()`, `close()`** - Forward to inner coroutine for ecosystem compatibility
35+
- **`__await__()`** - Enable `await` syntax
36+
- **`__aenter__()` / `__aexit__()`** - Enable async context manager usage
4737

48-
def __init__(self, coro: Coroutine[Any, Any, SnowflakeConnection]) -> None:
49-
self._coro = coro
50-
self._conn: SnowflakeConnection | None = None
38+
### preserve_metadata Decorator
5139

52-
def send(self, arg: Any) -> Any:
53-
"""Send a value into the wrapped coroutine."""
54-
return self._coro.send(arg)
40+
Copies `__wrapped__`, `__doc__`, `__module__`, `__annotations__` from a source class to instances during `__init__`. Allows instances to be introspected like the source class's `__init__`, enabling IDE support and tooling.
5541

56-
def throw(self, *args: Any, **kwargs: Any) -> Any:
57-
"""Throw an exception into the wrapped coroutine."""
58-
return self._coro.throw(*args, **kwargs)
42+
### _AsyncConnectWrapper
5943

60-
def close(self) -> None:
61-
"""Close the wrapped coroutine."""
62-
return self._coro.close()
44+
Callable wrapper decorated with `@preserve_metadata(SnowflakeConnection)` to preserve metadata. Its `__call__` method creates and returns a `_AsyncConnectContextManager` instance.
6345

64-
def __await__(self) -> Generator[Any, None, SnowflakeConnection]:
65-
"""Enable: conn = await connect(...)"""
66-
return self._coro.__await__()
46+
## Design Rationale
6747

68-
def __iter__(self) -> Generator[Any, None, SnowflakeConnection]:
69-
"""Make the wrapper iterable like a coroutine."""
70-
return self.__await__()
48+
### Why Full Coroutine Protocol?
7149

72-
async def __aenter__(self) -> SnowflakeConnection:
73-
"""Enable: async with connect(...) as conn:"""
74-
self._conn = await self._coro
75-
return await self._conn.__aenter__()
76-
77-
async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
78-
"""Exit async context manager."""
79-
if self._conn is not None:
80-
return await self._conn.__aexit__(exc_type, exc, tb)
50+
External async code (timeout handlers, async schedulers, introspection tools) expects `send()`, `throw()`, and `close()` methods. Without these, our wrapper breaks compatibility and may fail at runtime with code like:
51+
```python
52+
result = await asyncio.wait_for(aio.connect(...), timeout=5.0) # May call throw()
8153
```
8254

83-
#### Coroutine Protocol Methods
55+
### Why Preserve Metadata?
8456

85-
- **`send(arg)`**: Send a value into the wrapped coroutine (used for manual coroutine driving)
86-
- **`throw(*args, **kwargs)`**: Throw an exception into the wrapped coroutine
87-
- **`close()`**: Gracefully close the wrapped coroutine
88-
- **`__await__()`**: Return a generator to enable `await` syntax
89-
- **`__iter__()`**: Make the wrapper iterable (some async utilities require this)
57+
- IDEs show correct function signature when hovering over `connect`
58+
- `help(connect)` displays proper docstring
59+
- `inspect.signature(connect)` works correctly
60+
- Static type checkers recognize parameters
9061

91-
### Class: _AsyncConnectWrapper
62+
### Why Both Await and Context Manager?
9263

93-
Callable wrapper that preserves `SnowflakeConnection.__init__` metadata.
94-
95-
```python
96-
class _AsyncConnectWrapper:
97-
"""Preserves SnowflakeConnection.__init__ metadata for async connect function.
98-
99-
This wrapper enables introspection tools and IDEs to see the same signature
100-
as the synchronous snowflake.connector.connect function.
101-
"""
102-
103-
def __init__(self) -> None:
104-
self.__wrapped__ = SnowflakeConnection.__init__
105-
self.__name__ = "connect"
106-
self.__doc__ = SnowflakeConnection.__init__.__doc__
107-
self.__module__ = __name__
108-
self.__qualname__ = "connect"
109-
self.__annotations__ = getattr(SnowflakeConnection.__init__, "__annotations__", {})
110-
111-
@wraps(SnowflakeConnection.__init__)
112-
def __call__(self, **kwargs: Any) -> _AsyncConnectContextManager:
113-
"""Create and connect to a Snowflake connection asynchronously."""
114-
async def _connect_coro() -> SnowflakeConnection:
115-
conn = SnowflakeConnection(**kwargs)
116-
await conn.connect()
117-
return conn
118-
119-
return _AsyncConnectContextManager(_connect_coro())
120-
```
64+
- **Await pattern:** Simple, lightweight for one-off connections
65+
- **Context manager pattern:** Ensures cleanup via `__aexit__`, safer for resource management
12166

12267
## Usage Patterns
12368

124-
### Pattern 1: Simple Await (No Context Manager)
12569
```python
126-
conn = await aio.connect(
127-
account="myaccount",
128-
user="myuser",
129-
password="mypassword"
130-
)
70+
# Pattern 1: Simple await
71+
conn = await aio.connect(account="myaccount", user="user", password="pass")
13172
result = await conn.cursor().execute("SELECT 1")
13273
await conn.close()
133-
```
13474

135-
**Flow:**
136-
1. `aio.connect(...)` returns `_AsyncConnectContextManager`
137-
2. `await` calls `__await__()`, returns inner coroutine result
138-
3. `conn` is a `SnowflakeConnection` object
139-
4. Wrapper is garbage collected
140-
141-
### Pattern 2: Async Context Manager
142-
```python
143-
async with aio.connect(
144-
account="myaccount",
145-
user="myuser",
146-
password="mypassword"
147-
) as conn:
75+
# Pattern 2: Async context manager (recommended)
76+
async with aio.connect(account="myaccount", user="user", password="pass") as conn:
14877
result = await conn.cursor().execute("SELECT 1")
14978
# Auto-closes on exit
15079
```
15180

152-
**Flow:**
153-
1. `aio.connect(...)` returns `_AsyncConnectContextManager`
154-
2. `async with` calls `__aenter__()`, awaits coroutine, returns connection
155-
3. Code block executes
156-
4. `async with` calls `__aexit__()` on exit
81+
## Implementation Details
15782

158-
### Pattern 3: Manual Coroutine Driving (Advanced)
159-
```python
160-
# For advanced use cases with manual iteration
161-
coro_wrapper = aio.connect(account="myaccount", user="user", password="pass")
162-
# Can use send(), throw(), close() methods directly if needed
163-
```
83+
### __slots__ = ("_coro", "_conn")
84+
Memory efficient, especially when many connections are created.
16485

165-
## Key Design Decisions
166-
167-
### 1. Metadata Copying in `__init__`
168-
```python
169-
self.__wrapped__ = SnowflakeConnection.__init__
170-
self.__name__ = "connect"
171-
# ...
172-
```
173-
**Why:** Allows direct attribute access for introspection before calling `__call__`
174-
175-
### 2. `@wraps` Decorator on `__call__`
176-
```python
177-
@wraps(SnowflakeConnection.__init__)
178-
def __call__(self, **kwargs: Any) -> _AsyncConnectContextManager:
179-
```
180-
**Why:** IDEs and inspection tools that examine `__call__` see correct metadata
181-
182-
### 3. Inner Coroutine Function
86+
### Inner Coroutine Function
18387
```python
18488
async def _connect_coro() -> SnowflakeConnection:
18589
conn = SnowflakeConnection(**kwargs)
18690
await conn.connect()
18791
return conn
18892
```
189-
**Why:** Defers connection creation and establishment until await time, not at `connect()` call time
190-
191-
### 4. `__slots__` on Context Manager
192-
```python
193-
__slots__ = ("_coro", "_conn")
194-
```
195-
**Why:** Memory efficient, especially when many connections are created
196-
197-
### 5. Full Coroutine Protocol
198-
Following aiohttp's `_RequestContextManager` pattern, we implement `send()`, `throw()`, and `close()` methods. This ensures:
199-
- Maximum compatibility with async utilities and libraries
200-
- Support for manual coroutine driving if needed
201-
- Proper cleanup and exception handling
202-
203-
## Comparison with aiohttp's _RequestContextManager
204-
205-
Our implementation follows the same proven pattern used by aiohttp:
206-
207-
| Feature | aiohttp | Our Implementation |
208-
|---------|---------|-------------------|
209-
| `send()` method | ✓ | ✓ |
210-
| `throw()` method | ✓ | ✓ |
211-
| `close()` method | ✓ | ✓ |
212-
| `__await__()` method | ✓ | ✓ |
213-
| `__iter__()` method | ✓ | ✓ |
214-
| `__aenter__()` method | ✓ | ✓ |
215-
| `__aexit__()` method | ✓ | ✓ |
216-
| Metadata preservation | N/A | ✓ |
217-
218-
## Behavior Comparison
219-
220-
| Aspect | Pattern 1 | Pattern 2 | Pattern 3 |
221-
|--------|-----------|-----------|-----------|
222-
| Syntax | `conn = await aio.connect(...)` | `async with aio.connect(...) as conn:` | Manual iteration |
223-
| Connection Object | In local scope | In context scope | Via wrapper |
224-
| Cleanup | Manual (`await conn.close()`) | Automatic (`__aexit__`) | Manual |
225-
| Metadata Available | Yes | Yes | Yes |
226-
| Error Handling | Manual try/except | Automatic via context manager | Manual |
227-
| Use Case | Simple connections | Resource management | Advanced/testing |
93+
Defers connection creation and establishment until await time, not at `connect()` call time.
22894

22995
## Verification
23096

@@ -234,36 +100,20 @@ assert connect.__name__ == "connect"
234100
assert hasattr(connect, "__wrapped__")
235101
assert callable(connect)
236102

237-
# Return type validation
103+
# Return type is hybrid awaitable + context manager
238104
result = connect(account="test")
239-
assert hasattr(result, "__await__") # Awaitable
240-
assert hasattr(result, "__aenter__") # Async context manager
241-
assert hasattr(result, "__aexit__") # Async context manager
242-
243-
# Full coroutine protocol
244-
assert hasattr(result, "send") # Coroutine protocol
245-
assert hasattr(result, "throw") # Exception injection
246-
assert hasattr(result, "close") # Cleanup
247-
assert hasattr(result, "__iter__") # Iteration support
105+
assert hasattr(result, "__await__") # Awaitable
106+
assert hasattr(result, "__aenter__") # Async context manager
107+
assert hasattr(result, "send") # Full coroutine protocol
248108
```
249109

250-
## File Location
251-
252-
`src/snowflake/connector/aio/__init__.py`
253-
254110
## Backwards Compatibility
255111

256112
- ✅ Existing code using `await aio.connect(...)` works unchanged
257-
- ✅ New code can use `async with aio.connect(...) as conn:`
258-
- ✅ Metadata available for IDE tooltips and introspection tools
259-
- ✅ Signature matches synchronous version in tooling
260-
- ✅ Full coroutine protocol support for advanced use cases
261-
262-
## Benefits
263-
264-
✅ **Full Coroutine Protocol** - Compatible with all async utilities and libraries
265-
✅ **Flexible Usage** - Simple await or async context manager patterns
266-
✅ **Metadata Preservation** - IDE tooltips and introspection support
267-
✅ **Transparent & Efficient** - Minimal overhead, garbage collected after use
268-
✅ **Backwards Compatible** - No breaking changes to existing code
269-
✅ **Battle-tested Pattern** - Follows aiohttp's proven design
113+
- ✅ Metadata available for IDE tooltips and introspection
114+
- ✅ Full coroutine protocol support for advanced/external tooling
115+
- ✅ No breaking changes
116+
117+
## File Location
118+
119+
`src/snowflake/connector/aio/__init__.py`

0 commit comments

Comments
 (0)