Feat: add support for initializing vecs client with custom schema

src/tests/conftest.py

@@ -13,7 +13,7 @@
 import vecs
 
 PYTEST_DB = "postgresql://postgres:password@localhost:5611/vecs_db"
-
+PYTEST_SCHEMA = "test_schema"
 
 @pytest.fixture(scope="session")
 def maybe_start_pg() -> Generator[None, None, None]:
@@ -94,12 +94,12 @@ def maybe_start_pg() -> Generator[None, None, None]:
 def clean_db(maybe_start_pg: None) -> Generator[str, None, None]:
     eng = create_engine(PYTEST_DB)
     with eng.begin() as connection:
-        connection.execute(text("drop schema if exists vecs cascade;"))
+        connection.execute(text(f"drop schema if exists {PYTEST_SCHEMA} cascade;"))
     yield PYTEST_DB
     eng.dispose()
 
 
 @pytest.fixture(scope="function")
 def client(clean_db: str) -> Generator[vecs.Client, None, None]:
-    client_ = vecs.create_client(clean_db)
+    client_ = vecs.create_client(clean_db, PYTEST_SCHEMA)
     yield client_

src/tests/test_client.py

@@ -1,7 +1,12 @@
 import pytest
-
 import vecs
 
+def test_create_client(clean_db) -> None:
+    client = vecs.create_client(clean_db)
+    assert client.schema == "vecs"
+
+    client = vecs.create_client(clean_db, "my_schema")
+    assert client.schema == "my_schema"
 
 def test_extracts_vector_version(client: vecs.Client) -> None:
     # pgvector version is sucessfully extracted

src/vecs/__init__.py

@@ -23,6 +23,6 @@
 ]
 
 
-def create_client(connection_string: str) -> Client:
+def create_client(connection_string: str, schema: str="vecs") -> Client:
     """Creates a client from a Postgres connection string"""
-    return Client(connection_string)
+    return Client(connection_string=connection_string, schema=schema)

src/vecs/client.py

@@ -47,23 +47,24 @@ class Client:
         vx.disconnect()
     """
 
-    def __init__(self, connection_string: str):
+    def __init__(self, connection_string: str, schema: str):
         """
         Initialize a Client instance.
 
         Args:
             connection_string (str): A string representing the database connection information.
-
+            schema (str): A string representing the database schema to connect to.
         Returns:
             None
         """
+        self.schema = schema
         self.engine = create_engine(connection_string)
-        self.meta = MetaData(schema="vecs")
+        self.meta = MetaData(schema=self.schema)
         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(f"create schema if not exists {self.schema};"))
                 sess.execute(text("create extension if not exists vector;"))
                 self.vector_version: str = sess.execute(
                     text(
@@ -105,7 +106,7 @@ def get_or_create_collection(
             CollectionAlreadyExists: If a collection with the same name already exists
         """
         from vecs.collection import Collection
-
+        
         adapter_dimension = adapter.exported_dimension if adapter else None
 
         collection = Collection(
@@ -162,7 +163,7 @@ def get_collection(self, name: str) -> Collection:
             join pg_attribute pa
                 on pc.oid = pa.attrelid
         where
-            pc.relnamespace = 'vecs'::regnamespace
+            pc.relnamespace = '{self.schema}'::regnamespace
             and pc.relkind = 'r'
             and pa.attname = 'vec'
             and not pc.relname ^@ '_'

src/vecs/collection.py

@@ -235,7 +235,7 @@ def _create_if_not_exists(self):
             join pg_attribute pa
                 on pc.oid = pa.attrelid
         where
-            pc.relnamespace = 'vecs'::regnamespace
+            pc.relnamespace = '{self.client.schema}'::regnamespace
             and pc.relkind = 'r'
             and pa.attname = 'vec'
             and not pc.relname ^@ '_'
@@ -289,7 +289,7 @@ def _create(self):
                 text(
                     f"""
                     create index ix_meta_{unique_string}
-                      on vecs."{self.table.name}"
+                      on {self.client.schema}."{self.table.name}"
                       using gin ( metadata jsonb_path_ops )
                     """
                 )
@@ -576,7 +576,7 @@ def _list_collections(cls, client: "Client") -> List["Collection"]:
         """
 
         query = text(
-            """
+            f"""
         select
             relname as table_name,
             atttypmod as embedding_dim
@@ -585,7 +585,7 @@ def _list_collections(cls, client: "Client") -> List["Collection"]:
             join pg_attribute pa
                 on pc.oid = pa.attrelid
         where
-            pc.relnamespace = 'vecs'::regnamespace
+            pc.relnamespace = '{client.schema}'::regnamespace
             and pc.relkind = 'r'
             and pa.attname = 'vec'
             and not pc.relname ^@ '_'
@@ -636,13 +636,13 @@ def index(self) -> Optional[str]:
 
         if self._index is None:
             query = text(
-                """
+                f"""
             select
                 relname as table_name
             from
                 pg_class pc
             where
-                pc.relnamespace = 'vecs'::regnamespace
+                pc.relnamespace = '{self.client.schema}'::regnamespace
                 and relname ilike 'ix_vector%'
                 and pc.relkind = 'i'
             """
@@ -760,7 +760,7 @@ def create_index(
             with sess.begin():
                 if self.index is not None:
                     if replace:
-                        sess.execute(text(f'drop index vecs."{self.index}";'))
+                        sess.execute(text(f'drop index "{self.client.schema}"."{self.index}";'))
                         self._index = None
                     else:
                         raise ArgError("replace is set to False but an index exists")
@@ -787,7 +787,7 @@ def create_index(
                         text(
                             f"""
                             create index ix_{ops}_ivfflat_nl{n_lists}_{unique_string}
-                              on vecs."{self.table.name}"
+                              on {self.client.schema}."{self.table.name}"
                               using ivfflat (vec {ops}) with (lists={n_lists})
                             """
                         )
@@ -806,7 +806,7 @@ def create_index(
                         text(
                             f"""
                             create index ix_{ops}_hnsw_m{m}_efc{ef_construction}_{unique_string}
-                              on vecs."{self.table.name}"
+                              on {self.client.schema}."{self.table.name}"
                               using hnsw (vec {ops}) WITH (m={m}, ef_construction={ef_construction});
                             """
                         )