feat: Create Python integration testing (#320)

crystalxyz · web-flow · commit b5e65dff2eff · 2025-04-25T09:18:45.000+09:00
* Add a workflow step for python integration tests and add a duckdb test

* fix some bugs in workflow

* Fix a bug in pytest installation

* Implement flight test

* Replace unittest with pytest and update workflow script

* Add two integration tests and add roapi to git workflow

* Add ODBC integration test

* update path

* Add sqliteodbc to pipeline

* Split python test pipeline into a separate step

* Add sqlite odbc lib

* Update step name

* Fix bugs in mysql and postgres tests
diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
@@ -107,3 +107,34 @@ jobs:
 
       - name: Run tests
         run: make test
+
+  python-integration-test:
+    name: Python Tests
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: dtolnay/rust-toolchain@stable
+
+      - uses: astral-sh/setup-uv@v5
+        with:
+            enable-cache: true
+
+      - name: Install ODBC, Sqlite and Roapi
+        run: |
+          sudo apt-get install -y unixodbc-dev libsqliteodbc
+          sudo apt-get install -y libsqlite3-dev
+          cargo install --locked --git https://github.com/roapi/roapi --branch main --bins roapi
+
+      - name: Build Python package
+        run: |
+          cd python
+          uv sync --dev --no-install-package datafusion
+          uv run --no-project maturin develop --uv
+
+  
+      - name: Run Python tests
+        run: |
+          cd python/python/tests
+          uv run --no-project pytest -v .
diff --git a/python/pyproject.toml b/python/pyproject.toml
@@ -1,6 +1,6 @@
 [build-system]
 requires = ["maturin>=1.5.1,<1.6.0"]
-build-backend = "[maturin]"
+build-backend = "maturin"
 
 [project]
 name = "datafusion_table_providers"
@@ -57,3 +57,15 @@ max-doc-length = 88
 "dev/*" = ["D"]
 "benchmarks/*" = ["D", "F"]
 "docs/*" = ["D"]
+
+[dependency-groups]
+dev = [
+    "maturin>=1.8.1",
+    "numpy>1.25.0",
+    "pytest>=7.4.4",
+    "pytest-asyncio>=0.23.3",
+    "ruff>=0.9.1",
+    "toml>=0.10.2",
+    "pygithub==2.5.0",
+    "pyarrow>=19.0.1"
+]
diff --git a/python/python/tests/test_duckdb.py b/python/python/tests/test_duckdb.py
@@ -0,0 +1,68 @@
+import pytest
+import os
+from datafusion import SessionContext
+from datafusion_table_providers import duckdb
+
+class TestDuckDBIntegration:
+    def setup_method(self):
+        """Set up the test environment"""
+        self.ctx = SessionContext()
+        self.db_path = os.path.join(os.path.dirname(__file__), "..", "..", "..", "core", "examples", "duckdb_example.db")
+        self.pool_readonly = duckdb.DuckDBTableFactory(self.db_path, duckdb.AccessMode.ReadOnly)
+        self.pool_readwrite = duckdb.DuckDBTableFactory(self.db_path)
+
+    def test_get_tables(self):
+        """Test retrieving tables from the database"""
+        tables = self.pool_readonly.tables()
+        assert isinstance(tables, list)
+        assert len(tables) == 2
+        assert tables == ["companies", "projects"]
+        
+    def test_query_companies(self):
+        """Test querying companies table with SQL"""
+        self.ctx.register_table_provider("companies", self.pool_readonly.get_table("companies"))
+        
+        # Run SQL query to select Microsoft row
+        df = self.ctx.sql("SELECT name FROM companies WHERE ticker = 'MSFT'")
+        result = df.collect()
+        
+        # Verify single row returned with name = Microsoft
+        assert len(result) == 1
+        assert str(result[0]['name'][0]) == "Microsoft"
+        
+    def test_complex_query(self):
+        """Test querying companies table with SQL"""
+        self.ctx.register_table_provider("companies", self.pool_readonly.get_table("companies"))
+        self.ctx.register_table_provider("projects", self.pool_readonly.get_table("projects"))
+        
+        # Run SQL query to select Microsoft row
+        df = self.ctx.sql(
+            """SELECT companies.id, companies.name as company_name, projects.name as project_name
+            FROM companies, projects
+            WHERE companies.id = projects.id"""
+        )
+        result = df.collect()
+    
+        assert len(result) == 1
+        assert str(result[0]['company_name'][0]) == "Microsoft"
+        assert str(result[0]['project_name'][0]) == "DataFusion"
+        
+    def test_write_fails(self):
+        """Test that writing fails on read-only mode"""
+        table_name = "companies"
+        self.ctx.register_table_provider(table_name, self.pool_readonly.get_table("companies"))
+        
+        with pytest.raises(Exception):
+            tmp = self.ctx.sql("INSERT INTO companies VALUES (3, 'Test Corp', 'TEST')")
+            tmp.collect() # this will trigger the execution of the query
+
+    def test_write_fails_readwrite(self):
+        """Test that writing fails because it is not supported"""
+        # Insertion fails because duckdb does not implement write operations even when
+        # database is opened in read-write mode.
+        table_name = "companies"
+        self.ctx.register_table_provider(table_name, self.pool_readwrite.get_table("companies"))
+        
+        with pytest.raises(Exception):
+            tmp = self.ctx.sql("INSERT INTO companies VALUES (3, 'Test Corp', 'TEST')")
+            tmp.collect()
diff --git a/python/python/tests/test_flight.py b/python/python/tests/test_flight.py
@@ -0,0 +1,54 @@
+from datafusion import SessionContext
+from datafusion_table_providers import flight
+import pytest
+import subprocess
+import time
+
+
+class TestFlightIntegration:     
+    @classmethod
+    def setup_class(self):
+        """Called once before all test methods in the class"""
+        self.ctx = SessionContext()
+        self.pool = flight.FlightTableFactory()
+        self.process = subprocess.Popen(
+            ["roapi", "-t", "taxi=https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2024-01.parquet"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE
+        )
+        # 20s timeout is required to ensure the server is running and data is loaded
+        # The timeout is determined by empirical testing
+        time.sleep(20)
+
+    @classmethod 
+    def teardown_class(self):
+        """Called once after all test methods in the class"""
+        self.process.kill()
+
+    def test_query_companies(self):
+        """Test querying companies table with SQL"""
+        print("Running test_query_companies")
+        table_name = "taxi_flight_table"
+        self.ctx.register_table_provider(table_name, self.pool.get_table("http://localhost:32010", {
+            "flight.sql.query": "SELECT * FROM taxi"
+        }))
+        df = self.ctx.sql(f"""
+            SELECT "VendorID", COUNT(*) as count, SUM(passenger_count) as passenger_counts, SUM(total_amount) as total_amounts
+                FROM {table_name}
+                GROUP BY "VendorID"
+                ORDER BY COUNT(*) DESC
+            """)
+        result = df.collect()
+        
+        # Verify the results
+        vendor_ids = result[0]['VendorID'].tolist()
+        assert vendor_ids == [2, 1, 6]
+        
+        counts = result[0]['count'].tolist()
+        assert counts == [2234632, 729732, 260]
+        
+        passenger_counts = result[0]['passenger_counts'].tolist()
+        assert passenger_counts == [2971865, 810883, None]
+        
+        total_amounts = result[0]['total_amounts'].tolist()
+        assert total_amounts == pytest.approx([60602721.27, 18841261.98, 12401.03])
diff --git a/python/python/tests/test_mysql.py b/python/python/tests/test_mysql.py
@@ -0,0 +1,82 @@
+import subprocess
+import time
+
+from datafusion import SessionContext
+from datafusion_table_providers import mysql
+
+def run_docker_container():
+    """Run the Docker container with the MySQL image"""
+    result = subprocess.run(
+        ["docker", "run", "--name", "mysql", "-e", "MYSQL_ROOT_PASSWORD=password", "-e", "MYSQL_DATABASE=mysql_db", 
+         "-p", "3306:3306", "-d", "mysql:9.0"],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE
+    )
+    if result.returncode != 0:
+        print(f"Failed to start MySQL container: {result.stderr.decode()}")
+
+def create_table_and_insert_data():
+    """Create a table and insert data into MySQL"""
+    sql_commands = """
+        CREATE TABLE companies (
+            id INT PRIMARY KEY,
+            name VARCHAR(100)
+        );
+        
+        INSERT INTO companies (id, name) VALUES (1, 'Acme Corporation');
+    """
+    
+    # Execute the SQL commands inside the Docker container
+    result = subprocess.run(
+        ["docker", "exec", "-i", "mysql", "mysql", "-uroot", "-ppassword", "mysql_db"],
+        input=sql_commands.encode(),  # Pass SQL commands to stdin
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE
+    )
+    
+    # Check if the SQL execution was successful
+    if result.returncode != 0:
+        print(f"Error executing SQL commands: {result.stderr.decode()}")
+    else:
+        print(f"SQL commands executed successfully:\n{result.stdout.decode()}")
+
+def stop_and_remove_container():
+    """Stop and remove the MySQL container after use"""
+    subprocess.run(["docker", "stop", "mysql"])
+    subprocess.run(["docker", "rm", "mysql"])
+    print("MySQL container stopped and removed.")
+
+
+class TestMySQLIntegration:
+    @classmethod
+    def setup_class(self):
+        run_docker_container()
+        time.sleep(30)
+        create_table_and_insert_data()
+        time.sleep(10)
+        self.ctx = SessionContext()
+        connection_param = {
+            "connection_string": "mysql://root:password@localhost:3306/mysql_db",
+            "sslmode": "disabled"}
+        self.pool = mysql.MySQLTableFactory(connection_param)
+
+    @classmethod
+    def teardown_class(self):
+        stop_and_remove_container()
+
+    def test_get_tables(self):
+        """Test retrieving tables from the database"""
+        tables = self.pool.tables()
+        assert isinstance(tables, list)
+        assert len(tables) == 1
+        assert tables == ["companies"]
+
+    def test_query_companies(self):
+        """Test querying companies table with SQL"""
+        table_name = "companies"
+        self.ctx.register_table_provider(table_name, self.pool.get_table("companies"))
+        query = "SELECT * FROM companies"
+        df = self.ctx.sql(query).collect()
+        assert df is not None
+        name_column = df[0]['name']
+        assert str(name_column[0]) == "Acme Corporation"
diff --git a/python/python/tests/test_odbc.py b/python/python/tests/test_odbc.py
@@ -0,0 +1,49 @@
+import pytest
+import os
+from datafusion import SessionContext
+from datafusion_table_providers import odbc
+
+class TestOdbcIntegration:
+    def setup_method(self):
+        """Set up the test environment"""
+        self.ctx = SessionContext()
+        connection_param: dict = {'connection_string': 'driver=SQLite3;database=../../../core/examples/sqlite_example.db;'}
+        self.pool = odbc.ODBCTableFactory(connection_param)
+        
+    def test_query_companies(self):
+        """Test querying companies table with SQL"""
+        self.ctx.register_table_provider("companies", self.pool.get_table("companies"))
+        
+        # Run SQL query to select Microsoft row
+        df = self.ctx.sql("SELECT name FROM companies WHERE ticker = 'MSFT'")
+        result = df.collect()
+        
+        # Verify single row returned with name = Microsoft
+        assert len(result) == 1
+        assert str(result[0]['name'][0]) == "Microsoft"
+        
+    def test_complex_query(self):
+        """Test querying companies table with SQL"""
+        self.ctx.register_table_provider("companies", self.pool.get_table("companies"))
+        self.ctx.register_table_provider("projects", self.pool.get_table("projects"))
+        
+        # Run SQL query to select Microsoft row
+        df = self.ctx.sql(
+            """SELECT companies.id, companies.name as company_name, projects.name as project_name
+            FROM companies, projects
+            WHERE companies.id = projects.id"""
+        )
+        result = df.collect()
+    
+        assert len(result) == 1
+        assert str(result[0]['company_name'][0]) == "Microsoft"
+        assert str(result[0]['project_name'][0]) == "DataFusion"
+        
+    def test_write_fails(self):
+        """Test that writing fails because it is not supported"""
+        table_name = "companies"
+        self.ctx.register_table_provider(table_name, self.pool.get_table("companies"))
+        
+        with pytest.raises(Exception):
+            tmp = self.ctx.sql("INSERT INTO companies VALUES (3, 'Test Corp', 'TEST')")
+            tmp.collect() # this will trigger the execution of the query
diff --git a/python/python/tests/test_postgres.py b/python/python/tests/test_postgres.py
diff --git a/python/python/tests/test_sqlite.py b/python/python/tests/test_sqlite.py
diff --git a/python/uv.lock b/python/uv.lock
