Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import itertools
2import os
3from collections.abc import Mapping, Sequence
5import sqlalchemy as sa
6from sqlalchemy.engine.url import make_url
7from sqlalchemy.exc import OperationalError, ProgrammingError
9from ..utils import starts_with
10from .orm import quote
13def escape_like(string, escape_char='*'):
14 """
15 Escape the string parameter used in SQL LIKE expressions.
17 ::
19 from sqlalchemy_utils import escape_like
22 query = session.query(User).filter(
23 User.name.ilike(escape_like('John'))
24 )
27 :param string: a string to escape
28 :param escape_char: escape character
29 """
30 return (
31 string.replace(escape_char, escape_char * 2)
32 .replace('%', escape_char + '%')
33 .replace('_', escape_char + '_')
34 )
37def json_sql(value, scalars_to_json=True):
38 """
39 Convert python data structures to PostgreSQL specific SQLAlchemy JSON
40 constructs. This function is extremly useful if you need to build
41 PostgreSQL JSON on python side.
43 .. note::
45 This function needs PostgreSQL >= 9.4
47 Scalars are converted to to_json SQLAlchemy function objects
49 ::
51 json_sql(1) # Equals SQL: to_json(1)
53 json_sql('a') # to_json('a')
56 Mappings are converted to json_build_object constructs
58 ::
60 json_sql({'a': 'c', '2': 5}) # json_build_object('a', 'c', '2', 5)
63 Sequences (other than strings) are converted to json_build_array constructs
65 ::
67 json_sql([1, 2, 3]) # json_build_array(1, 2, 3)
70 You can also nest these data structures
72 ::
74 json_sql({'a': [1, 2, 3]})
75 # json_build_object('a', json_build_array[1, 2, 3])
78 :param value:
79 value to be converted to SQLAlchemy PostgreSQL function constructs
80 """
81 if scalars_to_json:
83 def scalar_convert(a):
84 return sa.func.to_json(sa.text(a))
85 else:
86 scalar_convert = sa.text
88 if isinstance(value, Mapping):
89 return sa.func.json_build_object(
90 *(
91 json_sql(v, scalars_to_json=False)
92 for v in itertools.chain(*value.items())
93 )
94 )
95 elif isinstance(value, str):
96 return scalar_convert(f"'{value}'")
97 elif isinstance(value, Sequence):
98 return sa.func.json_build_array(
99 *(json_sql(v, scalars_to_json=False) for v in value)
100 )
101 elif isinstance(value, (int, float)):
102 return scalar_convert(str(value))
103 return value
106def jsonb_sql(value, scalars_to_jsonb=True):
107 """
108 Convert python data structures to PostgreSQL specific SQLAlchemy JSONB
109 constructs. This function is extremly useful if you need to build
110 PostgreSQL JSONB on python side.
112 .. note::
114 This function needs PostgreSQL >= 9.4
116 Scalars are converted to to_jsonb SQLAlchemy function objects
118 ::
120 jsonb_sql(1) # Equals SQL: to_jsonb(1)
122 jsonb_sql('a') # to_jsonb('a')
125 Mappings are converted to jsonb_build_object constructs
127 ::
129 jsonb_sql({'a': 'c', '2': 5}) # jsonb_build_object('a', 'c', '2', 5)
132 Sequences (other than strings) converted to jsonb_build_array constructs
134 ::
136 jsonb_sql([1, 2, 3]) # jsonb_build_array(1, 2, 3)
139 You can also nest these data structures
141 ::
143 jsonb_sql({'a': [1, 2, 3]})
144 # jsonb_build_object('a', jsonb_build_array[1, 2, 3])
147 :param value:
148 value to be converted to SQLAlchemy PostgreSQL function constructs
149 :boolean jsonbb:
150 Flag to alternatively convert the return with a to_jsonb construct
151 """
152 if scalars_to_jsonb:
154 def scalar_convert(a):
155 return sa.func.to_jsonb(sa.text(a))
156 else:
157 scalar_convert = sa.text
159 if isinstance(value, Mapping):
160 return sa.func.jsonb_build_object(
161 *(
162 jsonb_sql(v, scalars_to_jsonb=False)
163 for v in itertools.chain(*value.items())
164 )
165 )
166 elif isinstance(value, str):
167 return scalar_convert(f"'{value}'")
168 elif isinstance(value, Sequence):
169 return sa.func.jsonb_build_array(
170 *(jsonb_sql(v, scalars_to_jsonb=False) for v in value)
171 )
172 elif isinstance(value, (int, float)):
173 return scalar_convert(str(value))
174 return value
177def has_index(column_or_constraint):
178 """
179 Return whether or not given column or the columns of given foreign key
180 constraint have an index. A column has an index if it has a single column
181 index or it is the first column in compound column index.
183 A foreign key constraint has an index if the constraint columns are the
184 first columns in compound column index.
186 :param column_or_constraint:
187 SQLAlchemy Column object or SA ForeignKeyConstraint object
189 .. versionadded: 0.26.2
191 .. versionchanged: 0.30.18
192 Added support for foreign key constaints.
194 ::
196 from sqlalchemy_utils import has_index
199 class Article(Base):
200 __tablename__ = 'article'
201 id = sa.Column(sa.Integer, primary_key=True)
202 title = sa.Column(sa.String(100))
203 is_published = sa.Column(sa.Boolean, index=True)
204 is_deleted = sa.Column(sa.Boolean)
205 is_archived = sa.Column(sa.Boolean)
207 __table_args__ = (
208 sa.Index('my_index', is_deleted, is_archived),
209 )
212 table = Article.__table__
214 has_index(table.c.is_published) # True
215 has_index(table.c.is_deleted) # True
216 has_index(table.c.is_archived) # False
219 Also supports primary key indexes
221 ::
223 from sqlalchemy_utils import has_index
226 class ArticleTranslation(Base):
227 __tablename__ = 'article_translation'
228 id = sa.Column(sa.Integer, primary_key=True)
229 locale = sa.Column(sa.String(10), primary_key=True)
230 title = sa.Column(sa.String(100))
233 table = ArticleTranslation.__table__
235 has_index(table.c.locale) # False
236 has_index(table.c.id) # True
239 This function supports foreign key constraints as well
241 ::
244 class User(Base):
245 __tablename__ = 'user'
246 first_name = sa.Column(sa.Unicode(255), primary_key=True)
247 last_name = sa.Column(sa.Unicode(255), primary_key=True)
249 class Article(Base):
250 __tablename__ = 'article'
251 id = sa.Column(sa.Integer, primary_key=True)
252 author_first_name = sa.Column(sa.Unicode(255))
253 author_last_name = sa.Column(sa.Unicode(255))
254 __table_args__ = (
255 sa.ForeignKeyConstraint(
256 [author_first_name, author_last_name],
257 [User.first_name, User.last_name]
258 ),
259 sa.Index(
260 'my_index',
261 author_first_name,
262 author_last_name
263 )
264 )
266 table = Article.__table__
267 constraint = list(table.foreign_keys)[0].constraint
269 has_index(constraint) # True
270 """
271 table = column_or_constraint.table
272 if not isinstance(table, sa.Table):
273 raise TypeError(
274 'Only columns belonging to Table objects are supported. Given '
275 'column belongs to %r.' % table
276 )
277 primary_keys = table.primary_key.columns.values()
278 if isinstance(column_or_constraint, sa.ForeignKeyConstraint):
279 columns = list(column_or_constraint.columns.values())
280 else:
281 columns = [column_or_constraint]
283 return (primary_keys and starts_with(primary_keys, columns)) or any(
284 starts_with(index.columns.values(), columns) for index in table.indexes
285 )
288def has_unique_index(column_or_constraint):
289 """
290 Return whether or not given column or given foreign key constraint has a
291 unique index.
293 A column has a unique index if it has a single column primary key index or
294 it has a single column UniqueConstraint.
296 A foreign key constraint has a unique index if the columns of the
297 constraint are the same as the columns of table primary key or the coluns
298 of any unique index or any unique constraint of the given table.
300 :param column: SQLAlchemy Column object
302 .. versionadded: 0.27.1
304 .. versionchanged: 0.30.18
305 Added support for foreign key constaints.
307 Fixed support for unique indexes (previously only worked for unique
308 constraints)
310 ::
312 from sqlalchemy_utils import has_unique_index
315 class Article(Base):
316 __tablename__ = 'article'
317 id = sa.Column(sa.Integer, primary_key=True)
318 title = sa.Column(sa.String(100))
319 is_published = sa.Column(sa.Boolean, unique=True)
320 is_deleted = sa.Column(sa.Boolean)
321 is_archived = sa.Column(sa.Boolean)
324 table = Article.__table__
326 has_unique_index(table.c.is_published) # True
327 has_unique_index(table.c.is_deleted) # False
328 has_unique_index(table.c.id) # True
331 This function supports foreign key constraints as well
333 ::
336 class User(Base):
337 __tablename__ = 'user'
338 first_name = sa.Column(sa.Unicode(255), primary_key=True)
339 last_name = sa.Column(sa.Unicode(255), primary_key=True)
341 class Article(Base):
342 __tablename__ = 'article'
343 id = sa.Column(sa.Integer, primary_key=True)
344 author_first_name = sa.Column(sa.Unicode(255))
345 author_last_name = sa.Column(sa.Unicode(255))
346 __table_args__ = (
347 sa.ForeignKeyConstraint(
348 [author_first_name, author_last_name],
349 [User.first_name, User.last_name]
350 ),
351 sa.Index(
352 'my_index',
353 author_first_name,
354 author_last_name,
355 unique=True
356 )
357 )
359 table = Article.__table__
360 constraint = list(table.foreign_keys)[0].constraint
362 has_unique_index(constraint) # True
365 :raises TypeError: if given column does not belong to a Table object
366 """
367 table = column_or_constraint.table
368 if not isinstance(table, sa.Table):
369 raise TypeError(
370 'Only columns belonging to Table objects are supported. Given '
371 'column belongs to %r.' % table
372 )
373 primary_keys = list(table.primary_key.columns.values())
374 if isinstance(column_or_constraint, sa.ForeignKeyConstraint):
375 columns = list(column_or_constraint.columns.values())
376 else:
377 columns = [column_or_constraint]
379 return (
380 (columns == primary_keys)
381 or any(
382 columns == list(constraint.columns.values())
383 for constraint in table.constraints
384 if isinstance(constraint, sa.sql.schema.UniqueConstraint)
385 )
386 or any(
387 columns == list(index.columns.values())
388 for index in table.indexes
389 if index.unique
390 )
391 )
394def is_auto_assigned_date_column(column):
395 """
396 Returns whether or not given SQLAlchemy Column object's is auto assigned
397 DateTime or Date.
399 :param column: SQLAlchemy Column object
400 """
401 return (
402 isinstance(column.type, sa.DateTime) or isinstance(column.type, sa.Date)
403 ) and (
404 column.default
405 or column.server_default
406 or column.onupdate
407 or column.server_onupdate
408 )
411def _set_url_database(url: sa.engine.url.URL, database):
412 """Set the database of an engine URL.
414 :param url: A SQLAlchemy engine URL.
415 :param database: New database to set.
417 """
418 # Cannot use URL.set() as database may need to be set to None.
419 ret = url._replace(database=database)
420 assert ret.database == database, ret
421 return ret
424def _get_scalar_result(engine, sql):
425 with engine.connect() as conn:
426 return conn.scalar(sql)
429def _sqlite_file_exists(database):
430 if not os.path.isfile(database) or os.path.getsize(database) < 100:
431 return False
433 with open(database, 'rb') as f:
434 header = f.read(100)
436 return header[:16] == b'SQLite format 3\x00'
439def database_exists(url):
440 """Check if a database exists.
442 :param url: A SQLAlchemy engine URL.
444 Performs backend-specific testing to quickly determine if a database
445 exists on the server. ::
447 database_exists('postgresql://postgres@localhost/name') #=> False
448 create_database('postgresql://postgres@localhost/name')
449 database_exists('postgresql://postgres@localhost/name') #=> True
451 Supports checking against a constructed URL as well. ::
453 engine = create_engine('postgresql://postgres@localhost/name')
454 database_exists(engine.url) #=> False
455 create_database(engine.url)
456 database_exists(engine.url) #=> True
458 """
460 url = make_url(url)
461 database = url.database
462 dialect_name = url.get_dialect().name
463 engine = None
464 try:
465 if dialect_name == 'postgresql':
466 text = "SELECT 1 FROM pg_database WHERE datname='%s'" % database
467 for db in (database, 'postgres', 'template1', 'template0', None):
468 url = _set_url_database(url, database=db)
469 engine = sa.create_engine(url, isolation_level='AUTOCOMMIT')
470 try:
471 return bool(_get_scalar_result(engine, sa.text(text)))
472 except (ProgrammingError, OperationalError):
473 pass
474 return False
476 elif dialect_name == 'mysql':
477 url = _set_url_database(url, database=None)
478 engine = sa.create_engine(url)
479 text = (
480 'SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA '
481 "WHERE SCHEMA_NAME = '%s'" % database
482 )
483 return bool(_get_scalar_result(engine, sa.text(text)))
485 elif dialect_name == 'sqlite':
486 url = _set_url_database(url, database=None)
487 engine = sa.create_engine(url)
488 if database:
489 return database == ':memory:' or _sqlite_file_exists(database)
490 else:
491 # The default SQLAlchemy database is in memory, and :memory: is
492 # not required, thus we should support that use case.
493 return True
494 elif dialect_name == 'mssql':
495 text = "SELECT 1 FROM sys.databases WHERE name = '%s'" % database
496 url = _set_url_database(url, database='master')
497 engine = sa.create_engine(url, isolation_level='AUTOCOMMIT')
498 try:
499 return bool(_get_scalar_result(engine, sa.text(text)))
500 except (ProgrammingError, OperationalError):
501 return False
502 else:
503 text = 'SELECT 1'
504 try:
505 engine = sa.create_engine(url)
506 return bool(_get_scalar_result(engine, sa.text(text)))
507 except (ProgrammingError, OperationalError):
508 return False
509 finally:
510 if engine:
511 engine.dispose()
514def create_database(url, encoding='utf8', template=None):
515 """Issue the appropriate CREATE DATABASE statement.
517 :param url: A SQLAlchemy engine URL.
518 :param encoding: The encoding to create the database as.
519 :param template:
520 The name of the template from which to create the new database. At the
521 moment only supported by PostgreSQL driver.
523 To create a database, you can pass a simple URL that would have
524 been passed to ``create_engine``. ::
526 create_database('postgresql://postgres@localhost/name')
528 You may also pass the url from an existing engine. ::
530 create_database(engine.url)
532 Has full support for mysql, postgres, and sqlite. In theory,
533 other database engines should be supported.
534 """
536 url = make_url(url)
537 database = url.database
538 dialect_name = url.get_dialect().name
539 dialect_driver = url.get_dialect().driver
541 if dialect_name == 'postgresql':
542 url = _set_url_database(url, database='postgres')
543 elif dialect_name == 'mssql':
544 url = _set_url_database(url, database='master')
545 elif dialect_name == 'cockroachdb':
546 url = _set_url_database(url, database='defaultdb')
547 elif not dialect_name == 'sqlite':
548 url = _set_url_database(url, database=None)
550 if (dialect_name == 'mssql' and dialect_driver in {'pymssql', 'pyodbc'}) or (
551 dialect_name == 'postgresql'
552 and dialect_driver
553 in {'asyncpg', 'pg8000', 'psycopg', 'psycopg2', 'psycopg2cffi'}
554 ):
555 engine = sa.create_engine(url, isolation_level='AUTOCOMMIT')
556 else:
557 engine = sa.create_engine(url)
559 if dialect_name == 'postgresql':
560 if not template:
561 template = 'template1'
563 with engine.begin() as conn:
564 text = "CREATE DATABASE {} ENCODING '{}' TEMPLATE {}".format(
565 quote(conn, database), encoding, quote(conn, template)
566 )
567 conn.execute(sa.text(text))
569 elif dialect_name == 'mysql':
570 with engine.begin() as conn:
571 text = "CREATE DATABASE {} CHARACTER SET = '{}'".format(
572 quote(conn, database), encoding
573 )
574 conn.execute(sa.text(text))
576 elif dialect_name == 'sqlite' and database != ':memory:':
577 if database:
578 with engine.begin() as conn:
579 conn.execute(sa.text('CREATE TABLE DB(id int)'))
580 conn.execute(sa.text('DROP TABLE DB'))
582 else:
583 with engine.begin() as conn:
584 text = f'CREATE DATABASE {quote(conn, database)}'
585 conn.execute(sa.text(text))
587 engine.dispose()
590def drop_database(url):
591 """Issue the appropriate DROP DATABASE statement.
593 :param url: A SQLAlchemy engine URL.
595 Works similar to the :func:`create_database` method in that both url text
596 and a constructed url are accepted.
598 ::
600 drop_database('postgresql://postgres@localhost/name')
601 drop_database(engine.url)
603 """
605 url = make_url(url)
606 database = url.database
607 dialect_name = url.get_dialect().name
608 dialect_driver = url.get_dialect().driver
610 if dialect_name == 'postgresql':
611 url = _set_url_database(url, database='postgres')
612 elif dialect_name == 'mssql':
613 url = _set_url_database(url, database='master')
614 elif dialect_name == 'cockroachdb':
615 url = _set_url_database(url, database='defaultdb')
616 elif not dialect_name == 'sqlite':
617 url = _set_url_database(url, database=None)
619 if (dialect_name == 'mssql' and dialect_driver in {'pymssql', 'pyodbc'}) or (
620 dialect_name == 'postgresql'
621 and dialect_driver
622 in {'asyncpg', 'pg8000', 'psycopg', 'psycopg2', 'psycopg2cffi'}
623 ):
624 engine = sa.create_engine(url, isolation_level='AUTOCOMMIT')
625 else:
626 engine = sa.create_engine(url)
628 if dialect_name == 'sqlite' and database != ':memory:':
629 if database:
630 os.remove(database)
631 else:
632 with engine.begin() as conn:
633 text = f'DROP DATABASE {quote(conn, database)}'
634 conn.execute(sa.text(text))
636 engine.dispose()