diff --git a/examples/workers/src/entry.py b/examples/workers/src/entry.py index 4edf443..5ae0ca3 100644 --- a/examples/workers/src/entry.py +++ b/examples/workers/src/entry.py @@ -41,6 +41,8 @@ async def fetch(self, request, env): return await self.test_sqlalchemy_select() elif path == "sqlalchemy-crud": return await self.test_sqlalchemy_crud() + elif path == "sqlalchemy-composite-pk": + return await self.test_sqlalchemy_composite_pk() elif path == "sqlalchemy-reflect": return await self.test_sqlalchemy_reflect() # Empty result set tests (GitHub issue #4) @@ -188,6 +190,7 @@ async def index(self): "/parameterized": "Test parameterized queries", "/sqlalchemy-select": "Test SQLAlchemy Core SELECT (no raw SQL)", "/sqlalchemy-crud": "Test SQLAlchemy Core CRUD (no raw SQL)", + "/sqlalchemy-composite-pk": "Test SQLAlchemy composite primary key DDL", "/sqlalchemy-reflect": "Test SQLAlchemy table reflection", "/empty-result": "Test empty result set description (issue #4)", "/empty-result-sqlalchemy": "Test SQLAlchemy empty result (issue #4)", @@ -548,6 +551,76 @@ async def test_sqlalchemy_crud(self): status=500, ) + async def test_sqlalchemy_composite_pk(self): + """Test SQLAlchemy composite primary key DDL against D1 binding.""" + from sqlalchemy import Column, MetaData, String, Table, select + + table_name = f"test_composite_pk_{uuid.uuid4().hex[:8]}" + + try: + engine = self.get_engine() + metadata = MetaData() + + test_table = Table( + table_name, + metadata, + Column("tenant_id", String, primary_key=True), + Column("record_key", String, primary_key=True), + Column("value", String), + ) + + metadata.create_all(engine) + + with engine.connect() as conn: + conn.execute( + test_table.insert().values( + tenant_id="tenant_a", + record_key="label_a", + value="value_a", + ) + ) + conn.commit() + + result = conn.execute( + select( + test_table.c.tenant_id, + test_table.c.record_key, + test_table.c.value, + ) + ) + row = result.fetchone() + columns = list(result.keys()) + + metadata.drop_all(engine) + + success = row is not None and tuple(row) == ( + "tenant_a", + "label_a", + "value_a", + ) + + return Response.json( + { + "test": "sqlalchemy_composite_pk", + "success": success, + "table_name": table_name, + "columns": columns, + "row": list(row) if row is not None else None, + } + ) + except Exception as e: + try: + engine = self.get_engine() + metadata = MetaData() + test_table = Table(table_name, metadata) + metadata.drop_all(engine) + except Exception: + pass + return Response.json( + {"test": "sqlalchemy_composite_pk", "success": False, "error": str(e)}, + status=500, + ) + async def test_sqlalchemy_reflect(self): """Test SQLAlchemy table reflection. diff --git a/src/sqlalchemy_cloudflare_d1/compiler.py b/src/sqlalchemy_cloudflare_d1/compiler.py index 68e22c6..dcdf3a0 100644 --- a/src/sqlalchemy_cloudflare_d1/compiler.py +++ b/src/sqlalchemy_cloudflare_d1/compiler.py @@ -157,10 +157,16 @@ def get_column_specification( if not column.nullable: colspec += " NOT NULL" - # Only add PRIMARY KEY inline if first_pk=True - # This prevents duplicate PRIMARY KEY constraints - # (one inline, one as separate constraint) - if column.primary_key and first_pk: + # Only add PRIMARY KEY inline for a SINGLE-column primary key. A + # composite primary key is emitted at the table level (see + # create_table_constraints), so inlining it here as well would produce + # two PRIMARY KEY clauses, which D1 rejects ("more than one primary + # key": SQLITE_ERROR). + if ( + column.primary_key + and first_pk + and len(column.table.primary_key.columns) == 1 + ): colspec += " PRIMARY KEY" if column.computed is not None: diff --git a/tests/integration/test_restapi_integration.py b/tests/integration/test_restapi_integration.py index 89b27d1..819da68 100644 --- a/tests/integration/test_restapi_integration.py +++ b/tests/integration/test_restapi_integration.py @@ -416,6 +416,47 @@ def test_engine_create_table_with_metadata(self, d1_engine, test_table_name): # Clean up metadata.drop_all(d1_engine) + def test_engine_create_table_with_composite_primary_key( + self, d1_engine, test_table_name + ): + """Test creating and using a composite primary key table on D1.""" + metadata = MetaData() + + test_table = Table( + test_table_name, + metadata, + Column("tenant_id", String, primary_key=True), + Column("record_key", String, primary_key=True), + Column("value", String), + ) + + metadata.create_all(d1_engine) + + try: + with d1_engine.connect() as conn: + conn.execute( + test_table.insert().values( + tenant_id="tenant_a", + record_key="label_a", + value="value_a", + ) + ) + conn.commit() + + result = conn.execute( + select( + test_table.c.tenant_id, + test_table.c.record_key, + test_table.c.value, + ) + ) + row = result.fetchone() + + assert row is not None + assert tuple(row) == ("tenant_a", "label_a", "value_a") + finally: + metadata.drop_all(d1_engine) + def test_engine_insert_and_select(self, d1_engine, test_table_name): """Test INSERT and SELECT using SQLAlchemy ORM-style.""" metadata = MetaData() diff --git a/tests/integration/test_worker_integration.py b/tests/integration/test_worker_integration.py index ebaa417..390b175 100644 --- a/tests/integration/test_worker_integration.py +++ b/tests/integration/test_worker_integration.py @@ -181,6 +181,19 @@ def test_sqlalchemy_crud(self, dev_server): assert "name" in data["columns"] assert "value" in data["columns"] + def test_sqlalchemy_composite_primary_key(self, dev_server): + """Test SQLAlchemy creates a composite primary key table in Workers.""" + port = dev_server + response = requests.get(f"http://localhost:{port}/sqlalchemy-composite-pk") + + assert response.status_code == 200 + data = response.json() + + assert data["test"] == "sqlalchemy_composite_pk" + assert data["success"] is True + assert data["row"] == ["tenant_a", "label_a", "value_a"] + assert data["columns"] == ["tenant_id", "record_key", "value"] + def test_sqlalchemy_reflect(self, dev_server): """Test SQLAlchemy table reflection with autoload_with.""" port = dev_server diff --git a/tests/unit/test_dialect.py b/tests/unit/test_dialect.py index 66104a7..0b2125b 100644 --- a/tests/unit/test_dialect.py +++ b/tests/unit/test_dialect.py @@ -176,6 +176,35 @@ def test_create_table_no_autoincrement_on_text(): assert "AUTOINCREMENT" not in sql.upper(), f"Unexpected AUTOINCREMENT in: {sql}" +def test_create_table_composite_primary_key(): + """Test that a composite PRIMARY KEY emits exactly one constraint. + + Regression: D1 rejects "more than one primary key" (SQLITE_ERROR). A + multi-column primary key must be rendered only as a table-level + PRIMARY KEY (...), never also inlined on the first column. + """ + from sqlalchemy import Column, MetaData, String, Table + from sqlalchemy.schema import CreateTable + + dialect = CloudflareD1Dialect() + metadata = MetaData() + + test_table = Table( + "test_table", + metadata, + Column("tenant_id", String, primary_key=True), + Column("key", String, primary_key=True), + Column("value", String), + ) + + sql = str(CreateTable(test_table).compile(dialect=dialect)) + + pk_count = sql.upper().count("PRIMARY KEY") + assert pk_count == 1, f"Expected 1 PRIMARY KEY, found {pk_count} in: {sql}" + # The single constraint must be the table-level composite form. + assert "PRIMARY KEY (" in sql.upper(), f"composite PK not table-level in: {sql}" + + def test_async_dialect_import(): """Test that the async dialect can be imported.""" from sqlalchemy_cloudflare_d1 import CloudflareD1Dialect_async