Skip to content

Commit 6a934c0

Browse files
author
Peng Ren
committed
Fixed bugs
1 parent a09fb7e commit 6a934c0

6 files changed

Lines changed: 527 additions & 142 deletions

File tree

‎pymongosql/__init__.py‎

Lines changed: 5 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -42,115 +42,13 @@ def connect(*args, **kwargs) -> "Connection":
4242
return Connection(*args, **kwargs)
4343

4444

45-
# SQLAlchemy integration
45+
# SQLAlchemy integration (optional)
46+
# For SQLAlchemy functionality, import from pymongosql.sqlalchemy_mongodb:
47+
# from pymongosql.sqlalchemy_mongodb import create_engine_url, create_engine_from_mongodb_uri
4648
try:
47-
# Import and register the dialect automatically
48-
from .sqlalchemy_compat import (
49-
get_sqlalchemy_version,
50-
is_sqlalchemy_2x,
51-
)
52-
53-
# Make compatibility info easily accessible
54-
__sqlalchemy_version__ = get_sqlalchemy_version()
55-
__supports_sqlalchemy__ = __sqlalchemy_version__ is not None
56-
__supports_sqlalchemy_2x__ = is_sqlalchemy_2x()
57-
49+
from .sqlalchemy_mongodb import __sqlalchemy_version__, __supports_sqlalchemy_2x__, __supports_sqlalchemy__
5850
except ImportError:
59-
# SQLAlchemy not available
51+
# SQLAlchemy integration not available
6052
__sqlalchemy_version__ = None
6153
__supports_sqlalchemy__ = False
6254
__supports_sqlalchemy_2x__ = False
63-
64-
65-
def create_engine_url(host: str = "localhost", port: int = 27017, database: str = "test", **kwargs) -> str:
66-
"""Create a SQLAlchemy engine URL for PyMongoSQL.
67-
68-
Args:
69-
host: MongoDB host
70-
port: MongoDB port
71-
database: Database name
72-
**kwargs: Additional connection parameters
73-
74-
Returns:
75-
SQLAlchemy URL string (uses mongodb:// format)
76-
77-
Example:
78-
>>> url = create_engine_url("localhost", 27017, "mydb")
79-
>>> engine = sqlalchemy.create_engine(url)
80-
"""
81-
params = []
82-
for key, value in kwargs.items():
83-
params.append(f"{key}={value}")
84-
85-
param_str = "&".join(params)
86-
if param_str:
87-
param_str = "?" + param_str
88-
89-
return f"mongodb://{host}:{port}/{database}{param_str}"
90-
91-
92-
def create_mongodb_url(mongodb_uri: str) -> str:
93-
"""Convert a standard MongoDB URI to work with PyMongoSQL SQLAlchemy dialect.
94-
95-
Args:
96-
mongodb_uri: Standard MongoDB connection string
97-
(e.g., 'mongodb://localhost:27017/mydb' or 'mongodb+srv://...')
98-
99-
Returns:
100-
SQLAlchemy-compatible URL for PyMongoSQL
101-
102-
Example:
103-
>>> url = create_mongodb_url("mongodb://user:pass@localhost:27017/mydb")
104-
>>> engine = sqlalchemy.create_engine(url)
105-
"""
106-
# Return the MongoDB URI as-is since the dialect now handles MongoDB URLs directly
107-
return mongodb_uri
108-
109-
110-
def create_engine_from_mongodb_uri(mongodb_uri: str, **engine_kwargs):
111-
"""Create a SQLAlchemy engine from any MongoDB connection string.
112-
113-
This function handles both mongodb:// and mongodb+srv:// URIs properly.
114-
Use this instead of create_engine() directly for mongodb+srv URIs.
115-
116-
Args:
117-
mongodb_uri: Standard MongoDB connection string
118-
**engine_kwargs: Additional arguments passed to create_engine
119-
120-
Returns:
121-
SQLAlchemy Engine object
122-
123-
Example:
124-
>>> # For SRV records (Atlas/Cloud)
125-
>>> engine = create_engine_from_mongodb_uri("mongodb+srv://user:pass@cluster.net/db")
126-
>>> # For standard MongoDB
127-
>>> engine = create_engine_from_mongodb_uri("mongodb://localhost:27017/mydb")
128-
"""
129-
try:
130-
from sqlalchemy import create_engine
131-
132-
if mongodb_uri.startswith("mongodb+srv://"):
133-
# For MongoDB+SRV, convert to standard mongodb:// for SQLAlchemy compatibility
134-
# SQLAlchemy doesn't handle the + character in scheme names well
135-
converted_uri = mongodb_uri.replace("mongodb+srv://", "mongodb://")
136-
137-
# Create engine with converted URI
138-
engine = create_engine(converted_uri, **engine_kwargs)
139-
140-
def custom_create_connect_args(url):
141-
# Use original SRV URI for actual MongoDB connection
142-
opts = {"host": mongodb_uri}
143-
return [], opts
144-
145-
engine.dialect.create_connect_args = custom_create_connect_args
146-
return engine
147-
else:
148-
# Standard mongodb:// URLs work fine with SQLAlchemy
149-
return create_engine(mongodb_uri, **engine_kwargs)
150-
151-
except ImportError:
152-
raise ImportError("SQLAlchemy is required for engine creation")
153-
154-
155-
# Note: PyMongoSQL now uses standard MongoDB connection strings directly
156-
# No need for PyMongoSQL-specific URL format

‎pymongosql/sqlalchemy_mongodb/__init__.py‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
# -*- coding: utf-8 -*-
2+
"""
3+
SQLAlchemy MongoDB dialect and integration for PyMongoSQL.
4+
5+
This package provides SQLAlchemy integration including:
6+
- MongoDB-specific dialect
7+
- Version compatibility utilities
8+
- Engine creation helpers
9+
- MongoDB URI handling
10+
"""
11+
12+
# SQLAlchemy integration
13+
try:
14+
# Import and register the dialect automatically
15+
from .sqlalchemy_compat import (
16+
get_sqlalchemy_version,
17+
is_sqlalchemy_2x,
18+
)
19+
20+
# Make compatibility info easily accessible
21+
__sqlalchemy_version__ = get_sqlalchemy_version()
22+
__supports_sqlalchemy__ = __sqlalchemy_version__ is not None
23+
__supports_sqlalchemy_2x__ = is_sqlalchemy_2x()
24+
25+
except ImportError:
26+
# SQLAlchemy not available
27+
__sqlalchemy_version__ = None
28+
__supports_sqlalchemy__ = False
29+
__supports_sqlalchemy_2x__ = False
30+
31+
32+
def create_engine_url(host: str = "localhost", port: int = 27017, database: str = "test", **kwargs) -> str:
33+
"""Create a SQLAlchemy engine URL for PyMongoSQL.
34+
35+
Args:
36+
host: MongoDB host
37+
port: MongoDB port
38+
database: Database name
39+
**kwargs: Additional connection parameters
40+
41+
Returns:
42+
SQLAlchemy URL string (uses mongodb:// format)
43+
44+
Example:
45+
>>> url = create_engine_url("localhost", 27017, "mydb")
46+
>>> engine = sqlalchemy.create_engine(url)
47+
"""
48+
params = []
49+
for key, value in kwargs.items():
50+
params.append(f"{key}={value}")
51+
52+
param_str = "&".join(params)
53+
if param_str:
54+
param_str = "?" + param_str
55+
56+
return f"mongodb://{host}:{port}/{database}{param_str}"
57+
58+
59+
def create_mongodb_url(mongodb_uri: str) -> str:
60+
"""Convert a standard MongoDB URI to work with PyMongoSQL SQLAlchemy dialect.
61+
62+
Args:
63+
mongodb_uri: Standard MongoDB connection string
64+
(e.g., 'mongodb://localhost:27017/mydb' or 'mongodb+srv://...')
65+
66+
Returns:
67+
SQLAlchemy-compatible URL for PyMongoSQL
68+
69+
Example:
70+
>>> url = create_mongodb_url("mongodb://user:pass@localhost:27017/mydb")
71+
>>> engine = sqlalchemy.create_engine(url)
72+
"""
73+
# Return the MongoDB URI as-is since the dialect now handles MongoDB URLs directly
74+
return mongodb_uri
75+
76+
77+
def create_engine_from_mongodb_uri(mongodb_uri: str, **engine_kwargs):
78+
"""Create a SQLAlchemy engine from any MongoDB connection string.
79+
80+
This function handles both mongodb:// and mongodb+srv:// URIs properly.
81+
Use this instead of create_engine() directly for mongodb+srv URIs.
82+
83+
Args:
84+
mongodb_uri: Standard MongoDB connection string
85+
**engine_kwargs: Additional arguments passed to create_engine
86+
87+
Returns:
88+
SQLAlchemy Engine object
89+
90+
Example:
91+
>>> # For SRV records (Atlas/Cloud)
92+
>>> engine = create_engine_from_mongodb_uri("mongodb+srv://user:pass@cluster.net/db")
93+
>>> # For standard MongoDB
94+
>>> engine = create_engine_from_mongodb_uri("mongodb://localhost:27017/mydb")
95+
"""
96+
try:
97+
from sqlalchemy import create_engine
98+
99+
if mongodb_uri.startswith("mongodb+srv://"):
100+
# For MongoDB+SRV, convert to standard mongodb:// for SQLAlchemy compatibility
101+
# SQLAlchemy doesn't handle the + character in scheme names well
102+
converted_uri = mongodb_uri.replace("mongodb+srv://", "mongodb://")
103+
104+
# Create engine with converted URI
105+
engine = create_engine(converted_uri, **engine_kwargs)
106+
107+
def custom_create_connect_args(url):
108+
# Use original SRV URI for actual MongoDB connection
109+
opts = {"host": mongodb_uri}
110+
return [], opts
111+
112+
engine.dialect.create_connect_args = custom_create_connect_args
113+
return engine
114+
else:
115+
# Standard mongodb:// URLs work fine with SQLAlchemy
116+
return create_engine(mongodb_uri, **engine_kwargs)
117+
118+
except ImportError:
119+
raise ImportError("SQLAlchemy is required for engine creation")
120+
121+
122+
def register_dialect():
123+
"""Register the PyMongoSQL dialect with SQLAlchemy.
124+
125+
This function handles registration for both SQLAlchemy 1.x and 2.x.
126+
Registers support for standard MongoDB connection strings only.
127+
"""
128+
try:
129+
from sqlalchemy.dialects import registry
130+
131+
# Register for standard MongoDB URLs only
132+
registry.register("mongodb", "pymongosql.sqlalchemy_mongodb.sqlalchemy_dialect", "PyMongoSQLDialect")
133+
# Note: mongodb+srv is handled by converting to mongodb in create_connect_args
134+
# SQLAlchemy doesn't support the + character in scheme names directly
135+
136+
return True
137+
except ImportError:
138+
# Fallback for versions without registry
139+
return False
140+
except Exception:
141+
# Handle other registration errors gracefully
142+
return False
143+
144+
145+
# Attempt registration on module import
146+
_registration_successful = register_dialect()
147+
148+
# Export all SQLAlchemy-related functionality
149+
__all__ = [
150+
"create_engine_url",
151+
"create_mongodb_url",
152+
"create_engine_from_mongodb_uri",
153+
"register_dialect",
154+
"__sqlalchemy_version__",
155+
"__supports_sqlalchemy__",
156+
"__supports_sqlalchemy_2x__",
157+
"_registration_successful",
158+
]
159+
160+
# Note: PyMongoSQL now uses standard MongoDB connection strings directly
161+
# No need for PyMongoSQL-specific URL format

pymongosql/sqlalchemy_compat.py renamed to pymongosql/sqlalchemy_mongodb/sqlalchemy_compat.py

File renamed without changes.

pymongosql/sqlalchemy_dialect.py renamed to pymongosql/sqlalchemy_mongodb/sqlalchemy_dialect.py

Lines changed: 12 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -410,7 +410,12 @@ def do_rollback(self, dbapi_connection):
410410
"""
411411
# PyMongoSQL should handle this
412412
if hasattr(dbapi_connection, "rollback"):
413-
dbapi_connection.rollback()
413+
try:
414+
dbapi_connection.rollback()
415+
except Exception:
416+
# MongoDB doesn't always support rollback - ignore errors
417+
# This is normal behavior for MongoDB connections without active transactions
418+
pass
414419

415420
def do_commit(self, dbapi_connection):
416421
"""Commit transaction.
@@ -419,37 +424,14 @@ def do_commit(self, dbapi_connection):
419424
"""
420425
# PyMongoSQL should handle this
421426
if hasattr(dbapi_connection, "commit"):
422-
dbapi_connection.commit()
427+
try:
428+
dbapi_connection.commit()
429+
except Exception:
430+
# MongoDB auto-commits most operations - ignore errors
431+
# This is normal behavior for MongoDB connections
432+
pass
423433

424434

425-
# Register the dialect with SQLAlchemy
426-
# This allows using MongoDB connection strings directly
427-
def register_dialect():
428-
"""Register the PyMongoSQL dialect with SQLAlchemy.
429-
430-
This function handles registration for both SQLAlchemy 1.x and 2.x.
431-
Registers support for standard MongoDB connection strings only.
432-
"""
433-
try:
434-
from sqlalchemy.dialects import registry
435-
436-
# Register for standard MongoDB URLs only
437-
registry.register("mongodb", "pymongosql.sqlalchemy_dialect", "PyMongoSQLDialect")
438-
# Note: mongodb+srv is handled by converting to mongodb in create_connect_args
439-
# SQLAlchemy doesn't support the + character in scheme names directly
440-
441-
return True
442-
except ImportError:
443-
# Fallback for versions without registry
444-
return False
445-
except Exception:
446-
# Handle other registration errors gracefully
447-
return False
448-
449-
450-
# Attempt registration on module import
451-
_registration_successful = register_dialect()
452-
453435
# Version information
454436
__sqlalchemy_version__ = SQLALCHEMY_VERSION
455437
__supports_sqlalchemy_2x__ = SQLALCHEMY_2X

‎tests/test_sqlalchemy_dialect.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,8 @@ class _TestBase(DeclarativeBase): # Prefix with _ to avoid pytest collection
4040
from sqlalchemy.ext.declarative import declarative_base
4141

4242
import pymongosql
43-
from pymongosql.sqlalchemy_dialect import (
43+
from pymongosql.sqlalchemy_mongodb import create_engine_url
44+
from pymongosql.sqlalchemy_mongodb.sqlalchemy_dialect import (
4445
PyMongoSQLDDLCompiler,
4546
PyMongoSQLDialect,
4647
PyMongoSQLIdentifierPreparer,
@@ -289,16 +290,16 @@ class TestSQLAlchemyIntegration(unittest.TestCase):
289290

290291
def test_create_engine_url_helper(self):
291292
"""Test the URL helper function."""
292-
url = pymongosql.create_engine_url("localhost", 27017, "testdb")
293+
url = create_engine_url("localhost", 27017, "testdb")
293294
self.assertEqual(url, "mongodb://localhost:27017/testdb")
294295

295296
# Test with additional parameters
296-
url_with_params = pymongosql.create_engine_url("localhost", 27017, "testdb", ssl=True, replicaSet="rs0")
297+
url_with_params = create_engine_url("localhost", 27017, "testdb", ssl=True, replicaSet="rs0")
297298
self.assertIn("mongodb://localhost:27017/testdb", url_with_params)
298299
self.assertIn("ssl=True", url_with_params)
299300
self.assertIn("replicaSet=rs0", url_with_params)
300301

301-
@patch("pymongosql.sqlalchemy_dialect.pymongosql.connect")
302+
@patch("pymongosql.sqlalchemy_mongodb.sqlalchemy_dialect.pymongosql.connect")
302303
def test_engine_creation(self, mock_connect):
303304
"""Test SQLAlchemy engine creation."""
304305
if not HAS_SQLALCHEMY:
@@ -360,7 +361,7 @@ def test_dialect_registration(self):
360361
try:
361362
from sqlalchemy.dialects import registry
362363

363-
from pymongosql.sqlalchemy_dialect import _registration_successful
364+
from pymongosql.sqlalchemy_mongodb import _registration_successful
364365

365366
# The dialect should be registered
366367
self.assertTrue(hasattr(registry, "load"))

0 commit comments

Comments
 (0)