Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: consolidate Postgres providers #494

Merged
merged 15 commits into from
Sep 22, 2024
Merged

feat: consolidate Postgres providers #494

merged 15 commits into from
Sep 22, 2024

Conversation

anubhav756
Copy link
Contributor

@anubhav756 anubhav756 commented Sep 18, 2024

To align with our existing AlloyDB and CloudSQL providers, we're updating the non-cloud Postgres provider to use SQLAlchemy. This also consolidates AlloyDB and CloudSQL Postgres interactions through the non-Cloud Postgres client by reusing its query handlers.

Also, the format_sql helper now accepts dict for its params argument. This change accommodates the AlloyDB and CloudSQL Postgres providers, and now, post consolidation, non-cloud Postres client as well, which pass params as dict through postgres.py. The updated format_sql helper is used in postgres.py query handlers to format SQL queries, populate them with variables, and return them as traces.

Create a common PostgresDatastore class to encapsulate all the DB query logic based on a given connection pool. Reuse the same logic across AlloyDB & CloudSQL postgres providers for better code maintainability and less redundancy.
@anubhav756 anubhav756 changed the base branch from consolidate-pg-providers to main September 18, 2024 19:58
The `format_sql` helper now supports both `tuple` and `dict` for its `params` argument. This change accommodates the AlloyDB and CloudSQL Postgres providers, which pass `params` as `dict` through `postgres_datastore.py`, unlike the non-cloud Postgres provider that uses `tuples`.  The updated `format_sql` helper is used in `postgres_datastore.py` query handlers to format SQL queries, populate them with variables, and return them as traces.
@anubhav756 anubhav756 changed the base branch from main to consolidate-pg-providers September 18, 2024 21:02
@anubhav756 anubhav756 marked this pull request as ready for review September 18, 2024 22:07
@anubhav756 anubhav756 requested a review from a team as a code owner September 18, 2024 22:07
Copy link
Member

@vishwarajanand vishwarajanand left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@vishwarajanand
Copy link
Member

I believe this needs to be merged before #493

Base automatically changed from consolidate-pg-providers to main September 19, 2024 09:55
anubhav756 and others added 6 commits September 19, 2024 21:07
Co-authored-by: Vishwaraj Anand <vishwaraj.anand00@gmail.com>
The `format_sql` helper now supports both `tuple` and `dict` for its `params` argument. This change accommodates the AlloyDB and CloudSQL Postgres providers, which pass `params` as `dict` through `postgres_datastore.py`, unlike the non-cloud Postgres provider that uses `tuples`.  The updated `format_sql` helper is used in `postgres_datastore.py` query handlers to format SQL queries, populate them with variables, and return them as traces.
Co-authored-by: Vishwaraj Anand <vishwaraj.anand00@gmail.com>
retrieval_service/datastore/helpers.py Outdated Show resolved Hide resolved
Copy link
Collaborator

@kurtisvg kurtisvg left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PR looks unchanged

To align with our existing AlloyDB and CloudSQL providers, we're updating the non-cloud Postgres provider to use SQLAlchemy. This consolidates all Postgres interactions through the common `postgres_datastore.py` reusing its query handlers.
This commit renames the `pool` variable to `async_engine` in all Postgres providers.

Previously, these providers used `asyncpg` directly for connection pool management.  They have now been migrated to SQLAlchemy, which uses `AsyncEngine` for connection management. This name change reflects the shift to SQLAlchemy and provides better clarity.
This is now possible since non-cloud `postgres.py` was the only file where tuple was being passed to `format_sql` prior to its consolidation, but after consolidation, it passes `dict` only.
@kurtisvg kurtisvg changed the title feat: Add tracing for AlloyDB and CloudSQL Postgres providers feat: consolidate Postgres providers Sep 22, 2024
@anubhav756 anubhav756 merged commit 2fa03bc into main Sep 22, 2024
15 checks passed
@anubhav756 anubhav756 deleted the tracing branch September 22, 2024 18:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants