feat: asyncio support (async wrapper, plugins, and SQLAlchemy dialects)#1257
Open
AhmadMasry wants to merge 56 commits into
Open
feat: asyncio support (async wrapper, plugins, and SQLAlchemy dialects)#1257AhmadMasry wants to merge 56 commits into
AhmadMasry wants to merge 56 commits into
Conversation
AhmadMasry
added a commit
to AhmadMasry/aws-advanced-python-wrapper
that referenced
this pull request
Jul 6, 2026
Open
2 tasks
….aio Async counterparts of the sync plugins and services (failover v2, RWS, EFM v2, IAM/Secrets/federated/Okta auth, tracker, topology monitor, custom endpoint, stale DNS, blue/green, limitless, fastest response), backed by psycopg async and aiomysql, plus async SQLAlchemy dialects and async-additive shared-source changes (accepted_strategies view, async failover rewrap mixin, get_async_dialect_cls hook, aiomysql/pymysql exception hardening).
…odes and version phrasing in docs
Global Aurora: parse the real 4-column topology-query shape (server_id, is_writer, lag, aws_region) and the documented region:host:port instance template format (both previously masked by synthetic test fixtures). MultiAz: honor the dialect's writer-host column index (MySQL SHOW REPLICA STATUS column 39). Base Aurora provider: return empty topology when no writer is present, matching sync. Wire panic-mode writer discovery (build_probe_host) into the topology monitor -- it was fully implemented and tested but never armed in production.
The recheck (rws_recheck_reader_role property, role re-probe on freshly opened reader connections, ReaderRoleMismatch message key) has no sync counterpart -- sync RWS selects by strategy and connects without re-verifying the live role. It originated in the SQLAlchemy-pool work that is out of scope (SA+RWS unsupported). Async RWS now mirrors sync standalone semantics; the failover-eviction mixin is retained as it is a faithful port of sync behavior.
…c plugin service Adds the missing AsyncPluginManager.notify_host_list_changed fan-out and ports sync's topology diffing into refresh/force_refresh_host_list (availability re-hydration, HostEvent computation, plugin notification) -- previously async plugins defined the hook but nothing ever dispatched it. set_current_connection now mirrors sync's switch lifecycle: session-state begin/complete bracket, ROLLBACK_ON_SWITCH for an in-transaction old connection, OldConnectionSuggestedAction aggregation across plugins, and pristine-state restore + close of the old connection unless a plugin votes PRESERVE (async RWS now votes PRESERVE during its swaps, protecting its cached reader/writer connections; previously old connections leaked).
…ilover Reader failover now data-plane-verifies each candidate's role and binds the verified role (sync failover_v2 parity): in STRICT_READER a topology-labeled reader that probes as the promoted writer is rejected, and the original writer is attempted as a demotion candidate with is_original_writer_still_writer bookkeeping. The failed host is marked UNAVAILABLE before target selection, and writer failover filters the new writer through the allow/block list (custom endpoints).
Factory weight table now matches sync PLUGIN_FACTORY_WEIGHTS exactly (srw 310, federated 900, okta 1000) and the connect-time/execute-time/ developer plugins use sync's WEIGHT_RELATIVE_TO_PRIOR_PLUGIN semantics -- they sort right after whatever the user placed before them instead of at fixed absolute positions. always_use_pipeline methods now include only SUBSCRIBED plugins in the chain (sync _make_pipeline parity).
set_read_only on a closed connection now raises SetReadOnlyOnClosedConnection; connect() validates the configured reader-selection strategy and fail-fasts on an unverifiable initial host role (sync read_write_splitting_plugin.py:188-195, :448-470 parity). Removed two async-only behaviors with no sync counterpart: the transient reader-less force_refresh re-probe (sync warns and stays on the current connection) and the sqlalchemy.pool module heuristic in _is_pool_connection (SA+RWS is unsupported; the provider-manager check remains).
Replaces the single-shot router picker with the sync LimitlessRouterService semantics: WAIT_FOR_ROUTER_INFO (default True) with synchronous get-routers-with-retry honoring GET_ROUTER_MAX_RETRIES / GET_ROUTER_RETRY_INTERVAL_MS raising NoRoutersAvailable; least-loaded retry via highest_weight up to MAX_RETRIES_MS with UNAVAILABLE marking and MaxRetriesExceeded; AuroraLimitlessDialect gating with refresh-then-recheck; router cache TTL + idle monitor disposal per LIMITLESS_MONITOR_DISPOSAL_TIME_MS; monitor probes via force_connect with prefix-stripped props; 5s query timeout; monitor telemetry context. 38 tests mirroring the sync suite.
Full port of the sync blue_green_plugin.py semantics: two role-scoped status monitors driven by the provider (interval ramping BASELINE/ INCREASED/HIGH per phase), topology/DNS-aware corresponding-host mapping (writer-to-writer, sorted readers with modulo wrap, cluster-DNS mapping), hash-based interim-status dedup, rollback detection, switchover timer, DNS-completion flags, POST-phase RejectConnectRouting + SuspendUntilCorrespondingHostFound emission, routing re-selection loops in connect/execute, IAM host substitution, monitoring-prefix property overrides, suspend TimeoutError semantics with BG_CONNECT_TIMEOUT_MS, non-swallowed substitute failures, telemetry contexts + hold-time, and message-key routing. Tests grow 33 -> 60 mirroring sync behaviors. Known gap (follow-up): per-monitor host-list-provider bootstrap needs an async supplier on DatabaseDialect; topology mapping logic is implemented and tested, production mapping is driven by host names / cluster DNS.
… and Okta ADFS SAML acquisition now performs the sync form flow over aiohttp (GET sign-in page, parse form action + hidden inputs, POST urlencoded, scrape SAMLResponse) with SamlUtils URL/response validation, replacing the BasicAuth GET. All AWS calls route through AwsCredentialsManager / IamAuthUtils inside asyncio.to_thread so aws_profile and custom credential providers apply. Region resolution uses RegionUtils/GdbRegionUtils with hostname auto-discovery and raises on unresolved regions; hosts validate via IamAuthUtils.get_iam_host. Federated/Okta honor IAM_TOKEN_EXPIRATION. Secrets cache key includes the endpoint; malformed secret JSON maps to JsonDecodeError instead of being swallowed. Error keys mirror sync (ConnectException on first terminal failure, UnhandledException on post-refetch retry failure, AwsConnectError on network exceptions -- recognized as network by both dialect exception handlers). Token-cache size gauges added. 83 tests.
Fastest-response: monitor probes via force_connect with frt-prefixed monitoring properties and a 10s connect-timeout default; response-time state resets on CONNECTION_OBJECT_CHANGED; response-time cache and idle monitors expire after 10 minutes; corrected the misleading unit/identity docstring. Host list: provider.stop() tears down its topology monitor (was a leak until global cleanup); instance patterns validate against RDS Proxy/Custom endpoints and invalid DNS patterns; the base Aurora provider queries LAST_UPDATE_TIMESTAMP and picks the latest-updated writer when multiple writer rows appear; a ProgrammingError from the topology query raises RdsHostListProvider.InvalidQuery instead of silently returning empty. Plugin service: current_host_info gets sync's fallback chain (initial host, topology writer, allowed-host validation with the sync error keys); fill_aliases ported and wired; get_host_role honors AUXILIARY_QUERY_TIMEOUT_SEC.
Rearchitects the per-connection watchdog into a faithful v2 port: shared monitors keyed by detection-params + host URL with 60s idle expiry; dedicated monitoring connections opened via force_connect with monitoring-prefixed properties (never in-band pings on the app connection); per-execute monitoring contexts (weakrefs) bounding probes to active calls; duration-based failure math (interval * (count-1)); supports_abort_connection guard raising ConfigurationNotSupported; cluster-endpoint alias resolution on connect; CONNECTION_OBJECT_CHANGED- gated notify with NO_OPINION; efm2.connections.aborted counter; message- key routing. Abort now severs the socket (socket.shutdown on the raw fd, mirroring sync pg_driver_dialect) instead of close(), waking a suspended read promptly without raising CancelledError across the failover boundary; aiomysql documents why close() suffices on a single loop. Fixes the process-lifetime leak: no per-instance shutdown hooks -- one module-level registry hook; plugins and app connections are collectable once their execute scope ends. 34 EFM tests. Known approximation: cluster alias fill uses identify_connection + as_aliases (async service has no full fill_aliases wiring yet); the sync MonitorResetEvent bus has no async analog -- disposal is idle-expiry + host-deleted + shutdown.
execute() opens a TOP_LEVEL context named for the method with a python_call attribute and success flag; connect() opens a NESTED method context; every plugin invocation in the chain is wrapped in a NESTED context named after the plugin class; force_connect matches sync in opening no method-level span. The factory is cached once in __init__ and all context operations are None-guarded, so the default Null factory adds no overhead. Dispatch semantics (pipeline construction, subscription gating, suggested-action aggregation, host-list fan-out) unchanged. 7 new tests with a recording telemetry double.
… keys Failover raises use FailoverPlugin.TransactionResolutionUnknownError / ConnectionChangedError / UnableToConnectToReader / UnableToConnectToWriter; RWS raises use ReadWriteSplittingPlugin.SetReadOnlyFalseInTransaction / NoReadersAvailable / NoWriterFound -- the same catalog entries sync uses, replacing inline English strings. Behavior-preserving.
…pings The async psycopg dialect dropped the wrapper-level database property after the base stripped it, so URL-style connects (postgresql://host/db) and database= kwargs landed on the DEFAULT database -- re-add it as psycopg's dbname (sync PgDriverDialect parity). transfer_session_state now also carries isolation_level (best-effort, like read_only). Both dialects' ping() probes are bounded by a 10s asyncio.wait_for matching sync DriverDialect.ping's exec_timeout, so a blackholed host cannot hang the liveness probe until the OS TCP timeout.
Routes the remaining user-facing strings through the sync message-key catalog: SRW config/transaction/connection errors, topology-monitor fetch-failure and refresh-timeout messages, no-writer topology log, plugin-service transaction/identify/role errors, wrapper closed-conn set_read_only, custom-endpoint wait timeout, stale-DNS writer-not-allowed and dynamic-provider errors, and the initial-connection-strategy unsupported-strategy raise. Behavior-preserving; 16 substitutions across 8 files; remaining unmatched strings are candidates for new shared keys (tracked for the PR body).
AsyncDefaultPlugin.execute now bounds every dispatched DBAPI operation with asyncio.wait_for(SOCKET_TIMEOUT_SEC) when the property is set; on expiry the connection socket is severed via abort_connection and QueryTimeoutError raised with the sync DriverDialect.ExecuteTimeout message -- the async analog of sync DriverDialect.execute's network-bound timeout + abort. The AsyncDriverDialect docstring now points at the real enforcement site. Also closes two review test gaps: an aio.psycopg submodule PEP 249 contract test (mirroring sync) and an async cache put-displaces-and-disposes test.
Adds the filtered hosts property to the async plugin service (allowed/ blocked view, sync plugin_service parity) and uses it for default-plugin strategy selection. The async wrapper gains plugin-routed tpc_begin/ prepare/commit/rollback/recover, async cursor iteration (__aiter__ via plugin-routed fetchone), routed get_read_only/get_autocommit coroutines (the sync-property reads keep back-compat and now document that they bypass plugins), a TOP_LEVEL telemetry span around connect with failure marking, and close() releasing plugin-service resources best-effort. AsyncDefaultPlugin marks hosts AVAILABLE and refreshes the driver dialect after successful connect/force_connect (sync default-plugin bookkeeping; database-dialect upgrade continues to live in the async connect flow). 19 new tests.
Minor plugins: connect/execute timing logs with the sync message keys, developer-plugin method-name targeting on the one-shot exception (with empty-name guard) and a force_connect injection path, sync-matching subscriptions (connect-time: CONNECT+FORCE_CONNECT; execute-time: ALL). Connection tracker: 3-minute post-failover settling window keeps refreshing topology after a FailoverError instead of once. Custom endpoint: refresh-rate property threaded into the monitor, excluded members propagated into AllowedAndBlockedHosts, the CLUSTER_ID gate that silently disabled membership enforcement removed, and execute now re-ensures the monitor and waits for endpoint info like sync. Stale DNS: network-bound execute subscription with best-effort topology refresh. Initial connection strategy: reader candidates filtered to the connect URL's region. ~30 new/updated tests.
Optional statuses and corresponding-host pairs are assert-narrowed before attribute/index access, and the suspend-until-found fixture builds a real ConcurrentDict instead of a plain dict (the field's declared type). Clears 18 pre-existing type errors that slipped in with the BG port (only the source file had been type-checked).
Ports sync failover_v2's connect path: with enable_connect_failover (and failover) enabled, a failover-worthy error on the initial dial marks the target host UNAVAILABLE and runs failover instead of surfacing the error, and a target already known UNAVAILABLE skips the doomed dial entirely -- topology is refreshed and failover runs directly. FailoverFailedError propagates like sync; stale-DNS verification remains the separate async plugin. Closes the last recognized-but-inert flag from the parity review.
…h the sync fixes With the parity-review sync fixes in the base, the async layer drops its two compatibility shims: limitless _is_login_exception now returns the classification verdict so login failures short-circuit out of the router retry loop (previously mirrored sync's discarded-verdict behavior), and the Blue/Green suspend-until-corresponding-host log uses the SuspendConnectRouting.WaitConnectUntilCorrespondingHostFound message key instead of the plain-string workaround for the formerly missing key.
…reader Port of the Aurora-validated dead-writer convergence fix (async-parity 83e22cc, never previously pushed) onto phase2, adapted to the diff-and-notify plugin_service and i18n'd monitor. After a writer failover the async path could not discover the promoted writer when the old writer was UNREACHABLE: get_writer_connection tried only the topology-labeled writer and, at the bottom of the loop, refreshed topology through the DEAD current connection -- which can never observe the new writer -- looping until the failover deadline (observed on real Aurora: 29 connection attempts to the dead old writer, 0 to the surviving reader, ~5-min timeout in test_fail_from_writer_with_session_states_readonly_async[PG_ASYNC-failover_v2]). Sync-parity gaps closed (reference: aws main failover_v2 + RdsHostListProvider): 1. retry_util.get_writer_connection: when the writer-labeled host is unreachable/absent, connect to a surviving READER from the topology and force_refresh THROUGH it (_refresh_topology_via_reader) before the last-ditch dead-connection refresh. Sync's monitor likewise discovers the writer through reader connections (HostMonitor fan-out); this performs the same reader-driven refresh in-band. Validated on Aurora (v8 run: test passed post-fix). 2. failover_plugin._do_failover: when the live refresh returns empty, fall back to plugin_service.all_hosts (cached full topology with surviving readers) before the single initial_connection_host_info seed. Sync failover always operates on the full topology, never a single seed host. 3. Non-empty topology guards (sync _get_topology "use live only if len > 0"): plugin_service.refresh/force_refresh_host_list adopt only non-empty results (keeping phase2's diff-and-notify path); cluster_topology_monitor._run and force_refresh_with_connection never overwrite _last_topology with an empty result, and _run drops the likely-dead owned connection on an empty read so the next tick reopens to a live host (also keeps panic mode armed, which gates on a non-empty _last_topology). Adds the dead-writer regression test (old writer unreachable + surviving reader -> converges on the new writer via the reader refresh). Unit suite: 2186 passed.
…c parity) Sync DriverDialect.execute applies SOCKET_TIMEOUT_SEC only to the dialect's network_bound_methods (driver_dialect.py:134); the async terminal plugin was wrapping EVERY dispatched method in asyncio.wait_for, so users who set socket_timeout could see a legitimately-slow LOCAL method spuriously aborted with QueryTimeoutError -- behavior sync never produces. - AsyncDefaultPlugin.execute now consults the driver dialect's network_bound_methods with sync's exact gate (ALL-in-set or method-in-set); non-network methods run unbounded. The cancel-then-abort mechanism for bounded methods is unchanged. - Align the async dialects' network_bound_methods membership with their sync counterparts (PgDriverDialect / MySQLDriverDialect) plus two async-dispatch additions kept from phase2: CONNECT (async connects flow through the plugin pipeline) and CURSOR_EXECUTEMANY. Like sync MySQL, aiomysql leaves CONNECTION_CLOSE unbounded. - Regression test: a slow non-network method with socket_timeout set completes instead of raising QueryTimeoutError; existing timeout test now declares its method network-bound. Unit suite: 2187 passed.
…arity) The async IAM / Secrets Manager / Federated / Okta plugins cached tokens and secrets in PLUGIN-INSTANCE attributes, but a fresh plugin set is built on every connect() -- so each new connection re-fetched: a Secrets Manager get_secret_value network call per connection, and a full SAML round-trip + STS AssumeRoleWithSAML per connection for federated/Okta. Sync avoids exactly this by caching in the process-wide services_container StorageService (iam_plugin.py:58-59, aws_secrets_manager_plugin.py:73-74, federated_plugin.py:65-66, okta_plugin.py:61-62). Refactor to the sync pattern, sharing the SAME cache entries as the sync plugins in the same process: - AsyncIamAuthPlugin: TokenInfo entries keyed by IamAuthUtils.get_cache_key, registered 15-min expiration; wall-clock TokenInfo.is_expired checks (the async-only 60s regeneration grace window is dropped -- sync has none). - AsyncAwsSecretsManagerPlugin: Secret(SimpleNamespace(**secret)) entries keyed by the sync 3-tuple (secret_id, region, endpoint), registered 30-min expiration; the async SECRETS_MANAGER_EXPIRATION override is honored via put's per-item expiration. Custom username/password field keys still apply on cache hits (full secret is stored, like sync). - _RdsTokenMixin (federated/Okta): TokenInfo entries with sync string keys, registered 30-min expiration. - Invalidation on login-failure retry uses StorageService.remove; cache-size telemetry gauges report storage size. Tests updated to assert via the storage service, with an autouse fixture clearing TokenInfo/Secret around each test (the cache is now process-global). Unit suite: 2187 passed.
…y at 1.8.5 AWS stays on the Poetry 1.8.x line: unit CI (main.yml) pins poetry 1.8.2 and pyproject is deliberately kept in the 1.8-compatible [tool.poetry] layout (precedent: "fix: pyproject in Poetry 1.8.2-compatible [tool.poetry] format"). This branch's lock had been regenerated by Poetry 2.4.1 (lock-version 2.1) and the integration container's poetry install bumped to 2.3.4, leaving three poetry generations in play (CI 1.8.2 / container 2.3.4 / lock 2.4.1) and a "lock file might not be compatible" warning on every 1.8.x install. The 2.3.4 bump's stated rationale (PEP 621 [project] table) does not apply -- pyproject still uses [tool.poetry]. - Regenerate poetry.lock with Poetry 1.8.2 (lock-version 2.0, --no-update; all async deps -- aiomysql, aiohttp, greenlet -- unchanged). Verified: 1.8.2 `check --lock` + `install --dry-run` clean, and Poetry 2.4.1 reads the 2.0 lock without issue. - Revert ContainerHelper.java's in-container poetry to 1.8.5 with main's original comment (last version supporting python 3.8).
…s (sync parity) Sync failover_v2 treats a role-probe error as "drop this candidate for the pass": the reader loop's surrounding except removes the candidate (failover_v2_plugin.py:288-289) and the original-writer probe's `except: pass` moves on without accepting (:307-308). The async `_probe_role` instead fell back to the topology's ASSUMED role on probe failure, which (a) ACCEPTED an unverifiable candidate in non-STRICT_READER modes, and (b) in STRICT_READER falsely set `is_original_writer_still_writer` when the original-writer probe failed (assumed WRITER), permanently excluding a possibly-demoted writer from subsequent passes. `_probe_role` now returns None on failure; both call sites reject the candidate (close the probe connection quietly -- additive over sync, which leaves it to GC) and continue. Tests that leaned on the assumed-role fallback now stub get_host_role explicitly, matching the file's existing pattern. Unit suite: 2187 passed.
…ilover now exists) The comment claimed async failover 'currently only triggers on the execute path' and that the flag 'only takes effect once async grows connect-time failover' -- contradicted by _connect_with_failover added on this branch. It also claimed the flag defaults to true; ENABLE_CONNECT_FAILOVER defaults to False (properties.py:296). Reworded to describe the actual gating: connect-time failover requires enable_failover AND enable_connect_failover; execute-path failover is gated by enable_failover alone.
Completes the half-wired panic mode so the async topology monitor is the discovery engine for writer failover, exactly like sync v2 (cluster_topology_monitor.py + failover_v2_plugin.py:333): - Winning panic probe now fetches topology THROUGH the verified-writer connection and PUBLISHES it (monitor state + provider cache + wake event) -- sync HostMonitor._fetch_topology_and_update_cache (:561). Reader probes publish a topology that shows a NEW writer before closing (sync _reader_thread_fetch_topology, :589-612). - The run loop HARVESTS the panic-found writer: its connection is promoted to the monitor's own monitoring connection (the monitor keeps it; failover opens a fresh connection off the published topology -- sync :262-284). claim_verified_writer(), which had no production consumer, is removed. - New blocking AsyncClusterTopologyMonitor.force_monitoring_refresh( should_verify_writer, timeout): drops the monitoring connection + verified flag (deliberately forcing panic fan-out on monitor-owned connections), wakes the run loop immediately, and awaits the next topology publication -- sync force_refresh(True, t) + _wait_till_topology_gets_updated (:136-178). - AsyncAuroraHostListProvider.force_monitoring_refresh + adopt_topology; AsyncPluginService.force_monitoring_refresh_host_list (Protocol + impl, with a plain-forced-refresh fallback for monitor-less providers). - AsyncFailoverPlugin._do_failover: STRICT_WRITER is now MONITOR-FIRST -- block on force_monitoring_refresh_host_list(True, remaining) before the retry loop (sync failover_v2_plugin.py:333). Unlike sync's hard raise, a monitor failure falls through to the existing fallback chain so monitor-less providers still fail over. - AsyncRetryUtil.get_writer_connection: the bottom-of-loop refresh NO LONGER queries through the (dead) current connection -- it reads the monitor-maintained provider cache (force_refresh(None)), matching sync RetryUtil's cached refresh_host_list per pass. The in-band reader-refresh remains as the fallback for monitor-less providers and when the cache is unchanged. Tests: winner-probe publication, blocking verify-writer refresh end-to-end (dead cluster endpoint -> panic -> publish -> unblock), harvest/promote semantics, plugin-service adopt/timeout/fallback. Unit suite: 2192 passed.
…get_running_loop Three async-only hardening items (no sync counterpart; thread model differs): 1. Thread-safe monitor cancellation. Module-level monitor registries (EFM host monitoring, limitless, blue/green) can hand a monitor created on loop A to a caller running on loop B; Task.cancel() is not thread-safe across loops. New aio.cleanup.cancel_task_threadsafe routes through the owner loop's call_soon_threadsafe when the caller is on a different loop; the limitless and blue/green monitors now record their owner loop at start() (EFM already did), and their async stop paths no longer await a foreign-loop task (invalid) -- they signal-and-return in that case. 2. Pool-permit leak on acquire timeout (Python 3.10). _AsyncPool.acquire used asyncio.wait_for(semaphore.acquire(), t); on 3.10 the permit can be granted in the same tick the waiter is cancelled and Semaphore doesn't roll it back, permanently shrinking pool capacity. Replaced with asyncio.wait (no cancel-on-timeout) + explicit cancel, releasing the permit if the acquire completed anyway; still raises asyncio.TimeoutError. 3. asyncio.get_event_loop() -> asyncio.get_running_loop() across aio/ (29 call sites, 8 files). All call sites run inside coroutines; get_event_loop is deprecated there and can silently create a new loop when misused. Unit suite: 2192 passed.
…nt aiohttp prerequisite - Okta step-1 parsed resp.json() before _validate_response_status: a non-2xx non-JSON error body raised ContentTypeError, masking the intended SamlUtils.RequestFailed message. Validate first, then parse the captured body text. - PluginChainCompatibility.md: note that async federated/Okta require aiohttp at runtime (dev-group dependency only; users install it themselves).
….14) asyncio.iscoroutinefunction emits a DeprecationWarning on Python 3.14 (slated for removal in 3.16), which fails the unit suite under -Werror as CI runs it.
…rovider)
Integration finding (Aurora PG multi-2, pooled RWS suite): panic-mode writer
probes routed through plugin_service.connect, i.e. the EFFECTIVE connection
provider. With AsyncPooledConnectionProvider installed, failover probes
created internal pools (cluster-URL invariant 'no pools' broken) and left
stale pooled connections to the demoted writer that later tests received
('the connection is closed' / SSL EOF).
Sync parity: ClusterTopologyMonitor opens ALL its connections via
plugin_service.force_connect (cluster_topology_monitor.py:344 main conn,
:519 per-host probes) - plugin pipeline runs (auth re-applies) but the
DEFAULT provider is used. Async force_connect already mirrors those
semantics; use it in build_probe_host.
Unit: probe-helper tests now assert force_connect and that connect is never
awaited (regression); 2192 passed.
…ance_id get_instance_id is deliberately best-effort (returns None on failure), but the bare swallow made the EFM cluster-endpoint identification failure seen on Aurora (test_wrapper_connection_reader_cluster_with_efm_enabled_async: 'Unable to identify the connected database instance') undiagnosable from logs. Log the exception (i18n'd) before returning None; behavior unchanged.
…t task) Integration finding (Aurora PG multi-5): test_writer_failover_in_idle_ connections_async failed 2/10 params - idle connections to the demoted writer stayed open after failover. The tracker spawned invalidate_all as a fire-and-forget asyncio.create_task, drained only by the release_resources_async shutdown hook; any task still pending when the event loop closes is cancelled, losing the invalidation. Sync's offload (daemon Thread, aurora_connection_tracker_plugin.py:231) carries a lifetime guarantee asyncio tasks don't have: the thread outlives the failover call and always completes. The async equivalent of that guarantee is awaiting the invalidation inline in the execute/failover path before the FailoverError propagates (sync parity :356-357, incl. log_opened_connections). create_task remains only for the sync-signature notify_host_list_changed path (still drained on release_resources_async). Also adds sync-parity logging (InvalidatingConnections / OpenedConnectionsTracked, reusing sync's i18n keys) and debug-logs best-effort close failures instead of silently swallowing them. Unit: new regression test asserts every tracked conn is closed and no pending task remains BEFORE the failover exception reaches the caller (no cooperative yield); 2193 passed, flake8/isort/mypy clean.
…ways failed) Integration finding (deterministic, both Aurora envs): test_wrapper_connection_reader_cluster_with_efm_enabled_async raised '[HostMonitoringV2Plugin] Unable to identify the connected database instance' on every cluster-RO connect. Root cause: AsyncAuroraHostListProvider._rows_to_topology built HostInfo rows without host_id, and plugin_service.identify_connection matches 'h.host_id == instance_id' -- so identification of ANY cluster-endpoint connection failed unconditionally: the instance-id query succeeded, but no topology row could ever match. Sync parity: create_host (host_list_provider.py:548-556) sets host_id and adds the bare instance id as an alias. The async MultiAz and GlobalAurora providers already do both; the main Aurora provider was the one omission. Unit: regression test asserts topology rows carry host_id + instance-id alias; 2194 passed, flake8/isort/mypy clean.
…etection race) Targeted re-run finding (fresh Aurora cluster, fix branch): 1/10 writer_failover_in_idle_connections_async params still failed -- the new invalidation logging proved that in the failing param NO invalidation was attempted at all: the FailoverError handler's plain refresh_host_list() served the monitor's cache from BEFORE the monitor observed the failover, so _update_writer_from_topology saw no writer change (detection race; the 92095a6 inline-await fix only guarantees completion once a change is detected). Sync tolerates that staleness: its daemon-thread and topology-notify paths re-detect the change after the failover call returns. In asyncio everything dies with the event loop, so the handler is the one deterministic shot at detection. Use force_refresh_host_list() there -- it queries topology through the post-failover current connection (the just-verified new writer), bypassing the lagging cache. The settling window's per-execute refresh stays a plain refresh (sync parity). Unit: settling-window test now asserts the handler refresh is a force refresh; 2194 passed, flake8/isort clean.
Second targeted re-run finding (fresh Aurora multi-5, code incl. 41efc70): 1/10 idle-connection params STILL missed invalidation -- zero invalidation events in the window even with the handler's force refresh. Root cause: right after promotion, aurora_replica_status itself can still report the OLD writer (server-side lag), so ANY topology-based comparison can miss the change no matter how it is queried. The failover plugin, however, just MOVED this connection: comparing plugin_service.current_host_info captured before the failing execute with its post-failover value is deterministic and needs no topology cooperation. On a provable move, invalidate the departed host's tracked connections and re-pin the writer to the connection's new host. Guarded: only fires when the departed host is the pinned writer (or no writer was pinned), so reader-mode failovers moving off a reader don't nuke that host; when the topology comparison already handled the change it is a no-op. Also debug-logs the handler's swallowed refresh failures. Sync does not need this: its daemon-thread and topology-notify paths re-detect the change once replica_status catches up; async's last chance is before the loop closes. Unit: stale-topology regression test (topology frozen on the old writer, detection via host move alone) + reader-departure guard test; 2196 passed, flake8/isort/mypy clean.
…ettles Comparing against sync exposed the real invariant my monitor port broke: sync failover can only succeed THROUGH the topology cache (it connects to the cache's writer entry and role-verifies it), so by the time FailoverSuccessError propagates, the cache provably names the new writer and the CONVERTED_TO_* notifications have fired. The async port's candidate-probing fallback (deliberate robustness, aws#1246 lineage) lets failover succeed while aurora_replica_status -- and therefore every publication built from it -- still names the demoted writer for several seconds after promotion. Consumers then act on stale roles: the connection tracker misses the writer change (idle connections survive), and RWS can switch to a stale-epoch writer. Fix at the source: the panic-mode winner was verified via a live is_reader probe (sync's own verification standard, HostMonitor :550-560), so while the post-panic settling window is open the monitor publishes a role-corrected view -- the verified writer's row is WRITER, any other row claiming WRITER is demoted, a missing row is appended. The window is armed at verification (sync arms it at WriterPickedUpFromHostMonitors, :272-274 -- previously missing from the port) and expires after HIGH_REFRESH_PERIOD_SEC so a genuinely newer failover is never masked for long. Publications outside the window are raw, as before. Complements (does not replace) the tracker's departed-host invalidation: that guards the tracker even when the monitor path is bypassed entirely. Unit: stale-fetch winner publication overlay + window-expiry regression tests; 2198 passed, flake8/isort/mypy clean.
…onnecting
Stabilizes the post-failover instance-reboot race observed in 3 consecutive
Aurora runs: pooled_connection__failover_async (block-2 reconnect) and
pooled_connection__different_users_async (setup) connect to the
just-demoted writer seconds after triggering failover; mid-reboot the
instance accepts the connection and kills the session at first use ('the
connection is closed' at cursor(), SSL EOF on the first DDL, 'FATAL: the
database system is starting up').
Bounded wait (rds_utils.wait_until_instance_has_desired_status, 5 min,
via asyncio.to_thread) before the racy reconnects; no assertion changes.
The sync twins share the identical ordering hazard and are deliberately
left untouched: the equivalent sync change is proposed upstream for AWS
review (see sync-failure report, suggestion 3).
Integration finding F5 (test_connect_to_writer__switch_read_only_async [srw], Aurora multi-5, 3 consecutive runs): the simple-RWS plugin had no notify_connection_changed hook, so set_current_connection's default old-connection handling CLOSED the cached reader connection every time the plugin switched back to the writer. The next set_read_only(True) then opened a brand-new read-endpoint connection and, on clusters with several readers, landed on a DIFFERENT reader than the first toggle -- breaking the reader-reuse contract the test asserts. Invisible on 1-reader clusters (any reconnect lands on the same instance), which is why multi-2 always passed. Sync parity: AbstractReadWriteSplittingPlugin.notify_connection_changed returns PRESERVE mid-split (read_write_splitting_plugin.py:99-107); the srw equivalent votes PRESERVE exactly when the outgoing current connection is one of the plugin's cached endpoint connections, and NO_OPINION otherwise (a failover replacing a dead connection must not be blocked from disposal). Unit: regression test for PRESERVE on cached reader/writer departure and NO_OPINION on foreign/none; 2199 passed, flake8/isort/mypy clean.
… topology v11 full-axis finding (multi-2, first two idle-connection params, both starting <30s after a prior failover): _pin_current_writer pinned _current_writer from the connect-time topology refresh, but right after a failover aurora_replica_status can still name the DEMOTED writer -- so the pin was stale. A wrong pin defeats BOTH failover-time detection paths at once: the topology comparison sees "no change" (lagged view agrees with the lagged pin) and the departed-host guard refuses because the pinned writer disagrees with the host the connection departed. Result: zero invalidation attempts, idle connections to the demoted writer survived. Keep the topology-based pin as the base, then override with ground truth: a live is_reader probe of THIS connection (the same verification standard sync's monitor uses for writer detection). Probe says WRITER -> pin the connection's own host over whatever topology claimed. Best-effort like the rest of the pin path. Unit: regression test (lagged topology names a stale writer; probe truth wins; departed-host invalidation fires on failover with topology still lagged); 2200 passed, flake8/isort clean.
The v11 run proved the status-based stabilization (9542cd8) ineffective: RDS keeps the demoted writer's instance status 'available' through an Aurora failover engine restart (the wait returned in <1s and the test still got a killed session -- different_users_async failed again with 'the connection is closed' at first cursor use). Replace it with a live probe: wait_until_endpoint_accepts_queries_async retries a RAW driver connect + SELECT 1 (5s connect timeout, 2s backoff, 180s bound, best-effort on timeout) until the endpoint actually serves queries -- the only signal that survives Aurora's status semantics. Applied at both racy reconnects (pooled_connection__failover_async block 2, different_users_async setup). Assertions unchanged; sync twins still untouched pending upstream review.
…ort it Finding F6 (deterministic: test_pooled_connection__different_users_async failed 10/10 in every multi-instance env across four Aurora runs; a live raw-connect probe proved the server healthy while the wrapper's pooled connection was BAD milliseconds later, falsifying the earlier env-timing classification). Root cause: AsyncAwsWrapperConnection.close() first returns the pooled connection ALIVE to the pool (proxy.close -> pool.release), then calls plugin_service.release_resources(), which ABORTED the current connection via driver_dialect.abort_connection -- a hard socket kill. The pool was left holding a corpse; the next acquirer (the test's finally-block cleanup connection, pool_size=1) received it: 'the connection is closed' at cursor() / 'SSL error: unexpected eof' at first execute. Sync parity: sync release_resources CLOSES current_connection (plugin_service.py:783-789) -- for SQLAlchemy pool fairies that is an idempotent pool-release, for raw connections an idempotent no-op. Match it: close (guarded by driver_dialect.is_closed) instead of abort, and make _PooledAsyncConnectionProxy.close() idempotent (double release corrupted the pool's checked-out count, semaphore, and idle queue -- SQLAlchemy's fairy close is idempotent, ours was not). Unit: two wrapper tests updated to the sync-parity contract (pipeline close + release close; second is a skipped/idempotent no-op on real drivers -- mocks keep closed=False so both awaits are counted); 2200 passed, flake8/isort/mypy clean.
Two multi-5 idle-connection params still missed invalidation entirely (no path fired) even with the live-probe pin; the existing logging cannot distinguish handler-not-entered from both-paths-bailed. Log pre/post/ pinned hosts at the decision point so the next integration run is decisive.
MySQL integration finding (idle-connection async params failed 7/7 in multi-2; the new handler diagnostic made it decisive: pre=old-writer, post=new-writer, pinned=NEW-writer): on MySQL the connect-time writer pin fails silently (the database dialect is not yet upgraded when the tracker's connect hook runs, so the connect-time topology refresh/role probe error out and are swallowed). At failover time the handler's own force-refresh then registers the NEW writer as a FIRST observation -- the pin equals post, no change is detected by the topology comparison, and the departed-host guard vetoed because the pin differed from pre. Every detection path missed; idle connections to the demoted writer survived. Relax the guard to veto only when the pin names a THIRD host distinct from both ends of the move: pinned == pre (normal), pinned is None (pin failed), and pinned == post (first-observation) all invalidate the provably-departed host; a reader-mode failover with the real writer pinned elsewhere still skips. Unit: MySQL-shaped regression (no connect pin, first-observation at handler time, departed host invalidated); 2201 passed, flake8/isort clean.
…twork Finding F8 (MySQL 3.14 confirm run, multi-5 on a cluster doing 60-90s failovers -- 4/4 identical tracebacks): during the long outage aiomysql's reader task sees EOF and tears the connection down locally; every later operation then raises pymysql.err.InterfaceError(0, 'Not connected') -- args (0, msg) with NO 2xxx client error code, so the existing pymysql shape-matching (args[0] in _NETWORK_ERRORS) missed it and the error escaped the async wrapper RAW instead of triggering failover. With fast failovers the first post-outage operation gets a classified 2xxx error instead, which is why multi-2 passed 10/10 on the same code. Narrow, additive match (code 0 + 'Not connected' message), following the same precedent as the existing pymysql-shape block in this shared handler: mysql.connector network errors are caught by earlier branches, so the sync verdict is unchanged; flagged for the maintainer regardless since the file is shared with sync. Unit: regression asserts the shape classifies as network and a code-0 error with an unrelated message does not; 2202 passed, flake8 clean.
…ected' F8 residual (1/10 idle params still failed with the handler fix active; full traceback proved _should_failover returned False and the raw pymysql.err.InterfaceError(0, 'Not connected') propagated): the handler-based classification path returns False whenever the plugin service's database dialect is transiently unset -- ExceptionManager resolves the handler FROM the dialect, and no dialect means no handler means False, regardless of the handler fix. Recognize the shape directly in the async failover plugin's _should_failover, dialect-independent, walking the __cause__ chain -- the same precedent as the existing _is_connection_os_error and 'connection is lost' escape hatches. Also debug-log any exception that is NOT classified as failover-worthy: a raw driver error propagating past failover was previously invisible in logs, which cost several diagnosis cycles. Unit: shape recognition (direct, cause-wrapped, and narrow-match negative); 2203 passed, flake8/isort/mypy clean.
The not-classified debug line (added one commit ago) caught it live on
the very next run: aiomysql does NOT raise the pymysql tuple form -- it
raises InterfaceError("(0, 'Not connected')"), the tuple's repr embedded
in a SINGLE string argument (aiomysql/connection.py:1123). Both the
handler branch and the failover plugin's dialect-independent hatch
matched only the two-arg tuple form and missed it.
Extend both to the single-string form, narrowly ("(0," prefix + message).
Still unreachable for mysql.connector, which always carries
(errno, msg, sqlstate) 3-tuples with errno normalized to -1 (verified
empirically, including adversarial errno=0 construction).
Unit: real-shape regressions on both paths (plus narrow-match negatives);
2203 passed, flake8 clean.
… connect hook Restores sync parity that the async connect flow was silently missing: sync upgrades the database dialect from the live connection INSIDE the terminal plugin's connect (default_plugin.py:82 -> plugin_service.update_dialect), so every outer plugin's post-connect logic already sees the corrected dialect. The async layer instead ran its upgrade (_upgrade_database_dialect_after_connect) only after the whole pipeline in AsyncAwsWrapperConnection.connect -- plugin connect hooks therefore worked with the pattern-guessed dialect, which on an Aurora MySQL INSTANCE endpoint is RdsMysqlDialect (no topology query; @@read_only role probe that does not flip on Aurora failover). That ordering gap was the root asymmetry behind the connection tracker's pin failures on MySQL (F7) and the transient dialect-dependent exception-classification miss (F8 residual); both defenses remain as belt-and-suspenders with comments updated to reflect the restored parity. Changes: - AsyncPluginServiceImpl.update_database_dialect(conn): the relocated upgrade logic (probe 'aurora_version' for base/RDS MySQL dialects, switch to AuroraMysqlDialect, rewire the Aurora provider's topology query/port), with a settled flag mirroring sync's can_update-goes- False semantics; transient probe failures leave it unsettled so the next connect retries (sync runs update_dialect on every connect). - AsyncDefaultPlugin.connect/force_connect call it right after the connection is obtained (sync ordering: availability -> driver dialect -> database dialect). - aio/wrapper.py drops the post-pipeline upgrade; the connect-time topology priming stays as belt-and-suspenders with its comment rewritten (its original justification -- hooks saw the un-upgraded dialect -- no longer applies). Unit: ordering regression (outer post-connect logic sees the upgrade already applied), settle/no-op semantics test; 2205 passed, flake8/isort/mypy clean.
AhmadMasry
force-pushed
the
feat/async-phase2
branch
from
030ffff to
b27ab83
Compare
Contributor
Author
|
Hi @karenc-bq |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
feat: asyncio support (async wrapper, plugins, and SQLAlchemy dialects)
Adds a full asyncio counterpart of the wrapper under
aws_advanced_python_wrapper.aio, targeting behavioral parity with the sync wrapper for the shipped plugin set: failover v2, read/write splitting, EFM v2, IAM auth, AWS Secrets Manager, federated + Okta auth, Aurora connection tracker, cluster topology monitor, custom endpoint, stale DNS, Aurora initial connection strategy, simple read/write splitting, developer plugin, blue/green deployment, limitless, and fastest-response strategy.Backed by psycopg (async) and aiomysql, with async SQLAlchemy support via
create_async_engine(postgresql+aws_wrapper_psycopg://serves both sync and async; MySQL async usesmysql+aws_wrapper_aiomysql://). IncludesAsyncConnectionProvider/AsyncPooledConnectionProvider,AsyncSessionStateService, an async IdP factory registry, andrelease_resources_async()for background-task teardown. Builds on the previously merged sync groundwork (#1252, #1255).Design principles
main; deviations exist only where asyncio's execution model requires them, and each is documented in code comments citing the sync file/line it mirrors or departs from.is_readerchecks, publish topology through the verified connection, and the monitor keeps that connection. During the post-promotion settling window the published topology is role-corrected against the probe-verified writer, becauseaurora_replica_statuscan lag the promotion and async failover (unlike sync) can succeed without cache convergence.release_resourcescloses (never aborts) the current connection, and the pooled-connection proxy close is idempotent — matching the SQLAlchemy pool-fairy semantics the sync wrapper relies on.update_database_dialect, mirroring syncdefault_plugin.py:82→update_dialect), so outer plugins' post-connect logic always sees the corrected dialect.Sync-visible file changes (called out for review)
The async layer is additive. The only non-
aio/product changes:utils/mysql_exception_handler.py: two narrow, additive branches classifying the pymysql/aiomysql "Not connected" shapes (InterfaceError(0, 'Not connected')and aiomysql's single-string variant) as network errors. Verified unreachable formysql.connector(the sync driver): its errors always carry(errno, msg, sqlstate)tuples witherrnonormalized to-1, soargs[0] == 0cannot match; pymysql is not a dependency of the sync wrapper.resources/…messages.properties: additive i18n keys read only by async code._is_login_exceptionreturns its verdict; session-state reset-callable fix; MySQLis_read_only_exceptionwidening — each flagged in its own commit.Validation
Full integration matrix on real Aurora (us-east-2), fresh clusters per axis, sync + async drivers in every axis (matrix ran at the branch head minus the final dialect-relocation commit; that commit was separately validated on Aurora with a targeted MySQL run — 106/0/0 on the affected paths — plus the full unit suite):
maincode paths (byte-identical files) and will be reported separately with evidence (pooled-connection eviction on topology-detected failover; SSL-EOF error shapes during failover recovery; one toxiproxy harness flake; one test assumption about failover never re-electing the same writer).host_idpropagation, pooled-provider lifecycle, srw connection preservation, MySQL error-shape classification), each with a regression unit test citing the integration failure it reproduces.By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.