add add_columns and drop_columns method to ObVecClient (#15)

Author: whhe

pyobvector/client/ob_vec_client.py

@@ -90,11 +90,19 @@ def __init__(
                     )
 
     def refresh_metadata(self, tables: Optional[list[str]] = None):
-        """Reload metadata from the database."""
-        if tables is None:
-            self.metadata_obj.reflect(bind=self.engine, extend_existing=True)
-        else:
+        """Reload metadata from the database.
+        
+        Args:
+            tables (Optional[list[str]]): names of the tables to refresh. If None, refresh all tables.
+        """
+        if tables is not None:
+            for table_name in tables:
+                if table_name in self.metadata_obj.tables:
+                    self.metadata_obj.remove(Table(table_name, self.metadata_obj))
             self.metadata_obj.reflect(bind=self.engine, only=tables, extend_existing=True)
+        else:
+            self.metadata_obj.clear()
+            self.metadata_obj.reflect(bind=self.engine, extend_existing=True)
 
     def _insert_partition_hint_for_query_sql(self, sql: str, partition_hint: str):
         from_index = sql.find("FROM")
@@ -808,3 +816,47 @@ def perform_raw_text_sql(
         with self.engine.connect() as conn:
             with conn.begin():
                 return conn.execute(text(text_sql))
+
+    def add_columns(
+        self,
+        table_name: str,
+        columns: list[Column],
+    ):
+        """Add multiple columns to an existing table.
+        
+        Args:
+            table_name (string): table name
+            columns (list[Column]): list of SQLAlchemy Column objects representing the new columns
+        """
+        compiler = self.engine.dialect.ddl_compiler(self.engine.dialect, None)
+        column_specs = [compiler.get_column_specification(column) for column in columns]
+        columns_ddl = ", ".join(f"ADD COLUMN {spec}" for spec in column_specs)
+        
+        with self.engine.connect() as conn:
+            with conn.begin():
+                conn.execute(
+                    text(f"ALTER TABLE `{table_name}` {columns_ddl}")
+                )
+                
+        self.refresh_metadata([table_name])
+
+    def drop_columns(
+        self,
+        table_name: str,
+        column_names: list[str],
+    ):
+        """Drop multiple columns from an existing table.
+        
+        Args:
+            table_name (string): table name
+            column_names (list[str]): names of the columns to drop
+        """
+        columns_ddl = ", ".join(f"DROP COLUMN `{name}`" for name in column_names)
+        
+        with self.engine.connect() as conn:
+            with conn.begin():
+                conn.execute(
+                    text(f"ALTER TABLE `{table_name}` {columns_ddl}")
+                )
+                
+        self.refresh_metadata([table_name])

tests/test_ob_vec_client.py

@@ -232,6 +232,51 @@ def test_refresh_metadata(self):
         ).fetchall()
         self.assertTrue(len(res) > 0)
 
+    def test_alter_table(self):
+        test_collection_name = "ob_alter_table_test"
+        self.client.drop_table_if_exist(test_collection_name)
+
+        self.client.create_table(
+            table_name=test_collection_name,
+            columns=[
+                Column("id", Integer, primary_key=True, autoincrement=True),
+            ]
+        )
+        self.client.add_columns(
+            table_name=test_collection_name,
+            columns=[
+                Column("name", String(64), nullable=True),
+                Column("age", Integer, nullable=True),
+            ]
+        )
+        self.client.insert(
+            table_name=test_collection_name,
+            data={
+                "id": 1,
+                "name": "Alice",
+                "age": 20,
+            },
+        )
+
+        res = self.client.get(
+            table_name=test_collection_name,
+            ids=[1],
+        ).fetchall()
+        self.assertEqual(len(res), 1)
+        self.assertEqual(len(res[0]), 3)
+
+        self.client.drop_columns(
+            table_name=test_collection_name,
+            column_names=["age"]
+        )
+
+        res = self.client.get(
+            table_name=test_collection_name,
+            ids=[1],
+        ).fetchall()
+        self.assertEqual(len(res), 1)
+        self.assertEqual(len(res[0]), 2)
+
 
 if __name__ == "__main__":
     unittest.main()