Merge pull request #34375 from def-/pr-scalability-password-sasl

scalability framework: Introduce password + sasl connection tests

Author: def-

ci/nightly/pipeline.template.yml

@@ -215,7 +215,8 @@ steps:
       - id: scalability-benchmark-connection
         label: "Scalability benchmark (connection) against merge base or 'latest'"
         depends_on: build-x86_64
-        timeout_in_minutes: 180
+        # TODO: Lower the timeout again once password/sasl auth is performant: https://github.com/MaterializeInc/database-issues/issues/9956
+        timeout_in_minutes: 1440
         agents:
           # Larger instance is more stable in performance
           queue: hetzner-x86-64-dedi-16cpu-64gb
@@ -232,9 +233,9 @@ steps:
                 - --workload-group-marker
                 - "ConnectionWorkload"
                 - --exponent-base
-                - "2.5"
+                - "2"
                 - --max-concurrency
-                - "2048"
+                - "512"
       - id: parallel-benchmark
         label: "Parallel Benchmark"
         depends_on: build-x86_64

misc/python/materialize/scalability/endpoint/endpoint.py

@@ -7,11 +7,18 @@
 # the Business Source License, use of this software will be governed
 # by the Apache License, Version 2.0.
 
+from enum import Enum, auto
 from typing import Any
 
 import psycopg
 
 
+class ConnectionKind(Enum):
+    Plain = auto()
+    Password = auto()
+    Sasl = auto()
+
+
 class Endpoint:
 
     _version: str | None = None
@@ -20,19 +27,26 @@ def __init__(self, specified_target: str):
         self._specified_target = specified_target
 
     def sql_connection(
-        self, quiet: bool = False
+        self, quiet: bool = False, kind: ConnectionKind = ConnectionKind.Plain
     ) -> psycopg.connection.Connection[tuple[Any, ...]]:
         if not quiet:
             print(f"Connecting to URL: {self.url()}")
 
-        conn = psycopg.connect(self.url())
+        conn = psycopg.connect(self.url(kind))
         conn.autocommit = True
         return conn
 
-    def url(self) -> str:
-        return (
-            f"postgresql://{self.user()}:{self.password()}@{self.host()}:{self.port()}"
+    def url(self, kind: ConnectionKind = ConnectionKind.Plain) -> str:
+        port = (
+            self.port()
+            if kind == ConnectionKind.Plain
+            else (
+                self.port_password()
+                if kind == ConnectionKind.Password
+                else self.port_sasl()
+            )
         )
+        return f"postgresql://{self.user()}:{self.password()}@{self.host()}:{port}/{self.database()}"
 
     def specified_target(self) -> str:
         return self._specified_target
@@ -49,9 +63,18 @@ def user(self) -> str:
     def password(self) -> str:
         raise NotImplementedError
 
+    def database(self) -> str:
+        raise NotImplementedError
+
     def port(self) -> int:
         raise NotImplementedError
 
+    def port_password(self) -> int:
+        raise NotImplementedError
+
+    def port_sasl(self) -> int:
+        raise NotImplementedError
+
     def up(self) -> None:
         raise NotImplementedError
 

misc/python/materialize/scalability/endpoint/endpoints.py

@@ -9,7 +9,7 @@
 
 import psycopg
 
-from materialize import git
+from materialize import MZ_ROOT, git
 from materialize.mzcompose import ADDITIONAL_BENCHMARKING_SYSTEM_PARAMETERS
 from materialize.mzcompose.composition import Composition
 from materialize.mzcompose.services.materialized import Materialized
@@ -60,6 +60,9 @@ def user(self) -> str:
     def password(self) -> str:
         return "postgres"
 
+    def database(self) -> str:
+        return "postgres"
+
     def up(self) -> None:
         self.composition.down(destroy_volumes=True)
         with self.composition.override(Postgres()):
@@ -83,12 +86,6 @@ def host(self) -> str:
     def internal_host(self) -> str:
         return "localhost"
 
-    def user(self) -> str:
-        return "materialize"
-
-    def password(self) -> str:
-        return "materialize"
-
     def internal_port(self) -> int:
         raise NotImplementedError
 
@@ -115,6 +112,15 @@ def port(self) -> int:
     def internal_port(self) -> int:
         return 6877
 
+    def user(self) -> str:
+        return "materialize"
+
+    def password(self) -> str:
+        return "materialize"
+
+    def database(self) -> str:
+        return "materialize"
+
     def up(self) -> None:
         self.lift_limits()
 
@@ -138,17 +144,36 @@ def __init__(
             alternative_image if image != alternative_image else None
         )
         self._port: int | None = None
+        self._port_password: int | None = None
+        self._port_sasl: int | None = None
         self._resolved_target = resolved_target
         self.use_balancerd = use_balancerd
         super().__init__(specified_target)
 
     def resolved_target(self) -> str | None:
         return self._resolved_target
 
+    def user(self) -> str:
+        return "user1"
+
+    def password(self) -> str:
+        return "password"
+
+    def database(self) -> str:
+        return "materialize"
+
     def port(self) -> int:
         assert self._port is not None
         return self._port
 
+    def port_password(self) -> int:
+        assert self._port_password is not None
+        return self._port_password
+
+    def port_sasl(self) -> int:
+        assert self._port_sasl is not None
+        return self._port_sasl
+
     def internal_port(self) -> int:
         return self.composition.port("materialized", 6877)
 
@@ -182,16 +207,24 @@ def up_internal(self) -> None:
                 additional_system_parameter_defaults=ADDITIONAL_BENCHMARKING_SYSTEM_PARAMETERS,
                 external_metadata_store=True,
                 metadata_store="cockroach",
+                listeners_config_path=f"{MZ_ROOT}/src/materialized/ci/listener_configs/testdrive_sasl.json",
             )
         ):
             self.composition.up("materialized")
             self.composition.verify_build_profile()
+            self.composition.sql(
+                "create role \"user1\" with login password 'password';"
+            )
+            self.composition.sql("grant materialize to user1;")
 
             if self.use_balancerd:
                 self.composition.up("balancerd")
                 self._port = self.composition.default_port("balancerd")
             else:
                 self._port = self.composition.default_port("materialized")
+            # TODO: Also make it work through balancerd
+            self._port_password = self.composition.port("materialized", 6880)
+            self._port_sasl = self.composition.port("materialized", 6881)
 
     def __str__(self) -> str:
         return f"MaterializeContainer ({self.image} specified as {self.specified_target()})"

misc/python/materialize/scalability/operation/operations/operations.py

@@ -8,7 +8,7 @@
 # by the Apache License, Version 2.0.
 from psycopg import Connection
 
-from materialize.scalability.endpoint.endpoint import Endpoint
+from materialize.scalability.endpoint.endpoint import ConnectionKind, Endpoint
 from materialize.scalability.operation.operation_data import OperationData
 from materialize.scalability.operation.scalability_operation import (
     Operation,
@@ -201,7 +201,55 @@ def _execute(self, data: OperationData) -> OperationData:
         endpoint: Endpoint = data.get("endpoint")
         schema: Schema = data.get("schema")
 
-        connection = endpoint.sql_connection(quiet=True)
+        connection = endpoint.sql_connection(quiet=True, kind=ConnectionKind.Plain)
+        connection.autocommit = True
+        cursor = connection.cursor()
+
+        # this sets the database schema
+        for connect_sql in schema.connect_sqls():
+            cursor.execute(connect_sql.encode("utf8"))
+
+        data.push("connection", connection)
+        data.push("cursor", cursor)
+        return data
+
+
+class ConnectPassword(Operation):
+    def required_keys(self) -> set[str]:
+        return {"endpoint", "schema"}
+
+    def produced_keys(self) -> set[str]:
+        return {"connection", "cursor"}
+
+    def _execute(self, data: OperationData) -> OperationData:
+        endpoint: Endpoint = data.get("endpoint")
+        schema: Schema = data.get("schema")
+
+        connection = endpoint.sql_connection(quiet=True, kind=ConnectionKind.Password)
+        connection.autocommit = True
+        cursor = connection.cursor()
+
+        # this sets the database schema
+        for connect_sql in schema.connect_sqls():
+            cursor.execute(connect_sql.encode("utf8"))
+
+        data.push("connection", connection)
+        data.push("cursor", cursor)
+        return data
+
+
+class ConnectSasl(Operation):
+    def required_keys(self) -> set[str]:
+        return {"endpoint", "schema"}
+
+    def produced_keys(self) -> set[str]:
+        return {"connection", "cursor"}
+
+    def _execute(self, data: OperationData) -> OperationData:
+        endpoint: Endpoint = data.get("endpoint")
+        schema: Schema = data.get("schema")
+
+        connection = endpoint.sql_connection(quiet=True, kind=ConnectionKind.Sasl)
         connection.autocommit = True
         cursor = connection.cursor()
 

misc/python/materialize/scalability/workload/workloads/connection_workloads.py

@@ -9,6 +9,8 @@
 from materialize.scalability.operation.operation_data import OperationData
 from materialize.scalability.operation.operations.operations import (
     Connect,
+    ConnectPassword,
+    ConnectSasl,
     Disconnect,
     SelectOne,
 )
@@ -28,3 +30,29 @@ def amend_data_before_execution(self, data: OperationData) -> None:
 
     def operations(self) -> list["Operation"]:
         return [OperationChainWithDataExchange([Connect(), SelectOne(), Disconnect()])]
+
+
+class EstablishPasswordConnectionWorkload(WorkloadWithContext, ConnectionWorkload):
+    def amend_data_before_execution(self, data: OperationData) -> None:
+        data.push("endpoint", self.endpoint)
+        data.push("schema", self.schema)
+        data.remove("cursor")
+
+    def operations(self) -> list["Operation"]:
+        return [
+            OperationChainWithDataExchange(
+                [ConnectPassword(), SelectOne(), Disconnect()]
+            )
+        ]
+
+
+class EstablishSaslConnectionWorkload(WorkloadWithContext, ConnectionWorkload):
+    def amend_data_before_execution(self, data: OperationData) -> None:
+        data.push("endpoint", self.endpoint)
+        data.push("schema", self.schema)
+        data.remove("cursor")
+
+    def operations(self) -> list["Operation"]:
+        return [
+            OperationChainWithDataExchange([ConnectSasl(), SelectOne(), Disconnect()])
+        ]

test/scalability/mzcompose.py

@@ -101,7 +101,7 @@
 def workflow_default(c: Composition, parser: WorkflowArgumentParser) -> None:
     parser.add_argument(
         "--target",
-        help="Target for the benchmark: 'HEAD', 'local', 'remote', 'common-ancestor', 'Postgres', or a DockerHub tag",
+        help="Target for the benchmark: 'HEAD', 'local', 'remote', 'common-ancestor', 'postgres', or a DockerHub tag",
         action="append",
         default=[],
     )