Skip to content

Commit 91bc8de

Browse files
amotlmatriv
andcommitted
get_table_names(): Fix reflection method to respect schema query argument
It did not respect the `schema` query argument in SQLAlchemy connection URLs. Co-authored-by: Marios Trivyzas <[email protected]>
1 parent 47ccba8 commit 91bc8de

File tree

2 files changed

+13
-0
lines changed

2 files changed

+13
-0
lines changed

CHANGES.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
- Added support for CrateDB's [FLOAT_VECTOR] data type and its accompanying
66
[KNN_MATCH] function, for HNSW matches. For SQLAlchemy column definitions,
77
you can use it like `FloatVector(dimensions=1536)`.
8+
- Fixed `get_table_names()` reflection method to respect the
9+
`schema` query argument in SQLAlchemy connection URLs.
810

911
[FLOAT_VECTOR]: https://cratedb.com/docs/crate/reference/en/latest/general/ddl/data-types.html#float-vector
1012
[KNN_MATCH]: https://cratedb.com/docs/crate/reference/en/latest/general/builtins/scalar-functions.html#scalar-knn-match

src/sqlalchemy_cratedb/dialect.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,15 @@ def connect(self, host=None, port=None, *args, **kwargs):
229229
def _get_default_schema_name(self, connection):
230230
return 'doc'
231231

232+
def _get_effective_schema_name(self, connection):
233+
schema_name_raw = connection.engine.url.query.get("schema")
234+
schema_name = None
235+
if isinstance(schema_name_raw, str):
236+
schema_name = schema_name_raw
237+
elif isinstance(schema_name_raw, tuple):
238+
schema_name = schema_name_raw[0]
239+
return schema_name
240+
232241
def _get_server_version_info(self, connection):
233242
return tuple(connection.connection.lowest_server_version.version)
234243

@@ -258,6 +267,8 @@ def get_schema_names(self, connection, **kw):
258267

259268
@reflection.cache
260269
def get_table_names(self, connection, schema=None, **kw):
270+
if schema is None:
271+
schema = self._get_effective_schema_name(connection)
261272
cursor = connection.exec_driver_sql(
262273
"SELECT table_name FROM information_schema.tables "
263274
"WHERE {0} = ? "

0 commit comments

Comments
 (0)