Merge pull request #97 from bmeares/chunking

Added `test_pandas.py` to the `pytest` suite.

Author: Mause

.pre-commit-config.yaml

@@ -42,8 +42,6 @@ repos:
     rev: '3.9.2'  # pick a git hash / tag to point to
     hooks:
     - id: flake8
-      additional_dependencies:
-        - flake8-aaa
       args:
         - --jobs=1
         - --extend-ignore=W503,E203,E501

duckdb_engine/__init__.py

@@ -20,6 +20,10 @@ class ConnectionWrapper:
     def cursor(self):
         return self
 
+    def fetchmany(self, size=None):
+        # TODO: remove this once duckdb supports fetchmany natively
+        return self.c.fetch_df_chunk().values.tolist()
+
     def __getattr__(self, name):
         return getattr(self.c, name)
 
@@ -31,8 +35,11 @@ def close(self):
         # duckdb doesn't support 'soft closes'
         pass
 
-    def execute(self, statement, parameters, context):
-        self.c.execute(statement, parameters)
+    def execute(self, statement, parameters=None, context=None):
+        if parameters is None:
+            self.c.execute(statement)
+        else:
+            self.c.execute(statement, parameters)
 
 
 class Dialect(postgres_dialect):

duckdb_engine/tests/test_pandas.py

@@ -0,0 +1,84 @@
+#! /usr/bin/env python3
+# vim:fenc=utf-8
+
+"""
+Test pandas functionality.
+"""
+
+from collections import OrderedDict
+from datetime import datetime
+from itertools import product
+from typing import Optional
+
+import pandas as pd
+from pytest import mark
+from sqlalchemy import create_engine
+
+_possible_args = OrderedDict(
+    {
+        "chunksize": [None, 1, 10000],
+        "if_exists": ["fail", "replace", "append"],
+        "method": [
+            None,
+            "multi",
+        ],  # TODO Implement a callable insert method?
+    }
+)
+
+args = {
+    "to_sql": ["chunksize", "if_exists", "method"],
+    "read_sql": ["chunksize"],
+}
+params = {
+    k: list(
+        product(*(_v for _k, _v in _possible_args.items() if _k in args[k]))  # type: ignore
+    )
+    for k in args
+}
+params_strings = {k: (",".join([str(_k) for _k in args[k]])) for k in args}
+
+sample_df = pd.DataFrame(
+    {
+        "datetime": [datetime.utcnow()],
+        "int": [1],
+        "float": [1.01],
+        "str": ["foo"],
+    }
+)
+
+
+@mark.parametrize(params_strings["to_sql"], params["to_sql"])
+def test_to_sql(
+    chunksize: Optional[int],
+    if_exists: str,
+    method: Optional[str],
+    index: bool = False,
+):
+    eng = create_engine("duckdb:///:memory:")
+    try:
+        sample_df.to_sql(
+            name="foo",
+            con=eng,
+            if_exists=if_exists,
+            chunksize=chunksize,
+            index=index,
+        )
+    except ValueError as e:
+        if if_exists != "fail":
+            raise e
+
+
+@mark.parametrize(params_strings["read_sql"], params["read_sql"])
+def test_read_sql(
+    chunksize: Optional[int],
+):
+    eng = create_engine("duckdb:///:memory:")
+    sample_df.to_sql(name="test_read", con=eng, if_exists="replace")
+    query = "SELECT * FROM test_read"
+    result = pd.read_sql(
+        query,
+        eng,
+        chunksize=chunksize,
+    )
+    chunks = [result] if chunksize is None else [chunk for chunk in result]
+    pd.concat(chunks).reset_index(drop=True)

poetry.lock

@@ -18,6 +18,7 @@ pre-commit = "^2.13.0"
 pdbpp = "^0.10.2"
 mypy = "^0.910"
 hypothesis = "^6.14.1"
+pandas = "^1"
 
 [tool.poetry.plugins."sqlalchemy.dialects"]
 duckdb = "duckdb_engine"