Add StarRocks Database Container Support

modules/starrocks/example_basic.py

@@ -0,0 +1,90 @@
+import pandas as pd
+import sqlalchemy
+from sqlalchemy import text
+
+from testcontainers.starrocks import StarRocksContainer
+
+from datetime import datetime
+
+def basic_example():
+    with StarRocksContainer() as starrocks:
+        # Get connection URL
+        connection_url = starrocks.get_connection_url()
+
+        # Create SQLAlchemy engine
+        engine = sqlalchemy.create_engine(connection_url)
+        print("Connected to StarRocks")
+
+        # Create a test table
+        create_table_sql = """
+            CREATE TABLE IF NOT EXISTS weatherdata (
+                DATE DATETIME,
+                NAME STRING,
+                Temperature STRING
+            )
+            ENGINE=olap
+            DUPLICATE KEY(DATE)
+            DISTRIBUTED BY HASH(DATE) BUCKETS 10;
+        """
+
+        with engine.begin() as connection:
+            connection.execute(text(create_table_sql))
+            print("Created test table")
+
+        # Insert test data
+        test_data = test_data = [
+                {"date": datetime(2025, 1, 1, 10, 0), "name": "New York", "temperature": "25°C"},
+                {"date": datetime(2025, 1, 2, 10, 0), "name": "London", "temperature": "15°C"},
+                {"date": datetime(2025, 1, 3, 10, 0), "name": "Tokyo", "temperature": "20°C"},
+                {"date": datetime(2025, 1, 4, 10, 0), "name": "Sydney", "temperature": "30°C"},
+                {"date": datetime(2025, 1, 5, 10, 0), "name": "Paris", "temperature": "18°C"},
+                {"date": datetime(2025, 1, 6, 10, 0), "name": "Berlin", "temperature": "16°C"},
+                {"date": datetime(2025, 1, 7, 10, 0), "name": "Moscow", "temperature": "5°C"},
+                {"date": datetime(2025, 1, 8, 10, 0), "name": "Dubai", "temperature": "35°C"},
+                {"date": datetime(2025, 1, 9, 10, 0), "name": "Singapore", "temperature": "28°C"},
+                {"date": datetime(2025, 1, 10, 10, 0), "name": "Toronto", "temperature": "10°C"},
+            ]
+
+        with engine.begin() as connection:
+            for data in test_data:
+                connection.execute(text("INSERT INTO weatherdata (DATE, NAME, Temperature) VALUES (:date, :name, :temperature)"), data)
+            print("Inserted test data")
+
+        # Query data
+        with engine.connect() as connection:
+            try:
+                result = connection.execute(text("SELECT * FROM weatherdata ORDER BY DATE"))
+                rows = result.fetchall()
+
+                print("\nQuery results:")
+                for row in rows:
+                    print(f"Date: {row[0]}, City: {row[1]}, Temperature: {row[2]}")
+            except sqlalchemy.exc.ProgrammingError as e:
+                print(f"Error querying data: {e}")
+
+        # Execute a more complex query
+        with engine.connect() as connection:
+            try:
+                result = connection.execute(
+                    text("""
+                        SELECT
+                            NAME,
+                            AVG(CAST(REGEXP_REPLACE(Temperature, '[^0-9.]', '') AS FLOAT)) as avg_temperature,
+                            COUNT(*) as count,
+                            MIN(DATE) as first_date,
+                            MAX(DATE) as last_date
+                        FROM weatherdata
+                        GROUP BY NAME
+                        ORDER BY avg_temperature DESC
+                    """)
+                )
+
+                print("\nAggregation results:")
+                for row in result:
+                    print(f"City: {row[0]}, Avg Temperature: {row[1]:.2f}°C, Count: {row[2]}, First: {row[3]}, Last: {row[4]}")
+            except sqlalchemy.exc.ProgrammingError as e:
+                print(f"Error executing query: {e}")
+
+
+if __name__ == "__main__":
+    basic_example()

modules/starrocks/testcontainers/starrocks/__init__.py

@@ -0,0 +1,111 @@
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+import os
+import re
+import string
+from typing import Optional
+
+from testcontainers.core.generic import DbContainer
+from testcontainers.core.utils import raise_for_deprecated_parameter
+from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
+
+_UNSET = object()
+
+
+class StarRocksContainer(DbContainer):
+    """
+    StarRocks database container.
+
+    To get a URL without a driver, pass in :code:`driver=None`.
+
+    Example:
+
+        The example spins up a StarRocks database and connects to it using the :code:`pymysql`
+        driver.
+
+        .. doctest::
+
+            >>> from testcontainers.starrocks import StarRocksContainer
+            >>> import sqlalchemy
+
+            >>> with StarRocksContainer("starrocks/allin1-ubuntu:3.5.3") as starrocks:
+            ...     engine = sqlalchemy.create_engine(starrocks.get_connection_url())
+            ...     with engine.begin() as connection:
+            ...         result = connection.execute(sqlalchemy.text("select version()"))
+            ...         version, = result.fetchone()
+            >>> version
+            'StarRocks...'
+    """
+
+    def __init__(
+        self,
+        image: str = "repo.smartech.ir/starrocks/allin1-ubuntu:3.5.3",
+        dbname: string = "test",
+        port: int = 9030,
+        http_port: int = 8030,
+        rpc_port: int = 9020,
+        edit_log_port:int = 9010,
+        heartbeat_port:int = 9050,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        driver: Optional[str] = "pymysql",
+        **kwargs,
+    ) -> None:
+        raise_for_deprecated_parameter(kwargs, "user", "username")
+        super().__init__(image=image, **kwargs)
+        self.dbname: str = dbname
+        self.username: str = username or os.environ.get("STARROCKS_USER", "root")
+        self.password: str = password or os.environ.get("STARROCKS_PASSWORD", "")
+        self.port = port
+        self.http_port = http_port
+        self.rpc_port = rpc_port
+        self.edit_log_port = edit_log_port
+        self.heartbeat_port = heartbeat_port
+        self.driver = f"+{driver}" if driver else ""
+
+        self.with_exposed_ports(self.port, self.http_port, self.rpc_port, self.edit_log_port, self.heartbeat_port)
+
+    def _configure(self) -> None:
+        self.with_env("STARROCKS_NODE_ROLE", "fe")
+        self.with_env("FE_QUERY_PORT", str(self.port))
+        self.with_env("FE_HTTP_PORT", str(self.http_port))
+        self.with_env("FE_RPC_PORT", str(self.rpc_port))
+        self.with_env("FE_EDIT_LOG_PORT", str(self.edit_log_port))
+        self.with_env("STARROCKS_ROOT_PASSWORD", self.password)
+
+    @wait_container_is_ready()
+    def _connect(self) -> None:
+        wait_for_logs(self,
+                      re.compile(".*Enjoy the journey to StarRocks blazing-fast lake-house engine!.*", flags=re.DOTALL | re.MULTILINE).search,)
+
+        # Create The Default Database If It Does Not Exist
+        command = f'mysql -P 9030 -h 127.0.0.1 -u root -e "CREATE DATABASE IF NOT EXISTS {self.dbname}"'
+        self.exec(command)
+
+    def get_connection_url(self, host: Optional[str] = None, driver: Optional[str] = _UNSET) -> str:
+        """Get a DB connection URL to connect to the StarRocks DB.
+
+        If a driver is set in the constructor (defaults to pymysql), the URL will contain the
+        driver. The optional driver argument to :code:`get_connection_url` overwrites the constructor
+        set value. Pass :code:`driver=None` to get URLs without a driver.
+        """
+        driver_str = "" if driver is None else self.driver if driver is _UNSET else f"+{driver}"
+
+        return super()._create_connection_url(
+            dialect=f"mysql{driver_str}",
+            username=self.username,
+            password=self.password,
+            dbname=self.dbname,
+            host=host,
+            port=self.port,
+        )

modules/starrocks/tests/test_starrocks.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+
+import requests
+
+import pytest
+import sqlalchemy
+
+from testcontainers.starrocks import StarRocksContainer
+
+
+def test_docker_run_starrocks():
+    starrocks_container = StarRocksContainer()
+    with starrocks_container as starrocks:
+        host = starrocks.get_container_host_ip()
+        port = starrocks.get_exposed_port(8030)
+
+        response = requests.get(f'http://{host}:{port}/api/health')
+
+        assert response.status_code == 200
+
+@pytest.mark.inside_docker_check
+def test_docker_run_starrocks_with_sqlalchemy():
+    starrocks_container = StarRocksContainer()
+    with starrocks_container as starrocks:
+        engine = sqlalchemy.create_engine(starrocks.get_connection_url())
+        with engine.begin() as connection:
+            result = connection.execute(sqlalchemy.text("select version()"))
+            for row in result:
+                assert row[0].lower().startswith("8.0.33")