-
Notifications
You must be signed in to change notification settings - Fork 41
Description
Currently SQLAlchemy allows the creation of an engine with a creator method which overrides the other engine creation arguments for example in my code:
from google.cloud.sql.connector import Connector
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
connector = Connector()
def getconn():
return connector.connect(
CLOUD_SQL_INSTANCE,
"pg8000",
user=DATABASE_USER,
password=DATABASE_PASSWORD,
db=DATABASE_NAME,
)
# SQLAlchemy engine using the connector
engine = create_engine("postgresql+pg8000://", creator=getconn, future=True)
This allows us to use SQLAlchemy with google cloud connector instead of passing a connection string. Extremely useful for cloud deployed databases to use the IAM authorization through the google cloud connector instead of exposing a connection to the internet.
Vecs also creates an SQLAlchemy engine under the hood when you call create_client or the Client() class itself. However both ONLY accept a connection string:
def create_client(connection_string: str) -> Client:
"""Creates a client from a Postgres connection string"""
return Client(connection_string)
This makes it impossible to use vecs client with google sql connector.
I managed to get this working just my patching the init method to accept an existing SQLAlchemy engine for example:
def _patched_init(self, engine):
self.engine = engine
from sqlalchemy import MetaData
from sqlalchemy.orm import sessionmaker
self.meta = MetaData(schema="vecs")
self.Session = sessionmaker(self.engine)
with self.Session() as sess:
with sess.begin():
sess.execute(text("create schema if not exists vecs;"))
sess.execute(text("create extension if not exists vector;"))
self.vector_version: str = sess.execute(
text(
"select installed_version from pg_available_extensions where name = 'vector' limit 1;"
)
).scalar_one()
Client.__init__ = _patched_init
so this is probably an easy addition that woould allow vecs client to work with cloud deployed DBs using the standard and more secure cloud client for example in this case the google cloud sql connector.