Initial commit

Author: Mause

.gitignore

@@ -0,0 +1,6 @@
+.mypy_cache
+.pytest_cache
+*.log
+dist/
+*.egg-info
+__pycache__/

.pre-commit-config.yaml

@@ -0,0 +1,49 @@
+exclude: .*\.snap
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v3.1.0
+    hooks:
+      - id: check-yaml
+      - id: check-json
+      - id: check-merge-conflict
+      - id: debug-statements
+      - id: check-case-conflict
+      - id: check-toml
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+  - repo: https://github.com/timothycrosley/isort
+    rev: 5.4.2
+    hooks:
+      - id: isort
+        args:
+          - --float-to-top
+  - repo: https://github.com/humitos/mirrors-autoflake
+    rev: v1.3
+    hooks:
+      - id: autoflake
+        args:
+          - --in-place
+          - --remove-all-unused-imports
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v2.7.0
+    hooks:
+      - id: pyupgrade
+        args:
+          - --py3-plus
+  - repo: https://github.com/psf/black
+    rev: b59a524
+    hooks:
+      - id: black
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v0.782
+    hooks:
+     - id: mypy
+  - repo: https://gitlab.com/pycqa/flake8
+    rev: ''  # 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

@@ -0,0 +1,108 @@
+import re
+from dataclasses import dataclass
+from typing import List
+
+import duckdb
+from sqlalchemy.dialects.postgresql import dialect
+
+
+class DBAPI:
+    paramstyle = "qmark"
+
+    class Error(Exception):
+        pass
+
+
+DNE = re.compile(r"Catalog: ([^ ]+) with name ([^ ]+) does not exist!")
+
+
+@dataclass
+class ConnectionWrapper:
+    c: duckdb.DuckDBPyConnection
+
+    def cursor(self):
+        return self
+
+    @property
+    def connection(self):
+        return self
+
+    def close(self):
+        pass
+
+    @property
+    def rowcount(self):
+        return self.c.rowcount
+
+    def execute(self, statement, parameters):
+        try:
+            self.c.execute(statement, parameters)
+            self.c.commit()
+            return self
+
+        except Exception as e:
+            m = DNE.match(e.args[0])
+            if m:
+                raise Exception(m.groups())
+            raise e
+
+    def fetchone(self):
+        return self.c.fetchone()
+
+    def commit(self):
+        self.c.commit()
+
+    description: List[str] = []
+
+    notices: List[str] = []
+
+
+class Dialect(dialect):
+    _has_events = False
+    identifier_preparer = None
+
+    def connect(self, *args, **kwargs):
+        return ConnectionWrapper(duckdb.connect(*args, **kwargs))
+
+    def on_connect(self):
+        pass
+
+    def do_execute(self, cursor, statement, parameters, context):
+        if "FOREIGN" in statement:
+            # doesn't support foreign key constraints
+            statement = statement.strip().splitlines()
+            statement = [line for line in statement if "FOREIGN" not in line]
+            statement[-2] = statement[-2].strip().strip(",")
+            statement = "\n".join(statement)
+
+        super().do_execute(cursor, statement, parameters, context)
+
+        # if "CREATE SEQUENCE" in statement:
+        #     seque = statement.split(" ")[2].strip('"')
+        #     print(cursor.execute(f"SELECT nextval('{seque}');", ()).fetchone())
+
+    def has_table(self, connection, table_name, schema=None):
+        return False
+
+    def has_sequence(self, connection, sequence_name, schema=None):
+        return False
+
+    def has_type(self, connection, type_name, schema=None):
+        return False
+
+    @staticmethod
+    def dbapi():
+        return DBAPI
+
+    def create_connect_args(self, u):
+        return (), {"database": str(u).split("///")[1]}
+
+    def initialize(self, connection):
+        pass
+
+    def do_rollback(self, connection):
+        pass
+
+    @classmethod
+    def get_dialect_cls(cls, u):
+        return cls

duckdb_engine/conftest.py

@@ -0,0 +1,43 @@
+from pytest import fixture
+from sqlalchemy import Column, Integer, Sequence, String, create_engine
+from sqlalchemy.engine.url import registry
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+
+@fixture
+def engine():
+    registry.register("duckdb", "duckdb_engine", "Dialect")
+
+    ewng = create_engine("duckdb:///:memory:")
+    Base.metadata.create_all(ewng)
+    return ewng
+
+
+Base = declarative_base()
+
+
+class FakeModel(Base):  # type: ignore
+    __tablename__ = "fake"
+
+    id = Column(Integer, Sequence("fakemodel_id_sequence"), primary_key=True)
+    name = Column(String)
+
+
+@fixture
+def Session(engine):
+    return sessionmaker(bind=engine)
+
+
+@fixture
+def session(Session):
+    return Session()
+
+
+def test_basic(session):
+    session.add(FakeModel(name="Frank"))
+    session.commit()
+
+    frank = session.query(FakeModel).one()  # act
+
+    assert frank.name == "Frank"