Skip to content

Allow vecs Client creation with prebuild SQLAlchemy Engine or with creator method #111

@kziovas

Description

@kziovas

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions