Skip to content

Commit

Permalink
Fix get_table_names() reflection method
Browse files Browse the repository at this point in the history
It did not respect the `schema` query argument in SQLAlchemy connection
URLs.
  • Loading branch information
amotl committed Jan 15, 2024
1 parent 5ddc8e1 commit 43c45fb
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 0 deletions.
3 changes: 3 additions & 0 deletions CHANGES.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@
- Added support for CrateDB's [FLOAT_VECTOR] data type. For SQLAlchemy
column definitions, you can use it like `FloatVector(dimensions=1024)`.

- Fixed `get_table_names()` reflection method, it did not respect the
`schema` query argument in SQLAlchemy connection URLs.

[FLOAT_VECTOR]: https://crate.io/docs/crate/reference/en/master/general/ddl/data-types.html#float-vector


Expand Down
11 changes: 11 additions & 0 deletions src/sqlalchemy_cratedb/dialect.py
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,15 @@ def connect(self, host=None, port=None, *args, **kwargs):
def _get_default_schema_name(self, connection):
return 'doc'

def _get_effective_schema_name(self, connection):
schema_name_raw = connection.engine.url.query.get("schema")
schema_name = None
if isinstance(schema_name_raw, str):
schema_name = schema_name_raw
elif isinstance(schema_name_raw, tuple):
schema_name = schema_name_raw[0]
return schema_name

def _get_server_version_info(self, connection):
return tuple(connection.connection.lowest_server_version.version)

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

@reflection.cache
def get_table_names(self, connection, schema=None, **kw):
if schema is None:
schema = self._get_effective_schema_name(connection)
cursor = connection.exec_driver_sql(
"SELECT table_name FROM information_schema.tables "
"WHERE {0} = ? "
Expand Down

0 comments on commit 43c45fb

Please sign in to comment.