Class RemoteRegistry

sdk/python/feast/infra/registry/remote.py:97–794

class RemoteRegistry(BaseRegistry):
    def __init__(
        self,
        registry_config: Union[RegistryConfig, RemoteRegistryConfig],
        project: str,
        repo_path: Optional[Path],
        auth_config: AuthConfig = NoAuthConfig(),
    ):
        self.auth_config = auth_config
        assert isinstance(registry_config, RemoteRegistryConfig)
        self.channel = self._create_grpc_channel(registry_config)
        weakref.finalize(self, self.channel.close)

        auth_header_interceptor = GrpcClientAuthHeaderInterceptor(auth_config)
        self.channel = grpc.intercept_channel(self.channel, auth_header_interceptor)
        self.stub = RegistryServer_pb2_grpc.RegistryServerStub(self.channel)

    def _create_grpc_channel(self, registry_config):
        assert isinstance(registry_config, RemoteRegistryConfig)
        if registry_config.cert or registry_config.is_tls:
            cafile = (
                registry_config.cert
                or os.getenv("SSL_CERT_FILE")
                or os.getenv("REQUESTS_CA_BUNDLE")
            )
            if not cafile:
                raise EnvironmentError(
                    "SSL_CERT_FILE or REQUESTS_CA_BUNDLE environment variable must be set to use secure TLS or set the cert parameter in feature_Store.yaml file under remote registry configuration."
                )
            if (registry_config.client_cert and not registry_config.client_key) or (
                not registry_config.client_cert and registry_config.client_key
            ):
                raise ValueError(
                    "Both client_cert and client_key must be provided for mTLS. "
                    "Only one was set in the remote registry configuration."
                )

            with open(cafile, "rb") as cert_file:
                trusted_certs = cert_file.read()
            private_key: Optional[bytes] = None
            certificate_chain: Optional[bytes] = None
            if registry_config.client_cert and registry_config.client_key:
                with open(registry_config.client_key, "rb") as key_file:
                    private_key = key_file.read()
                with open(registry_config.client_cert, "rb") as cert_file:
                    certificate_chain = cert_file.read()
            tls_credentials = grpc.ssl_channel_credentials(
                root_certificates=trusted_certs,
                private_key=private_key,
                certificate_chain=certificate_chain,
            )

            options = []
            if registry_config.authority:
                options.append(("grpc.default_authority", registry_config.authority))

            return grpc.secure_channel(
                registry_config.path, tls_credentials, options=options