Skip to content

Macrostrat Database

Macrostrat's Database module provides a simplified wrapper over SQLAlchemy databases, making it easier to build common database management functionality.

macrostrat.database.Database

Bases: object

Source code in database/macrostrat/database/core.py 
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
class Database(object):
    mapper: Optional[DatabaseMapper] = None
    metadata: MetaData
    session: Session
    instance_params: dict

    __inspector__ = None

    def __init__(self, db_conn: DatabaseInput, *, echo_sql=False, **kwargs):
        """
        Wrapper for interacting with a database using SQLAlchemy.
        Optimized for use with PostgreSQL, but usable with SQLite
        as well.

        Args:
            db_conn (str | URL | Engine): Connection string or engine for the database.

        Keyword Args:
            echo_sql (bool): If True, will echo SQL commands to the
                console. Default is False.
            instance_params (dict): Parameters to
                pass to queries and other database operations.
        """

        compiles(Insert, "postgresql")(prefix_inserts)

        self.instance_params = kwargs.pop("instance_params", {})

        if echo_sql:
            kwargs["echo"] = True

        self.engine = create_engine(db_conn, **kwargs)

        self.metadata = kwargs.get("metadata", metadata)

        # Scoped session for database
        # https://docs.sqlalchemy.org/en/13/orm/contextual.html#unitofwork-contextual
        # https://docs.sqlalchemy.org/en/13/orm/session_basics.html#session-faq-whentocreate
        self._session_factory = sessionmaker(bind=self.engine)
        self.session = scoped_session(self._session_factory)
        # Use the self.session_scope function to more explicitly manage sessions.
        self._table_cache: dict = {}

    def create_tables(self):
        """
        Create all tables described by the database's metadata instance.
        """
        metadata.create_all(bind=self.engine)

    def automap(self, **kwargs):
        log.info("Automapping the database")
        self.mapper = DatabaseMapper(self)
        self.mapper.reflect_database(**kwargs)

    def get_server_version(self):
        with self.engine.connect():
            return self.engine.dialect.server_version_info

    @contextmanager
    def session_scope(self, commit=True):
        """Provide a transactional scope around a series of operations."""
        # self.__old_session = self.session
        # session = self._session_factory()
        session = self.session
        try:
            yield session
            if commit:
                session.commit()
        except Exception as err:
            session.rollback()
            raise err
        finally:
            session.close()

    def _flush_nested_objects(self, session):
        """
        Flush objects remaining in a session (generally these are objects loaded
        during schema-based importing).
        """
        for object in session:
            try:
                session.flush(objects=[object])
                log.debug(f"Successfully flushed instance {object}")
            except IntegrityError as err:
                session.rollback()
                log.debug(err)

    def run_sql(self, fn, params=None, **kwargs):
        """Executes SQL files or query strings using the run_sql function.

        Args:
            fn (str|Path): SQL file or query string to execute.
            params (dict): Parameters to pass to the query.

        Keyword Args:
            use_instance_params (bool): If True, will use the instance_params set on
                the Database object. Default is True.

        Returns: Iterator of results from the query.
        """
        params = self._setup_params(params, kwargs)
        return run_sql(self.session, fn, params, **kwargs)

    def run_query(self, sql, params=None, **kwargs):
        """Run a single query on the database object, returning the result.

        Args:
            sql (str): SQL file or query to execute.
            params (dict): Parameters to pass to the query.

        Keyword Args:
            use_instance_params (bool): If True, will use the instance_params set on
                the Database object. Default is True.
        """
        params = self._setup_params(params, kwargs)
        return run_query(self.session, sql, params, **kwargs)

    def run_fixtures(self, fixtures: Union[Path, list[Path]], params=None, **kwargs):
        """Run a set of fixtures on the database object.

        Args:
            fixtures (Path|list[Path]): Path to a directory of fixtures or a list of paths to fixture files.
            params (dict): Parameters to pass to the query.

        Keyword Args:
            use_instance_params (bool): If True, will use the instance_params set on
                the Database object. Default is True.
        """
        params = self._setup_params(params, kwargs)
        return run_fixtures(self.session, fixtures, params, **kwargs)

    def _setup_params(self, params, kwargs):
        use_instance_params = kwargs.pop("use_instance_params", True)
        if params is None:
            params = {}
        if use_instance_params:
            if isinstance(params, dict):
                params.update(self.instance_params)
                return params
            if isinstance(params, list):
                if all(isinstance(p, dict) for p in params):
                    params = [dict(p, **self.instance_params) for p in params]
                    return params
            warnings.warn(
                "Could not apply shared instance params to %", params, stacklevel=2
            )
        return params

    def exec_sql(self, sql, params=None, **kwargs):
        """Executes SQL files passed"""
        warnings.warn(
            "exec_sql is deprecated and will be removed in version 4.0. Use run_sql instead",
            DeprecationWarning,
        )
        return self.run_sql(sql, params, **kwargs)

    def get_dataframe(self, *args):
        """Returns a Pandas DataFrame from a SQL query"""
        return get_dataframe(self.engine, *args)

    @property
    def inspector(self):
        if self.__inspector__ is None:
            self.__inspector__ = inspect(self.engine)
        return self.__inspector__

    def refresh_schema(self, *, automap=None):
        """
        Refresh the current database connection

        - closes the session and flushes
        - removes the inspector

        If automap is True, will automap the database after refreshing.
        If automap is False, will not automap the database after refreshing.
        If automap is None, it will re-map the database if it was previously mapped.
        """
        # Close the session
        self.session.flush()
        self.session.close()
        # Remove the inspector
        self.__inspector__ = None

        if automap is None:
            automap = self.mapper is not None

        if automap:
            self.automap()

    def entity_names(self, **kwargs):
        """
        Returns an iterator of names of *schema objects*
        (both tables and views) from a the database.
        """
        yield from self.inspector.get_table_names(**kwargs)
        yield from self.inspector.get_view_names(**kwargs)

    def get(self, model, *args, **kwargs):
        if isinstance(model, str):
            model = getattr(self.model, model)
        return self.session.query(model).get(*args, **kwargs)

    def get_or_create(self, model, **kwargs):
        """
        Get an instance of a model, or create it if it doesn't
        exist.
        """
        if isinstance(model, str):
            model = getattr(self.model, model)
        return get_or_create(self.session, model, **kwargs)

    def reflect_table(self, *args, **kwargs):
        """
        One-off reflection of a database table or view. Note: for most purposes,
        it will be better to use the database tables automapped at runtime using
        `self.automap()`. Then, tables can be accessed using the
        `self.table` object. However, this function can be useful for views (which
        are not reflected automatically), or to customize type definitions for mapped
        tables.

        A set of `column_args` can be used to pass columns to override with the mapper, for
        instance to set up foreign and primary key constraints.
        https://docs.sqlalchemy.org/en/13/core/reflection.html#reflecting-views
        """
        warnings.warn(
            "reflect_table is deprecated and will be removed in version 4.0. Shift away from table refection, or use reflect_table from the macrostrat.database.utils module.",
            DeprecationWarning,
        )

        return reflect_table(self.engine, *args, **kwargs)

    @property
    def table(self):
        """
        Map of all tables in the database as SQLAlchemy table objects
        """
        if self.mapper is None or self.mapper._tables is None:
            self.automap()
        return self.mapper._tables

    @property
    def model(self):
        """
        Map of all tables in the database as SQLAlchemy models

        https://docs.sqlalchemy.org/en/latest/orm/extensions/automap.html
        """
        if self.mapper is None or self.mapper._models is None:
            self.automap()
        return self.mapper._models

    @property
    def mapped_classes(self):
        return self.model

    @contextmanager
    def transaction(self, *, rollback="on-error", connection=None, raise_errors=True):
        """Create a database session that can be rolled back after use.
        This is similar to the `session_scope` method but includes
        more fine-grained control over transactions. The two methods may be integrated
        in the future.

        This is based on the Sparrow's implementation:
        https://github.com/EarthCubeGeochron/Sparrow/blob/main/backend/conftest.py

        It can be effectively used in a Pytest fixture like so:
        ```
        @fixture(scope="class")
        def db(base_db):
            with base_db.transaction(rollback=True):
                yield base_db
        """
        if connection is None:
            connection = self.engine.connect()
        transaction = connection.begin()
        session = Session(bind=connection)
        prev_session = self.session
        self.session = session

        should_rollback = rollback == "always"

        try:
            yield self
        except Exception as e:
            should_rollback = rollback != "never"
            if raise_errors:
                raise e
        finally:
            if should_rollback:
                transaction.rollback()
            else:
                transaction.commit()
            session.close()
            self.session = prev_session

    savepoint_counter = 0

    @contextmanager
    def savepoint(self, name=None, rollback="on-error", connection=None):
        """A PostgreSQL-specific savepoint context manager. This is similar to the
        `transaction` context manager but uses savepoints directly for simpler operation.
        Notably, it supports nested savepoints, a feature that is difficult in SQLAlchemy's `transaction`
        model.

        This function is not yet drop-in compatible with the `transaction` context manager, but that
        is a future goal.
        """
        if name is None:
            name = f"sp_{self.savepoint_counter}"
            self.savepoint_counter += 1

        _prev_session = self.session

        if connection is None:
            connection = self.session.connection()

        params = {"name": Identifier(name)}
        run_query(connection, "SAVEPOINT {name}", params)
        should_rollback = rollback == "always"
        self.session = Session(bind=connection)
        try:
            yield name
        except Exception as e:
            should_rollback = rollback != "never"
            raise e
        finally:
            _clear_savepoint(connection, name, rollback=should_rollback)
            self.session.close()
            self.session = _prev_session

    def get_table(self, name, *, schema=None):
        """Return a reflected SQLAlchemy Table object, with per-instance caching.

        After the first call the result is cached; subsequent calls for the
        same table are instant.  If automap has already been run the mapper's
        existing Table is reused, avoiding a second round-trip.

        Args:
            name: Table name as ``"table"``, ``"schema.table"``, or
                  ``("schema", "table")``.
            schema: Explicit schema override (default ``"public"``).
        """
        schema_, table_name = _parse_table_name(name, schema)
        cache_key = (schema_, table_name)
        if cache_key in self._table_cache:
            return self._table_cache[cache_key]

        # Reuse the already-reflected Table from automap when available
        if self.mapper is not None:
            model_key = _model_key(schema_, table_name)
            if model_key in self.mapper._models:
                tbl = self.mapper._models[model_key].__table__
                self._table_cache[cache_key] = tbl
                return tbl

        # Per-table reflection; "public" → None matches how automap stores it
        reflect_schema = None if schema_ == "public" else schema_
        tbl = reflect_table(self.engine, table_name, schema=reflect_schema)
        self._table_cache[cache_key] = tbl
        return tbl

    def get_model(self, name, *, schema=None, automap=True):
        """Return the ORM model class for a table.

        If the target schema has not yet been reflected and ``automap=True``
        (the default), it is reflected lazily before the lookup.  Set
        ``automap=False`` to raise ``LookupError`` instead, which is useful
        when you want strict control over when reflection happens.

        Args:
            name: Table name as ``"table"``, ``"schema.table"``, or
                  ``("schema", "table")``.
            schema: Explicit schema override (default ``"public"``).
            automap: Lazily reflect the schema if not yet mapped.

        Raises:
            LookupError: When the model is not found.
        """
        schema_, table_name = _parse_table_name(name, schema)
        model_key = _model_key(schema_, table_name)

        if self.mapper is not None and model_key in self.mapper._models:
            return self.mapper._models[model_key]

        if not automap:
            raise LookupError(
                f"No ORM model found for {schema_}.{table_name}. "
                "Call db.automap() first, or use get_table() for a Table object."
            )

        # Lazy automap: reflect only the needed schema
        if self.mapper is None:
            self.automap(schemas=[schema_])
        elif schema_ not in self.mapper._reflected_schemas:
            self.mapper.reflect_schema(schema_)

        if model_key in self.mapper._models:
            return self.mapper._models[model_key]

        raise LookupError(
            f"No ORM model found for {schema_}.{table_name} after reflecting "
            f"schema '{schema_}'. Verify the table exists, or use get_table()."
        )

    def __getitem__(self, name):
        """Subscript shorthand for get_table()."""
        return self.get_table(name)

    # Destroy engine on cleanup
    def cleanup(self):
        try:
            self.session.close()
        except OperationalError:
            pass
        self.engine.dispose()

    def __del__(self):
        self.cleanup()

model property

Map of all tables in the database as SQLAlchemy models

https://docs.sqlalchemy.org/en/latest/orm/extensions/automap.html

table property

Map of all tables in the database as SQLAlchemy table objects

__getitem__(name)

Subscript shorthand for get_table().

Source code in database/macrostrat/database/core.py 
457
458
459
def __getitem__(self, name):
    """Subscript shorthand for get_table()."""
    return self.get_table(name)

__init__(db_conn, *, echo_sql=False, **kwargs)

Wrapper for interacting with a database using SQLAlchemy. Optimized for use with PostgreSQL, but usable with SQLite as well.

Parameters:

Name Type Description Default
db_conn str | URL | Engine

Connection string or engine for the database.

 required

Other Parameters:

Name Type Description
echo_sql bool

If True, will echo SQL commands to the console. Default is False.
instance_params dict

Parameters to pass to queries and other database operations.
Source code in database/macrostrat/database/core.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def __init__(self, db_conn: DatabaseInput, *, echo_sql=False, **kwargs):
    """
    Wrapper for interacting with a database using SQLAlchemy.
    Optimized for use with PostgreSQL, but usable with SQLite
    as well.

    Args:
        db_conn (str | URL | Engine): Connection string or engine for the database.

    Keyword Args:
        echo_sql (bool): If True, will echo SQL commands to the
            console. Default is False.
        instance_params (dict): Parameters to
            pass to queries and other database operations.
    """

    compiles(Insert, "postgresql")(prefix_inserts)

    self.instance_params = kwargs.pop("instance_params", {})

    if echo_sql:
        kwargs["echo"] = True

    self.engine = create_engine(db_conn, **kwargs)

    self.metadata = kwargs.get("metadata", metadata)

    # Scoped session for database
    # https://docs.sqlalchemy.org/en/13/orm/contextual.html#unitofwork-contextual
    # https://docs.sqlalchemy.org/en/13/orm/session_basics.html#session-faq-whentocreate
    self._session_factory = sessionmaker(bind=self.engine)
    self.session = scoped_session(self._session_factory)
    # Use the self.session_scope function to more explicitly manage sessions.
    self._table_cache: dict = {}

create_tables()

Create all tables described by the database's metadata instance.

Source code in database/macrostrat/database/core.py 
 96
 97
 98
 99
100
def create_tables(self):
    """
    Create all tables described by the database's metadata instance.
    """
    metadata.create_all(bind=self.engine)

entity_names(**kwargs)

Returns an iterator of names of schema objects (both tables and views) from a the database.

Source code in database/macrostrat/database/core.py 
242
243
244
245
246
247
248
def entity_names(self, **kwargs):
    """
    Returns an iterator of names of *schema objects*
    (both tables and views) from a the database.
    """
    yield from self.inspector.get_table_names(**kwargs)
    yield from self.inspector.get_view_names(**kwargs)

exec_sql(sql, params=None, **kwargs)

Executes SQL files passed

Source code in database/macrostrat/database/core.py 
201
202
203
204
205
206
207
def exec_sql(self, sql, params=None, **kwargs):
    """Executes SQL files passed"""
    warnings.warn(
        "exec_sql is deprecated and will be removed in version 4.0. Use run_sql instead",
        DeprecationWarning,
    )
    return self.run_sql(sql, params, **kwargs)

get_dataframe(*args)

Returns a Pandas DataFrame from a SQL query

Source code in database/macrostrat/database/core.py 
209
210
211
def get_dataframe(self, *args):
    """Returns a Pandas DataFrame from a SQL query"""
    return get_dataframe(self.engine, *args)

get_model(name, *, schema=None, automap=True)

Return the ORM model class for a table.

If the target schema has not yet been reflected and automap=True (the default), it is reflected lazily before the lookup. Set automap=False to raise LookupError instead, which is useful when you want strict control over when reflection happens.

Parameters:

Name Type Description Default
name

Table name as "table", "schema.table", or ("schema", "table").

 required
schema

Explicit schema override (default "public").

 None
automap

Lazily reflect the schema if not yet mapped.

 True

Raises:

Type Description
LookupError

When the model is not found.
Source code in database/macrostrat/database/core.py 
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
def get_model(self, name, *, schema=None, automap=True):
    """Return the ORM model class for a table.

    If the target schema has not yet been reflected and ``automap=True``
    (the default), it is reflected lazily before the lookup.  Set
    ``automap=False`` to raise ``LookupError`` instead, which is useful
    when you want strict control over when reflection happens.

    Args:
        name: Table name as ``"table"``, ``"schema.table"``, or
              ``("schema", "table")``.
        schema: Explicit schema override (default ``"public"``).
        automap: Lazily reflect the schema if not yet mapped.

    Raises:
        LookupError: When the model is not found.
    """
    schema_, table_name = _parse_table_name(name, schema)
    model_key = _model_key(schema_, table_name)

    if self.mapper is not None and model_key in self.mapper._models:
        return self.mapper._models[model_key]

    if not automap:
        raise LookupError(
            f"No ORM model found for {schema_}.{table_name}. "
            "Call db.automap() first, or use get_table() for a Table object."
        )

    # Lazy automap: reflect only the needed schema
    if self.mapper is None:
        self.automap(schemas=[schema_])
    elif schema_ not in self.mapper._reflected_schemas:
        self.mapper.reflect_schema(schema_)

    if model_key in self.mapper._models:
        return self.mapper._models[model_key]

    raise LookupError(
        f"No ORM model found for {schema_}.{table_name} after reflecting "
        f"schema '{schema_}'. Verify the table exists, or use get_table()."
    )

get_or_create(model, **kwargs)

Get an instance of a model, or create it if it doesn't exist.

Source code in database/macrostrat/database/core.py 
255
256
257
258
259
260
261
262
def get_or_create(self, model, **kwargs):
    """
    Get an instance of a model, or create it if it doesn't
    exist.
    """
    if isinstance(model, str):
        model = getattr(self.model, model)
    return get_or_create(self.session, model, **kwargs)

get_table(name, *, schema=None)

Return a reflected SQLAlchemy Table object, with per-instance caching.

After the first call the result is cached; subsequent calls for the same table are instant. If automap has already been run the mapper's existing Table is reused, avoiding a second round-trip.

Parameters:

Name Type Description Default
name

Table name as "table", "schema.table", or ("schema", "table").

 required
schema

Explicit schema override (default "public").

 None
Source code in database/macrostrat/database/core.py 
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
def get_table(self, name, *, schema=None):
    """Return a reflected SQLAlchemy Table object, with per-instance caching.

    After the first call the result is cached; subsequent calls for the
    same table are instant.  If automap has already been run the mapper's
    existing Table is reused, avoiding a second round-trip.

    Args:
        name: Table name as ``"table"``, ``"schema.table"``, or
              ``("schema", "table")``.
        schema: Explicit schema override (default ``"public"``).
    """
    schema_, table_name = _parse_table_name(name, schema)
    cache_key = (schema_, table_name)
    if cache_key in self._table_cache:
        return self._table_cache[cache_key]

    # Reuse the already-reflected Table from automap when available
    if self.mapper is not None:
        model_key = _model_key(schema_, table_name)
        if model_key in self.mapper._models:
            tbl = self.mapper._models[model_key].__table__
            self._table_cache[cache_key] = tbl
            return tbl

    # Per-table reflection; "public" → None matches how automap stores it
    reflect_schema = None if schema_ == "public" else schema_
    tbl = reflect_table(self.engine, table_name, schema=reflect_schema)
    self._table_cache[cache_key] = tbl
    return tbl

reflect_table(*args, **kwargs)

One-off reflection of a database table or view. Note: for most purposes, it will be better to use the database tables automapped at runtime using self.automap(). Then, tables can be accessed using the self.table object. However, this function can be useful for views (which are not reflected automatically), or to customize type definitions for mapped tables.

A set of column_args can be used to pass columns to override with the mapper, for instance to set up foreign and primary key constraints. https://docs.sqlalchemy.org/en/13/core/reflection.html#reflecting-views

Source code in database/macrostrat/database/core.py 
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
def reflect_table(self, *args, **kwargs):
    """
    One-off reflection of a database table or view. Note: for most purposes,
    it will be better to use the database tables automapped at runtime using
    `self.automap()`. Then, tables can be accessed using the
    `self.table` object. However, this function can be useful for views (which
    are not reflected automatically), or to customize type definitions for mapped
    tables.

    A set of `column_args` can be used to pass columns to override with the mapper, for
    instance to set up foreign and primary key constraints.
    https://docs.sqlalchemy.org/en/13/core/reflection.html#reflecting-views
    """
    warnings.warn(
        "reflect_table is deprecated and will be removed in version 4.0. Shift away from table refection, or use reflect_table from the macrostrat.database.utils module.",
        DeprecationWarning,
    )

    return reflect_table(self.engine, *args, **kwargs)

refresh_schema(*, automap=None)

Refresh the current database connection

  • closes the session and flushes
  • removes the inspector

If automap is True, will automap the database after refreshing. If automap is False, will not automap the database after refreshing. If automap is None, it will re-map the database if it was previously mapped.

Source code in database/macrostrat/database/core.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
def refresh_schema(self, *, automap=None):
    """
    Refresh the current database connection

    - closes the session and flushes
    - removes the inspector

    If automap is True, will automap the database after refreshing.
    If automap is False, will not automap the database after refreshing.
    If automap is None, it will re-map the database if it was previously mapped.
    """
    # Close the session
    self.session.flush()
    self.session.close()
    # Remove the inspector
    self.__inspector__ = None

    if automap is None:
        automap = self.mapper is not None

    if automap:
        self.automap()

run_fixtures(fixtures, params=None, **kwargs)

Run a set of fixtures on the database object.

Parameters:

Name Type Description Default
fixtures Path | list[Path]

Path to a directory of fixtures or a list of paths to fixture files.

 required
params dict

Parameters to pass to the query.

 None

Other Parameters:

Name Type Description
use_instance_params bool

If True, will use the instance_params set on the Database object. Default is True.
Source code in database/macrostrat/database/core.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
def run_fixtures(self, fixtures: Union[Path, list[Path]], params=None, **kwargs):
    """Run a set of fixtures on the database object.

    Args:
        fixtures (Path|list[Path]): Path to a directory of fixtures or a list of paths to fixture files.
        params (dict): Parameters to pass to the query.

    Keyword Args:
        use_instance_params (bool): If True, will use the instance_params set on
            the Database object. Default is True.
    """
    params = self._setup_params(params, kwargs)
    return run_fixtures(self.session, fixtures, params, **kwargs)

run_query(sql, params=None, **kwargs)

Run a single query on the database object, returning the result.

Parameters:

Name Type Description Default
sql str

SQL file or query to execute.

 required
params dict

Parameters to pass to the query.

 None

Other Parameters:

Name Type Description
use_instance_params bool

If True, will use the instance_params set on the Database object. Default is True.
Source code in database/macrostrat/database/core.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
def run_query(self, sql, params=None, **kwargs):
    """Run a single query on the database object, returning the result.

    Args:
        sql (str): SQL file or query to execute.
        params (dict): Parameters to pass to the query.

    Keyword Args:
        use_instance_params (bool): If True, will use the instance_params set on
            the Database object. Default is True.
    """
    params = self._setup_params(params, kwargs)
    return run_query(self.session, sql, params, **kwargs)

run_sql(fn, params=None, **kwargs)

Executes SQL files or query strings using the run_sql function.

Parameters:

Name Type Description Default
fn str | Path

SQL file or query string to execute.

 required
params dict

Parameters to pass to the query.

 None

Other Parameters:

Name Type Description
use_instance_params bool

If True, will use the instance_params set on the Database object. Default is True.

Returns: Iterator of results from the query.

Source code in database/macrostrat/database/core.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
def run_sql(self, fn, params=None, **kwargs):
    """Executes SQL files or query strings using the run_sql function.

    Args:
        fn (str|Path): SQL file or query string to execute.
        params (dict): Parameters to pass to the query.

    Keyword Args:
        use_instance_params (bool): If True, will use the instance_params set on
            the Database object. Default is True.

    Returns: Iterator of results from the query.
    """
    params = self._setup_params(params, kwargs)
    return run_sql(self.session, fn, params, **kwargs)

savepoint(name=None, rollback='on-error', connection=None)

A PostgreSQL-specific savepoint context manager. This is similar to the transaction context manager but uses savepoints directly for simpler operation. Notably, it supports nested savepoints, a feature that is difficult in SQLAlchemy's transaction model.

This function is not yet drop-in compatible with the transaction context manager, but that is a future goal.

Source code in database/macrostrat/database/core.py 
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
@contextmanager
def savepoint(self, name=None, rollback="on-error", connection=None):
    """A PostgreSQL-specific savepoint context manager. This is similar to the
    `transaction` context manager but uses savepoints directly for simpler operation.
    Notably, it supports nested savepoints, a feature that is difficult in SQLAlchemy's `transaction`
    model.

    This function is not yet drop-in compatible with the `transaction` context manager, but that
    is a future goal.
    """
    if name is None:
        name = f"sp_{self.savepoint_counter}"
        self.savepoint_counter += 1

    _prev_session = self.session

    if connection is None:
        connection = self.session.connection()

    params = {"name": Identifier(name)}
    run_query(connection, "SAVEPOINT {name}", params)
    should_rollback = rollback == "always"
    self.session = Session(bind=connection)
    try:
        yield name
    except Exception as e:
        should_rollback = rollback != "never"
        raise e
    finally:
        _clear_savepoint(connection, name, rollback=should_rollback)
        self.session.close()
        self.session = _prev_session

session_scope(commit=True)

Provide a transactional scope around a series of operations.

Source code in database/macrostrat/database/core.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
@contextmanager
def session_scope(self, commit=True):
    """Provide a transactional scope around a series of operations."""
    # self.__old_session = self.session
    # session = self._session_factory()
    session = self.session
    try:
        yield session
        if commit:
            session.commit()
    except Exception as err:
        session.rollback()
        raise err
    finally:
        session.close()

transaction(*, rollback='on-error', connection=None, raise_errors=True)

Create a database session that can be rolled back after use. This is similar to the session_scope method but includes more fine-grained control over transactions. The two methods may be integrated in the future.

This is based on the Sparrow's implementation: https://github.com/EarthCubeGeochron/Sparrow/blob/main/backend/conftest.py

It can be effectively used in a Pytest fixture like so: ``` @fixture(scope="class") def db(base_db): with base_db.transaction(rollback=True): yield base_db

Source code in database/macrostrat/database/core.py 
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
@contextmanager
def transaction(self, *, rollback="on-error", connection=None, raise_errors=True):
    """Create a database session that can be rolled back after use.
    This is similar to the `session_scope` method but includes
    more fine-grained control over transactions. The two methods may be integrated
    in the future.

    This is based on the Sparrow's implementation:
    https://github.com/EarthCubeGeochron/Sparrow/blob/main/backend/conftest.py

    It can be effectively used in a Pytest fixture like so:
    ```
    @fixture(scope="class")
    def db(base_db):
        with base_db.transaction(rollback=True):
            yield base_db
    """
    if connection is None:
        connection = self.engine.connect()
    transaction = connection.begin()
    session = Session(bind=connection)
    prev_session = self.session
    self.session = session

    should_rollback = rollback == "always"

    try:
        yield self
    except Exception as e:
        should_rollback = rollback != "never"
        if raise_errors:
            raise e
    finally:
        if should_rollback:
            transaction.rollback()
        else:
            transaction.commit()
        session.close()
        self.session = prev_session