Merge pull request #26 from juneHQ/vinayak/2024-12-06-add-port-checks

Add CLICKHOUSE_SECURE env var

HISTORY.rst

@@ -2,6 +2,11 @@
 History
 =======
 
+0.2.1 (2024-12-06)
+------------------
+
+* Add `CLICKHOUSE_SECURE` environment variable.
+
 0.2.0 (2024-12-06)
 ------------------
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "houseplant"
-version = "0.2.0"
+version = "0.2.1"
 description = "Database Migrations for ClickHouse."
 readme = "README.md"
 authors = [{name = "June", email = "[email protected]"}]

src/houseplant/__version__.py

@@ -1,4 +1,4 @@
-VERSION = (0, 2, 0)
+VERSION = (0, 2, 1)
 PRERELEASE = None  # alpha, beta or rc
 REVISION = None
 

src/houseplant/clickhouse_client.py

@@ -41,20 +41,35 @@ def __init__(self, database):
 
 
 class ClickHouseClient:
-    def __init__(self, host=None, database=None, port=None):
+    def __init__(
+        self, host=None, port=None, database=None, user=None, password=None, secure=None
+    ):
         self.host = host or os.getenv("CLICKHOUSE_HOST", "localhost")
+        # Parse port from host:port string if present, otherwise use port parameter or default
         if ":" in self.host:
-            self.host, self.port = self.host.split(":")
+            self.host, port_str = self.host.split(":")
+            self.port = int(port_str)
+        else:
+            self.port = int(port or os.getenv("CLICKHOUSE_PORT", 9000))
 
-        self.port = port or os.getenv("CLICKHOUSE_PORT", 9000)
         self.database = database or os.getenv("CLICKHOUSE_DB", "development")
 
+        self.user = user or os.getenv("CLICKHOUSE_USER", "default")
+        self.password = password or os.getenv("CLICKHOUSE_PASSWORD", "")
+
+        # Use SSL port by default if secure
+        self.secure = secure or os.getenv("CLICKHOUSE_SECURE", "n").lower()
+        self.secure = self.secure in ("true", "t", "yes", "y", "1")
+        self.port = 9440 if self.secure else self.port
+
         self.client = Client(
             host=self.host,
             port=self.port,
             database=self.database,
-            user=os.getenv("CLICKHOUSE_USER", "default"),
-            password=os.getenv("CLICKHOUSE_PASSWORD", ""),
+            user=self.user,
+            password=self.password,
+            secure=self.secure,
+            verify=False,
         )
 
         self._cluster = None

tests/test_clickhouse_client.py

@@ -15,6 +15,51 @@ def migrations_table(ch_client):
     return ch_client
 
 
+def test_port_parsing_from_host():
+    """Test port parsing from host string."""
+    from houseplant.clickhouse_client import ClickHouseClient
+
+    client = ClickHouseClient(host="localhost:1234")
+    assert client.host == "localhost"
+    assert client.port == 1234
+
+
+def test_default_port(monkeypatch):
+    """Test default port selection."""
+    from houseplant.clickhouse_client import ClickHouseClient
+
+    # Test default non-secure port
+    client = ClickHouseClient()
+    assert client.port == 9000
+
+    # Test default secure port
+    monkeypatch.setenv("CLICKHOUSE_SECURE", "true")
+    client = ClickHouseClient()
+    assert client.port == 9440
+
+
+def test_port_precedence(monkeypatch):
+    """Test port parameter precedence."""
+    from houseplant.clickhouse_client import ClickHouseClient
+
+    # Host:port should override port parameter and env var
+    monkeypatch.setenv("CLICKHOUSE_PORT", "8000")
+    client = ClickHouseClient(host="localhost:7000", port=6000)
+    assert client.port == 7000
+
+    # Host:port should override env var
+    client = ClickHouseClient(host="localhost:7000")
+    assert client.port == 7000
+
+    # Port parameter should override env var
+    client = ClickHouseClient(host="localhost", port=7000)
+    assert client.port == 7000
+
+    # Env var should be used if no other port specified
+    client = ClickHouseClient(host="localhost")
+    assert client.port == 8000
+
+
 def test_connection_error(monkeypatch):
     """Test connection error handling."""
     monkeypatch.setenv("CLICKHOUSE_HOST", "invalid_host")