How to handle same entity on different schemas

[Nahuel92]

Hi there!

I'm using Micronaut Data JDBC and I have the following situation:

There is an entity that exists in more than one DB schema. Since the entity's definition is the same one among schemas, I would like to have a unique JdbcRepository that allows me to pass a parameter to determine which schema to use at runtime. I'm not sure this can be done at this time.

A workaround I could use is to create a base entity with all the fields and several entities extending the base one and annotated with @MappedEntity(value = "my_common_entity", schema = "schema_name") with their corresponding repositories.

Another workaround is to use JdbcOperations and do what I need all by hand, but I'm trying to see if there is a better approach for this.

I'll appreciate any suggestions on this. Thanks in advance!

[dstepanov]

Some databases support defining the schema in the context. Maybe we can support overriding the schema in JdbcRepository

[Nahuel92]

Yes, I thought of that too, but I don't think it's possible to do so on my particular project. Let me give you more context:

The app I'm working on has (so far) 3 cron-job methods (annotated with @Scheduled) that run in different threads (and each uses a virtual thread executor to further parallelize processing). Each "pipeline", if you will, has a bunch of repositories for different entities (and from different schemas) injected.

On the other hand, I can get the schema from the Connection, but if I modify the schema and change it back via JdbcSchemaHandler, wouldn't that cause errors in other repositories executing code during the schema changing process? I know my situation is very complex, but I can't do much about it.

On a different subject, I tried setting tenant manually in the entity, but it doesn't work as I expected because Micronaut Data can't find the table at runtime. Using above's entity and repository:

final var entity = new MySharedEntity();
entity.setTenant("schema1");
myRepository.save(entity);

final var entity2 = new MySharedEntity();
entity2.setTenant("schema2");
myRepository.save(entity2);

final var entity3 = new MySharedEntity();
entity3.setTenant("schema3");
myRepository.save(entity3);

Previous code causes Micronaut to error with io.micronaut.data.exceptions.DataAccessException: Error executing PERSIST: ERROR: relation "my_shared_entity" does not exist.

Not sure if that is an allowed use of tenancy. but I wish it is, as it could save me a lot of pain and code. I'd love to hear from you any other suggestion. Anyway, I really appreciate your help and patience.

[dstepanov]

The connection is only shared after your job is done with so if you reset the schema it should be fine

[Nahuel92]

Interesting, I didn't know that. In that case, you solved my problem then (I already have a PoC, I was just waiting to hear from you on the job/threading concern).

Is there any documentation/place that states what you mentioned? I would like to document my code with it for future developers. If not, I'll link your comment above.

I'm planning to use a lock to protect method calls to MyRepository because I'll use virtual threads to process each DTO in parallel. My concern is about any race condition that may change the schema in the wrong moment and cause sneaky bugs.

[dstepanov]

Well that the behaviour of every connection pool. If you modify any kind of state there (read only etc) you need to put the previous state after you are done with it. Once again, for every connection work you can a connection assigned as after it's being put back into the pool.

What do you mean by lock protect methods? I don't think it's needed.

[Nahuel92]

I forgot to mention that my connection pool has only one connection (hard rule by our DBA). I'll review how a connection pool works, thanks.

Regarding the lock to protect methods calls to MyRepository, I meant this. Suppose my cron job calls an endpoint that returns this payload:

[
  {
    "id": "1",
    "schema": "schema1",
    "firstName": "name1"
  },
  {
    "id": "2",
    "schema": "schema3",
    "firstName": "name2"
  },
  {
    "id": "3",
    "schema": "schema2",
    "firstName": "name3"
  },
  {
    "id": "4",
    "schema": "schema1",
    "firstName": "name4"
  }
]

For each one, I'm creating a virtual thread to process it in parallel.

My Java code looks like this:

// constructor omitted for brevity
@Singleton
class MyCronJob {
    private final MyService myService;

    @Scheduled(fixedRate = "10s")
    void run() {
        myService.process();
    }
}

// constructor omitted for brevity
@Singleton
class MyService { 
    private final MyDeclarativeApi myDeclarativeApi;
    private final MyMapper myMapper;
    private final MyRepository myRepository;
    private final Connection connection;
    private final JdbcSchemaHandler jdbcSchemaHandler;

    public void process() {
        List<MyDto> dtos = myDeclarativeApi.getAll();

        List<Runnable> entitiesToProcess = dtos.stream()
                .map(e -> createRunnable())
                .toList();

        try (var executor = Executors.newVirtualThreadPerTaskExecutor()) {
            entitiesToProcess.forEach(e -> executor.submit(e));
        }
    }
    
    private Runnable createRunnable() {
        return dto -> {
            String originalSchema;
            try {
                 originalSchema = connection.getSchema();
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            jdbcSchemaHandler.useSchema(connection, Dialect.POSTGRES, dto.getSchema()); // consider schema valid at this point

            MySharedEntity mySharedEntity = myMapper.map(dto);
            MySharedEntity existingEntity = myRepository.findById(mySharedEntity.getId());
            if (existingEntity.isPresent()) {
                myRepository.update(mySharedEntity);
            } else {
                myRepository.save(mySharedEntity);
            }
            jdbcSchemaHandler.useSchema(connection, Dialect.POSTGRES, originalSchema);
        };
    }
}

I'm concerned that most of the Runnable code needs to be protected by a Lock to avoid any race condition (such as DTO 1 changing schema for DTO 3, for instance). I appreciate any input on this. Thanks a lot :)

[Nahuel92]

Hey, I wanted to give an update on this discussion.

In the end, I opted for using JdbcOperations and pass it a query created dynamically by hand (after validating parameters to avoid a SQL-injection vulnerability). This is because schemas can't be passed as a parameter to PreparedStatement. It's not ideal but it saves me from more complexity and code duplication that, otherwise, I would have ended up with.

It would be awesome if I could use Micronaut Data JDBC to map the same entity to different schemas, but as far as I know it seems not to be possible (I couldn't, at least).