Batch improvement

[chishankar-work]

The idea is too emulate the way JDBCIO writes to SQL.

When calling write_record it does the .execute() and then .commit() in sequence, writing and committing each record to disk one at a time. The proposed change allows for more effective batching while better managing connection pooling to CSQL.

• Removed the session.commit() from the write_record which is called on every element in the _WriteToRelationalDBFn ParDo. Instead, we just call .execute() on each record, and then commit it to disk all at once.

• Instead of building the engine at the start of each bundle, move self._db = SqlAlchemyDB(self.source_config) to the .setup() method so it's only created once for the object and handles for connection pooling for the sessions that are opened and closed at the start and finish of each bundle.

• Handled the .commit() logic in the DoFn. In the start_bundle create a record_counter = 0 and records = []. This will allow us to build the commits up to sizes and ensure that they don't get too big.

• In cases where the bundles are small or divide unevenly leaving a chunk with less than 1000 records, we can directly call commit_records in the finish_bundle() to take care of the remaining elements in the bundle and flush the buffer.

• Made max_batch_size a configurable value with a default value of 1000. This can be changed easily by the user by doing something such as:

relational_db.Write(
    source_config=source_config,
    table_config=table_config
    table_config=table_config,
    max_batch_size=1500
 )

[chishankar-work]

GH-36

[mohaseeb]

Thanks @chishankar-work for this PR (and sorry for being late to review it). I had just minor comments.

P.S. you might want to check #38 if you will be running with SQLAlchemy 1.4.

[mohaseeb]

Shouldn't we remove this call now?

[mohaseeb]

The parentheses are redundant.

[mohaseeb]

self.record_counter() to self.record_counter

[mohaseeb]

Let's add this to the class documentation and describe the new behavior.

[mohaseeb]

Not sure about SQLAlchemy behavior here. If the behavior is sending each stmt directly the DB, then another potential improvement could be to create a single insert stmt for the whole batch using, e.g., something like this or this, significantly reducing the number of DB calls.

[octopop]

@mohaseeb Correct. I ended up creating my own batch insert mod based on your project and the only thing I changed was around removing assumptions that record is a dict. SQLAlchemy while generating insert statement supports both a single row (dict ) or multiple rows (list of dicts).