Repository: simonsobs/sotodlib

Issue:
bookbinder should not add G3Timestream to G3TimesampleMap
In recent spt3g, '0.3-175-g5d56a18', bookbinder tests fail with this:
```
>           output[c] = core.G3Timestream(np.array(self.data[c][:len(output.times)]))
E           ValueError: Cannot add member (Azimuth_Corrected): not a supported vector type.

sotodlib/io/bookbinder.py:82: ValueError
```

`output` is a G3TimesampleMap.  Those were only intended to house simple vectors; G3VectorInt, G3VectorDouble, G3VectorBool ...  The fact that it accepted G3Timestream was an accident, I guess because G3Timestream used to inherit from G3VectorDouble?

In any case, it doesn't work now (I guess because of recent interesting change to G3Timestream) ... probably just change it to G3VectorDouble.

Discussion:

Patch:
diff --git a/sotodlib/io/bookbinder.py b/sotodlib/io/bookbinder.py
index 01b6f53c..e9ad6062 100644
--- a/sotodlib/io/bookbinder.py
+++ b/sotodlib/io/bookbinder.py
@@ -81,7 +81,7 @@ class _HKBundle():
         self.times = [t for t in self.times if t >= flush_time]
 
         for c in self.data.keys():
-            output[c] = core.G3Timestream(np.array(self.data[c][:len(output.times)]))
+            output[c] = core.G3VectorDouble(np.array(self.data[c][:len(output.times)]))
             if len(output.times) < len(self.data[c]):
                 self.data[c] = self.data[c][len(output.times):]
             else:
@@ -747,7 +747,6 @@ class Bookbinder(object):
             frames_list = list(frames_list)
 
         if len(frames_list) == 0: return
-        if self._verbose: print(f"=> Writing {len(frames_list)} frames")
 
         for f in frames_list:
             # If the number of samples (per channel) exceeds the max allowed, create a new output file
diff --git a/sotodlib/io/imprinter.py b/sotodlib/io/imprinter.py
index cf9cc432..531db4c9 100644
--- a/sotodlib/io/imprinter.py
+++ b/sotodlib/io/imprinter.py
@@ -19,9 +19,9 @@ from .bookbinder import Bookbinder
 ####################
 
 # book status
-UNBOUND = 0
-BOUND = 1
-FAILED = 2
+UNBOUND = "Unbound"
+BOUND = "Bound"
+FAILED = "Failed"
 
 # tel tube, stream_id, slot mapping
 VALID_OBSTYPES = ['obs', 'oper', 'smurf', 'hk', 'stray', 'misc']
@@ -72,6 +72,11 @@ class Books(Base):
     obs: list of observations within the book
     type: type of book, e.g., "oper", "hk", "obs"
     status: integer, stage of processing, 0 is unbound, 1 is bound,
+    message: error message if book failed
+    tel_tube: telescope tube
+    slots: slots in comma separated string
+    created_at: time when book was created
+    updated_at: time when book was updated
 
     """
     __tablename__ = 'books'
@@ -81,51 +86,58 @@ class Books(Base):
     max_channels = db.Column(db.Integer)
     obs = relationship("Observations", back_populates='book')  # one to many
     type = db.Column(db.String)
-    status = db.Column(db.Integer)
+    status = db.Column(db.String, default=UNBOUND)
+    message = db.Column(db.String, default="")
+    tel_tube = db.Column(db.String)
+    slots = db.Column(db.String)
+    created_at = db.Column(db.DateTime, default=dt.datetime.utcnow)
+    updated_at = db.Column(db.DateTime, default=dt.datetime.utcnow, onupdate=dt.datetime.utcnow)
 
     def __repr__(self):
         return f"<Book: {self.bid}>"
 
+
 ##############
 # main logic #
 ##############
 
+# convenient decorator to repeat a method over all data sources
+def loop_over_sources(method):
+    def wrapper(self, *args, **kwargs):
+        for source in self.sources:
+            method(self, source, *args, **kwargs)
+    return wrapper
+
 class Imprinter:
-    def __init__(self, config, g3tsmurf_config=None, echo=False):
+    def __init__(self, im_config=None, db_args={}):
         """Imprinter manages the book database.
 
         Parameters
         ----------
-        db_path: str
-            path to the book database
-        g3tsmurf_config: str
-            path to config file for the g3tsmurf database
-        echo: boolean
-            if True, print all SQL statements
-
+        im_config: str
+            path to imprinter configuration file
+        db_args: dict
+            arguments to pass to sqlalchemy.create_engine
         """
         # load config file and parameters
-        with open(config, "r") as f:
+        with open(im_config, "r") as f:
             self.config = yaml.safe_load(f)
+
         self.db_path = self.config['db_path']
-        self.slots = self.config['slots']
-        self.tel_tube = self.config['tel_tube']
+        self.sources = self.config['sources']
 
         # check whether db_path directory exists
         if not op.exists(op.dirname(self.db_path)):
             # to create the database, we first make sure that the folder exists
             os.makedirs(op.abspath(op.dirname(self.db_path)), exist_ok=True)
-
-        self.engine = db.create_engine(f"sqlite:///{self.db_path}", echo=echo,
-                                       connect_args={'check_same_thread': False})
+        self.engine = db.create_engine(f"sqlite:///{self.db_path}", **db_args)
 
         # create all tables or (do nothing if tables exist)
         Base.metadata.create_all(self.engine)
 
         self.session = None
-        self.g3tsmurf_session = None
-        self.archive = None
-        self.g3tsmurf_config = g3tsmurf_config
+        self.g3tsmurf_sessions = {}
+        self.archives = {}
 
     def get_session(self):
         """Get a new session or return the existing one
@@ -140,11 +152,13 @@ class Imprinter:
             self.session = Session()
         return self.session
 
-    def get_g3tsmurf_session(self, return_archive=False):
+    def get_g3tsmurf_session(self, source, return_archive=False):
         """Get a new g3tsmurf session or return an existing one.
 
         Parameter
         ---------
+        source: str
+            data source, e.g., "sat1", "tsat", "latrt"
         return_archive: bool
             whether to return SMURF archive object
 
@@ -154,18 +168,18 @@ class Imprinter:
         archive: if return_archive is true
 
         """
-        if self.g3tsmurf_session is None:
-            self.g3tsmurf_session, self.archive = create_g3tsmurf_session(self.g3tsmurf_config)
+        if source not in self.g3tsmurf_sessions:
+            self.g3tsmurf_sessions[source], self.archives[source] = create_g3tsmurf_session(self.sources[source]['g3tsmurf'])
         if not return_archive:
-            return self.g3tsmurf_session
-        return self.g3tsmurf_session, self.archive
+            return self.g3tsmurf_sessions[source]
+        return self.g3tsmurf_sessions[source], self.archives[source]
 
     def register_book(self, obsset, bid=None, commit=True, session=None, verbose=True):
         """Register book to database
 
         Parameters
         ----------
-        obs_list: ObsSet object
+        obsset: ObsSet object
             thin wrapper of a list of observations
         bid: str
             book id
@@ -202,12 +216,17 @@ class Imprinter:
 
         # create observation and book objects
         observations = [Observations(obs_id=obs_id) for obs_id in obsset.obs_ids]
-        book = Books(bid=bid, obs=observations, type=obsset.mode, status=UNBOUND)
-
-        # update book details
-        book.start = start_t
-        book.stop = stop_t
-        book.max_channels = max_channels
+        book = Books(
+            bid=bid,
+            obs=observations,
+            type=obsset.mode,
+            status=UNBOUND,
+            start=start_t,
+            stop=stop_t,
+            max_channels=max_channels,
+            tel_tube=obsset.tel_tube,
+            slots=','.join([s for s in obsset.slots if obsset.contains_stream(s)]),  # not worth having a extra table
+        )
 
         # add book to database
         session.add(book)
@@ -266,10 +285,7 @@ class Imprinter:
         # a dictionary of {stream_id: [file_paths]}
         filedb = self.get_files_for_book(book)
         # get readout ids
-        try:
-            readout_ids = self.get_readout_ids_for_book(book)
-        except ValueError:
-            pass
+        readout_ids = self.get_readout_ids_for_book(book)
         hkfiles = []  # fixme: add housekeeping files support
 
         start_t = int(book.start.timestamp()*1e8)
@@ -316,11 +332,26 @@ class Imprinter:
 
         Returns
         -------
-        book: book object
+        book: book object or None if not found
 
         """
         if session is None: session = self.get_session()
-        return session.query(Books).filter(Books.bid == bid).first()
+        return session.query(Books).filter(Books.bid == bid).one_or_none()
+
+    def get_books(self, session=None):
+        """Get all books from database.
+
+        Parameters
+        ----------
+        session: BookDB session
+
+        Returns
+        -------
+        books: list of book objects
+
+        """
+        if session is None: session = self.get_session()
+        return session.query(Books).all()
 
     def get_unbound_books(self, session=None):
         """Get all unbound books from database.
@@ -410,13 +441,16 @@ class Imprinter:
         if session is None: session = self.get_session()
         session.rollback()
 
-    def update_bookdb_from_g3tsmurf(self, min_ctime=None, max_ctime=None,
+    @loop_over_sources
+    def update_bookdb_from_g3tsmurf(self, source, min_ctime=None, max_ctime=None,
                                     min_overlap=30, ignore_singles=False,
                                     stream_ids=None, force_single_stream=False):
         """Update bdb with new observations from g3tsmurf db.
 
         Parameters
         ----------
+        source: str
+            g3tsmurf db source. e.g., 'sat1', 'latrt', 'tsat'
         min_ctime: float
             minimum ctime to consider (in seconds since unix epoch)
         max_ctime: float
@@ -430,8 +464,7 @@ class Imprinter:
         force_single_stream: boolean
             if True, treat observations from different streams separately
         """
-
-        session = self.get_g3tsmurf_session()
+        session = self.get_g3tsmurf_session(source)
         # set sensible ctime range is none is given
         if min_ctime is None:
             min_ctime = session.query(G3tObservations.timestamp).order_by(G3tObservations.timestamp).first()[0]
@@ -463,7 +496,10 @@ class Imprinter:
                 # distinguish different types of books
                 # operation book
                 if get_obs_type(str_obs) == 'oper':
-                    output.append(ObsSet([str_obs], mode="oper", slots=self.slots, tel_tube=self.tel_tube))
+                    output.append(ObsSet([str_obs],
+                        mode="oper",
+                        slots=self.sources[source]['slots'],
+                        tel_tube=source))
                 elif get_obs_type(str_obs) == 'obs':
                     # force each observation to be its own book
                     if force_single_stream:
@@ -494,7 +530,7 @@ class Imprinter:
                             if overlap_time.total_seconds() < 0: continue
                             if overlap_time.total_seconds() < min_overlap: continue
                             # add all of the possible overlaps
-                            output.append(ObsSet(obs_list, mode="obs", slots=self.slots, tel_tube=self.tel_tube))
+                            output.append(ObsSet(obs_list, mode="obs", slots=self.sources[source]['slots'], tel_tube=source))
 
         # remove exact duplicates in output
         output = drop_duplicates(output)
@@ -528,7 +564,7 @@ class Imprinter:
             {obs_id: [file_paths...]}
 
         """
-        session = self.get_g3tsmurf_session()
+        session = self.get_g3tsmurf_session(book.tel_tube)
         obs_ids = [o.obs_id for o in book.obs]
         obs = session.query(G3tObservations).filter(G3tObservations.obs_id.in_(obs_ids)).all()
 
@@ -551,7 +587,7 @@ class Imprinter:
             {obs_id: [readout_ids...]}}
 
         """
-        _, SMURF = self.get_g3tsmurf_session(return_archive=True)
+        _, SMURF = self.get_g3tsmurf_session(book.tel_tube, return_archive=True)
         out = {}
         # load all obs and associated files
         for obs_id, files in self.get_files_for_book(book).items():
@@ -591,6 +627,9 @@ def create_g3tsmurf_session(config):
 
     """
     # create database connection
+    with open(config, "r") as f:
+        config = yaml.safe_load(f)
+    config['db_args'] = {'connect_args': {'check_same_thread': False}}
     SMURF = G3tSmurf.from_configs(config)
     session = SMURF.Session()
     return session, SMURF
@@ -637,7 +676,7 @@ class ObsSet(list):
     def __init__(self, *args, mode="obs", slots=None, tel_tube=None, **kwargs):
         super().__init__(*args, **kwargs)
         self.mode = mode
-        self.slots = slots
+        self.slots = slots  # all possible slots (not those in the ObsSet)
         self.tel_tube = tel_tube
 
     @property
@@ -693,7 +732,7 @@ class ObsSet(list):
         if mode in ['obs', 'oper']:
             # get slot flags
             slot_flags = ''
-            for slot in self.slots[self.tel_tube]:
+            for slot in self.slots:
                 slot_flags += '1' if self.contains_stream(slot) else '0'
             bid = f"{mode}_{timestamp}_{self.tel_tube}_{slot_flags}"
         else:
@@ -722,3 +761,19 @@ def drop_duplicates(obsset_list: List[ObsSet]):
         (tuple(ids), oset) for (ids, oset) in zip(ids_list, obsset_list)).values()
     )
     return new_obsset_list
+
+def require(dct, keys):
+    """
+    Check whether a dictionary contains all the required keys.
+
+    Parameters
+    ----------
+    dct: dict
+        dictionary to check
+    keys: list
+        list of required keys
+    """
+    if dict is None: raise ValueError("Missing required dictionary")
+    for k in keys:
+        if k not in dct:
+            raise ValueError(f"Missing required key {k}")
diff --git a/sotodlib/io/load_smurf.py b/sotodlib/io/load_smurf.py
index f9fcfbf9..09542015 100644
--- a/sotodlib/io/load_smurf.py
+++ b/sotodlib/io/load_smurf.py
@@ -150,7 +150,7 @@ def _file_has_end_frames(filename):
 
 
 class G3tSmurf:
-    def __init__(self, archive_path, db_path=None, meta_path=None, echo=False):
+    def __init__(self, archive_path, db_path=None, meta_path=None, echo=False, db_args={}):
         """
         Class to manage a smurf data archive.
 
@@ -166,14 +166,15 @@ class G3tSmurf:
                 assignments). Required for full functionality.
             echo: bool, optional
                 If true, all sql statements will print to stdout.
+            db_args: dict, optional
+                Additional arguments to pass to sqlalchemy.create_engine
         """
         if db_path is None:
             db_path = os.path.join(archive_path, "frames.db")
         self.archive_path = archive_path
         self.meta_path = meta_path
         self.db_path = db_path
-        self.engine = db.create_engine(f"sqlite:///{db_path}", echo=echo,
-                                       connect_args={'check_same_thread': False})
+        self.engine = db.create_engine(f"sqlite:///{db_path}", echo=echo, **db_args)
         Session.configure(bind=self.engine)
         self.Session = sessionmaker(bind=self.engine)
         Base.metadata.create_all(self.engine)
@@ -203,7 +204,8 @@ class G3tSmurf:
             configs = yaml.safe_load( open(configs, "r"))
         return cls(os.path.join(configs["data_prefix"], "timestreams"),
                    configs["g3tsmurf_db"],
-                   meta_path=os.path.join(configs["data_prefix"],"smurf"))
+                   meta_path=os.path.join(configs["data_prefix"],"smurf"),
+                   db_args=configs.get("db_args", {}))
 
     @staticmethod
     def _make_datetime(x):
diff --git a/sotodlib/site_pipeline/make_book.py b/sotodlib/site_pipeline/make_book.py
index e9765276..d10b3190 100755
--- a/sotodlib/site_pipeline/make_book.py
+++ b/sotodlib/site_pipeline/make_book.py
@@ -1,11 +1,14 @@
+from typing import Optional
 import typer
+import traceback
+
 from ..io.imprinter import Imprinter
 
 
 def main(
     im_config: str,
-    g3tsmurf_config: str,
     output_root: str,
+    source: Optional[str],
     ):
     """Make books based on imprinter db
     
@@ -13,15 +16,18 @@ def main(
     ----------
     im_config : str
         path to imprinter configuration file
-    g3tsmurf_config : str
-        path to the g3tsmurf db configuration file
     output_root : str
         root path of the books 
+    source: str, optional
+        data source to use, e.g., sat1, latrt, tsat. If None, use all sources
     """
-    imprinter = Imprinter(im_config, g3tsmurf_config)
+    imprinter = Imprinter(im_config, db_args={'connect_args': {'check_same_thread': False}})
     # get unbound books
     unbound_books = imprinter.get_unbound_books()
     failed_books = imprinter.get_failed_books()
+    if source is not None:
+        unbound_books = [book for book in unbound_books if book.tel_tube == source]
+        failed_books = [book for book in failed_books if book.tel_tube == source]
     print(f"Found {len(unbound_books)} unbound books and {len(failed_books)} failed books")
     for book in unbound_books:
         print(f"Binding book {book.bid}")
@@ -29,6 +35,7 @@ def main(
             imprinter.bind_book(book, output_root=output_root)
         except Exception as e:
             print(f"Error binding book {book.bid}: {e}")
+            print(traceback.format_exc())
 
     print("Retrying previously failed books") 
     for book in failed_books:
@@ -37,6 +44,7 @@ def main(
             imprinter.bind_book(book, output_root=output_root)
         except Exception as e:
             print(f"Error binding book {book.bid}: {e}")
+            print(traceback.format_exc())
             # it has failed twice, ideally we want people to look at it now
             # do something here
 
diff --git a/sotodlib/site_pipeline/update_book_plan.py b/sotodlib/site_pipeline/update_book_plan.py
index 73898c41..309343ed 100644
--- a/sotodlib/site_pipeline/update_book_plan.py
+++ b/sotodlib/site_pipeline/update_book_plan.py
@@ -7,7 +7,6 @@ from ..io.imprinter import Imprinter
 
 def main(
     config: str,
-    g3tsmurf_config: str,
     min_ctime: Optional[float] = None,
     max_ctime: Optional[float] = None,
     stream_ids: Optional[str] = None,
@@ -23,8 +22,6 @@ def main(
     ----------
     config : str
         Path to config file for imprinter
-    g3tsmurf_config : str
-        Path to config file for g3tsmurf database.
     min_ctime : Optional[float], optional
         The minimum ctime to include in the book plan, by default None
     max_ctime : Optional[float], optional
@@ -42,7 +39,7 @@ def main(
     """
     if stream_ids is not None:
         stream_ids = stream_ids.split(",")
-    imprinter = Imprinter(config, g3tsmurf_config)
+    imprinter = Imprinter(config, db_args={'connect_args': {'check_same_thread': False}})
     # leaving min_ctime and max_ctime as None will go through all available data,
     # so preferreably set them to a reasonable range based on update_delay
     if not from_scratch:
@@ -55,6 +52,5 @@ def main(
                                           ignore_singles=False,
                                           stream_ids=stream_ids,
                                           force_single_stream=force_single_stream)
-
 if __name__ == "__main__":
-    typer.run(main)
+    typer.run(main)
\ No newline at end of file


Test patch:
diff --git a/tests/test_imprinter.py b/tests/test_imprinter.py
new file mode 100644
index 00000000..dfca03e1
--- /dev/null
+++ b/tests/test_imprinter.py
@@ -0,0 +1,121 @@
+import pytest
+from unittest.mock import create_autospec
+import tempfile
+import yaml
+import os
+from datetime import datetime
+
+from sotodlib.io.imprinter import *
+from sotodlib.io.g3tsmurf_db import Files
+
+@pytest.fixture(scope="session", autouse=True)
+def imprinter():
+    with tempfile.NamedTemporaryFile(mode='w') as f:
+        im_config = {
+            'db_path': '_pytest_imprinter.db',
+            'sources': {
+                'lat': {
+                    'slots': ['slot1', 'slot2']
+                },
+            }
+        }
+        yaml.dump(im_config, f)
+        f.flush()
+        im = Imprinter(f.name)
+    return im
+
+# use this to clean up db file afterwards
+@pytest.fixture(scope="session", autouse=True)
+def delete_imprinter_db(imprinter):
+    yield
+    os.remove(imprinter.db_path)
+
+@pytest.fixture
+def obsset():
+    t0 = 1674090159
+    t1 = 1674090259
+    dt_t0 = datetime.utcfromtimestamp(t0)
+    dt_t1 = datetime.utcfromtimestamp(t1)
+
+    file = create_autospec(Files)
+    file.start = dt_t0
+    file.stop = dt_t1
+    file.n_channels = 10
+    obs1 = create_autospec(G3tObservations)
+    obs1.obs_id = f"slot1_{t0}"
+    obs1.files = [file]*2
+    obs2 = create_autospec(G3tObservations)
+    obs2.obs_id = f"slot2_{t0}"
+    obs2.files = [file]*2
+    obsset = ObsSet([obs1, obs2], mode="oper", slots=["slot1", "slot2", "slot3"], tel_tube="lat")
+    return obsset
+
+def test_ObsSet(obsset):
+    assert obsset.mode == "oper"
+    assert obsset.slots == ["slot1", "slot2", "slot3"]
+    assert obsset.tel_tube == "lat"
+    assert obsset.obs_ids == ["slot1_1674090159", "slot2_1674090159"]
+    assert obsset.get_id() == "oper_1674090159_lat_110"
+    assert obsset.contains_stream("slot1") == True
+    assert obsset.contains_stream("slot2") == True
+    assert obsset.contains_stream("slot3") == False
+
+def test_register_book(imprinter, obsset):
+    book = imprinter.register_book(obsset)
+    assert book.bid == "oper_1674090159_lat_110"
+    assert book.status == UNBOUND
+    assert book.obs[0].obs_id == "slot1_1674090159"
+    assert book.obs[1].obs_id == "slot2_1674090159"
+    assert book.tel_tube == "lat"
+    assert book.type == "oper"
+    assert book.start == datetime.utcfromtimestamp(1674090159)
+    assert book.stop == datetime.utcfromtimestamp(1674090259)
+    assert book.max_channels == 10
+    assert book.message == ""
+    assert book.slots == "slot1,slot2"
+    assert imprinter.book_bound("oper_1674090159_lat_110") == False
+    with pytest.raises(BookExistsError):
+        imprinter.register_book(obsset)
+
+def test_book_exists(imprinter):
+    assert imprinter.book_exists("oper_1674090159_lat_110") == True
+    assert imprinter.book_exists("oper_1674090159_lat_111") == False
+
+def test_register_book_raises_BookExistsError(obsset, imprinter):
+    # try to register the same ObsSet again, should raise BookExistsError
+    with pytest.raises(BookExistsError):
+        assert imprinter.register_book(obsset)
+
+def test_get_session(imprinter):
+    session = imprinter.get_session()
+    assert session is not None
+
+def test_get_book(imprinter):
+    book = imprinter.get_book('oper_1674090159_lat_110')
+    assert book is not None
+    assert book.bid == 'oper_1674090159_lat_110'
+    assert book.status == UNBOUND
+
+def test_get_unbound_books(imprinter):
+    # Create a new unbound book and add it to the database
+    unbound_books = imprinter.get_unbound_books()
+    # Assert that the book we just created is in the list of unbound books
+    assert any([book.bid == 'oper_1674090159_lat_110' and book.status == UNBOUND for book in unbound_books])
+    unbound_books[0].status = BOUND
+    imprinter.session.commit()
+    unbound_books = imprinter.get_unbound_books()
+    # Assert that the book we just created is NOT in the list of unbound books
+    assert all([book.bid != 'oper_1674090159_lat_110' or book.status != UNBOUND for book in unbound_books])
+
+def test_book_bound(imprinter):
+    books = imprinter.get_books()
+    assert imprinter.book_bound(books[0].bid) == True
+    books[0].status = UNBOUND
+    imprinter.session.commit()
+    assert imprinter.book_bound(books[0].bid) == False
+
+def test_stream_timestamp():
+    obs_id = 'stream1_1674090159'
+    stream, timestamp = stream_timestamp(obs_id)
+    assert stream == 'stream1'
+    assert timestamp == '1674090159'
\ No newline at end of file