feat(tests): add Starrocks TestContainer support

Author: HosseinNejatiJavaremi

.env.test

@@ -39,6 +39,15 @@ POSTGRES_SQLALCHEMY__DATABASE=test_db
 POSTGRES_SQLALCHEMY__USERNAME=test_user
 POSTGRES_SQLALCHEMY__PASSWORD=test_password
 
+# StarRocks Configuration
+STARROCKS__IMAGE=starrocks/allin1-ubuntu:4.0.3
+STARROCKS_SQLALCHEMY__DRIVER_NAME=starrocks
+STARROCKS_SQLALCHEMY__HOST=localhost
+STARROCKS_SQLALCHEMY__PORT=9030
+STARROCKS_SQLALCHEMY__DATABASE=test_db
+STARROCKS_SQLALCHEMY__USERNAME=root
+STARROCKS_SQLALCHEMY__PASSWORD=
+
 # Keycloak Configuration
 KEYCLOAK__IMAGE=quay.io/keycloak/keycloak:26.4.5
 KEYCLOAK__SERVER_URL=http://localhost:8080

archipy/adapters/starrocks/sqlalchemy/session_managers.py

@@ -1,7 +1,10 @@
 from typing import override
 
 from sqlalchemy import URL
+from sqlalchemy.dialects.postgresql import UUID as PostgresUUID
 from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.sql.functions import GenericFunction
+from starrocks.dialect import StarRocksSQLCompiler, StarRocksTypeCompiler
 
 from archipy.adapters.base.sqlalchemy.session_managers import (
     AsyncBaseSQLAlchemySessionManager,
@@ -13,6 +16,52 @@
 from archipy.models.errors import DatabaseConnectionError
 
 
+# Patch the StarRocks type compiler to map UUID to VARCHAR at module level
+# This ensures the patch is applied before any engines are created
+def _patch_starrocks_uuid_mapping() -> None:
+    """Patch the StarRocks type compiler to map UUID to VARCHAR.
+
+    StarRocks doesn't support UUID type natively, so we need to map it to VARCHAR(36).
+    This is patched at module level to ensure it's applied before engine creation.
+    """
+
+    def visit_UUID(self: StarRocksTypeCompiler, type_: PostgresUUID, **kw: object) -> str:  # noqa: ARG001
+        """Map PostgreSQL UUID to VARCHAR(36) for StarRocks."""
+        return "VARCHAR(36)"
+
+    # Patch the type compiler class
+    StarRocksTypeCompiler.visit_UUID = visit_UUID  # ty:ignore[invalid-assignment]
+
+
+def _patch_starrocks_now_function() -> None:
+    """Patch the StarRocks SQL compiler to map func.now() to CURRENT_TIMESTAMP.
+
+    StarRocks doesn't support now() function, it requires CURRENT_TIMESTAMP instead.
+    This is patched at module level to ensure it's applied before engine creation.
+    """
+    # Store original visit_function if it exists
+    original_visit_function = getattr(StarRocksSQLCompiler, "visit_function", None)
+
+    def visit_function(self: StarRocksSQLCompiler, func_: GenericFunction, **kw: object) -> str:
+        """Map func.now() to CURRENT_TIMESTAMP for StarRocks."""
+        # Check if this is func.now()
+        if func_.name == "now":
+            return "CURRENT_TIMESTAMP"
+        # For other functions, use the original handler if it exists
+        if original_visit_function:
+            return original_visit_function(self, func_, **kw)
+        # Fallback to default behavior
+        return f"{func_.name}()"
+
+    # Patch the SQL compiler class
+    StarRocksSQLCompiler.visit_function = visit_function  # ty:ignore[invalid-assignment]
+
+
+# Apply the patches when the module is imported
+_patch_starrocks_uuid_mapping()
+_patch_starrocks_now_function()
+
+
 class StarRocksSQlAlchemySessionManager(BaseSQLAlchemySessionManager[StarRocksSQLAlchemyConfig], metaclass=Singleton):
     """Synchronous SQLAlchemy session manager for StarRocks.
 
@@ -122,18 +171,27 @@ def _get_database_name(self) -> str:
     def _create_url(self, configs: StarRocksSQLAlchemyConfig) -> URL:
         """Create an async StarRocks connection URL.
 
+        For async operations, StarRocks requires an async driver (mysql+aiomysql)
+        instead of the sync driver (mysql+pymysql).
+
         Args:
             configs: StarRocks configuration.
 
         Returns:
-            A SQLAlchemy URL object for StarRocks.
+            A SQLAlchemy URL object for StarRocks with async driver.
 
         Raises:
             DatabaseConnectionError: If there's an error creating the URL.
         """
         try:
+            # For async operations, use mysql+aiomysql driver
+            # If the driver is mysql+pymysql or starrocks, replace with mysql+aiomysql
+            async_driver = configs.DRIVER_NAME
+            if async_driver in ("mysql+pymysql", "starrocks", "mysql"):
+                async_driver = "mysql+aiomysql"
+
             return URL.create(
-                drivername=configs.DRIVER_NAME,
+                drivername=async_driver,
                 username=configs.USERNAME,
                 password=configs.PASSWORD,
                 host=configs.HOST,
@@ -144,3 +202,15 @@ def _create_url(self, configs: StarRocksSQLAlchemyConfig) -> URL:
             raise DatabaseConnectionError(
                 database=self._get_database_name(),
             ) from e
+
+    @override
+    def _get_connect_args(self) -> dict:
+        """Return connection arguments for async StarRocks to ensure proper transaction support.
+
+        StarRocks (using MySQL protocol) requires autocommit to be explicitly disabled
+        to ensure transactions work properly with rollback support.
+
+        Returns:
+            A dictionary with autocommit=False to ensure transaction support.
+        """
+        return {"autocommit": False}

archipy/configs/config_template.py

@@ -11,7 +11,7 @@
 from typing import Literal, Self
 from urllib.parse import urlparse
 
-from pydantic import BaseModel, Field, PostgresDsn, SecretStr, model_validator
+from pydantic import BaseModel, Field, PostgresDsn, SecretStr, field_validator, model_validator
 
 from archipy.models.errors import FailedPreconditionError, InvalidArgumentError
 
@@ -541,9 +541,39 @@ class StarRocksSQLAlchemyConfig(SQLAlchemyConfig):
     """Configuration settings for Starrocks SQLAlchemy ORM.
 
     Extends SQLAlchemyConfig with Starrocks-specific settings.
+
+    Note: StarRocks only supports READ COMMITTED isolation level.
     """
 
+    DRIVER_NAME: str = Field(default="starrocks", description="StarRocks driver name")
     CATALOG: str | None = Field(default=None, description="Starrocks catalog name")
+    ISOLATION_LEVEL: str = Field(
+        default="READ COMMITTED",
+        description="Transaction isolation level (StarRocks only supports READ COMMITTED)",
+    )
+
+    @field_validator("ISOLATION_LEVEL")
+    @classmethod
+    def validate_isolation_level(cls, v: str) -> str:
+        """Validate that isolation level is READ COMMITTED for StarRocks.
+
+        Args:
+            v: The isolation level value to validate.
+
+        Returns:
+            The validated isolation level.
+
+        Raises:
+            ValueError: If the isolation level is not READ COMMITTED.
+        """
+        # Normalize the value (handle case variations and underscores)
+        normalized = v.upper().replace("_", " ").strip()
+        if normalized != "READ COMMITTED":
+            raise ValueError(
+                f"StarRocks only supports READ COMMITTED isolation level. Got: {v}. "
+                "StarRocks does not support other isolation levels like REPEATABLE READ or SERIALIZABLE.",
+            )
+        return "READ COMMITTED"
 
 
 class PrometheusConfig(BaseModel):

features/environment.py

@@ -34,6 +34,7 @@ class TestConfig(BaseConfig):
     MINIO__IMAGE: str
     KEYCLOAK__IMAGE: str
     SCYLLADB__IMAGE: str
+    STARROCKS__IMAGE: str
     TESTCONTAINERS_RYUK_CONTAINER_IMAGE: str | None = None
 
     def __init__(self, **kwargs) -> None:

features/test_containers.py

@@ -19,6 +19,7 @@
     PostgresSQLAlchemyConfig,
     RedisConfig,
     ScyllaDBConfig,
+    StarRocksSQLAlchemyConfig,
 )
 from archipy.helpers.metaclasses.singleton import Singleton
 
@@ -33,6 +34,7 @@
     "needs-keycloak": "keycloak",
     "needs-redis": "redis",
     "needs-scylladb": "scylladb",
+    "needs-starrocks": "starrocks",
 }
 
 
@@ -642,3 +644,83 @@ def stop(self) -> None:
         self.port = None
 
         logger.info("ScyllaDB container stopped")
+
+
+@ContainerManager.register("starrocks")
+class StarRocksTestContainer(metaclass=Singleton, thread_safe=True):
+    """Test container for StarRocks."""
+
+    def __init__(self, config: StarRocksSQLAlchemyConfig | None = None, image: str | None = None) -> None:
+        """Initialize StarRocks test container.
+
+        Args:
+            config (StarRocksSQLAlchemyConfig | None): Configuration for StarRocks. Defaults to None.
+            image (str | None): Docker image to use. Defaults to None (uses STARROCKS__IMAGE from config).
+        """
+        self.name = "starrocks"
+        self.config = config or BaseConfig.global_config().STARROCKS_SQLALCHEMY
+        self.image = image or BaseConfig.global_config().STARROCKS__IMAGE
+        self._is_running: bool = False
+
+        # Container properties
+        self.host: str | None = None
+        self.port: int | None = None
+        self.database: str | None = self.config.DATABASE
+        self.username: str | None = self.config.USERNAME
+        self.password: str | None = self.config.PASSWORD
+
+        # Set up the container
+        self._container = DockerContainer(self.image)
+        # Expose ports: 9030 (MySQL protocol), 8030 (FE HTTP), 8040 (BE HTTP)
+        # These will be mapped to random available host ports automatically
+        self._container.with_exposed_ports(9030, 8030, 8040)
+
+    def start(self) -> DockerContainer:
+        """Start the StarRocks container.
+
+        Returns:
+            DockerContainer: The running container instance.
+        """
+        if self._is_running:
+            return self._container
+
+        # Start the container
+        self._container.start()
+
+        # Wait for StarRocks to be ready
+        # StarRocks logs "cluster initialization DONE!" when the cluster is fully initialized
+        # This appears after "FE service query port:9030 is alive!" and indicates readiness
+        wait_for_logs(self._container, "Enjoy the journey to StarRocks blazing-fast lake-house engine!", timeout=120)
+
+        self._is_running = True
+
+        # Get dynamic host and random port (container port 9030 mapped to random host port)
+        self.host = self._container.get_container_host_ip()
+        # get_exposed_port returns the random host port that maps to container port 9030
+        self.port = int(self._container.get_exposed_port(9030))
+
+        # Update global config with actual container endpoint
+        global_config = BaseConfig.global_config()
+        global_config.STARROCKS_SQLALCHEMY.HOST = self.host
+        global_config.STARROCKS_SQLALCHEMY.PORT = self.port
+
+        logger.info("StarRocks container started on %s:%s", self.host, self.port)
+
+        return self._container
+
+    def stop(self) -> None:
+        """Stop the StarRocks container."""
+        if not self._is_running:
+            return
+
+        if self._container:
+            self._container.stop()
+
+        self._container = None
+        self._is_running = False
+
+        # Reset container properties
+        self.host = None
+        self.port = None
+
+        logger.info("StarRocks container stopped")

features/test_entity.py

@@ -29,6 +29,13 @@ class TestEntity(UpdatableDeletableEntity):
 
     __tablename__ = "test_entities"
 
+    __table_args__ = {
+        "comment": "Test entity table",
+        "starrocks_primary_key": "test_uuid",
+        "starrocks_distributed_by": "HASH(test_uuid) BUCKETS 10",
+        "starrocks_properties": {"replication_num": "1"},
+    }
+
     test_uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
     pk_uuid = Synonym("test_uuid")
 
@@ -85,6 +92,13 @@ class TestManagerEntity(UpdatableManagerEntity):
 
     __tablename__ = "test_manager_entities"
 
+    __table_args__ = {
+        "comment": "Test manager entity table",
+        "starrocks_primary_key": "test_uuid",
+        "starrocks_distributed_by": "HASH(test_uuid) BUCKETS 10",
+        "starrocks_properties": {"replication_num": "1"},
+    }
+
     test_uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
     pk_uuid = Synonym("test_uuid")
 
@@ -143,6 +157,13 @@ class TestAdminEntity(UpdatableAdminEntity):
 
     __tablename__ = "test_admin_entities"
 
+    __table_args__ = {
+        "comment": "Test admin entity table",
+        "starrocks_primary_key": "test_uuid",
+        "starrocks_distributed_by": "HASH(test_uuid) BUCKETS 10",
+        "starrocks_properties": {"replication_num": "1"},
+    }
+
     test_uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
     pk_uuid = Synonym("test_uuid")
 
@@ -202,6 +223,13 @@ class RelatedTestEntity(BaseEntity):
 
     __tablename__ = "related_test_entities"
 
+    __table_args__ = {
+        "comment": "Related test entity table",
+        "starrocks_primary_key": "related_uuid",
+        "starrocks_distributed_by": "HASH(related_uuid) BUCKETS 10",
+        "starrocks_properties": {"replication_num": "1"},
+    }
+
     related_uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
     pk_uuid = Synonym("related_uuid")
 

pyproject.toml

@@ -360,7 +360,7 @@ invalid-type-variable-constraints = "error"
 missing-argument = "error"
 missing-typed-dict-key = "error"
 no-matching-overload = "error"
-non-subscriptable = "error"
+not-subscriptable = "error"
 not-iterable = "error"
 override-of-final-method = "error"
 parameter-already-assigned = "error"