Search code examples
apache-beamapache-beam-io

Is it possible to get the generated key using Apache Beam JdbcIO.Write?


If I call a stored procedure using JdbcIO.Write is it possible to capture the ID (primary key) if the stored procedure returns this data?

public JdbcIO.Write<MyObject> writeMyObject() {
      final String UPSERT_MY_OBJECT = "EXEC [MySchema].[UspertMyObject] ?,?,?";
        
      // If my stored procedure returns the generated or existing ID
      // is it possible to update the object I'm writing with the ID?
      return JdbcIO.<MyObject>write()
            .withDataSourceConfiguration(myDataSourceConfig)
            .withStatement(UPSERT_MY_OBJECT)
            .withPreparedStatementSetter((JdbcIO.PreparedStatementSetter<MyObject>) (myObject, ps) -> {
                  ps.setInt(1, myObject.getFieldOne());
                  ps.setString(2, myObject.getFieldTwo());
                  ps.setString(3, myObject.getFieldThree());
      });
}

Solution

  • I don't think it's possible but, as a workaround, you can wait for write's finish (with Wait transform, see an example there) and then read them from database.