SingleStore Integration: Add SingleStore store functionality with tests

neomatamune · neomatamune · commit 77de267b718b · 2025-03-24T14:12:10.000+01:00
diff --git a/cosmotech/coal/singlestore/__init__.py b/cosmotech/coal/singlestore/__init__.py
@@ -0,0 +1,17 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+"""
+SingleStore integration module.
+
+This module provides functions for interacting with SingleStore databases.
+"""
+
+# Re-export functions from the store module
+from cosmotech.coal.singlestore.store import (
+    load_from_singlestore,
+)
diff --git a/cosmotech/coal/singlestore/store.py b/cosmotech/coal/singlestore/store.py
@@ -0,0 +1,100 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+"""
+SingleStore store operations module.
+
+This module provides functions for interacting with SingleStore databases
+for store operations.
+"""
+
+import pathlib
+import time
+import csv
+import singlestoredb as s2
+
+from cosmotech.coal.store.csv import store_csv_file
+from cosmotech.coal.store.store import Store
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def _get_data(table_name: str, output_directory: str, cursor) -> None:
+    """
+    Run a SQL query to fetch all data from a table and write them in csv files.
+
+    Args:
+        table_name: Table name
+        output_directory: Output directory
+        cursor: SingleStore cursor
+    """
+    start_time = time.perf_counter()
+    cursor.execute(f"SELECT * FROM {table_name}")
+    rows = cursor.fetchall()
+    end_time = time.perf_counter()
+    LOGGER.info(
+        T("coal.logs.database.rows_fetched").format(
+            table=table_name, count=len(rows), time=round(end_time - start_time, 2)
+        )
+    )
+    with open(f"{output_directory}/{table_name}.csv", "w", newline="") as csv_stock:
+        w = csv.DictWriter(csv_stock, rows[0].keys())
+        w.writeheader()
+        w.writerows(rows)
+
+
+def load_from_singlestore(
+    single_store_host: str,
+    single_store_port: int,
+    single_store_db: str,
+    single_store_user: str,
+    single_store_password: str,
+    store_folder: str,
+    single_store_tables: str = "",
+) -> None:
+    """
+    Load data from SingleStore and store it in the Store.
+
+    Args:
+        single_store_host: SingleStore host
+        single_store_port: SingleStore port
+        single_store_db: SingleStore database name
+        single_store_user: SingleStore username
+        single_store_password: SingleStore password
+        store_folder: Store folder
+        single_store_tables: Comma-separated list of tables to load
+    """
+    single_store_working_dir = store_folder + "/singlestore"
+    if not pathlib.Path.exists(single_store_working_dir):
+        pathlib.Path.mkdir(single_store_working_dir)
+
+    start_full = time.perf_counter()
+
+    conn = s2.connect(
+        host=single_store_host,
+        port=single_store_port,
+        database=single_store_db,
+        user=single_store_user,
+        password=single_store_password,
+        results_type="dicts",
+    )
+    with conn:
+        with conn.cursor() as cur:
+            if single_store_tables == "":
+                cur.execute("SHOW TABLES")
+                table_names = cur.fetchall()
+            else:
+                table_names = single_store_tables.split(",")
+            LOGGER.info(T("coal.logs.database.tables_to_fetch").format(tables=table_names))
+            for name in table_names:
+                _get_data(name, single_store_working_dir, cur)
+    end_full = time.perf_counter()
+    LOGGER.info(T("coal.logs.database.full_dataset").format(time=round(end_full - start_full, 2)))
+
+    for csv_path in pathlib.Path(single_store_working_dir).glob("*.csv"):
+        LOGGER.info(T("coal.logs.storage.found_file").format(file=csv_path.name))
+        store_csv_file(csv_path.name[:-4], csv_path, store=Store(False, store_folder))
diff --git a/tests/unit/coal/test_singlestore/test_singlestore_store.py b/tests/unit/coal/test_singlestore/test_singlestore_store.py
@@ -0,0 +1,204 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+import csv
+import os
+import pathlib
+import tempfile
+from unittest.mock import MagicMock, patch, mock_open, call
+
+import pytest
+import singlestoredb as s2
+
+from cosmotech.coal.singlestore.store import load_from_singlestore, _get_data
+
+
+class TestStoreFunctions:
+    """Tests for top-level functions in the store module."""
+
+    @patch("cosmotech.coal.singlestore.store.Store")
+    @patch("cosmotech.coal.singlestore.store.store_csv_file")
+    @patch("cosmotech.coal.singlestore.store.s2.connect")
+    @patch("pathlib.Path.mkdir")
+    @patch("pathlib.Path.exists")
+    @patch("pathlib.Path.glob")
+    @patch("os.path.exists")
+    def test_load_from_singlestore(
+        self, mock_os_exists, mock_glob, mock_exists, mock_mkdir, mock_connect, mock_store_csv_file, mock_store
+    ):
+        """Test the load_from_singlestore function."""
+        # Arrange
+        single_store_host = "localhost"
+        single_store_port = 3306
+        single_store_db = "test_db"
+        single_store_user = "user"
+        single_store_password = "password"
+        store_folder = "/tmp/store"
+        single_store_tables = "table1,table2"
+
+        # Mock Path.exists to return False so that mkdir is called
+        mock_exists.return_value = False
+
+        # Mock os.path.exists to return True for the CSV files
+        mock_os_exists.return_value = True
+
+        # Mock the connection and cursor
+        mock_conn = MagicMock()
+        mock_cursor = MagicMock()
+        mock_conn.__enter__.return_value = mock_conn
+        mock_conn.cursor.return_value.__enter__.return_value = mock_cursor
+        mock_connect.return_value = mock_conn
+
+        # Mock the cursor.fetchall to return table names
+        mock_cursor.fetchall.return_value = [{"TABLE_NAME": "table1"}, {"TABLE_NAME": "table2"}]
+
+        # Mock Path.glob to return paths to CSV files
+        mock_csv_path1 = MagicMock()
+        mock_csv_path1.name = "table1.csv"
+        mock_csv_path2 = MagicMock()
+        mock_csv_path2.name = "table2.csv"
+        mock_glob.return_value = [mock_csv_path1, mock_csv_path2]
+
+        # Mock Store instance
+        mock_store_instance = MagicMock()
+        mock_store.return_value = mock_store_instance
+
+        # Act
+        with patch("builtins.open", mock_open()) as mock_file:
+            load_from_singlestore(
+                single_store_host=single_store_host,
+                single_store_port=single_store_port,
+                single_store_db=single_store_db,
+                single_store_user=single_store_user,
+                single_store_password=single_store_password,
+                store_folder=store_folder,
+                single_store_tables=single_store_tables,
+            )
+
+        # Assert
+        # Verify that the directory was created
+        mock_exists.assert_called_once_with("/tmp/store/singlestore")
+        mock_mkdir.assert_called_once()
+
+        # Verify that the connection was established with the correct parameters
+        mock_connect.assert_called_once_with(
+            host=single_store_host,
+            port=single_store_port,
+            database=single_store_db,
+            user=single_store_user,
+            password=single_store_password,
+            results_type="dicts",
+        )
+
+        # Verify that _get_data was called for each table
+        assert mock_cursor.execute.call_count >= 2  # At least one call per table
+
+        # Verify that store_csv_file was called for each CSV file
+        assert mock_store_csv_file.call_count == 2
+        mock_store_csv_file.assert_any_call("table1", mock_csv_path1, store=mock_store_instance)
+        mock_store_csv_file.assert_any_call("table2", mock_csv_path2, store=mock_store_instance)
+
+        # Verify that Store was initialized correctly
+        mock_store.assert_called_with(False, store_folder)
+
+    def test_get_data(self):
+        """Test the _get_data function."""
+        # Arrange
+        table_name = "test_table"
+        output_directory = "/tmp/output"
+        cursor = MagicMock()
+
+        # Mock data returned from the database
+        mock_rows = [{"id": 1, "name": "John"}, {"id": 2, "name": "Jane"}]
+        cursor.fetchall.return_value = mock_rows
+
+        # Act
+        with patch("builtins.open", mock_open()) as mock_file:
+            with patch("csv.DictWriter") as mock_dict_writer:
+                # Mock the DictWriter
+                mock_writer = MagicMock()
+                mock_dict_writer.return_value = mock_writer
+
+                # Ensure the directory exists
+                with patch("os.path.exists") as mock_exists:
+                    mock_exists.return_value = True
+
+                    _get_data(table_name, output_directory, cursor)
+
+        # Assert
+        # Verify that the SQL query was executed
+        cursor.execute.assert_called_once_with("SELECT * FROM test_table")
+
+        # Verify that fetchall was called
+        cursor.fetchall.assert_called_once()
+
+    @patch("cosmotech.coal.singlestore.store.Store")
+    @patch("cosmotech.coal.singlestore.store.store_csv_file")
+    @patch("cosmotech.coal.singlestore.store.s2.connect")
+    @patch("pathlib.Path.mkdir")
+    @patch("pathlib.Path.exists")
+    @patch("pathlib.Path.glob")
+    @patch("os.path.exists")
+    def test_load_from_singlestore_no_tables_specified(
+        self, mock_os_exists, mock_glob, mock_exists, mock_mkdir, mock_connect, mock_store_csv_file, mock_store
+    ):
+        """Test the load_from_singlestore function when no tables are specified."""
+        # Arrange
+        single_store_host = "localhost"
+        single_store_port = 3306
+        single_store_db = "test_db"
+        single_store_user = "user"
+        single_store_password = "password"
+        store_folder = "/tmp/store"
+
+        # Mock Path.exists to return False so that mkdir is called
+        mock_exists.return_value = False
+
+        # Mock os.path.exists to return True for the CSV files
+        mock_os_exists.return_value = True
+
+        # Mock the connection and cursor
+        mock_conn = MagicMock()
+        mock_cursor = MagicMock()
+        mock_conn.__enter__.return_value = mock_conn
+        mock_conn.cursor.return_value.__enter__.return_value = mock_cursor
+        mock_connect.return_value = mock_conn
+
+        # Mock the cursor.fetchall to return table names for SHOW TABLES
+        mock_cursor.fetchall.return_value = [
+            {"TABLE_NAME": "table1"},
+            {"TABLE_NAME": "table2"},
+            {"TABLE_NAME": "table3"},
+        ]
+
+        # Mock Path.glob to return paths to CSV files
+        mock_csv_paths = [MagicMock() for _ in range(3)]
+        for i, path in enumerate(mock_csv_paths):
+            path.name = f"table{i+1}.csv"
+        mock_glob.return_value = mock_csv_paths
+
+        # Mock Store instance
+        mock_store_instance = MagicMock()
+        mock_store.return_value = mock_store_instance
+
+        # Act
+        with patch("builtins.open", mock_open()) as mock_file:
+            load_from_singlestore(
+                single_store_host=single_store_host,
+                single_store_port=single_store_port,
+                single_store_db=single_store_db,
+                single_store_user=single_store_user,
+                single_store_password=single_store_password,
+                store_folder=store_folder,
+            )
+
+        # Assert
+        # Verify that the cursor was used to execute a query
+        assert mock_cursor.execute.called
+
+        # Verify that store_csv_file was called for each CSV file
+        assert mock_store_csv_file.call_count == 3
