feat: fork database (#185)

Poseidon-fan · web-flow · commit 8a68b6c96caa · 2026-03-18T17:32:44.000+08:00
Fixes #184  ## Summary by CodeRabbit * **New Features** * Admin-level database fork: create independent duplicates of databases (preserves data and attributes). Requires server >= 1.2.0. * **Tests** * Added integration tests covering successful forks, empty sources, attribute preservation, destination-exists errors, independence, and repeated forks. * **Chores** * Updated pre-commit Python version and lint configuration for tests. * **Style** * Minor test cleanup (no functional change).
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,3 +1,6 @@
+default_language_version:
+  python: python3.11
+
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: "v6.0.0"
diff --git a/pyproject.toml b/pyproject.toml
@@ -107,7 +107,7 @@ ignore = [
 ]
 
 [tool.ruff.lint.per-file-ignores]
-"tests/*" = ["S101"]
+"tests/*" = ["S101", "S608"]
 "src/pyseekdb/client/client_base.py" = ["S608"]
 
 [tool.ruff.format]
diff --git a/src/pyseekdb/client/admin_client.py b/src/pyseekdb/client/admin_client.py
@@ -88,6 +88,29 @@ def list_databases(
         """
         pass
 
+    @abstractmethod
+    def fork_database(self, source_name: str, destination_name: str, tenant: str = DEFAULT_TENANT) -> Database:
+        """
+        Fork (duplicate) a database to create a new independent copy.
+
+        The destination database is logically equivalent to the source database at the
+        fork snapshot moment, containing all user tables and their data. It can be used
+        as an independent database for subsequent read/write operations.
+
+        Args:
+            source_name: source database name
+            destination_name: destination database name (must not already exist)
+            tenant: tenant name (for OceanBase)
+
+        Returns:
+            Database object for the newly created destination database
+
+        Raises:
+            ValueError: If fork is not supported (requires seekdb >= 1.2.0),
+                        or if the destination database already exists.
+        """
+        pass
+
 
 class _AdminClientProxy(AdminAPI):
     """
@@ -129,6 +152,10 @@ def list_databases(
         """Proxy to server implementation"""
         return self._server.list_databases(limit=limit, offset=offset, tenant=tenant)
 
+    def fork_database(self, source_name: str, destination_name: str, tenant: str = DEFAULT_TENANT) -> Database:
+        """Proxy to server implementation"""
+        return self._server.fork_database(source_name=source_name, destination_name=destination_name, tenant=tenant)
+
     def __repr__(self):
         return f"<AdminClient server={self._server}>"
 
diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py
@@ -577,6 +577,28 @@ def list_databases(
         logger.debug(f"✅ Found {len(databases)} databases{self._database_context(effective_tenant)}")
         return databases
 
+    def fork_database(self, source_name: str, destination_name: str, tenant: str = DEFAULT_TENANT) -> Database:
+        """
+        Fork (duplicate) a database to create a new independent copy.
+
+        Args:
+            source_name: source database name
+            destination_name: destination database name (must not already exist)
+            tenant: tenant name (for OceanBase)
+
+        Returns:
+            Database object for the newly created destination database
+        """
+        if not self._fork_database_enabled():
+            raise ValueError("Fork database is not enabled (requires seekdb >= 1.2.0)")
+
+        effective_tenant = self._database_tenant(tenant)
+        logger.debug(f"Forking database: {source_name} -> {destination_name}{self._database_context(effective_tenant)}")
+        sql = f"FORK DATABASE `{source_name}` TO `{destination_name}`"
+        self._execute(sql)
+        logger.debug(f"✅ Successfully forked database '{source_name}' to '{destination_name}'")
+        return self.get_database(destination_name, tenant=tenant)
+
     # ==================== Collection Management (User-facing) ====================
 
     def _prepare_schema_parameters(  # noqa: C901
@@ -1460,12 +1482,18 @@ def _get_collection_table_name(self, collection_id: str | None, collection_name:
             return CollectionNames.table_name_v2(collection_id)
         return CollectionNames.table_name(collection_name)
 
-    def _fork_enabled(self) -> bool:
+    def _fork_table_enabled(self) -> bool:
         db_type, version = self.detect_db_type_and_version()
         version_110 = Version("1.1.0.0")
         logger.debug(f"db_type: {db_type}, version: {version}")
         return db_type.lower() == "seekdb" and version >= version_110
 
+    def _fork_database_enabled(self) -> bool:
+        db_type, version = self.detect_db_type_and_version()
+        version_120 = Version("1.2.0.0")
+        logger.debug(f"db_type: {db_type}, version: {version}")
+        return db_type.lower() == "seekdb" and version >= version_120
+
     def _get_collection_id(self, collection_name: str) -> str:
         collection_id_query_sql = f"SELECT COLLECTION_ID FROM `{CollectionNames.sdk_collections_table_name()}` WHERE COLLECTION_NAME = '{collection_name}'"
         collection_id_query_result = self._execute(collection_id_query_sql)
@@ -1487,7 +1515,7 @@ def _collection_fork(self, collection: Collection, forked_name: str) -> None:
             collection: Collection to fork
             forked_name: Forked collection name
         """
-        if not self._fork_enabled():
+        if not self._fork_table_enabled():
             raise ValueError("Fork is not enabled for this database")
 
         _validate_collection_name(forked_name)
diff --git a/tests/integration_tests/test_collection_hybrid_search_source_inference.py b/tests/integration_tests/test_collection_hybrid_search_source_inference.py
@@ -70,7 +70,7 @@ def _insert_test_data(self, collection, dimension: int):
     def _get_sql_query(self, client, table_name: str, search_parm: dict) -> str:
         search_parm_json = json.dumps(search_parm, ensure_ascii=False)
         client._server._execute(f"SET @search_parm = '{escape_string(search_parm_json)}'")
-        get_sql_query = f"SELECT DBMS_HYBRID_SEARCH.GET_SQL('{table_name}', @search_parm) as query_sql FROM dual"  # noqa: S608
+        get_sql_query = f"SELECT DBMS_HYBRID_SEARCH.GET_SQL('{table_name}', @search_parm) as query_sql FROM dual"
         rows = client._server._execute(get_sql_query)
         assert rows and rows[0].get("query_sql")
         query_sql = rows[0]["query_sql"]
diff --git a/tests/integration_tests/test_database_fork.py b/tests/integration_tests/test_database_fork.py
@@ -0,0 +1,220 @@
+"""
+Integration tests for AdminClient.fork_database method.
+
+Tests the fork_database functionality against real databases, including:
+- Successful fork operations
+- Empty database fork
+- Fork preserves database attributes (charset, collation)
+- Fork creates independent databases
+- Error handling when destination already exists
+- Error handling when fork is not enabled
+"""
+
+import contextlib
+import logging
+import time
+
+import pytest
+
+logger = logging.getLogger(__name__)
+
+
+class TestDatabaseFork:
+    """Tests for AdminClient.fork_database() using real database connections."""
+
+    def _is_fork_database_enabled(self, admin_client) -> bool:
+        """Check if fork_database is enabled for the given admin client."""
+        try:
+            return admin_client._server._fork_database_enabled()
+        except Exception:
+            logger.exception("Failed to check if fork_database is enabled")
+            return False
+
+    def _unique_name(self, prefix: str) -> str:
+        return f"{prefix}_{int(time.time() * 1000)}"
+
+    def test_fork_database_success(self, admin_client):
+        """
+        Test successful fork_database operation.
+
+        Automatically runs for: embedded, server, oceanbase
+        Skips if fork_database is not enabled.
+        """
+        if not self._is_fork_database_enabled(admin_client):
+            pytest.skip("Fork database is not enabled for this database")
+
+        source_name = self._unique_name("test_forkdb_src")
+        dest_name = self._unique_name("test_forkdb_dst")
+
+        try:
+            admin_client.create_database(source_name)
+
+            forked_db = admin_client.fork_database(source_name, dest_name)
+
+            assert forked_db is not None
+            assert forked_db.name == dest_name
+
+            retrieved_db = admin_client.get_database(dest_name)
+            assert retrieved_db.name == dest_name
+
+        finally:
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(dest_name)
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(source_name)
+
+    def test_fork_empty_database(self, admin_client):
+        """
+        Test fork of an empty database (no tables).
+
+        Automatically runs for: embedded, server, oceanbase
+        Skips if fork_database is not enabled.
+        """
+        if not self._is_fork_database_enabled(admin_client):
+            pytest.skip("Fork database is not enabled for this database")
+
+        source_name = self._unique_name("test_forkdb_empty_src")
+        dest_name = self._unique_name("test_forkdb_empty_dst")
+
+        try:
+            admin_client.create_database(source_name)
+
+            forked_db = admin_client.fork_database(source_name, dest_name)
+            assert forked_db is not None
+            assert forked_db.name == dest_name
+
+        finally:
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(dest_name)
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(source_name)
+
+    def test_fork_database_preserves_attributes(self, admin_client):
+        """
+        Test that fork preserves database attributes (charset, collation).
+
+        Automatically runs for: embedded, server, oceanbase
+        Skips if fork_database is not enabled.
+        """
+        if not self._is_fork_database_enabled(admin_client):
+            pytest.skip("Fork database is not enabled for this database")
+
+        source_name = self._unique_name("test_forkdb_attr_src")
+        dest_name = self._unique_name("test_forkdb_attr_dst")
+
+        try:
+            admin_client.create_database(source_name)
+            source_db = admin_client.get_database(source_name)
+
+            forked_db = admin_client.fork_database(source_name, dest_name)
+
+            assert forked_db.charset == source_db.charset
+            assert forked_db.collation == source_db.collation
+
+        finally:
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(dest_name)
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(source_name)
+
+    def test_fork_database_destination_already_exists(self, admin_client):
+        """
+        Test that forking to an existing database raises an error.
+
+        Automatically runs for: embedded, server, oceanbase
+        Skips if fork_database is not enabled.
+        """
+        if not self._is_fork_database_enabled(admin_client):
+            pytest.skip("Fork database is not enabled for this database")
+
+        source_name = self._unique_name("test_forkdb_dup_src")
+        dest_name = self._unique_name("test_forkdb_dup_dst")
+
+        try:
+            admin_client.create_database(source_name)
+            admin_client.create_database(dest_name)
+
+            with pytest.raises(ValueError):
+                admin_client.fork_database(source_name, dest_name)
+
+        finally:
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(dest_name)
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(source_name)
+
+    def test_fork_database_independent_operations(self, admin_client):
+        """
+        Test that forked database is independent from the source.
+
+        Creates a table in source before fork, then verifies that modifications
+        to the forked database do not affect the source and vice versa.
+
+        Automatically runs for: embedded, server, oceanbase
+        Skips if fork_database is not enabled.
+        """
+        if not self._is_fork_database_enabled(admin_client):
+            pytest.skip("Fork database is not enabled for this database")
+
+        source_name = self._unique_name("test_forkdb_indep_src")
+        dest_name = self._unique_name("test_forkdb_indep_dst")
+
+        try:
+            admin_client.create_database(source_name)
+
+            admin_client._server._execute(f"CREATE TABLE `{source_name}`.`t1` (id INT PRIMARY KEY, val VARCHAR(100))")
+            admin_client._server._execute(f"INSERT INTO `{source_name}`.`t1` VALUES (1, 'original')")
+
+            admin_client.fork_database(source_name, dest_name)
+
+            admin_client._server._execute(f"INSERT INTO `{dest_name}`.`t1` VALUES (2, 'forked_only')")
+
+            source_rows = admin_client._server._execute(f"SELECT COUNT(*) as cnt FROM `{source_name}`.`t1`")
+            dest_rows = admin_client._server._execute(f"SELECT COUNT(*) as cnt FROM `{dest_name}`.`t1`")
+
+            source_count = source_rows[0]["cnt"] if isinstance(source_rows[0], dict) else source_rows[0][0]
+            dest_count = dest_rows[0]["cnt"] if isinstance(dest_rows[0], dict) else dest_rows[0][0]
+
+            assert source_count == 1, f"Source should have 1 row, got {source_count}"
+            assert dest_count == 2, f"Destination should have 2 rows, got {dest_count}"
+
+        finally:
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(dest_name)
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(source_name)
+
+    def test_fork_database_multiple_times(self, admin_client):
+        """
+        Test that the same source database can be forked multiple times.
+
+        Automatically runs for: embedded, server, oceanbase
+        Skips if fork_database is not enabled.
+        """
+        if not self._is_fork_database_enabled(admin_client):
+            pytest.skip("Fork database is not enabled for this database")
+
+        source_name = self._unique_name("test_forkdb_multi_src")
+        dest_name_1 = self._unique_name("test_forkdb_multi_d1")
+        dest_name_2 = self._unique_name("test_forkdb_multi_d2")
+
+        try:
+            admin_client.create_database(source_name)
+
+            forked_1 = admin_client.fork_database(source_name, dest_name_1)
+            forked_2 = admin_client.fork_database(source_name, dest_name_2)
+
+            assert forked_1.name == dest_name_1
+            assert forked_2.name == dest_name_2
+
+        finally:
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(dest_name_2)
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(dest_name_1)
+            with contextlib.suppress(Exception):
+                admin_client.delete_database(source_name)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

Original file line number	Diff line number	Diff line change
`@@ -107,7 +107,7 @@ ignore = [`
`107`	`107`	`]`
`108`	`108`
`109`	`109`	`[tool.ruff.lint.per-file-ignores]`
`110`		`-"tests/*" = ["S101"]`
	`110`	`+"tests/*" = ["S101", "S608"]`
`111`	`111`	`"src/pyseekdb/client/client_base.py" = ["S608"]`
`112`	`112`
`113`	`113`	`[tool.ruff.format]`