sqlbenchdag 0.1.0__py3-none-any.whl

sql_benchmarks/resources/typedb_client.py

@@ -0,0 +1,370 @@
+import time
+import polars as pl
+from typing import Dict, Any, Optional
+from typedb.driver import TypeDB, Credentials, DriverOptions, TransactionType
+
+from .base_schema_client import SchemaFirstClient
+
+
+class TypeDBClient(SchemaFirstClient):
+    """
+    Concrete SchemaFirstClient for TypeDB.
+
+    Implements the three loading steps using the TypeDB Python driver:
+        _prepare_store  → drop / create TypeDB database
+        _define_schema  → SCHEMA transaction with a TypeQL ``define`` block
+        _insert_batch   → WRITE transaction with a TypeQL ``insert`` block
+
+    Query execution uses a READ transaction and forces full result
+    materialisation via ``list(answer.as_concept_rows())``.
+    """
+
+    def __init__(self, address: str, db_name: str):
+        self.address = address
+        self.db_name = db_name
+
+    # ------------------------------------------------------------------
+    # SchemaFirstClient — abstract step implementations
+    # ------------------------------------------------------------------
+
+    def _prepare_store(self, entity_type: str) -> None:
+        """Drop the TypeDB database if it exists, then recreate it."""
+        with self._driver() as driver:
+            if driver.databases.contains(self.db_name):
+                driver.databases.get(self.db_name).delete()
+            driver.databases.create(self.db_name)
+
+    def _define_schema(self, entity_type: str, df: pl.DataFrame) -> None:
+        """Run a SCHEMA transaction containing a TypeQL ``define`` block."""
+        schema_tql = self._build_schema(entity_type, df)
+        with self._driver() as driver:
+            with driver.transaction(self.db_name, TransactionType.SCHEMA) as tx:
+                tx.query(schema_tql).resolve()
+                tx.commit()
+
+    def _insert_batch(self, entity_type: str, rows: list) -> None:
+        """Run a WRITE transaction containing a TypeQL ``insert`` block."""
+        insert_tql = self._build_insert(entity_type, rows)
+        with self._driver() as driver:
+            with driver.transaction(self.db_name, TransactionType.WRITE) as tx:
+                tx.query(insert_tql).resolve()
+                tx.commit()
+
+    # ------------------------------------------------------------------
+    # SchemaFirstClient — run_query
+    # ------------------------------------------------------------------
+
+    def run_query(self, typeql: str, scenario_params: Dict[str, Any]) -> Optional[float]:
+        """
+        Execute a TypeQL query and return wall-clock duration including
+        full result consumption (consistent with DuckDB / Actian timing).
+
+        Returns ``None`` if the TypeDB server crashes during evaluation (e.g.
+        stack overflow from deep recursive function calls on Zipf hub graphs).
+        Crashes are logged but do not propagate so that the benchmark
+        infrastructure can record a null result and continue.
+        """
+        from typedb.common.exception import TypeDBDriverException  # local import — optional dep
+        try:
+            with self._driver() as driver:
+                with driver.transaction(self.db_name, TransactionType.READ) as tx:
+                    start = time.time()
+                    answer = tx.query(typeql).resolve()
+                    list(answer.as_concept_rows())  # force full materialisation
+                    end = time.time()
+            return end - start
+        except TypeDBDriverException as exc:
+            print(
+                f"[TypeDBClient] Query failed on '{self.db_name}' "
+                f"(TypeDB server error — likely stack overflow or OOM "
+                f"during recursive evaluation): {exc}"
+            )
+            return None
+        except Exception as exc:
+            print(f"[TypeDBClient] Query failed on '{self.db_name}' (unexpected): {exc}")
+            return None
+
+    # ------------------------------------------------------------------
+    # TypeQL builders (TypeDB-specific, not part of base contract)
+    # ------------------------------------------------------------------
+
+    def _build_schema(self, entity_type: str, df: pl.DataFrame) -> str:
+        """
+        Produce a TypeQL ``define`` block from a Polars schema.
+
+        Example output:
+            define
+              entity skewed_data_ssd_small,
+                owns id @key,
+                owns selectivity_code,
+                owns payload;
+              attribute id, value integer;
+              attribute selectivity_code, value string;
+              attribute payload, value string;
+        """
+        attr_defs = []
+        owns_clauses = []
+
+        for col_name, dtype in df.schema.items():
+            typeql_value = self._polars_to_typeql_value(dtype)
+            attr_defs.append(f"  attribute {col_name}, value {typeql_value};")
+            if col_name == "id":
+                owns_clauses.append(f"    owns {col_name} @key")
+            else:
+                owns_clauses.append(f"    owns {col_name}")
+
+        owns_str = ",\n".join(owns_clauses) + ";"
+        attr_str = "\n".join(attr_defs)
+
+        return f"define\n  entity {entity_type},\n{owns_str}\n{attr_str}"
+
+    def _build_insert(self, entity_type: str, rows: list) -> str:
+        """
+        Produce a single TypeQL ``insert`` block for a batch of rows.
+
+        Each variable is uniquely named ($x0, $x1, …) to avoid clashes
+        within the same transaction.
+        """
+        stmts = []
+        for i, row in enumerate(rows):
+            parts = [f"$x{i} isa {entity_type}"]
+            for col, val in row.items():
+                if val is None:
+                    continue
+                if isinstance(val, bool):
+                    parts.append(f"has {col} {str(val).lower()}")
+                elif isinstance(val, str):
+                    escaped = val.replace("\\", "\\\\").replace('"', '\\"')
+                    parts.append(f'has {col} "{escaped}"')
+                else:
+                    parts.append(f"has {col} {val}")
+            stmts.append(", ".join(parts) + ";")
+
+        return "insert\n  " + "\n  ".join(stmts)
+
+    @staticmethod
+    def _polars_to_typeql_value(dtype) -> str:
+        """Map a Polars dtype to the appropriate TypeQL value type."""
+        if dtype in (
+            pl.Int8, pl.Int16, pl.Int32, pl.Int64,
+            pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
+        ):
+            return "integer"
+        if dtype in (pl.Float32, pl.Float64):
+            return "double"
+        if dtype == pl.Boolean:
+            return "boolean"
+        if dtype == pl.Date or isinstance(dtype, pl.Datetime):
+            return "datetime"
+        return "string"
+
+    # ------------------------------------------------------------------
+    # Multi-table loading (hypergraph / relation support)
+    # ------------------------------------------------------------------
+
+    def initialize_db(self) -> None:
+        """
+        Drop and recreate the TypeDB database.
+
+        Called once per partition at the start of a bulk-load phase so that
+        multiple entity types can be loaded into the same database without
+        each one wiping the previous data (the original ``_prepare_store``
+        behaviour).
+        """
+        with self._driver() as driver:
+            if driver.databases.contains(self.db_name):
+                driver.databases.get(self.db_name).delete()
+            driver.databases.create(self.db_name)
+
+    def load_entity(self, filepath: str, entity_type: str) -> None:
+        """
+        Define schema for an entity type and load its data without dropping
+        the database first.
+
+        Use this instead of ``bulk_load()`` when multiple entity types must
+        coexist in the same database (e.g. a supply-chain hypergraph where
+        supplier, buyer, product and supply_contract all live in one DB).
+        """
+        df = pl.read_parquet(filepath)
+        self._define_schema(entity_type, df)
+        rows = df.to_dicts()
+        for i in range(0, len(rows), self.BATCH_SIZE):
+            self._insert_batch(entity_type, rows[i : i + self.BATCH_SIZE])
+        print(f"[TypeDBClient] Loaded {len(rows)} rows into '{entity_type}'")
+
+    def bulk_load_relation(
+        self,
+        filepath: str,
+        relation_type: str,
+        role_map: Dict[str, list],
+        attributes: list,
+    ) -> None:
+        """
+        Define a TypeDB n-ary relation schema and load data into it.
+
+        All referenced entity types must already exist in the database
+        (i.e. ``load_entity()`` must have been called for each role player).
+
+        Args:
+            filepath:      Path to the Parquet file whose rows become relations.
+            relation_type: TypeDB relation type name (e.g. ``supply_contract_small``).
+            role_map:      Mapping of DataFrame column name → [entity_type, role_name].
+                           Example: ``{"supplier_id": ["supplier_small", "supplier_role"]}``.
+            attributes:    Column names that become relation-owned attributes
+                           (everything not in ``role_map`` and not ``id``).
+        """
+        df = pl.read_parquet(filepath)
+        schema_tql = self._build_relation_schema(relation_type, role_map, attributes, df)
+        with self._driver() as driver:
+            with driver.transaction(self.db_name, TransactionType.SCHEMA) as tx:
+                tx.query(schema_tql).resolve()
+                tx.commit()
+
+        rows = df.to_dicts()
+        for i in range(0, len(rows), self.BATCH_SIZE):
+            insert_tql = self._build_relation_insert(
+                relation_type, role_map, attributes, rows[i : i + self.BATCH_SIZE]
+            )
+            with self._driver() as driver:
+                with driver.transaction(self.db_name, TransactionType.WRITE) as tx:
+                    tx.query(insert_tql).resolve()
+                    tx.commit()
+
+        print(f"[TypeDBClient] Loaded {len(rows)} relations into '{relation_type}'")
+
+    def _build_relation_schema(
+        self,
+        relation_type: str,
+        role_map: Dict[str, list],
+        attributes: list,
+        df: pl.DataFrame,
+    ) -> str:
+        """
+        Produce a TypeQL ``define`` block that augments an existing schema with:
+          - attribute types for the relation's own attributes
+          - the relation type declaration (roles + owned attributes)
+          - ``plays`` declarations for each role-player entity type
+
+        Example output (supply_contract_small with 3 roles)::
+
+            define
+              attribute volume, value integer;
+              attribute price_per_unit, value double;
+              relation supply_contract_small,
+                relates supplier_role,
+                relates buyer_role,
+                relates product_role,
+                owns volume,
+                owns price_per_unit;
+              supplier_small plays supply_contract_small:supplier_role;
+              buyer_small plays supply_contract_small:buyer_role;
+              product_small plays supply_contract_small:product_role;
+        """
+        # Attribute type definitions (relation-owned columns only)
+        attr_defs = []
+        for col in attributes:
+            typeql_val = self._polars_to_typeql_value(df.schema[col])
+            attr_defs.append(f"  attribute {col}, value {typeql_val};")
+
+        # Unique role names (preserve declaration order)
+        role_names = list(dict.fromkeys(entry[1] for entry in role_map.values()))
+
+        # Relation type body
+        relates_clauses = [f"    relates {r}" for r in role_names]
+        owns_clauses = [f"    owns {col}" for col in attributes]
+        relation_body = ",\n".join(relates_clauses + owns_clauses) + ";"
+        relation_def = f"  relation {relation_type},\n{relation_body}"
+
+        # Plays declarations (one per role player)
+        plays_lines = [
+            f"  {entry[0]} plays {relation_type}:{entry[1]};"
+            for entry in role_map.values()
+        ]
+
+        return "define\n" + "\n".join(attr_defs + [relation_def] + plays_lines)
+
+    def _build_relation_insert(
+        self,
+        relation_type: str,
+        role_map: Dict[str, list],
+        attributes: list,
+        rows: list,
+    ) -> str:
+        """
+        Produce a single TypeQL ``match … insert`` block for a batch of
+        relation rows.
+
+        Each row uses unique variable names (``$e{i}_{col}``) so that all
+        match constraints are satisfied in a single solution, and the insert
+        phase creates one relation per row in one round trip.
+
+        Example (one row, supply_contract)::
+
+            match
+              $e0_supplier_id isa supplier_small, has id 3;
+              $e0_buyer_id isa buyer_small, has id 7;
+              $e0_product_id isa product_small, has id 2;
+            insert
+              (supplier_role: $e0_supplier_id, buyer_role: $e0_buyer_id,
+               product_role: $e0_product_id) isa supply_contract_small,
+               has volume 450, has price_per_unit 12.5;
+        """
+        match_parts: list[str] = []
+        insert_parts: list[str] = []
+
+        for i, row in enumerate(rows):
+            # Match: bind each role-player entity by its integer key
+            for col, (entity_type, _) in role_map.items():
+                var = f"$e{i}_{col}"
+                match_parts.append(f"  {var} isa {entity_type}, has id {row[col]};")
+
+            # Insert: build the relation with roles and attribute values
+            roles_str = ", ".join(
+                f"{entry[1]}: $e{i}_{col}" for col, entry in role_map.items()
+            )
+            attr_parts: list[str] = []
+            for col in attributes:
+                val = row.get(col)
+                if val is None:
+                    continue
+                if isinstance(val, bool):
+                    attr_parts.append(f"has {col} {str(val).lower()}")
+                elif isinstance(val, str):
+                    escaped = val.replace("\\", "\\\\").replace('"', '\\"')
+                    attr_parts.append(f'has {col} "{escaped}"')
+                else:
+                    attr_parts.append(f"has {col} {val}")
+
+            attr_str = (", " + ", ".join(attr_parts)) if attr_parts else ""
+            insert_parts.append(f"  ({roles_str}) isa {relation_type}{attr_str};")
+
+        return "match\n" + "\n".join(match_parts) + "\ninsert\n" + "\n".join(insert_parts)
+
+    def apply_inference_schema(self, tql: str) -> None:
+        """
+        Apply a TypeQL ``define`` block as a SCHEMA transaction.
+
+        Used to augment an existing database with inferred relation types and
+        inference rules *after* the base data has been loaded.  Calling this
+        multiple times with non-conflicting definitions is safe (TypeDB is
+        idempotent for ``define`` of already-existing types/rules).
+
+        Args:
+            tql: A complete TypeQL ``define`` block, e.g. the transitive
+                 reachability rule generated by
+                 ``TypeDBEngine._build_transitive_inference_schema()``.
+        """
+        with self._driver() as driver:
+            with driver.transaction(self.db_name, TransactionType.SCHEMA) as tx:
+                tx.query(tql).resolve()
+                tx.commit()
+        print(f"[TypeDBClient] Applied inference schema to '{self.db_name}'")
+
+    # ------------------------------------------------------------------
+    # Driver factory
+    # ------------------------------------------------------------------
+
+    def _driver(self):
+        creds = Credentials("admin", "password")
+        opts = DriverOptions(is_tls_enabled=False)
+        return TypeDB.driver(self.address, creds, opts)

sql_benchmarks/resources/typedb_engine.py

@@ -0,0 +1,303 @@
+import time
+import socket
+import docker
+from docker.errors import NotFound, APIError
+from dagster import ConfigurableResource
+from typing import Dict, Any, Optional
+from pydantic import ConfigDict
+from .typedb_client import TypeDBClient
+from ..utils.system import thrash_os_cache
+
+# ---------------------------------------------------------------------------
+# Module-level partition tracker
+#
+# Records which TypeDB database names have already been initialised (drop +
+# create) in the current process.  This lets bulk_load be called once per
+# table while only wiping the database on the very first call for a given
+# partition, so multiple entity types can coexist in the same database.
+#
+# The set is naturally cleared between benchmark runs because execute_run.py
+# starts a fresh Python process each time.
+# ---------------------------------------------------------------------------
+_INITIALIZED_PARTITIONS: set = set()
+
+
+class TypeDBEngine(ConfigurableResource):
+    """
+    Dagster Resource for TypeDB (local Docker).
+
+    Manages the Docker container lifecycle and provides cold-cache enforcement
+    by killing and recreating the container (data is preserved in a named
+    Docker volume, identical to the PostgresEngine approach).
+
+    Multi-table support
+    -------------------
+    When an experiment loads several tables into the same partition (e.g. a
+    supply-chain hypergraph with supplier, buyer, product, supply_contract),
+    this engine initialises the TypeDB database *once* on the first
+    ``bulk_load`` call and then adds each subsequent table additively.
+
+    Relation tables are declared via ``relation_configs``::
+
+        relation_configs = {
+            "supply_contract": {          # base table name (without partition key)
+                "roles": {
+                    "supplier_id": ["supplier", "supplier_role"],
+                    "buyer_id":    ["buyer",    "buyer_role"],
+                    "product_id":  ["product",  "product_role"],
+                },
+                "attributes": ["volume", "price_per_unit"],
+            }
+        }
+
+    Entity tables (not in ``relation_configs``) are loaded via
+    ``TypeDBClient.load_entity()``.  Relation tables are loaded via
+    ``TypeDBClient.bulk_load_relation()`` using a ``match … insert`` pattern.
+    """
+
+    # --- CONFIGURATION ---
+    address: str = "127.0.0.1:1729"
+    container_name: str = "bench_typedb"
+    docker_image: str = "typedb/typedb:latest"
+    relation_configs: Dict[str, Any] = {}
+
+    model_config = ConfigDict(extra="forbid")
+
+    # ------------------------------------------------------------------
+    # IBenchmarkEngine interface
+    # ------------------------------------------------------------------
+
+    def get_engine_name(self) -> str:
+        return "typedb"
+
+    def clear_cache(self, settings: dict = None):
+        """
+        Enforces a cold-cache run:
+        1. Flood host RAM to evict the OS page cache.
+        2. Kill and recreate the TypeDB container (clears DBMS buffer pool).
+           Data survives because it lives in a named Docker volume.
+        """
+        thrash_os_cache()
+        self._restart_container()
+        self._wait_for_ready()
+
+    def run_query(self, sql: str, partition_key: str, engine_params: Dict[str, Any] = None) -> Optional[float]:
+        """engine_params is the 'typedb' namespace — accepted but not yet applied."""
+        self.clear_cache()
+        client = self._get_client(partition_key)
+        return client.run_query(sql, {})
+
+    def bulk_load(self, filepath: str, table_name: str, partition_key: str) -> None:
+        """
+        Load one table into TypeDB for the given partition.
+
+        On the *first* call for a partition the database is dropped and
+        recreated (clean slate).  Subsequent calls for the same partition add
+        to the existing database so that multiple entity types and relations
+        can coexist.
+
+        Entity tables use ``TypeDBClient.load_entity()``.
+        Relation tables (declared in ``relation_configs``) use
+        ``TypeDBClient.bulk_load_relation()`` with a match-insert pattern.
+        """
+        self._ensure_container()
+        self._wait_for_ready()
+        client = self._get_client(partition_key)
+        db_key = self._db_name(partition_key)
+
+        # Initialise DB once per partition per process
+        if db_key not in _INITIALIZED_PARTITIONS:
+            client.initialize_db()
+            _INITIALIZED_PARTITIONS.add(db_key)
+
+        # Dispatch: relation vs entity
+        base_name = self._base_table_name(table_name, partition_key)
+        if base_name in self.relation_configs:
+            config = self.relation_configs[base_name]
+            role_map = self._resolve_role_map(config["roles"], partition_key)
+            client.bulk_load_relation(
+                filepath,
+                table_name,
+                role_map,
+                config.get("attributes", []),
+            )
+            # After loading the relation, optionally define an inference rule
+            # so that subsequent read queries can traverse the graph transitively.
+            if config.get("inference") == "transitive":
+                role_values = list(role_map.values())
+                entity_type = role_values[0][0]  # same for both roles (self-referential)
+                from_role   = role_values[0][1]
+                to_role     = role_values[1][1]
+                tql = self._build_transitive_inference_schema(
+                    table_name, entity_type, from_role, to_role
+                )
+                client.apply_inference_schema(tql)
+        else:
+            client.load_entity(filepath, table_name)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _get_client(self, partition_key: str) -> TypeDBClient:
+        db_name = self._db_name(partition_key)
+        return TypeDBClient(address=self.address, db_name=db_name)
+
+    @staticmethod
+    def _db_name(partition_key: str) -> str:
+        """Convert a partition key into a valid TypeDB database name."""
+        return f"bench_{partition_key}".replace("-", "_").replace("~", "_")
+
+    @staticmethod
+    def _base_table_name(table_name: str, partition_key: str) -> str:
+        """
+        Strip the partition-key suffix to obtain the base table name.
+
+        Example: ``supply_contract_small`` with ``small`` → ``supply_contract``.
+        """
+        suffix = f"_{partition_key}"
+        return table_name[: -len(suffix)] if table_name.endswith(suffix) else table_name
+
+    @staticmethod
+    def _build_transitive_inference_schema(
+        relation_type: str,
+        entity_type: str,
+        from_role: str,
+        to_role: str,
+    ) -> str:
+        """
+        Generate a TypeQL ``define`` block containing a recursive stream function
+        that computes the transitive closure over ``relation_type``.
+
+        TypeDB 3.x does not support inference rules (``rule`` keyword) — functions
+        are the supported mechanism.  The generated ``fun reachable`` uses the
+        ``{ base } or { recursive }`` pattern inside its ``match`` body; TypeDB
+        evaluates it lazily with tabling so cycles terminate automatically.
+
+        Args:
+            relation_type: Fully qualified relation name, e.g.
+                           ``supplies_small_small``.
+            entity_type:   Entity type that plays both roles (self-referential),
+                           e.g. ``company_small_small``.
+            from_role:     Name of the "source" role, e.g. ``seller_role``.
+            to_role:       Name of the "destination" role, e.g. ``buyer_role``.
+
+        Returns:
+            TypeQL define block as a string ready to pass to
+            ``TypeDBClient.apply_inference_schema()``.
+
+        Example output::
+
+            define
+              fun reachable($from: company_small_small) -> { company_small_small }:
+              match
+                { (seller_role: $from, buyer_role: $to) isa supplies_small_small; } or
+                { (seller_role: $from, buyer_role: $via) isa supplies_small_small;
+                  let $to in reachable($via); };
+              return { $to };
+        """
+        return (
+            "define\n"
+            f"  fun reachable($from: {entity_type}) -> {{ {entity_type} }}:\n"
+            "  match\n"
+            f"    {{ ({from_role}: $from, {to_role}: $to) isa {relation_type}; }} or\n"
+            f"    {{ ({from_role}: $from, {to_role}: $via) isa {relation_type};\n"
+            "      let $to in reachable($via); };\n"
+            "  return { $to };"
+        )
+
+    @staticmethod
+    def _resolve_role_map(role_configs: dict, partition_key: str) -> dict:
+        """
+        Expand base entity type names in role configs to full runtime names
+        by appending the partition key.
+
+        Input:  ``{"supplier_id": ["supplier", "supplier_role"]}``
+        Output: ``{"supplier_id": ["supplier_small", "supplier_role"]}``
+        """
+        return {
+            col: [f"{entry[0]}_{partition_key}", entry[1]]
+            for col, entry in role_configs.items()
+        }
+
+    # ------------------------------------------------------------------
+    # Docker lifecycle
+    # ------------------------------------------------------------------
+
+    def _ensure_container(self):
+        """Start the container if it is not already running."""
+        client = docker.from_env()
+        try:
+            container = client.containers.get(self.container_name)
+            if container.status != "running":
+                container.start()
+        except NotFound:
+            self._start_container(client)
+
+    def _restart_container(self):
+        """Kill any existing container and start a fresh one."""
+        client = docker.from_env()
+        self._kill_container(client)
+        self._start_container(client)
+
+    def _kill_container(self, client=None):
+        client = client or docker.from_env()
+        try:
+            container = client.containers.get(self.container_name)
+            container.remove(force=True)
+        except NotFound:
+            pass
+        except APIError as e:
+            raise RuntimeError(f"Failed to remove TypeDB container '{self.container_name}': {e}")
+
+    def _start_container(self, client=None):
+        client = client or docker.from_env()
+
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                client.containers.run(
+                    image=self.docker_image,
+                    name=self.container_name,
+                    detach=True,
+                    ports={"1729/tcp": 1729},
+                    volumes={
+                        "typedb_bench_data": {
+                            "bind": "/var/lib/typedb/data",
+                            "mode": "rw",
+                        }
+                    },
+                )
+                return
+            except APIError as e:
+                if ("port is already allocated" in str(e) or "Conflict" in str(e)) and attempt < max_retries - 1:
+                    time.sleep(2)
+                    self._kill_container(client)
+                    continue
+                raise RuntimeError(f"TypeDB container failed to start: {e}")
+
+    def _wait_for_ready(self, timeout: int = 60):
+        """
+        Poll port 1729 until TypeDB accepts TCP connections, then attempt
+        a lightweight driver ping to confirm the server is fully up.
+        """
+        start = time.time()
+        while time.time() - start < timeout:
+            if self._port_open():
+                try:
+                    from typedb.driver import TypeDB, Credentials, DriverOptions  # noqa: PLC0415
+                    creds = Credentials("admin", "password")
+                    opts = DriverOptions(is_tls_enabled=False)
+                    with TypeDB.driver(self.address, creds, opts) as driver:
+                        driver.databases.all()
+                    return
+                except Exception:
+                    pass
+            time.sleep(1)
+
+        raise TimeoutError(f"TypeDB at {self.address} did not become ready within {timeout}s")
+
+    def _port_open(self) -> bool:
+        host, port = self.address.split(":")
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            return s.connect_ex((host, int(port))) == 0

sql_benchmarks/scripts/sql/acid_test/duckdb/test.sql

@@ -0,0 +1 @@
+SELECT 1;

sql_benchmarks/scripts/sql/acid_test/postgres/test.sql

@@ -0,0 +1 @@
+SELECT 1;